Skip to main content

ferrox_server/
lib.rs

1//! ferrox-server: OpenAI-compatible HTTP surface (`/health`,
2//! `/v1/models`, `/v1/chat/completions`, `/v1/completions`,
3//! `/v1/tokenize`, `/v1/detokenize`, `/v1/embeddings`) over the
4//! ferrox-models decoder, plus a whole-response cache for exact-repeat
5//! requests (see `cache` module). Loads a real GGUF checkpoint and its
6//! own real tokenizer when `-m`/`--model` or `FERROX_MODEL_PATH` is set
7//! (see `model` module). Supports sampling
8//! (temperature/top_p/top_k/repetition_penalty), stop sequences, and SSE
9//! streaming (see `generate` module).
10//!
11//! Concurrency: the loaded model
12//! (`Model`) is immutable once loaded and shared via `Arc`, not locked
13//! behind a `Mutex` -- there is no shared mutable decoder state for
14//! concurrent requests to contend on or for one panicking request to
15//! poison. The *pointer* to it is swappable (`AppState::active`, behind
16//! an `RwLock` held only long enough to clone one `Arc`), which is what
17//! `/admin/models/load` swaps; a request that has already cloned its
18//! handle finishes against the exact weights it started on, and the old
19//! model is freed when the last such request lets go.
20//! Each request builds its own KV cache (see `generate::generate`)
21//! and runs its decode loop on tokio's blocking-thread pool via
22//! `spawn_blocking`, so CPU-bound generation no longer blocks the async
23//! reactor threads -- multiple requests can decode genuinely
24//! concurrently, bounded by that pool rather than serialized through one
25//! lock. Only the small whole-response cache is still mutable shared
26//! state, and it's locked only for the brief get/put around it, never
27//! across a decode.
28//!
29//! Streaming scope: when `stream: true` and tools are inactive, each
30//! decoded chunk is pushed through a bounded `mpsc` channel from the
31//! blocking generate task into the SSE writer so time-to-first-byte
32//! overlaps with ongoing decode. Tool-call requests still buffer the
33//! full response first (detection needs the stop-bounded text).
34//! Continuous-batching streaming also buffers (batcher returns one
35//! string).
36
37mod admin;
38mod anthropic;
39mod attribution;
40mod batch_scheduler;
41mod budget;
42mod cache;
43mod cancel;
44mod chat_template;
45mod generate;
46mod health;
47mod hub;
48mod journal;
49mod json_mode;
50mod limits;
51mod mcp;
52mod model;
53mod openai_extra;
54mod resume;
55mod security;
56mod session;
57mod sse;
58mod stats;
59mod stop;
60mod tasks;
61
62use std::convert::Infallible;
63use std::fmt;
64use std::net::{IpAddr, Ipv4Addr, SocketAddr};
65use std::path::PathBuf;
66use std::str::FromStr;
67use std::sync::{Arc, Mutex, MutexGuard};
68use std::time::Duration;
69
70use axum::{
71    extract::State,
72    http::StatusCode,
73    response::sse::{Event, KeepAlive, Sse},
74    response::{IntoResponse, Response},
75    routing::{get, post},
76    Json, Router,
77};
78use clap::{Parser, ValueEnum};
79use serde::{Deserialize, Serialize};
80
81use cache::{CacheKey, ResponseCache};
82use ferrox_core::cache::KvBlockPool;
83use ferrox_models::kimi_tokenizer::KimiTokenizer;
84use ferrox_models::sampling::SamplingParams;
85use ferrox_models::tokenizer::StopTokens;
86use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
87use generate::{FinishReason, GenerationParams};
88use model::ServerTokenizer;
89
90// `PartialEq` so ferrox-cli's serve tests can assert that both front
91// ends parse a command line into the SAME arguments, rather than
92// asserting field by field and missing whichever one is added next.
93#[derive(Parser, Debug, PartialEq)]
94// No `version` here on purpose. This struct is both `ferrox-server`'s
95// own argv and the body of ferrox-cli's `serve` subcommand, and clap
96// gives an embedded subcommand its own `--version` derived from the
97// variant name: `ferrox serve --version` printed `ferrox-serve 0.10.0`,
98// naming a binary nobody ships. The front end's own `--version` is the
99// truth, and both report the same workspace version anyway.
100#[command(
101    name = "ferrox-server",
102    about = "OpenAI-compatible Ferrox inference server"
103)]
104pub struct ServerArgs {
105    /// Model path (GGUF file or Kimi checkpoint directory).
106    #[arg(short = 'm', long = "model", value_name = "FILE")]
107    model: Option<String>,
108
109    /// IP address to listen on.
110    #[arg(long, value_name = "HOST")]
111    host: Option<IpAddr>,
112
113    /// Port to listen on. `0` asks the kernel for a free one; the
114    /// actually-bound address is then announced on stdout (see
115    /// [`announce_ready`]), which is how a supervising process is meant
116    /// to learn it.
117    #[arg(long, value_name = "PORT")]
118    port: Option<u16>,
119
120    /// CPU threads (sets FERROX_CPU_THREADS and RAYON_NUM_THREADS).
121    #[arg(short = 't', long = "threads", value_name = "N")]
122    threads: Option<usize>,
123
124    /// Device used for offloading (`none` disables GPU use).
125    #[arg(
126        long = "device",
127        visible_alias = "dev",
128        value_name = "DEVICE",
129        ignore_case = true
130    )]
131    device: Option<OffloadDevice>,
132
133    /// Print available offload devices and exit.
134    #[arg(long = "list-devices", default_value_t = false)]
135    list_devices: bool,
136
137    /// GPU layers: `0`, a positive number, `auto`, or `all`.
138    ///
139    /// Partial placement is not implemented yet; any value above zero
140    /// currently enables all supported operations on the selected backend.
141    #[arg(
142        long = "n-gpu-layers",
143        visible_aliases = ["gpu-layers", "ngl"],
144        value_name = "N"
145    )]
146    n_gpu_layers: Option<GpuLayers>,
147
148    /// MCP tool-server config JSON (stub: listed in `/v1/models` metadata).
149    #[arg(long = "mcp-config", value_name = "PATH")]
150    mcp_config: Option<PathBuf>,
151
152    /// Exit when stdin reaches EOF (for a supervising parent process).
153    ///
154    /// Opt-in on purpose: a server started with stdin redirected from
155    /// `/dev/null` -- systemd, cron, `nohup` -- sees EOF immediately,
156    /// and making this the default would turn those into a server that
157    /// exits the moment it starts. A parent that *wants* the guarantee
158    /// (the desktop shell) passes the flag and keeps the pipe open.
159    #[arg(long = "exit-on-stdin-close", default_value_t = false)]
160    exit_on_stdin_close: bool,
161
162    /// Start even though another ferrox process is already holding a
163    /// model. Off by default: two models on one box do not share it,
164    /// they thrash it, and both serve slower than either would alone.
165    /// `FERROX_ALLOW_MULTIPLE_INSTANCES=1` does the same.
166    #[arg(long = "allow-multiple-instances", default_value_t = false)]
167    allow_multiple_instances: bool,
168}
169
170impl ServerArgs {
171    /// Parses `ferrox-server`'s own argv, including the llama.cpp-style
172    /// multi-character short options (`-ngl`, `-dev`) that clap cannot
173    /// express and which are rewritten to their long forms first.
174    ///
175    /// Public because ferrox-cli's `serve` subcommand hands the same
176    /// arguments to the same parser rather than reimplementing it.
177    pub fn parse_llama_style<I>(argv: I) -> Self
178    where
179        I: IntoIterator<Item = String>,
180    {
181        Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
182    }
183}
184
185#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
186enum OffloadDevice {
187    Auto,
188    None,
189    Cpu,
190    Metal,
191    Cuda,
192}
193
194#[derive(Debug, Clone, Copy, PartialEq, Eq)]
195enum GpuLayers {
196    Auto,
197    All,
198    Count(u32),
199}
200
201impl GpuLayers {
202    fn offload_enabled(self) -> bool {
203        !matches!(self, Self::Count(0))
204    }
205}
206
207impl FromStr for GpuLayers {
208    type Err = String;
209
210    fn from_str(value: &str) -> Result<Self, Self::Err> {
211        match value {
212            "auto" => Ok(Self::Auto),
213            "all" => Ok(Self::All),
214            _ => value
215                .parse::<u32>()
216                .map(Self::Count)
217                .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
218        }
219    }
220}
221
222impl fmt::Display for GpuLayers {
223    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
224        match self {
225            Self::Auto => f.write_str("auto"),
226            Self::All => f.write_str("all"),
227            Self::Count(value) => value.fmt(f),
228        }
229    }
230}
231
232/// Whether this build of the server has the Metal kernels compiled in.
233///
234/// Exists for the front ends that link this library: ferrox-cli's
235/// `metal` feature has to forward into ferrox-server
236/// (`ferrox-server?/metal`) or `ferrox serve --device metal` refuses on
237/// a Metal host while `ferrox run` on the same binary uses it. That
238/// mismatch is one Cargo manifest edit away and compiles cleanly, so
239/// ferrox-cli asserts on this constant at compile time.
240pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
241
242/// Whether this build of the server has the CUDA kernels compiled in.
243/// See [`BUILT_WITH_METAL`].
244pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
245
246fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
247    args.into_iter()
248        .map(|arg| match arg.as_str() {
249            "-ngl" => "--n-gpu-layers".into(),
250            "-dev" => "--device".into(),
251            _ => arg,
252        })
253        .collect()
254}
255
256fn print_available_devices() {
257    println!("Available devices:");
258    println!("  CPU");
259
260    let metal = ferrox_metal::MetalProfile::detect();
261    if let Some(name) = metal.device_name {
262        println!("  Metal: {name}");
263    }
264
265    let cuda = ferrox_cuda::HardwareProfile::detect();
266    if cuda.cuda_available {
267        let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
268        println!("  CUDA: {name}");
269        if cuda.cuda_device_count > 1 {
270            println!("        ({} devices detected)", cuda.cuda_device_count);
271        }
272    }
273}
274
275fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
276    if args.host.is_none() && args.port.is_none() {
277        return None;
278    }
279
280    let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
281    let host = args
282        .host
283        .or_else(|| existing.map(|addr| addr.ip()))
284        .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
285    let port = args
286        .port
287        .or_else(|| existing.map(|addr| addr.port()))
288        .unwrap_or(8383);
289    Some(SocketAddr::new(host, port).to_string())
290}
291
292fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
293    if let Some(model) = &args.model {
294        // SAFETY: called before the runtime starts worker threads.
295        unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
296    }
297
298    if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
299        // SAFETY: called before the runtime starts worker threads.
300        unsafe { std::env::set_var("FERROX_ADDR", addr) };
301    }
302
303    if let Some(threads) = args.threads {
304        if threads == 0 {
305            anyhow::bail!("--threads must be greater than zero");
306        }
307        // SAFETY: called before the runtime starts worker threads.
308        unsafe {
309            std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
310            std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
311        }
312    }
313
314    if args.device.is_none() && args.n_gpu_layers.is_none() {
315        // device overrides skipped
316    } else {
317        let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
318        let device = if layers.offload_enabled() {
319            args.device.unwrap_or(OffloadDevice::Auto)
320        } else {
321            OffloadDevice::None
322        };
323
324        match device {
325            OffloadDevice::None | OffloadDevice::Cpu => unsafe {
326                std::env::set_var("FERROX_METAL", "0");
327                std::env::set_var("FERROX_METAL_ATTN", "0");
328                std::env::set_var("FERROX_CUDA", "0");
329            },
330            OffloadDevice::Auto => unsafe {
331                std::env::set_var("FERROX_METAL", "auto");
332                std::env::set_var("FERROX_CUDA", "auto");
333                if std::env::var_os("FERROX_METAL_ATTN").is_none() {
334                    std::env::set_var("FERROX_METAL_ATTN", "1");
335                }
336            },
337            OffloadDevice::Metal => {
338                #[cfg(not(feature = "metal"))]
339                {
340                    anyhow::bail!(
341                        "Metal requested but this binary was built without --features metal"
342                    );
343                }
344                #[cfg(feature = "metal")]
345                {
346                    if !ferrox_metal::MetalProfile::detect().available {
347                        anyhow::bail!("Metal requested but no Metal device is available");
348                    }
349                    unsafe {
350                        std::env::set_var("FERROX_METAL", "1");
351                        if std::env::var_os("FERROX_METAL_ATTN").is_none() {
352                            std::env::set_var("FERROX_METAL_ATTN", "1");
353                        }
354                        std::env::set_var("FERROX_CUDA", "0");
355                    }
356                }
357            }
358            OffloadDevice::Cuda => {
359                #[cfg(not(feature = "cuda"))]
360                {
361                    anyhow::bail!(
362                        "CUDA requested but this binary was built without --features cuda"
363                    );
364                }
365                #[cfg(feature = "cuda")]
366                {
367                    if !ferrox_cuda::HardwareProfile::detect().cuda_available {
368                        anyhow::bail!("CUDA requested but no CUDA device is available");
369                    }
370                    unsafe {
371                        std::env::set_var("FERROX_CUDA", "1");
372                        std::env::set_var("FERROX_METAL", "0");
373                        std::env::set_var("FERROX_METAL_ATTN", "0");
374                    }
375                }
376            }
377        }
378    }
379
380    Ok(())
381}
382
383/// The loaded model: immutable once built, so it needs no lock at all --
384/// just cheap `Arc` sharing across concurrent request tasks. Two real
385/// checkpoint shapes exist (see `model::LoadedModel`'s doc comment for
386/// why `FERROX_MODEL_PATH` picks between them); everything that isn't
387/// engine-specific (chat template, tokenizer kind reporting, whether
388/// this is the synthetic demo) goes through the small inherent methods
389/// below rather than being matched on ad hoc at every call site.
390#[allow(clippy::large_enum_variant)] // KimiEngine/MlaEngine dwarf Arc<Decoder>; boxing would churn call sites
391pub(crate) enum Model {
392    Gguf(GgufModel),
393    Kimi(KimiModel),
394    Mla(MlaModel),
395    Gemma4(Gemma4Model),
396    Glm52(Glm52Model),
397}
398
399pub(crate) struct GgufModel {
400    decoder: Arc<Decoder>,
401    tokenizer: Arc<ServerTokenizer>,
402    stop_tokens: StopTokens,
403    bos_id: Option<usize>,
404    is_synthetic: bool,
405    chat_template: chat_template::ChatTemplate,
406}
407
408pub(crate) struct KimiModel {
409    engine: KimiEngine,
410    tokenizer: KimiTokenizer,
411    stop_tokens: StopTokens,
412    chat_template: chat_template::ChatTemplate,
413}
414
415pub(crate) struct MlaModel {
416    engine: MlaEngine,
417    tokenizer: ServerTokenizer,
418    stop_tokens: StopTokens,
419    bos_id: Option<usize>,
420    name: String,
421    chat_template: chat_template::ChatTemplate,
422}
423
424pub(crate) struct Gemma4Model {
425    engine: Gemma4Engine,
426    tokenizer: ServerTokenizer,
427    stop_tokens: StopTokens,
428    bos_id: Option<usize>,
429    name: String,
430    chat_template: chat_template::ChatTemplate,
431}
432
433pub(crate) struct Glm52Model {
434    engine: ferrox_models::Glm52Engine,
435    tokenizer: ServerTokenizer,
436    stop_tokens: StopTokens,
437    bos_id: Option<usize>,
438    name: String,
439    chat_template: chat_template::ChatTemplate,
440}
441
442impl Model {
443    pub(crate) fn chat_template(&self) -> chat_template::ChatTemplate {
444        match self {
445            Model::Gguf(m) => m.chat_template,
446            Model::Kimi(m) => m.chat_template,
447            Model::Mla(m) => m.chat_template,
448            Model::Gemma4(m) => m.chat_template,
449            Model::Glm52(m) => m.chat_template,
450        }
451    }
452
453    /// Kimi K3 / MLA / GLM-5.2 have no synthetic-weight demo path through this
454    /// server (unlike GGUF, which falls back to one when
455    /// `FERROX_MODEL_PATH` is unset) -- a loaded `Model::Kimi` /
456    /// `Model::Mla` / `Model::Glm52` is always a real checkpoint.
457    fn is_synthetic(&self) -> bool {
458        match self {
459            Model::Gguf(m) => m.is_synthetic,
460            Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
461        }
462    }
463
464    fn tokenizer_kind(&self) -> &'static str {
465        match self {
466            Model::Gguf(m) => m.tokenizer.kind(),
467            Model::Kimi(_) => "kimi-tiktoken-bpe",
468            Model::Mla(m) => m.tokenizer.kind(),
469            Model::Gemma4(m) => m.tokenizer.kind(),
470            Model::Glm52(m) => m.tokenizer.kind(),
471        }
472    }
473
474    /// Live counters of the bounded expert cache, when the model
475    /// streams routed experts (`FERROX_EXPERT_CACHE_BYTES`); `None`
476    /// for fully resident models.
477    fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
478        match self {
479            Model::Gguf(m) => m.decoder.expert_store_stats(),
480            Model::Kimi(m) => m.engine.weights.expert_store_stats(),
481            Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
482        }
483    }
484
485    pub(crate) fn name(&self) -> &str {
486        match self {
487            Model::Gguf(m) => m.decoder.config.name,
488            Model::Kimi(_) => "kimi-k3",
489            Model::Mla(m) => m.name.as_str(),
490            Model::Gemma4(m) => m.name.as_str(),
491            Model::Glm52(m) => m.name.as_str(),
492        }
493    }
494
495    pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
496        match self {
497            Model::Gguf(m) => m.tokenizer.encode(text),
498            Model::Kimi(m) => m
499                .tokenizer
500                .encode(text)
501                .into_iter()
502                .map(|id| id as usize)
503                .collect(),
504            Model::Mla(m) => m.tokenizer.encode(text),
505            Model::Gemma4(m) => m.tokenizer.encode(text),
506            Model::Glm52(m) => m.tokenizer.encode(text),
507        }
508    }
509
510    pub(crate) fn decode(&self, ids: &[usize]) -> String {
511        match self {
512            Model::Gguf(m) => m.tokenizer.decode(ids),
513            Model::Kimi(m) => {
514                let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
515                m.tokenizer.decode(&ids32)
516            }
517            Model::Mla(m) => m.tokenizer.decode(ids),
518            Model::Gemma4(m) => m.tokenizer.decode(ids),
519            Model::Glm52(m) => m.tokenizer.decode(ids),
520        }
521    }
522
523    /// Final-normed last-layer hidden states for GGUF Decoder only.
524    /// Returns `None` for engines without a hidden-state hook (e.g. Kimi/MLA/GLM).
525    pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
526        match self {
527            Model::Gguf(m) => {
528                let mut caches: Vec<_> = (0..m.decoder.layers.len())
529                    .map(|_| {
530                        ferrox_core::cache::KvCache::new(
531                            m.decoder.config.n_kv_heads,
532                            m.decoder.config.head_dim,
533                        )
534                    })
535                    .collect();
536                Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
537            }
538            Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
539        }
540    }
541
542    pub(crate) fn vocab_size(&self) -> Option<usize> {
543        match self {
544            Model::Gguf(m) => Some(m.decoder.config.vocab_size),
545            Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
546            Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
547            Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
548            Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
549        }
550    }
551}
552
553/// The model the server is serving *right now*, together with the
554/// pieces that are built from it and must be replaced with it.
555///
556/// The continuous batcher owns a worker thread holding an
557/// `Arc<Decoder>`, so it belongs to one specific model: keeping it in a
558/// separate field would let a swap leave a batcher decoding against the
559/// old weights while `Model` named the new ones. Bundling them means
560/// one `Arc` swap replaces a consistent pair.
561pub(crate) struct ActiveModel {
562    /// Admin-surface id (see `admin::discover`), or `None` for a model
563    /// that was not discovered through it -- the synthetic fallback, or
564    /// a `FERROX_MODEL_PATH` outside the scanned directory.
565    pub(crate) id: Option<String>,
566    pub(crate) model: Arc<Model>,
567    /// Opt-in continuous-batching decode worker (`FERROX_CONTINUOUS_BATCHING=1`).
568    /// Shares `forward_multi_seq` across concurrent GGUF requests. Disabled
569    /// when a KV pool or prefix cache is configured (those keep the
570    /// private-loop `generate` path).
571    pub(crate) batcher: Option<batch_scheduler::ContinuousBatcher>,
572    /// The per-request context ceiling this model was priced for, or
573    /// `None` when it could not be priced (see `crate::budget`).
574    ///
575    /// Lives on the *model* rather than on `AppState` because it is a
576    /// property of the checkpoint plus the machine: `/admin/models/load`
577    /// swapping in a different model must swap in its ceiling too,
578    /// never keep the old model's arithmetic. The same `Arc` is inside
579    /// this model's `batcher`, so the batched and private decode paths
580    /// admit on one object.
581    pub(crate) ceiling: Option<Arc<budget::ContextCeiling>>,
582}
583
584pub(crate) struct AppState {
585    /// The swappable active model.
586    ///
587    /// **A reader clones the `Arc` under the read lock and then runs;
588    /// the lock is never held across a decode.** That is the whole
589    /// design: `RwLock` guards the *pointer*, not the model, so
590    /// `/admin/models/load` swapping in a new `Arc` cannot stall a
591    /// request that is already generating, and a request that started
592    /// against the old model keeps decoding against the exact weights
593    /// it began with until it finishes -- the old `ActiveModel` (and
594    /// its batcher thread) is dropped only when the last in-flight
595    /// holder releases it, not when the swap happens. Requests that
596    /// arrive after the swap see the new model. There is deliberately
597    /// no attempt to migrate an in-flight request: half a completion
598    /// from one checkpoint and half from another is worse than either.
599    ///
600    /// `None` means nothing is loaded (after `/admin/models/unload`, or
601    /// a failed startup load): generation endpoints answer 503 rather
602    /// than pretending, and `/health` reports `unavailable`.
603    active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
604    /// Set while a load task is in flight, so a second load request is
605    /// rejected instead of racing the first. A load is not cheap and
606    /// two concurrent ones would fight for the same memory.
607    pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
608    /// Long-running jobs (download, load) -- see the `tasks` module.
609    pub(crate) tasks: Arc<tasks::TaskRegistry>,
610    /// Generations that can currently be stopped by `POST /v1/cancel`
611    /// -- see the `cancel` module for why a dropped socket alone is not
612    /// enough.
613    pub(crate) cancels: Arc<cancel::CancelRegistry>,
614    /// Recent-request ring buffer and the counters behind
615    /// `/admin/stats` -- see the `stats` module.
616    pub(crate) stats: stats::Stats,
617    /// Replay buffers for streams started with `stream_resumable`.
618    /// See the `resume` module.
619    pub(crate) streams: resume::StreamRegistry,
620    /// The directory `/admin/models` scans, when one is configured.
621    pub(crate) model_dir: Option<PathBuf>,
622    /// The only shared *mutable* state in the server. Locked only for
623    /// the brief get/put around a cache lookup, never held across a
624    /// decode -- see the module doc comment.
625    response_cache: Mutex<ResponseCache>,
626    /// `Some` when `FERROX_KV_POOL_BLOCKS`/`FERROX_KV_POOL_BLOCK_SIZE`
627    /// are set: every request's per-layer KV caches then draw from
628    /// this one shared, bounded pool instead of each growing
629    /// unboundedly. A request whose caches can't get their first block
630    /// retries for up to `FERROX_KV_POOL_QUEUE_TIMEOUT_MS` (zero by
631    /// default -- reject immediately) before being rejected with 503,
632    /// rather than being admitted regardless of how many other
633    /// requests are already decoding -- see
634    /// `ferrox_core::cache::KvBlockPool` and `generate::KvPoolConfig`.
635    /// `None` (the default) preserves the
636    /// original unbounded-per-request behavior exactly.
637    pub(crate) kv_pool: Option<generate::KvPoolConfig>,
638    /// `Some` when `FERROX_PREFIX_CACHE_ENTRIES` is set: a shared,
639    /// LRU-bounded store of previously processed prompt+KV-state
640    /// snapshots (see `ferrox_models::PrefixCache`), consulted so a
641    /// request that *extends* an earlier one -- the common multi-turn-
642    /// chat case -- can skip recomputing the shared part. Mutually
643    /// exclusive with `kv_pool` (see `generate::generate`'s doc
644    /// comment for why); `None` (the default) means every request
645    /// processes its full prompt from scratch, exactly as before this
646    /// existed.
647    pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
648    /// Server-side per-session conversation history -- see
649    /// `session::SessionStore`'s doc comment.
650    /// Always present (unlike `kv_pool`/`prefix_cache`, it's not
651    /// opt-in): a request that never sends `session_id` simply never
652    /// touches it, at negligible cost (one empty `HashMap`).
653    sessions: session::SessionStore,
654    requests_total: std::sync::atomic::AtomicU64,
655    request_errors_total: std::sync::atomic::AtomicU64,
656    started_at: std::time::Instant,
657    /// Milliseconds after `started_at` at which the last request
658    /// finished; 0 means none has. Reported by `/health` as an age, so a
659    /// client that sees a slow health poll from a GPU-saturated server
660    /// has positive evidence of liveness instead of declaring it dead.
661    last_request_ms: std::sync::atomic::AtomicU64,
662    /// Backend capability probe behind `/health` (see `health` module).
663    detection: Arc<health::Detection>,
664    /// Loaded MCP config (`--mcp-config`); tool invocation not wired yet.
665    mcp: Option<mcp::LoadedMcpConfig>,
666    /// Whether a swapped-in GGUF model should get a continuous-batching
667    /// worker, decided once at startup from the same env var and
668    /// exclusions as the initial load.
669    pub(crate) continuous_batching_enabled: bool,
670    /// The model id a load task is currently working on, so
671    /// `/admin/models` can report `loading` for it. Separate from
672    /// `load_in_progress` because that is a gate and this is a label.
673    loading_model: Mutex<Option<String>>,
674    /// The last failed load, as `(model id, message)`. Sticky until the
675    /// next successful load so `/admin/models` can say *why* an entry
676    /// is in `error` without the user retrying to find out.
677    last_load_error: Mutex<Option<(String, String)>>,
678}
679
680impl AppState {
681    /// Clones the active model's `Arc` and releases the lock before
682    /// returning. Every caller then runs against its own handle, so no
683    /// decode ever holds this lock -- see [`AppState::active`].
684    pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
685        self.active
686            .read()
687            .unwrap_or_else(|p| p.into_inner())
688            .clone()
689    }
690
691    /// [`AppState::active`] for a request that cannot proceed without a
692    /// model. 503 with a `Retry-After`-shaped explanation is the honest
693    /// answer while nothing is loaded; the alternative -- keeping a
694    /// stale model around so the endpoint never fails -- would serve
695    /// tokens from a checkpoint the operator explicitly unloaded.
696    pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
697        self.active().ok_or_else(|| {
698            (
699                StatusCode::SERVICE_UNAVAILABLE,
700                Json(serde_json::json!({"error": {
701                    "message": "no model is loaded; POST /admin/models/load with an id from \
702                                GET /admin/models",
703                    "type": "model_not_loaded"
704                }})),
705            )
706        })
707    }
708
709    /// [`AppState::active`]'s model only, for the many call sites that
710    /// do not care about the batcher.
711    pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
712        Ok(Arc::clone(&self.require_active()?.model))
713    }
714
715    /// Publishes a new active model (or `None` to unload) and returns
716    /// the previous one.
717    ///
718    /// The write lock is held only for the pointer swap. The returned
719    /// value is the caller's to drop *outside* the lock: dropping a
720    /// multi-gigabyte model can take a moment, and doing it under the
721    /// lock would block every reader for exactly as long.
722    pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
723        let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
724        std::mem::replace(&mut *guard, next)
725    }
726
727    /// Stamps "a request just finished" for `/health`'s liveness
728    /// vouching. Relaxed: this is a freshness hint, not a
729    /// synchronization point.
730    fn mark_request_finished(&self) {
731        let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
732        self.last_request_ms
733            .store(ms, std::sync::atomic::Ordering::Relaxed);
734    }
735
736    pub(crate) fn uptime(&self) -> Duration {
737        self.started_at.elapsed()
738    }
739
740    pub(crate) fn requests_total(&self) -> u64 {
741        self.requests_total
742            .load(std::sync::atomic::Ordering::Relaxed)
743    }
744
745    pub(crate) fn errors_total(&self) -> u64 {
746        self.request_errors_total
747            .load(std::sync::atomic::Ordering::Relaxed)
748    }
749
750    pub(crate) fn cache_stats(&self) -> cache::CacheStats {
751        lock_cache(&self.response_cache).stats()
752    }
753
754    /// Seconds since the last request finished, or `None` when none
755    /// has. Same derivation `/health` uses, so the two agree.
756    pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
757        let last = self
758            .last_request_ms
759            .load(std::sync::atomic::Ordering::Relaxed);
760        (last > 0)
761            .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
762            .map(|age| age.max(0.0))
763    }
764
765    pub(crate) fn loading_model_id(&self) -> Option<String> {
766        self.loading_model
767            .lock()
768            .unwrap_or_else(|p| p.into_inner())
769            .clone()
770    }
771
772    pub(crate) fn set_loading_model(&self, id: Option<String>) {
773        *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
774    }
775
776    pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
777        self.last_load_error
778            .lock()
779            .unwrap_or_else(|p| p.into_inner())
780            .clone()
781    }
782
783    pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
784        *self
785            .last_load_error
786            .lock()
787            .unwrap_or_else(|p| p.into_inner()) = error;
788    }
789
790    /// Records one finished request in the `/admin/stats` ring buffer.
791    ///
792    /// `attribution` is threaded from the request's own headers rather
793    /// than looked up here: by the time a generation task finishes, the
794    /// request parts are long gone, and reconstructing "who was that"
795    /// afterwards is exactly the guessing the monitor exists to avoid.
796    /// The model that would serve a request right now, as `/v1/models`
797    /// names it. `None` when nothing is loaded.
798    pub(crate) fn active_model_name(&self) -> Option<String> {
799        self.active().map(|a| a.model.name().to_string())
800    }
801
802    pub(crate) fn record_request(&self, record: stats::Record<'_>) {
803        self.stats.record(stats::entry(record));
804    }
805}
806
807/// Defense in depth: if a panic ever happened while this lock was held
808/// (none of the CPU-bound decode work runs under it, so this should be
809/// very unlikely), recovering the inner state on poison rather than
810/// `.unwrap()`ing keeps the cache from permanently bricking the server.
811fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
812    cache
813        .lock()
814        .unwrap_or_else(|poisoned| poisoned.into_inner())
815}
816
817#[derive(Debug, Clone, Deserialize)]
818#[serde(untagged)]
819pub(crate) enum MessageContent {
820    Text(String),
821    Parts(Vec<ContentPart>),
822}
823
824#[derive(Debug, Clone, Deserialize)]
825struct ContentPart {
826    #[serde(rename = "type")]
827    kind: String,
828    #[serde(default)]
829    text: Option<String>,
830    #[serde(default)]
831    image_url: Option<serde_json::Value>,
832}
833
834impl MessageContent {
835    fn as_text(&self) -> String {
836        match self {
837            Self::Text(s) => s.clone(),
838            Self::Parts(parts) => parts
839                .iter()
840                .filter_map(|p| p.text.as_deref())
841                .collect::<Vec<_>>()
842                .join(""),
843        }
844    }
845
846    fn has_image(&self) -> bool {
847        match self {
848            Self::Text(_) => false,
849            Self::Parts(parts) => parts
850                .iter()
851                .any(|p| p.kind == "image_url" || p.image_url.is_some()),
852        }
853    }
854}
855
856#[derive(Debug, Clone, Deserialize)]
857pub(crate) struct ChatMessage {
858    pub(crate) role: String,
859    /// `None` for an assistant message that made tool calls instead of
860    /// replying with text (the real OpenAI convention: `content` and
861    /// `tool_calls` are mutually exclusive on an assistant message).
862    #[serde(default)]
863    pub(crate) content: Option<MessageContent>,
864    /// Present on a replayed assistant message that previously made
865    /// one or more tool calls (conversation history a client sends
866    /// back on a follow-up request).
867    #[serde(default)]
868    pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
869    /// Present on a `"tool"`-role message carrying a call's result
870    /// (unused by rendering today -- `role` alone already
871    /// distinguishes it -- but accepted so real OpenAI-shaped tool-
872    /// result messages deserialize without error).
873    #[serde(default)]
874    #[allow(dead_code)]
875    pub(crate) tool_call_id: Option<String>,
876}
877
878impl ChatMessage {
879    /// The text this message actually contributes to a rendered
880    /// prompt: `content` verbatim for an ordinary message, or (for a
881    /// replayed assistant message carrying `tool_calls`) each call
882    /// re-rendered as the same `<tool_call>{...}</tool_call>` marker
883    /// text a model is asked to produce for a *new* call -- see
884    /// `chat_template`'s module doc comment for why.
885    fn rendered_content(&self) -> String {
886        let mut out = self
887            .content
888            .as_ref()
889            .map(MessageContent::as_text)
890            .unwrap_or_default();
891        if let Some(calls) = &self.tool_calls {
892            for call in calls {
893                out.push_str(&format!(
894                    "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
895                    call.function.name, call.function.arguments
896                ));
897            }
898        }
899        out
900    }
901}
902
903#[derive(Debug, Clone, Deserialize)]
904pub(crate) struct ToolCallIn {
905    #[serde(default)]
906    #[allow(dead_code)]
907    id: String,
908    #[serde(rename = "type", default)]
909    #[allow(dead_code)]
910    kind: String,
911    function: ToolCallFunctionIn,
912}
913
914#[derive(Debug, Clone, Deserialize)]
915struct ToolCallFunctionIn {
916    name: String,
917    /// A JSON-encoded string (the real OpenAI convention for
918    /// `tool_calls[].function.arguments`), not a nested object --
919    /// spliced directly into the re-rendered `<tool_call>{...}` marker
920    /// text since it's already valid JSON.
921    arguments: String,
922}
923
924/// A tool definition in the real OpenAI request shape:
925/// `{"type": "function", "function": {"name", "description", "parameters"}}`.
926#[derive(Debug, Clone, Deserialize)]
927struct ToolDef {
928    #[serde(rename = "type", default)]
929    #[allow(dead_code)]
930    kind: String,
931    function: ToolFunctionDef,
932}
933
934#[derive(Debug, Clone, Deserialize)]
935struct ToolFunctionDef {
936    name: String,
937    #[serde(default)]
938    description: Option<String>,
939    #[serde(default)]
940    parameters: Option<serde_json::Value>,
941}
942
943/// OpenAI's `tool_choice`: `"auto"`/`"none"`/`"required"`, or an object
944/// pinning one specific function. Only whether it's literally
945/// `"none"` is actually consulted (to suppress tool-calling prompting
946/// entirely) -- forcing a *specific* named call isn't implementable
947/// honestly without grammar-constrained decoding (which doesn't exist
948/// in this server), so `"required"` and a
949/// specific-function choice are both treated the same as `"auto"`:
950/// offered, not forced. A real, disclosed simplification, not silently
951/// wrong behavior.
952#[derive(Debug, Clone, Deserialize)]
953#[serde(untagged)]
954enum ToolChoice {
955    Mode(String),
956    #[allow(dead_code)]
957    Specific(serde_json::Value),
958}
959
960/// OpenAI's `stop` field accepts either a single string or an array of
961/// strings.
962#[derive(Deserialize)]
963#[serde(untagged)]
964enum StopParam {
965    One(String),
966    Many(Vec<String>),
967}
968
969#[derive(Deserialize)]
970struct ChatCompletionRequest {
971    model: String,
972    messages: Vec<ChatMessage>,
973    #[serde(default = "default_max_tokens")]
974    max_tokens: usize,
975    #[serde(default)]
976    temperature: Option<f32>,
977    #[serde(default)]
978    top_p: Option<f32>,
979    #[serde(default)]
980    top_k: Option<usize>,
981    #[serde(default)]
982    repetition_penalty: Option<f32>,
983    #[serde(default)]
984    seed: Option<u64>,
985    #[serde(default)]
986    stop: Option<StopParam>,
987    #[serde(default)]
988    stream: Option<bool>,
989    /// Ferrox extension. `true` asks the server to keep a replay buffer
990    /// for this stream so a dropped connection can be resumed from the
991    /// last `id:` seen, or drained over the JSON polling fallback.
992    ///
993    /// It also changes what a dropped socket *means*. Without it, the
994    /// connection closing cancels the generation (see the `cancel`
995    /// module). With it, the generation keeps running into the replay
996    /// buffer -- which is the entire point, and the reason this is the
997    /// caller's decision rather than the server's: a tab that navigated
998    /// away wants the CPU back, and a tab whose proxy dropped a
999    /// 90-second answer wants the answer. `POST /v1/cancel` stops a
1000    /// resumable stream either way.
1001    #[serde(default)]
1002    stream_resumable: Option<bool>,
1003    #[serde(default)]
1004    tools: Vec<ToolDef>,
1005    #[serde(default)]
1006    tool_choice: Option<ToolChoice>,
1007    /// Server-side conversation history key (see the `session`
1008    /// module): when set, `messages` is treated as
1009    /// *only the new turn(s)* to append to this session's stored
1010    /// history, not the whole conversation.
1011    #[serde(default)]
1012    session_id: Option<String>,
1013    /// OpenAI fields we explicitly reject rather than silently ignore.
1014    #[serde(default)]
1015    logprobs: Option<bool>,
1016    #[serde(default)]
1017    top_logprobs: Option<u32>,
1018    #[serde(default)]
1019    n: Option<u32>,
1020    #[serde(default)]
1021    presence_penalty: Option<f32>,
1022    #[serde(default)]
1023    frequency_penalty: Option<f32>,
1024    #[serde(default)]
1025    response_format: Option<serde_json::Value>,
1026}
1027
1028fn default_max_tokens() -> usize {
1029    16
1030}
1031
1032impl ChatCompletionRequest {
1033    fn sampling_params(&self) -> SamplingParams {
1034        SamplingParams {
1035            temperature: self.temperature.unwrap_or(0.0),
1036            top_p: self.top_p.unwrap_or(1.0),
1037            top_k: self.top_k.unwrap_or(0),
1038            repetition_penalty: self.repetition_penalty.unwrap_or(1.0),
1039            presence_penalty: self.presence_penalty.unwrap_or(0.0),
1040            frequency_penalty: self.frequency_penalty.unwrap_or(0.0),
1041        }
1042    }
1043
1044    fn stop_sequences(&self) -> Vec<String> {
1045        self.stop
1046            .as_ref()
1047            .map(|s| match s {
1048                StopParam::One(v) => vec![v.clone()],
1049                StopParam::Many(v) => v.clone(),
1050            })
1051            .unwrap_or_default()
1052    }
1053
1054    /// Real tool-calling is only offered when `tools` is non-empty AND
1055    /// the client hasn't explicitly disabled it via `tool_choice:
1056    /// "none"` -- see `ToolChoice`'s doc comment for what the other
1057    /// values do (nothing different from `"auto"`).
1058    fn tools_active(&self) -> bool {
1059        !self.tools.is_empty()
1060            && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1061    }
1062
1063    /// Reject OpenAI fields we do not implement, and `tool_choice`
1064    /// values that would silently lie (required / named function).
1065    fn validate_supported_fields(&self) -> Result<(), ApiError> {
1066        for msg in &self.messages {
1067            if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1068                return Err(unsupported_feature(
1069                    "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1070                ));
1071            }
1072        }
1073        if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1074            return Err(unsupported_feature(
1075                "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1076            ));
1077        }
1078        if self.n.is_some_and(|n| n > 1) {
1079            return Err(unsupported_feature(
1080                "n > 1 is not implemented (single completion only)",
1081            ));
1082        }
1083        if let Some(fmt) = &self.response_format {
1084            match fmt.get("type").and_then(|v| v.as_str()) {
1085                Some("json_object") => {}
1086                Some(other) => {
1087                    return Err((
1088                        StatusCode::BAD_REQUEST,
1089                        Json(serde_json::json!({
1090                            "error": {
1091                                "message": format!(
1092                                    "response_format type {other:?} is not supported (only json_object)"
1093                                )
1094                            }
1095                        })),
1096                    ));
1097                }
1098                None => {
1099                    return Err((
1100                        StatusCode::BAD_REQUEST,
1101                        Json(serde_json::json!({
1102                            "error": {
1103                                "message": "response_format must include \"type\" (only json_object is supported)"
1104                            }
1105                        })),
1106                    ));
1107                }
1108            }
1109        }
1110        match &self.tool_choice {
1111            Some(ToolChoice::Mode(m)) if m == "required" => {
1112                return Err(unsupported_feature(
1113                    "tool_choice=required needs constrained decoding (not implemented)",
1114                ));
1115            }
1116            Some(ToolChoice::Specific(_)) => {
1117                return Err(unsupported_feature(
1118                    "named tool_choice is not implemented (use auto/none)",
1119                ));
1120            }
1121            _ => {}
1122        }
1123        Ok(())
1124    }
1125
1126    /// `stop_sequences()` plus `</tool_call>` when tool-calling is
1127    /// active -- reusing the existing stop-sequence machinery
1128    /// (`generate::generate`'s `earliest_stop_match`) to end generation
1129    /// right after a tool call's JSON body, rather than adding any new
1130    /// decode-time logic. See `tool_preamble`'s doc comment for the
1131    /// full real, disclosed approach.
1132    fn effective_stop_sequences(&self) -> Vec<String> {
1133        let mut stop = self.stop_sequences();
1134        if self.tools_active() {
1135            stop.push("</tool_call>".to_string());
1136        }
1137        stop
1138    }
1139
1140    fn json_object_mode(&self) -> bool {
1141        self.response_format
1142            .as_ref()
1143            .and_then(|v| v.get("type"))
1144            .and_then(|v| v.as_str())
1145            == Some("json_object")
1146    }
1147
1148    fn generation_params(&self) -> GenerationParams {
1149        GenerationParams {
1150            max_tokens: self.max_tokens,
1151            sampling: self.sampling_params(),
1152            seed: self.resolved_seed(),
1153            stop: self.effective_stop_sequences(),
1154            // Resolved by `run_generation_emit`, the layer that holds a
1155            // tokenizer: a request body names stop strings, and only
1156            // the model can say which of them are single tokens.
1157            stop_token_ids: Vec::new(),
1158            json_object: self.json_object_mode(),
1159            // Filled in by the handler that owns the request id --
1160            // the request body cannot name its own cancel token.
1161            cancel: None,
1162        }
1163    }
1164
1165    /// Like [`Self::generation_params`], plus architecture-default stop
1166    /// strings (Gemma IT emits `<end_of_turn>` before `<eos>`).
1167    fn generation_params_for_template(
1168        &self,
1169        template: chat_template::ChatTemplate,
1170    ) -> GenerationParams {
1171        let mut params = self.generation_params();
1172        if matches!(
1173            template,
1174            chat_template::ChatTemplate::Gemma | chat_template::ChatTemplate::Gemma4
1175        ) {
1176            let stop = match template {
1177                chat_template::ChatTemplate::Gemma => "<end_of_turn>",
1178                chat_template::ChatTemplate::Gemma4 => "<turn|>",
1179                _ => unreachable!(),
1180            };
1181            if !params.stop.iter().any(|s| s == stop) {
1182                params.stop.push(stop.to_string());
1183            }
1184        }
1185        params
1186    }
1187
1188    /// A request only has a deterministic outcome -- and therefore is
1189    /// only safe to serve from or populate into the whole-response
1190    /// cache -- when it's plain greedy decode (temperature <= 0) or an
1191    /// explicit seed was given. Anything else must always regenerate:
1192    /// a "cache hit" for an unseeded sampled request would silently
1193    /// replay one random draw forever, defeating the purpose of
1194    /// sampling and surprising any client expecting fresh output per
1195    /// call.
1196    fn is_cacheable(&self) -> bool {
1197        self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1198    }
1199
1200    fn cache_key(&self, prompt: &str) -> CacheKey {
1201        CacheKey {
1202            model: self.model.clone(),
1203            prompt: prompt.to_string(),
1204            max_tokens: self.max_tokens,
1205            temperature_bits: self.temperature.unwrap_or(0.0).to_bits(),
1206            top_p_bits: self.top_p.unwrap_or(1.0).to_bits(),
1207            top_k: self.top_k.unwrap_or(0),
1208            repetition_penalty_bits: self.repetition_penalty.unwrap_or(1.0).to_bits(),
1209            presence_penalty_bits: self.presence_penalty.unwrap_or(0.0).to_bits(),
1210            frequency_penalty_bits: self.frequency_penalty.unwrap_or(0.0).to_bits(),
1211            seed: self.seed,
1212            stop: self.effective_stop_sequences(),
1213        }
1214    }
1215
1216    fn resolved_seed(&self) -> u64 {
1217        self.seed.unwrap_or_else(|| {
1218            std::time::SystemTime::now()
1219                .duration_since(std::time::UNIX_EPOCH)
1220                .map(|d| d.as_nanos() as u64)
1221                .unwrap_or(0xDEFA017)
1222        })
1223    }
1224}
1225
1226#[derive(Serialize)]
1227struct ChatCompletionChoice {
1228    index: usize,
1229    message: ChatCompletionResponseMessage,
1230    finish_reason: &'static str,
1231}
1232
1233#[derive(Serialize)]
1234struct ChatCompletionResponseMessage {
1235    role: &'static str,
1236    #[serde(skip_serializing_if = "Option::is_none")]
1237    content: Option<String>,
1238    #[serde(skip_serializing_if = "Option::is_none")]
1239    tool_calls: Option<Vec<ToolCallOut>>,
1240}
1241
1242#[derive(Serialize, Clone)]
1243struct ToolCallOut {
1244    id: String,
1245    #[serde(rename = "type")]
1246    kind: &'static str,
1247    function: ToolCallFunctionOut,
1248}
1249
1250#[derive(Serialize, Clone)]
1251struct ToolCallFunctionOut {
1252    name: String,
1253    /// A JSON-encoded string, matching the real OpenAI
1254    /// `tool_calls[].function.arguments` convention (see
1255    /// `ToolCallFunctionIn::arguments`'s doc comment).
1256    arguments: String,
1257}
1258
1259#[derive(Serialize)]
1260struct ChatCompletionResponse {
1261    id: String,
1262    /// Non-standard extension: the same value as `id`, stated under the
1263    /// name the rest of ferrox keys by (metrics, logs, `POST /cancel`
1264    /// once it exists). `id` is OpenAI's completion id and a client has
1265    /// no way to know ferrox also uses it as the request key -- saying
1266    /// so costs one field and removes the guess.
1267    request_id: String,
1268    object: &'static str,
1269    model: String,
1270    choices: Vec<ChatCompletionChoice>,
1271    /// OpenAI-convention token accounting (prompt/completion/total),
1272    /// counted from the exact ids the generation loop processed. On a
1273    /// whole-response cache hit, this is the original computation's
1274    /// accounting (same prompt, same deterministic outcome).
1275    usage: generate::Usage,
1276    /// Non-standard extension field (not part of the OpenAI API
1277    /// contract, but additive and harmless to OpenAI-compatible
1278    /// clients that ignore unknown fields): "hit" if this exact
1279    /// cacheable request was already computed, "miss" if this request
1280    /// just computed and cached a fresh completion, or "skip" if the
1281    /// request wasn't cacheable at all (sampling without a seed --
1282    /// see `ChatCompletionRequest::is_cacheable`).
1283    ferrox_cache: &'static str,
1284}
1285
1286#[derive(Serialize)]
1287struct ChatCompletionChunkDelta {
1288    #[serde(skip_serializing_if = "Option::is_none")]
1289    role: Option<&'static str>,
1290    #[serde(skip_serializing_if = "Option::is_none")]
1291    content: Option<String>,
1292    #[serde(skip_serializing_if = "Option::is_none")]
1293    tool_calls: Option<Vec<ToolCallOut>>,
1294}
1295
1296#[derive(Serialize)]
1297struct ChatCompletionChunkChoice {
1298    index: usize,
1299    delta: ChatCompletionChunkDelta,
1300    finish_reason: Option<&'static str>,
1301}
1302
1303#[derive(Serialize)]
1304struct ChatCompletionChunk {
1305    id: String,
1306    /// Present on the **first** chunk of a stream (see
1307    /// `ChatCompletionResponse::request_id`). A client learns the key
1308    /// for this generation before any content arrives, so a live view
1309    /// can correlate metrics with the stream it is rendering instead of
1310    /// guessing which in-flight request is "probably mine" -- a guess
1311    /// that mis-attributes the moment two chats run at once.
1312    #[serde(skip_serializing_if = "Option::is_none")]
1313    request_id: Option<String>,
1314    object: &'static str,
1315    model: String,
1316    choices: Vec<ChatCompletionChunkChoice>,
1317    /// Present only on the final chunk (the one carrying
1318    /// `finish_reason`), mirroring OpenAI's stream `usage` shape.
1319    #[serde(skip_serializing_if = "Option::is_none")]
1320    usage: Option<generate::Usage>,
1321}
1322
1323/// Liveness, readiness and capabilities in one cheap answer (see the
1324/// `health` module for why detection is a visible state rather than a
1325/// gap). Never behind auth or rate limiting, and never blocking: this is
1326/// the endpoint a supervisor asks when it is deciding whether to kill
1327/// the process.
1328async fn health(State(state): State<Arc<AppState>>) -> Response {
1329    let snapshot = state.detection.snapshot();
1330    let mut capabilities = snapshot.capabilities;
1331    let active = state.active();
1332
1333    // Model-derived capabilities need no probing, so they are answered
1334    // even while backend detection is still running.
1335    capabilities.push(match active.as_deref() {
1336        // `unavailable` was defined in Phase 1 but unreachable, because
1337        // the server only bound the port after a successful load. With
1338        // `/admin/models/unload` it is a state a client can actually
1339        // observe, and it must not read as "loaded but synthetic".
1340        None => ferrox_api::Capability::unavailable(
1341            ferrox_api::health::capability::REAL_WEIGHTS,
1342            ferrox_api::health::reason::MODEL_NOT_LOADED,
1343            "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1344        ),
1345        Some(active) if active.model.is_synthetic() => ferrox_api::Capability::unavailable(
1346            ferrox_api::health::capability::REAL_WEIGHTS,
1347            ferrox_api::health::reason::MODEL_NOT_LOADED,
1348            "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1349             checkpoint. Output from this model is noise.",
1350        ),
1351        Some(active) => ferrox_api::Capability::available(
1352            ferrox_api::health::capability::REAL_WEIGHTS,
1353            format!("Serving the real checkpoint '{}'.", active.model.name()),
1354        ),
1355    });
1356    capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1357        ferrox_api::Capability::available(
1358            ferrox_api::health::capability::CONTINUOUS_BATCHING,
1359            "Concurrent requests share one batched decode step.",
1360        )
1361    } else {
1362        ferrox_api::Capability::unavailable(
1363            ferrox_api::health::capability::CONTINUOUS_BATCHING,
1364            ferrox_api::health::reason::DISABLED,
1365            "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1366        )
1367    });
1368
1369    let last_request_ms = state
1370        .last_request_ms
1371        .load(std::sync::atomic::Ordering::Relaxed);
1372    let uptime = state.started_at.elapsed();
1373    // Readiness is "can this server generate", and with nothing loaded
1374    // it cannot -- so `unavailable` (503) wins over whatever the backend
1375    // probe concluded. Phase 1 defined this state but nothing could
1376    // reach it, because the process only bound the port after a
1377    // successful load; `/admin/models/unload` makes it reachable, and a
1378    // 200 `ready` here would tell a supervisor to send traffic that is
1379    // guaranteed to 503.
1380    let health_state = if active.is_none() {
1381        ferrox_api::HealthState::Unavailable
1382    } else {
1383        snapshot.state
1384    };
1385    let body = ferrox_api::HealthResponse {
1386        state: health_state,
1387        reason: match health_state {
1388            ferrox_api::HealthState::Ready => None,
1389            ferrox_api::HealthState::Unavailable => {
1390                Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
1391            }
1392            ferrox_api::HealthState::Detecting => {
1393                Some(ferrox_api::health::reason::DETECTING.to_string())
1394            }
1395        },
1396        detail: match health_state {
1397            ferrox_api::HealthState::Ready => None,
1398            ferrox_api::HealthState::Unavailable => Some(
1399                "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1400                    .to_string(),
1401            ),
1402            ferrox_api::HealthState::Detecting => {
1403                Some("Probing available compute backends.".to_string())
1404            }
1405        },
1406        model: active
1407            .as_deref()
1408            .map(|active| ferrox_api::health::ModelSummary {
1409                id: active.model.name().to_string(),
1410                tokenizer: active.model.tokenizer_kind().to_string(),
1411                synthetic_weights: active.model.is_synthetic(),
1412            }),
1413        capabilities,
1414        version: env!("CARGO_PKG_VERSION").to_string(),
1415        pid: std::process::id(),
1416        uptime_seconds: uptime.as_secs_f64(),
1417        server_time_unix_ms: std::time::SystemTime::now()
1418            .duration_since(std::time::UNIX_EPOCH)
1419            .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1420            .unwrap_or(0),
1421        last_request_age_seconds: (last_request_ms > 0)
1422            .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1423            .map(|age| age.max(0.0)),
1424    };
1425
1426    let status =
1427        StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1428    (status, Json(body)).into_response()
1429}
1430
1431async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1432    // OpenAI's `/v1/models` lists what can be *used* right now, which
1433    // after an unload is nothing. The inventory of what is on disk is a
1434    // different question and lives at `/admin/models`.
1435    let Some(active) = state.active() else {
1436        return Json(serde_json::json!({ "object": "list", "data": [] }));
1437    };
1438    let mut model_entry = serde_json::json!({
1439        "id": active.model.name(),
1440        "object": "model",
1441        "ferrox_synthetic_weights": active.model.is_synthetic(),
1442        "ferrox_tokenizer": active.model.tokenizer_kind(),
1443    });
1444    if let Some(mcp) = &state.mcp {
1445        model_entry["ferrox_mcp"] = mcp.models_metadata();
1446    }
1447    Json(serde_json::json!({
1448        "object": "list",
1449        "data": [model_entry]
1450    }))
1451}
1452
1453#[derive(Serialize)]
1454struct CombinedCacheStats {
1455    response_cache: cache::CacheStats,
1456    /// `None` when `FERROX_PREFIX_CACHE_ENTRIES` isn't set.
1457    prefix_cache: Option<ferrox_models::PrefixCacheStats>,
1458}
1459
1460async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1461    Json(CombinedCacheStats {
1462        response_cache: lock_cache(&state.response_cache).stats(),
1463        prefix_cache: state
1464            .prefix_cache
1465            .as_ref()
1466            .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1467    })
1468}
1469
1470/// Prometheus text-exposition format (`# HELP`/`# TYPE` plus
1471/// `name value` lines), so this endpoint can be scraped directly by a
1472/// Prometheus server or anything compatible with that format without
1473/// ferrox needing to speak any particular metrics client library.
1474async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1475    use std::sync::atomic::Ordering;
1476
1477    let cache_stats = lock_cache(&state.response_cache).stats();
1478    let active = state.active();
1479    let requests_total = state.requests_total.load(Ordering::Relaxed);
1480    let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1481    let uptime = state.started_at.elapsed().as_secs_f64();
1482
1483    let body = format!(
1484        "# HELP ferrox_requests_total Total chat completion requests received.\n\
1485         # TYPE ferrox_requests_total counter\n\
1486         ferrox_requests_total {requests_total}\n\
1487         # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
1488         # TYPE ferrox_request_errors_total counter\n\
1489         ferrox_request_errors_total {errors_total}\n\
1490         # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
1491         # TYPE ferrox_cache_hits_total counter\n\
1492         ferrox_cache_hits_total {}\n\
1493         # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
1494         # TYPE ferrox_cache_misses_total counter\n\
1495         ferrox_cache_misses_total {}\n\
1496         # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
1497         # TYPE ferrox_cache_entries gauge\n\
1498         ferrox_cache_entries {}\n\
1499         # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1500         # TYPE ferrox_synthetic_weights gauge\n\
1501         ferrox_synthetic_weights {}\n\
1502         # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
1503         # TYPE ferrox_uptime_seconds gauge\n\
1504         ferrox_uptime_seconds {uptime}\n",
1505        cache_stats.hits,
1506        cache_stats.misses,
1507        cache_stats.entries,
1508        // With nothing loaded there are no weights at all, synthetic or
1509        // otherwise; 0 is the reading that keeps the gauge meaning
1510        // "serving noise" rather than "serving nothing".
1511        active
1512            .as_ref()
1513            .map(|a| a.model.is_synthetic() as u8)
1514            .unwrap_or(0),
1515    );
1516
1517    // Expert-store counters, present only when the model streams
1518    // routed experts through the bounded cache
1519    // (FERROX_EXPERT_CACHE_BYTES).
1520    let body = match active
1521        .as_ref()
1522        .and_then(|a| a.model.expert_store_stats())
1523    {
1524        Some(es) => format!(
1525            "{body}\
1526             # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
1527             # TYPE ferrox_expert_cache_hits_total counter\n\
1528             ferrox_expert_cache_hits_total {}\n\
1529             # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1530             # TYPE ferrox_expert_cache_misses_total counter\n\
1531             ferrox_expert_cache_misses_total {}\n\
1532             # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
1533             # TYPE ferrox_expert_cache_evictions_total counter\n\
1534             ferrox_expert_cache_evictions_total {}\n\
1535             # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1536             # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
1537             ferrox_expert_cache_pass_throughs_total {}\n\
1538             # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1539             # TYPE ferrox_expert_cache_bytes_read_total counter\n\
1540             ferrox_expert_cache_bytes_read_total {}\n\
1541             # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1542             # TYPE ferrox_expert_cache_resident_bytes gauge\n\
1543             ferrox_expert_cache_resident_bytes {}\n",
1544            es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1545        ),
1546        None => body,
1547    };
1548
1549    // Scheduler counters, present only under continuous batching
1550    // (FERROX_CONTINUOUS_BATCHING=1). `prefill_chunks` next to
1551    // `prefill_tokens` is what makes chunked prefill observable: their
1552    // ratio is the effective chunk size the worker actually ran.
1553    let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
1554        Some(batcher) => {
1555            let sched = batcher.stats();
1556            format!(
1557                "{body}\
1558                 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
1559                 # TYPE ferrox_prefill_chunks_total counter\n\
1560                 ferrox_prefill_chunks_total {}\n\
1561                 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
1562                 # TYPE ferrox_prefill_tokens_total counter\n\
1563                 ferrox_prefill_tokens_total {}\n\
1564                 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
1565                 # TYPE ferrox_decode_steps_total counter\n\
1566                 ferrox_decode_steps_total {}\n\
1567                 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
1568                 # TYPE ferrox_scheduler_queue_depth gauge\n\
1569                 ferrox_scheduler_queue_depth {}\n\
1570                 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
1571                 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
1572                 ferrox_scheduler_queue_rejected_total {}\n\
1573                 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
1574                 # TYPE ferrox_kv_blocks_total gauge\n\
1575                 ferrox_kv_blocks_total {}\n\
1576                 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
1577                 # TYPE ferrox_kv_blocks_free gauge\n\
1578                 ferrox_kv_blocks_free {}\n\
1579                 # HELP ferrox_kv_block_size Token positions per KV block.\n\
1580                 # TYPE ferrox_kv_block_size gauge\n\
1581                 ferrox_kv_block_size {}\n\
1582                 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
1583                 # TYPE ferrox_kv_rejected_too_large_total counter\n\
1584                 ferrox_kv_rejected_too_large_total {}\n\
1585                 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
1586                 # TYPE ferrox_kv_rejected_context_length_total counter\n\
1587                 ferrox_kv_rejected_context_length_total {}\n\
1588                 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
1589                 # TYPE ferrox_scheduler_aborted_total counter\n\
1590                 ferrox_scheduler_aborted_total {}\n",
1591                sched.prefill_chunks,
1592                sched.prefill_tokens,
1593                sched.decode_steps,
1594                sched.queue_depth,
1595                sched.queue_rejected,
1596                sched.kv_blocks_total,
1597                sched.kv_blocks_free,
1598                sched.kv_block_size,
1599                sched.kv_rejected_too_large,
1600                sched.kv_rejected_context_length,
1601                sched.aborted,
1602            )
1603        }
1604        None => body,
1605    };
1606
1607    (
1608        [(
1609            axum::http::header::CONTENT_TYPE,
1610            "text/plain; version=0.0.4",
1611        )],
1612        body,
1613    )
1614        .into_response()
1615}
1616
1617pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
1618
1619pub(crate) fn unsupported_feature(message: &str) -> ApiError {
1620    (
1621        StatusCode::NOT_IMPLEMENTED,
1622        Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
1623    )
1624}
1625
1626pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
1627    let status = match e {
1628        generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
1629        // The request is bigger than the server can ever serve. That
1630        // is a property of the request, so it is the client's 400 --
1631        // answering 503 would send it into a retry loop that cannot
1632        // succeed.
1633        generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
1634        // Not the client's fault, and true of the exact same request a
1635        // moment later once capacity frees up -- 503, not 400. The
1636        // `Retry-After` header these need is stamped centrally by
1637        // `limits::retry_after`; see that function for why it lives in a
1638        // layer rather than here.
1639        generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
1640            StatusCode::SERVICE_UNAVAILABLE
1641        }
1642    };
1643    tracing::warn!("decode error: {e}");
1644    let mut body = serde_json::json!({"error": {"message": e.to_string()}});
1645    // A refusal against a ceiling names the ceiling and both sides of
1646    // the arithmetic. "Out of memory" (or a bare 400) tells a caller
1647    // that something did not fit; it does not tell them whether to
1648    // shorten the prompt or to run a bigger box, and those are the only
1649    // two actions available.
1650    if let generate::DecodeError::KvBudgetExceeded {
1651        binding,
1652        estimated_bytes,
1653        limit_bytes,
1654        positions,
1655        positions_limit,
1656        ..
1657    } = &e
1658    {
1659        body["error"]["type"] = serde_json::json!("invalid_request_error");
1660        body["error"]["code"] = serde_json::json!(binding);
1661        body["error"]["binding"] = serde_json::json!(binding);
1662        body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
1663        body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
1664        body["error"]["positions"] = serde_json::json!(positions);
1665        body["error"]["positions_limit"] = serde_json::json!(positions_limit);
1666    }
1667    // The header carries the same hint (stamped by `limits::retry_after`);
1668    // repeating it in the body is for clients that read JSON and never
1669    // look at headers, which is most of them.
1670    if let Some(secs) = e.retry_after_secs() {
1671        body["error"]["retry_after_seconds"] = serde_json::json!(secs);
1672    }
1673    (status, Json(body))
1674}
1675
1676pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
1677    tracing::error!("generation task panicked: {e}");
1678    (
1679        StatusCode::INTERNAL_SERVER_ERROR,
1680        Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
1681    )
1682}
1683
1684/// Runs generation for `params` against `model`, calling `emit` for each
1685/// decoded text chunk. Returns finish reason, usage, and the concatenated
1686/// text (for sessions / tool-call detection). Pure CPU-bound work with
1687/// no I/O and no shared lock: safe to run on `spawn_blocking`.
1688#[allow(clippy::too_many_arguments)] // one clear parameter per concern:
1689                                     // model + prompt + params, then the three optional shared
1690                                     // facilities (KV pool, prefix cache, batcher), the context
1691                                     // ceiling, and the sink. Bundling them would only move the
1692                                     // same list behind a struct at two call sites.
1693fn run_generation_emit(
1694    model: &Model,
1695    prompt: &str,
1696    params: &GenerationParams,
1697    kv_pool: Option<&generate::KvPoolConfig>,
1698    prefix_cache: Option<&Mutex<PrefixCache>>,
1699    continuous_batcher: Option<&batch_scheduler::ContinuousBatcher>,
1700    ceiling: Option<&budget::ContextCeiling>,
1701    mut emit: impl FnMut(&str),
1702) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
1703    let synthetic = model.is_synthetic();
1704    let mut chunks = Vec::new();
1705    // Layer 1 of the stop machinery is resolved exactly here, because
1706    // this is the one place that has both the request's stop strings
1707    // and the model's tokenizer. Both the batched and the private
1708    // decode paths below read the result off the params, so there is
1709    // one answer rather than two that can drift.
1710    let params = &{
1711        let mut resolved = params.clone();
1712        resolved.stop_token_ids =
1713            crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
1714        resolved
1715    };
1716    let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
1717    let (finish, usage) = match model {
1718        Model::Gguf(m) => {
1719            if let Some(batcher) = continuous_batcher {
1720                let mut tokens = m.tokenizer.encode(prompt);
1721                ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
1722                let (finish, _generated_ids, text, usage) =
1723                    batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?;
1724                if !text.is_empty() {
1725                    chunks.push(text);
1726                }
1727                (finish, usage)
1728            } else {
1729                generate::generate(
1730                    &m.decoder,
1731                    m.tokenizer.as_ref(),
1732                    &m.stop_tokens,
1733                    m.bos_id,
1734                    prompt,
1735                    params,
1736                    kv_pool,
1737                    prefix_cache,
1738                    ceiling,
1739                    |chunk| {
1740                        chunks.push(chunk.to_string());
1741                        if !synthetic {
1742                            emit(chunk);
1743                        }
1744                    },
1745                )?
1746            }
1747        }
1748        Model::Kimi(m) => generate::generate_engine(
1749            &m.engine,
1750            &m.tokenizer,
1751            &m.stop_tokens,
1752            None,
1753            prompt,
1754            params,
1755            |chunk| {
1756                chunks.push(chunk.to_string());
1757                if !synthetic {
1758                    emit(chunk);
1759                }
1760            },
1761        )?,
1762        Model::Mla(m) => generate::generate_engine(
1763            &m.engine,
1764            &m.tokenizer,
1765            &m.stop_tokens,
1766            m.bos_id,
1767            prompt,
1768            params,
1769            |chunk| {
1770                chunks.push(chunk.to_string());
1771                if !synthetic {
1772                    emit(chunk);
1773                }
1774            },
1775        )?,
1776        Model::Gemma4(m) => generate::generate_engine(
1777            &m.engine,
1778            &m.tokenizer,
1779            &m.stop_tokens,
1780            m.bos_id,
1781            prompt,
1782            params,
1783            |chunk| {
1784                chunks.push(chunk.to_string());
1785                if !synthetic {
1786                    emit(chunk);
1787                }
1788            },
1789        )?,
1790        Model::Glm52(m) => generate::generate_engine(
1791            &m.engine,
1792            &m.tokenizer,
1793            &m.stop_tokens,
1794            m.bos_id,
1795            prompt,
1796            params,
1797            |chunk| {
1798                chunks.push(chunk.to_string());
1799                if !synthetic {
1800                    emit(chunk);
1801                }
1802            },
1803        )?,
1804    };
1805
1806    let mut full = chunks.concat();
1807    if synthetic {
1808        full = format!(
1809            "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
1810             to serve a real model. Decoded ids -> {full:?}]"
1811        );
1812        emit(&full);
1813    } else if used_batcher && !full.is_empty() {
1814        emit(&full);
1815    }
1816
1817    Ok((finish, usage, full))
1818}
1819
1820/// Collecting wrapper around [`run_generation_emit`] for non-streaming
1821/// paths and tests.
1822pub(crate) fn run_generation(
1823    model: &Model,
1824    prompt: &str,
1825    params: &GenerationParams,
1826    kv_pool: Option<&generate::KvPoolConfig>,
1827    prefix_cache: Option<&Mutex<PrefixCache>>,
1828    continuous_batcher: Option<&batch_scheduler::ContinuousBatcher>,
1829    ceiling: Option<&budget::ContextCeiling>,
1830) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
1831    let (finish, usage, full) = run_generation_emit(
1832        model,
1833        prompt,
1834        params,
1835        kv_pool,
1836        prefix_cache,
1837        continuous_batcher,
1838        ceiling,
1839        |_| {},
1840    )?;
1841    Ok((
1842        if full.is_empty() {
1843            Vec::new()
1844        } else {
1845            vec![full]
1846        },
1847        finish,
1848        usage,
1849    ))
1850}
1851
1852pub(crate) fn prompt_from_messages(
1853    messages: &[ChatMessage],
1854    template: chat_template::ChatTemplate,
1855    tools: &[ToolDef],
1856) -> String {
1857    if tools.is_empty() {
1858        template.render(messages)
1859    } else {
1860        let mut with_preamble = Vec::with_capacity(messages.len() + 1);
1861        with_preamble.push(ChatMessage {
1862            role: "system".to_string(),
1863            content: Some(MessageContent::Text(tool_preamble(tools))),
1864            tool_calls: None,
1865            tool_call_id: None,
1866        });
1867        with_preamble.extend_from_slice(messages);
1868        template.render(&with_preamble)
1869    }
1870}
1871
1872/// Real, disclosed approach for tool-calling without grammar-
1873/// constrained decoding (which doesn't exist in this server):
1874/// describe each tool in plain text and ask the
1875/// model to wrap a call in a literal `<tool_call>{...}</tool_call>`
1876/// marker, then reuse the existing stop-sequence machinery (see
1877/// `ChatCompletionRequest::effective_stop_sequences`) to end
1878/// generation right after it, and parse the captured text for that
1879/// marker afterward (`extract_tool_call`). This is stop-bounded,
1880/// prompt-engineered JSON extraction, not enforced-valid-JSON output --
1881/// a real limitation, not overclaimed.
1882fn tool_preamble(tools: &[ToolDef]) -> String {
1883    let mut out = String::from(
1884        "You can call tools to help answer the user. To call a tool, respond with \
1885         EXACTLY one line in this format and nothing else:\n\
1886         <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
1887         object matching that tool's parameters>}}</tool_call>\n\n\
1888         Available tools:\n",
1889    );
1890    for t in tools {
1891        out.push_str(&format!(
1892            "- {}: {}\n  parameters (JSON schema): {}\n",
1893            t.function.name,
1894            t.function.description.as_deref().unwrap_or(""),
1895            t.function
1896                .parameters
1897                .as_ref()
1898                .map(|v| v.to_string())
1899                .unwrap_or_else(|| "{}".to_string()),
1900        ));
1901    }
1902    out
1903}
1904
1905/// Looks for a `<tool_call>{...}</tool_call>` marker (see
1906/// `tool_preamble`) and parses its JSON body into a `(name,
1907/// arguments)` pair -- `arguments` kept as a JSON-encoded *string*,
1908/// matching OpenAI's real `tool_calls[].function.arguments`
1909/// convention (a string, even though the model itself writes it as
1910/// literal JSON).
1911fn extract_tool_call(text: &str) -> Option<(String, String)> {
1912    const START: &str = "<tool_call>";
1913    const END: &str = "</tool_call>";
1914    let start = text.find(START)? + START.len();
1915    let end = start + text[start..].find(END)?;
1916    let body = text[start..end].trim();
1917    let value: serde_json::Value = serde_json::from_str(body).ok()?;
1918    let name = value.get("name")?.as_str()?.to_string();
1919    let arguments = value
1920        .get("arguments")
1921        .cloned()
1922        .unwrap_or_else(|| serde_json::json!({}));
1923    Some((name, arguments.to_string()))
1924}
1925
1926/// Builds the final response message + finish reason from raw
1927/// generated text: promotes `base_finish` to `"tool_calls"` (moving
1928/// the text into a structured `tool_calls` entry instead of `content`)
1929/// when tool-calling was active and the text actually contains a real
1930/// `<tool_call>{...}</tool_call>` marker -- a model can still just
1931/// answer in plain text despite tools being offered, which must fall
1932/// through to an ordinary text response, not an error.
1933fn build_response_message(
1934    content: String,
1935    tools_active: bool,
1936    base_finish: &'static str,
1937) -> (ChatCompletionResponseMessage, &'static str) {
1938    if tools_active {
1939        if let Some((name, arguments)) = extract_tool_call(&content) {
1940            return (
1941                ChatCompletionResponseMessage {
1942                    role: "assistant",
1943                    content: None,
1944                    tool_calls: Some(vec![ToolCallOut {
1945                        id: "call_0".to_string(),
1946                        kind: "function",
1947                        function: ToolCallFunctionOut { name, arguments },
1948                    }]),
1949                },
1950                "tool_calls",
1951            );
1952        }
1953    }
1954    (
1955        ChatCompletionResponseMessage {
1956            role: "assistant",
1957            content: Some(content),
1958            tool_calls: None,
1959        },
1960        base_finish,
1961    )
1962}
1963
1964/// Resolves the full message history a prompt should be rendered
1965/// from: `req.messages` verbatim when no session is in play, or (see
1966/// `session` module) `req.messages` appended to `session_id`'s stored
1967/// history, returning the accumulated whole.
1968fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
1969    let mut history = match &req.session_id {
1970        Some(id) => state.sessions.extend_and_get(id, &req.messages),
1971        None => req.messages.clone(),
1972    };
1973    if req.json_object_mode() {
1974        inject_json_object_system_hint(&mut history);
1975    }
1976    history
1977}
1978
1979fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
1980    const HINT: &str =
1981        "You must respond with valid JSON only (a single JSON object, no markdown fences).";
1982    if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
1983        match &mut sys.content {
1984            Some(MessageContent::Text(s)) if !s.contains("JSON") => {
1985                s.push_str("\n\n");
1986                s.push_str(HINT);
1987            }
1988            None => {
1989                sys.content = Some(MessageContent::Text(HINT.to_string()));
1990            }
1991            _ => {}
1992        }
1993    } else {
1994        messages.insert(
1995            0,
1996            ChatMessage {
1997                role: "system".to_string(),
1998                content: Some(MessageContent::Text(HINT.to_string())),
1999                tool_calls: None,
2000                tool_call_id: None,
2001            },
2002        );
2003    }
2004}
2005
2006async fn chat_completions(
2007    State(state): State<Arc<AppState>>,
2008    headers: axum::http::HeaderMap,
2009    Json(req): Json<ChatCompletionRequest>,
2010) -> Response {
2011    let attribution = attribution::Attribution::from_headers(&headers);
2012    state
2013        .requests_total
2014        .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2015    let started = std::time::Instant::now();
2016
2017    // One id per request, assigned before any work starts -- including
2018    // before validation -- so the streaming and non-streaming paths
2019    // agree and a rejected request is still nameable in the monitor.
2020    let request_id = ferrox_api::next_request_id();
2021    let stream = req.stream.unwrap_or(false);
2022
2023    if let Err(err) = req.validate_supported_fields() {
2024        state
2025            .request_errors_total
2026            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2027        let response = err.into_response();
2028        state.record_request(stats::Record {
2029            request_id: &request_id,
2030            route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2031            model: state.active_model_name(),
2032            status: response.status().as_u16(),
2033            stream,
2034            duration_ms: started.elapsed().as_millis() as u64,
2035            usage: None,
2036            attribution: &attribution,
2037        });
2038        return response;
2039    }
2040
2041    let response = if stream {
2042        chat_completions_stream(
2043            Arc::clone(&state),
2044            req,
2045            request_id.clone(),
2046            started,
2047            attribution.clone(),
2048        )
2049        .await
2050        .into_response()
2051    } else {
2052        chat_completions_full(
2053            Arc::clone(&state),
2054            req,
2055            request_id.clone(),
2056            started,
2057            attribution.clone(),
2058        )
2059        .await
2060        .into_response()
2061    };
2062
2063    if response.status().is_client_error() || response.status().is_server_error() {
2064        state
2065            .request_errors_total
2066            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2067        // Only failures are recorded here. A success has already
2068        // recorded itself from the path that knows the token counts --
2069        // and, for a stream, that has not even happened yet.
2070        state.record_request(stats::Record {
2071            request_id: &request_id,
2072            route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2073            // `None` here is the 503 case and says so: nothing was
2074            // loaded, so nothing served it.
2075            model: state.active_model_name(),
2076            status: response.status().as_u16(),
2077            stream,
2078            duration_ms: started.elapsed().as_millis() as u64,
2079            usage: None,
2080            attribution: &attribution,
2081        });
2082    }
2083    state.mark_request_finished();
2084
2085    response
2086}
2087
2088async fn chat_completions_full(
2089    state: Arc<AppState>,
2090    req: ChatCompletionRequest,
2091    request_id: String,
2092    started: std::time::Instant,
2093    attribution: attribution::Attribution,
2094) -> Result<Json<ChatCompletionResponse>, ApiError> {
2095    let tools_active = req.tools_active();
2096    // Cloned once, up front: this request decodes against exactly this
2097    // model even if `/admin/models/load` swaps a different one in
2098    // halfway through (see `AppState::active`).
2099    let active = state.require_active()?;
2100    let history = resolve_history(&state, &req);
2101    let prompt = prompt_from_messages(&history, active.model.chat_template(), &req.tools);
2102    let key = req.is_cacheable().then(|| req.cache_key(&prompt));
2103
2104    let (completion, cache_status) = if let Some(cached) = key
2105        .as_ref()
2106        .and_then(|key| lock_cache(&state.response_cache).get(key))
2107    {
2108        tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2109        (cached, "hit")
2110    } else {
2111        let model = Arc::clone(&active.model);
2112        let kv_pool = state.kv_pool.clone();
2113        let prefix_cache = state.prefix_cache.clone();
2114        let batcher = active.batcher.clone();
2115        let ceiling = active.ceiling.clone();
2116        let params = req.generation_params_for_template(active.model.chat_template());
2117        let prompt_for_task = prompt.clone();
2118        let (chunks, finish, usage) = tokio::task::spawn_blocking(move || {
2119            run_generation(
2120                &model,
2121                &prompt_for_task,
2122                &params,
2123                kv_pool.as_ref(),
2124                prefix_cache.as_deref(),
2125                batcher.as_ref(),
2126                ceiling.as_deref(),
2127            )
2128        })
2129        .await
2130        .map_err(join_error_response)?
2131        .map_err(decode_error_response)?;
2132
2133        let completion = cache::CachedCompletion {
2134            content: chunks.concat(),
2135            finish,
2136            usage,
2137        };
2138        let cache_status = if let Some(key) = key {
2139            tracing::debug!("cache miss for key {}", key.digest());
2140            lock_cache(&state.response_cache).put(key, completion.clone());
2141            "miss"
2142        } else {
2143            "skip"
2144        };
2145        (completion, cache_status)
2146    };
2147    let content = completion.content;
2148
2149    if req.json_object_mode() {
2150        json_mode::validate_json_object_output(&content)?;
2151    }
2152
2153    // Stored regardless of cache hit/miss, so a session's history is
2154    // always consistent with what a client would see, whether or not
2155    // this exact prompt happened to be served from cache.
2156    if let Some(id) = &req.session_id {
2157        state.sessions.store_reply(
2158            id,
2159            ChatMessage {
2160                role: "assistant".to_string(),
2161                content: Some(MessageContent::Text(content.clone())),
2162                tool_calls: None,
2163                tool_call_id: None,
2164            },
2165        );
2166    }
2167
2168    let (message, finish_reason) =
2169        build_response_message(content, tools_active, completion.finish.as_str());
2170
2171    state.record_request(stats::Record {
2172        request_id: &request_id,
2173        route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2174        // The handle this request decoded against, not `req.model`: a
2175        // swap mid-flight does not change which weights answered.
2176        model: Some(active.model.name().to_string()),
2177        status: 200,
2178        stream: false,
2179        duration_ms: started.elapsed().as_millis() as u64,
2180        usage: Some(&completion.usage),
2181        attribution: &attribution,
2182    });
2183
2184    Ok(Json(ChatCompletionResponse {
2185        id: request_id.clone(),
2186        request_id,
2187        object: "chat.completion",
2188        model: req.model,
2189        choices: vec![ChatCompletionChoice {
2190            index: 0,
2191            message,
2192            finish_reason,
2193        }],
2194        usage: completion.usage,
2195        ferrox_cache: cache_status,
2196    }))
2197}
2198
2199async fn chat_completions_stream(
2200    state: Arc<AppState>,
2201    req: ChatCompletionRequest,
2202    request_id: String,
2203    started: std::time::Instant,
2204    attribution: attribution::Attribution,
2205) -> Result<Response, ApiError> {
2206    // Streaming requests are never served from or written to the response cache.
2207    let tools_active = req.tools_active();
2208    // See `chat_completions_full`: the handle is taken once and the
2209    // whole stream runs against it, so a mid-stream model swap cannot
2210    // splice two checkpoints into one completion.
2211    let active = state.require_active()?;
2212    let history = resolve_history(&state, &req);
2213    let prompt = prompt_from_messages(&history, active.model.chat_template(), &req.tools);
2214    let model_name = req.model.clone();
2215    let session_id = req.session_id.clone();
2216    let sessions = state.sessions.clone();
2217
2218    let model = Arc::clone(&active.model);
2219    let kv_pool = state.kv_pool.clone();
2220    let prefix_cache = state.prefix_cache.clone();
2221    let batcher = active.batcher.clone();
2222    let ceiling = active.ceiling.clone();
2223    let mut params = req.generation_params_for_template(active.model.chat_template());
2224    let stats_state = Arc::clone(&state);
2225    // Read now, off the handle this stream will decode against. Read
2226    // later it would name whatever a swap had made current by then.
2227    let served_model = active.model.name().to_string();
2228
2229    // Tier two of cancellation: the id is already on the wire, so the
2230    // client can name it. The guard rides with the generation task and
2231    // deregisters however that task ends, panic included -- see the
2232    // `cancel` module.
2233    let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2234    params.cancel = Some(cancel_token.clone());
2235
2236    // Tool-call detection needs the full stop-bounded text; continuous
2237    // batching returns one string. Both stay buffered. Otherwise each
2238    // decoded chunk is pushed on a channel for overlapped SSE delivery.
2239    let overlap = !tools_active && batcher.is_none();
2240
2241    // Opt-in replay. Registering a buffer is also what decides whether a
2242    // dropped socket cancels this generation -- see `resume`'s module
2243    // doc for why that is the caller's call and not the server's.
2244    let slot = req
2245        .stream_resumable
2246        .unwrap_or(false)
2247        .then(|| state.streams.register(&request_id));
2248    let emitter = resume::Emitter::new(slot);
2249
2250    let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2251
2252    tokio::task::spawn_blocking(move || {
2253        // Held for the whole generation; dropping it is what takes the
2254        // id back out of the cancel registry.
2255        let _cancel_guard = cancel_guard;
2256        let tx_chunks = tx.clone();
2257        // The orphan deadline (see `crate::sse`): a client that is
2258        // neither reading nor disconnected must not park this blocking
2259        // thread -- and the model handle and cancel guard it holds --
2260        // for the life of the process.
2261        let orphan_timeout = sse::orphan_timeout_from_env();
2262        let mut first = true;
2263        let head_request_id = request_id.clone();
2264        let result = run_generation_emit(
2265            &model,
2266            &prompt,
2267            &params,
2268            kv_pool.as_ref(),
2269            prefix_cache.as_deref(),
2270            batcher.as_ref(),
2271            ceiling.as_deref(),
2272            |chunk| {
2273                if !overlap || chunk.is_empty() {
2274                    return;
2275                }
2276                let role = if first { Some("assistant") } else { None };
2277                let request_id = first.then(|| head_request_id.clone());
2278                first = false;
2279                let payload = ChatCompletionChunk {
2280                    id: head_request_id.clone(),
2281                    request_id,
2282                    object: "chat.completion.chunk",
2283                    model: model_name.clone(),
2284                    choices: vec![ChatCompletionChunkChoice {
2285                        index: 0,
2286                        delta: ChatCompletionChunkDelta {
2287                            role,
2288                            content: Some(chunk.to_string()),
2289                            tool_calls: None,
2290                        },
2291                        finish_reason: None,
2292                    }],
2293                    usage: None,
2294                };
2295                // Tier one of cancellation. A failed send means the SSE
2296                // receiver is gone -- the browser tab closed, the
2297                // client aborted, the connection dropped -- and until
2298                // this was checked the return value was discarded and
2299                // the decode loop happily generated the remaining
2300                // hundreds of tokens into nothing. Flipping the same
2301                // flag `/v1/cancel` sets means there is one stop path,
2302                // not two.
2303                if let Err(why) =
2304                    sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
2305                {
2306                    if why == sse::SendFailure::Orphaned {
2307                        tracing::warn!(
2308                            "SSE stream {head_request_id} accepted nothing for the orphan \
2309                             deadline; treating it as abandoned"
2310                        );
2311                    }
2312                    // Two features met here and only one of them may
2313                    // win. The orphan deadline exists to stop work
2314                    // nobody is reading. A resumable stream is exactly
2315                    // the case where a gone receiver must NOT stop the
2316                    // work: the client said it may come back, the
2317                    // buffer is still being filled for it, and
2318                    // cancelling would make every reconnect resume into
2319                    // a truncated answer. So the deadline still detects
2320                    // and logs, and only a non-resumable stream is
2321                    // cancelled by it. `POST /v1/cancel` is the stop
2322                    // path for the resumable ones.
2323                    if !emitter.is_resumable() {
2324                        cancel_token.cancel();
2325                    }
2326                }
2327            },
2328        );
2329
2330        // `first` is still true when nothing was streamed from the emit
2331        // closure (the buffered tool-call/batching path, or an empty
2332        // generation), so the id has not gone out yet. `take()` on the
2333        // way into each payload below guarantees it is announced
2334        // exactly once, on whichever chunk really is first.
2335        let mut pending_request_id = first.then(|| request_id.clone());
2336
2337        match result {
2338            Ok((finish, usage, full_text)) => {
2339                if let Some(id) = &session_id {
2340                    sessions.store_reply(
2341                        id,
2342                        ChatMessage {
2343                            role: "assistant".to_string(),
2344                            content: Some(MessageContent::Text(full_text.clone())),
2345                            tool_calls: None,
2346                            tool_call_id: None,
2347                        },
2348                    );
2349                }
2350                let tool_call = if tools_active {
2351                    extract_tool_call(&full_text)
2352                } else {
2353                    None
2354                };
2355                if !overlap {
2356                    if let Some((name, arguments)) = &tool_call {
2357                        let payload = ChatCompletionChunk {
2358                            id: request_id.clone(),
2359                            request_id: pending_request_id.take(),
2360                            object: "chat.completion.chunk",
2361                            model: model_name.clone(),
2362                            choices: vec![ChatCompletionChunkChoice {
2363                                index: 0,
2364                                delta: ChatCompletionChunkDelta {
2365                                    role: Some("assistant"),
2366                                    content: None,
2367                                    tool_calls: Some(vec![ToolCallOut {
2368                                        id: "call_0".to_string(),
2369                                        kind: "function",
2370                                        function: ToolCallFunctionOut {
2371                                            name: name.clone(),
2372                                            arguments: arguments.clone(),
2373                                        },
2374                                    }]),
2375                                },
2376                                finish_reason: None,
2377                            }],
2378                            usage: None,
2379                        };
2380                        let _ =
2381                            sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
2382                    } else if !full_text.is_empty() {
2383                        let payload = ChatCompletionChunk {
2384                            id: request_id.clone(),
2385                            request_id: pending_request_id.take(),
2386                            object: "chat.completion.chunk",
2387                            model: model_name.clone(),
2388                            choices: vec![ChatCompletionChunkChoice {
2389                                index: 0,
2390                                delta: ChatCompletionChunkDelta {
2391                                    role: Some("assistant"),
2392                                    content: Some(full_text),
2393                                    tool_calls: None,
2394                                },
2395                                finish_reason: None,
2396                            }],
2397                            usage: None,
2398                        };
2399                        let _ =
2400                            sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
2401                    }
2402                }
2403                let final_finish_reason = if tool_call.is_some() {
2404                    "tool_calls"
2405                } else {
2406                    finish.as_str()
2407                };
2408                let final_payload = ChatCompletionChunk {
2409                    id: request_id.clone(),
2410                    request_id: pending_request_id.take(),
2411                    object: "chat.completion.chunk",
2412                    model: model_name,
2413                    choices: vec![ChatCompletionChunkChoice {
2414                        index: 0,
2415                        delta: ChatCompletionChunkDelta {
2416                            role: None,
2417                            content: None,
2418                            tool_calls: None,
2419                        },
2420                        finish_reason: Some(final_finish_reason),
2421                    }],
2422                    usage: Some(usage.clone()),
2423                };
2424                let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
2425                let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
2426                // Recorded here rather than where the handler returned:
2427                // the handler returns as soon as the SSE headers go out,
2428                // which is before a single token exists, so timing it
2429                // there would report every stream as instant.
2430                stats_state.record_request(stats::Record {
2431                    request_id: &request_id,
2432                    route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2433                    model: Some(served_model.clone()),
2434                    status: 200,
2435                    stream: true,
2436                    duration_ms: started.elapsed().as_millis() as u64,
2437                    usage: Some(&usage),
2438                    attribution: &attribution,
2439                });
2440            }
2441            Err(e) => {
2442                tracing::warn!("decode error on streamed request {request_id}: {e}");
2443                // The socket carried 200 -- SSE headers precede the
2444                // first token -- but the request produced no completion.
2445                // The monitor records outcomes, and a 200 row with zero
2446                // tokens would read as a successful empty answer, so the
2447                // failure is stated as 500 here and only here.
2448                stats_state.record_request(stats::Record {
2449                    request_id: &request_id,
2450                    route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2451                    model: Some(served_model.clone()),
2452                    status: 500,
2453                    stream: true,
2454                    duration_ms: started.elapsed().as_millis() as u64,
2455                    usage: None,
2456                    attribution: &attribution,
2457                });
2458                let payload = ChatCompletionChunk {
2459                    id: request_id.clone(),
2460                    request_id: pending_request_id.take(),
2461                    object: "chat.completion.chunk",
2462                    model: model_name,
2463                    choices: vec![ChatCompletionChunkChoice {
2464                        index: 0,
2465                        delta: ChatCompletionChunkDelta {
2466                            role: Some("assistant"),
2467                            content: Some(format!("[error: {e}]")),
2468                            tool_calls: None,
2469                        },
2470                        finish_reason: Some("stop"),
2471                    }],
2472                    usage: None,
2473                };
2474                let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
2475                let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
2476            }
2477        }
2478        // The buffer is closed by dropping `emitter` here -- including
2479        // on a panic, which is the case an explicit call would miss.
2480        // See `resume::Emitter`'s `Drop`.
2481        drop(emitter);
2482    });
2483
2484    let stream =
2485        futures_util::stream::unfold(
2486            rx,
2487            |mut rx| async move { rx.recv().await.map(|ev| (ev, rx)) },
2488        );
2489    // `X-Accel-Buffering: no` is the one header that actually reaches
2490    // the problem the plan names: nginx (and the proxies that copied
2491    // its convention) buffer `text/event-stream` by default, which
2492    // turns a token-by-token stream into one silent wait followed by
2493    // the whole answer at once -- indistinguishable, from the browser,
2494    // from a hung backend. axum already sets `Cache-Control: no-cache`
2495    // on an `Sse` response, so that half is covered.
2496    //
2497    // The keep-alive comment every 15s is the other half: it gives an
2498    // idle-but-healthy stream something to send, so a client's stall
2499    // timeout measures the *connection* rather than the model's
2500    // time-to-first-token on a long prompt.
2501    Ok((
2502        [(
2503            axum::http::HeaderName::from_static("x-accel-buffering"),
2504            axum::http::HeaderValue::from_static("no"),
2505        )],
2506        Sse::new(stream).keep_alive(KeepAlive::default()),
2507    )
2508        .into_response())
2509}
2510
2511/// The axum pattern for one of the published stream templates.
2512///
2513/// `ferrox_api::routes` writes placeholders in the OpenAPI style
2514/// because it is imported by clients that have never heard of this
2515/// server's router; axum 0.7 wants `:name`. Converting here keeps one
2516/// published spelling and one router spelling, and the test below fails
2517/// if they ever stop describing the same path.
2518fn resume_route(template: &str) -> String {
2519    template.replace("{request_id}", ":request_id")
2520}
2521
2522/// `POST /v1/cancel` -- the explicit half of two-tier cancellation.
2523///
2524/// Answers `200` when a live generation was signalled and `404` when
2525/// the id names nothing that is running. That difference is the whole
2526/// point of the endpoint returning a body at all: "already finished"
2527/// and "stopped it" are both fine outcomes, but only one of them saved
2528/// any work, and a UI told `ok: true` for both will claim it stopped
2529/// something it did not.
2530async fn cancel_generation(
2531    State(state): State<Arc<AppState>>,
2532    Json(req): Json<ferrox_api::CancelGenerationRequest>,
2533) -> Response {
2534    let cancelled = state.cancels.cancel(&req.request_id);
2535    let status = if cancelled {
2536        StatusCode::OK
2537    } else {
2538        StatusCode::NOT_FOUND
2539    };
2540    let detail = if cancelled {
2541        "the generation was asked to stop; it ends at its next token".to_string()
2542    } else {
2543        "no generation with that request_id is running -- it has already \
2544         finished, was never issued, or was served by a path that does \
2545         not register for cancellation"
2546            .to_string()
2547    };
2548    (
2549        status,
2550        Json(ferrox_api::CancelGenerationResponse {
2551            request_id: req.request_id,
2552            cancelled,
2553            detail,
2554        }),
2555    )
2556        .into_response()
2557}
2558
2559/// What a freshly loaded checkpoint becomes when it is published as the
2560/// active model: the model itself, its optional continuous-batching
2561/// worker, and the context ceiling both decode paths admit on.
2562type Activated = (
2563    Model,
2564    Option<batch_scheduler::ContinuousBatcher>,
2565    Option<Arc<budget::ContextCeiling>>,
2566);
2567
2568/// The scheduler config for a freshly loaded GGUF, with the ceilings an
2569/// operator did not configure *derived* from the checkpoint instead of
2570/// left absent.
2571///
2572/// This is the server half of `mem-preload-kv-budget`: `ferrox run`
2573/// already priced weights + `n_ctx * per_token_kv` + headroom against
2574/// the device budget before loading, while `ferrox-server` admitted on
2575/// whatever `FERROX_CB_*` happened to be set and otherwise on nothing.
2576///
2577/// Precedence is one-directional and deliberate: an explicit
2578/// `FERROX_CB_MAX_CONTEXT` / `FERROX_CB_KV_BLOCKS` is never overridden,
2579/// because an operator who names a number has information this
2580/// arithmetic does not. Derivation only ever fills an *absent* ceiling,
2581/// where the alternative is no ceiling at all.
2582///
2583/// `path` is `None` for the synthetic-weights fallback, which has no
2584/// checkpoint on disk to price.
2585fn price_batcher_config(path: Option<&str>) -> batch_scheduler::BatcherConfig {
2586    let mut batcher = batch_scheduler::BatcherConfig::from_env();
2587    if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
2588        // Nothing left to derive, and pricing the checkpoint would only
2589        // print arithmetic that decides nothing.
2590        return batcher;
2591    }
2592    let Some(path) = path else {
2593        return batcher;
2594    };
2595    // `ferrox_core::cache::KvCache` is `Vec<f32>` on both decode paths,
2596    // so f32 is the width really kept, even under Metal attention where
2597    // the *device* also holds an f16 copy. Budgeting the host store is
2598    // the conservative reading: it over-charges KV and therefore
2599    // under-states the context that fits.
2600    let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
2601    let Some((priced, gguf_ctx, source)) = priced else {
2602        return batcher;
2603    };
2604    let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
2605        // See `budget`'s module doc: a fit of zero tokens is not a
2606        // ceiling of zero, it is an estimate saying this model should
2607        // not have loaded -- and it did. Say so and admit as before.
2608        tracing::warn!(
2609            "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
2610             bytes against a {} byte budget. Serving with no derived context ceiling -- set \
2611             FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
2612             admit on a number you choose.",
2613            priced.weights_bytes,
2614            priced.device_budget_bytes,
2615        );
2616        return batcher;
2617    };
2618    tracing::info!("{source}");
2619    tracing::info!("{}", derived.fit);
2620    let adopted = budget::apply_derived(&mut batcher, &derived);
2621    if adopted.max_context {
2622        tracing::info!(
2623            "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
2624             override with FERROX_CB_MAX_CONTEXT",
2625            derived.max_context
2626        );
2627    }
2628    if adopted.kv_blocks {
2629        tracing::info!(
2630            "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
2631            derived.kv_blocks,
2632            batcher.kv_block_size
2633        );
2634    }
2635    batcher
2636}
2637
2638/// Turns a freshly loaded checkpoint into the parts that get published
2639/// as the active model.
2640///
2641/// Extracted from `build_app_state` so `/admin/models/load` builds its
2642/// replacement exactly the way startup builds the first one -- a second
2643/// copy of this match would be a second place for a new engine variant
2644/// to be forgotten, and the difference would only show up as a model
2645/// that silently loses continuous batching after a swap.
2646pub(crate) fn activate_loaded_model(
2647    loaded: model::LoadedModel,
2648    enable_continuous_batching: bool,
2649    path: Option<&str>,
2650) -> Activated {
2651    match loaded {
2652        model::LoadedModel::Gguf(g) => {
2653            let decoder = Arc::new(g.decoder);
2654            let tokenizer = Arc::new(g.tokenizer);
2655            let config = price_batcher_config(path);
2656            // Prefill is still a per-token `forward_token` loop on both
2657            // paths (see `sched-chunked-prefill`: chunking bought
2658            // fairness, not a batched prefill kernel), so a sliding
2659            // layer really does need only `window + 1 - 1` positions
2660            // live. `chunk = 1` here is the truth, not a simplification.
2661            let shape =
2662                ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
2663            let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
2664            let batcher = if enable_continuous_batching {
2665                tracing::info!(
2666                    "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
2667                     (stop sequences use the same pending-buffer trim as the private generate loop)"
2668                );
2669                let tok = Arc::clone(&tokenizer);
2670                let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
2671                Some(batch_scheduler::ContinuousBatcher::spawn_with_ceiling(
2672                    Arc::clone(&decoder),
2673                    decode,
2674                    config,
2675                    Arc::clone(&ceiling),
2676                ))
2677            } else {
2678                None
2679            };
2680            (
2681                Model::Gguf(GgufModel {
2682                    decoder,
2683                    tokenizer,
2684                    stop_tokens: g.stop_tokens,
2685                    bos_id: g.bos_id,
2686                    is_synthetic: g.is_synthetic,
2687                    chat_template: g.chat_template,
2688                }),
2689                batcher,
2690                Some(ceiling),
2691            )
2692        }
2693        model::LoadedModel::Kimi(k) => (
2694            Model::Kimi(KimiModel {
2695                engine: k.engine,
2696                tokenizer: k.tokenizer,
2697                stop_tokens: k.stop_tokens,
2698                chat_template: k.chat_template,
2699            }),
2700            None,
2701            None,
2702        ),
2703        model::LoadedModel::Mla(m) => (
2704            Model::Mla(MlaModel {
2705                engine: m.engine,
2706                tokenizer: m.tokenizer,
2707                stop_tokens: m.stop_tokens,
2708                bos_id: m.bos_id,
2709                name: m.name,
2710                chat_template: m.chat_template,
2711            }),
2712            None,
2713            None,
2714        ),
2715        model::LoadedModel::Gemma4(m) => (
2716            Model::Gemma4(Gemma4Model {
2717                engine: m.engine,
2718                tokenizer: m.tokenizer,
2719                stop_tokens: m.stop_tokens,
2720                bos_id: m.bos_id,
2721                name: m.name,
2722                chat_template: m.chat_template,
2723            }),
2724            None,
2725            None,
2726        ),
2727        model::LoadedModel::Glm52(g) => (
2728            Model::Glm52(Glm52Model {
2729                engine: g.engine,
2730                tokenizer: g.tokenizer,
2731                stop_tokens: g.stop_tokens,
2732                bos_id: g.bos_id,
2733                name: g.name,
2734                chat_template: g.chat_template,
2735            }),
2736            None,
2737            None,
2738        ),
2739    }
2740}
2741
2742fn build_app_state(
2743    loaded: model::LoadedModel,
2744    kv_pool: Option<generate::KvPoolConfig>,
2745    prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
2746    enable_continuous_batching: bool,
2747    mcp: Option<mcp::LoadedMcpConfig>,
2748    detection: Arc<health::Detection>,
2749) -> AppState {
2750    let (model, batcher, ceiling) = activate_loaded_model(
2751        loaded,
2752        enable_continuous_batching,
2753        std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
2754    );
2755    // The startup model's admin id is whichever discovered entry sits
2756    // at the configured path; `None` when it was not discovered (the
2757    // synthetic fallback, or a path outside the scanned directories),
2758    // in which case `/admin/models` reports nothing as active rather
2759    // than inventing an id no `load` request could name.
2760    let id = startup_model_id();
2761    AppState {
2762        active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
2763            id,
2764            model: Arc::new(model),
2765            batcher,
2766            ceiling,
2767        }))),
2768        load_in_progress: std::sync::atomic::AtomicBool::new(false),
2769        tasks: Arc::new(tasks::TaskRegistry::new()),
2770        cancels: Arc::new(cancel::CancelRegistry::new()),
2771        stats: stats::Stats::new(),
2772        streams: resume::StreamRegistry::new(),
2773        model_dir: admin::model_dirs().into_iter().next(),
2774        response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
2775        kv_pool,
2776        prefix_cache,
2777        sessions: session::SessionStore::new(),
2778        requests_total: std::sync::atomic::AtomicU64::new(0),
2779        request_errors_total: std::sync::atomic::AtomicU64::new(0),
2780        started_at: std::time::Instant::now(),
2781        last_request_ms: std::sync::atomic::AtomicU64::new(0),
2782        detection,
2783        mcp,
2784        continuous_batching_enabled: enable_continuous_batching,
2785        loading_model: Mutex::new(None),
2786        last_load_error: Mutex::new(None),
2787    }
2788}
2789
2790/// The `/admin/models` id of the checkpoint `FERROX_MODEL_PATH` names,
2791/// when discovery finds it. Matching on the resolved path rather than
2792/// on the filename keeps two same-named files in different directories
2793/// from claiming each other's id.
2794fn startup_model_id() -> Option<String> {
2795    let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
2796    let configured = std::fs::canonicalize(&configured).ok()?;
2797    admin::discover(&admin::model_dirs())
2798        .into_iter()
2799        .find(|d| {
2800            std::fs::canonicalize(&d.path)
2801                .map(|p| p == configured)
2802                .unwrap_or(false)
2803        })
2804        .map(|d| d.id)
2805}
2806
2807/// Builds the global rayon pool up front, on the main thread, with an
2808/// explicit width and QoS (see [`ferrox_core::threads`]).
2809///
2810/// Doing this from `main` rather than letting rayon build lazily is the
2811/// point: the first rayon call inside this server happens on a Tokio
2812/// `spawn_blocking` thread, so the workers used to inherit that thread's
2813/// QoS class -- which on macOS decides whether they land on performance
2814/// or efficiency cores.
2815fn init_cpu_pool() {
2816    match ferrox_core::threads::init_cpu_pool() {
2817        Some(n) => eprintln!(
2818            "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
2819            ferrox_core::threads::perf_core_count()
2820        ),
2821        None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
2822    }
2823}
2824
2825/// Prints the machine-readable ready line (see `ferrox_api::lifecycle`)
2826/// on stdout and flushes it.
2827///
2828/// This one line is what makes `--port 0` usable, and it deletes a whole
2829/// feature from any supervising process: no "is the port free" probe, no
2830/// `lsof` to work out whether an existing listener is a stale copy of
2831/// ourselves or a stranger's server, no dialog to explain the result.
2832/// The kernel picks the port and the child says what it got.
2833///
2834/// Shares stdout with the tracing subscriber on purpose -- a parent
2835/// reads stdout line by line and ignores anything that is not the ready
2836/// event, which `ServerReady::from_line` does for it.
2837fn announce_ready(addr: SocketAddr, scheme: &str) {
2838    use std::io::Write;
2839    let ready =
2840        ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
2841    let mut stdout = std::io::stdout().lock();
2842    let _ = writeln!(stdout, "{}", ready.to_line());
2843    let _ = stdout.flush();
2844}
2845
2846/// Resolves when the server should stop serving.
2847///
2848/// Stdin-close is the one orphan-prevention mechanism that behaves
2849/// identically on macOS, Windows and Linux and survives a parent that
2850/// dies rather than exiting cleanly: the kernel closes the pipe either
2851/// way. The POSIX alternative -- a signal handler plus an exit hook plus
2852/// a reaper -- has no Windows equivalent at all, since there is no
2853/// SIGTERM there.
2854///
2855/// When disabled this future never resolves, which is exactly the
2856/// previous behaviour: serve until the process is stopped externally.
2857async fn shutdown_signal(exit_on_stdin_close: bool) {
2858    if !exit_on_stdin_close {
2859        std::future::pending::<()>().await;
2860        return;
2861    }
2862    let _ = tokio::task::spawn_blocking(|| {
2863        use std::io::Read;
2864        let mut sink = [0u8; 256];
2865        let mut stdin = std::io::stdin().lock();
2866        loop {
2867            match stdin.read(&mut sink) {
2868                // EOF: the parent is gone, or closed the pipe.
2869                Ok(0) => break,
2870                // Input on stdin is not a protocol here; drain it.
2871                Ok(_) => continue,
2872                Err(e) => {
2873                    tracing::warn!("stdin read failed ({e}); treating it as closed");
2874                    break;
2875                }
2876            }
2877        }
2878    })
2879    .await;
2880    tracing::info!("stdin closed; shutting down");
2881}
2882
2883/// Tokio worker threads. The default is one per logical core, which on a
2884/// 10-core M2 Pro means 10 async workers oversubscribing the same cores
2885/// the rayon decode pool needs. Serving work here is almost entirely I/O
2886/// plus `spawn_blocking` handoff, so a small fixed pool is enough.
2887fn tokio_worker_threads() -> usize {
2888    std::env::var("FERROX_TOKIO_WORKERS")
2889        .ok()
2890        .and_then(|v| v.trim().parse::<usize>().ok())
2891        .filter(|n| *n > 0)
2892        .unwrap_or(2)
2893}
2894
2895/// Parses llama-server-style options and applies their environment
2896/// overrides before creating Tokio or Rayon worker threads. It then
2897/// brackets the async server lifecycle with journal records.
2898/// Install rustls' `ring` crypto provider as the process default.
2899///
2900/// `axum-server` is built with `tls-rustls-no-provider`, which
2901/// deliberately does NOT pick a backend -- see the comment on the
2902/// dependency in `Cargo.toml`. rustls then has no default provider, and
2903/// building a `ServerConfig` without one fails at ACCEPT time rather
2904/// than at compile time, which is the worst place for it to surface: a
2905/// server that started cleanly and refuses every TLS connection.
2906///
2907/// So this runs unconditionally at startup, not lazily in the TLS arm.
2908/// `install_default` returns `Err` if a provider is already installed,
2909/// which is not a failure -- it means something else got there first
2910/// and the invariant we care about (there IS a provider) already holds.
2911fn install_ring_crypto_provider() {
2912    let _ = rustls::crypto::ring::default_provider().install_default();
2913}
2914
2915/// Runs the server to completion.
2916///
2917/// Takes already-parsed arguments so the same library backs both the
2918/// `ferrox-server` binary and ferrox-cli's optional `serve` feature,
2919/// and neither front end can drift into its own startup logic.
2920pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
2921    if args.list_devices {
2922        print_available_devices();
2923        return Ok(());
2924    }
2925    apply_cli_overrides(&args)?;
2926
2927    // Before the model is loaded and before the port is bound: refuse
2928    // to be the second process holding weights on this host. Held for
2929    // the life of the process -- dropping it deregisters us.
2930    let _instance = {
2931        use ferrox_core::instance::{register, InstancePolicy};
2932        let policy = if args.allow_multiple_instances {
2933            InstancePolicy::Multi
2934        } else {
2935            InstancePolicy::from_env_or(InstancePolicy::Single)
2936        };
2937        let model = std::env::var("FERROX_MODEL_PATH").ok();
2938        register(
2939            "server",
2940            model.as_deref(),
2941            ferrox_core::instance::current_backend(),
2942            policy,
2943        )
2944        .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
2945    };
2946
2947    let journal = journal::Journal::from_env();
2948    eprintln!(
2949        "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
2950        journal.path()
2951    );
2952    journal.append(&journal::Record::session_start(
2953        env!("CARGO_PKG_VERSION"),
2954        std::process::id(),
2955    ));
2956    journal::install_panic_hook(journal.clone());
2957
2958    let mcp_config_path = args.mcp_config.clone();
2959    let exit_on_stdin_close = args.exit_on_stdin_close
2960        || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
2961            .map(|v| v == "1")
2962            .unwrap_or(false);
2963
2964    // Before Tokio exists, so the decode pool's threads are not spawned
2965    // from (and do not inherit the QoS of) a blocking-pool thread.
2966    // SAFETY: still single-threaded here.
2967    unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
2968    init_cpu_pool();
2969
2970    let runtime = tokio::runtime::Builder::new_multi_thread()
2971        .worker_threads(tokio_worker_threads())
2972        .enable_all()
2973        .build()?;
2974    let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
2975
2976    let reason = match &result {
2977        Ok(()) => "normal".to_string(),
2978        Err(e) => e.to_string(),
2979    };
2980    journal.append(&journal::Record::session_exit(reason));
2981
2982    // Dropping the runtime instead would wait for blocking tasks, and
2983    // the stdin watcher parks in a blocking read that may never return
2984    // (a terminal keeps stdin open forever). The serving future has
2985    // already finished by here, so nothing useful is being abandoned.
2986    runtime.shutdown_background();
2987
2988    result
2989}
2990
2991async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
2992    // `try_init`, not `init`. As a library this runs inside a process
2993    // that may already have a subscriber: ferrox-cli installs one
2994    // before it dispatches, so `ferrox serve` would panic on startup
2995    // with "a global default trace dispatcher has already been set".
2996    // Losing the race is not an error, it means logging is configured.
2997    let _ = tracing_subscriber::fmt::try_init();
2998
2999    // Fail-closed listener check, before anything else (including
3000    // loading the model, so a misconfigured bind fails fast rather than
3001    // after however long that takes): refuse to start bound to a
3002    // non-loopback address with no API key configured, unless the
3003    // operator has explicitly opted into that via
3004    // FERROX_ALLOW_UNAUTHENTICATED_REMOTE=1 -- see
3005    // `security::check_bind_authorization`'s doc comment for why an
3006    // address that doesn't even parse as loopback is treated the same
3007    // as a confirmed non-loopback one.
3008    let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
3009    let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
3010    let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
3011        .map(|v| v == "1")
3012        .unwrap_or(false);
3013    if let Err(msg) =
3014        security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
3015    {
3016        anyhow::bail!(msg);
3017    }
3018
3019    let mut loaded = model::load()?;
3020    match &loaded {
3021        model::LoadedModel::Gguf(g) => tracing::info!(
3022            "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
3023            g.decoder.config.name,
3024            g.is_synthetic,
3025            g.tokenizer.kind()
3026        ),
3027        model::LoadedModel::Kimi(k) => tracing::info!(
3028            "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
3029            k.tokenizer.vocab_size()
3030        ),
3031        model::LoadedModel::Mla(m) => tracing::info!(
3032            "loaded MLA GGUF '{}' (tokenizer={})",
3033            m.name,
3034            m.tokenizer.kind()
3035        ),
3036        model::LoadedModel::Gemma4(m) => tracing::info!(
3037            "loaded Gemma4 GGUF '{}' (tokenizer={})",
3038            m.name,
3039            m.tokenizer.kind()
3040        ),
3041        model::LoadedModel::Glm52(g) => tracing::info!(
3042            "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
3043            g.name,
3044            g.tokenizer.kind()
3045        ),
3046    }
3047    // Opt-in VRAM budget for GPU-resident MoE experts. When unset but
3048    // Metal is active, default to a large budget so routed experts that
3049    // have Metal-capable quants run via `run_expert_placed` (Metal
3050    // matvec) instead of staying on CPU after Metal attention. Explicit
3051    // `FERROX_GPU_VRAM_BUDGET_BYTES=0` keeps the historical all-CPU MoE
3052    // placement. CUDA builds still require an explicit budget (Vast /
3053    // multi-GPU hosts vary too much for a safe default).
3054    let metal_default_moe_budget = {
3055        #[cfg(feature = "metal")]
3056        {
3057            ferrox_core::metal_dense_enabled()
3058                && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
3059        }
3060        #[cfg(not(feature = "metal"))]
3061        {
3062            false
3063        }
3064    };
3065    if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
3066        let budget: u64 = budget_str
3067            .parse()
3068            .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
3069        match &mut loaded {
3070            model::LoadedModel::Gguf(g) => {
3071                tracing::info!(
3072                    "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
3073                     (CUDA and/or Metal matvecs when built with the matching feature)"
3074                );
3075                g.decoder.gpu_vram_budget_bytes = Some(budget);
3076            }
3077            model::LoadedModel::Kimi(_) => {
3078                tracing::warn!(
3079                    "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
3080                     supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
3081                );
3082            }
3083            model::LoadedModel::Mla(_) => {
3084                tracing::warn!(
3085                    "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
3086                     FFN path only today; ignoring expert VRAM budget"
3087                );
3088            }
3089            model::LoadedModel::Gemma4(_) => {
3090                tracing::warn!(
3091                    "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
3092                     ignoring expert VRAM budget"
3093                );
3094            }
3095            model::LoadedModel::Glm52(_) => {
3096                tracing::warn!(
3097                    "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
3098                     GPU expert placement not wired yet; ignoring"
3099                );
3100            }
3101        }
3102    } else if metal_default_moe_budget {
3103        // ~64 GiB sentinel: place as many experts as the planner allows;
3104        // Metal unified memory makes a hard VRAM split less meaningful
3105        // than on discrete CUDA cards.
3106        const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
3107        if let model::LoadedModel::Gguf(g) = &mut loaded {
3108            tracing::info!(
3109                "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
3110                 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
3111            );
3112            g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
3113        }
3114    }
3115    #[cfg(feature = "cuda")]
3116    {
3117        if ferrox_core::cuda_dense_enabled() {
3118            tracing::info!(
3119                "CUDA dense matvec enabled for WeightMatrix::apply \
3120                 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
3121            );
3122        } else {
3123            tracing::info!(
3124                "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
3125            );
3126        }
3127    }
3128    #[cfg(feature = "metal")]
3129    {
3130        if ferrox_core::metal_dense_enabled() {
3131            tracing::info!(
3132                "Metal dense matvec enabled for WeightMatrix::apply \
3133                 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
3134            );
3135            match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
3136                Some("1") | Some("true") | Some("on") | Some("attn") => {
3137                    tracing::info!(
3138                        "Metal fused attention requested (FERROX_METAL_ATTN): \
3139                         QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
3140                    );
3141                }
3142                _ => {}
3143            }
3144            match std::env::var("FERROX_METAL_LOGITS").ok().as_deref() {
3145                Some("1") | Some("true") | Some("on") | Some("logits") => {
3146                    tracing::info!(
3147                        "Metal logits-in-stack enabled (FERROX_METAL_LOGITS): \
3148                         final_norm+lm_head on-GPU (often slower than host lm_head)"
3149                    );
3150                }
3151                _ => {}
3152            }
3153            match std::env::var("FERROX_METAL_GREEDY_GPU").ok().as_deref() {
3154                Some("0") | Some("false") | Some("off") | Some("no") => {
3155                    tracing::info!(
3156                        "Metal greedy GPU argmax disabled (FERROX_METAL_GREEDY_GPU=0); \
3157                         temperature<=0 uses host lm_head"
3158                    );
3159                }
3160                _ => {
3161                    tracing::info!(
3162                        "Metal greedy GPU argmax on by default (opt out FERROX_METAL_GREEDY_GPU=0): \
3163                         temperature<=0 folds final_norm+lm_head+argmax into dense stack"
3164                    );
3165                }
3166            }
3167        } else {
3168            tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
3169        }
3170    }
3171    // Both env vars are required together to enable pooling; unset ->
3172    // caches keep their original unbounded-per-request growth. This
3173    // mirrors the FERROX_API_KEY / FERROX_RATE_LIMIT_PER_MINUTE
3174    // pattern below: opt-in, off by default.
3175    //
3176    // Block count can be set explicitly (`FERROX_KV_POOL_BLOCKS` +
3177    // `FERROX_KV_POOL_BLOCK_SIZE`) or derived from a byte budget
3178    // (`FERROX_KV_BYTE_BUDGET` + `FERROX_KV_POOL_BLOCK_SIZE`, GGUF
3179    // models only). `FERROX_KV_POOL_BLOCKS` and
3180    // `FERROX_KV_BYTE_BUDGET` are mutually exclusive.
3181    let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
3182    let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
3183    let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
3184    if blocks_env.is_ok() && byte_budget_env.is_ok() {
3185        panic!(
3186            "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
3187             (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
3188        );
3189    }
3190    let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
3191        (Ok(blocks), Ok(block_size), Err(_)) => {
3192            let total_blocks: usize = blocks
3193                .parse()
3194                .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
3195            let block_size: usize = block_size
3196                .parse()
3197                .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
3198            // Optional and independent of the two above: how long a
3199            // request retries before giving up when the pool is
3200            // momentarily exhausted, instead of rejecting on the very
3201            // first failed attempt. Zero (the default if unset)
3202            // preserves the original reject-immediately behavior.
3203            let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
3204                .ok()
3205                .map(|v| {
3206                    v.parse()
3207                        .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
3208                })
3209                .unwrap_or(0);
3210            tracing::info!(
3211                "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
3212                 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
3213            );
3214            Some(generate::KvPoolConfig {
3215                pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
3216                queue_wait: Duration::from_millis(queue_wait_ms),
3217            })
3218        }
3219        (Err(_), Ok(block_size), Ok(byte_budget)) => {
3220            let block_size: usize = block_size
3221                .parse()
3222                .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
3223            let budget: u64 = byte_budget
3224                .parse()
3225                .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
3226            let cfg = match &loaded {
3227                model::LoadedModel::Gguf(g) => &g.decoder.config,
3228                model::LoadedModel::Kimi(_)
3229                | model::LoadedModel::Mla(_)
3230                | model::LoadedModel::Gemma4(_)
3231                | model::LoadedModel::Glm52(_) => {
3232                    panic!(
3233                        "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
3234                         (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
3235                    );
3236                }
3237            };
3238            let bytes_per_block = block_size
3239                * cfg.n_layers
3240                * cfg.n_kv_heads
3241                * cfg.head_dim
3242                * 2
3243                * std::mem::size_of::<f32>();
3244            assert!(
3245                bytes_per_block > 0,
3246                "derived KV block byte size must be positive (check model config and block size)"
3247            );
3248            let total_blocks = (budget as usize / bytes_per_block).max(1);
3249            let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
3250                .ok()
3251                .map(|v| {
3252                    v.parse()
3253                        .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
3254                })
3255                .unwrap_or(0);
3256            tracing::info!(
3257                "KV cache block pool enabled from byte budget: {budget} bytes / \
3258                 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
3259                 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
3260                cfg.n_layers
3261            );
3262            Some(generate::KvPoolConfig {
3263                pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
3264                queue_wait: Duration::from_millis(queue_wait_ms),
3265            })
3266        }
3267        (Err(_), Err(_), Err(_)) => None,
3268        (Err(_), Ok(_), Err(_)) => panic!(
3269            "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
3270             (or unset all three to disable KV cache pooling)"
3271        ),
3272        (Ok(_), Ok(_), Ok(_)) => {
3273            unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
3274        }
3275        (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
3276            "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
3277             set together (or neither, to disable KV cache pooling)"
3278        ),
3279    };
3280    // Mutually exclusive with kv_pool (see generate::generate's doc
3281    // comment on why a pool-backed cache can't safely be restored from
3282    // a prefix-cache clone): if both are set, the KV pool wins and
3283    // prefix caching is simply never consulted -- generate() already
3284    // enforces this per-request, so this is a heads-up for the
3285    // operator, not a hard failure.
3286    let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
3287        let max_entries: usize = v
3288            .parse()
3289            .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
3290        if kv_pool.is_some() {
3291            tracing::warn!(
3292                "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
3293                     caching will never be consulted while a KV pool is configured"
3294            );
3295        }
3296        tracing::info!(
3297            "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
3298                 all requests"
3299        );
3300        Arc::new(Mutex::new(PrefixCache::new(max_entries)))
3301    });
3302    if matches!(
3303        loaded,
3304        model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
3305    ) && (kv_pool.is_some() || prefix_cache.is_some())
3306    {
3307        tracing::warn!(
3308            "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
3309             neither is consulted for those engines (state shapes differ from Decoder KV); see \
3310             ferrox_models::engine's module docs"
3311        );
3312    }
3313    let enable_cb = std::env::var("FERROX_CONTINUOUS_BATCHING")
3314        .map(|v| v == "1")
3315        .unwrap_or(false)
3316        && kv_pool.is_none()
3317        && prefix_cache.is_none()
3318        && matches!(loaded, model::LoadedModel::Gguf(_));
3319    if std::env::var("FERROX_CONTINUOUS_BATCHING")
3320        .map(|v| v == "1")
3321        .unwrap_or(false)
3322        && (kv_pool.is_some() || prefix_cache.is_some())
3323    {
3324        tracing::warn!(
3325            "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
3326             (those modes keep the private generate path)"
3327        );
3328    }
3329    if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
3330        if let Ok(chunk) = n.parse::<usize>() {
3331            if chunk > 0 {
3332                tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
3333            }
3334        }
3335    }
3336    if matches!(
3337        std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
3338        Some("1")
3339    ) {
3340        tracing::warn!(
3341            "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
3342             (minimal spill; full layer offload still planned)"
3343        );
3344    }
3345
3346    let mcp = match mcp_config_path {
3347        Some(path) => {
3348            let loaded = mcp::load_mcp_config(&path)?;
3349            tracing::info!(
3350                "MCP config loaded from {} ({} server(s); invocation not wired yet)",
3351                loaded.path,
3352                loaded.servers.len()
3353            );
3354            Some(loaded)
3355        }
3356        None => None,
3357    };
3358
3359    // Started before the router is built so the probe overlaps with
3360    // binding the port: by the time a client can ask, it has usually
3361    // already landed.
3362    let detection = health::Detection::spawn();
3363
3364    let state = Arc::new(build_app_state(
3365        loaded,
3366        kv_pool,
3367        prefix_cache,
3368        enable_cb,
3369        mcp,
3370        detection,
3371    ));
3372
3373    // Paths come from `ferrox_api::routes` rather than string literals
3374    // so the UI, `ferrox chat` and this router cannot disagree about
3375    // what the surface is.
3376    use ferrox_api::routes;
3377
3378    // Ferrox Studio is a separate app served by its own dev/static
3379    // server (see `ui/` at the repository root); it reaches this
3380    // process over the public HTTP API like any other client, so there
3381    // is nothing to mount here and `/` stays a 404.
3382    let public = Router::new().route(routes::HEALTH, get(health));
3383
3384    let mut protected = Router::new()
3385        .route(routes::V1_MODELS, get(list_models))
3386        .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3387        // Behind the same key as the endpoint that started the work:
3388        // an unauthenticated caller must not be able to stop someone
3389        // else's generation by guessing at request ids.
3390        .route(routes::V1_CANCEL, post(cancel_generation))
3391        // Reconnect and the polling fallback, both behind the same key
3392        // as the request that filled the buffer: the replay window holds
3393        // the model's output, so reading it must cost what producing it
3394        // cost.
3395        .route(&resume_route(routes::V1_STREAM), get(resume::resume))
3396        .route(&resume_route(routes::V1_STREAM_POLL), get(resume::poll))
3397        .route(routes::V1_MESSAGES, post(anthropic::messages))
3398        .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3399        .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3400        .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3401        .route(routes::V1_EMBEDDINGS, post(openai_extra::embeddings))
3402        .route(routes::CACHE_STATS, get(cache_stats))
3403        .route(routes::METRICS, get(metrics))
3404        // The control surface. Registered inside `protected` on
3405        // purpose: these routes change what the server serves and write
3406        // to disk, so they get the same FERROX_API_KEY gate as /v1/*
3407        // and never the unauthenticated treatment /health has.
3408        .route(routes::ADMIN_MODELS, get(admin::models))
3409        .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3410        .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3411        .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3412        .route(routes::ADMIN_TASKS, get(admin::tasks))
3413        .route(&admin::cancel_route(), post(admin::cancel_task))
3414        .route(routes::ADMIN_STATS, get(admin::stats));
3415
3416    // Both off by default; set the corresponding env var to enable.
3417    // route_layer (not layer) so these apply only to the routes above,
3418    // never to /health, which stays reachable for liveness/readiness
3419    // probes regardless of auth or rate-limit configuration.
3420    if let Ok(key) = std::env::var("FERROX_API_KEY") {
3421        tracing::info!("API key auth enabled");
3422        let auth = limits::AuthConfig {
3423            api_key: Arc::new(key),
3424        };
3425        protected = protected.route_layer(axum::middleware::from_fn_with_state(
3426            auth,
3427            limits::require_api_key,
3428        ));
3429    }
3430    if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
3431        let rpm: u32 = rpm
3432            .parse()
3433            .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
3434        tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
3435        let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
3436        protected = protected.route_layer(axum::middleware::from_fn_with_state(
3437            limiter,
3438            limits::rate_limit,
3439        ));
3440    }
3441    // Off by default; set FERROX_CORS_ORIGINS (comma-separated exact
3442    // origins) to enable. No wildcard support by design -- see
3443    // `security::parse_cors_origins`'s doc comment. Added last (so it's
3444    // the outermost route_layer, run before auth/rate-limiting): a CORS
3445    // preflight (OPTIONS) request carries no Authorization header and
3446    // is answered directly by `CorsLayer` itself, so it must not be
3447    // blocked by the auth/rate-limit layers underneath.
3448    if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
3449        let origins = security::parse_cors_origins(&spec)
3450            .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
3451        tracing::info!(
3452            "CORS enabled: {} allow-listed origin(s) ({})",
3453            origins.len(),
3454            spec
3455        );
3456        let cors = tower_http::cors::CorsLayer::new()
3457            .allow_origin(tower_http::cors::AllowOrigin::list(origins))
3458            .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
3459            .allow_headers([
3460                axum::http::header::CONTENT_TYPE,
3461                axum::http::header::AUTHORIZATION,
3462                // The self-declared client label the monitor records
3463                // (see `attribution`). A custom header makes every
3464                // cross-origin call preflighted, so omitting it here
3465                // would not merely drop the label -- it would fail the
3466                // request outright.
3467                axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
3468                // Set by hand rather than by `EventSource`, because
3469                // this API needs POST and a bearer token. Same
3470                // consequence if it is missing.
3471                axum::http::HeaderName::from_static("last-event-id"),
3472            ]);
3473        protected = protected.route_layer(cors);
3474    }
3475
3476    // Outermost on purpose: every 503 this server can emit -- from a
3477    // handler, from `require_active`, or from the batch scheduler's
3478    // queue cap -- leaves with a `Retry-After` a client can act on.
3479    let app = public
3480        .merge(protected)
3481        .layer(axum::middleware::from_fn(limits::retry_after))
3482        .with_state(state);
3483
3484    // TLS is off by default -- set FERROX_TLS_CERT and FERROX_TLS_KEY
3485    // together to serve HTTPS instead of plain HTTP; unset (either or
3486    // both) preserves the original plain-HTTP behavior exactly. See
3487    // `security::tls_paths_from_env`'s doc comment for why this can't
3488    // be meaningfully unit-tested here.
3489    let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
3490    install_ring_crypto_provider();
3491    // Both arms bind first and read the address back off the socket
3492    // rather than trusting the requested one: with `--port 0` the
3493    // requested port is a lie by construction, and the ready line has
3494    // to carry what the kernel actually handed out.
3495    match tls_paths {
3496        Some(paths) => {
3497            let config =
3498                axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
3499                    .await
3500                    .map_err(|e| {
3501                        anyhow::anyhow!(
3502                            "failed to load TLS cert/key ({:?}, {:?}): {e}",
3503                            paths.cert,
3504                            paths.key
3505                        )
3506                    })?;
3507            let socket_addr: std::net::SocketAddr = addr
3508                .parse()
3509                .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
3510            let listener = std::net::TcpListener::bind(socket_addr)?;
3511            // Tokio panics outright when handed a BLOCKING socket
3512            // ("Registering a blocking socket with the tokio runtime is
3513            // unsupported"), and axum-server registers this one
3514            // internally. Without this the TLS arm binds, prints its
3515            // ready line, and then panics on the first accept -- so the
3516            // failure looks like a healthy start followed by a server
3517            // that answers nothing.
3518            listener.set_nonblocking(true)?;
3519            let bound = listener.local_addr()?;
3520            tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
3521            announce_ready(bound, "https");
3522
3523            let handle = axum_server::Handle::new();
3524            let shutdown_handle = handle.clone();
3525            tokio::spawn(async move {
3526                shutdown_signal(exit_on_stdin_close).await;
3527                shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
3528            });
3529            axum_server::from_tcp_rustls(listener, config)?
3530                .handle(handle)
3531                .serve(app.into_make_service())
3532                .await?;
3533        }
3534        None => {
3535            let listener = tokio::net::TcpListener::bind(&addr).await?;
3536            let bound = listener.local_addr()?;
3537            tracing::info!("ferrox-server listening on {bound}");
3538            announce_ready(bound, "http");
3539            axum::serve(listener, app)
3540                .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
3541                .await?;
3542        }
3543    }
3544    Ok(())
3545}
3546
3547#[cfg(test)]
3548mod tests {
3549    use super::*;
3550    use ferrox_models::config::test_dense_fixture;
3551
3552    #[test]
3553    fn parses_llama_server_style_options() {
3554        let argv = [
3555            "ferrox-server",
3556            "-m",
3557            "model.gguf",
3558            "--host",
3559            "::1",
3560            "--port",
3561            "9000",
3562            "-t",
3563            "4",
3564            "-dev",
3565            "Metal",
3566            "-ngl",
3567            "all",
3568        ]
3569        .into_iter()
3570        .map(String::from)
3571        .collect();
3572        let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
3573
3574        assert_eq!(args.model.as_deref(), Some("model.gguf"));
3575        assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
3576        assert_eq!(args.port, Some(9000));
3577        assert_eq!(args.threads, Some(4));
3578        assert_eq!(args.device, Some(OffloadDevice::Metal));
3579        assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
3580        assert_eq!(
3581            cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
3582            Some("[::1]:9000")
3583        );
3584    }
3585
3586    #[test]
3587    fn port_zero_survives_argument_parsing_as_a_real_request() {
3588        // `--port 0` must reach the bind call intact: it is a request
3589        // for a kernel-assigned port, not a missing value to default to
3590        // 8383. The address it produces is deliberately provisional --
3591        // the ready line reports what was actually bound.
3592        let argv = ["ferrox-server", "--port", "0"]
3593            .into_iter()
3594            .map(String::from)
3595            .collect();
3596        let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
3597        assert_eq!(args.port, Some(0));
3598        assert_eq!(
3599            cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
3600            Some("127.0.0.1:0")
3601        );
3602    }
3603
3604    #[test]
3605    fn stdin_close_exit_is_opt_in() {
3606        // Default off: a server whose stdin is /dev/null (systemd, cron,
3607        // nohup) would otherwise exit the instant it started.
3608        let args =
3609            ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
3610        assert!(!args.exit_on_stdin_close);
3611        let args = ServerArgs::try_parse_from(
3612            ["ferrox-server", "--exit-on-stdin-close"]
3613                .into_iter()
3614                .map(String::from),
3615        )
3616        .unwrap();
3617        assert!(args.exit_on_stdin_close);
3618    }
3619
3620    #[test]
3621    fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
3622        let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
3623        let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
3624        let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
3625        assert_eq!(parsed.port, 51999);
3626        assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
3627        // A parent reads stdout line by line; tracing shares the stream.
3628        assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
3629    }
3630
3631    fn test_model() -> Model {
3632        // Tiny vocab (32): raw byte ids ≥32 (e.g. ASCII "hello") are OOV.
3633        // HTTP/chat-template tests that need full ASCII use
3634        // `test_model_full_byte_vocab` instead.
3635        let cfg = test_dense_fixture();
3636        Model::Gguf(GgufModel {
3637            decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
3638            tokenizer: Arc::new(ServerTokenizer::Byte),
3639            stop_tokens: StopTokens::default(),
3640            bos_id: None,
3641            is_synthetic: true,
3642            chat_template: chat_template::ChatTemplate::Plain,
3643        })
3644    }
3645
3646    fn greedy_params(max_tokens: usize) -> GenerationParams {
3647        GenerationParams {
3648            max_tokens,
3649            sampling: SamplingParams::default(),
3650            seed: 1,
3651            stop: Vec::new(),
3652            stop_token_ids: Vec::new(),
3653            json_object: false,
3654            cancel: None,
3655        }
3656    }
3657
3658    /// Declares a full 0..255 byte-compatible vocab so HTTP-level tests
3659    /// that render chat templates (ASCII role names) do not spuriously
3660    /// reject their own prompt prefixes.
3661    fn test_model_full_byte_vocab() -> Model {
3662        let mut cfg = test_dense_fixture();
3663        cfg.vocab_size = 256;
3664        Model::Gguf(GgufModel {
3665            decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
3666            tokenizer: Arc::new(ServerTokenizer::Byte),
3667            stop_tokens: StopTokens::default(),
3668            bos_id: None,
3669            is_synthetic: true,
3670            chat_template: chat_template::ChatTemplate::Plain,
3671        })
3672    }
3673
3674    /// One `AppState` for the HTTP-level tests, so a new field on the
3675    /// struct is added in one place rather than in every test that
3676    /// builds one.
3677    fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
3678        AppState {
3679            active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3680                id: None,
3681                model: Arc::new(model),
3682                batcher: None,
3683                ceiling: None,
3684            }))),
3685            load_in_progress: std::sync::atomic::AtomicBool::new(false),
3686            tasks: Arc::new(tasks::TaskRegistry::new()),
3687            cancels: Arc::new(cancel::CancelRegistry::new()),
3688            stats: stats::Stats::new(),
3689            streams: resume::StreamRegistry::new(),
3690            model_dir: None,
3691            response_cache: Mutex::new(response_cache),
3692            kv_pool: None,
3693            prefix_cache: None,
3694            sessions: session::SessionStore::new(),
3695            requests_total: std::sync::atomic::AtomicU64::new(0),
3696            request_errors_total: std::sync::atomic::AtomicU64::new(0),
3697            started_at: std::time::Instant::now(),
3698            last_request_ms: std::sync::atomic::AtomicU64::new(0),
3699            detection: Arc::new(health::Detection::ready(health::probe_backends())),
3700            mcp: None,
3701            continuous_batching_enabled: false,
3702            loading_model: Mutex::new(None),
3703            last_load_error: Mutex::new(None),
3704        }
3705    }
3706
3707    /// A real axum `Router` wired exactly like `main()`'s (minus auth/
3708    /// rate-limiting, which are orthogonal and already covered by
3709    /// `limits`'s own tests), backed by a fresh
3710    /// `test_model_full_byte_vocab()` -- so tool-calling/session tests
3711    /// exercise the real HTTP request/response path (JSON
3712    /// (de)serialization, routing, handler wiring, chat-template
3713    /// rendering) via `tower::ServiceExt::oneshot`, not just the inner
3714    /// functions directly.
3715    fn test_app() -> Router {
3716        test_app_with_state(Arc::new(test_state(
3717            test_model_full_byte_vocab(),
3718            ResponseCache::new(1000, Duration::from_secs(3600)),
3719        )))
3720    }
3721
3722    /// [`test_app`] over a caller-owned state, so a test can reach in
3723    /// and swap or unload the model behind a live router.
3724    fn test_app_with_state(state: Arc<AppState>) -> Router {
3725        Router::new()
3726            .route(ferrox_api::routes::HEALTH, get(health))
3727            .route(ferrox_api::routes::V1_MODELS, get(list_models))
3728            .route("/v1/chat/completions", post(chat_completions))
3729            .route("/v1/tokenize", post(openai_extra::tokenize))
3730            .route("/v1/detokenize", post(openai_extra::detokenize))
3731            .route("/v1/embeddings", post(openai_extra::embeddings))
3732            .route("/v1/completions", post(openai_extra::completions))
3733            .route(
3734                ferrox_api::routes::ADMIN_MODELS_UNLOAD,
3735                post(admin::unload_model),
3736            )
3737            .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
3738            .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
3739            .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
3740            .route(
3741                &resume_route(ferrox_api::routes::V1_STREAM),
3742                get(resume::resume),
3743            )
3744            .route(
3745                &resume_route(ferrox_api::routes::V1_STREAM_POLL),
3746                get(resume::poll),
3747            )
3748            .with_state(state)
3749    }
3750
3751    fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
3752        let mut cfg = test_dense_fixture();
3753        cfg.name = name;
3754        cfg.vocab_size = vocab_size;
3755        Model::Gguf(GgufModel {
3756            decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
3757            tokenizer: Arc::new(ServerTokenizer::Byte),
3758            stop_tokens: StopTokens::default(),
3759            bos_id: None,
3760            is_synthetic: true,
3761            chat_template: chat_template::ChatTemplate::Plain,
3762        })
3763    }
3764
3765    fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
3766        Arc::new(ActiveModel {
3767            id: Some(name.to_string()),
3768            model: Arc::new(named_test_model(name, 256)),
3769            batcher: None,
3770            ceiling: None,
3771        })
3772        .tap_into(state)
3773    }
3774
3775    /// Small helper so the swap tests read as "publish this model".
3776    trait TapInto {
3777        fn tap_into(self, state: &AppState) -> Self;
3778    }
3779    impl TapInto for Arc<ActiveModel> {
3780        fn tap_into(self, state: &AppState) -> Self {
3781            state.swap_active(Some(Arc::clone(&self)));
3782            self
3783        }
3784    }
3785
3786    /// The load-order guarantee the whole swap design exists to make:
3787    /// a request that has already taken its handle finishes against the
3788    /// weights it started on, even though a different model has since
3789    /// been published. Anything else would splice two checkpoints into
3790    /// one completion.
3791    #[test]
3792    fn an_in_flight_request_keeps_the_model_it_started_on() {
3793        let state = test_state(
3794            named_test_model("model-a", 256),
3795            ResponseCache::new(4, Duration::from_secs(60)),
3796        );
3797
3798        // A request that has begun: it has cloned the handle and is
3799        // about to decode against it.
3800        let in_flight = state.active().expect("a model is loaded");
3801        assert_eq!(in_flight.model.name(), "model-a");
3802
3803        active_model(&state, "model-b");
3804
3805        // The swap is visible to anything that asks *now*...
3806        assert_eq!(state.active().unwrap().model.name(), "model-b");
3807        // ...and completely invisible to the request already running.
3808        assert_eq!(in_flight.model.name(), "model-a");
3809        let (_chunks, finish, _usage) = run_generation(
3810            &in_flight.model,
3811            "hi",
3812            &greedy_params(3),
3813            None,
3814            None,
3815            None,
3816            None,
3817        )
3818        .expect("the old model must still decode after being swapped out");
3819        assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
3820    }
3821
3822    /// The other half of the same guarantee: the old model is not freed
3823    /// at swap time, it is freed when the last holder lets go. A design
3824    /// that dropped it eagerly would free weights out from under a
3825    /// decode loop.
3826    #[test]
3827    fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
3828        let state = test_state(
3829            named_test_model("model-a", 256),
3830            ResponseCache::new(4, Duration::from_secs(60)),
3831        );
3832        let in_flight = state.active().expect("a model is loaded");
3833        let weights = Arc::clone(&in_flight.model);
3834        assert!(Arc::strong_count(&weights) >= 2);
3835
3836        let previous = state.swap_active(Some(Arc::new(ActiveModel {
3837            id: Some("model-b".to_string()),
3838            model: Arc::new(named_test_model("model-b", 256)),
3839            batcher: None,
3840            ceiling: None,
3841        })));
3842        drop(previous);
3843        // The registry has let go; the in-flight request has not.
3844        assert!(Arc::strong_count(&weights) >= 2);
3845        drop(in_flight);
3846        assert_eq!(Arc::strong_count(&weights), 1);
3847    }
3848
3849    /// Unload is not "keep serving the last thing loaded". A request
3850    /// that arrives afterwards must be told there is no model, not
3851    /// quietly served by a checkpoint the operator dropped.
3852    #[tokio::test]
3853    async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
3854        let state = Arc::new(test_state(
3855            named_test_model("model-a", 256),
3856            ResponseCache::new(4, Duration::from_secs(60)),
3857        ));
3858        let app = test_app_with_state(Arc::clone(&state));
3859
3860        let (status, body) = post_json_uri(
3861            &app,
3862            ferrox_api::routes::ADMIN_MODELS_UNLOAD,
3863            serde_json::json!({}),
3864        )
3865        .await;
3866        assert_eq!(status, StatusCode::OK);
3867        assert_eq!(body["ok"], true);
3868        assert!(body["active"].is_null());
3869        assert!(state.active().is_none());
3870
3871        let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
3872        assert_eq!(status, StatusCode::OK);
3873        let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
3874        assert_eq!(models["data"].as_array().unwrap().len(), 0);
3875
3876        let (status, body) = post_json_uri(
3877            &app,
3878            "/v1/chat/completions",
3879            serde_json::json!({
3880                "model": "x",
3881                "messages": [{"role": "user", "content": "hi"}]
3882            }),
3883        )
3884        .await;
3885        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
3886        assert_eq!(body["error"]["type"], "model_not_loaded");
3887    }
3888
3889    /// `/health` must keep answering with nothing loaded -- a supervisor
3890    /// polls it to decide whether to kill the process, and "no model"
3891    /// is not "no server".
3892    #[tokio::test]
3893    async fn health_reports_the_unloaded_state_rather_than_going_silent() {
3894        let state = Arc::new(test_state(
3895            named_test_model("model-a", 256),
3896            ResponseCache::new(4, Duration::from_secs(60)),
3897        ));
3898        let app = test_app_with_state(Arc::clone(&state));
3899        state.swap_active(None);
3900
3901        let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
3902        // Not `ready`: a supervisor reading 200 here would route traffic
3903        // that is guaranteed to 503 on arrival.
3904        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
3905        assert_eq!(body["state"], "unavailable");
3906        assert_eq!(body["reason"], "model_not_loaded");
3907        assert!(body["model"].is_null());
3908        let real_weights = body["capabilities"]
3909            .as_array()
3910            .unwrap()
3911            .iter()
3912            .find(|c| c["id"] == "real_weights")
3913            .cloned()
3914            .expect("real_weights is always reported");
3915        assert_eq!(real_weights["available"], false);
3916        assert_eq!(real_weights["reason"], "model_not_loaded");
3917    }
3918
3919    /// The API-monitor contract: a finished request lands in the ring
3920    /// buffer keyed by the id the response carried, with the two
3921    /// durations reported separately.
3922    #[tokio::test]
3923    async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
3924        let app = test_app();
3925
3926        let (status, completion) = post_json_uri(
3927            &app,
3928            "/v1/chat/completions",
3929            serde_json::json!({
3930                "model": "x",
3931                "messages": [{"role": "user", "content": "hi"}],
3932                "max_tokens": 4
3933            }),
3934        )
3935        .await;
3936        assert_eq!(status, StatusCode::OK);
3937        let request_id = completion["request_id"].as_str().unwrap().to_string();
3938
3939        let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
3940        assert_eq!(status, StatusCode::OK);
3941        let recent = stats["recent"].as_array().unwrap();
3942        assert_eq!(recent.len(), 1);
3943        let row = &recent[0];
3944        assert_eq!(row["request_id"], request_id);
3945        assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
3946        assert_eq!(row["status"], 200);
3947        assert_eq!(row["stream"], false);
3948        // Separate fields, and the decode phase is a real measurement
3949        // rather than a copy of the total.
3950        assert!(row["duration_ms"].is_number());
3951        assert!(row["decode_ms"].is_number());
3952        assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
3953        assert_eq!(
3954            stats["tokens_prompt_total"].as_u64().unwrap(),
3955            row["prompt_tokens"].as_u64().unwrap()
3956        );
3957    }
3958
3959    /// A rejected request is still a request the monitor should show;
3960    /// otherwise the screen quietly omits exactly the traffic someone
3961    /// is debugging.
3962    #[tokio::test]
3963    async fn a_rejected_request_is_recorded_too() {
3964        let state = Arc::new(test_state(
3965            named_test_model("model-a", 256),
3966            ResponseCache::new(4, Duration::from_secs(60)),
3967        ));
3968        let app = test_app_with_state(Arc::clone(&state));
3969        state.swap_active(None);
3970
3971        let (status, _) = post_json_uri(
3972            &app,
3973            "/v1/chat/completions",
3974            serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
3975        )
3976        .await;
3977        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
3978
3979        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
3980        let recent = stats["recent"].as_array().unwrap();
3981        assert_eq!(recent.len(), 1);
3982        assert_eq!(recent[0]["status"], 503);
3983        assert_eq!(recent[0]["completion_tokens"], 0);
3984        assert!(recent[0]["decode_ms"].is_null());
3985        assert_eq!(stats["errors_total"], 1);
3986    }
3987
3988    /// POSTs with caller-supplied headers, so the attribution tests
3989    /// exercise the same header parsing a real client's request goes
3990    /// through rather than calling `Attribution::from_headers` twice.
3991    async fn post_json_with_headers(
3992        app: &Router,
3993        uri: &str,
3994        body: serde_json::Value,
3995        headers: &[(&str, &str)],
3996    ) -> (StatusCode, serde_json::Value) {
3997        use http_body_util::BodyExt;
3998        use tower::ServiceExt;
3999
4000        let mut builder = axum::http::Request::builder()
4001            .method("POST")
4002            .uri(uri)
4003            .header("content-type", "application/json");
4004        for (name, value) in headers {
4005            builder = builder.header(*name, *value);
4006        }
4007        let response = app
4008            .clone()
4009            .oneshot(
4010                builder
4011                    .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
4012                    .unwrap(),
4013            )
4014            .await
4015            .unwrap();
4016        let status = response.status();
4017        let bytes = response.into_body().collect().await.unwrap().to_bytes();
4018        let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
4019        (status, json)
4020    }
4021
4022    /// The three small endpoints used to be served and never recorded,
4023    /// which made the monitor wrong rather than incomplete: an editor
4024    /// hammering `/v1/embeddings` showed up as an idle server.
4025    #[tokio::test]
4026    async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
4027        let app = test_app();
4028
4029        let (status, _) = post_json_uri(
4030            &app,
4031            ferrox_api::routes::V1_TOKENIZE,
4032            serde_json::json!({"prompt": "hello"}),
4033        )
4034        .await;
4035        assert_eq!(status, StatusCode::OK);
4036        let (status, _) = post_json_uri(
4037            &app,
4038            ferrox_api::routes::V1_DETOKENIZE,
4039            serde_json::json!({"tokens": [104, 105]}),
4040        )
4041        .await;
4042        assert_eq!(status, StatusCode::OK);
4043        let (status, _) = post_json_uri(
4044            &app,
4045            ferrox_api::routes::V1_EMBEDDINGS,
4046            serde_json::json!({"input": "hello"}),
4047        )
4048        .await;
4049        assert_eq!(status, StatusCode::OK);
4050
4051        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4052        let routes: Vec<&str> = stats["recent"]
4053            .as_array()
4054            .unwrap()
4055            .iter()
4056            .map(|row| row["route"].as_str().unwrap())
4057            .collect();
4058        for expected in [
4059            ferrox_api::routes::V1_TOKENIZE,
4060            ferrox_api::routes::V1_DETOKENIZE,
4061            ferrox_api::routes::V1_EMBEDDINGS,
4062        ] {
4063            assert!(
4064                routes.contains(&expected),
4065                "{expected} is missing: {routes:?}"
4066            );
4067        }
4068
4069        let row = |route: &str| {
4070            stats["recent"]
4071                .as_array()
4072                .unwrap()
4073                .iter()
4074                .find(|r| r["route"] == route)
4075                .cloned()
4076                .unwrap()
4077        };
4078        // Embeddings run a forward pass, so their prompt tokens are
4079        // real prompt tokens. There is no decode loop, so `decode_ms`
4080        // stays null instead of borrowing the total.
4081        let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
4082        assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
4083        assert!(embed["decode_ms"].is_null());
4084        assert_eq!(embed["completion_tokens"], 0);
4085        // Tokenizing runs the tokenizer and not the model, so it
4086        // contributes nothing to the token counters those counters
4087        // claim to measure.
4088        assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
4089        assert_eq!(
4090            stats["tokens_prompt_total"].as_u64().unwrap(),
4091            embed["prompt_tokens"].as_u64().unwrap(),
4092            "only the forward pass counted"
4093        );
4094    }
4095
4096    /// A failed small-endpoint call is still traffic. A 400 that leaves
4097    /// no row is indistinguishable from a request that was never sent.
4098    #[tokio::test]
4099    async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
4100        let app = test_app();
4101        let (status, _) = post_json_uri(
4102            &app,
4103            ferrox_api::routes::V1_EMBEDDINGS,
4104            serde_json::json!({"input": "hi", "encoding_format": "base64"}),
4105        )
4106        .await;
4107        assert_eq!(status, StatusCode::BAD_REQUEST);
4108
4109        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4110        let recent = stats["recent"].as_array().unwrap();
4111        assert_eq!(recent.len(), 1);
4112        assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
4113        assert_eq!(recent[0]["status"], 400);
4114        assert_eq!(
4115            recent[0]["prompt_tokens"], 0,
4116            "a rejected call embedded nothing"
4117        );
4118    }
4119
4120    /// Attribution: which key served a request, and what the caller
4121    /// says it is. The key itself must never appear.
4122    #[tokio::test]
4123    async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
4124        let app = test_app();
4125        let key = "sk-monitor-secret";
4126        let (status, _) = post_json_with_headers(
4127            &app,
4128            "/v1/chat/completions",
4129            serde_json::json!({
4130                "model": "x",
4131                "messages": [{"role": "user", "content": "hi"}],
4132                "max_tokens": 2
4133            }),
4134            &[
4135                ("authorization", &format!("Bearer {key}")),
4136                ("x-ferrox-client", "ferrox-studio"),
4137            ],
4138        )
4139        .await;
4140        assert_eq!(status, StatusCode::OK);
4141
4142        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4143        let row = stats["recent"].as_array().unwrap()[0].clone();
4144        let fingerprint = row["via_api_key"]
4145            .as_str()
4146            .expect("the row names the key that served it")
4147            .to_string();
4148        assert_eq!(fingerprint, attribution::key_fingerprint(key));
4149        assert!(!fingerprint.contains(key));
4150        assert!(
4151            !serde_json::to_string(&stats).unwrap().contains(key),
4152            "the stats payload must not carry the key in any form"
4153        );
4154        assert_eq!(row["client"], "ferrox-studio");
4155    }
4156
4157    /// Two different keys are two different callers, and no key at all
4158    /// is a third answer -- not a copy of either.
4159    #[tokio::test]
4160    async fn different_keys_are_different_callers_and_no_key_is_null() {
4161        let app = test_app();
4162        let body = serde_json::json!({
4163            "model": "x",
4164            "messages": [{"role": "user", "content": "hi"}],
4165            "max_tokens": 1
4166        });
4167        for headers in [
4168            vec![("authorization", "Bearer key-one")],
4169            vec![("authorization", "Bearer key-two")],
4170            vec![],
4171        ] {
4172            let (status, _) =
4173                post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
4174            assert_eq!(status, StatusCode::OK);
4175        }
4176
4177        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4178        let recent = stats["recent"].as_array().unwrap();
4179        assert_eq!(recent.len(), 3);
4180        let one = recent[0]["via_api_key"].as_str().unwrap();
4181        let two = recent[1]["via_api_key"].as_str().unwrap();
4182        assert_ne!(one, two, "two keys must not collapse into one caller");
4183        assert!(
4184            recent[2]["via_api_key"].is_null(),
4185            "an unauthenticated call is null, not a fingerprint of nothing"
4186        );
4187        assert!(recent[2]["client"].is_null());
4188    }
4189
4190    /// The row names the model that SERVED the request. `req.model` is
4191    /// ignored by this server -- it decodes against whatever is loaded
4192    /// -- so echoing that string back would make the log agree with the
4193    /// caller's belief instead of with what happened.
4194    #[tokio::test]
4195    async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
4196        let state = Arc::new(test_state(
4197            named_test_model("really-loaded", 256),
4198            ResponseCache::new(4, Duration::from_secs(60)),
4199        ));
4200        let app = test_app_with_state(Arc::clone(&state));
4201
4202        let (status, _) = post_json_uri(
4203            &app,
4204            "/v1/chat/completions",
4205            serde_json::json!({
4206                "model": "gpt-4-turbo-that-is-not-here",
4207                "messages": [{"role": "user", "content": "hi"}],
4208                "max_tokens": 2
4209            }),
4210        )
4211        .await;
4212        assert_eq!(status, StatusCode::OK);
4213
4214        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4215        assert_eq!(stats["recent"][0]["model"], "really-loaded");
4216
4217        // Nothing loaded: nothing served it, and the row says so rather
4218        // than repeating what the request asked for.
4219        state.swap_active(None);
4220        let (status, _) = post_json_uri(
4221            &app,
4222            "/v1/chat/completions",
4223            serde_json::json!({
4224                "model": "gpt-4-turbo-that-is-not-here",
4225                "messages": [{"role": "user", "content": "hi"}]
4226            }),
4227        )
4228        .await;
4229        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
4230        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4231        let recent = stats["recent"].as_array().unwrap();
4232        assert!(recent[recent.len() - 1]["model"].is_null());
4233    }
4234
4235    /// A streamed request names its model too, and names the handle it
4236    /// decoded against rather than whatever a swap made current while it
4237    /// was running.
4238    #[tokio::test]
4239    async fn a_streamed_row_names_the_model_it_decoded_against() {
4240        let state = Arc::new(test_state(
4241            named_test_model("model-before", 256),
4242            ResponseCache::new(4, Duration::from_secs(60)),
4243        ));
4244        let app = test_app_with_state(Arc::clone(&state));
4245        let _ = post_sse_raw(&app, resumable_request()).await;
4246        // The stream has finished; a swap now must not rewrite history.
4247        active_model(&state, "model-after");
4248
4249        let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4250        assert_eq!(stats["recent"][0]["model"], "model-before");
4251    }
4252
4253    /// The queue gauge reports a queue that exists or says there is
4254    /// none. `0` would claim an empty queue was measured.
4255    #[tokio::test]
4256    async fn the_queue_gauge_is_null_when_nothing_can_queue() {
4257        let app = test_app();
4258        let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4259        assert_eq!(status, StatusCode::OK);
4260        assert!(
4261            stats["queue_depth"].is_null(),
4262            "without continuous batching nothing queues, so there is nothing to measure"
4263        );
4264        assert!(stats["queue_rejected_total"].is_null());
4265        assert_eq!(
4266            stats["generating_now"], 0,
4267            "work in progress is measured and really is zero here"
4268        );
4269    }
4270
4271    /// The raw SSE body, so the tests below can assert on the `id:` and
4272    /// `retry:` fields themselves rather than only on the JSON inside
4273    /// `data:`. Those two fields are the whole of the replay contract
4274    /// on the wire.
4275    async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
4276        use http_body_util::BodyExt;
4277        use tower::ServiceExt;
4278
4279        let response = app
4280            .clone()
4281            .oneshot(
4282                axum::http::Request::builder()
4283                    .method("POST")
4284                    .uri("/v1/chat/completions")
4285                    .header("content-type", "application/json")
4286                    .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
4287                    .unwrap(),
4288            )
4289            .await
4290            .unwrap();
4291        let bytes = response.into_body().collect().await.unwrap().to_bytes();
4292        String::from_utf8(bytes.to_vec()).unwrap()
4293    }
4294
4295    async fn get_json_with_headers(
4296        app: &Router,
4297        uri: &str,
4298        headers: &[(&str, &str)],
4299    ) -> (StatusCode, serde_json::Value) {
4300        use http_body_util::BodyExt;
4301        use tower::ServiceExt;
4302
4303        let mut builder = axum::http::Request::builder().method("GET").uri(uri);
4304        for (name, value) in headers {
4305            builder = builder.header(*name, *value);
4306        }
4307        let response = app
4308            .clone()
4309            .oneshot(builder.body(axum::body::Body::empty()).unwrap())
4310            .await
4311            .unwrap();
4312        let status = response.status();
4313        let bytes = response.into_body().collect().await.unwrap().to_bytes();
4314        (
4315            status,
4316            serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
4317        )
4318    }
4319
4320    fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
4321        body.lines()
4322            .filter_map(|line| line.strip_prefix(field))
4323            .map(str::trim)
4324            .collect()
4325    }
4326
4327    fn resumable_request() -> serde_json::Value {
4328        serde_json::json!({
4329            "model": "m",
4330            "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
4331            "max_tokens": 4,
4332            "temperature": 0,
4333            "stream": true,
4334            "stream_resumable": true,
4335        })
4336    }
4337
4338    /// The wire half of the replay contract: every event is numbered,
4339    /// the numbers are qualified by the request so a `Last-Event-ID`
4340    /// cannot be mistaken for a position in another stream, and the
4341    /// reconnect delay is stated once.
4342    #[tokio::test]
4343    async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
4344        let app = test_app();
4345        let body = post_sse_raw(&app, resumable_request()).await;
4346
4347        let request_id = body
4348            .lines()
4349            .find_map(|l| l.strip_prefix("data: "))
4350            .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
4351            .and_then(|v| v["request_id"].as_str().map(str::to_string))
4352            .expect("the first chunk names the request");
4353
4354        let ids = sse_field(&body, "id:");
4355        let datas = sse_field(&body, "data:");
4356        assert_eq!(
4357            ids.len(),
4358            datas.len(),
4359            "every event carries an id, or a reconnect cannot name where it stopped"
4360        );
4361        for (i, id) in ids.iter().enumerate() {
4362            assert_eq!(*id, format!("{request_id}:{i}"));
4363        }
4364        let retries = sse_field(&body, "retry:");
4365        assert_eq!(
4366            retries.len(),
4367            1,
4368            "the reconnect delay is stated once, not on every event"
4369        );
4370        assert_eq!(retries[0], "1500");
4371        assert!(
4372            body.contains("data: [DONE]"),
4373            "the end of stream is still stated"
4374        );
4375    }
4376
4377    /// The refusal this feature was written around: an `id:` with no
4378    /// replay buffer behind it tells a client it may reconnect into
4379    /// something that does not exist.
4380    #[tokio::test]
4381    async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
4382        let app = test_app();
4383        let mut request = resumable_request();
4384        request["stream_resumable"] = serde_json::json!(false);
4385        let body = post_sse_raw(&app, request).await;
4386        assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
4387        assert!(
4388            sse_field(&body, "id:").is_empty(),
4389            "an id promises a replay this stream cannot serve"
4390        );
4391        assert!(sse_field(&body, "retry:").is_empty());
4392    }
4393
4394    /// The polling fallback, which is the answer to the proxy that
4395    /// buffers `text/event-stream`: the same events, over a short JSON
4396    /// response nothing can hold back.
4397    #[tokio::test]
4398    async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
4399        let app = test_app();
4400        let body = post_sse_raw(&app, resumable_request()).await;
4401        let request_id = sse_field(&body, "id:")[0]
4402            .rsplit_once(':')
4403            .unwrap()
4404            .0
4405            .to_string();
4406        let streamed: Vec<String> = sse_field(&body, "data:")
4407            .iter()
4408            .map(|d| d.to_string())
4409            .collect();
4410
4411        let (status, polled) = get_json(
4412            &app,
4413            &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
4414        )
4415        .await;
4416        assert_eq!(status, StatusCode::OK);
4417        let events: Vec<String> = polled["events"]
4418            .as_array()
4419            .unwrap()
4420            .iter()
4421            .map(|e| e["data"].as_str().unwrap().to_string())
4422            .collect();
4423        assert_eq!(
4424            events, streamed,
4425            "the fallback must deliver the same answer, not a re-run of it"
4426        );
4427        assert_eq!(polled["request_id"], request_id);
4428        assert_eq!(
4429            polled["done"], false,
4430            "events were still being handed out, so the client must ask again"
4431        );
4432
4433        // Drained: only now is it done, so a client that stops on
4434        // `done` never discards events it was not given.
4435        let next = polled["next_index"].as_u64().unwrap();
4436        let (_, drained) = get_json(
4437            &app,
4438            &format!(
4439                "{}?from={next}",
4440                ferrox_api::routes::v1_stream_poll(&request_id)
4441            ),
4442        )
4443        .await;
4444        assert_eq!(drained["done"], true);
4445        assert_eq!(drained["events"].as_array().unwrap().len(), 0);
4446    }
4447
4448    /// A resume returns what was missed and not what was already
4449    /// rendered -- repeating delivered tokens would make replay worse
4450    /// than starting over.
4451    #[tokio::test]
4452    async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
4453        let app = test_app();
4454        let body = post_sse_raw(&app, resumable_request()).await;
4455        let ids = sse_field(&body, "id:");
4456        let datas: Vec<String> = sse_field(&body, "data:")
4457            .iter()
4458            .map(|d| d.to_string())
4459            .collect();
4460        assert!(
4461            ids.len() >= 3,
4462            "need a few events to resume into the middle"
4463        );
4464        let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
4465
4466        let (status, resumed) = get_json_with_headers(
4467            &app,
4468            &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
4469            &[],
4470        )
4471        .await;
4472        assert_eq!(status, StatusCode::OK);
4473        assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
4474
4475        // Now from the middle, the way a reconnect would.
4476        let (_, tail) = get_json(
4477            &app,
4478            &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
4479        )
4480        .await;
4481        let tail_events: Vec<String> = tail["events"]
4482            .as_array()
4483            .unwrap()
4484            .iter()
4485            .map(|e| e["data"].as_str().unwrap().to_string())
4486            .collect();
4487        assert_eq!(tail_events, datas[2..].to_vec());
4488    }
4489
4490    /// Reconnecting over SSE picks up where the last id left off, with
4491    /// the ids still attached so a second drop can be resumed too.
4492    #[tokio::test]
4493    async fn an_sse_reconnect_resumes_from_the_last_event_id() {
4494        use http_body_util::BodyExt;
4495        use tower::ServiceExt;
4496
4497        let app = test_app();
4498        let body = post_sse_raw(&app, resumable_request()).await;
4499        let ids = sse_field(&body, "id:");
4500        let datas: Vec<String> = sse_field(&body, "data:")
4501            .iter()
4502            .map(|d| d.to_string())
4503            .collect();
4504        let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
4505
4506        let response = app
4507            .clone()
4508            .oneshot(
4509                axum::http::Request::builder()
4510                    .method("GET")
4511                    .uri(ferrox_api::routes::v1_stream(&request_id))
4512                    .header("last-event-id", format!("{request_id}:0"))
4513                    .body(axum::body::Body::empty())
4514                    .unwrap(),
4515            )
4516            .await
4517            .unwrap();
4518        assert_eq!(response.status(), StatusCode::OK);
4519        assert_eq!(
4520            response
4521                .headers()
4522                .get("x-accel-buffering")
4523                .and_then(|v| v.to_str().ok()),
4524            Some("no"),
4525            "the reconnect needs the same anti-buffering header as the stream"
4526        );
4527        let bytes = response.into_body().collect().await.unwrap().to_bytes();
4528        let resumed = String::from_utf8(bytes.to_vec()).unwrap();
4529        assert_eq!(
4530            sse_field(&resumed, "data:")
4531                .iter()
4532                .map(|d| d.to_string())
4533                .collect::<Vec<_>>(),
4534            datas[1..].to_vec()
4535        );
4536        assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
4537    }
4538
4539    /// A `Last-Event-ID` from another stream is refused rather than
4540    /// rounded down to zero: replaying a whole different answer would
4541    /// be a silent, confident lie.
4542    #[tokio::test]
4543    async fn a_last_event_id_from_another_stream_is_refused() {
4544        let app = test_app();
4545        let body = post_sse_raw(&app, resumable_request()).await;
4546        let request_id = sse_field(&body, "id:")[0]
4547            .rsplit_once(':')
4548            .unwrap()
4549            .0
4550            .to_string();
4551
4552        let (status, err) = get_json_with_headers(
4553            &app,
4554            &ferrox_api::routes::v1_stream(&request_id),
4555            &[("last-event-id", "chatcmpl-someone-else:3")],
4556        )
4557        .await;
4558        assert_eq!(status, StatusCode::BAD_REQUEST);
4559        assert_eq!(err["error"]["code"], "bad_last_event_id");
4560    }
4561
4562    /// A stream that was never resumable, or has been forgotten, is a
4563    /// 404 that says which -- not an empty stream that reads as an
4564    /// answer with no tokens in it.
4565    #[tokio::test]
4566    async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
4567        let app = test_app();
4568        let mut request = resumable_request();
4569        request["stream_resumable"] = serde_json::json!(false);
4570        let body = post_sse_raw(&app, request).await;
4571        let request_id = body
4572            .lines()
4573            .find_map(|l| l.strip_prefix("data: "))
4574            .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
4575            .and_then(|v| v["request_id"].as_str().map(str::to_string))
4576            .unwrap();
4577
4578        let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
4579        assert_eq!(status, StatusCode::NOT_FOUND);
4580        assert_eq!(err["error"]["code"], "stream_not_found");
4581        assert!(err["error"]["message"]
4582            .as_str()
4583            .unwrap()
4584            .contains("stream_resumable"));
4585    }
4586
4587    /// The published template and the router's pattern must describe
4588    /// the same path, or a client built from `ferrox_api::routes` asks
4589    /// for something this server does not serve.
4590    #[test]
4591    fn the_axum_stream_patterns_match_the_published_templates() {
4592        assert_eq!(
4593            resume_route(ferrox_api::routes::V1_STREAM),
4594            "/v1/stream/:request_id"
4595        );
4596        assert_eq!(
4597            resume_route(ferrox_api::routes::V1_STREAM_POLL),
4598            "/v1/stream/:request_id/poll"
4599        );
4600        assert_eq!(
4601            ferrox_api::routes::v1_stream("abc"),
4602            resume_route(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
4603        );
4604    }
4605
4606    /// An empty task list is a list, not a missing key -- the UI renders
4607    /// "no jobs" from it rather than from an error.
4608    #[tokio::test]
4609    async fn the_task_list_starts_empty_rather_than_absent() {
4610        let app = test_app();
4611        let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
4612        assert_eq!(status, StatusCode::OK);
4613        assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
4614    }
4615
4616    async fn post_json_uri(
4617        app: &Router,
4618        uri: &str,
4619        body: serde_json::Value,
4620    ) -> (StatusCode, serde_json::Value) {
4621        use http_body_util::BodyExt;
4622        use tower::ServiceExt;
4623
4624        let response = app
4625            .clone()
4626            .oneshot(
4627                axum::http::Request::builder()
4628                    .method("POST")
4629                    .uri(uri)
4630                    .header("content-type", "application/json")
4631                    .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
4632                    .unwrap(),
4633            )
4634            .await
4635            .unwrap();
4636        let status = response.status();
4637        let bytes = response.into_body().collect().await.unwrap().to_bytes();
4638        let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
4639        (status, json)
4640    }
4641
4642    async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
4643        post_json_uri(app, "/v1/chat/completions", body).await.1
4644    }
4645
4646    /// Cancelling an id that is not generating must not answer `200`.
4647    /// A UI told "ok" for an already-finished request would report that
4648    /// it stopped work it did not stop, and the two outcomes are the
4649    /// only thing this endpoint exists to distinguish.
4650    #[tokio::test]
4651    async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
4652        let app = test_app();
4653        let (status, body) = post_json_uri(
4654            &app,
4655            ferrox_api::routes::V1_CANCEL,
4656            serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
4657        )
4658        .await;
4659        assert_eq!(status, StatusCode::NOT_FOUND);
4660        assert_eq!(body["cancelled"], serde_json::json!(false));
4661        assert_eq!(body["request_id"], "chatcmpl-never-issued");
4662        assert!(
4663            body["detail"].as_str().is_some_and(|d| !d.is_empty()),
4664            "the verdict must carry a human reason: {body}"
4665        );
4666    }
4667
4668    /// The endpoint reaches the registry the streaming path registers
4669    /// into -- not a second, parallel one. Registered by hand here
4670    /// because a `oneshot` router cannot hold a stream open.
4671    #[tokio::test]
4672    async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
4673        let state = Arc::new(test_state(
4674            test_model_full_byte_vocab(),
4675            ResponseCache::new(1000, Duration::from_secs(3600)),
4676        ));
4677        let app = test_app_with_state(Arc::clone(&state));
4678        let (token, _guard) = state.cancels.register("chatcmpl-live");
4679
4680        let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4681        assert_eq!(status, StatusCode::OK);
4682        assert_eq!(before["generating_now"], serde_json::json!(1));
4683
4684        let (status, body) = post_json_uri(
4685            &app,
4686            ferrox_api::routes::V1_CANCEL,
4687            serde_json::json!({ "request_id": "chatcmpl-live" }),
4688        )
4689        .await;
4690        assert_eq!(status, StatusCode::OK);
4691        assert_eq!(body["cancelled"], serde_json::json!(true));
4692        assert!(
4693            token.is_cancelled(),
4694            "the endpoint answered ok without setting the flag the decode loop reads"
4695        );
4696    }
4697
4698    #[tokio::test]
4699    async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
4700        let app = test_app();
4701        let (status, tok) =
4702            post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
4703        assert_eq!(status, StatusCode::OK);
4704        let tokens = tok["tokens"].as_array().unwrap();
4705        assert_eq!(tok["count"], tokens.len());
4706        assert!(!tokens.is_empty());
4707
4708        let (status, detok) = post_json_uri(
4709            &app,
4710            "/v1/detokenize",
4711            serde_json::json!({ "tokens": tokens }),
4712        )
4713        .await;
4714        assert_eq!(status, StatusCode::OK);
4715        assert_eq!(detok["text"], "Hi");
4716
4717        let (status, emb) = post_json_uri(
4718            &app,
4719            "/v1/embeddings",
4720            serde_json::json!({
4721                "input": "Hi",
4722                "embedding_type": "mean"
4723            }),
4724        )
4725        .await;
4726        assert_eq!(status, StatusCode::OK);
4727        let vec = emb["data"][0]["embedding"].as_array().unwrap();
4728        assert!(!vec.is_empty());
4729        assert!(vec.iter().all(|v| v.as_f64().is_some()));
4730    }
4731
4732    /// The /metrics endpoint must expose the bounded expert cache's
4733    /// counters when the model streams routed experts, and the
4734    /// counters must reflect real decode activity (a forward pass
4735    /// through store-backed MoE layers produces misses/hits).
4736    #[tokio::test]
4737    async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
4738        use http_body_util::BodyExt;
4739        use tower::ServiceExt;
4740
4741        let fixture = concat!(
4742            "../ferrox-models/tests/fixtures/",
4743            "ferrox_real_moe_test.gguf"
4744        );
4745        let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
4746        let decoder = Decoder::from_gguf_with_expert_cache(
4747            &fixture,
4748            ferrox_models::config::test_moe_fixture(),
4749            Some(1024 * 1024),
4750        )
4751        .expect("MoE fixture must load store-backed");
4752
4753        // Drive one real forward pass so the store sees decode
4754        // activity (the fixture's tiny vocab can't survive the HTTP
4755        // path's template text, so decode directly).
4756        let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
4757            .layers
4758            .iter()
4759            .map(|_| {
4760                ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
4761            })
4762            .collect();
4763        decoder.forward_token(1, 0, &mut caches);
4764
4765        let model = Model::Gguf(GgufModel {
4766            decoder: Arc::new(decoder),
4767            tokenizer: Arc::new(ServerTokenizer::Byte),
4768            stop_tokens: StopTokens::default(),
4769            bos_id: None,
4770            is_synthetic: false,
4771            chat_template: chat_template::ChatTemplate::Plain,
4772        });
4773        let state = Arc::new(test_state(
4774            model,
4775            ResponseCache::new(16, Duration::from_secs(60)),
4776        ));
4777        let app = Router::new()
4778            .route("/metrics", axum::routing::get(metrics))
4779            .route("/v1/chat/completions", post(chat_completions))
4780            .with_state(state);
4781
4782        let fetch_metrics = |app: Router| async move {
4783            let resp = app
4784                .oneshot(
4785                    axum::http::Request::builder()
4786                        .method("GET")
4787                        .uri("/metrics")
4788                        .body(axum::body::Body::empty())
4789                        .unwrap(),
4790                )
4791                .await
4792                .unwrap();
4793            let bytes = resp.into_body().collect().await.unwrap().to_bytes();
4794            String::from_utf8(bytes.to_vec()).unwrap()
4795        };
4796
4797        let after = fetch_metrics(app.clone()).await;
4798        assert!(
4799            after.contains("ferrox_expert_cache_misses_total"),
4800            "streaming model must expose expert-cache metrics: {after}"
4801        );
4802        let misses: u64 = after
4803            .lines()
4804            .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
4805            .and_then(|l| l.split_whitespace().nth(1))
4806            .and_then(|v| v.parse().ok())
4807            .expect("misses metric line must parse");
4808        assert!(
4809            misses > 0,
4810            "decode must have read experts through the store: {after}"
4811        );
4812    }
4813
4814    fn weather_tool() -> serde_json::Value {
4815        serde_json::json!({
4816            "type": "function",
4817            "function": {
4818                "name": "get_weather",
4819                "description": "Get the current weather for a location.",
4820                "parameters": {
4821                    "type": "object",
4822                    "properties": {"location": {"type": "string"}},
4823                    "required": ["location"]
4824                }
4825            }
4826        })
4827    }
4828
4829    fn weather_tool_def() -> ToolDef {
4830        ToolDef {
4831            kind: "function".to_string(),
4832            function: ToolFunctionDef {
4833                name: "get_weather".to_string(),
4834                description: Some("Get the current weather for a location.".to_string()),
4835                parameters: Some(serde_json::json!({
4836                    "type": "object",
4837                    "properties": {"location": {"type": "string"}},
4838                    "required": ["location"]
4839                })),
4840            },
4841        }
4842    }
4843
4844    #[test]
4845    fn tool_preamble_mentions_every_tool_name_and_description() {
4846        let preamble = tool_preamble(&[weather_tool_def()]);
4847        assert!(preamble.contains("get_weather"));
4848        assert!(preamble.contains("Get the current weather for a location."));
4849        assert!(preamble.contains("<tool_call>"));
4850        assert!(preamble.contains("</tool_call>"));
4851    }
4852
4853    #[test]
4854    fn extract_tool_call_parses_a_real_marker() {
4855        let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
4856        let (name, arguments) = extract_tool_call(text).expect("must find the marker");
4857        assert_eq!(name, "get_weather");
4858        let parsed: serde_json::Value = serde_json::from_str(&arguments).unwrap();
4859        assert_eq!(parsed["location"], "Paris");
4860    }
4861
4862    #[test]
4863    fn extract_tool_call_returns_none_when_no_marker_present() {
4864        assert_eq!(
4865            extract_tool_call("just a plain answer, no markers here"),
4866            None
4867        );
4868    }
4869
4870    #[test]
4871    fn extract_tool_call_returns_none_on_malformed_json_inside_the_marker() {
4872        assert_eq!(
4873            extract_tool_call("<tool_call>not valid json at all</tool_call>"),
4874            None
4875        );
4876    }
4877
4878    #[test]
4879    fn extract_tool_call_defaults_to_empty_arguments_when_the_field_is_absent() {
4880        let (name, arguments) =
4881            extract_tool_call("<tool_call>{\"name\": \"ping\"}</tool_call>").unwrap();
4882        assert_eq!(name, "ping");
4883        assert_eq!(arguments, "{}");
4884    }
4885
4886    #[test]
4887    fn build_response_message_promotes_a_real_tool_call_to_tool_calls() {
4888        let (message, finish) = build_response_message(
4889            "<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Rome\"}}</tool_call>"
4890                .to_string(),
4891            true,
4892            "stop",
4893        );
4894        assert_eq!(finish, "tool_calls");
4895        assert!(message.content.is_none());
4896        let calls = message.tool_calls.expect("must carry a tool call");
4897        assert_eq!(calls.len(), 1);
4898        assert_eq!(calls[0].function.name, "get_weather");
4899    }
4900
4901    #[test]
4902    fn build_response_message_falls_back_to_plain_text_when_tools_are_inactive() {
4903        let (message, finish) = build_response_message(
4904            "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
4905            false,
4906            "stop",
4907        );
4908        assert_eq!(finish, "stop");
4909        assert!(message.tool_calls.is_none());
4910        assert!(message.content.is_some());
4911    }
4912
4913    #[test]
4914    fn build_response_message_falls_back_to_plain_text_when_no_marker_is_present() {
4915        let (message, finish) = build_response_message("just an answer".to_string(), true, "stop");
4916        assert_eq!(finish, "stop");
4917        assert!(message.tool_calls.is_none());
4918        assert_eq!(message.content.as_deref(), Some("just an answer"));
4919    }
4920
4921    /// Zero-regression proof: an ordinary request with no `tools`/
4922    /// `session_id` produces the plain response shape -- `content` a
4923    /// string, no `tool_calls` field -- with an honest finish reason:
4924    /// this 4-token greedy request truncates at `max_tokens`, so
4925    /// `finish_reason` must be "length" (an earlier version hardcoded
4926    /// "stop" for every non-streaming response), and `usage` counts
4927    /// exactly the generated tokens.
4928    #[tokio::test]
4929    async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
4930        let app = test_app();
4931        let body = serde_json::json!({
4932            "model": "m",
4933            "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
4934            "max_tokens": 4,
4935            "temperature": 0,
4936        });
4937        let resp = post_json(&app, body).await;
4938        let message = &resp["choices"][0]["message"];
4939        assert!(message["content"].is_string());
4940        assert!(message.get("tool_calls").is_none());
4941        assert_eq!(resp["choices"][0]["finish_reason"], "length");
4942        assert_eq!(resp["usage"]["completion_tokens"], 4);
4943        assert_eq!(
4944            resp["usage"]["total_tokens"],
4945            resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
4946        );
4947    }
4948
4949    async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
4950        use http_body_util::BodyExt;
4951        use tower::ServiceExt;
4952
4953        let response = app
4954            .clone()
4955            .oneshot(
4956                axum::http::Request::builder()
4957                    .method("GET")
4958                    .uri(uri)
4959                    .body(axum::body::Body::empty())
4960                    .unwrap(),
4961            )
4962            .await
4963            .unwrap();
4964        let status = response.status();
4965        let bytes = response.into_body().collect().await.unwrap().to_bytes();
4966        (status, serde_json::from_slice(&bytes).unwrap())
4967    }
4968
4969    #[tokio::test]
4970    async fn health_answers_a_capability_handshake_not_a_boolean() {
4971        let app = test_app();
4972        let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
4973        assert_eq!(status, StatusCode::OK);
4974
4975        let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
4976        assert_eq!(health.state, ferrox_api::HealthState::Ready);
4977        assert!(health.pid > 0);
4978        assert!(health.server_time_unix_ms > 0);
4979        // Nothing has been served yet: the field is absent rather than
4980        // claiming a request happened at time zero.
4981        assert_eq!(health.last_request_age_seconds, None);
4982
4983        // Every control the UI might grey out has a code it can switch
4984        // on and a sentence it can show.
4985        for id in [
4986            ferrox_api::health::capability::CPU,
4987            ferrox_api::health::capability::METAL,
4988            ferrox_api::health::capability::CUDA,
4989            ferrox_api::health::capability::REAL_WEIGHTS,
4990            ferrox_api::health::capability::CONTINUOUS_BATCHING,
4991        ] {
4992            let cap = health
4993                .capability(id)
4994                .unwrap_or_else(|| panic!("{id} missing"));
4995            assert!(!cap.reason.is_empty(), "{cap:?}");
4996            assert!(!cap.detail.is_empty(), "{cap:?}");
4997        }
4998        // The test app serves synthetic random weights, and health must
4999        // say so: a UI that presents noise as a model invites a bug
5000        // report about "quality".
5001        let weights = health
5002            .capability(ferrox_api::health::capability::REAL_WEIGHTS)
5003            .unwrap();
5004        assert!(!weights.available);
5005        assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
5006        assert!(health.model.as_ref().unwrap().synthetic_weights);
5007    }
5008
5009    #[tokio::test]
5010    async fn health_vouches_for_liveness_after_a_request_has_been_served() {
5011        let app = test_app();
5012        let _ = post_json(
5013            &app,
5014            serde_json::json!({
5015                "model": "m",
5016                "messages": [{"role": "user", "content": "\u{1}"}],
5017                "max_tokens": 1,
5018                "temperature": 0,
5019            }),
5020        )
5021        .await;
5022        let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5023        let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
5024        let age = health
5025            .last_request_age_seconds
5026            .expect("a served request is evidence of liveness");
5027        assert!((0.0..5.0).contains(&age), "implausible age {age}");
5028    }
5029
5030    /// Every `data:` payload of an SSE response body, `[DONE]` excluded.
5031    async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
5032        use http_body_util::BodyExt;
5033        use tower::ServiceExt;
5034
5035        let response = app
5036            .clone()
5037            .oneshot(
5038                axum::http::Request::builder()
5039                    .method("POST")
5040                    .uri("/v1/chat/completions")
5041                    .header("content-type", "application/json")
5042                    .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5043                    .unwrap(),
5044            )
5045            .await
5046            .unwrap();
5047        let bytes = response.into_body().collect().await.unwrap().to_bytes();
5048        String::from_utf8(bytes.to_vec())
5049            .unwrap()
5050            .lines()
5051            .filter_map(|line| line.strip_prefix("data: "))
5052            .filter(|payload| *payload != "[DONE]")
5053            .map(|payload| serde_json::from_str(payload).unwrap())
5054            .collect()
5055    }
5056
5057    #[tokio::test]
5058    async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
5059        let app = test_app();
5060        let chunks = post_sse_chunks(
5061            &app,
5062            serde_json::json!({
5063                "model": "m",
5064                "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5065                "max_tokens": 4,
5066                "temperature": 0,
5067                "stream": true,
5068            }),
5069        )
5070        .await;
5071
5072        assert!(!chunks.is_empty());
5073        let request_id = chunks[0]["request_id"]
5074            .as_str()
5075            .expect("the first chunk names the request")
5076            .to_string();
5077        assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
5078        // Once, and before any content: a client that reads the id from
5079        // chunk zero never has to correlate by heuristic.
5080        for (i, chunk) in chunks.iter().enumerate().skip(1) {
5081            assert!(
5082                chunk.get("request_id").is_none(),
5083                "chunk {i} repeats request_id"
5084            );
5085        }
5086        // Every chunk of one stream carries the same `id`, and it is
5087        // that request id -- not a shared constant.
5088        for chunk in &chunks {
5089            assert_eq!(chunk["id"], serde_json::json!(request_id));
5090        }
5091
5092        let other = post_sse_chunks(
5093            &app,
5094            serde_json::json!({
5095                "model": "m",
5096                "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5097                "max_tokens": 4,
5098                "temperature": 0,
5099                "stream": true,
5100            }),
5101        )
5102        .await;
5103        assert_ne!(
5104            other[0]["request_id"].as_str().unwrap(),
5105            request_id,
5106            "two concurrent chats must not share an id"
5107        );
5108    }
5109
5110    #[tokio::test]
5111    async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
5112        let app = test_app();
5113        let resp = post_json(
5114            &app,
5115            serde_json::json!({
5116                "model": "m",
5117                "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5118                "max_tokens": 2,
5119                "temperature": 0,
5120            }),
5121        )
5122        .await;
5123        assert_eq!(resp["id"], resp["request_id"]);
5124        assert!(resp["request_id"]
5125            .as_str()
5126            .unwrap()
5127            .starts_with("chatcmpl-"));
5128    }
5129
5130    /// The whole point of server-reported timings: a client can tell
5131    /// prefill from decode without a stopwatch (see `ferrox_api::usage`).
5132    #[tokio::test]
5133    async fn usage_carries_separate_prefill_and_decode_timings() {
5134        let app = test_app();
5135        let resp = post_json(
5136            &app,
5137            serde_json::json!({
5138                "model": "m",
5139                "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5140                "max_tokens": 4,
5141                "temperature": 0,
5142            }),
5143        )
5144        .await;
5145        let usage = &resp["usage"];
5146        assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
5147        assert!(usage["generation_duration_ms"].is_number(), "{usage}");
5148        assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
5149        assert!(usage["predicted_per_second"].is_number(), "{usage}");
5150        // No prefix cache in this app: the field must be absent, not 0.
5151        assert!(usage.get("cached_tokens").is_none(), "{usage}");
5152    }
5153
5154    /// A real, deterministic small model with random weights will not
5155    /// spontaneously produce a `<tool_call>{...}</tool_call>` marker
5156    /// (whether a real deployed model does is a property of that
5157    /// model, not of ferrox's plumbing) -- so the real, testable
5158    /// end-to-end property here is that a `tools`-bearing request
5159    /// whose output does NOT contain the marker falls through cleanly
5160    /// to an ordinary text response instead of erroring or panicking.
5161    #[tokio::test]
5162    async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
5163        let app = test_app();
5164        let body = serde_json::json!({
5165            "model": "m",
5166            "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5167            "max_tokens": 4,
5168            "temperature": 0,
5169            "tools": [weather_tool()],
5170        });
5171        let resp = post_json(&app, body).await;
5172        let message = &resp["choices"][0]["message"];
5173        assert!(
5174            message["content"].is_string(),
5175            "must fall back to plain content when no real tool-call marker is present: {resp:?}"
5176        );
5177        assert!(message.get("tool_calls").is_none());
5178        // Truncated at max_tokens, so the honest finish reason is
5179        // "length" -- the point here is only that it is NOT
5180        // "tool_calls".
5181        assert_eq!(resp["choices"][0]["finish_reason"], "length");
5182    }
5183
5184    /// A whole-response cache hit must be indistinguishable from
5185    /// recomputing: same content, same (honest) finish_reason, same
5186    /// usage counts -- only the `ferrox_cache` marker may differ.
5187    #[tokio::test]
5188    async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
5189        let app = test_app();
5190        let body = serde_json::json!({
5191            "model": "m",
5192            "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
5193            "max_tokens": 3,
5194            "temperature": 0,
5195        });
5196        let first = post_json(&app, body.clone()).await;
5197        assert_eq!(first["ferrox_cache"], "miss");
5198        let second = post_json(&app, body).await;
5199        assert_eq!(second["ferrox_cache"], "hit");
5200        assert_eq!(
5201            first["choices"][0]["message"]["content"],
5202            second["choices"][0]["message"]["content"]
5203        );
5204        assert_eq!(
5205            first["choices"][0]["finish_reason"],
5206            second["choices"][0]["finish_reason"]
5207        );
5208        assert_eq!(first["usage"], second["usage"]);
5209        assert_eq!(second["usage"]["completion_tokens"], 3);
5210    }
5211
5212    /// The real proof for session reuse:
5213    /// a two-request session where the second request sends only its
5214    /// new message must produce exactly the same output as manually
5215    /// resending the full history (built from the *real* first reply,
5216    /// not an assumed one) with no `session_id` at all.
5217    #[tokio::test]
5218    async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
5219        let session_app = test_app();
5220        let manual_app = test_app();
5221
5222        // Turn 1, via session.
5223        let turn1 = post_json(
5224            &session_app,
5225            serde_json::json!({
5226                "model": "m",
5227                "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5228                "session_id": "s1",
5229                "max_tokens": 5,
5230                "temperature": 0,
5231            }),
5232        )
5233        .await;
5234        let reply1 = turn1["choices"][0]["message"]["content"]
5235            .as_str()
5236            .unwrap()
5237            .to_string();
5238
5239        // Turn 1, manually, for comparison -- must match exactly
5240        // (trivially, since it's the literal same single-turn
5241        // request), confirming the session path's first turn isn't
5242        // doing anything different from a plain request.
5243        let manual_turn1 = post_json(
5244            &manual_app,
5245            serde_json::json!({
5246                "model": "m",
5247                "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5248                "max_tokens": 5,
5249                "temperature": 0,
5250            }),
5251        )
5252        .await;
5253        assert_eq!(
5254            manual_turn1["choices"][0]["message"]["content"]
5255                .as_str()
5256                .unwrap(),
5257            reply1
5258        );
5259
5260        // Turn 2, via session: sends ONLY the new message.
5261        let turn2 = post_json(
5262            &session_app,
5263            serde_json::json!({
5264                "model": "m",
5265                "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
5266                "session_id": "s1",
5267                "max_tokens": 5,
5268                "temperature": 0,
5269            }),
5270        )
5271        .await;
5272        let reply2 = turn2["choices"][0]["message"]["content"]
5273            .as_str()
5274            .unwrap()
5275            .to_string();
5276
5277        // Turn 2, manually: the full three-message history
5278        // reconstructed using the REAL reply1 text, with no
5279        // session_id -- must produce byte-identical output.
5280        let manual_turn2 = post_json(
5281            &manual_app,
5282            serde_json::json!({
5283                "model": "m",
5284                "messages": [
5285                    {"role": "user", "content": "\u{1}\u{2}\u{3}"},
5286                    {"role": "assistant", "content": reply1},
5287                    {"role": "user", "content": "\u{4}\u{5}"},
5288                ],
5289                "max_tokens": 5,
5290                "temperature": 0,
5291            }),
5292        )
5293        .await;
5294        assert_eq!(
5295            manual_turn2["choices"][0]["message"]["content"]
5296                .as_str()
5297                .unwrap(),
5298            reply2,
5299            "resuming a session must produce identical output to manually resending the full history"
5300        );
5301    }
5302
5303    /// `lock_cache` must return a usable guard even after the mutex was
5304    /// poisoned by a panic elsewhere.
5305    #[test]
5306    fn lock_cache_recovers_from_a_poisoned_mutex() {
5307        let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
5308
5309        let poison_cache = Arc::clone(&cache);
5310        let _ = std::thread::spawn(move || {
5311            let _guard = poison_cache.lock().unwrap();
5312            panic!("simulated panic while holding the lock");
5313        })
5314        .join();
5315
5316        // A plain `.lock().unwrap()` would panic here; lock_cache must not.
5317        let recovered = lock_cache(&cache);
5318        assert_eq!(recovered.stats().entries, 0);
5319    }
5320
5321    #[test]
5322    fn is_cacheable_true_for_greedy_or_seeded_requests() {
5323        let mut req_body = serde_json::json!({
5324            "model": "m",
5325            "messages": [{"role": "user", "content": "hi"}],
5326        });
5327        let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
5328        assert!(
5329            req.is_cacheable(),
5330            "default (temperature 0) must be cacheable"
5331        );
5332
5333        req_body["temperature"] = serde_json::json!(0.8);
5334        let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
5335        assert!(
5336            !req.is_cacheable(),
5337            "unseeded sampling must never be cacheable"
5338        );
5339
5340        req_body["seed"] = serde_json::json!(42);
5341        let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
5342        assert!(
5343            req.is_cacheable(),
5344            "sampling with an explicit seed is deterministic and must be cacheable"
5345        );
5346    }
5347
5348    #[test]
5349    fn stop_param_accepts_both_single_string_and_array() {
5350        let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
5351            "model": "m",
5352            "messages": [{"role": "user", "content": "hi"}],
5353            "stop": "END",
5354        }))
5355        .unwrap();
5356        assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
5357
5358        let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
5359            "model": "m",
5360            "messages": [{"role": "user", "content": "hi"}],
5361            "stop": ["A", "B"],
5362        }))
5363        .unwrap();
5364        assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
5365    }
5366
5367    #[test]
5368    fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
5369        let model = test_model();
5370        let result = run_generation(&model, "hello", &greedy_params(4), None, None, None, None);
5371        assert!(matches!(
5372            result,
5373            Err(generate::DecodeError::TokenOutOfVocab { .. })
5374        ));
5375    }
5376
5377    /// A pool that *could* serve this request but is momentarily fully
5378    /// held is the server being behind: 503, and retrying is honest
5379    /// advice because the blocks really do come back.
5380    #[test]
5381    fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
5382        let model = test_model(); // 2 layers -> 2 blocks
5383        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5384        let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
5385
5386        let holder_pool = Arc::clone(&pool);
5387        let holder = std::thread::spawn(move || {
5388            let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
5389            held.push(&[0.0], &[0.0]).unwrap(); // crosses into the second block
5390            std::thread::sleep(Duration::from_millis(200));
5391            drop(held);
5392        });
5393        std::thread::sleep(Duration::from_millis(15));
5394
5395        let config = generate::KvPoolConfig {
5396            pool,
5397            queue_wait: Duration::ZERO,
5398        };
5399        let result = run_generation(
5400            &model,
5401            &prompt,
5402            &greedy_params(4),
5403            Some(&config),
5404            None,
5405            None,
5406            None,
5407        );
5408        assert!(matches!(
5409            result,
5410            Err(generate::DecodeError::KvPoolExhausted)
5411        ));
5412
5413        let (status, _body) = decode_error_response(result.unwrap_err());
5414        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5415        holder.join().unwrap();
5416    }
5417
5418    /// The same endpoint, the same pool size, a request too big for the
5419    /// *whole* pool: a 400 rather than a 503, because an idle server
5420    /// refuses it identically and `Retry-After` would be a promise
5421    /// nothing can keep.
5422    ///
5423    /// Confirmed to FAIL when `generate`'s `pool_immovable_refusal`
5424    /// check is removed: the status comes back 503.
5425    #[test]
5426    fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
5427        let model = test_model(); // 2 layers
5428        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5429        // One block, two layers: no schedule ever serves this.
5430        let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
5431        let config = generate::KvPoolConfig {
5432            pool,
5433            queue_wait: Duration::ZERO,
5434        };
5435
5436        let result = run_generation(
5437            &model,
5438            &prompt,
5439            &greedy_params(4),
5440            Some(&config),
5441            None,
5442            None,
5443            None,
5444        );
5445        let err = result.expect_err("one block cannot hold two layers' caches");
5446        assert!(
5447            matches!(
5448                &err,
5449                generate::DecodeError::KvBudgetExceeded { binding, .. }
5450                    if *binding == ferrox_models::Ceiling::DeviceMemory.code()
5451            ),
5452            "expected an immovable device-memory refusal, got {err:?}"
5453        );
5454        let (status, _body) = decode_error_response(err);
5455        assert_eq!(status, StatusCode::BAD_REQUEST);
5456    }
5457
5458    /// A full admission queue is the server being behind, not the
5459    /// client being wrong: 503, with the wait hint in the body (and the
5460    /// `Retry-After` header stamped by `limits::retry_after`) and the
5461    /// depth and cap named so an operator can tell a retry storm from a
5462    /// single oversized request.
5463    #[test]
5464    fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
5465        let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
5466            queued: 512,
5467            cap: 512,
5468        });
5469        assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5470        assert_eq!(body["error"]["retry_after_seconds"], 1);
5471        let message = body["error"]["message"].as_str().expect("message");
5472        assert!(message.contains("512"), "{message}");
5473    }
5474
5475    #[test]
5476    fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
5477        let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
5478            token: 99,
5479            vocab_size: 32,
5480        });
5481        assert!(
5482            body["error"]["retry_after_seconds"].is_null(),
5483            "retrying a prompt this model cannot tokenize never helps"
5484        );
5485    }
5486
5487    #[test]
5488    fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
5489        let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
5490            token: 99,
5491            vocab_size: 32,
5492        });
5493        assert_eq!(status, StatusCode::BAD_REQUEST);
5494    }
5495
5496    #[test]
5497    fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
5498        let model = test_model(); // 2 layers
5499        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5500        let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
5501        let config = generate::KvPoolConfig {
5502            pool: pool.clone(),
5503            queue_wait: Duration::ZERO,
5504        };
5505
5506        let (_, finish, _usage) = run_generation(
5507            &model,
5508            &prompt,
5509            &greedy_params(4),
5510            Some(&config),
5511            None,
5512            None,
5513            None,
5514        )
5515        .unwrap();
5516        assert_eq!(finish, FinishReason::Length);
5517        assert_eq!(
5518            pool.lock().unwrap().free_blocks(),
5519            2,
5520            "a completed request must return its blocks to the pool"
5521        );
5522    }
5523
5524    /// The core concurrency claim: two requests using the *same* `Arc<Model>`
5525    /// must be able to run their (independent, per-call) KV caches
5526    /// concurrently without interfering with each other or needing any
5527    /// shared lock around the model itself.
5528    #[tokio::test]
5529    async fn concurrent_requests_against_the_same_model_do_not_interfere() {
5530        let model = Arc::new(test_model());
5531        let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5532
5533        let mut handles = Vec::new();
5534        for _ in 0..8 {
5535            let model = Arc::clone(&model);
5536            let prompt = prompt.clone();
5537            handles.push(tokio::task::spawn_blocking(move || {
5538                run_generation(&model, &prompt, &greedy_params(6), None, None, None, None).unwrap()
5539            }));
5540        }
5541
5542        let mut results = Vec::new();
5543        for h in handles {
5544            results.push(h.await.unwrap());
5545        }
5546        // Same prompt, same seed, same (greedy) sampling, same
5547        // immutable model -> every concurrent run must produce
5548        // identical output, proving no request's KV cache leaked into
5549        // another's.
5550        for r in &results[1..] {
5551            assert_eq!(r.0, results[0].0, "decoded chunks must match");
5552            assert_eq!(r.1, results[0].1, "finish reason must match");
5553            assert_eq!(
5554                r.2.prompt_tokens, results[0].2.prompt_tokens,
5555                "prompt token count must match"
5556            );
5557            assert_eq!(
5558                r.2.completion_tokens, results[0].2.completion_tokens,
5559                "completion token count must match"
5560            );
5561        }
5562    }
5563
5564    /// A real, minimal safetensors shard: JSON header (name -> real
5565    /// dtype/shape/`data_offsets`) followed by the concatenated raw
5566    /// F32 bytes -- exactly the format `ShardedSafetensors::open_index`
5567    /// parses, hand-built here rather than depending on
5568    /// `ferrox-models::kimi_loader`'s own private test helpers (not
5569    /// visible across the crate boundary).
5570    fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
5571        let mut header_entries = Vec::new();
5572        let mut data = Vec::new();
5573        for (name, shape, values) in tensors {
5574            let start = data.len();
5575            for v in values {
5576                data.extend_from_slice(&v.to_le_bytes());
5577            }
5578            let end = data.len();
5579            let shape_str = shape
5580                .iter()
5581                .map(|d| d.to_string())
5582                .collect::<Vec<_>>()
5583                .join(",");
5584            header_entries.push(format!(
5585                "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
5586            ));
5587        }
5588        let header = format!("{{{}}}", header_entries.join(","));
5589        let header_bytes = header.as_bytes();
5590        let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
5591        out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
5592        out.extend_from_slice(header_bytes);
5593        out.extend_from_slice(&data);
5594        out
5595    }
5596
5597    /// Builds a small but completely real Kimi K3 checkpoint directory
5598    /// on disk (real `model.safetensors.index.json` + shard bytes +
5599    /// `tiktoken.model`, the exact file layout `ferrox-cli`'s
5600    /// `run-kimi` command expects) and loads it through
5601    /// `model::load_kimi_checkpoint_with_config` (the same real loading
5602    /// logic `model::load()` uses for `FERROX_MODEL_PATH` pointing at a
5603    /// directory, parametrized here only so the checkpoint can be small
5604    /// -- see that function's doc comment). Shared by every test that
5605    /// needs a real, loaded `KimiLoaded` rather than duplicating this
5606    /// setup per test.
5607    fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
5608        use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
5609        use ferrox_models::kimi_loader::KimiRealHparams;
5610        use ferrox_moe::{GatingFunction, MoeLayerConfig};
5611
5612        let hidden_dim = 8;
5613        let kda_num_heads = 2;
5614        let kda_head_dim = 3;
5615        let kda_proj = kda_num_heads * kda_head_dim;
5616        let conv_kernel = 4;
5617        let dense_intermediate = 5;
5618        // One token per byte value -- enough to round-trip a simple
5619        // ASCII prompt through the real tiktoken-format vocab below,
5620        // matching `kimi_generate`'s own test convention.
5621        let vocab_size = 256;
5622        let mla_num_heads = 1;
5623        let mla_q_lora_rank = 2;
5624        let mla_kv_lora_rank = 2;
5625        let mla_qk_nope_head_dim = 2;
5626        let mla_qk_rope_head_dim = 2;
5627        let mla_v_head_dim = 2;
5628
5629        let model_cfg = ferrox_models::ModelConfig {
5630            name: "synthetic-kimi-server-test",
5631            n_layers: 1,
5632            hidden_dim,
5633            n_heads: 1,
5634            n_kv_heads: 1,
5635            head_dim: 4,
5636            vocab_size,
5637            rope_theta: 10000.0,
5638            rms_norm_eps: 1e-5,
5639            sliding_window: None,
5640            moe: MoeLayerConfig {
5641                expert_weights_scale: 1.0,
5642                n_experts: 1,
5643                n_experts_active: 1,
5644                n_shared_experts: 0,
5645                hidden_dim,
5646                expert_ffn_dim: 4,
5647                gating: GatingFunction::Sigmoid,
5648                norm_topk_prob: true,
5649                expert_group_count: None,
5650                expert_group_used_count: None,
5651            },
5652            // Layer 0 is the sole dense leading layer, using KDA
5653            // attention (real Kimi K3's own layer-0 shape) -- the
5654            // 1-indexed `kda_layers`/`full_attn_layers` convention is
5655            // `ModelConfig::layer_attention_kind`'s, not this test's.
5656            n_dense_leading_layers: 1,
5657            attention: AttentionKind::KimiHybrid(KimiHybridAttention {
5658                kda_layers: vec![1],
5659                full_attn_layers: vec![],
5660                mla: MlaConfig {
5661                    num_heads: mla_num_heads,
5662                    q_lora_rank: mla_q_lora_rank,
5663                    kv_lora_rank: mla_kv_lora_rank,
5664                    qk_nope_head_dim: mla_qk_nope_head_dim,
5665                    qk_rope_head_dim: mla_qk_rope_head_dim,
5666                    v_head_dim: mla_v_head_dim,
5667                    use_output_gate: true,
5668                    rope: None,
5669                },
5670                kda: KdaConfig {
5671                    num_heads: kda_num_heads,
5672                    head_dim: kda_head_dim,
5673                    short_conv_kernel_size: conv_kernel,
5674                    gate_lower_bound: -5.0,
5675                    use_full_rank_gate: true,
5676                },
5677            }),
5678            rope_freqs: None,
5679            rope_attn_factor: 1.0,
5680            rope_dim: None,
5681            rope_freqs_long: None,
5682            rope_freqs_short: None,
5683            rope_orig_ctx: None,
5684            rope_layout: ferrox_models::config::RopeLayout::Neox,
5685            qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
5686            swa_pattern: None,
5687            attn_logit_softcap: None,
5688            final_logit_softcap: None,
5689            embedding_scale: None,
5690            attention_scale: None,
5691            rope_theta_swa: None,
5692            ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
5693            best_effort_fields: &["synthetic test config, not a real preset"],
5694        };
5695        let hp = KimiRealHparams {
5696            hidden_dim,
5697            kda_num_heads,
5698            kda_head_dim,
5699            mla_num_heads,
5700            mla_q_lora_rank,
5701            mla_kv_lora_rank,
5702            mla_qk_nope_head_dim,
5703            mla_qk_rope_head_dim,
5704            mla_v_head_dim,
5705            dense_intermediate_dim: dense_intermediate,
5706            moe_hidden_dim: hidden_dim,
5707            moe_intermediate_dim: 4,
5708            n_experts: 1,
5709            num_shared_experts: 0,
5710        };
5711
5712        // Every real tensor name `kimi_loader::load_kimi_layer` (dense
5713        // FFN + KDA attention + block residual) and
5714        // `load_kimi_checkpoint` (top-level) actually read.
5715        let prefix = "language_model.model.layers.0";
5716        let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
5717        let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
5718                    name: String,
5719                    shape: Vec<usize>,
5720                    n: usize| {
5721            tensors.push((name, shape, vec![0.01f32; n]));
5722        };
5723        push(
5724            &mut tensors,
5725            format!("{prefix}.input_layernorm.weight"),
5726            vec![hidden_dim],
5727            hidden_dim,
5728        );
5729        push(
5730            &mut tensors,
5731            format!("{prefix}.post_attention_layernorm.weight"),
5732            vec![hidden_dim],
5733            hidden_dim,
5734        );
5735        push(
5736            &mut tensors,
5737            format!("{prefix}.self_attention_res_norm.weight"),
5738            vec![hidden_dim],
5739            hidden_dim,
5740        );
5741        push(
5742            &mut tensors,
5743            format!("{prefix}.self_attention_res_proj.weight"),
5744            vec![1, hidden_dim],
5745            hidden_dim,
5746        );
5747        push(
5748            &mut tensors,
5749            format!("{prefix}.mlp_res_norm.weight"),
5750            vec![hidden_dim],
5751            hidden_dim,
5752        );
5753        push(
5754            &mut tensors,
5755            format!("{prefix}.mlp_res_proj.weight"),
5756            vec![1, hidden_dim],
5757            hidden_dim,
5758        );
5759        push(
5760            &mut tensors,
5761            format!("{prefix}.self_attn.q_proj.weight"),
5762            vec![kda_proj, hidden_dim],
5763            kda_proj * hidden_dim,
5764        );
5765        push(
5766            &mut tensors,
5767            format!("{prefix}.self_attn.k_proj.weight"),
5768            vec![kda_proj, hidden_dim],
5769            kda_proj * hidden_dim,
5770        );
5771        push(
5772            &mut tensors,
5773            format!("{prefix}.self_attn.v_proj.weight"),
5774            vec![kda_proj, hidden_dim],
5775            kda_proj * hidden_dim,
5776        );
5777        push(
5778            &mut tensors,
5779            format!("{prefix}.self_attn.q_conv1d.weight"),
5780            vec![kda_proj, 1, conv_kernel],
5781            kda_proj * conv_kernel,
5782        );
5783        push(
5784            &mut tensors,
5785            format!("{prefix}.self_attn.k_conv1d.weight"),
5786            vec![kda_proj, 1, conv_kernel],
5787            kda_proj * conv_kernel,
5788        );
5789        push(
5790            &mut tensors,
5791            format!("{prefix}.self_attn.v_conv1d.weight"),
5792            vec![kda_proj, 1, conv_kernel],
5793            kda_proj * conv_kernel,
5794        );
5795        push(
5796            &mut tensors,
5797            format!("{prefix}.self_attn.A_log"),
5798            vec![kda_num_heads],
5799            kda_num_heads,
5800        );
5801        push(
5802            &mut tensors,
5803            format!("{prefix}.self_attn.f_a_proj.weight"),
5804            vec![kda_head_dim, hidden_dim],
5805            kda_head_dim * hidden_dim,
5806        );
5807        push(
5808            &mut tensors,
5809            format!("{prefix}.self_attn.f_b_proj.weight"),
5810            vec![kda_proj, kda_head_dim],
5811            kda_proj * kda_head_dim,
5812        );
5813        push(
5814            &mut tensors,
5815            format!("{prefix}.self_attn.dt_bias"),
5816            vec![kda_proj],
5817            kda_proj,
5818        );
5819        push(
5820            &mut tensors,
5821            format!("{prefix}.self_attn.b_proj.weight"),
5822            vec![kda_num_heads, hidden_dim],
5823            kda_num_heads * hidden_dim,
5824        );
5825        push(
5826            &mut tensors,
5827            format!("{prefix}.self_attn.g_proj.weight"),
5828            vec![kda_proj, hidden_dim],
5829            kda_proj * hidden_dim,
5830        );
5831        push(
5832            &mut tensors,
5833            format!("{prefix}.self_attn.o_norm.weight"),
5834            vec![kda_head_dim],
5835            kda_head_dim,
5836        );
5837        push(
5838            &mut tensors,
5839            format!("{prefix}.self_attn.o_proj.weight"),
5840            vec![hidden_dim, kda_proj],
5841            hidden_dim * kda_proj,
5842        );
5843        push(
5844            &mut tensors,
5845            format!("{prefix}.mlp.gate_proj.weight"),
5846            vec![dense_intermediate, hidden_dim],
5847            dense_intermediate * hidden_dim,
5848        );
5849        push(
5850            &mut tensors,
5851            format!("{prefix}.mlp.up_proj.weight"),
5852            vec![dense_intermediate, hidden_dim],
5853            dense_intermediate * hidden_dim,
5854        );
5855        push(
5856            &mut tensors,
5857            format!("{prefix}.mlp.down_proj.weight"),
5858            vec![hidden_dim, dense_intermediate],
5859            hidden_dim * dense_intermediate,
5860        );
5861        push(
5862            &mut tensors,
5863            "language_model.model.embed_tokens.weight".to_string(),
5864            vec![vocab_size, hidden_dim],
5865            vocab_size * hidden_dim,
5866        );
5867        push(
5868            &mut tensors,
5869            "language_model.lm_head.weight".to_string(),
5870            vec![vocab_size, hidden_dim],
5871            vocab_size * hidden_dim,
5872        );
5873        push(
5874            &mut tensors,
5875            "language_model.model.norm.weight".to_string(),
5876            vec![hidden_dim],
5877            hidden_dim,
5878        );
5879        push(
5880            &mut tensors,
5881            "language_model.model.output_attn_res_norm.weight".to_string(),
5882            vec![hidden_dim],
5883            hidden_dim,
5884        );
5885        push(
5886            &mut tensors,
5887            "language_model.model.output_attn_res_proj.weight".to_string(),
5888            vec![1, hidden_dim],
5889            hidden_dim,
5890        );
5891
5892        let dir = std::env::temp_dir().join(format!(
5893            "ferrox_server_kimi_e2e_test_{}_{}",
5894            std::process::id(),
5895            vocab_size
5896        ));
5897        std::fs::create_dir_all(&dir).unwrap();
5898        let shard_bytes = write_safetensors_shard(&tensors);
5899        std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
5900        let map_entries: Vec<String> = tensors
5901            .iter()
5902            .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
5903            .collect();
5904        let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
5905        std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
5906
5907        // A real tiktoken-format vocab file: one base64-encoded byte
5908        // plus its rank per line -- enough to round-trip an ASCII
5909        // prompt without needing the real 163584-entry Kimi K3 vocab.
5910        use base64::Engine;
5911        let vocab_lines: Vec<String> = (0..vocab_size as u32)
5912            .map(|b| {
5913                let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
5914                format!("{b64} {b}")
5915            })
5916            .collect();
5917        std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
5918
5919        let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
5920            .expect("must load the synthetic Kimi checkpoint end to end");
5921        std::fs::remove_dir_all(&dir).ok();
5922        loaded
5923    }
5924
5925    /// The real end-to-end proof for Kimi-through-the-server: a real
5926    /// synthetic Kimi K3 checkpoint served through the exact same
5927    /// `run_generation` entry point the HTTP handlers call for the
5928    /// GGUF path. Proves the whole new plumbing end to end: directory-
5929    /// shaped checkpoint loading, `KimiEngine`/`KimiTokenizer` wired
5930    /// through the `Model` enum, and `generate::generate_engine`
5931    /// producing real, bounded generated text.
5932    #[test]
5933    fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
5934        let loaded = build_synthetic_kimi_loaded();
5935        let state = build_app_state(
5936            model::LoadedModel::Kimi(loaded),
5937            None,
5938            None,
5939            false,
5940            None,
5941            Arc::new(health::Detection::ready(health::probe_backends())),
5942        );
5943        let active = state.active().expect("a freshly built state has a model");
5944        assert_eq!(active.model.tokenizer_kind(), "kimi-tiktoken-bpe");
5945        assert!(!active.model.is_synthetic());
5946
5947        let (_chunks, finish, _usage) = run_generation(
5948            &active.model,
5949            "hi",
5950            &greedy_params(5),
5951            None,
5952            None,
5953            None,
5954            None,
5955        )
5956        .expect("a real Kimi checkpoint must generate without error");
5957        assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5958    }
5959
5960    /// Explicit proof of the "gate, don't paper over" design decision
5961    /// (see `ferrox_models::engine`'s module docs): even when an operator configures
5962    /// a KV block pool and/or prefix cache, a Kimi request must never
5963    /// consult either -- `generate_engine`'s signature has no
5964    /// parameter for them at all, so this isn't just an unexercised
5965    /// code path, it's structurally impossible for a Kimi request to
5966    /// touch them. Confirmed here by observing both are completely
5967    /// untouched (pool blocks unchanged, cache stats unchanged) after a
5968    /// real Kimi generation runs alongside both.
5969    #[test]
5970    fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
5971        let loaded = build_synthetic_kimi_loaded();
5972        let state = build_app_state(
5973            model::LoadedModel::Kimi(loaded),
5974            None,
5975            None,
5976            false,
5977            None,
5978            Arc::new(health::Detection::ready(health::probe_backends())),
5979        );
5980
5981        let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
5982        let kv_pool_config = generate::KvPoolConfig {
5983            pool: pool.clone(),
5984            queue_wait: Duration::ZERO,
5985        };
5986        let pc = Mutex::new(PrefixCache::new(4));
5987
5988        run_generation(
5989            &state
5990                .active()
5991                .expect("a freshly built state has a model")
5992                .model,
5993            "hi",
5994            &greedy_params(5),
5995            Some(&kv_pool_config),
5996            Some(&pc),
5997            None,
5998            None,
5999        )
6000        .expect("a real Kimi checkpoint must generate without error");
6001
6002        assert_eq!(
6003            pool.lock().unwrap().free_blocks(),
6004            4,
6005            "the KV pool must be completely untouched by a Kimi request"
6006        );
6007        let stats = pc.lock().unwrap().stats();
6008        assert_eq!(
6009            stats.hits + stats.misses,
6010            0,
6011            "the prefix cache must never be consulted for a Kimi request"
6012        );
6013    }
6014}