1mod admin;
36mod anthropic;
37mod attribution;
38mod budget;
39mod cache_admin;
40mod cancel;
41mod chat_template;
42mod completion;
43mod conversations;
44mod decode_task;
45mod embeddings;
46mod generate;
47mod grammar_request;
48mod health;
49mod journal;
50mod json_mode;
51mod limits;
52mod loaded;
53mod mcp;
54mod model;
55mod openai_extra;
56mod output;
57mod policy;
58mod rerank;
59mod response_cache;
60pub(crate) mod responses;
61mod resume;
62mod sample_step;
63mod sampling_knobs;
64mod security;
65mod serving;
66mod session;
67mod sse;
68mod stats;
69mod stop;
70mod stream_events;
71mod tasks;
72mod tool_grammar;
73mod unsupported_sampling;
74
75use std::cell::RefCell;
76use std::convert::Infallible;
77use std::fmt;
78use std::net::{IpAddr, Ipv4Addr, SocketAddr};
79use std::path::PathBuf;
80use std::rc::Rc;
81use std::str::FromStr;
82use std::sync::{Arc, Mutex, MutexGuard};
83use std::time::Duration;
84
85use axum::{
86 extract::State,
87 http::StatusCode,
88 response::sse::{Event, Sse},
89 response::{IntoResponse, Response},
90 routing::{get, post},
91 Json, Router,
92};
93use clap::{Parser, ValueEnum};
94use serde::{Deserialize, Serialize};
95
96use ferrox_core::cache::KvBlockPool;
97use ferrox_models::kimi_tokenizer::KimiTokenizer;
98use ferrox_models::sampling::SamplingParams;
99use ferrox_models::tokenizer::StopTokens;
100use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
101use generate::{FinishReason, GenerationParams};
102pub(crate) use loaded::{ActiveModel, Loaded};
103use model::ServerTokenizer;
104use rerank::encoder_endpoints;
105use response_cache::{CacheKey, ResponseCache};
106use sampling_knobs::SamplingKnobs;
107
108#[derive(Parser, Debug, PartialEq)]
112#[command(
119 name = "ferrox-server",
120 about = "OpenAI-compatible Ferrox inference server"
121)]
122pub struct ServerArgs {
123 #[arg(short = 'm', long = "model", value_name = "FILE")]
125 model: Option<String>,
126
127 #[arg(
135 long = "hf-repo",
136 visible_alias = "hf",
137 value_name = "REPO[:QUANT]",
138 conflicts_with = "model"
139 )]
140 hf_repo: Option<String>,
141
142 #[arg(long = "hf-file", value_name = "FILE", requires = "hf_repo")]
147 hf_file: Option<String>,
148
149 #[arg(short = 'c', long = "ctx-size", value_name = "N")]
155 ctx_size: Option<usize>,
156
157 #[arg(long = "api-key", value_name = "KEY")]
160 api_key: Option<String>,
161
162 #[arg(long = "api-key-file", value_name = "PATH", conflicts_with = "api_key")]
167 api_key_file: Option<std::path::PathBuf>,
168
169 #[arg(long = "alias", visible_alias = "model-alias", value_name = "NAME")]
172 alias: Option<String>,
173
174 #[arg(long = "ctk", visible_alias = "cache-type-k", value_name = "TYPE")]
177 ctk: Option<String>,
178
179 #[arg(long = "jinja", default_value_t = false)]
185 jinja: bool,
186
187 #[arg(long = "no-jinja", default_value_t = false)]
190 no_jinja: bool,
191
192 #[arg(long = "no-warmup", default_value_t = false)]
194 no_warmup: bool,
195
196 #[arg(long = "flash-attn", visible_alias = "fa", value_name = "MODE", num_args = 0..=1, default_missing_value = "auto")]
200 flash_attn: Option<String>,
201
202 #[arg(long, value_name = "HOST")]
204 host: Option<IpAddr>,
205
206 #[arg(long, value_name = "PORT")]
211 port: Option<u16>,
212
213 #[arg(short = 't', long = "threads", value_name = "N")]
215 threads: Option<usize>,
216
217 #[arg(
219 long = "device",
220 visible_alias = "dev",
221 value_name = "DEVICE",
222 ignore_case = true
223 )]
224 device: Option<OffloadDevice>,
225
226 #[arg(long = "list-devices", default_value_t = false)]
228 list_devices: bool,
229
230 #[arg(
235 long = "n-gpu-layers",
236 visible_aliases = ["gpu-layers", "ngl"],
237 value_name = "N"
238 )]
239 n_gpu_layers: Option<GpuLayers>,
240
241 #[arg(long = "mcp-config", value_name = "PATH")]
243 mcp_config: Option<PathBuf>,
244
245 #[arg(long = "exit-on-stdin-close", default_value_t = false)]
253 exit_on_stdin_close: bool,
254
255 #[arg(
258 long = "cont-batching",
259 visible_aliases = ["continuous-batching", "cb"],
260 default_value_t = false
261 )]
262 cont_batching: bool,
263
264 #[arg(
267 long = "no-cont-batching",
268 default_value_t = false,
269 conflicts_with = "cont_batching"
270 )]
271 no_cont_batching: bool,
272
273 #[arg(long = "parallel", visible_alias = "np", value_name = "N")]
277 parallel: Option<usize>,
278
279 #[arg(long = "allow-multiple-instances", default_value_t = false)]
284 allow_multiple_instances: bool,
285}
286
287impl ServerArgs {
288 pub fn parse_llama_style<I>(argv: I) -> Self
295 where
296 I: IntoIterator<Item = String>,
297 {
298 Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
299 }
300}
301
302#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
303enum OffloadDevice {
304 Auto,
305 None,
306 Cpu,
307 Metal,
308 Cuda,
309}
310
311#[derive(Debug, Clone, Copy, PartialEq, Eq)]
312enum GpuLayers {
313 Auto,
314 All,
315 Count(u32),
316}
317
318impl GpuLayers {
319 fn offload_enabled(self) -> bool {
320 !matches!(self, Self::Count(0))
321 }
322}
323
324impl FromStr for GpuLayers {
325 type Err = String;
326
327 fn from_str(value: &str) -> Result<Self, Self::Err> {
328 match value {
329 "auto" => Ok(Self::Auto),
330 "all" => Ok(Self::All),
331 _ => value
332 .parse::<u32>()
333 .map(Self::Count)
334 .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
335 }
336 }
337}
338
339impl fmt::Display for GpuLayers {
340 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
341 match self {
342 Self::Auto => f.write_str("auto"),
343 Self::All => f.write_str("all"),
344 Self::Count(value) => value.fmt(f),
345 }
346 }
347}
348
349pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
358
359pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
362
363fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
364 args.into_iter()
365 .map(|arg| match arg.as_str() {
366 "-ngl" => "--n-gpu-layers".into(),
367 "-dev" => "--device".into(),
368 "-cb" => "--cont-batching".into(),
369 "-np" => "--parallel".into(),
370 "-hf" => "--hf-repo".into(),
375 "-hff" => "--hf-file".into(),
376 _ => arg,
377 })
378 .collect()
379}
380
381fn print_available_devices() {
382 println!("Available devices:");
383 println!(" CPU");
384
385 let metal = ferrox_metal::MetalProfile::detect();
386 if let Some(name) = metal.device_name {
387 println!(" Metal: {name}");
388 }
389
390 let cuda = ferrox_cuda::HardwareProfile::detect();
391 if cuda.cuda_available {
392 let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
393 println!(" CUDA: {name}");
394 if cuda.cuda_device_count > 1 {
395 println!(" ({} devices detected)", cuda.cuda_device_count);
396 }
397 }
398}
399
400fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
401 if args.host.is_none() && args.port.is_none() {
402 return None;
403 }
404
405 let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
406 let host = args
407 .host
408 .or_else(|| existing.map(|addr| addr.ip()))
409 .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
410 let port = args
411 .port
412 .or_else(|| existing.map(|addr| addr.port()))
413 .unwrap_or(8383);
414 Some(SocketAddr::new(host, port).to_string())
415}
416
417fn resolve_hf_repo(spec: &str, file: Option<&str>) -> anyhow::Result<String> {
423 let mut hf = ferrox_models::hub::HfRef::parse(spec);
424 if let Some(f) = file {
425 hf.file = Some(f.to_string());
426 }
427 eprintln!(
428 "ferrox: resolving {} on the Hub{}",
429 hf.repo,
430 hf.quant
431 .as_deref()
432 .map(|q| format!(" ({q})"))
433 .unwrap_or_default()
434 );
435
436 let mut last = std::time::Instant::now();
437 let mut draw = move |done: u64, total: Option<u64>| {
438 if last.elapsed() < std::time::Duration::from_millis(200) {
439 return;
440 }
441 last = std::time::Instant::now();
442 let mib = done as f64 / 1024.0 / 1024.0;
443 match total {
444 Some(t) if t > 0 => {
445 eprint!(
446 "\r {mib:>9.1} MiB {:5.1}%",
447 (done as f64 / t as f64) * 100.0
448 )
449 }
450 _ => eprint!("\r {mib:>9.1} MiB"),
451 }
452 };
453
454 let (path, downloaded) = hf
455 .ensure_local(&mut draw)
456 .map_err(|e| anyhow::anyhow!("{e}"))?;
457 if downloaded {
458 eprintln!();
459 eprintln!("ferrox: downloaded {}", path.display());
460 } else {
461 eprintln!("ferrox: using cached {}", path.display());
462 }
463 Ok(path.to_string_lossy().into_owned())
464}
465
466fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
467 if let Some(model) = &args.model {
468 unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
470 }
471 if let Some(spec) = &args.hf_repo {
472 let path = resolve_hf_repo(spec, args.hf_file.as_deref())?;
473 unsafe { std::env::set_var("FERROX_MODEL_PATH", &path) };
475 }
476 if let Some(n) = args.ctx_size {
477 if n == 0 {
478 anyhow::bail!("--ctx-size must be greater than zero");
479 }
480 unsafe { std::env::set_var("FERROX_CB_MAX_CONTEXT", n.to_string()) };
482 }
483 if let Some(key) = &args.api_key {
484 unsafe { std::env::set_var("FERROX_API_KEY", key) };
486 }
487 if let Some(path) = &args.api_key_file {
488 let key = std::fs::read_to_string(path)
489 .map_err(|e| anyhow::anyhow!("reading --api-key-file {}: {e}", path.display()))?;
490 let key = key.trim();
491 if key.is_empty() {
492 anyhow::bail!(
493 "--api-key-file {} is empty: an empty key would leave every route open, \
494 which is the opposite of what passing the flag asked for",
495 path.display()
496 );
497 }
498 unsafe { std::env::set_var("FERROX_API_KEY", key) };
500 }
501 if let Some(alias) = &args.alias {
502 unsafe { std::env::set_var("FERROX_MODEL_NAME", alias) };
504 }
505 if let Some(ctk) = &args.ctk {
506 unsafe { std::env::set_var("FERROX_CTK", ctk.trim()) };
508 }
509 if args.no_jinja {
514 anyhow::bail!(
515 "--no-jinja: ferrox has no template-free mode. It compiles and evaluates the GGUF's \
516 own tokenizer.chat_template, which is what llama.cpp's --jinja turns on, and there \
517 is no sniffing fallback to switch to. Use --no-cnv on `ferrox run` for a raw \
518 completion"
519 );
520 }
521 if let Some(mode) = &args.flash_attn {
522 let mode = mode.trim().to_ascii_lowercase();
523 if mode == "off" || mode == "disabled" || mode == "0" {
524 anyhow::bail!(
525 "--flash-attn off: fused attention is a backend property here, not a per-run \
526 switch. Set FERROX_METAL_ATTN=0 to take the unfused Metal path, or --device cpu"
527 );
528 }
529 }
530
531 if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
532 unsafe { std::env::set_var("FERROX_ADDR", addr) };
534 }
535
536 if let Some(threads) = args.threads {
537 if threads == 0 {
538 anyhow::bail!("--threads must be greater than zero");
539 }
540 unsafe {
542 std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
543 std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
544 }
545 }
546
547 if args.device.is_none() && args.n_gpu_layers.is_none() {
548 } else {
550 let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
551 let device = if layers.offload_enabled() {
552 args.device.unwrap_or(OffloadDevice::Auto)
553 } else {
554 OffloadDevice::None
555 };
556
557 match device {
558 OffloadDevice::None | OffloadDevice::Cpu => unsafe {
559 std::env::set_var("FERROX_METAL", "0");
560 std::env::set_var("FERROX_METAL_ATTN", "0");
561 std::env::set_var("FERROX_CUDA", "0");
562 },
563 OffloadDevice::Auto => unsafe {
564 std::env::set_var("FERROX_METAL", "auto");
565 std::env::set_var("FERROX_CUDA", "auto");
566 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
567 std::env::set_var("FERROX_METAL_ATTN", "1");
568 }
569 },
570 OffloadDevice::Metal => {
571 #[cfg(not(feature = "metal"))]
572 {
573 anyhow::bail!(
574 "Metal requested but this binary was built without --features metal"
575 );
576 }
577 #[cfg(feature = "metal")]
578 {
579 if !ferrox_metal::MetalProfile::detect().available {
580 anyhow::bail!("Metal requested but no Metal device is available");
581 }
582 unsafe {
583 std::env::set_var("FERROX_METAL", "1");
584 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
585 std::env::set_var("FERROX_METAL_ATTN", "1");
586 }
587 std::env::set_var("FERROX_CUDA", "0");
588 }
589 }
590 }
591 OffloadDevice::Cuda => {
592 #[cfg(not(feature = "cuda"))]
593 {
594 anyhow::bail!(
595 "CUDA requested but this binary was built without --features cuda"
596 );
597 }
598 #[cfg(feature = "cuda")]
599 {
600 if !ferrox_cuda::HardwareProfile::detect().cuda_available {
601 anyhow::bail!("CUDA requested but no CUDA device is available");
602 }
603 unsafe {
604 std::env::set_var("FERROX_CUDA", "1");
605 std::env::set_var("FERROX_METAL", "0");
606 std::env::set_var("FERROX_METAL_ATTN", "0");
607 }
608 }
609 }
610 }
611 }
612
613 if let Some(n) = args.parallel {
614 if n == 0 {
615 anyhow::bail!("--parallel must be greater than zero");
616 }
617 unsafe { std::env::set_var("FERROX_CB_MAX_SEQS", n.to_string()) };
619 }
620
621 if args.cont_batching {
622 unsafe { std::env::set_var("FERROX_CONTINUOUS_BATCHING", "1") };
624 } else if args.no_cont_batching {
625 unsafe { std::env::set_var("FERROX_CONTINUOUS_BATCHING", "0") };
627 } else if args.parallel.is_some() {
628 unsafe { std::env::set_var("FERROX_CONTINUOUS_BATCHING", "1") };
631 }
632
633 Ok(())
634}
635
636#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
645 Gguf(GgufModel),
646 Kimi(KimiModel),
647 Mla(MlaModel),
648 Gemma4(Gemma4Model),
649 Glm52(Glm52Model),
650}
651
652pub(crate) struct GgufModel {
653 decoder: Arc<Decoder>,
654 tokenizer: Arc<ServerTokenizer>,
655 stop_tokens: StopTokens,
656 bos_id: Option<usize>,
657 is_synthetic: bool,
658 chat_template: chat_template::PromptTemplate,
659}
660
661pub(crate) struct KimiModel {
662 engine: KimiEngine,
663 tokenizer: KimiTokenizer,
664 stop_tokens: StopTokens,
665 chat_template: chat_template::PromptTemplate,
666}
667
668pub(crate) struct MlaModel {
669 engine: MlaEngine,
670 tokenizer: ServerTokenizer,
671 stop_tokens: StopTokens,
672 bos_id: Option<usize>,
673 name: String,
674 chat_template: chat_template::PromptTemplate,
675}
676
677pub(crate) struct Gemma4Model {
678 engine: Gemma4Engine,
679 tokenizer: ServerTokenizer,
680 stop_tokens: StopTokens,
681 bos_id: Option<usize>,
682 name: String,
683 chat_template: chat_template::PromptTemplate,
684}
685
686pub(crate) struct Glm52Model {
687 engine: ferrox_models::Glm52Engine,
688 tokenizer: ServerTokenizer,
689 stop_tokens: StopTokens,
690 bos_id: Option<usize>,
691 name: String,
692 chat_template: chat_template::PromptTemplate,
693}
694
695impl Model {
696 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
697 match self {
698 Model::Gguf(m) => m.chat_template.clone(),
699 Model::Kimi(m) => m.chat_template.clone(),
700 Model::Mla(m) => m.chat_template.clone(),
701 Model::Gemma4(m) => m.chat_template.clone(),
702 Model::Glm52(m) => m.chat_template.clone(),
703 }
704 }
705
706 fn is_synthetic(&self) -> bool {
711 match self {
712 Model::Gguf(m) => m.is_synthetic,
713 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
714 }
715 }
716
717 fn tokenizer_kind(&self) -> &'static str {
718 match self {
719 Model::Gguf(m) => m.tokenizer.kind(),
720 Model::Kimi(_) => "kimi-tiktoken-bpe",
721 Model::Mla(m) => m.tokenizer.kind(),
722 Model::Gemma4(m) => m.tokenizer.kind(),
723 Model::Glm52(m) => m.tokenizer.kind(),
724 }
725 }
726
727 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
731 match self {
732 Model::Gguf(m) => m.decoder.expert_store_stats(),
733 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
734 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
735 }
736 }
737
738 pub(crate) fn name(&self) -> &str {
739 match self {
740 Model::Gguf(m) => m.decoder.config.name,
741 Model::Kimi(_) => "kimi-k3",
742 Model::Mla(m) => m.name.as_str(),
743 Model::Gemma4(m) => m.name.as_str(),
744 Model::Glm52(m) => m.name.as_str(),
745 }
746 }
747
748 pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
749 match self {
750 Model::Gguf(m) => m.tokenizer.encode(text),
751 Model::Kimi(m) => m
752 .tokenizer
753 .encode(text)
754 .into_iter()
755 .map(|id| id as usize)
756 .collect(),
757 Model::Mla(m) => m.tokenizer.encode(text),
758 Model::Gemma4(m) => m.tokenizer.encode(text),
759 Model::Glm52(m) => m.tokenizer.encode(text),
760 }
761 }
762
763 pub(crate) fn bos_id(&self) -> Option<usize> {
772 match self {
773 Model::Gguf(m) => m.bos_id,
774 Model::Kimi(_) => None,
775 Model::Mla(m) => m.bos_id,
776 Model::Gemma4(m) => m.bos_id,
777 Model::Glm52(m) => m.bos_id,
778 }
779 }
780
781 pub(crate) fn decode(&self, ids: &[usize]) -> String {
782 match self {
783 Model::Gguf(m) => m.tokenizer.decode(ids),
784 Model::Kimi(m) => {
785 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
786 m.tokenizer.decode(&ids32)
787 }
788 Model::Mla(m) => m.tokenizer.decode(ids),
789 Model::Gemma4(m) => m.tokenizer.decode(ids),
790 Model::Glm52(m) => m.tokenizer.decode(ids),
791 }
792 }
793
794 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
797 match self {
798 Model::Gguf(m) => {
799 let mut caches: Vec<_> = (0..m.decoder.layers.len())
800 .map(|_| {
801 ferrox_core::cache::KvCache::new(
802 m.decoder.config.n_kv_heads,
803 m.decoder.config.head_dim,
804 )
805 })
806 .collect();
807 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
808 }
809 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
810 }
811 }
812
813 pub(crate) fn vocab_size(&self) -> Option<usize> {
814 match self {
815 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
816 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
817 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
818 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
819 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
820 }
821 }
822}
823
824pub(crate) struct AppState {
825 pub(crate) embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
839 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
858 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
862 pub(crate) tasks: Arc<tasks::TaskRegistry>,
864 pub(crate) cancels: Arc<cancel::CancelRegistry>,
868 pub(crate) stats: stats::Stats,
871 pub(crate) streams: resume::StreamRegistry,
874 pub(crate) model_dir: Option<PathBuf>,
876 response_cache: Mutex<ResponseCache>,
880 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
892 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
905 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
915 sessions: session::SessionStore,
921 requests_total: std::sync::atomic::AtomicU64,
922 request_errors_total: std::sync::atomic::AtomicU64,
923 started_at: std::time::Instant,
924 last_request_ms: std::sync::atomic::AtomicU64,
929 detection: Arc<health::Detection>,
931 mcp: Option<mcp::LoadedMcpConfig>,
933 pub(crate) continuous_batching_enabled: bool,
937 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
941 loading_model: Mutex<Option<String>>,
945 last_load_error: Mutex<Option<(String, String)>>,
949 pub(crate) serving: Mutex<crate::stats::ServingStats>,
954 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
959 pub(crate) footprint:
965 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
966 pub(crate) started_unix: u64,
976}
977
978pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
984
985impl AppState {
986 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
990 self.active
991 .read()
992 .unwrap_or_else(|p| p.into_inner())
993 .clone()
994 }
995
996 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
1002 self.active().ok_or_else(|| {
1003 (
1004 StatusCode::SERVICE_UNAVAILABLE,
1005 Json(serde_json::json!({"error": {
1006 "message": "no model is loaded; POST /admin/models/load with an id from \
1007 GET /admin/models",
1008 "type": "model_not_loaded"
1009 }})),
1010 )
1011 })
1012 }
1013
1014 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
1022 Ok(Arc::clone(self.require_active()?.generative()?))
1023 }
1024
1025 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
1033 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
1034 std::mem::replace(&mut *guard, next)
1035 }
1036
1037 fn mark_request_finished(&self) {
1041 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
1042 self.last_request_ms
1043 .store(ms, std::sync::atomic::Ordering::Relaxed);
1044 }
1045
1046 pub(crate) fn uptime(&self) -> Duration {
1047 self.started_at.elapsed()
1048 }
1049
1050 pub(crate) fn requests_total(&self) -> u64 {
1051 self.requests_total
1052 .load(std::sync::atomic::Ordering::Relaxed)
1053 }
1054
1055 pub(crate) fn errors_total(&self) -> u64 {
1056 self.request_errors_total
1057 .load(std::sync::atomic::Ordering::Relaxed)
1058 }
1059
1060 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
1061 lock_cache(&self.response_cache).stats()
1062 }
1063
1064 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
1067 let last = self
1068 .last_request_ms
1069 .load(std::sync::atomic::Ordering::Relaxed);
1070 (last > 0)
1071 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
1072 .map(|age| age.max(0.0))
1073 }
1074
1075 pub(crate) fn loading_model_id(&self) -> Option<String> {
1076 self.loading_model
1077 .lock()
1078 .unwrap_or_else(|p| p.into_inner())
1079 .clone()
1080 }
1081
1082 pub(crate) fn set_loading_model(&self, id: Option<String>) {
1083 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
1084 }
1085
1086 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
1087 self.last_load_error
1088 .lock()
1089 .unwrap_or_else(|p| p.into_inner())
1090 .clone()
1091 }
1092
1093 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
1094 *self
1095 .last_load_error
1096 .lock()
1097 .unwrap_or_else(|p| p.into_inner()) = error;
1098 }
1099
1100 pub(crate) fn active_model_name(&self) -> Option<String> {
1109 self.active().map(|a| a.name().to_string())
1110 }
1111
1112 pub(crate) fn embedding_model(&self) -> Option<Arc<ferrox_models::EmbeddingModel>> {
1122 self.embedding
1123 .clone()
1124 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
1125 }
1126
1127 pub(crate) fn embedding_model_name(&self) -> Option<String> {
1131 match self.embedding_model() {
1132 Some(e) => Some(e.name().to_string()),
1133 None => self.active_model_name(),
1134 }
1135 }
1136
1137 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
1138 self.stats.record(stats::entry(record));
1139 }
1140}
1141
1142fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
1147 cache
1148 .lock()
1149 .unwrap_or_else(|poisoned| poisoned.into_inner())
1150}
1151
1152#[derive(Debug, Clone, Deserialize)]
1153#[serde(untagged)]
1154pub(crate) enum MessageContent {
1155 Text(String),
1156 Parts(Vec<ContentPart>),
1157}
1158
1159#[derive(Debug, Clone, Deserialize)]
1160struct ContentPart {
1161 #[serde(rename = "type")]
1162 kind: String,
1163 #[serde(default)]
1164 text: Option<String>,
1165 #[serde(default)]
1166 image_url: Option<serde_json::Value>,
1167}
1168
1169impl MessageContent {
1170 fn as_text(&self) -> String {
1171 match self {
1172 Self::Text(s) => s.clone(),
1173 Self::Parts(parts) => parts
1174 .iter()
1175 .filter_map(|p| p.text.as_deref())
1176 .collect::<Vec<_>>()
1177 .join(""),
1178 }
1179 }
1180
1181 fn has_image(&self) -> bool {
1182 match self {
1183 Self::Text(_) => false,
1184 Self::Parts(parts) => parts
1185 .iter()
1186 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
1187 }
1188 }
1189}
1190
1191#[derive(Debug, Clone, Deserialize)]
1192pub(crate) struct ChatMessage {
1193 pub(crate) role: String,
1194 #[serde(default)]
1198 pub(crate) content: Option<MessageContent>,
1199 #[serde(default)]
1203 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
1204 #[serde(default)]
1209 #[allow(dead_code)]
1210 pub(crate) tool_call_id: Option<String>,
1211 #[serde(default, alias = "reasoning")]
1227 pub(crate) reasoning_content: Option<String>,
1228}
1229
1230impl ChatMessage {
1231 fn rendered_content(&self) -> String {
1238 let mut out = self
1239 .content
1240 .as_ref()
1241 .map(MessageContent::as_text)
1242 .unwrap_or_default();
1243 if let Some(calls) = &self.tool_calls {
1244 for call in calls {
1245 out.push_str(&format!(
1246 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
1247 call.function.name, call.function.arguments
1248 ));
1249 }
1250 }
1251 out
1252 }
1253}
1254
1255#[derive(Debug, Clone, Deserialize)]
1256pub(crate) struct ToolCallIn {
1257 #[serde(default)]
1258 #[allow(dead_code)]
1259 id: String,
1260 #[serde(rename = "type", default)]
1261 #[allow(dead_code)]
1262 kind: String,
1263 function: ToolCallFunctionIn,
1264}
1265
1266#[derive(Debug, Clone, Deserialize)]
1267struct ToolCallFunctionIn {
1268 name: String,
1269 arguments: String,
1274}
1275
1276#[derive(Debug, Clone, Deserialize)]
1279struct ToolDef {
1280 #[serde(rename = "type", default)]
1281 #[allow(dead_code)]
1282 kind: String,
1283 function: ToolFunctionDef,
1284}
1285
1286#[derive(Debug, Clone, Deserialize)]
1287struct ToolFunctionDef {
1288 name: String,
1289 #[serde(default)]
1290 description: Option<String>,
1291 #[serde(default)]
1292 parameters: Option<serde_json::Value>,
1293}
1294
1295#[derive(Debug, Clone, Deserialize)]
1306#[serde(untagged)]
1307enum ToolChoice {
1308 Mode(String),
1309 Specific(serde_json::Value),
1310}
1311
1312#[derive(Deserialize)]
1315#[serde(untagged)]
1316enum StopParam {
1317 One(String),
1318 Many(Vec<String>),
1319}
1320
1321#[derive(Deserialize)]
1322struct ChatCompletionRequest {
1323 model: String,
1324 messages: Vec<ChatMessage>,
1325 #[serde(default = "default_max_tokens")]
1326 max_tokens: usize,
1327 #[serde(default)]
1328 temperature: Option<f32>,
1329 #[serde(default)]
1330 top_p: Option<f32>,
1331 #[serde(default)]
1336 min_p: Option<f32>,
1337 #[serde(default)]
1338 top_k: Option<usize>,
1339 #[serde(default)]
1340 repetition_penalty: Option<f32>,
1341 #[serde(default)]
1342 seed: Option<u64>,
1343 #[serde(default)]
1344 stop: Option<StopParam>,
1345 #[serde(default)]
1346 stream: Option<bool>,
1347 #[serde(default)]
1360 stream_resumable: Option<bool>,
1361 #[serde(default)]
1371 ignore_eos: Option<bool>,
1372 #[serde(default)]
1373 tools: Vec<ToolDef>,
1374 #[serde(default)]
1375 tool_choice: Option<ToolChoice>,
1376 #[serde(default)]
1384 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
1385 #[serde(default)]
1393 reasoning_effort: Option<String>,
1394 #[serde(default)]
1398 thinking: Option<ThinkingSwitch>,
1399 #[serde(default)]
1404 session_id: Option<String>,
1405 #[serde(default)]
1407 logprobs: Option<bool>,
1408 #[serde(default)]
1409 top_logprobs: Option<u32>,
1410 #[serde(default)]
1411 n: Option<u32>,
1412 #[serde(default)]
1413 presence_penalty: Option<f32>,
1414 #[serde(default)]
1415 frequency_penalty: Option<f32>,
1416 #[serde(default)]
1417 response_format: Option<serde_json::Value>,
1418 #[serde(default)]
1425 logit_bias: Option<serde_json::Value>,
1426 #[serde(default)]
1436 grammar: Option<String>,
1437}
1438
1439const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1451
1452#[derive(Debug, Clone, Deserialize)]
1454pub(crate) struct ThinkingSwitch {
1455 #[serde(rename = "type")]
1456 pub(crate) kind: String,
1457}
1458
1459const THINKING_KWARG_KEYS: [&str; 4] = [
1463 "enable_thinking",
1464 "thinking",
1465 "thinking_mode",
1466 "reasoning_effort",
1467];
1468
1469const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1473
1474fn default_max_tokens() -> usize {
1475 DEFAULT_CHAT_MAX_TOKENS
1476}
1477
1478impl ChatCompletionRequest {
1479 fn sampling_knobs(&self) -> SamplingKnobs {
1484 SamplingKnobs {
1485 temperature: self.temperature,
1486 top_p: self.top_p,
1487 min_p: self.min_p,
1488 top_k: self.top_k,
1489 repetition_penalty: self.repetition_penalty,
1490 presence_penalty: self.presence_penalty,
1491 frequency_penalty: self.frequency_penalty,
1492 penalty_last_n: None,
1496 }
1497 }
1498
1499 fn sampling_params(&self) -> SamplingParams {
1500 self.sampling_knobs().resolve()
1501 }
1502
1503 fn stop_sequences(&self) -> Vec<String> {
1504 self.stop
1505 .as_ref()
1506 .map(|s| match s {
1507 StopParam::One(v) => vec![v.clone()],
1508 StopParam::Many(v) => v.clone(),
1509 })
1510 .unwrap_or_default()
1511 }
1512
1513 fn tools_active(&self) -> bool {
1518 !self.tools.is_empty()
1519 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1520 }
1521
1522 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1534 match &self.tool_choice {
1535 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1536 Some(ToolChoice::Specific(value)) => {
1537 let name = value
1541 .get("function")
1542 .and_then(|f| f.get("name"))
1543 .or_else(|| value.get("name"))
1544 .and_then(|n| n.as_str());
1545 match name {
1546 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1547 None => Err(invalid_request(
1548 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1549 function.name",
1550 "tool_choice",
1551 )),
1552 }
1553 }
1554 _ => Ok(None),
1555 }
1556 }
1557
1558 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1560 self.tools
1561 .iter()
1562 .map(|t| tool_grammar::ToolSpec {
1563 name: &t.function.name,
1564 parameters: t.function.parameters.as_ref(),
1565 })
1566 .collect()
1567 }
1568
1569 fn thinking_direction(&self) -> Option<bool> {
1612 if let Some(switch) = &self.thinking {
1613 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1614 "disabled" => Some(false),
1615 "enabled" => Some(true),
1616 _ => None,
1619 };
1620 }
1621 let effort = self.reasoning_effort.as_ref()?;
1622 DISABLE_EFFORTS
1623 .contains(&effort.trim().to_ascii_lowercase().as_str())
1624 .then_some(false)
1625 }
1626
1627 fn resolve_template_kwargs(
1628 &self,
1629 template: &chat_template::PromptTemplate,
1630 ) -> serde_json::Map<String, serde_json::Value> {
1631 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1632 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1636
1637 if !caller_steered {
1638 match self.thinking_direction() {
1639 Some(false) => {
1640 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1641 kwargs.insert(k, v);
1642 }
1643 return kwargs;
1646 }
1647 Some(true) => {
1648 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1649 kwargs.insert(k, v);
1650 }
1651 }
1652 None => {}
1653 }
1654 if let Some(effort) = &self.reasoning_effort {
1655 kwargs
1656 .entry("reasoning_effort".to_string())
1657 .or_insert_with(|| serde_json::json!(effort));
1658 }
1659 }
1660
1661 let offered: Vec<serde_json::Value> = if self.tools_active() {
1662 self.tools.iter().map(chat_template::tool_json).collect()
1663 } else {
1664 Vec::new()
1665 };
1666 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1667 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1668 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1669 kwargs.entry(k).or_insert(v);
1670 }
1671 }
1672 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1673 crate::policy::effort::EffortMapping::Mapped(to) => {
1674 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1675 }
1676 crate::policy::effort::EffortMapping::Dropped => {
1677 tracing::debug!(
1678 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1679 enough, so its own default applies"
1680 );
1681 }
1682 crate::policy::effort::EffortMapping::Unchanged => {}
1683 }
1684 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1685 kwargs
1686 }
1687
1688 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1691 if self.max_tokens == 0 {
1698 return Err(invalid_request(
1699 "max_tokens must be at least 1",
1700 "max_tokens",
1701 ));
1702 }
1703 if let Some(switch) = &self.thinking {
1708 let kind = switch.kind.trim().to_ascii_lowercase();
1709 if kind != "enabled" && kind != "disabled" {
1710 return Err(invalid_request(
1711 "thinking.type must be \"enabled\" or \"disabled\"",
1712 "thinking.type",
1713 ));
1714 }
1715 }
1716 for msg in &self.messages {
1717 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1718 return Err(unsupported_feature(
1719 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1720 ));
1721 }
1722 }
1723 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1724 return Err(unsupported_feature(
1725 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1726 ));
1727 }
1728 if self.n.is_some_and(|n| n > 1) {
1729 return Err(unsupported_feature(
1730 "n > 1 is not implemented (single completion only)",
1731 ));
1732 }
1733 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1734 let stated_grammar =
1749 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1750 if let Some(forced) = self.forced_tool_choice()? {
1758 if self.tools.is_empty() {
1759 return Err(invalid_request(
1760 "tool_choice forces a tool call, but no tools were offered",
1761 "tool_choice",
1762 ));
1763 }
1764 if let tool_grammar::Forced::Named(name) = forced {
1765 if !self.tools.iter().any(|t| t.function.name == name) {
1766 return Err(invalid_request(
1767 &format!(
1768 "tool_choice names {name:?}, which is not one of the tools offered"
1769 ),
1770 "tool_choice",
1771 ));
1772 }
1773 }
1774 if stated_grammar.is_some() {
1784 return Err(invalid_request(
1785 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1786 are two different constraints on the same generation; send one",
1787 "tool_choice",
1788 ));
1789 }
1790 if self.json_object_mode() {
1791 return Err(invalid_request(
1792 "a forced tool_choice cannot be combined with response_format json_object: \
1793 the tool-call markers are not JSON",
1794 "tool_choice",
1795 ));
1796 }
1797 }
1798 Ok(())
1799 }
1800
1801 fn effective_stop_sequences(&self) -> Vec<String> {
1808 let mut stop = self.stop_sequences();
1809 if self.tools_active() {
1810 stop.push("</tool_call>".to_string());
1811 }
1812 stop
1813 }
1814
1815 fn json_object_mode(&self) -> bool {
1816 self.response_format
1817 .as_ref()
1818 .and_then(|v| v.get("type"))
1819 .and_then(|v| v.as_str())
1820 == Some("json_object")
1821 }
1822
1823 fn generation_params(&self) -> Result<GenerationParams, ApiError> {
1828 Ok(GenerationParams {
1829 max_tokens: self.max_tokens,
1830 sampling: self.sampling_params(),
1831 seed: self.resolved_seed(),
1832 stop: self.effective_stop_sequences(),
1833 stop_token_ids: Vec::new(),
1837 json_object: self.json_object_mode(),
1838 grammar: grammar_request::for_request(
1839 self.grammar.as_deref(),
1840 self.response_format.as_ref(),
1841 )?,
1842 cancel: None,
1845 ignore_eos: self.ignore_eos.unwrap_or(false),
1846 })
1847 }
1848
1849 fn generation_params_for_template(
1860 &self,
1861 template: &chat_template::PromptTemplate,
1862 served_model: &str,
1863 ) -> Result<GenerationParams, ApiError> {
1864 let mut params = self.generation_params()?;
1865 if let Some(stop) = template.end_of_turn() {
1866 if !params.stop.iter().any(|s| s == stop) {
1867 params.stop.push(stop.to_string());
1868 }
1869 }
1870 if let Some(forced) = self.forced_tool_choice()? {
1871 params.grammar = Some(tool_grammar::build(
1875 forced,
1876 &self.tool_specs(),
1877 policy::parser::ToolCallFormat::infer(served_model),
1878 )?);
1879 }
1880 Ok(params)
1881 }
1882
1883 fn is_cacheable(&self) -> bool {
1892 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1893 }
1894
1895 fn cache_key(&self, prompt: &str, params: &GenerationParams) -> CacheKey {
1911 CacheKey {
1912 model: self.model.clone(),
1913 prompt: prompt.to_string(),
1914 generation: response_cache::generation_key(params),
1915 seed: self.seed,
1916 }
1917 }
1918
1919 fn resolved_seed(&self) -> u64 {
1920 self.seed.unwrap_or_else(|| {
1921 std::time::SystemTime::now()
1922 .duration_since(std::time::UNIX_EPOCH)
1923 .map(|d| d.as_nanos() as u64)
1924 .unwrap_or(0xDEFA017)
1925 })
1926 }
1927}
1928
1929#[derive(Serialize)]
1930struct ChatCompletionChoice {
1931 index: usize,
1932 message: ChatCompletionResponseMessage,
1933 finish_reason: &'static str,
1934}
1935
1936#[derive(Serialize)]
1937struct ChatCompletionResponseMessage {
1938 role: &'static str,
1939 #[serde(skip_serializing_if = "Option::is_none")]
1940 content: Option<String>,
1941 #[serde(skip_serializing_if = "Option::is_none")]
1945 reasoning_content: Option<String>,
1946 #[serde(skip_serializing_if = "Option::is_none")]
1947 tool_calls: Option<Vec<ToolCallOut>>,
1948}
1949
1950#[derive(Serialize, Clone)]
1951struct ToolCallOut {
1952 id: String,
1953 #[serde(rename = "type")]
1954 kind: &'static str,
1955 function: ToolCallFunctionOut,
1956}
1957
1958#[derive(Serialize, Clone)]
1966struct ToolCallDelta {
1967 index: usize,
1968 #[serde(skip_serializing_if = "Option::is_none")]
1969 id: Option<String>,
1970 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1971 kind: Option<&'static str>,
1972 function: ToolCallFunctionDelta,
1973}
1974
1975#[derive(Serialize, Clone, Default)]
1976struct ToolCallFunctionDelta {
1977 #[serde(skip_serializing_if = "Option::is_none")]
1978 name: Option<String>,
1979 #[serde(skip_serializing_if = "Option::is_none")]
1982 arguments: Option<String>,
1983}
1984
1985impl ToolCallDelta {
1986 fn whole(index: usize, name: String, arguments: String) -> Self {
1988 ToolCallDelta {
1989 index,
1990 id: Some(format!("call_{index}")),
1991 kind: Some("function"),
1992 function: ToolCallFunctionDelta {
1993 name: Some(name),
1994 arguments: Some(arguments),
1995 },
1996 }
1997 }
1998
1999 fn opening(index: usize, name: String) -> Self {
2001 ToolCallDelta {
2002 index,
2003 id: Some(format!("call_{index}")),
2004 kind: Some("function"),
2005 function: ToolCallFunctionDelta {
2006 name: Some(name),
2007 arguments: Some(String::new()),
2008 },
2009 }
2010 }
2011
2012 fn arguments(index: usize, fragment: String) -> Self {
2014 ToolCallDelta {
2015 index,
2016 id: None,
2017 kind: None,
2018 function: ToolCallFunctionDelta {
2019 name: None,
2020 arguments: Some(fragment),
2021 },
2022 }
2023 }
2024}
2025
2026#[derive(Serialize, Clone)]
2027struct ToolCallFunctionOut {
2028 name: String,
2029 arguments: String,
2033}
2034
2035#[derive(Serialize)]
2036struct ChatCompletionResponse {
2037 id: String,
2038 request_id: String,
2044 object: &'static str,
2045 model: String,
2046 choices: Vec<ChatCompletionChoice>,
2047 usage: generate::Usage,
2052 ferrox_cache: &'static str,
2063}
2064
2065#[derive(Serialize)]
2066struct ChatCompletionChunkDelta {
2067 #[serde(skip_serializing_if = "Option::is_none")]
2068 role: Option<&'static str>,
2069 #[serde(skip_serializing_if = "Option::is_none")]
2070 content: Option<String>,
2071 #[serde(skip_serializing_if = "Option::is_none")]
2073 reasoning_content: Option<String>,
2074 #[serde(skip_serializing_if = "Option::is_none")]
2075 tool_calls: Option<Vec<ToolCallDelta>>,
2076}
2077
2078#[derive(Serialize)]
2079struct ChatCompletionChunkChoice {
2080 index: usize,
2081 delta: ChatCompletionChunkDelta,
2082 finish_reason: Option<&'static str>,
2083}
2084
2085#[derive(Serialize)]
2086struct ChatCompletionChunk {
2087 id: String,
2088 #[serde(skip_serializing_if = "Option::is_none")]
2095 request_id: Option<String>,
2096 object: &'static str,
2097 model: String,
2098 choices: Vec<ChatCompletionChunkChoice>,
2099 #[serde(skip_serializing_if = "Option::is_none")]
2102 usage: Option<generate::Usage>,
2103}
2104
2105async fn health(State(state): State<Arc<AppState>>) -> Response {
2111 let snapshot = state.detection.snapshot();
2112 let mut capabilities = snapshot.capabilities;
2113 let active = state.active();
2114
2115 capabilities.push(match active.as_deref() {
2118 None => ferrox_api::Capability::unavailable(
2123 ferrox_api::health::capability::REAL_WEIGHTS,
2124 ferrox_api::health::reason::MODEL_NOT_LOADED,
2125 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
2126 ),
2127 Some(active) if active.is_synthetic() => ferrox_api::Capability::unavailable(
2128 ferrox_api::health::capability::REAL_WEIGHTS,
2129 ferrox_api::health::reason::MODEL_NOT_LOADED,
2130 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
2131 checkpoint. Output from this model is noise.",
2132 ),
2133 Some(active) if active.encoder().is_some() => {
2143 let endpoints = active
2144 .encoder()
2145 .map(|e| encoder_endpoints(e))
2146 .unwrap_or_default();
2147 let served_by = match endpoints.is_empty() {
2148 true => "no endpoint in this build serves it".to_string(),
2149 false => format!("served by {}", endpoints.join(" and ")),
2150 };
2151 ferrox_api::Capability::available(
2152 ferrox_api::health::capability::REAL_WEIGHTS,
2153 format!(
2154 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
2155 {served_by}; generation endpoints refuse it.",
2156 active.name(),
2157 ),
2158 )
2159 }
2160 Some(active) => ferrox_api::Capability::available(
2161 ferrox_api::health::capability::REAL_WEIGHTS,
2162 format!("Serving the real checkpoint '{}'.", active.name()),
2163 ),
2164 });
2165 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
2166 ferrox_api::Capability::available(
2167 ferrox_api::health::capability::CONTINUOUS_BATCHING,
2168 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
2169 "On by default on Metal. Concurrent requests share one batched decode worker."
2170 } else {
2171 "Concurrent requests share one batched decode step."
2172 },
2173 )
2174 } else if state.metal_private_decode_gate.is_some() {
2175 ferrox_api::Capability::unavailable(
2176 ferrox_api::health::capability::CONTINUOUS_BATCHING,
2177 ferrox_api::health::reason::DISABLED,
2178 "Off; private Metal decodes serialize (one at a time). Set FERROX_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
2179 )
2180 } else {
2181 ferrox_api::Capability::unavailable(
2182 ferrox_api::health::capability::CONTINUOUS_BATCHING,
2183 ferrox_api::health::reason::DISABLED,
2184 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
2185 )
2186 });
2187
2188 let last_request_ms = state
2189 .last_request_ms
2190 .load(std::sync::atomic::Ordering::Relaxed);
2191 let uptime = state.started_at.elapsed();
2192 let health_state = if active.is_none() {
2200 ferrox_api::HealthState::Unavailable
2201 } else {
2202 snapshot.state
2203 };
2204 let body = ferrox_api::HealthResponse {
2205 state: health_state,
2206 reason: match health_state {
2207 ferrox_api::HealthState::Ready => None,
2208 ferrox_api::HealthState::Unavailable => {
2209 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
2210 }
2211 ferrox_api::HealthState::Detecting => {
2212 Some(ferrox_api::health::reason::DETECTING.to_string())
2213 }
2214 },
2215 detail: match health_state {
2216 ferrox_api::HealthState::Ready => None,
2217 ferrox_api::HealthState::Unavailable => Some(
2218 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
2219 .to_string(),
2220 ),
2221 ferrox_api::HealthState::Detecting => {
2222 Some("Probing available compute backends.".to_string())
2223 }
2224 },
2225 model: active
2226 .as_deref()
2227 .map(|active| ferrox_api::health::ModelSummary {
2228 id: active.name().to_string(),
2229 tokenizer: active.tokenizer_kind().to_string(),
2230 synthetic_weights: active.is_synthetic(),
2231 }),
2232 capabilities,
2233 version: env!("CARGO_PKG_VERSION").to_string(),
2234 pid: std::process::id(),
2235 uptime_seconds: uptime.as_secs_f64(),
2236 server_time_unix_ms: std::time::SystemTime::now()
2237 .duration_since(std::time::UNIX_EPOCH)
2238 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
2239 .unwrap_or(0),
2240 last_request_age_seconds: (last_request_ms > 0)
2241 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
2242 .map(|age| age.max(0.0)),
2243 };
2244
2245 let status =
2246 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
2247 (status, Json(body)).into_response()
2248}
2249
2250async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2251 let Some(active) = state.active() else {
2255 return Json(serde_json::json!({ "object": "list", "data": [] }));
2256 };
2257 let mut model_entry = serde_json::json!({
2258 "id": active.name(),
2259 "object": "model",
2260 "ferrox_synthetic_weights": active.is_synthetic(),
2261 "ferrox_tokenizer": active.tokenizer_kind(),
2262 });
2263 if let Some(encoder) = active.encoder() {
2269 model_entry["ferrox_model_kind"] = serde_json::json!("embedding");
2270 model_entry["ferrox_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
2271 model_entry["ferrox_n_embd"] = serde_json::json!(encoder.n_embd());
2272 model_entry["ferrox_pooling"] = serde_json::json!(encoder.pooling_type().name());
2273 model_entry["ferrox_context_length"] = serde_json::json!(encoder.n_ctx_train());
2274 }
2275 if let Some(model) = active.generative_opt() {
2283 let parser_configured =
2284 crate::policy::parser::ReasoningFormat::infer(active.name()).is_some();
2285 let gears = model.chat_template().think_gears(parser_configured);
2286 if !gears.is_empty() {
2287 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
2288 if let Some(default) = &gears.default {
2289 model_entry["default_reasoning_effort"] = serde_json::json!(default);
2290 }
2291 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
2294 }
2295 }
2296 if let Some(mcp) = &state.mcp {
2297 model_entry["ferrox_mcp"] = mcp.models_metadata();
2298 }
2299 Json(serde_json::json!({
2300 "object": "list",
2301 "data": [model_entry]
2302 }))
2303}
2304
2305async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2319 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
2320 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
2321 let active = state.active();
2322 Json(serde_json::json!({
2323 "model": active.as_ref().map(|a| a.name()),
2324 "state": state
2325 .maintenance
2326 .lock()
2327 .unwrap_or_else(|p| p.into_inner())
2328 .state()
2329 .as_str(),
2330 "uptime_s": state.uptime().as_secs(),
2331 "throughput": {
2332 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2333 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2334 },
2335 "requests": {
2336 "active": state.cancels.live_count(),
2337 "completed": state.stats.recorded_total(),
2338 "p95_ms": state.stats.p95_duration_ms(),
2339 "ttft_mean_ms": state.stats.ttft_mean_ms(),
2340 "prompt_tokens_total": state.stats.tokens_prompt_total(),
2341 "completion_tokens_total": state.stats.tokens_generated_total(),
2342 },
2343 "pools": cache_admin::pool_gauges(&state),
2349 "memory": cache_admin::footprint_json(&state),
2352 }))
2353}
2354
2355#[derive(Deserialize)]
2356struct RequestsQuery {
2357 #[serde(default)]
2358 since: u64,
2359 #[serde(default = "default_requests_limit")]
2360 limit: usize,
2361}
2362
2363fn default_requests_limit() -> usize {
2364 stats::MAX_PAGE
2365}
2366
2367async fn recent_requests(
2375 State(state): State<Arc<AppState>>,
2376 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
2377) -> Json<serde_json::Value> {
2378 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
2379 Json(serde_json::json!({
2380 "requests": rows,
2381 "next_cursor": cursor,
2382 "missed": missed,
2383 "total": state.stats.recorded_total(),
2384 }))
2385}
2386
2387#[derive(Serialize)]
2388struct CombinedCacheStats {
2389 response_cache: response_cache::CacheStats,
2390 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
2392}
2393
2394async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
2395 Json(CombinedCacheStats {
2396 response_cache: lock_cache(&state.response_cache).stats(),
2397 prefix_cache: state
2398 .prefix_cache
2399 .as_ref()
2400 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
2401 })
2402}
2403
2404async fn metrics(State(state): State<Arc<AppState>>) -> Response {
2409 use std::sync::atomic::Ordering;
2410
2411 let cache_stats = lock_cache(&state.response_cache).stats();
2412 let active = state.active();
2413 let requests_total = state.requests_total.load(Ordering::Relaxed);
2414 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2415 let uptime = state.started_at.elapsed().as_secs_f64();
2416
2417 let body = format!(
2418 "# HELP ferrox_requests_total Total chat completion requests received.\n\
2419 # TYPE ferrox_requests_total counter\n\
2420 ferrox_requests_total {requests_total}\n\
2421 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
2422 # TYPE ferrox_request_errors_total counter\n\
2423 ferrox_request_errors_total {errors_total}\n\
2424 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
2425 # TYPE ferrox_cache_hits_total counter\n\
2426 ferrox_cache_hits_total {}\n\
2427 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
2428 # TYPE ferrox_cache_misses_total counter\n\
2429 ferrox_cache_misses_total {}\n\
2430 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
2431 # TYPE ferrox_cache_entries gauge\n\
2432 ferrox_cache_entries {}\n\
2433 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2434 # TYPE ferrox_synthetic_weights gauge\n\
2435 ferrox_synthetic_weights {}\n\
2436 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
2437 # TYPE ferrox_uptime_seconds gauge\n\
2438 ferrox_uptime_seconds {uptime}\n",
2439 cache_stats.hits,
2440 cache_stats.misses,
2441 cache_stats.entries,
2442 active
2446 .as_ref()
2447 .map(|a| a.is_synthetic() as u8)
2448 .unwrap_or(0),
2449 );
2450
2451 let body = match active
2455 .as_ref()
2456 .and_then(|a| a.expert_store_stats())
2457 {
2458 Some(es) => format!(
2459 "{body}\
2460 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
2461 # TYPE ferrox_expert_cache_hits_total counter\n\
2462 ferrox_expert_cache_hits_total {}\n\
2463 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2464 # TYPE ferrox_expert_cache_misses_total counter\n\
2465 ferrox_expert_cache_misses_total {}\n\
2466 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
2467 # TYPE ferrox_expert_cache_evictions_total counter\n\
2468 ferrox_expert_cache_evictions_total {}\n\
2469 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2470 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
2471 ferrox_expert_cache_pass_throughs_total {}\n\
2472 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2473 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
2474 ferrox_expert_cache_bytes_read_total {}\n\
2475 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2476 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
2477 ferrox_expert_cache_resident_bytes {}\n",
2478 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2479 ),
2480 None => body,
2481 };
2482
2483 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2488 Some(batcher) => {
2489 let sched = batcher.stats();
2490 format!(
2491 "{body}\
2492 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2493 # TYPE ferrox_prefill_chunks_total counter\n\
2494 ferrox_prefill_chunks_total {}\n\
2495 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2496 # TYPE ferrox_prefill_tokens_total counter\n\
2497 ferrox_prefill_tokens_total {}\n\
2498 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2499 # TYPE ferrox_decode_steps_total counter\n\
2500 ferrox_decode_steps_total {}\n\
2501 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2502 # TYPE ferrox_scheduler_queue_depth gauge\n\
2503 ferrox_scheduler_queue_depth {}\n\
2504 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2505 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2506 ferrox_scheduler_queue_rejected_total {}\n\
2507 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2508 # TYPE ferrox_kv_blocks_total gauge\n\
2509 ferrox_kv_blocks_total {}\n\
2510 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2511 # TYPE ferrox_kv_blocks_free gauge\n\
2512 ferrox_kv_blocks_free {}\n\
2513 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2514 # TYPE ferrox_kv_block_size gauge\n\
2515 ferrox_kv_block_size {}\n\
2516 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2517 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2518 ferrox_kv_rejected_too_large_total {}\n\
2519 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2520 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2521 ferrox_kv_rejected_context_length_total {}\n\
2522 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2523 # TYPE ferrox_scheduler_aborted_total counter\n\
2524 ferrox_scheduler_aborted_total {}\n",
2525 sched.prefill_chunks,
2526 sched.prefill_tokens,
2527 sched.decode_steps,
2528 sched.queue_depth,
2529 sched.queue_rejected,
2530 sched.kv_blocks_total,
2531 sched.kv_blocks_free,
2532 sched.kv_block_size,
2533 sched.kv_rejected_too_large,
2534 sched.kv_rejected_context_length,
2535 sched.aborted,
2536 )
2537 }
2538 None => body,
2539 };
2540
2541 (
2542 [(
2543 axum::http::header::CONTENT_TYPE,
2544 "text/plain; version=0.0.4",
2545 )],
2546 body,
2547 )
2548 .into_response()
2549}
2550
2551pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2552
2553pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2559 (
2560 StatusCode::BAD_REQUEST,
2561 Json(serde_json::json!({"error": {
2562 "message": message,
2563 "type": "invalid_request_error",
2564 "param": param,
2565 "code": null,
2566 }})),
2567 )
2568}
2569
2570pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2571 (
2572 StatusCode::NOT_IMPLEMENTED,
2573 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2574 )
2575}
2576
2577pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2578 let status = match e {
2579 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2580 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2585 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2591 StatusCode::SERVICE_UNAVAILABLE
2592 }
2593 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2597 };
2598 tracing::warn!("decode error: {e}");
2599 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2600 if let generate::DecodeError::KvBudgetExceeded {
2606 binding,
2607 estimated_bytes,
2608 limit_bytes,
2609 positions,
2610 positions_limit,
2611 ..
2612 } = &e
2613 {
2614 body["error"]["type"] = serde_json::json!("invalid_request_error");
2615 body["error"]["code"] = serde_json::json!(binding);
2616 body["error"]["binding"] = serde_json::json!(binding);
2617 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2618 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2619 body["error"]["positions"] = serde_json::json!(positions);
2620 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2621 }
2622 if let Some(secs) = e.retry_after_secs() {
2626 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2627 }
2628 (status, Json(body))
2629}
2630
2631pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2632 tracing::error!("generation task panicked: {e}");
2633 (
2634 StatusCode::INTERNAL_SERVER_ERROR,
2635 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2636 )
2637}
2638
2639#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2649 model: &Model,
2650 prompt: &str,
2651 params: &GenerationParams,
2652 kv_pool: Option<&generate::KvPoolConfig>,
2653 paged_kv: Option<&generate::PagedKvConfig>,
2654 prefix_cache: Option<&Mutex<PrefixCache>>,
2655 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2656 ceiling: Option<&budget::ContextCeiling>,
2657 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2658 mut emit: impl FnMut(&str),
2659) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2660 let synthetic = model.is_synthetic();
2661 let mut chunks = Vec::new();
2662 let params = &{
2668 let mut resolved = params.clone();
2669 resolved.stop_token_ids =
2670 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
2671 resolved
2672 };
2673 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2674 let _metal_private_guard =
2675 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2676 let (finish, usage) = match model {
2677 Model::Gguf(m) => {
2678 if let Some(batcher) = continuous_batcher {
2679 let mut tokens = m.tokenizer.encode(prompt);
2680 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2681 let (finish, _generated_ids, text, usage) = if synthetic {
2682 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2683 } else {
2684 batcher.generate_streaming(
2685 tokens,
2686 params.clone(),
2687 m.stop_tokens.clone(),
2688 Some(|chunk: &str| {
2689 if !chunk.is_empty() {
2690 chunks.push(chunk.to_string());
2691 emit(chunk);
2692 }
2693 }),
2694 )?
2695 };
2696 if !text.is_empty() && chunks.is_empty() {
2697 chunks.push(text);
2698 }
2699 (finish, usage)
2700 } else {
2701 generate::generate(
2702 &m.decoder,
2703 m.tokenizer.as_ref(),
2704 &m.stop_tokens,
2705 m.bos_id,
2706 prompt,
2707 params,
2708 kv_pool,
2709 paged_kv,
2710 prefix_cache,
2711 ceiling,
2712 |chunk| {
2713 chunks.push(chunk.to_string());
2714 if !synthetic {
2715 emit(chunk);
2716 }
2717 },
2718 )?
2719 }
2720 }
2721 Model::Kimi(m) => generate::generate_engine(
2722 &m.engine,
2723 &m.tokenizer,
2724 &m.stop_tokens,
2725 None,
2726 prompt,
2727 params,
2728 |chunk| {
2729 chunks.push(chunk.to_string());
2730 if !synthetic {
2731 emit(chunk);
2732 }
2733 },
2734 )?,
2735 Model::Mla(m) => generate::generate_engine(
2736 &m.engine,
2737 &m.tokenizer,
2738 &m.stop_tokens,
2739 m.bos_id,
2740 prompt,
2741 params,
2742 |chunk| {
2743 chunks.push(chunk.to_string());
2744 if !synthetic {
2745 emit(chunk);
2746 }
2747 },
2748 )?,
2749 Model::Gemma4(m) => generate::generate_engine(
2750 &m.engine,
2751 &m.tokenizer,
2752 &m.stop_tokens,
2753 m.bos_id,
2754 prompt,
2755 params,
2756 |chunk| {
2757 chunks.push(chunk.to_string());
2758 if !synthetic {
2759 emit(chunk);
2760 }
2761 },
2762 )?,
2763 Model::Glm52(m) => generate::generate_engine(
2764 &m.engine,
2765 &m.tokenizer,
2766 &m.stop_tokens,
2767 m.bos_id,
2768 prompt,
2769 params,
2770 |chunk| {
2771 chunks.push(chunk.to_string());
2772 if !synthetic {
2773 emit(chunk);
2774 }
2775 },
2776 )?,
2777 };
2778
2779 let mut full = chunks.concat();
2780 if synthetic {
2781 full = format!(
2782 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2783 to serve a real model. Decoded ids -> {full:?}]"
2784 );
2785 emit(&full);
2786 } else if used_batcher && !full.is_empty() && chunks.is_empty() {
2787 emit(&full);
2788 }
2789
2790 Ok((finish, usage, full))
2791}
2792
2793#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2798 model: &Model,
2799 prompt: &str,
2800 params: &GenerationParams,
2801 kv_pool: Option<&generate::KvPoolConfig>,
2802 paged_kv: Option<&generate::PagedKvConfig>,
2803 prefix_cache: Option<&Mutex<PrefixCache>>,
2804 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2805 ceiling: Option<&budget::ContextCeiling>,
2806 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2807) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2808 let (finish, usage, full) = run_generation_emit(
2809 model,
2810 prompt,
2811 params,
2812 kv_pool,
2813 paged_kv,
2814 prefix_cache,
2815 continuous_batcher,
2816 ceiling,
2817 metal_private_decode_gate,
2818 |_| {},
2819 )?;
2820 Ok((
2821 if full.is_empty() {
2822 Vec::new()
2823 } else {
2824 vec![full]
2825 },
2826 finish,
2827 usage,
2828 ))
2829}
2830
2831pub(crate) fn prompt_from_messages(
2842 messages: &[ChatMessage],
2843 template: &chat_template::PromptTemplate,
2844 tools: &[ToolDef],
2845 extra: serde_json::Map<String, serde_json::Value>,
2846) -> Result<String, ApiError> {
2847 let rendered = if tools.is_empty() || template.handles_tools() {
2848 template.render(messages, tools, extra)
2849 } else {
2850 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2851 with_preamble.push(ChatMessage {
2852 role: "system".to_string(),
2853 content: Some(MessageContent::Text(tool_preamble(tools))),
2854 tool_calls: None,
2855 tool_call_id: None,
2856 reasoning_content: None,
2857 });
2858 with_preamble.extend_from_slice(messages);
2859 template.render(&with_preamble, &[], extra)
2860 };
2861 rendered.map_err(template_error_response)
2862}
2863
2864fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2869 (
2870 StatusCode::BAD_REQUEST,
2871 Json(serde_json::json!({
2872 "error": {
2873 "message": format!("chat template failed to render: {err}"),
2874 "type": "invalid_request_error",
2875 "param": "messages",
2876 "code": null,
2877 }
2878 })),
2879 )
2880}
2881
2882fn tool_preamble(tools: &[ToolDef]) -> String {
2895 let mut out = String::from(
2896 "You can call tools to help answer the user. To call a tool, respond with \
2897 EXACTLY one line in this format and nothing else:\n\
2898 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2899 object matching that tool's parameters>}}</tool_call>\n\n\
2900 Available tools:\n",
2901 );
2902 for t in tools {
2903 out.push_str(&format!(
2904 "- {}: {}\n parameters (JSON schema): {}\n",
2905 t.function.name,
2906 t.function.description.as_deref().unwrap_or(""),
2907 t.function
2908 .parameters
2909 .as_ref()
2910 .map(|v| v.to_string())
2911 .unwrap_or_else(|| "{}".to_string()),
2912 ));
2913 }
2914 out
2915}
2916
2917fn tool_call_deltas(
2927 events: Vec<crate::policy::parser::ToolCallEvent>,
2928 opened: &std::cell::Cell<usize>,
2929) -> (String, Vec<ToolCallDelta>) {
2930 let mut text = String::new();
2931 let mut deltas = Vec::new();
2932 for event in events {
2933 match event {
2934 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2935 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2936 opened.set(opened.get().max(index + 1));
2937 deltas.push(ToolCallDelta::opening(index, name));
2938 }
2939 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2940 if !fragment.is_empty() {
2941 deltas.push(ToolCallDelta::arguments(index, fragment));
2942 }
2943 }
2944 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2945 }
2946 }
2947 (text, deltas)
2948}
2949
2950fn build_response_message(
2961 text: String,
2962 tools: &[ToolDef],
2963 posture: output::OutputPosture,
2964 base_finish: &'static str,
2965) -> (ChatCompletionResponseMessage, &'static str) {
2966 let parsed = output::parse_output(&text, tools, posture);
2967 let calls: Vec<ToolCallOut> = parsed
2968 .calls
2969 .into_iter()
2970 .enumerate()
2971 .map(|(index, call)| ToolCallOut {
2972 id: format!("call_{index}"),
2973 kind: "function",
2974 function: ToolCallFunctionOut {
2975 name: call.name,
2976 arguments: call.arguments,
2977 },
2978 })
2979 .collect();
2980 if !calls.is_empty() {
2981 return (
2982 ChatCompletionResponseMessage {
2983 role: "assistant",
2984 content: None,
2985 reasoning_content: parsed.reasoning,
2986 tool_calls: Some(calls),
2987 },
2988 "tool_calls",
2989 );
2990 }
2991 (
2992 ChatCompletionResponseMessage {
2993 role: "assistant",
2994 content: Some(parsed.content),
2995 reasoning_content: parsed.reasoning,
2996 tool_calls: None,
2997 },
2998 base_finish,
2999 )
3000}
3001
3002fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
3007 let mut history = match &req.session_id {
3008 Some(id) => state.sessions.extend_and_get(id, &req.messages),
3009 None => req.messages.clone(),
3010 };
3011 if req.json_object_mode() {
3012 inject_json_object_system_hint(&mut history);
3013 }
3014 history
3015}
3016
3017fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
3018 const HINT: &str =
3019 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
3020 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
3021 match &mut sys.content {
3022 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
3023 s.push_str("\n\n");
3024 s.push_str(HINT);
3025 }
3026 None => {
3027 sys.content = Some(MessageContent::Text(HINT.to_string()));
3028 }
3029 _ => {}
3030 }
3031 } else {
3032 messages.insert(
3033 0,
3034 ChatMessage {
3035 role: "system".to_string(),
3036 content: Some(MessageContent::Text(HINT.to_string())),
3037 tool_calls: None,
3038 tool_call_id: None,
3039 reasoning_content: None,
3040 },
3041 );
3042 }
3043}
3044
3045async fn chat_completions(
3046 State(state): State<Arc<AppState>>,
3047 headers: axum::http::HeaderMap,
3048 Json(req): Json<ChatCompletionRequest>,
3049) -> Response {
3050 let attribution = attribution::Attribution::from_headers(&headers);
3051 state
3052 .requests_total
3053 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3054 let started = std::time::Instant::now();
3055
3056 let request_id = ferrox_api::next_request_id();
3060 let stream = req.stream.unwrap_or(false);
3061
3062 let refusal = cache_admin::check_admission(&state)
3068 .err()
3069 .or_else(|| req.validate_supported_fields().err());
3070 if let Some(err) = refusal {
3071 state
3072 .request_errors_total
3073 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3074 let response = err.into_response();
3075 state.record_request(stats::Record {
3076 request_id: &request_id,
3077 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3078 model: state.active_model_name(),
3079 status: response.status().as_u16(),
3080 stream,
3081 duration_ms: started.elapsed().as_millis() as u64,
3082 usage: None,
3083 attribution: &attribution,
3084 });
3085 return response;
3086 }
3087
3088 let response = if stream {
3089 chat_completions_stream(
3090 Arc::clone(&state),
3091 req,
3092 request_id.clone(),
3093 started,
3094 attribution.clone(),
3095 )
3096 .await
3097 .into_response()
3098 } else {
3099 chat_completions_full(
3100 Arc::clone(&state),
3101 req,
3102 request_id.clone(),
3103 started,
3104 attribution.clone(),
3105 )
3106 .await
3107 .into_response()
3108 };
3109
3110 if response.status().is_client_error() || response.status().is_server_error() {
3111 state
3112 .request_errors_total
3113 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3114 state.record_request(stats::Record {
3118 request_id: &request_id,
3119 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3120 model: state.active_model_name(),
3123 status: response.status().as_u16(),
3124 stream,
3125 duration_ms: started.elapsed().as_millis() as u64,
3126 usage: None,
3127 attribution: &attribution,
3128 });
3129 }
3130 state.mark_request_finished();
3131
3132 response
3133}
3134
3135async fn chat_completions_full(
3136 state: Arc<AppState>,
3137 req: ChatCompletionRequest,
3138 request_id: String,
3139 started: std::time::Instant,
3140 attribution: attribution::Attribution,
3141) -> Result<Json<ChatCompletionResponse>, ApiError> {
3142 let tools_active = req.tools_active();
3143 let active = state.require_active()?;
3147 let history = resolve_history(&state, &req);
3148 let template = active.generative()?.chat_template();
3149 let kwargs = req.resolve_template_kwargs(&template);
3150 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
3151 let params = req.generation_params_for_template(&template, active.name())?;
3159 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
3160
3161 let (completion, cache_status) = if let Some(cached) = key
3162 .as_ref()
3163 .and_then(|key| lock_cache(&state.response_cache).get(key))
3164 {
3165 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
3166 (cached, "hit")
3167 } else {
3168 let (chunks, finish, usage) = decode_task::buffered(
3169 decode_task::DecodeHandles::take(&state, &active)?,
3170 prompt.clone(),
3171 params,
3172 )
3173 .await?;
3174
3175 let completion = response_cache::CachedCompletion {
3176 content: chunks.concat(),
3177 finish,
3178 usage,
3179 };
3180 let cache_status = match key {
3186 Some(key) => match completion.clone().cacheable() {
3190 Some(cacheable) => {
3191 tracing::debug!("cache miss for key {}", key.digest());
3192 lock_cache(&state.response_cache).put(key, cacheable);
3193 "miss"
3194 }
3195 None => "skip",
3196 },
3197 None => "skip",
3198 };
3199 (completion, cache_status)
3200 };
3201 let content = completion.content;
3202
3203 if req.json_object_mode() {
3204 json_mode::validate_json_object_output(&content)?;
3205 }
3206
3207 if let Some(id) = &req.session_id {
3211 state.sessions.store_reply(
3212 id,
3213 ChatMessage {
3214 role: "assistant".to_string(),
3215 content: Some(MessageContent::Text(content.clone())),
3216 tool_calls: None,
3217 tool_call_id: None,
3218 reasoning_content: None,
3219 },
3220 );
3221 }
3222
3223 let (message, finish_reason) = build_response_message(
3224 content,
3225 if tools_active { &req.tools } else { &[] },
3226 output::OutputPosture::resolve(active.name(), &prompt),
3227 completion.finish.as_str(),
3228 );
3229
3230 state.record_request(stats::Record {
3231 request_id: &request_id,
3232 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3233 model: Some(active.name().to_string()),
3236 status: 200,
3237 stream: false,
3238 duration_ms: started.elapsed().as_millis() as u64,
3239 usage: Some(&completion.usage),
3240 attribution: &attribution,
3241 });
3242
3243 Ok(Json(ChatCompletionResponse {
3244 id: request_id.clone(),
3245 request_id,
3246 object: "chat.completion",
3247 model: req.model,
3248 choices: vec![ChatCompletionChoice {
3249 index: 0,
3250 message,
3251 finish_reason,
3252 }],
3253 usage: completion.usage,
3254 ferrox_cache: cache_status,
3255 }))
3256}
3257
3258async fn chat_completions_stream(
3259 state: Arc<AppState>,
3260 req: ChatCompletionRequest,
3261 request_id: String,
3262 started: std::time::Instant,
3263 attribution: attribution::Attribution,
3264) -> Result<Response, ApiError> {
3265 let tools_active = req.tools_active();
3267 let active = state.require_active()?;
3271 let history = resolve_history(&state, &req);
3272 let template = active.generative()?.chat_template();
3273 let kwargs = req.resolve_template_kwargs(&template);
3274 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
3275 let model_name = req.model.clone();
3276 let session_id = req.session_id.clone();
3277 let sessions = state.sessions.clone();
3278
3279 let model = Arc::clone(active.generative()?);
3280 let kv_pool = state.kv_pool.clone();
3281 let paged_kv = state.paged_kv.clone();
3282 let prefix_cache = state.prefix_cache.clone();
3283 let batcher = active.batcher.clone();
3284 let ceiling = active.ceiling.clone();
3285 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3286 let mut params = req.generation_params_for_template(&template, active.name())?;
3287 let stats_state = Arc::clone(&state);
3288 let served_model = active.name().to_string();
3291 let posture = output::OutputPosture::resolve(&served_model, &prompt);
3295 let offered_tools: Vec<ToolDef> = if tools_active {
3298 req.tools.clone()
3299 } else {
3300 Vec::new()
3301 };
3302
3303 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3308 params.cancel = Some(cancel_token.clone());
3309
3310 let overlap = true;
3319
3320 let slot = req
3324 .stream_resumable
3325 .unwrap_or(false)
3326 .then(|| state.streams.register(&request_id));
3327 let emitter = resume::Emitter::new(slot);
3328
3329 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3330 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3334 id: request_id.clone(),
3335 request_id: None,
3336 object: "chat.completion.chunk",
3337 model: model_name.clone(),
3338 choices: vec![ChatCompletionChunkChoice {
3339 index: 0,
3340 delta: ChatCompletionChunkDelta {
3341 role: None,
3342 content: None,
3343 reasoning_content: None,
3344 tool_calls: None,
3345 },
3346 finish_reason: None,
3347 }],
3348 usage: None,
3349 });
3350
3351 tokio::task::spawn_blocking(move || {
3352 let _cancel_guard = cancel_guard;
3355 let tx_chunks = tx.clone();
3356 let orphan_timeout = sse::orphan_timeout_from_env();
3361 let mut first = true;
3362 let head_request_id = request_id.clone();
3363 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3373 Rc::new(RefCell::new(posture.reasoning_parser()));
3374 let emit_reasoning = Rc::clone(&stream_reasoning);
3375 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3379 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3380 );
3381 let emit_tools = Rc::clone(&stream_tools);
3382 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3386 let emit_streamed_calls = Rc::clone(&streamed_calls);
3387 let result = run_generation_emit(
3388 &model,
3389 &prompt,
3390 ¶ms,
3391 kv_pool.as_ref(),
3392 paged_kv.as_ref(),
3393 prefix_cache.as_deref(),
3394 batcher.as_ref(),
3395 ceiling.as_deref(),
3396 metal_private_decode_gate.as_deref(),
3397 |chunk| {
3398 if !overlap || chunk.is_empty() {
3399 return;
3400 }
3401 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3402 Some(parser) => {
3403 let delta = parser.push(chunk);
3404 (delta.reasoning, delta.content)
3405 }
3406 None => (String::new(), chunk.to_string()),
3407 };
3408 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3412 Some(parser) => {
3413 let (text, calls) =
3414 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3415 (text, calls)
3416 }
3417 None => (content, Vec::new()),
3418 };
3419 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3422 return;
3423 }
3424 let role = if first { Some("assistant") } else { None };
3425 let request_id = first.then(|| head_request_id.clone());
3426 first = false;
3427 let payload = ChatCompletionChunk {
3428 id: head_request_id.clone(),
3429 request_id,
3430 object: "chat.completion.chunk",
3431 model: model_name.clone(),
3432 choices: vec![ChatCompletionChunkChoice {
3433 index: 0,
3434 delta: ChatCompletionChunkDelta {
3435 role,
3436 content: (!content.is_empty()).then_some(content),
3437 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3438 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3439 },
3440 finish_reason: None,
3441 }],
3442 usage: None,
3443 };
3444 if let Err(why) =
3453 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3454 {
3455 if why == sse::SendFailure::Orphaned {
3456 tracing::warn!(
3457 "SSE stream {head_request_id} accepted nothing for the orphan \
3458 deadline; treating it as abandoned"
3459 );
3460 }
3461 if !emitter.is_resumable() {
3473 cancel_token.cancel();
3474 }
3475 }
3476 },
3477 );
3478
3479 let mut pending_request_id = first.then(|| request_id.clone());
3485
3486 match result {
3487 Ok((finish, usage, full_text)) => {
3488 if let Some(id) = &session_id {
3489 sessions.store_reply(
3490 id,
3491 ChatMessage {
3492 role: "assistant".to_string(),
3493 content: Some(MessageContent::Text(full_text.clone())),
3494 tool_calls: None,
3495 tool_call_id: None,
3496 reasoning_content: None,
3497 },
3498 );
3499 }
3500 let mut streamed_finish: Option<&'static str> = None;
3506 if overlap {
3507 let tail = stream_reasoning
3508 .borrow_mut()
3509 .as_mut()
3510 .map(|parser| parser.flush())
3511 .unwrap_or_default();
3512 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3513 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3514 let mut events = parser.push(&content);
3515 events.extend(parser.finish());
3516 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3517 content = text;
3518 tool_calls = calls;
3519 }
3520 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3521 let payload = ChatCompletionChunk {
3522 id: request_id.clone(),
3523 request_id: pending_request_id.take(),
3524 object: "chat.completion.chunk",
3525 model: model_name.clone(),
3526 choices: vec![ChatCompletionChunkChoice {
3527 index: 0,
3528 delta: ChatCompletionChunkDelta {
3529 role: None,
3530 content: (!content.is_empty()).then_some(content),
3531 reasoning_content: (!tail.reasoning.is_empty())
3532 .then_some(tail.reasoning),
3533 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3534 },
3535 finish_reason: None,
3536 }],
3537 usage: None,
3538 };
3539 let _ =
3540 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3541 }
3542 if streamed_calls.get() > 0 {
3543 streamed_finish = Some("tool_calls");
3544 }
3545 } else {
3546 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3549 let tool_calls: Vec<ToolCallDelta> = parsed
3550 .calls
3551 .iter()
3552 .enumerate()
3553 .map(|(index, call)| {
3554 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3555 })
3556 .collect();
3557 if !tool_calls.is_empty() {
3558 streamed_finish = Some("tool_calls");
3559 }
3560 if !tool_calls.is_empty()
3561 || !parsed.content.is_empty()
3562 || parsed.reasoning.is_some()
3563 {
3564 let payload = ChatCompletionChunk {
3565 id: request_id.clone(),
3566 request_id: pending_request_id.take(),
3567 object: "chat.completion.chunk",
3568 model: model_name.clone(),
3569 choices: vec![ChatCompletionChunkChoice {
3570 index: 0,
3571 delta: ChatCompletionChunkDelta {
3572 role: Some("assistant"),
3573 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3574 .then(|| parsed.content.clone()),
3575 reasoning_content: parsed.reasoning.clone(),
3576 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3577 },
3578 finish_reason: None,
3579 }],
3580 usage: None,
3581 };
3582 let _ =
3583 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3584 }
3585 }
3586 let final_finish_reason = match streamed_finish {
3590 Some(reason) if finish.as_str() != "length" => reason,
3591 _ => finish.as_str(),
3592 };
3593 let final_payload = ChatCompletionChunk {
3594 id: request_id.clone(),
3595 request_id: pending_request_id.take(),
3596 object: "chat.completion.chunk",
3597 model: model_name,
3598 choices: vec![ChatCompletionChunkChoice {
3599 index: 0,
3600 delta: ChatCompletionChunkDelta {
3601 role: None,
3602 content: None,
3603 reasoning_content: None,
3604 tool_calls: None,
3605 },
3606 finish_reason: Some(final_finish_reason),
3607 }],
3608 usage: Some(usage.clone()),
3609 };
3610 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3611 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3612 stats_state.record_request(stats::Record {
3617 request_id: &request_id,
3618 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3619 model: Some(served_model.clone()),
3620 status: 200,
3621 stream: true,
3622 duration_ms: started.elapsed().as_millis() as u64,
3623 usage: Some(&usage),
3624 attribution: &attribution,
3625 });
3626 }
3627 Err(e) => {
3628 tracing::warn!("decode error on streamed request {request_id}: {e}");
3629 stats_state.record_request(stats::Record {
3635 request_id: &request_id,
3636 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3637 model: Some(served_model.clone()),
3638 status: 500,
3639 stream: true,
3640 duration_ms: started.elapsed().as_millis() as u64,
3641 usage: None,
3642 attribution: &attribution,
3643 });
3644 let payload = ChatCompletionChunk {
3645 id: request_id.clone(),
3646 request_id: pending_request_id.take(),
3647 object: "chat.completion.chunk",
3648 model: model_name,
3649 choices: vec![ChatCompletionChunkChoice {
3650 index: 0,
3651 delta: ChatCompletionChunkDelta {
3652 role: Some("assistant"),
3653 content: Some(format!("[error: {e}]")),
3654 reasoning_content: None,
3655 tool_calls: None,
3656 },
3657 finish_reason: Some("stop"),
3658 }],
3659 usage: None,
3660 };
3661 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3662 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3663 }
3664 }
3665 drop(emitter);
3669 });
3670
3671 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3672 Ok((
3695 [(
3696 axum::http::HeaderName::from_static("x-accel-buffering"),
3697 axum::http::HeaderValue::from_static("no"),
3698 )],
3699 Sse::new(stream),
3700 )
3701 .into_response())
3702}
3703
3704fn axum_path(template: &str) -> String {
3721 let mut out = String::with_capacity(template.len());
3722 let mut rest = template;
3723 while let Some(open) = rest.find('{') {
3724 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3725 break;
3726 };
3727 out.push_str(&rest[..open]);
3728 out.push(':');
3729 out.push_str(&rest[open + 1..close]);
3730 rest = &rest[close + 1..];
3731 }
3732 out.push_str(rest);
3733 out
3734}
3735
3736async fn cancel_generation(
3745 State(state): State<Arc<AppState>>,
3746 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3747) -> Response {
3748 let cancelled = state.cancels.cancel(&req.request_id);
3749 let status = if cancelled {
3750 StatusCode::OK
3751 } else {
3752 StatusCode::NOT_FOUND
3753 };
3754 let detail = if cancelled {
3755 "the generation was asked to stop; it ends at its next token".to_string()
3756 } else {
3757 "no generation with that request_id is running -- it has already \
3758 finished, was never issued, or was served by a path that does \
3759 not register for cancellation"
3760 .to_string()
3761 };
3762 (
3763 status,
3764 Json(ferrox_api::CancelGenerationResponse {
3765 request_id: req.request_id,
3766 cancelled,
3767 detail,
3768 }),
3769 )
3770 .into_response()
3771}
3772
3773type Activated = (
3777 Loaded,
3778 Option<serving::batch::ContinuousBatcher>,
3779 Option<Arc<budget::ContextCeiling>>,
3780);
3781
3782fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3800 let mut batcher = serving::batch::BatcherConfig::from_env();
3801 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3802 return batcher;
3805 }
3806 let Some(path) = path else {
3807 return batcher;
3808 };
3809 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1);
3815 let Some((priced, gguf_ctx, source)) = priced else {
3816 return batcher;
3817 };
3818 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3819 tracing::warn!(
3823 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3824 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3825 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3826 admit on a number you choose.",
3827 priced.weights_bytes,
3828 priced.device_budget_bytes,
3829 );
3830 return batcher;
3831 };
3832 tracing::info!("{source}");
3833 tracing::info!("{}", derived.fit);
3834 let adopted = budget::apply_derived(&mut batcher, &derived);
3835 if adopted.max_context {
3836 tracing::info!(
3837 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3838 override with FERROX_CB_MAX_CONTEXT",
3839 derived.max_context
3840 );
3841 }
3842 if adopted.kv_blocks {
3843 tracing::info!(
3844 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3845 derived.kv_blocks,
3846 batcher.kv_block_size
3847 );
3848 }
3849 batcher
3850}
3851
3852pub(crate) fn activate_loaded_model(
3861 loaded: model::LoadedModel,
3862 enable_continuous_batching: bool,
3863 path: Option<&str>,
3864 paged_kv: Option<&generate::PagedKvConfig>,
3865) -> Activated {
3866 match loaded {
3867 model::LoadedModel::Gguf(g) => {
3868 let decoder = Arc::new(g.decoder);
3869 let tokenizer = Arc::new(g.tokenizer);
3870 let config = price_batcher_config(path);
3871 let shape =
3877 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32);
3878 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3879 let batcher = if enable_continuous_batching {
3880 tracing::info!(
3881 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3882 (stop sequences use the same pending-buffer trim as the private generate loop)"
3883 );
3884 let tok = Arc::clone(&tokenizer);
3885 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
3886 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3887 Arc::clone(&decoder),
3888 decode,
3889 config,
3890 Arc::clone(&ceiling),
3891 paged_kv.cloned(),
3892 ))
3893 } else {
3894 None
3895 };
3896 (
3897 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3898 decoder,
3899 tokenizer,
3900 stop_tokens: g.stop_tokens,
3901 bos_id: g.bos_id,
3902 is_synthetic: g.is_synthetic,
3903 chat_template: g.chat_template,
3904 }))),
3905 batcher,
3906 Some(ceiling),
3907 )
3908 }
3909 model::LoadedModel::Kimi(k) => (
3910 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3911 engine: k.engine,
3912 tokenizer: k.tokenizer,
3913 stop_tokens: k.stop_tokens,
3914 chat_template: k.chat_template,
3915 }))),
3916 None,
3917 None,
3918 ),
3919 model::LoadedModel::Mla(m) => (
3920 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3921 engine: m.engine,
3922 tokenizer: m.tokenizer,
3923 stop_tokens: m.stop_tokens,
3924 bos_id: m.bos_id,
3925 name: m.name,
3926 chat_template: m.chat_template,
3927 }))),
3928 None,
3929 None,
3930 ),
3931 model::LoadedModel::Gemma4(m) => (
3932 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3933 engine: m.engine,
3934 tokenizer: m.tokenizer,
3935 stop_tokens: m.stop_tokens,
3936 bos_id: m.bos_id,
3937 name: m.name,
3938 chat_template: m.chat_template,
3939 }))),
3940 None,
3941 None,
3942 ),
3943 model::LoadedModel::Glm52(g) => (
3944 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3945 engine: g.engine,
3946 tokenizer: g.tokenizer,
3947 stop_tokens: g.stop_tokens,
3948 bos_id: g.bos_id,
3949 name: g.name,
3950 chat_template: g.chat_template,
3951 }))),
3952 None,
3953 None,
3954 ),
3955 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3960 }
3961}
3962
3963struct StartupModels {
3970 loaded: model::LoadedModel,
3971 embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
3972}
3973
3974fn continuous_batching_env() -> Option<bool> {
3975 match std::env::var("FERROX_CONTINUOUS_BATCHING")
3976 .ok()
3977 .map(|v| v.trim().to_ascii_lowercase())
3978 .as_deref()
3979 {
3980 None => None,
3981 Some("1" | "true" | "yes" | "on") => Some(true),
3982 Some("0" | "false" | "no" | "off") => Some(false),
3983 _ => None,
3984 }
3985}
3986
3987fn metal_private_decode_active() -> bool {
3988 #[cfg(feature = "metal")]
3989 {
3990 BUILT_WITH_METAL
3991 && ferrox_metal::attn::metal_attn_enabled()
3992 && std::env::var("FERROX_METAL").ok().as_deref() != Some("0")
3993 }
3994 #[cfg(not(feature = "metal"))]
3995 {
3996 false
3997 }
3998}
3999
4000fn continuous_batching_compatible(
4001 loaded: &model::LoadedModel,
4002 kv_pool: &Option<generate::KvPoolConfig>,
4003 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
4004 paged_kv: &Option<generate::PagedKvConfig>,
4005) -> bool {
4006 matches!(loaded, model::LoadedModel::Gguf(_))
4007 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
4008}
4009
4010fn resolve_continuous_batching_enabled(
4011 loaded: &model::LoadedModel,
4012 kv_pool: &Option<generate::KvPoolConfig>,
4013 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
4014 paged_kv: &Option<generate::PagedKvConfig>,
4015) -> bool {
4016 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
4017 return false;
4018 }
4019 match continuous_batching_env() {
4020 Some(true) => true,
4021 Some(false) => false,
4022 None => metal_private_decode_active(),
4023 }
4024}
4025
4026fn acquire_metal_private_decode_gate(
4027 gate: Option<&std::sync::Mutex<()>>,
4028 used_batcher: bool,
4029) -> Option<std::sync::MutexGuard<'_, ()>> {
4030 if used_batcher {
4031 None
4032 } else {
4033 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
4034 }
4035}
4036
4037fn build_app_state(
4038 models: StartupModels,
4039 kv_pool: Option<generate::KvPoolConfig>,
4040 paged_kv: Option<generate::PagedKvConfig>,
4041 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
4042 enable_continuous_batching: bool,
4043 mcp: Option<mcp::LoadedMcpConfig>,
4044 detection: Arc<health::Detection>,
4045) -> AppState {
4046 let StartupModels { loaded, embedding } = models;
4047 let (loaded, batcher, ceiling) = activate_loaded_model(
4048 loaded,
4049 enable_continuous_batching,
4050 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
4051 paged_kv.as_ref(),
4052 );
4053 let id = startup_model_id();
4059 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
4060 {
4061 None
4062 } else {
4063 tracing::info!(
4064 "Metal private-loop decode will serialize concurrent requests until \
4065 continuous batching is enabled (FERROX_CONTINUOUS_BATCHING=1 or --cont-batching)"
4066 );
4067 Some(Arc::new(std::sync::Mutex::new(())))
4068 };
4069 AppState {
4070 embedding,
4071 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4072 id,
4073 loaded,
4074 batcher,
4075 ceiling,
4076 }))),
4077 paged_kv,
4078 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4079 tasks: Arc::new(tasks::TaskRegistry::new()),
4080 cancels: Arc::new(cancel::CancelRegistry::new()),
4081 stats: stats::Stats::new(),
4082 streams: resume::StreamRegistry::new(),
4083 model_dir: admin::model_dirs().into_iter().next(),
4084 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
4085 kv_pool,
4086 prefix_cache,
4087 sessions: session::SessionStore::new(),
4088 requests_total: std::sync::atomic::AtomicU64::new(0),
4089 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4090 started_at: std::time::Instant::now(),
4091 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4092 detection,
4093 mcp,
4094 continuous_batching_enabled: enable_continuous_batching,
4095 metal_private_decode_gate,
4096 loading_model: Mutex::new(None),
4097 last_load_error: Mutex::new(None),
4098 serving: Mutex::new(crate::stats::ServingStats::default()),
4099 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4100 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4101 started_unix: unix_now(),
4102 }
4103}
4104
4105fn load_embedding_model() -> anyhow::Result<Option<Arc<ferrox_models::EmbeddingModel>>> {
4113 let Ok(path) = std::env::var("FERROX_EMBEDDING_MODEL_PATH") else {
4114 return Ok(None);
4115 };
4116 let model = ferrox_models::EmbeddingModel::from_gguf_path(&path)
4117 .map_err(|e| anyhow::anyhow!("FERROX_EMBEDDING_MODEL_PATH={path}: {e}"))?;
4118 tracing::info!(
4119 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
4120 model.name(),
4121 model.architecture(),
4122 model.n_embd(),
4123 model.pooling_type().name(),
4124 model.n_ctx_train(),
4125 );
4126 Ok(Some(Arc::new(model)))
4127}
4128
4129fn unix_now() -> u64 {
4134 std::time::SystemTime::now()
4135 .duration_since(std::time::UNIX_EPOCH)
4136 .map(|d| d.as_secs())
4137 .unwrap_or(0)
4138}
4139
4140fn startup_model_id() -> Option<String> {
4145 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
4146 let configured = std::fs::canonicalize(&configured).ok()?;
4147 admin::discover(&admin::model_dirs())
4148 .into_iter()
4149 .find(|d| {
4150 std::fs::canonicalize(&d.path)
4151 .map(|p| p == configured)
4152 .unwrap_or(false)
4153 })
4154 .map(|d| d.id)
4155}
4156
4157fn init_cpu_pool() {
4166 match ferrox_core::threads::init_cpu_pool() {
4167 Some(n) => eprintln!(
4168 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
4169 ferrox_core::threads::perf_core_count()
4170 ),
4171 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
4172 }
4173}
4174
4175fn announce_ready(addr: SocketAddr, scheme: &str) {
4188 use std::io::Write;
4189 let ready =
4190 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
4191 let mut stdout = std::io::stdout().lock();
4192 let _ = writeln!(stdout, "{}", ready.to_line());
4193 let _ = stdout.flush();
4194}
4195
4196async fn shutdown_signal(exit_on_stdin_close: bool) {
4208 if !exit_on_stdin_close {
4209 std::future::pending::<()>().await;
4210 return;
4211 }
4212 let _ = tokio::task::spawn_blocking(|| {
4213 use std::io::Read;
4214 let mut sink = [0u8; 256];
4215 let mut stdin = std::io::stdin().lock();
4216 loop {
4217 match stdin.read(&mut sink) {
4218 Ok(0) => break,
4220 Ok(_) => continue,
4222 Err(e) => {
4223 tracing::warn!("stdin read failed ({e}); treating it as closed");
4224 break;
4225 }
4226 }
4227 }
4228 })
4229 .await;
4230 tracing::info!("stdin closed; shutting down");
4231}
4232
4233fn tokio_worker_threads() -> usize {
4238 std::env::var("FERROX_TOKIO_WORKERS")
4239 .ok()
4240 .and_then(|v| v.trim().parse::<usize>().ok())
4241 .filter(|n| *n > 0)
4242 .unwrap_or(2)
4243}
4244
4245fn install_ring_crypto_provider() {
4262 let _ = rustls::crypto::ring::default_provider().install_default();
4263}
4264
4265pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4271 if args.list_devices {
4272 print_available_devices();
4273 return Ok(());
4274 }
4275 apply_cli_overrides(&args)?;
4276
4277 let _instance = {
4281 use ferrox_core::instance::{register, InstancePolicy};
4282 let policy = if args.allow_multiple_instances {
4283 InstancePolicy::Multi
4284 } else {
4285 InstancePolicy::from_env_or(InstancePolicy::Single)
4286 };
4287 let model = std::env::var("FERROX_MODEL_PATH").ok();
4288 register(
4289 "server",
4290 model.as_deref(),
4291 ferrox_core::instance::current_backend(),
4292 policy,
4293 )
4294 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4295 };
4296
4297 let journal = journal::Journal::from_env();
4298 eprintln!(
4299 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
4300 journal.path()
4301 );
4302 journal.append(&journal::Record::session_start(
4303 env!("CARGO_PKG_VERSION"),
4304 std::process::id(),
4305 ));
4306 journal::install_panic_hook(journal.clone());
4307
4308 let mcp_config_path = args.mcp_config.clone();
4309 let exit_on_stdin_close = args.exit_on_stdin_close
4310 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
4311 .map(|v| v == "1")
4312 .unwrap_or(false);
4313
4314 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
4318 init_cpu_pool();
4319
4320 let runtime = tokio::runtime::Builder::new_multi_thread()
4321 .worker_threads(tokio_worker_threads())
4322 .enable_all()
4323 .build()?;
4324 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4325
4326 let reason = match &result {
4327 Ok(()) => "normal".to_string(),
4328 Err(e) => e.to_string(),
4329 };
4330 journal.append(&journal::Record::session_exit(reason));
4331
4332 runtime.shutdown_background();
4337
4338 result
4339}
4340
4341async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4342 let _ = tracing_subscriber::fmt::try_init();
4348
4349 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4359 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
4360 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
4361 .map(|v| v == "1")
4362 .unwrap_or(false);
4363 if let Err(msg) =
4364 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4365 {
4366 anyhow::bail!(msg);
4367 }
4368
4369 let embedding_model = load_embedding_model()?;
4376
4377 let mut loaded = model::load()?;
4378 match &loaded {
4379 model::LoadedModel::Gguf(g) => tracing::info!(
4380 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4381 g.decoder.config.name,
4382 g.is_synthetic,
4383 g.tokenizer.kind()
4384 ),
4385 model::LoadedModel::Kimi(k) => tracing::info!(
4386 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4387 k.tokenizer.vocab_size()
4388 ),
4389 model::LoadedModel::Mla(m) => tracing::info!(
4390 "loaded MLA GGUF '{}' (tokenizer={})",
4391 m.name,
4392 m.tokenizer.kind()
4393 ),
4394 model::LoadedModel::Gemma4(m) => tracing::info!(
4395 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4396 m.name,
4397 m.tokenizer.kind()
4398 ),
4399 model::LoadedModel::Glm52(g) => tracing::info!(
4400 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4401 g.name,
4402 g.tokenizer.kind()
4403 ),
4404 model::LoadedModel::Encoder(_) => {}
4407 }
4408 let metal_default_moe_budget = {
4416 #[cfg(feature = "metal")]
4417 {
4418 ferrox_core::metal_dense_enabled()
4419 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
4420 }
4421 #[cfg(not(feature = "metal"))]
4422 {
4423 false
4424 }
4425 };
4426 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
4427 let budget: u64 = budget_str
4428 .parse()
4429 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4430 match &mut loaded {
4431 model::LoadedModel::Gguf(g) => {
4432 tracing::info!(
4433 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4434 (CUDA and/or Metal matvecs when built with the matching feature)"
4435 );
4436 g.decoder.gpu_vram_budget_bytes = Some(budget);
4437 }
4438 model::LoadedModel::Kimi(_) => {
4439 tracing::warn!(
4440 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4441 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4442 );
4443 }
4444 model::LoadedModel::Mla(_) => {
4445 tracing::warn!(
4446 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4447 FFN path only today; ignoring expert VRAM budget"
4448 );
4449 }
4450 model::LoadedModel::Gemma4(_) => {
4451 tracing::warn!(
4452 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4453 ignoring expert VRAM budget"
4454 );
4455 }
4456 model::LoadedModel::Glm52(_) => {
4457 tracing::warn!(
4458 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4459 GPU expert placement not wired yet; ignoring"
4460 );
4461 }
4462 model::LoadedModel::Encoder(_) => {
4463 tracing::warn!(
4464 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4465 it has no routed experts to place; ignoring"
4466 );
4467 }
4468 }
4469 } else if metal_default_moe_budget {
4470 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4474 if let model::LoadedModel::Gguf(g) = &mut loaded {
4475 tracing::info!(
4476 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4477 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4478 );
4479 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4480 }
4481 }
4482 #[cfg(feature = "cuda")]
4483 {
4484 if ferrox_core::cuda_dense_enabled() {
4485 tracing::info!(
4486 "CUDA dense matvec enabled for WeightMatrix::apply \
4487 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4488 );
4489 } else {
4490 tracing::info!(
4491 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
4492 );
4493 }
4494 }
4495 #[cfg(feature = "metal")]
4496 {
4497 if ferrox_core::metal_dense_enabled() {
4498 tracing::info!(
4499 "Metal dense matvec enabled for WeightMatrix::apply \
4500 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4501 );
4502 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
4503 Some("1") | Some("true") | Some("on") | Some("attn") => {
4504 tracing::info!(
4505 "Metal fused attention requested (FERROX_METAL_ATTN): \
4506 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4507 );
4508 }
4509 _ => {}
4510 }
4511 tracing::info!(
4512 "Metal greedy GPU argmax: temperature<=0 folds \
4513 final_norm+lm_head+argmax into the dense stack"
4514 );
4515 } else {
4516 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
4517 }
4518 }
4519 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
4530 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
4531 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
4532 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4533 panic!(
4534 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
4535 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
4536 );
4537 }
4538 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4539 (Ok(blocks), Ok(block_size), Err(_)) => {
4540 let total_blocks: usize = blocks
4541 .parse()
4542 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
4543 let block_size: usize = block_size
4544 .parse()
4545 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4546 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4552 .ok()
4553 .map(|v| {
4554 v.parse()
4555 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4556 })
4557 .unwrap_or(0);
4558 tracing::info!(
4559 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4560 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4561 );
4562 Some(generate::KvPoolConfig {
4563 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4564 queue_wait: Duration::from_millis(queue_wait_ms),
4565 })
4566 }
4567 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4568 let block_size: usize = block_size
4569 .parse()
4570 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4571 let budget: u64 = byte_budget
4572 .parse()
4573 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
4574 let cfg = match &loaded {
4575 model::LoadedModel::Gguf(g) => &g.decoder.config,
4576 model::LoadedModel::Kimi(_)
4577 | model::LoadedModel::Mla(_)
4578 | model::LoadedModel::Gemma4(_)
4579 | model::LoadedModel::Glm52(_)
4580 | model::LoadedModel::Encoder(_) => {
4581 panic!(
4582 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
4583 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4584 );
4585 }
4586 };
4587 let bytes_per_block = block_size
4588 * cfg.n_layers
4589 * cfg.n_kv_heads
4590 * cfg.head_dim
4591 * 2
4592 * std::mem::size_of::<f32>();
4593 assert!(
4594 bytes_per_block > 0,
4595 "derived KV block byte size must be positive (check model config and block size)"
4596 );
4597 let total_blocks = (budget as usize / bytes_per_block).max(1);
4598 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4599 .ok()
4600 .map(|v| {
4601 v.parse()
4602 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4603 })
4604 .unwrap_or(0);
4605 tracing::info!(
4606 "KV cache block pool enabled from byte budget: {budget} bytes / \
4607 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4608 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4609 cfg.n_layers
4610 );
4611 Some(generate::KvPoolConfig {
4612 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4613 queue_wait: Duration::from_millis(queue_wait_ms),
4614 })
4615 }
4616 (Err(_), Err(_), Err(_)) => None,
4617 (Err(_), Ok(_), Err(_)) => panic!(
4618 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4619 (or unset all three to disable KV cache pooling)"
4620 ),
4621 (Ok(_), Ok(_), Ok(_)) => {
4622 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4623 }
4624 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4625 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4626 set together (or neither, to disable KV cache pooling)"
4627 ),
4628 };
4629 let paged_kv = match (
4635 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4636 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4637 ) {
4638 (Ok(blocks), Ok(block_size)) => {
4639 assert!(
4640 kv_pool.is_none(),
4641 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4642 mutually exclusive: both bound the same KV memory, by different means. \
4643 Set one."
4644 );
4645 let blocks_per_layer: usize = blocks
4659 .parse()
4660 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4661 let block_size: usize = block_size
4662 .parse()
4663 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4664 let gguf = match &loaded {
4665 model::LoadedModel::Gguf(g) => g,
4666 _ => panic!(
4667 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4668 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4669 ),
4670 };
4671 let cfg = &gguf.decoder.config;
4672 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4673 .ok()
4674 .map(|v| {
4675 v.parse()
4676 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4677 })
4678 .unwrap_or(0);
4679 tracing::info!(
4680 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4681 layer across {} layers, shared by all concurrent requests, \
4682 {queue_wait_ms}ms admission queue wait",
4683 cfg.n_layers
4684 );
4685 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4690 block_size,
4691 ))));
4692 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4701 crate::policy::parser::ToolCallFormat::infer(
4702 &std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
4703 )
4704 .opener(),
4705 |text| {
4706 gguf.tokenizer
4707 .encode(text)
4708 .into_iter()
4709 .map(|t| t as u32)
4710 .collect()
4711 },
4712 );
4713 if let Some(id) = anchor_token {
4714 tracing::info!(
4715 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4716 window stops short of where its next turn rejoins"
4717 );
4718 }
4719 let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
4720 .ok()
4721 .map(|v| {
4722 v.parse()
4723 .expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4724 })
4725 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4726 if let Some(window) = cfg.uniform_sliding_window() {
4727 tracing::info!(
4728 "Paged KV window slide enabled: every layer slides by {window} every \
4729 {slide_interval} decode steps, so a request holds its prompt and a \
4730 window rather than its whole context"
4731 );
4732 } else if cfg.kv_block_window().is_some() {
4733 tracing::info!(
4734 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4735 and a page group holds one block in every layer"
4736 );
4737 }
4738 Some(generate::PagedKvConfig {
4739 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4740 cfg.n_layers,
4741 block_size,
4742 blocks_per_layer,
4743 cfg.n_kv_heads,
4744 cfg.head_dim,
4745 )),
4746 queue_wait: Duration::from_millis(queue_wait_ms),
4747 radix,
4748 anchor_token,
4749 slide_interval,
4750 })
4751 }
4752 (Err(_), Err(_)) => None,
4753 _ => panic!(
4754 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4755 (or neither, to disable paged KV)"
4756 ),
4757 };
4758 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4765 let max_entries: usize = v
4766 .parse()
4767 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4768 if kv_pool.is_some() {
4769 tracing::warn!(
4770 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4771 caching will never be consulted while a KV pool is configured"
4772 );
4773 }
4774 assert!(
4781 paged_kv.is_none(),
4782 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4783 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4784 produce, so the cache could never hit. Set one."
4785 );
4786 tracing::info!(
4787 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4788 all requests"
4789 );
4790 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4791 });
4792 if matches!(
4793 loaded,
4794 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4795 ) && (kv_pool.is_some() || prefix_cache.is_some())
4796 {
4797 tracing::warn!(
4798 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4799 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4800 ferrox_models::engine's module docs"
4801 );
4802 }
4803 let enable_cb =
4804 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4805 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4806 tracing::info!(
4807 "continuous batching enabled by default on Metal for safe parallel serving \
4808 (set FERROX_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4809 );
4810 }
4811 if continuous_batching_env() == Some(true)
4812 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4813 && (kv_pool.is_some() || prefix_cache.is_some())
4814 {
4815 tracing::warn!(
4816 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4817 (those modes keep the private generate path)"
4818 );
4819 }
4820 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4821 if let Ok(chunk) = n.parse::<usize>() {
4822 if chunk > 0 {
4823 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4824 }
4825 }
4826 }
4827 if matches!(
4828 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4829 Some("1")
4830 ) {
4831 tracing::warn!(
4832 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4833 (minimal spill; full layer offload still planned)"
4834 );
4835 }
4836
4837 let mcp = match mcp_config_path {
4838 Some(path) => {
4839 let loaded = mcp::load_mcp_config(&path)?;
4840 tracing::info!(
4841 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4842 loaded.path,
4843 loaded.servers.len()
4844 );
4845 Some(loaded)
4846 }
4847 None => None,
4848 };
4849
4850 let detection = health::Detection::spawn();
4854
4855 let state = Arc::new(build_app_state(
4856 StartupModels {
4857 loaded,
4858 embedding: embedding_model,
4859 },
4860 kv_pool,
4861 paged_kv,
4862 prefix_cache,
4863 enable_cb,
4864 mcp,
4865 detection,
4866 ));
4867
4868 use ferrox_api::routes;
4872
4873 let public = Router::new().route(routes::HEALTH, get(health));
4878
4879 let mut protected = Router::new()
4880 .route(routes::V1_MODELS, get(list_models))
4881 .route(routes::V1_RESPONSES, post(responses::responses))
4885 .route(
4886 &axum_path(routes::V1_RESPONSE),
4887 get(responses::responses_get),
4888 )
4889 .route(
4890 &axum_path(routes::V1_RESPONSE_CANCEL),
4891 post(responses::responses_cancel),
4892 )
4893 .route(routes::V1_STATS, get(serving_stats))
4894 .route(routes::V1_REQUESTS, get(recent_requests))
4895 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
4896 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
4897 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
4898 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
4899 .route(routes::V1_CANCEL, post(cancel_generation))
4903 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
4908 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
4909 .route(routes::V1_MESSAGES, post(anthropic::messages))
4910 .route(
4911 routes::V1_MESSAGES_COUNT_TOKENS,
4912 post(anthropic::count_tokens),
4913 )
4914 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
4915 .route(routes::COMPLETION, post(completion::completion))
4920 .route(routes::COMPLETIONS, post(completion::completion))
4921 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
4922 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
4923 .route(routes::TOKENIZE, post(openai_extra::tokenize))
4929 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
4930 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
4931 .route(routes::V1_RERANK, post(rerank::rerank))
4935 .route(routes::RERANK, post(rerank::rerank))
4936 .route(routes::CACHE_STATS, get(cache_stats))
4937 .route(routes::METRICS, get(metrics))
4938 .route(routes::ADMIN_MODELS, get(admin::models))
4943 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
4944 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
4945 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
4946 .route(routes::ADMIN_TASKS, get(admin::tasks))
4947 .route(&admin::cancel_route(), post(admin::cancel_task))
4948 .route(routes::ADMIN_STATS, get(admin::stats))
4949 .merge(conversations::router());
4953
4954 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4959 tracing::info!("API key auth enabled");
4960 let auth = limits::AuthConfig {
4961 api_key: Arc::new(key),
4962 };
4963 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4964 auth,
4965 limits::require_api_key,
4966 ));
4967 }
4968 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4969 let rpm: u32 = rpm
4970 .parse()
4971 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4972 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4973 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4974 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4975 limiter,
4976 limits::rate_limit,
4977 ));
4978 }
4979 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
4987 let origins = security::parse_cors_origins(&spec)
4988 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
4989 tracing::info!(
4990 "CORS enabled: {} allow-listed origin(s) ({})",
4991 origins.len(),
4992 spec
4993 );
4994 let cors = tower_http::cors::CorsLayer::new()
4995 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4996 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4997 .allow_headers([
4998 axum::http::header::CONTENT_TYPE,
4999 axum::http::header::AUTHORIZATION,
5000 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
5006 axum::http::HeaderName::from_static("last-event-id"),
5010 ]);
5011 protected = protected.route_layer(cors);
5012 }
5013
5014 let app = public
5018 .merge(protected)
5019 .layer(axum::middleware::from_fn(limits::retry_after))
5020 .with_state(state);
5021
5022 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
5028 install_ring_crypto_provider();
5029 match tls_paths {
5034 Some(paths) => {
5035 let config =
5036 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
5037 .await
5038 .map_err(|e| {
5039 anyhow::anyhow!(
5040 "failed to load TLS cert/key ({:?}, {:?}): {e}",
5041 paths.cert,
5042 paths.key
5043 )
5044 })?;
5045 let socket_addr: std::net::SocketAddr = addr
5046 .parse()
5047 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
5048 let listener = std::net::TcpListener::bind(socket_addr)?;
5049 listener.set_nonblocking(true)?;
5057 let bound = listener.local_addr()?;
5058 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
5059 announce_ready(bound, "https");
5060
5061 let handle = axum_server::Handle::new();
5062 let shutdown_handle = handle.clone();
5063 tokio::spawn(async move {
5064 shutdown_signal(exit_on_stdin_close).await;
5065 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
5066 });
5067 axum_server::from_tcp_rustls(listener, config)?
5068 .handle(handle)
5069 .serve(app.into_make_service())
5070 .await?;
5071 }
5072 None => {
5073 let listener = tokio::net::TcpListener::bind(&addr).await?;
5074 let bound = listener.local_addr()?;
5075 tracing::info!("ferrox-server listening on {bound}");
5076 announce_ready(bound, "http");
5077 axum::serve(listener, app)
5078 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
5079 .await?;
5080 }
5081 }
5082 Ok(())
5083}
5084
5085#[cfg(test)]
5086mod tests {
5087 use super::*;
5088 use ferrox_models::config::test_dense_fixture;
5089
5090 #[test]
5091 fn parses_llama_server_style_options() {
5092 let argv = [
5093 "ferrox-server",
5094 "-m",
5095 "model.gguf",
5096 "--host",
5097 "::1",
5098 "--port",
5099 "9000",
5100 "-t",
5101 "4",
5102 "-dev",
5103 "Metal",
5104 "-ngl",
5105 "all",
5106 ]
5107 .into_iter()
5108 .map(String::from)
5109 .collect();
5110 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
5111
5112 assert_eq!(args.model.as_deref(), Some("model.gguf"));
5113 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
5114 assert_eq!(args.port, Some(9000));
5115 assert_eq!(args.threads, Some(4));
5116 assert_eq!(args.device, Some(OffloadDevice::Metal));
5117 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
5118 assert_eq!(
5119 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
5120 Some("[::1]:9000")
5121 );
5122 }
5123
5124 #[test]
5125 fn port_zero_survives_argument_parsing_as_a_real_request() {
5126 let argv = ["ferrox-server", "--port", "0"]
5131 .into_iter()
5132 .map(String::from)
5133 .collect();
5134 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
5135 assert_eq!(args.port, Some(0));
5136 assert_eq!(
5137 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
5138 Some("127.0.0.1:0")
5139 );
5140 }
5141
5142 #[test]
5143 fn parallel_flag_parses_and_rewrites_np() {
5144 let argv = ["ferrox-server", "-np", "4"]
5145 .into_iter()
5146 .map(String::from)
5147 .collect();
5148 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
5149 assert_eq!(args.parallel, Some(4));
5150 }
5151
5152 #[test]
5153 fn stdin_close_exit_is_opt_in() {
5154 let args =
5157 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
5158 assert!(!args.exit_on_stdin_close);
5159 let args = ServerArgs::try_parse_from(
5160 ["ferrox-server", "--exit-on-stdin-close"]
5161 .into_iter()
5162 .map(String::from),
5163 )
5164 .unwrap();
5165 assert!(args.exit_on_stdin_close);
5166 }
5167
5168 #[test]
5169 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
5170 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
5171 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
5172 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
5173 assert_eq!(parsed.port, 51999);
5174 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
5175 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
5177 }
5178
5179 fn test_model() -> Model {
5180 let cfg = test_dense_fixture();
5184 Model::Gguf(GgufModel {
5185 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
5186 tokenizer: Arc::new(ServerTokenizer::Byte),
5187 stop_tokens: StopTokens::default(),
5188 bos_id: None,
5189 is_synthetic: true,
5190 chat_template: chat_template::PromptTemplate::plain(),
5191 })
5192 }
5193
5194 fn greedy_params(max_tokens: usize) -> GenerationParams {
5195 GenerationParams {
5196 max_tokens,
5197 sampling: SamplingParams::default(),
5198 seed: 1,
5199 stop: Vec::new(),
5200 stop_token_ids: Vec::new(),
5201 json_object: false,
5202 grammar: None,
5203 cancel: None,
5204 ignore_eos: false,
5205 }
5206 }
5207
5208 fn test_model_full_byte_vocab() -> Model {
5212 test_model_full_byte_vocab_with_eos(None)
5213 }
5214
5215 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
5222 let mut cfg = test_dense_fixture();
5223 cfg.vocab_size = 256;
5224 Model::Gguf(GgufModel {
5225 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5226 tokenizer: Arc::new(ServerTokenizer::Byte),
5227 stop_tokens: StopTokens::from_eos(eos),
5228 bos_id: None,
5229 is_synthetic: true,
5230 chat_template: chat_template::PromptTemplate::plain(),
5231 })
5232 }
5233
5234 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
5238 AppState {
5239 embedding: None,
5240 paged_kv: None,
5241 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
5242 id: None,
5243 loaded: Loaded::Generative(Arc::new(model)),
5244 batcher: None,
5245 ceiling: None,
5246 }))),
5247 load_in_progress: std::sync::atomic::AtomicBool::new(false),
5248 tasks: Arc::new(tasks::TaskRegistry::new()),
5249 cancels: Arc::new(cancel::CancelRegistry::new()),
5250 stats: stats::Stats::new(),
5251 streams: resume::StreamRegistry::new(),
5252 model_dir: None,
5253 response_cache: Mutex::new(response_cache),
5254 kv_pool: None,
5255 prefix_cache: None,
5256 sessions: session::SessionStore::new(),
5257 requests_total: std::sync::atomic::AtomicU64::new(0),
5258 request_errors_total: std::sync::atomic::AtomicU64::new(0),
5259 started_at: std::time::Instant::now(),
5260 last_request_ms: std::sync::atomic::AtomicU64::new(0),
5261 detection: Arc::new(health::Detection::ready(health::probe_backends())),
5262 mcp: None,
5263 continuous_batching_enabled: false,
5264 metal_private_decode_gate: None,
5265 loading_model: Mutex::new(None),
5266 last_load_error: Mutex::new(None),
5267 serving: Mutex::new(crate::stats::ServingStats::default()),
5268 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
5269 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
5270 started_unix: unix_now(),
5271 }
5272 }
5273
5274 fn test_app() -> Router {
5283 test_app_with_state(Arc::new(test_state(
5284 test_model_full_byte_vocab(),
5285 ResponseCache::new(1000, Duration::from_secs(3600)),
5286 )))
5287 }
5288
5289 fn test_app_with_state(state: Arc<AppState>) -> Router {
5292 Router::new()
5293 .route(ferrox_api::routes::HEALTH, get(health))
5294 .route(ferrox_api::routes::V1_MODELS, get(list_models))
5295 .route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
5296 .route(
5297 &axum_path(ferrox_api::routes::V1_RESPONSE),
5298 get(responses::responses_get),
5299 )
5300 .route(
5301 &axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
5302 post(responses::responses_cancel),
5303 )
5304 .route(ferrox_api::routes::V1_STATS, get(serving_stats))
5305 .route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
5306 .route(
5307 ferrox_api::routes::V1_CACHE_STATUS,
5308 get(cache_admin::cache_status),
5309 )
5310 .route(
5311 ferrox_api::routes::V1_CACHE_REBUILD,
5312 post(cache_admin::cache_rebuild),
5313 )
5314 .route(
5315 ferrox_api::routes::ADMIN_PREPARE_STOP,
5316 post(cache_admin::prepare_stop),
5317 )
5318 .route("/v1/chat/completions", post(chat_completions))
5319 .route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
5320 .route(
5321 ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
5322 post(anthropic::count_tokens),
5323 )
5324 .route("/v1/tokenize", post(openai_extra::tokenize))
5325 .route("/v1/detokenize", post(openai_extra::detokenize))
5326 .route(ferrox_api::routes::TOKENIZE, post(openai_extra::tokenize))
5330 .route(
5331 ferrox_api::routes::DETOKENIZE,
5332 post(openai_extra::detokenize),
5333 )
5334 .route("/v1/embeddings", post(embeddings::embeddings))
5335 .route("/v1/completions", post(openai_extra::completions))
5336 .route(ferrox_api::routes::COMPLETION, post(completion::completion))
5338 .route(
5339 ferrox_api::routes::COMPLETIONS,
5340 post(completion::completion),
5341 )
5342 .route(
5343 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5344 post(admin::unload_model),
5345 )
5346 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
5347 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
5348 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
5349 .route(
5350 &axum_path(ferrox_api::routes::V1_STREAM),
5351 get(resume::resume),
5352 )
5353 .route(
5354 &axum_path(ferrox_api::routes::V1_STREAM_POLL),
5355 get(resume::poll),
5356 )
5357 .with_state(state)
5358 }
5359
5360 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
5361 let mut cfg = test_dense_fixture();
5362 cfg.name = name;
5363 cfg.vocab_size = vocab_size;
5364 Model::Gguf(GgufModel {
5365 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5366 tokenizer: Arc::new(ServerTokenizer::Byte),
5367 stop_tokens: StopTokens::default(),
5368 bos_id: None,
5369 is_synthetic: true,
5370 chat_template: chat_template::PromptTemplate::plain(),
5371 })
5372 }
5373
5374 fn model_with_template(name: &'static str, source: &str) -> Model {
5378 let mut cfg = test_dense_fixture();
5379 cfg.name = name;
5380 cfg.vocab_size = 256;
5381 Model::Gguf(GgufModel {
5382 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5383 tokenizer: Arc::new(ServerTokenizer::Byte),
5384 stop_tokens: StopTokens::default(),
5385 bos_id: None,
5386 is_synthetic: true,
5387 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5388 Some(source),
5389 Some("qwen3"),
5390 false,
5391 true,
5392 None,
5393 None,
5394 ),
5395 })
5396 }
5397
5398 #[tokio::test]
5407 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5408 let strict = "{% if messages | length > 1 %}\
5411 {{ raise_exception('this template takes one turn') }}\
5412 {% endif %}{{ messages[0].content }}";
5413 let state = Arc::new(test_state(
5414 model_with_template("strict", strict),
5415 ResponseCache::new(4, Duration::from_secs(60)),
5416 ));
5417 let app = test_app_with_state(state);
5418
5419 let (status, body) = post_json_uri(
5420 &app,
5421 "/v1/chat/completions",
5422 serde_json::json!({
5423 "model": "strict",
5424 "stream": true,
5425 "messages": [
5426 {"role": "user", "content": "one"},
5427 {"role": "user", "content": "two"},
5428 ],
5429 }),
5430 )
5431 .await;
5432 assert_eq!(status, StatusCode::BAD_REQUEST);
5433 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5434 assert!(
5435 body["error"]["message"]
5436 .as_str()
5437 .unwrap()
5438 .contains("one turn"),
5439 "the template's own message must reach the caller: {body}"
5440 );
5441
5442 let (status, _) = post_json_uri(
5444 &app,
5445 "/v1/chat/completions",
5446 serde_json::json!({
5447 "model": "strict",
5448 "stream": true,
5449 "max_tokens": 1,
5450 "messages": [{"role": "user", "content": "one"}],
5451 }),
5452 )
5453 .await;
5454 assert_eq!(status, StatusCode::OK);
5455 }
5456
5457 #[tokio::test]
5459 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5460 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5461 {% if reasoning_effort %}\
5462 {% if reasoning_effort not in ['low','medium','high'] %}\
5463 {{ raise_exception('bad effort') }}\
5464 {% endif %}[{{ reasoning_effort }}]\
5465 {% endif %}{{ messages[0].content }}";
5466 let state = Arc::new(test_state(
5467 model_with_template("thinker", reasoning),
5468 ResponseCache::new(4, Duration::from_secs(60)),
5469 ));
5470 let app = test_app_with_state(state);
5471 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5472 assert_eq!(status, StatusCode::OK);
5473 let entry = &models["data"][0];
5474 assert_eq!(
5475 entry["supported_reasoning_efforts"],
5476 serde_json::json!(["off", "low", "medium", "high"])
5477 );
5478 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5479 }
5480
5481 #[tokio::test]
5486 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5487 let app = test_app();
5488 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5489 let entry = &models["data"][0];
5490 assert!(entry.get("supported_reasoning_efforts").is_none());
5491 assert!(entry.get("default_reasoning_effort").is_none());
5492 }
5493
5494 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5495 Arc::new(ActiveModel {
5496 id: Some(name.to_string()),
5497 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5498 batcher: None,
5499 ceiling: None,
5500 })
5501 .tap_into(state)
5502 }
5503
5504 trait TapInto {
5506 fn tap_into(self, state: &AppState) -> Self;
5507 }
5508 impl TapInto for Arc<ActiveModel> {
5509 fn tap_into(self, state: &AppState) -> Self {
5510 state.swap_active(Some(Arc::clone(&self)));
5511 self
5512 }
5513 }
5514
5515 #[test]
5521 fn an_in_flight_request_keeps_the_model_it_started_on() {
5522 let state = test_state(
5523 named_test_model("model-a", 256),
5524 ResponseCache::new(4, Duration::from_secs(60)),
5525 );
5526
5527 let in_flight = state.active().expect("a model is loaded");
5530 assert_eq!(in_flight.name(), "model-a");
5531
5532 active_model(&state, "model-b");
5533
5534 assert_eq!(state.active().unwrap().name(), "model-b");
5536 assert_eq!(in_flight.name(), "model-a");
5538 let (_chunks, finish, _usage) = run_generation(
5539 in_flight.generative().unwrap(),
5540 "hi",
5541 &greedy_params(3),
5542 None,
5543 None,
5544 None,
5545 None,
5546 None,
5547 None,
5548 )
5549 .expect("the old model must still decode after being swapped out");
5550 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5551 }
5552
5553 #[test]
5558 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5559 let state = test_state(
5560 named_test_model("model-a", 256),
5561 ResponseCache::new(4, Duration::from_secs(60)),
5562 );
5563 let in_flight = state.active().expect("a model is loaded");
5564 let weights = Arc::clone(in_flight.generative().unwrap());
5565 assert!(Arc::strong_count(&weights) >= 2);
5566
5567 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5568 id: Some("model-b".to_string()),
5569 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5570 batcher: None,
5571 ceiling: None,
5572 })));
5573 drop(previous);
5574 assert!(Arc::strong_count(&weights) >= 2);
5576 drop(in_flight);
5577 assert_eq!(Arc::strong_count(&weights), 1);
5578 }
5579
5580 #[tokio::test]
5584 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5585 let state = Arc::new(test_state(
5586 named_test_model("model-a", 256),
5587 ResponseCache::new(4, Duration::from_secs(60)),
5588 ));
5589 let app = test_app_with_state(Arc::clone(&state));
5590
5591 let (status, body) = post_json_uri(
5592 &app,
5593 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5594 serde_json::json!({}),
5595 )
5596 .await;
5597 assert_eq!(status, StatusCode::OK);
5598 assert_eq!(body["ok"], true);
5599 assert!(body["active"].is_null());
5600 assert!(state.active().is_none());
5601
5602 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5603 assert_eq!(status, StatusCode::OK);
5604 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5605 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5606
5607 let (status, body) = post_json_uri(
5608 &app,
5609 "/v1/chat/completions",
5610 serde_json::json!({
5611 "model": "x",
5612 "messages": [{"role": "user", "content": "hi"}]
5613 }),
5614 )
5615 .await;
5616 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5617 assert_eq!(body["error"]["type"], "model_not_loaded");
5618 }
5619
5620 #[tokio::test]
5624 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5625 let state = Arc::new(test_state(
5626 named_test_model("model-a", 256),
5627 ResponseCache::new(4, Duration::from_secs(60)),
5628 ));
5629 let app = test_app_with_state(Arc::clone(&state));
5630 state.swap_active(None);
5631
5632 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5633 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5636 assert_eq!(body["state"], "unavailable");
5637 assert_eq!(body["reason"], "model_not_loaded");
5638 assert!(body["model"].is_null());
5639 let real_weights = body["capabilities"]
5640 .as_array()
5641 .unwrap()
5642 .iter()
5643 .find(|c| c["id"] == "real_weights")
5644 .cloned()
5645 .expect("real_weights is always reported");
5646 assert_eq!(real_weights["available"], false);
5647 assert_eq!(real_weights["reason"], "model_not_loaded");
5648 }
5649
5650 #[tokio::test]
5654 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5655 let app = test_app();
5656
5657 let (status, completion) = post_json_uri(
5658 &app,
5659 "/v1/chat/completions",
5660 serde_json::json!({
5661 "model": "x",
5662 "messages": [{"role": "user", "content": "hi"}],
5663 "max_tokens": 4
5664 }),
5665 )
5666 .await;
5667 assert_eq!(status, StatusCode::OK);
5668 let request_id = completion["request_id"].as_str().unwrap().to_string();
5669
5670 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5671 assert_eq!(status, StatusCode::OK);
5672 let recent = stats["recent"].as_array().unwrap();
5673 assert_eq!(recent.len(), 1);
5674 let row = &recent[0];
5675 assert_eq!(row["request_id"], request_id);
5676 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
5677 assert_eq!(row["status"], 200);
5678 assert_eq!(row["stream"], false);
5679 assert!(row["duration_ms"].is_number());
5682 assert!(row["decode_ms"].is_number());
5683 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5684 assert_eq!(
5685 stats["tokens_prompt_total"].as_u64().unwrap(),
5686 row["prompt_tokens"].as_u64().unwrap()
5687 );
5688 }
5689
5690 #[tokio::test]
5694 async fn a_rejected_request_is_recorded_too() {
5695 let state = Arc::new(test_state(
5696 named_test_model("model-a", 256),
5697 ResponseCache::new(4, Duration::from_secs(60)),
5698 ));
5699 let app = test_app_with_state(Arc::clone(&state));
5700 state.swap_active(None);
5701
5702 let (status, _) = post_json_uri(
5703 &app,
5704 "/v1/chat/completions",
5705 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5706 )
5707 .await;
5708 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5709
5710 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5711 let recent = stats["recent"].as_array().unwrap();
5712 assert_eq!(recent.len(), 1);
5713 assert_eq!(recent[0]["status"], 503);
5714 assert_eq!(recent[0]["completion_tokens"], 0);
5715 assert!(recent[0]["decode_ms"].is_null());
5716 assert_eq!(stats["errors_total"], 1);
5717 }
5718
5719 async fn post_json_with_headers(
5723 app: &Router,
5724 uri: &str,
5725 body: serde_json::Value,
5726 headers: &[(&str, &str)],
5727 ) -> (StatusCode, serde_json::Value) {
5728 use http_body_util::BodyExt;
5729 use tower::ServiceExt;
5730
5731 let mut builder = axum::http::Request::builder()
5732 .method("POST")
5733 .uri(uri)
5734 .header("content-type", "application/json");
5735 for (name, value) in headers {
5736 builder = builder.header(*name, *value);
5737 }
5738 let response = app
5739 .clone()
5740 .oneshot(
5741 builder
5742 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5743 .unwrap(),
5744 )
5745 .await
5746 .unwrap();
5747 let status = response.status();
5748 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5749 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5750 (status, json)
5751 }
5752
5753 #[tokio::test]
5757 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5758 let app = test_app();
5759
5760 let (status, _) = post_json_uri(
5761 &app,
5762 ferrox_api::routes::V1_TOKENIZE,
5763 serde_json::json!({"prompt": "hello"}),
5764 )
5765 .await;
5766 assert_eq!(status, StatusCode::OK);
5767 let (status, _) = post_json_uri(
5768 &app,
5769 ferrox_api::routes::V1_DETOKENIZE,
5770 serde_json::json!({"tokens": [104, 105]}),
5771 )
5772 .await;
5773 assert_eq!(status, StatusCode::OK);
5774 let (status, _) = post_json_uri(
5775 &app,
5776 ferrox_api::routes::V1_EMBEDDINGS,
5777 serde_json::json!({"input": "hello"}),
5778 )
5779 .await;
5780 assert_eq!(status, StatusCode::OK);
5781
5782 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5783 let routes: Vec<&str> = stats["recent"]
5784 .as_array()
5785 .unwrap()
5786 .iter()
5787 .map(|row| row["route"].as_str().unwrap())
5788 .collect();
5789 for expected in [
5790 ferrox_api::routes::V1_TOKENIZE,
5791 ferrox_api::routes::V1_DETOKENIZE,
5792 ferrox_api::routes::V1_EMBEDDINGS,
5793 ] {
5794 assert!(
5795 routes.contains(&expected),
5796 "{expected} is missing: {routes:?}"
5797 );
5798 }
5799
5800 let row = |route: &str| {
5801 stats["recent"]
5802 .as_array()
5803 .unwrap()
5804 .iter()
5805 .find(|r| r["route"] == route)
5806 .cloned()
5807 .unwrap()
5808 };
5809 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5813 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5814 assert!(embed["decode_ms"].is_null());
5815 assert_eq!(embed["completion_tokens"], 0);
5816 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5820 assert_eq!(
5821 stats["tokens_prompt_total"].as_u64().unwrap(),
5822 embed["prompt_tokens"].as_u64().unwrap(),
5823 "only the forward pass counted"
5824 );
5825 }
5826
5827 fn streaming_test_app() -> Router {
5832 let mut cfg = test_dense_fixture();
5833 cfg.vocab_size = 256;
5834 let model = Model::Gguf(GgufModel {
5835 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5836 tokenizer: Arc::new(ServerTokenizer::Byte),
5837 stop_tokens: StopTokens::default(),
5838 bos_id: None,
5839 is_synthetic: false,
5840 chat_template: chat_template::PromptTemplate::plain(),
5841 });
5842 test_app_with_state(Arc::new(test_state(
5843 model,
5844 ResponseCache::new(1000, Duration::from_secs(3600)),
5845 )))
5846 }
5847
5848 #[tokio::test]
5852 async fn the_native_completion_wire_is_not_the_openai_one() {
5853 let app = test_app();
5854
5855 let (status, native) = post_json_uri(
5856 &app,
5857 ferrox_api::routes::COMPLETION,
5858 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5859 )
5860 .await;
5861 assert_eq!(status, StatusCode::OK, "{native}");
5862 assert!(native["content"].is_string(), "{native}");
5863 assert_eq!(native["stop"], true);
5864 assert_eq!(native["stop_type"], "limit");
5865 assert_eq!(native["stopping_word"], "");
5866 assert_eq!(native["truncated"], false);
5867 assert_eq!(native["id_slot"], -1);
5868 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5869 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5870 assert!(
5871 native.get("choices").is_none(),
5872 "the native shape has no `choices`: {native}"
5873 );
5874
5875 let (status, openai) = post_json_uri(
5876 &app,
5877 ferrox_api::routes::V1_COMPLETIONS,
5878 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5879 )
5880 .await;
5881 assert_eq!(status, StatusCode::OK);
5882 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5883 assert!(
5884 openai.get("content").is_none(),
5885 "the OpenAI shape has no top-level `content`: {openai}"
5886 );
5887 }
5888
5889 #[tokio::test]
5893 async fn both_native_spellings_reach_the_same_handler() {
5894 let app = test_app();
5895 for route in [
5896 ferrox_api::routes::COMPLETION,
5897 ferrox_api::routes::COMPLETIONS,
5898 ] {
5899 let (status, body) = post_json_uri(
5900 &app,
5901 route,
5902 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5903 )
5904 .await;
5905 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5906 assert_eq!(body["stop"], true, "{route}");
5907 assert!(body["content"].is_string(), "{route}");
5908 }
5909
5910 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5913 let routes: Vec<&str> = stats["recent"]
5914 .as_array()
5915 .unwrap()
5916 .iter()
5917 .map(|row| row["route"].as_str().unwrap())
5918 .collect();
5919 assert!(
5920 routes.contains(&ferrox_api::routes::COMPLETION),
5921 "{routes:?}"
5922 );
5923 assert!(
5924 routes.contains(&ferrox_api::routes::COMPLETIONS),
5925 "{routes:?}"
5926 );
5927 }
5928
5929 #[tokio::test]
5935 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5936 let app = streaming_test_app();
5937 let raw = post_sse_raw_uri(
5938 &app,
5939 ferrox_api::routes::COMPLETION,
5940 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5941 )
5942 .await;
5943
5944 assert!(
5945 !raw.contains("[DONE]"),
5946 "llama.cpp's native stream has no sentinel: {raw}"
5947 );
5948 let frames: Vec<serde_json::Value> = raw
5949 .lines()
5950 .filter_map(|line| line.strip_prefix("data: "))
5951 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5952 .collect();
5953 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5954
5955 let (last, partials) = frames.split_last().unwrap();
5956 assert_eq!(last["stop"], true, "the last frame closes the stream");
5957 assert!(last["timings"].is_object(), "{last}");
5958 assert!(last["stop_type"].is_string(), "{last}");
5959 for partial in partials {
5960 assert_eq!(partial["stop"], false, "{partial}");
5961 assert!(partial["content"].is_string(), "{partial}");
5962 assert!(partial.get("timings").is_none(), "{partial}");
5966 assert!(partial.get("generation_settings").is_none(), "{partial}");
5967 }
5968 let streamed: String = partials
5971 .iter()
5972 .filter_map(|p| p["content"].as_str())
5973 .collect();
5974 assert_eq!(last["content"].as_str().unwrap(), streamed);
5975 }
5976
5977 #[tokio::test]
5982 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5983 let app = test_app();
5984 for body in [
5985 serde_json::json!({"prompt": "hi"}),
5986 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5987 ] {
5988 let (status, refusal) =
5989 post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
5990 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5991 assert!(
5992 refusal["error"]["message"]
5993 .as_str()
5994 .unwrap()
5995 .contains("n_predict"),
5996 "{refusal}"
5997 );
5998 }
5999 let (status, _) = post_json_uri(
6002 &app,
6003 ferrox_api::routes::COMPLETION,
6004 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6005 )
6006 .await;
6007 assert_eq!(status, StatusCode::OK);
6008 }
6009
6010 #[tokio::test]
6019 async fn a_stop_string_halts_the_answer_and_is_named_back() {
6020 let app = streaming_test_app();
6021 let ask = |stop: serde_json::Value| {
6022 let app = app.clone();
6023 async move {
6024 post_json_uri(
6025 &app,
6026 ferrox_api::routes::COMPLETION,
6027 serde_json::json!({
6028 "prompt": "hi",
6029 "n_predict": 64,
6030 "ignore_eos": true,
6031 "stop": stop,
6032 }),
6033 )
6034 .await
6035 .1
6036 }
6037 };
6038
6039 let baseline = ask(serde_json::json!([])).await;
6040 assert_eq!(baseline["stop_type"], "limit");
6041 assert_eq!(baseline["stopping_word"], "");
6042 let text = baseline["content"].as_str().unwrap().to_string();
6043 let sentinel: String = text.chars().skip(1).take(2).collect();
6049 assert_eq!(
6050 sentinel.chars().count(),
6051 2,
6052 "the fixture must produce enough output to cut: {text:?}"
6053 );
6054 let cut = text.find(&sentinel).expect("it came out of this text");
6055
6056 let stopped = ask(serde_json::json!([sentinel])).await;
6057 assert_eq!(stopped["stop_type"], "word", "{stopped}");
6058 assert_eq!(stopped["stopping_word"], sentinel);
6059 assert_eq!(
6060 stopped["content"].as_str().unwrap(),
6061 &text[..cut],
6062 "the answer must be cut at the sentinel, not run past it"
6063 );
6064 }
6065
6066 #[tokio::test]
6072 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
6073 let app = test_app();
6074
6075 let (v1_status, v1) = post_json_uri(
6076 &app,
6077 ferrox_api::routes::V1_TOKENIZE,
6078 serde_json::json!({"prompt": "hello"}),
6079 )
6080 .await;
6081 let (alias_status, alias) = post_json_uri(
6082 &app,
6083 ferrox_api::routes::TOKENIZE,
6084 serde_json::json!({"content": "hello"}),
6085 )
6086 .await;
6087 assert_eq!(v1_status, StatusCode::OK);
6088 assert_eq!(alias_status, StatusCode::OK, "{alias}");
6089 assert_eq!(v1["tokens"], alias["tokens"]);
6090 assert!(!alias["tokens"].as_array().unwrap().is_empty());
6091
6092 let (status, both_ways) = post_json_uri(
6095 &app,
6096 ferrox_api::routes::TOKENIZE,
6097 serde_json::json!({"prompt": "hello"}),
6098 )
6099 .await;
6100 assert_eq!(status, StatusCode::OK);
6101 assert_eq!(both_ways["tokens"], v1["tokens"]);
6102 }
6103
6104 #[tokio::test]
6109 async fn detokenize_answers_under_both_dialects_keys() {
6110 let app = test_app();
6111 for route in [
6112 ferrox_api::routes::DETOKENIZE,
6113 ferrox_api::routes::V1_DETOKENIZE,
6114 ] {
6115 let (status, body) =
6116 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
6117 assert_eq!(status, StatusCode::OK, "{route}");
6118 assert_eq!(body["text"], "hi", "{route}");
6119 assert_eq!(body["content"], body["text"], "{route}");
6120 }
6121 }
6122
6123 #[tokio::test]
6127 async fn the_alias_is_recorded_under_the_path_the_client_called() {
6128 let app = test_app();
6129 let (status, _) = post_json_uri(
6130 &app,
6131 ferrox_api::routes::TOKENIZE,
6132 serde_json::json!({"content": "hello"}),
6133 )
6134 .await;
6135 assert_eq!(status, StatusCode::OK);
6136
6137 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6138 let routes: Vec<&str> = stats["recent"]
6139 .as_array()
6140 .unwrap()
6141 .iter()
6142 .map(|row| row["route"].as_str().unwrap())
6143 .collect();
6144 assert!(
6145 routes.contains(&ferrox_api::routes::TOKENIZE),
6146 "the alias must be its own row: {routes:?}"
6147 );
6148 assert!(
6149 !routes.contains(&ferrox_api::routes::V1_TOKENIZE),
6150 "nothing called /v1/tokenize: {routes:?}"
6151 );
6152 }
6153
6154 #[tokio::test]
6159 async fn add_special_prepends_the_same_bos_the_decoder_would() {
6160 let mut cfg = test_dense_fixture();
6161 cfg.vocab_size = 256;
6162 let model = Model::Gguf(GgufModel {
6163 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
6164 tokenizer: Arc::new(ServerTokenizer::Byte),
6165 stop_tokens: StopTokens::default(),
6166 bos_id: Some(7),
6167 is_synthetic: true,
6168 chat_template: chat_template::PromptTemplate::plain(),
6169 });
6170 let app = test_app_with_state(Arc::new(test_state(
6171 model,
6172 ResponseCache::new(1000, Duration::from_secs(3600)),
6173 )));
6174
6175 let (_, plain) = post_json_uri(
6176 &app,
6177 ferrox_api::routes::TOKENIZE,
6178 serde_json::json!({"content": "hi"}),
6179 )
6180 .await;
6181 let (_, special) = post_json_uri(
6182 &app,
6183 ferrox_api::routes::TOKENIZE,
6184 serde_json::json!({"content": "hi", "add_special": true}),
6185 )
6186 .await;
6187
6188 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6189 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6190 assert_eq!(special["count"], 3);
6191 }
6192
6193 #[tokio::test]
6196 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6197 let app = test_app();
6198 let (status, _) = post_json_uri(
6199 &app,
6200 ferrox_api::routes::V1_EMBEDDINGS,
6201 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6202 )
6203 .await;
6204 assert_eq!(status, StatusCode::BAD_REQUEST);
6205
6206 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6207 let recent = stats["recent"].as_array().unwrap();
6208 assert_eq!(recent.len(), 1);
6209 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
6210 assert_eq!(recent[0]["status"], 400);
6211 assert_eq!(
6212 recent[0]["prompt_tokens"], 0,
6213 "a rejected call embedded nothing"
6214 );
6215 }
6216
6217 #[tokio::test]
6220 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6221 let app = test_app();
6222 let key = "sk-monitor-secret";
6223 let (status, _) = post_json_with_headers(
6224 &app,
6225 "/v1/chat/completions",
6226 serde_json::json!({
6227 "model": "x",
6228 "messages": [{"role": "user", "content": "hi"}],
6229 "max_tokens": 2
6230 }),
6231 &[
6232 ("authorization", &format!("Bearer {key}")),
6233 ("x-ferrox-client", "ferrox-studio"),
6234 ],
6235 )
6236 .await;
6237 assert_eq!(status, StatusCode::OK);
6238
6239 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6240 let row = stats["recent"].as_array().unwrap()[0].clone();
6241 let fingerprint = row["via_api_key"]
6242 .as_str()
6243 .expect("the row names the key that served it")
6244 .to_string();
6245 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6246 assert!(!fingerprint.contains(key));
6247 assert!(
6248 !serde_json::to_string(&stats).unwrap().contains(key),
6249 "the stats payload must not carry the key in any form"
6250 );
6251 assert_eq!(row["client"], "ferrox-studio");
6252 }
6253
6254 #[tokio::test]
6257 async fn different_keys_are_different_callers_and_no_key_is_null() {
6258 let app = test_app();
6259 let body = serde_json::json!({
6260 "model": "x",
6261 "messages": [{"role": "user", "content": "hi"}],
6262 "max_tokens": 1
6263 });
6264 for headers in [
6265 vec![("authorization", "Bearer key-one")],
6266 vec![("authorization", "Bearer key-two")],
6267 vec![],
6268 ] {
6269 let (status, _) =
6270 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6271 assert_eq!(status, StatusCode::OK);
6272 }
6273
6274 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6275 let recent = stats["recent"].as_array().unwrap();
6276 assert_eq!(recent.len(), 3);
6277 let one = recent[0]["via_api_key"].as_str().unwrap();
6278 let two = recent[1]["via_api_key"].as_str().unwrap();
6279 assert_ne!(one, two, "two keys must not collapse into one caller");
6280 assert!(
6281 recent[2]["via_api_key"].is_null(),
6282 "an unauthenticated call is null, not a fingerprint of nothing"
6283 );
6284 assert!(recent[2]["client"].is_null());
6285 }
6286
6287 #[tokio::test]
6292 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6293 let state = Arc::new(test_state(
6294 named_test_model("really-loaded", 256),
6295 ResponseCache::new(4, Duration::from_secs(60)),
6296 ));
6297 let app = test_app_with_state(Arc::clone(&state));
6298
6299 let (status, _) = post_json_uri(
6300 &app,
6301 "/v1/chat/completions",
6302 serde_json::json!({
6303 "model": "gpt-4-turbo-that-is-not-here",
6304 "messages": [{"role": "user", "content": "hi"}],
6305 "max_tokens": 2
6306 }),
6307 )
6308 .await;
6309 assert_eq!(status, StatusCode::OK);
6310
6311 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6312 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6313
6314 state.swap_active(None);
6317 let (status, _) = post_json_uri(
6318 &app,
6319 "/v1/chat/completions",
6320 serde_json::json!({
6321 "model": "gpt-4-turbo-that-is-not-here",
6322 "messages": [{"role": "user", "content": "hi"}]
6323 }),
6324 )
6325 .await;
6326 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6327 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6328 let recent = stats["recent"].as_array().unwrap();
6329 assert!(recent[recent.len() - 1]["model"].is_null());
6330 }
6331
6332 #[tokio::test]
6336 async fn a_streamed_row_names_the_model_it_decoded_against() {
6337 let state = Arc::new(test_state(
6338 named_test_model("model-before", 256),
6339 ResponseCache::new(4, Duration::from_secs(60)),
6340 ));
6341 let app = test_app_with_state(Arc::clone(&state));
6342 let _ = post_sse_raw(&app, resumable_request()).await;
6343 active_model(&state, "model-after");
6345
6346 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6347 assert_eq!(stats["recent"][0]["model"], "model-before");
6348 }
6349
6350 #[tokio::test]
6353 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6354 let app = test_app();
6355 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6356 assert_eq!(status, StatusCode::OK);
6357 assert!(
6358 stats["queue_depth"].is_null(),
6359 "without continuous batching nothing queues, so there is nothing to measure"
6360 );
6361 assert!(stats["queue_rejected_total"].is_null());
6362 assert_eq!(
6363 stats["generating_now"], 0,
6364 "work in progress is measured and really is zero here"
6365 );
6366 }
6367
6368 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6373 post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
6374 }
6375
6376 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6380 use http_body_util::BodyExt;
6381 use tower::ServiceExt;
6382
6383 let response = app
6384 .clone()
6385 .oneshot(
6386 axum::http::Request::builder()
6387 .method("POST")
6388 .uri(uri)
6389 .header("content-type", "application/json")
6390 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6391 .unwrap(),
6392 )
6393 .await
6394 .unwrap();
6395 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6396 String::from_utf8(bytes.to_vec()).unwrap()
6397 }
6398
6399 async fn get_json_with_headers(
6400 app: &Router,
6401 uri: &str,
6402 headers: &[(&str, &str)],
6403 ) -> (StatusCode, serde_json::Value) {
6404 use http_body_util::BodyExt;
6405 use tower::ServiceExt;
6406
6407 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6408 for (name, value) in headers {
6409 builder = builder.header(*name, *value);
6410 }
6411 let response = app
6412 .clone()
6413 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6414 .await
6415 .unwrap();
6416 let status = response.status();
6417 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6418 (
6419 status,
6420 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6421 )
6422 }
6423
6424 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6425 body.lines()
6426 .filter_map(|line| line.strip_prefix(field))
6427 .map(str::trim)
6428 .collect()
6429 }
6430
6431 fn resumable_request() -> serde_json::Value {
6432 serde_json::json!({
6433 "model": "m",
6434 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6435 "max_tokens": 4,
6436 "temperature": 0,
6437 "stream": true,
6438 "stream_resumable": true,
6439 })
6440 }
6441
6442 #[tokio::test]
6447 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6448 let app = test_app();
6449 let body = post_sse_raw(&app, resumable_request()).await;
6450
6451 let request_id = body
6452 .lines()
6453 .find_map(|l| l.strip_prefix("data: "))
6454 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6455 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6456 .expect("the first chunk names the request");
6457
6458 let ids = sse_field(&body, "id:");
6459 let datas = sse_field(&body, "data:");
6460 assert_eq!(
6461 ids.len(),
6462 datas.len(),
6463 "every event carries an id, or a reconnect cannot name where it stopped"
6464 );
6465 for (i, id) in ids.iter().enumerate() {
6466 assert_eq!(*id, format!("{request_id}:{i}"));
6467 }
6468 let retries = sse_field(&body, "retry:");
6469 assert_eq!(
6470 retries.len(),
6471 1,
6472 "the reconnect delay is stated once, not on every event"
6473 );
6474 assert_eq!(retries[0], "1500");
6475 assert!(
6476 body.contains("data: [DONE]"),
6477 "the end of stream is still stated"
6478 );
6479 }
6480
6481 #[tokio::test]
6485 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6486 let app = test_app();
6487 let mut request = resumable_request();
6488 request["stream_resumable"] = serde_json::json!(false);
6489 let body = post_sse_raw(&app, request).await;
6490 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6491 assert!(
6492 sse_field(&body, "id:").is_empty(),
6493 "an id promises a replay this stream cannot serve"
6494 );
6495 assert!(sse_field(&body, "retry:").is_empty());
6496 }
6497
6498 #[tokio::test]
6502 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6503 let app = test_app();
6504 let body = post_sse_raw(&app, resumable_request()).await;
6505 let request_id = sse_field(&body, "id:")[0]
6506 .rsplit_once(':')
6507 .unwrap()
6508 .0
6509 .to_string();
6510 let streamed: Vec<String> = sse_field(&body, "data:")
6511 .iter()
6512 .map(|d| d.to_string())
6513 .collect();
6514
6515 let (status, polled) = get_json(
6516 &app,
6517 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
6518 )
6519 .await;
6520 assert_eq!(status, StatusCode::OK);
6521 let events: Vec<String> = polled["events"]
6522 .as_array()
6523 .unwrap()
6524 .iter()
6525 .map(|e| e["data"].as_str().unwrap().to_string())
6526 .collect();
6527 assert_eq!(
6528 events, streamed,
6529 "the fallback must deliver the same answer, not a re-run of it"
6530 );
6531 assert_eq!(polled["request_id"], request_id);
6532 assert_eq!(
6533 polled["done"], false,
6534 "events were still being handed out, so the client must ask again"
6535 );
6536
6537 let next = polled["next_index"].as_u64().unwrap();
6540 let (_, drained) = get_json(
6541 &app,
6542 &format!(
6543 "{}?from={next}",
6544 ferrox_api::routes::v1_stream_poll(&request_id)
6545 ),
6546 )
6547 .await;
6548 assert_eq!(drained["done"], true);
6549 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6550 }
6551
6552 #[tokio::test]
6556 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6557 let app = test_app();
6558 let body = post_sse_raw(&app, resumable_request()).await;
6559 let ids = sse_field(&body, "id:");
6560 let datas: Vec<String> = sse_field(&body, "data:")
6561 .iter()
6562 .map(|d| d.to_string())
6563 .collect();
6564 assert!(
6565 ids.len() >= 3,
6566 "need a few events to resume into the middle"
6567 );
6568 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6569
6570 let (status, resumed) = get_json_with_headers(
6571 &app,
6572 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
6573 &[],
6574 )
6575 .await;
6576 assert_eq!(status, StatusCode::OK);
6577 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6578
6579 let (_, tail) = get_json(
6581 &app,
6582 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
6583 )
6584 .await;
6585 let tail_events: Vec<String> = tail["events"]
6586 .as_array()
6587 .unwrap()
6588 .iter()
6589 .map(|e| e["data"].as_str().unwrap().to_string())
6590 .collect();
6591 assert_eq!(tail_events, datas[2..].to_vec());
6592 }
6593
6594 #[tokio::test]
6597 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6598 use http_body_util::BodyExt;
6599 use tower::ServiceExt;
6600
6601 let app = test_app();
6602 let body = post_sse_raw(&app, resumable_request()).await;
6603 let ids = sse_field(&body, "id:");
6604 let datas: Vec<String> = sse_field(&body, "data:")
6605 .iter()
6606 .map(|d| d.to_string())
6607 .collect();
6608 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6609
6610 let response = app
6611 .clone()
6612 .oneshot(
6613 axum::http::Request::builder()
6614 .method("GET")
6615 .uri(ferrox_api::routes::v1_stream(&request_id))
6616 .header("last-event-id", format!("{request_id}:0"))
6617 .body(axum::body::Body::empty())
6618 .unwrap(),
6619 )
6620 .await
6621 .unwrap();
6622 assert_eq!(response.status(), StatusCode::OK);
6623 assert_eq!(
6624 response
6625 .headers()
6626 .get("x-accel-buffering")
6627 .and_then(|v| v.to_str().ok()),
6628 Some("no"),
6629 "the reconnect needs the same anti-buffering header as the stream"
6630 );
6631 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6632 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6633 assert_eq!(
6634 sse_field(&resumed, "data:")
6635 .iter()
6636 .map(|d| d.to_string())
6637 .collect::<Vec<_>>(),
6638 datas[1..].to_vec()
6639 );
6640 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6641 }
6642
6643 #[tokio::test]
6647 async fn a_last_event_id_from_another_stream_is_refused() {
6648 let app = test_app();
6649 let body = post_sse_raw(&app, resumable_request()).await;
6650 let request_id = sse_field(&body, "id:")[0]
6651 .rsplit_once(':')
6652 .unwrap()
6653 .0
6654 .to_string();
6655
6656 let (status, err) = get_json_with_headers(
6657 &app,
6658 &ferrox_api::routes::v1_stream(&request_id),
6659 &[("last-event-id", "chatcmpl-someone-else:3")],
6660 )
6661 .await;
6662 assert_eq!(status, StatusCode::BAD_REQUEST);
6663 assert_eq!(err["error"]["code"], "bad_last_event_id");
6664 }
6665
6666 #[tokio::test]
6670 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6671 let app = test_app();
6672 let mut request = resumable_request();
6673 request["stream_resumable"] = serde_json::json!(false);
6674 let body = post_sse_raw(&app, request).await;
6675 let request_id = body
6676 .lines()
6677 .find_map(|l| l.strip_prefix("data: "))
6678 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6679 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6680 .unwrap();
6681
6682 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
6683 assert_eq!(status, StatusCode::NOT_FOUND);
6684 assert_eq!(err["error"]["code"], "stream_not_found");
6685 assert!(err["error"]["message"]
6686 .as_str()
6687 .unwrap()
6688 .contains("stream_resumable"));
6689 }
6690
6691 #[test]
6695 fn the_axum_stream_patterns_match_the_published_templates() {
6696 assert_eq!(
6697 axum_path(ferrox_api::routes::V1_STREAM),
6698 "/v1/stream/:request_id"
6699 );
6700 assert_eq!(
6701 axum_path(ferrox_api::routes::V1_STREAM_POLL),
6702 "/v1/stream/:request_id/poll"
6703 );
6704 assert_eq!(
6705 ferrox_api::routes::v1_stream("abc"),
6706 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
6707 );
6708 }
6709
6710 #[test]
6720 fn no_published_template_reaches_the_router_with_its_braces() {
6721 for template in [
6722 ferrox_api::routes::V1_STREAM,
6723 ferrox_api::routes::V1_STREAM_POLL,
6724 ferrox_api::routes::V1_RESPONSE,
6725 ferrox_api::routes::V1_RESPONSE_CANCEL,
6726 ferrox_api::routes::ADMIN_TASK_CANCEL,
6727 ] {
6728 assert!(
6729 template.contains('{'),
6730 "{template} is in the template list but has no placeholder"
6731 );
6732 let mounted = axum_path(template);
6733 assert!(
6734 !mounted.contains('{') && !mounted.contains('}'),
6735 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6736 );
6737 assert!(
6738 mounted.contains(':'),
6739 "{template} lost its placeholder entirely and would match one path only"
6740 );
6741 }
6742 }
6743
6744 #[tokio::test]
6752 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6753 let app = test_app();
6754 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6755 assert_eq!(status, StatusCode::NOT_FOUND);
6756 assert!(
6757 !body.is_null(),
6758 "empty body means axum never matched the route, so the id was read as a literal segment"
6759 );
6760 }
6761
6762 #[tokio::test]
6765 async fn the_task_list_starts_empty_rather_than_absent() {
6766 let app = test_app();
6767 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
6768 assert_eq!(status, StatusCode::OK);
6769 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6770 }
6771
6772 async fn post_json_uri(
6773 app: &Router,
6774 uri: &str,
6775 body: serde_json::Value,
6776 ) -> (StatusCode, serde_json::Value) {
6777 use http_body_util::BodyExt;
6778 use tower::ServiceExt;
6779
6780 let response = app
6781 .clone()
6782 .oneshot(
6783 axum::http::Request::builder()
6784 .method("POST")
6785 .uri(uri)
6786 .header("content-type", "application/json")
6787 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6788 .unwrap(),
6789 )
6790 .await
6791 .unwrap();
6792 let status = response.status();
6793 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6794 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6795 (status, json)
6796 }
6797
6798 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6799 post_json_uri(app, "/v1/chat/completions", body).await.1
6800 }
6801
6802 #[tokio::test]
6811 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6812 let app = test_app();
6813 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6814 assert_eq!(status, StatusCode::OK);
6815
6816 let memory = &body["memory"];
6817 if memory.is_null() {
6818 return;
6821 }
6822 assert!(
6823 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6824 "a read that produced a zero is a broken read, not an idle \
6825 engine: {memory}"
6826 );
6827 assert!(
6828 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6829 "the quantity must travel with the number: {memory}"
6830 );
6831 }
6832
6833 #[tokio::test]
6839 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6840 let app = test_app();
6841 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6842 assert_eq!(status, StatusCode::OK);
6843 for pool in ["kv_pages", "window_slots", "state_slots"] {
6844 assert!(
6845 body["pools"][pool].is_null(),
6846 "{pool} must be null rather than a zero row: {}",
6847 body["pools"]
6848 );
6849 }
6850 }
6851
6852 #[tokio::test]
6860 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6861 use http_body_util::BodyExt;
6862 use tower::ServiceExt;
6863
6864 let app = test_app();
6865 let send = |body: serde_json::Value| {
6866 let app = app.clone();
6867 async move {
6868 app.oneshot(
6869 axum::http::Request::builder()
6870 .method("POST")
6871 .uri(ferrox_api::routes::V1_MESSAGES)
6872 .header("content-type", "application/json")
6873 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6874 .unwrap(),
6875 )
6876 .await
6877 .unwrap()
6878 }
6879 };
6880
6881 let ok = send(serde_json::json!({
6882 "model": "test",
6883 "max_tokens": 1,
6884 "messages": [{"role": "user", "content": "hi"}],
6885 }))
6886 .await;
6887 assert_eq!(ok.status(), StatusCode::OK);
6888 let id = ok
6889 .headers()
6890 .get("request-id")
6891 .expect("a served message names its id")
6892 .to_str()
6893 .unwrap()
6894 .to_string();
6895 assert!(!id.is_empty());
6896
6897 let bad = send(serde_json::json!({"model": "test"})).await;
6900 assert!(bad.status().is_client_error());
6901 let other = bad.headers().get("request-id").expect("errors too");
6902 assert_ne!(other.to_str().unwrap(), id);
6903 let _ = bad.into_body().collect().await.unwrap();
6904 }
6905
6906 #[tokio::test]
6914 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6915 let state = Arc::new(test_state(
6916 test_model_full_byte_vocab(),
6917 ResponseCache::new(1000, Duration::from_secs(3600)),
6918 ));
6919 let app = test_app_with_state(Arc::clone(&state));
6920 let body = serde_json::json!({
6921 "model": "test",
6922 "messages": [{"role": "user", "content": "hi"}],
6923 "max_tokens": 1,
6924 });
6925
6926 state
6927 .maintenance
6928 .lock()
6929 .unwrap()
6930 .begin_rebuild()
6931 .expect("a fresh server is serving, so the rebuild starts");
6932 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6933 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6934 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6935
6936 state.maintenance.lock().unwrap().finish_rebuild(true);
6937 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6938 assert_eq!(
6939 status,
6940 StatusCode::OK,
6941 "the gate reopens; a rebuild is not a latch"
6942 );
6943 }
6944
6945 #[tokio::test]
6950 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6951 let app = test_app();
6952 let (status, body) = post_json_uri(
6953 &app,
6954 ferrox_api::routes::V1_CANCEL,
6955 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6956 )
6957 .await;
6958 assert_eq!(status, StatusCode::NOT_FOUND);
6959 assert_eq!(body["cancelled"], serde_json::json!(false));
6960 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6961 assert!(
6962 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6963 "the verdict must carry a human reason: {body}"
6964 );
6965 }
6966
6967 #[tokio::test]
6971 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6972 let state = Arc::new(test_state(
6973 test_model_full_byte_vocab(),
6974 ResponseCache::new(1000, Duration::from_secs(3600)),
6975 ));
6976 let app = test_app_with_state(Arc::clone(&state));
6977 let (token, _guard) = state.cancels.register("chatcmpl-live");
6978
6979 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6980 assert_eq!(status, StatusCode::OK);
6981 assert_eq!(before["generating_now"], serde_json::json!(1));
6982
6983 let (status, body) = post_json_uri(
6984 &app,
6985 ferrox_api::routes::V1_CANCEL,
6986 serde_json::json!({ "request_id": "chatcmpl-live" }),
6987 )
6988 .await;
6989 assert_eq!(status, StatusCode::OK);
6990 assert_eq!(body["cancelled"], serde_json::json!(true));
6991 assert!(
6992 token.is_cancelled(),
6993 "the endpoint answered ok without setting the flag the decode loop reads"
6994 );
6995 }
6996
6997 #[tokio::test]
6998 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6999 let app = test_app();
7000 let (status, tok) =
7001 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
7002 assert_eq!(status, StatusCode::OK);
7003 let tokens = tok["tokens"].as_array().unwrap();
7004 assert_eq!(tok["count"], tokens.len());
7005 assert!(!tokens.is_empty());
7006
7007 let (status, detok) = post_json_uri(
7008 &app,
7009 "/v1/detokenize",
7010 serde_json::json!({ "tokens": tokens }),
7011 )
7012 .await;
7013 assert_eq!(status, StatusCode::OK);
7014 assert_eq!(detok["text"], "Hi");
7015
7016 let (status, emb) = post_json_uri(
7017 &app,
7018 "/v1/embeddings",
7019 serde_json::json!({
7020 "input": "Hi",
7021 "embedding_type": "mean"
7022 }),
7023 )
7024 .await;
7025 assert_eq!(status, StatusCode::OK);
7026 let vec = emb["data"][0]["embedding"].as_array().unwrap();
7027 assert!(!vec.is_empty());
7028 assert!(vec.iter().all(|v| v.as_f64().is_some()));
7029 }
7030
7031 #[tokio::test]
7037 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
7038 let app = test_app();
7039 let (status, body) = post_json_uri(
7040 &app,
7041 "/v1/embeddings",
7042 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
7043 )
7044 .await;
7045 assert_eq!(status, StatusCode::BAD_REQUEST);
7046 let msg = body["error"]["message"].as_str().unwrap();
7047 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
7048 }
7049
7050 #[tokio::test]
7055 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7056 async fn a_real_embedding_model_serves_v1_embeddings() {
7057 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7058 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7059 if !path.exists() {
7060 eprintln!("SKIP: {} not present", path.display());
7061 return;
7062 }
7063 let encoder = ferrox_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
7064 let mut state = test_state(
7065 test_model_full_byte_vocab(),
7066 ResponseCache::new(1000, Duration::from_secs(3600)),
7067 );
7068 state.embedding = Some(Arc::new(encoder));
7069 let app = test_app_with_state(Arc::new(state));
7070
7071 let (status, body) = post_json_uri(
7072 &app,
7073 "/v1/embeddings",
7074 serde_json::json!({ "input": ["Hello world", "a second input"] }),
7075 )
7076 .await;
7077 assert_eq!(status, StatusCode::OK, "{body}");
7078 assert_eq!(body["model"], "bge-small-en-v1.5");
7079 let data = body["data"].as_array().unwrap();
7080 assert_eq!(data.len(), 2);
7081 for (i, row) in data.iter().enumerate() {
7082 assert_eq!(row["index"], i);
7083 let v: Vec<f64> = row["embedding"]
7084 .as_array()
7085 .unwrap()
7086 .iter()
7087 .map(|x| x.as_f64().unwrap())
7088 .collect();
7089 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
7090 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
7091 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
7092 }
7093 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
7096
7097 let (status, mean) = post_json_uri(
7101 &app,
7102 "/v1/embeddings",
7103 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
7104 )
7105 .await;
7106 assert_eq!(status, StatusCode::OK);
7107 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
7108 }
7109
7110 #[tokio::test]
7122 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7123 async fn an_encoder_can_be_the_loaded_model() {
7124 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7125 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7126 if !path.exists() {
7127 eprintln!("SKIP: {} not present", path.display());
7128 return;
7129 }
7130
7131 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
7135 assert!(
7136 matches!(loaded, model::LoadedModel::Encoder(_)),
7137 "an encoder-only GGUF reached a decoder loader"
7138 );
7139 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
7140 assert!(
7141 matches!(loaded, Loaded::Encoder(_)),
7142 "the encoder did not stay an encoder through activation"
7143 );
7144 assert!(
7145 batcher.is_none() && ceiling.is_none(),
7146 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7147 );
7148
7149 let state = test_state(
7150 test_model_full_byte_vocab(),
7151 ResponseCache::new(1000, Duration::from_secs(3600)),
7152 );
7153 state.swap_active(Some(Arc::new(ActiveModel {
7154 id: None,
7155 loaded,
7156 batcher,
7157 ceiling,
7158 })));
7159 let app = test_app_with_state(Arc::new(state));
7160
7161 let (status, body) = post_json_uri(
7163 &app,
7164 "/v1/embeddings",
7165 serde_json::json!({ "input": "Hello world" }),
7166 )
7167 .await;
7168 assert_eq!(status, StatusCode::OK, "{body}");
7169 assert_eq!(body["model"], "bge-small-en-v1.5");
7170 let v = body["data"][0]["embedding"].as_array().unwrap();
7171 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
7172
7173 let (status, body) = post_json_uri(
7175 &app,
7176 "/v1/chat/completions",
7177 serde_json::json!({
7178 "model": "bge-small-en-v1.5",
7179 "messages": [{"role": "user", "content": "hi"}],
7180 }),
7181 )
7182 .await;
7183 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
7184 let msg = body["error"]["message"].as_str().unwrap();
7185 for fact in [
7186 "bge-small-en-v1.5",
7187 "bert",
7188 "embedding model",
7189 "/v1/embeddings",
7190 ] {
7191 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
7192 }
7193
7194 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
7196 assert_eq!(status, StatusCode::OK);
7197 let entry = &models["data"][0];
7198 assert_eq!(entry["id"], "bge-small-en-v1.5");
7199 assert_eq!(entry["ferrox_model_kind"], "embedding");
7200 assert_eq!(entry["ferrox_tokenizer"], "gguf-wordpiece");
7201 assert_eq!(entry["ferrox_n_embd"], 384);
7202 assert_eq!(entry["ferrox_pooling"], "CLS");
7203 assert_eq!(
7204 entry["ferrox_endpoints"],
7205 serde_json::json!(["/v1/embeddings"])
7206 );
7207 assert!(entry.get("supported_reasoning_efforts").is_none());
7210
7211 let (status, health) = get_json(&app, ferrox_api::routes::HEALTH).await;
7213 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
7214 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
7215 assert_eq!(health["model"]["synthetic_weights"], false);
7216 let weights = health["capabilities"]
7217 .as_array()
7218 .unwrap()
7219 .iter()
7220 .find(|c| c["id"] == ferrox_api::health::capability::REAL_WEIGHTS)
7221 .expect("a real-weights capability row");
7222 let detail = weights["detail"].as_str().unwrap_or_default();
7223 assert!(detail.contains("ENCODER"), "{detail}");
7224 let (status, body) = post_json_uri(
7232 &app,
7233 ferrox_api::routes::V1_TOKENIZE,
7234 serde_json::json!({ "content": "hello world" }),
7235 )
7236 .await;
7237 assert_eq!(
7238 status,
7239 StatusCode::OK,
7240 "an encoder has a real tokenizer: {body}"
7241 );
7242 let tokens = body["tokens"].as_array().expect("tokens array").clone();
7243 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
7244
7245 let (status, body) = post_json_uri(
7246 &app,
7247 ferrox_api::routes::V1_DETOKENIZE,
7248 serde_json::json!({ "tokens": tokens }),
7249 )
7250 .await;
7251 assert_eq!(status, StatusCode::OK, "{body}");
7252 let round_tripped = body["content"].as_str().expect("content").to_string();
7253 assert!(
7254 round_tripped.contains("hello") && round_tripped.contains("world"),
7255 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
7256 );
7257
7258 let (status, _) = post_json_uri(
7261 &app,
7262 "/v1/completions",
7263 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
7264 )
7265 .await;
7266 assert_eq!(
7267 status,
7268 StatusCode::NOT_IMPLEMENTED,
7269 "tokenizing an encoder must not have opened a path to generating with one"
7270 );
7271 }
7272
7273 #[tokio::test]
7278 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7279 use http_body_util::BodyExt;
7280 use tower::ServiceExt;
7281
7282 let fixture = concat!(
7283 "../ferrox-models/tests/fixtures/",
7284 "ferrox_real_moe_test.gguf"
7285 );
7286 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7287 let decoder = Decoder::from_gguf_with_expert_cache(
7288 &fixture,
7289 ferrox_models::config::test_moe_fixture(),
7290 Some(1024 * 1024),
7291 )
7292 .expect("MoE fixture must load store-backed");
7293
7294 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
7298 .layers
7299 .iter()
7300 .map(|_| {
7301 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
7302 })
7303 .collect();
7304 decoder.forward_token(1, 0, &mut caches);
7305
7306 let model = Model::Gguf(GgufModel {
7307 decoder: Arc::new(decoder),
7308 tokenizer: Arc::new(ServerTokenizer::Byte),
7309 stop_tokens: StopTokens::default(),
7310 bos_id: None,
7311 is_synthetic: false,
7312 chat_template: chat_template::PromptTemplate::plain(),
7313 });
7314 let state = Arc::new(test_state(
7315 model,
7316 ResponseCache::new(16, Duration::from_secs(60)),
7317 ));
7318 let app = Router::new()
7319 .route("/metrics", axum::routing::get(metrics))
7320 .route("/v1/chat/completions", post(chat_completions))
7321 .with_state(state);
7322
7323 let fetch_metrics = |app: Router| async move {
7324 let resp = app
7325 .oneshot(
7326 axum::http::Request::builder()
7327 .method("GET")
7328 .uri("/metrics")
7329 .body(axum::body::Body::empty())
7330 .unwrap(),
7331 )
7332 .await
7333 .unwrap();
7334 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7335 String::from_utf8(bytes.to_vec()).unwrap()
7336 };
7337
7338 let after = fetch_metrics(app.clone()).await;
7339 assert!(
7340 after.contains("ferrox_expert_cache_misses_total"),
7341 "streaming model must expose expert-cache metrics: {after}"
7342 );
7343 let misses: u64 = after
7344 .lines()
7345 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
7346 .and_then(|l| l.split_whitespace().nth(1))
7347 .and_then(|v| v.parse().ok())
7348 .expect("misses metric line must parse");
7349 assert!(
7350 misses > 0,
7351 "decode must have read experts through the store: {after}"
7352 );
7353 }
7354
7355 fn weather_tool() -> serde_json::Value {
7356 serde_json::json!({
7357 "type": "function",
7358 "function": {
7359 "name": "get_weather",
7360 "description": "Get the current weather for a location.",
7361 "parameters": {
7362 "type": "object",
7363 "properties": {"location": {"type": "string"}},
7364 "required": ["location"]
7365 }
7366 }
7367 })
7368 }
7369
7370 fn weather_tool_def() -> ToolDef {
7371 ToolDef {
7372 kind: "function".to_string(),
7373 function: ToolFunctionDef {
7374 name: "get_weather".to_string(),
7375 description: Some("Get the current weather for a location.".to_string()),
7376 parameters: Some(serde_json::json!({
7377 "type": "object",
7378 "properties": {"location": {"type": "string"}},
7379 "required": ["location"]
7380 })),
7381 },
7382 }
7383 }
7384
7385 #[test]
7386 fn tool_preamble_mentions_every_tool_name_and_description() {
7387 let preamble = tool_preamble(&[weather_tool_def()]);
7388 assert!(preamble.contains("get_weather"));
7389 assert!(preamble.contains("Get the current weather for a location."));
7390 assert!(preamble.contains("<tool_call>"));
7391 assert!(preamble.contains("</tool_call>"));
7392 }
7393
7394 #[test]
7395 fn a_real_marker_becomes_a_structured_tool_call() {
7396 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7397 let (message, finish) = build_response_message(
7398 text.to_string(),
7399 &[weather_tool_def()],
7400 output::OutputPosture::for_model("test-model"),
7401 "stop",
7402 );
7403 assert_eq!(finish, "tool_calls");
7404 let calls = message.tool_calls.expect("must carry a tool call");
7405 assert_eq!(calls[0].function.name, "get_weather");
7406 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7407 assert_eq!(parsed["location"], "Paris");
7408 }
7409
7410 #[test]
7411 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7412 let (message, finish) = build_response_message(
7413 "just an answer".to_string(),
7414 &[weather_tool_def()],
7415 output::OutputPosture::for_model("test-model"),
7416 "stop",
7417 );
7418 assert_eq!(finish, "stop");
7419 assert!(message.tool_calls.is_none());
7420 assert_eq!(message.content.as_deref(), Some("just an answer"));
7421 }
7422
7423 #[test]
7426 fn a_malformed_payload_is_not_a_tool_call() {
7427 let (message, finish) = build_response_message(
7428 "<tool_call>not valid json at all</tool_call>".to_string(),
7429 &[weather_tool_def()],
7430 output::OutputPosture::for_model("test-model"),
7431 "stop",
7432 );
7433 assert_eq!(finish, "stop");
7434 assert!(message.tool_calls.is_none());
7435 }
7436
7437 #[test]
7440 fn a_tool_that_was_never_offered_is_not_returned() {
7441 let (message, finish) = build_response_message(
7442 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7443 &[weather_tool_def()],
7444 output::OutputPosture::for_model("test-model"),
7445 "stop",
7446 );
7447 assert_eq!(finish, "stop");
7448 assert!(message.tool_calls.is_none());
7449 }
7450
7451 #[test]
7453 fn marker_text_with_no_tools_offered_stays_content() {
7454 let (message, finish) = build_response_message(
7455 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7456 &[],
7457 output::OutputPosture::for_model("test-model"),
7458 "stop",
7459 );
7460 assert_eq!(finish, "stop");
7461 assert!(message.tool_calls.is_none());
7462 assert!(message.content.is_some());
7463 }
7464
7465 #[test]
7469 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7470 let opened = std::cell::Cell::new(0usize);
7471 let mut parser = crate::policy::parser::ToolCallParser::new(
7472 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7473 vec![
7474 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7475 "write_file",
7476 serde_json::json!({"type": "object", "properties": {
7477 "path": {"type": "string"},
7478 "contents": {"type": "string"}
7479 }}),
7480 ),
7481 ],
7482 );
7483 let wire = "<tool_call><function=write_file>\
7484 <parameter=path>\n/tmp/x\n</parameter>\
7485 <parameter=contents>\nhello world\n</parameter>\
7486 </function></tool_call>";
7487
7488 let mut deltas = Vec::new();
7489 let mut text = String::new();
7490 for piece in wire.as_bytes().chunks(7) {
7491 let chunk = String::from_utf8_lossy(piece).into_owned();
7492 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7493 text.push_str(&more_text);
7494 deltas.extend(more);
7495 }
7496 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7497 text.push_str(&more_text);
7498 deltas.extend(more);
7499
7500 assert_eq!(opened.get(), 1, "one call opened");
7501 assert!(text.is_empty(), "the markers are not content: {text:?}");
7502
7503 let first = &deltas[0];
7504 assert_eq!(first.index, 0);
7505 assert_eq!(first.id.as_deref(), Some("call_0"));
7506 assert_eq!(first.kind, Some("function"));
7507 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7508
7509 let joined: String = deltas
7512 .iter()
7513 .filter_map(|d| d.function.arguments.clone())
7514 .collect();
7515 let parsed: serde_json::Value =
7516 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7517 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7518 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7519 assert!(
7520 deltas.len() >= 3,
7521 "the arguments arrived in pieces, not whole: {}",
7522 deltas.len()
7523 );
7524 assert!(
7525 deltas[1..].iter().all(|d| d.function.name.is_none()),
7526 "only the opening delta carries identity"
7527 );
7528 }
7529
7530 #[test]
7532 fn text_around_a_streamed_call_is_still_content() {
7533 let opened = std::cell::Cell::new(0usize);
7534 let mut parser = crate::policy::parser::ToolCallParser::new(
7535 crate::policy::parser::ToolCallFormat::Qwen25,
7536 vec![crate::policy::parser::tool_call::ToolSchema::new(
7537 "get_weather",
7538 )],
7539 );
7540 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7541 \"arguments\": {}}</tool_call> done";
7542 let mut text = String::new();
7543 for piece in wire.as_bytes().chunks(5) {
7544 let chunk = String::from_utf8_lossy(piece).into_owned();
7545 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7546 text.push_str(&more);
7547 }
7548 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7549 text.push_str(&more);
7550
7551 assert_eq!(opened.get(), 1);
7552 assert!(text.starts_with("let me check. "), "{text:?}");
7553 assert!(text.ends_with(" done"), "{text:?}");
7554 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7555 }
7556
7557 #[test]
7560 fn a_reasoning_block_is_split_out_of_the_answer() {
7561 let (message, finish) = build_response_message(
7562 "<think>weighing it up</think>The answer is 4.".to_string(),
7563 &[],
7564 output::OutputPosture::for_model("Qwen3-8B"),
7565 "stop",
7566 );
7567 assert_eq!(finish, "stop");
7568 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7569 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7570 }
7571
7572 #[test]
7575 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7576 let (message, _) = build_response_message(
7577 "Use the <think> tag like this.".to_string(),
7578 &[],
7579 output::OutputPosture::for_model("llama-3.1-8b"),
7580 "stop",
7581 );
7582 assert_eq!(
7583 message.content.as_deref(),
7584 Some("Use the <think> tag like this.")
7585 );
7586 assert!(message.reasoning_content.is_none());
7587 }
7588
7589 #[tokio::test]
7597 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7598 let app = test_app();
7599 let body = serde_json::json!({
7600 "model": "m",
7601 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7602 "max_tokens": 4,
7603 "temperature": 0,
7604 });
7605 let resp = post_json(&app, body).await;
7606 let message = &resp["choices"][0]["message"];
7607 assert!(message["content"].is_string());
7608 assert!(message.get("tool_calls").is_none());
7609 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7610 assert_eq!(resp["usage"]["completion_tokens"], 4);
7611 assert_eq!(
7612 resp["usage"]["total_tokens"],
7613 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7614 );
7615 }
7616
7617 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7618 use http_body_util::BodyExt;
7619 use tower::ServiceExt;
7620
7621 let response = app
7622 .clone()
7623 .oneshot(
7624 axum::http::Request::builder()
7625 .method("GET")
7626 .uri(uri)
7627 .body(axum::body::Body::empty())
7628 .unwrap(),
7629 )
7630 .await
7631 .unwrap();
7632 let status = response.status();
7633 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7634 (status, serde_json::from_slice(&bytes).unwrap())
7635 }
7636
7637 #[tokio::test]
7638 async fn health_answers_a_capability_handshake_not_a_boolean() {
7639 let app = test_app();
7640 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7641 assert_eq!(status, StatusCode::OK);
7642
7643 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7644 assert_eq!(health.state, ferrox_api::HealthState::Ready);
7645 assert!(health.pid > 0);
7646 assert!(health.server_time_unix_ms > 0);
7647 assert_eq!(health.last_request_age_seconds, None);
7650
7651 for id in [
7654 ferrox_api::health::capability::CPU,
7655 ferrox_api::health::capability::METAL,
7656 ferrox_api::health::capability::CUDA,
7657 ferrox_api::health::capability::REAL_WEIGHTS,
7658 ferrox_api::health::capability::CONTINUOUS_BATCHING,
7659 ] {
7660 let cap = health
7661 .capability(id)
7662 .unwrap_or_else(|| panic!("{id} missing"));
7663 assert!(!cap.reason.is_empty(), "{cap:?}");
7664 assert!(!cap.detail.is_empty(), "{cap:?}");
7665 }
7666 let weights = health
7670 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
7671 .unwrap();
7672 assert!(!weights.available);
7673 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
7674 assert!(health.model.as_ref().unwrap().synthetic_weights);
7675 }
7676
7677 #[tokio::test]
7678 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7679 let app = test_app();
7680 let _ = post_json(
7681 &app,
7682 serde_json::json!({
7683 "model": "m",
7684 "messages": [{"role": "user", "content": "\u{1}"}],
7685 "max_tokens": 1,
7686 "temperature": 0,
7687 }),
7688 )
7689 .await;
7690 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7691 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7692 let age = health
7693 .last_request_age_seconds
7694 .expect("a served request is evidence of liveness");
7695 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7696 }
7697
7698 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7700 use http_body_util::BodyExt;
7701 use tower::ServiceExt;
7702
7703 let response = app
7704 .clone()
7705 .oneshot(
7706 axum::http::Request::builder()
7707 .method("POST")
7708 .uri("/v1/chat/completions")
7709 .header("content-type", "application/json")
7710 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7711 .unwrap(),
7712 )
7713 .await
7714 .unwrap();
7715 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7716 String::from_utf8(bytes.to_vec())
7717 .unwrap()
7718 .lines()
7719 .filter_map(|line| line.strip_prefix("data: "))
7720 .filter(|payload| *payload != "[DONE]")
7721 .map(|payload| serde_json::from_str(payload).unwrap())
7722 .collect()
7723 }
7724
7725 #[tokio::test]
7726 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7727 let app = test_app();
7728 let chunks = post_sse_chunks(
7729 &app,
7730 serde_json::json!({
7731 "model": "m",
7732 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7733 "max_tokens": 4,
7734 "temperature": 0,
7735 "stream": true,
7736 }),
7737 )
7738 .await;
7739
7740 assert!(!chunks.is_empty());
7741 let request_id = chunks[0]["request_id"]
7742 .as_str()
7743 .expect("the first chunk names the request")
7744 .to_string();
7745 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7746 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7749 assert!(
7750 chunk.get("request_id").is_none(),
7751 "chunk {i} repeats request_id"
7752 );
7753 }
7754 for chunk in &chunks {
7757 assert_eq!(chunk["id"], serde_json::json!(request_id));
7758 }
7759
7760 let other = post_sse_chunks(
7761 &app,
7762 serde_json::json!({
7763 "model": "m",
7764 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7765 "max_tokens": 4,
7766 "temperature": 0,
7767 "stream": true,
7768 }),
7769 )
7770 .await;
7771 assert_ne!(
7772 other[0]["request_id"].as_str().unwrap(),
7773 request_id,
7774 "two concurrent chats must not share an id"
7775 );
7776 }
7777
7778 #[tokio::test]
7779 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7780 let app = test_app();
7781 let resp = post_json(
7782 &app,
7783 serde_json::json!({
7784 "model": "m",
7785 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7786 "max_tokens": 2,
7787 "temperature": 0,
7788 }),
7789 )
7790 .await;
7791 assert_eq!(resp["id"], resp["request_id"]);
7792 assert!(resp["request_id"]
7793 .as_str()
7794 .unwrap()
7795 .starts_with("chatcmpl-"));
7796 }
7797
7798 #[tokio::test]
7801 async fn usage_carries_separate_prefill_and_decode_timings() {
7802 let app = test_app();
7803 let resp = post_json(
7804 &app,
7805 serde_json::json!({
7806 "model": "m",
7807 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7808 "max_tokens": 4,
7809 "temperature": 0,
7810 }),
7811 )
7812 .await;
7813 let usage = &resp["usage"];
7814 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7815 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7816 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7817 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7818 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7820 }
7821
7822 #[tokio::test]
7830 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7831 let app = test_app();
7832 let body = serde_json::json!({
7833 "model": "m",
7834 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7835 "max_tokens": 4,
7836 "temperature": 0,
7837 "tools": [weather_tool()],
7838 });
7839 let resp = post_json(&app, body).await;
7840 let message = &resp["choices"][0]["message"];
7841 assert!(
7842 message["content"].is_string(),
7843 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7844 );
7845 assert!(message.get("tool_calls").is_none());
7846 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7850 }
7851
7852 #[tokio::test]
7856 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7857 let app = test_app();
7858 let body = serde_json::json!({
7859 "model": "m",
7860 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7861 "max_tokens": 3,
7862 "temperature": 0,
7863 });
7864 let first = post_json(&app, body.clone()).await;
7865 assert_eq!(first["ferrox_cache"], "miss");
7866 let second = post_json(&app, body).await;
7867 assert_eq!(second["ferrox_cache"], "hit");
7868 assert_eq!(
7869 first["choices"][0]["message"]["content"],
7870 second["choices"][0]["message"]["content"]
7871 );
7872 assert_eq!(
7873 first["choices"][0]["finish_reason"],
7874 second["choices"][0]["finish_reason"]
7875 );
7876 assert_eq!(first["usage"], second["usage"]);
7877 assert_eq!(second["usage"]["completion_tokens"], 3);
7878 }
7879
7880 #[tokio::test]
7891 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7892 let app = test_app();
7893 let plain = serde_json::json!({
7894 "model": "m",
7895 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7896 "max_tokens": 3,
7897 "temperature": 0,
7898 });
7899
7900 let first = post_json(&app, plain.clone()).await;
7901 assert_eq!(first["ferrox_cache"], "miss");
7902 let unconstrained = first["choices"][0]["message"]["content"]
7903 .as_str()
7904 .expect("content")
7905 .to_string();
7906
7907 let mut constrained = plain.clone();
7908 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7909 let second = post_json(&app, constrained).await;
7910 assert_eq!(
7911 second["ferrox_cache"], "miss",
7912 "a grammar is part of the key, so this body has never been answered"
7913 );
7914 let constrained_answer = second["choices"][0]["message"]["content"]
7918 .as_str()
7919 .expect("content")
7920 .to_string();
7921 assert!(
7922 constrained_answer.contains("-> \"yes\"]"),
7923 "the grammar must have been compiled AND applied, not skipped \
7924 by a cache hit: {constrained_answer}"
7925 );
7926 assert_ne!(
7927 constrained_answer, unconstrained,
7928 "the constrained request was served the unconstrained answer"
7929 );
7930
7931 let third = post_json(&app, plain).await;
7935 assert_eq!(third["ferrox_cache"], "hit");
7936 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7937 }
7938
7939 #[tokio::test]
7961 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7962 let state = Arc::new(test_state(
7963 test_model_full_byte_vocab(),
7964 ResponseCache::new(1000, Duration::from_secs(3600)),
7965 ));
7966 let app = test_app_with_state(state.clone());
7967 let plain = serde_json::json!({
7968 "model": "m",
7969 "messages": [
7970 {"role": "system", "content": "Answer in JSON when it helps."},
7971 {"role": "user", "content": "\u{1}\u{2}"},
7972 ],
7973 "max_tokens": 3,
7974 "temperature": 0,
7975 });
7976
7977 let first = post_json(&app, plain.clone()).await;
7978 assert_eq!(first["ferrox_cache"], "miss");
7979 assert_eq!(state.cache_stats().entries, 1);
7980
7981 let mut as_json = plain.clone();
7982 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7983 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7984 assert_eq!(
7985 status,
7986 StatusCode::BAD_REQUEST,
7987 "the demo banner is not a JSON object, whoever generated it"
7988 );
7989 assert_eq!(
7990 state.cache_stats().hits,
7991 0,
7992 "a json_object request must not be answered from an entry the \
7993 JSON mask never produced"
7994 );
7995 assert_eq!(
7996 state.cache_stats().entries,
7997 2,
7998 "json_object must key its own entry, not reuse the unconstrained \
7999 one it happens to render the same prompt as"
8000 );
8001 }
8002
8003 #[tokio::test]
8015 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
8016 let app = test_app_with_state(Arc::new(test_state(
8017 test_model_full_byte_vocab_with_eos(Some(0x77)),
8018 ResponseCache::new(1000, Duration::from_secs(3600)),
8019 )));
8020 let body = serde_json::json!({
8021 "model": "m",
8022 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8023 "max_tokens": 6,
8024 "temperature": 0,
8025 });
8026
8027 let stopped = post_json(&app, body.clone()).await;
8028 assert_eq!(stopped["ferrox_cache"], "miss");
8029 assert_eq!(
8030 stopped["choices"][0]["finish_reason"], "stop",
8031 "the fixture is only meaningful if the model's EOS really fires here"
8032 );
8033 assert_eq!(stopped["usage"]["completion_tokens"], 1);
8034
8035 let mut ignoring = body.clone();
8036 ignoring["ignore_eos"] = serde_json::json!(true);
8037 let ran_on = post_json(&app, ignoring).await;
8038 assert_eq!(
8039 ran_on["ferrox_cache"], "miss",
8040 "ignore_eos is part of the key, so this body has never been answered"
8041 );
8042 assert_eq!(
8043 ran_on["usage"]["completion_tokens"], 6,
8044 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
8045 );
8046 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
8047 assert_ne!(
8048 ran_on["choices"][0]["message"]["content"],
8049 stopped["choices"][0]["message"]["content"]
8050 );
8051 }
8052
8053 #[tokio::test]
8059 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
8060 let session_app = test_app();
8061 let manual_app = test_app();
8062
8063 let turn1 = post_json(
8065 &session_app,
8066 serde_json::json!({
8067 "model": "m",
8068 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8069 "session_id": "s1",
8070 "max_tokens": 5,
8071 "temperature": 0,
8072 }),
8073 )
8074 .await;
8075 let reply1 = turn1["choices"][0]["message"]["content"]
8076 .as_str()
8077 .unwrap()
8078 .to_string();
8079
8080 let manual_turn1 = post_json(
8085 &manual_app,
8086 serde_json::json!({
8087 "model": "m",
8088 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8089 "max_tokens": 5,
8090 "temperature": 0,
8091 }),
8092 )
8093 .await;
8094 assert_eq!(
8095 manual_turn1["choices"][0]["message"]["content"]
8096 .as_str()
8097 .unwrap(),
8098 reply1
8099 );
8100
8101 let turn2 = post_json(
8103 &session_app,
8104 serde_json::json!({
8105 "model": "m",
8106 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
8107 "session_id": "s1",
8108 "max_tokens": 5,
8109 "temperature": 0,
8110 }),
8111 )
8112 .await;
8113 let reply2 = turn2["choices"][0]["message"]["content"]
8114 .as_str()
8115 .unwrap()
8116 .to_string();
8117
8118 let manual_turn2 = post_json(
8122 &manual_app,
8123 serde_json::json!({
8124 "model": "m",
8125 "messages": [
8126 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
8127 {"role": "assistant", "content": reply1},
8128 {"role": "user", "content": "\u{4}\u{5}"},
8129 ],
8130 "max_tokens": 5,
8131 "temperature": 0,
8132 }),
8133 )
8134 .await;
8135 assert_eq!(
8136 manual_turn2["choices"][0]["message"]["content"]
8137 .as_str()
8138 .unwrap(),
8139 reply2,
8140 "resuming a session must produce identical output to manually resending the full history"
8141 );
8142 }
8143
8144 #[test]
8147 fn lock_cache_recovers_from_a_poisoned_mutex() {
8148 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
8149
8150 let poison_cache = Arc::clone(&cache);
8151 let _ = std::thread::spawn(move || {
8152 let _guard = poison_cache.lock().unwrap();
8153 panic!("simulated panic while holding the lock");
8154 })
8155 .join();
8156
8157 let recovered = lock_cache(&cache);
8159 assert_eq!(recovered.stats().entries, 0);
8160 }
8161
8162 #[test]
8163 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8164 let mut req_body = serde_json::json!({
8165 "model": "m",
8166 "messages": [{"role": "user", "content": "hi"}],
8167 });
8168 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8169 assert!(
8170 req.is_cacheable(),
8171 "default (temperature 0) must be cacheable"
8172 );
8173
8174 req_body["temperature"] = serde_json::json!(0.8);
8175 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8176 assert!(
8177 !req.is_cacheable(),
8178 "unseeded sampling must never be cacheable"
8179 );
8180
8181 req_body["seed"] = serde_json::json!(42);
8182 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
8183 assert!(
8184 req.is_cacheable(),
8185 "sampling with an explicit seed is deterministic and must be cacheable"
8186 );
8187 }
8188
8189 const GRADED: &str = "{% if reasoning_effort %}\
8193 {% if reasoning_effort not in ['low','medium','high'] %}\
8194 {{ raise_exception('unsupported effort') }}\
8195 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
8196 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
8197
8198 fn graded_template() -> chat_template::PromptTemplate {
8199 chat_template::PromptTemplate::from_gguf_metadata(
8200 Some(GRADED),
8201 Some("qwen3"),
8202 false,
8203 true,
8204 None,
8205 None,
8206 )
8207 }
8208
8209 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
8210 serde_json::from_value(value).expect("request")
8211 }
8212
8213 #[test]
8220 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
8221 let req = chat_request(serde_json::json!({
8222 "model": "m",
8223 "messages": [{"role": "user", "content": "hi"}],
8224 "grammar": "root ::= \"a\"+",
8225 }));
8226 req.validate_supported_fields()
8227 .expect("a valid grammar is a valid request");
8228 let params = req
8229 .generation_params()
8230 .expect("a valid grammar compiles at params time too");
8231 assert!(
8232 params.grammar.is_some(),
8233 "the grammar was dropped between the wire and the sampler"
8234 );
8235 assert!(
8236 params.needs_vocab_logits(),
8237 "a grammar request that may fold lm_head into a GPU argmax is \
8238 a grammar request served unconstrained"
8239 );
8240
8241 let plain = chat_request(serde_json::json!({
8242 "model": "m",
8243 "messages": [{"role": "user", "content": "hi"}],
8244 }));
8245 assert!(plain.generation_params().unwrap().grammar.is_none());
8246 }
8247
8248 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
8249 chat_request(serde_json::json!({
8250 "model": "m",
8251 "messages": [{"role": "user", "content": "weather in Rome?"}],
8252 "tools": [weather_tool()],
8253 "tool_choice": tool_choice,
8254 }))
8255 }
8256
8257 #[test]
8261 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
8262 for choice in [
8263 serde_json::json!("required"),
8264 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
8265 ] {
8266 let req = tool_request(choice.clone());
8267 req.validate_supported_fields()
8268 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
8269 let params = req
8270 .generation_params_for_template(&graded_template(), "Qwen3-8B")
8271 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
8272 let grammar = params
8273 .grammar
8274 .as_ref()
8275 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
8276 assert!(
8277 grammar.is_awaiting_trigger(),
8278 "the model must be free to think before it calls"
8279 );
8280 assert!(
8281 !grammar.allows_eog(),
8282 "{choice} must not be able to end the turn without a call"
8283 );
8284 assert!(
8290 params.needs_vocab_logits(),
8291 "{choice} would let a backend return a token id instead of logits"
8292 );
8293 assert!(
8294 !generate::greedy_gpu_fold_allowed(¶ms),
8295 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8296 );
8297 }
8298 }
8299
8300 #[test]
8302 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8303 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8304 let req = tool_request(choice.clone());
8305 req.validate_supported_fields().expect("still supported");
8306 let params = match req.generation_params_for_template(&graded_template(), "Qwen3-8B") {
8307 Ok(p) => p,
8308 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
8309 };
8310 assert!(
8311 params.grammar.is_none(),
8312 "{choice} does not force a call and must not be constrained"
8313 );
8314 }
8315 }
8316
8317 #[test]
8320 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8321 let req = chat_request(serde_json::json!({
8323 "model": "m",
8324 "messages": [{"role": "user", "content": "hi"}],
8325 "tool_choice": "required",
8326 }));
8327 let (status, _) = req
8328 .validate_supported_fields()
8329 .expect_err("nothing to call");
8330 assert_eq!(status, StatusCode::BAD_REQUEST);
8331
8332 let req =
8334 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8335 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8336 assert_eq!(status, StatusCode::BAD_REQUEST);
8337 assert_eq!(body["error"]["param"], "tool_choice");
8338
8339 let req = tool_request(serde_json::json!({"type": "function"}));
8341 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8342 assert_eq!(status, StatusCode::BAD_REQUEST);
8343
8344 let req = chat_request(serde_json::json!({
8346 "model": "m",
8347 "messages": [{"role": "user", "content": "hi"}],
8348 "tools": [weather_tool()],
8349 "tool_choice": "required",
8350 "grammar": "root ::= \"a\"+",
8351 }));
8352 let (status, _) = req
8353 .validate_supported_fields()
8354 .expect_err("a grammar and a forced call are two constraints");
8355 assert_eq!(status, StatusCode::BAD_REQUEST);
8356
8357 let req = tool_request(serde_json::json!("required"));
8363 let (status, Json(body)) =
8364 match req.generation_params_for_template(&graded_template(), "Gemma4-27B") {
8365 Err(e) => e,
8366 Ok(_) => panic!("a gemma4 call's arguments are not an object rule"),
8367 };
8368 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8369 assert!(
8370 body["error"]["message"]
8371 .as_str()
8372 .unwrap()
8373 .contains("gemma4"),
8374 "{body}"
8375 );
8376 }
8377
8378 #[test]
8381 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8382 let req = chat_request(serde_json::json!({
8383 "model": "m",
8384 "messages": [{"role": "user", "content": "hi"}],
8385 "grammar": "root ::= \"a",
8386 }));
8387 let (status, Json(body)) = req
8388 .validate_supported_fields()
8389 .expect_err("this does not parse");
8390 assert_eq!(status, StatusCode::BAD_REQUEST);
8391 assert_eq!(body["error"]["param"], "grammar");
8392 assert!(req.generation_params().is_err(), "and again at params time");
8393 }
8394
8395 #[test]
8402 fn response_format_json_schema_becomes_the_requests_grammar() {
8403 let req = chat_request(serde_json::json!({
8404 "model": "m",
8405 "messages": [{"role": "user", "content": "hi"}],
8406 "response_format": {
8407 "type": "json_schema",
8408 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8409 },
8410 }));
8411 req.validate_supported_fields()
8412 .expect("a boolean schema converts");
8413 let params = req.generation_params().expect("and compiles");
8414 let grammar = params.grammar.expect("the schema is the grammar");
8415 let mut g = (*grammar).clone();
8416 g.accept_token(0, b"true").expect("a boolean is accepted");
8417 assert!(g.allows_eog(), "and completes the parse");
8418 assert!(
8419 !params.json_object,
8420 "a schema is not the json_object character-class mask"
8421 );
8422 }
8423
8424 #[test]
8428 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8429 let req = chat_request(serde_json::json!({
8430 "model": "m",
8431 "messages": [{"role": "user", "content": "hi"}],
8432 "response_format": {
8433 "type": "json_schema",
8434 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8435 },
8436 }));
8437 let (status, Json(body)) = req
8438 .validate_supported_fields()
8439 .expect_err("minimum has no grammar in this port");
8440 assert_eq!(status, StatusCode::BAD_REQUEST);
8441 assert!(
8442 body["error"]["message"]
8443 .as_str()
8444 .expect("a message")
8445 .contains("minimum"),
8446 "the refusal must name the keyword: {body}"
8447 );
8448 assert!(req.generation_params().is_err(), "and again at params time");
8449 }
8450
8451 #[test]
8457 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8458 let req = chat_request(serde_json::json!({
8459 "model": "m",
8460 "messages": [{"role": "user", "content": "hi"}],
8461 "tool_choice": "required",
8462 "tools": [{
8463 "type": "function",
8464 "function": {"name": "f", "parameters": {"type": "object"}},
8465 }],
8466 "response_format": {
8467 "type": "json_schema",
8468 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8469 },
8470 }));
8471 let (status, Json(body)) = req
8472 .validate_supported_fields()
8473 .expect_err("two constraints, one generation");
8474 assert_eq!(status, StatusCode::BAD_REQUEST);
8475 assert_eq!(body["error"]["param"], "tool_choice");
8476 }
8477
8478 #[test]
8481 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8482 let req = chat_request(serde_json::json!({
8483 "model": "m",
8484 "messages": [{"role": "user", "content": "hi"}],
8485 }));
8486 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8487 }
8488
8489 #[test]
8494 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8495 let asked = chat_request(serde_json::json!({
8496 "model": "m",
8497 "messages": [{"role": "user", "content": "hi"}],
8498 "min_p": 0.07,
8499 }));
8500 assert_eq!(asked.sampling_params().min_p, 0.07);
8501
8502 let silent = chat_request(serde_json::json!({
8503 "model": "m",
8504 "messages": [{"role": "user", "content": "hi"}],
8505 }));
8506 assert_eq!(
8507 silent.sampling_params().min_p,
8508 0.0,
8509 "an unset min_p must be off, not llama.cpp's CLI default"
8510 );
8511 }
8512
8513 #[test]
8522 fn no_sampler_knob_is_missing_from_the_cache_key() {
8523 let base = serde_json::json!({
8524 "model": "m",
8525 "messages": [{"role": "user", "content": "hi"}],
8526 "seed": 1,
8527 });
8528 let key_for = |body: serde_json::Value| {
8529 let req = chat_request(body);
8530 let params = req.generation_params().expect("params");
8531 req.cache_key("prompt", ¶ms)
8532 };
8533 let baseline = key_for(base.clone());
8534 for (knob, value) in [
8535 ("temperature", serde_json::json!(0.5)),
8536 ("top_p", serde_json::json!(0.9)),
8537 ("min_p", serde_json::json!(0.05)),
8538 ("top_k", serde_json::json!(40)),
8539 ("repetition_penalty", serde_json::json!(1.1)),
8540 ("presence_penalty", serde_json::json!(0.3)),
8541 ("frequency_penalty", serde_json::json!(0.3)),
8542 ] {
8543 let mut body = base.clone();
8544 body[knob] = value;
8545 assert_ne!(
8546 key_for(body),
8547 baseline,
8548 "`{knob}` is not in the cache key: two requests differing \
8549 only in it would share one cached answer"
8550 );
8551 }
8552 }
8553
8554 #[test]
8564 fn no_constraint_is_missing_from_the_cache_key() {
8565 let base = serde_json::json!({
8566 "model": "m",
8567 "messages": [{"role": "user", "content": "pick one"}],
8568 });
8569 let key_for = |body: serde_json::Value| {
8570 let req = chat_request(body);
8571 let params = req.generation_params().expect("params");
8572 req.cache_key("prompt", ¶ms)
8573 };
8574 let baseline = key_for(base.clone());
8575 for (field, value) in [
8576 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8577 (
8578 "response_format",
8579 serde_json::json!({"type": "json_object"}),
8580 ),
8581 (
8582 "response_format",
8583 serde_json::json!({"type": "json_schema", "json_schema": {
8584 "name": "answer",
8585 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8586 }}),
8587 ),
8588 ("ignore_eos", serde_json::json!(true)),
8589 ("stop", serde_json::json!(["\n"])),
8590 ("max_tokens", serde_json::json!(7)),
8591 ] {
8592 let mut body = base.clone();
8593 body[field] = value.clone();
8594 assert_ne!(
8595 key_for(body),
8596 baseline,
8597 "`{field}: {value}` is not in the cache key: two requests \
8598 differing only in it would share one cached answer"
8599 );
8600 }
8601 }
8602
8603 #[test]
8607 fn an_explicit_zero_output_budget_is_a_client_error() {
8608 let req = chat_request(serde_json::json!({
8609 "model": "m",
8610 "messages": [{"role": "user", "content": "hi"}],
8611 "max_tokens": 0,
8612 }));
8613 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8614 assert_eq!(status, StatusCode::BAD_REQUEST);
8615 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8616 }
8617
8618 #[test]
8621 fn a_request_can_turn_thinking_off() {
8622 let template = graded_template();
8623 for body in [
8624 serde_json::json!({
8625 "model": "m",
8626 "messages": [{"role": "user", "content": "hi"}],
8627 "reasoning_effort": "none",
8628 }),
8629 serde_json::json!({
8630 "model": "m",
8631 "messages": [{"role": "user", "content": "hi"}],
8632 "thinking": {"type": "disabled"},
8633 }),
8634 ] {
8635 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8636 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8637 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8638 assert!(!kwargs.contains_key("reasoning_effort"));
8641 }
8642 }
8643
8644 #[test]
8647 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8648 let template = graded_template();
8649 let kwargs = chat_request(serde_json::json!({
8650 "model": "m",
8651 "messages": [{"role": "user", "content": "hi"}],
8652 "reasoning_effort": "high",
8653 "thinking": {"type": "disabled"},
8654 }))
8655 .resolve_template_kwargs(&template);
8656 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8657 assert!(!kwargs.contains_key("reasoning_effort"));
8658 }
8659
8660 #[test]
8663 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8664 let req = chat_request(serde_json::json!({
8665 "model": "m",
8666 "messages": [{"role": "user", "content": "hi"}],
8667 "thinking": {"type": "disable"},
8668 }));
8669 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8670 assert_eq!(status, StatusCode::BAD_REQUEST);
8671 }
8672
8673 #[test]
8676 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8677 let template = graded_template();
8678 let kwargs = chat_request(serde_json::json!({
8679 "model": "m",
8680 "messages": [{"role": "user", "content": "hi"}],
8681 "reasoning_effort": "none",
8682 "chat_template_kwargs": {"enable_thinking": true},
8683 }))
8684 .resolve_template_kwargs(&template);
8685 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8686 }
8687
8688 #[test]
8692 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8693 let template = graded_template();
8694 let req = chat_request(serde_json::json!({
8695 "model": "m",
8696 "messages": [{"role": "user", "content": "hi"}],
8697 "reasoning_effort": "minimal",
8698 }));
8699 let kwargs = req.resolve_template_kwargs(&template);
8700 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8701 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8702 assert!(prompt.starts_with("E:low|"), "{prompt}");
8703 }
8704
8705 #[test]
8709 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8710 let template = graded_template();
8711 let req = chat_request(serde_json::json!({
8712 "model": "m",
8713 "messages": [{"role": "user", "content": "hi"}],
8714 "chat_template_kwargs": {"reasoning_effort": "none"},
8715 }));
8716 let kwargs = req.resolve_template_kwargs(&template);
8717 assert!(!kwargs.contains_key("reasoning_effort"));
8718 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8719 assert_eq!(prompt, "hi");
8720 }
8721
8722 #[test]
8725 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8726 let template = graded_template();
8727 let req = chat_request(serde_json::json!({
8728 "model": "m",
8729 "messages": [{"role": "user", "content": "hi"}],
8730 "reasoning_effort": "low",
8731 "chat_template_kwargs": {"reasoning_effort": "high"},
8732 }));
8733 assert_eq!(
8734 req.resolve_template_kwargs(&template)["reasoning_effort"],
8735 serde_json::json!("high")
8736 );
8737 }
8738
8739 #[test]
8743 fn offering_tools_turns_thinking_on_by_itself() {
8744 let template = graded_template();
8745 let quiet = chat_request(serde_json::json!({
8746 "model": "m",
8747 "messages": [{"role": "user", "content": "hi"}],
8748 }));
8749 assert!(!quiet
8750 .resolve_template_kwargs(&template)
8751 .contains_key("enable_thinking"));
8752
8753 let with_tools = chat_request(serde_json::json!({
8754 "model": "m",
8755 "messages": [{"role": "user", "content": "hi"}],
8756 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8757 }));
8758 let kwargs = with_tools.resolve_template_kwargs(&template);
8759 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8760 let prompt =
8761 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8762 assert!(prompt.starts_with("THINK|"), "{prompt}");
8763 }
8764
8765 #[test]
8770 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8771 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8772 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8773 Some("qwen3"),
8774 false,
8775 true,
8776 None,
8777 None,
8778 );
8779 let req = chat_request(serde_json::json!({
8780 "model": "m",
8781 "messages": [{"role": "user", "content": "hi"}],
8782 "chat_template_kwargs": {"enable_thinking": true},
8783 }));
8784 let kwargs = req.resolve_template_kwargs(&opener);
8785 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8786 assert!(prompt.ends_with("<think>"), "{prompt}");
8787
8788 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8791 let (message, _) = build_response_message(
8792 "weighing it up</think>Paris.".to_string(),
8793 &[],
8794 posture,
8795 "stop",
8796 );
8797 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8798 assert_eq!(message.content.as_deref(), Some("Paris."));
8799
8800 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8803 let (message, _) = build_response_message(
8804 "weighing it up</think>Paris.".to_string(),
8805 &[],
8806 closed,
8807 "stop",
8808 );
8809 assert_eq!(message.reasoning_content, None);
8810 }
8811
8812 #[test]
8813 fn stop_param_accepts_both_single_string_and_array() {
8814 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8815 "model": "m",
8816 "messages": [{"role": "user", "content": "hi"}],
8817 "stop": "END",
8818 }))
8819 .unwrap();
8820 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8821
8822 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8823 "model": "m",
8824 "messages": [{"role": "user", "content": "hi"}],
8825 "stop": ["A", "B"],
8826 }))
8827 .unwrap();
8828 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8829 }
8830
8831 #[test]
8832 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8833 let model = test_model();
8834 let result = run_generation(
8835 &model,
8836 "hello",
8837 &greedy_params(4),
8838 None,
8839 None,
8840 None,
8841 None,
8842 None,
8843 None,
8844 );
8845 assert!(matches!(
8846 result,
8847 Err(generate::DecodeError::TokenOutOfVocab { .. })
8848 ));
8849 }
8850
8851 #[test]
8855 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8856 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8858 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8859
8860 let holder_pool = Arc::clone(&pool);
8861 let holder = std::thread::spawn(move || {
8862 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8863 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8865 drop(held);
8866 });
8867 std::thread::sleep(Duration::from_millis(15));
8868
8869 let config = generate::KvPoolConfig {
8870 pool,
8871 queue_wait: Duration::ZERO,
8872 };
8873 let result = run_generation(
8874 &model,
8875 &prompt,
8876 &greedy_params(4),
8877 Some(&config),
8878 None,
8879 None,
8880 None,
8881 None,
8882 None,
8883 );
8884 assert!(matches!(
8885 result,
8886 Err(generate::DecodeError::KvPoolExhausted)
8887 ));
8888
8889 let (status, _body) = decode_error_response(result.unwrap_err());
8890 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8891 holder.join().unwrap();
8892 }
8893
8894 #[test]
8902 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8903 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8905 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
8907 let config = generate::KvPoolConfig {
8908 pool,
8909 queue_wait: Duration::ZERO,
8910 };
8911
8912 let result = run_generation(
8913 &model,
8914 &prompt,
8915 &greedy_params(4),
8916 Some(&config),
8917 None,
8918 None,
8919 None,
8920 None,
8921 None,
8922 );
8923 let err = result.expect_err("one block cannot hold two layers' caches");
8924 assert!(
8925 matches!(
8926 &err,
8927 generate::DecodeError::KvBudgetExceeded { binding, .. }
8928 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
8929 ),
8930 "expected an immovable device-memory refusal, got {err:?}"
8931 );
8932 let (status, _body) = decode_error_response(err);
8933 assert_eq!(status, StatusCode::BAD_REQUEST);
8934 }
8935
8936 #[test]
8942 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8943 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8944 queued: 512,
8945 cap: 512,
8946 });
8947 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8948 assert_eq!(body["error"]["retry_after_seconds"], 1);
8949 let message = body["error"]["message"].as_str().expect("message");
8950 assert!(message.contains("512"), "{message}");
8951 }
8952
8953 #[test]
8954 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8955 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8956 token: 99,
8957 vocab_size: 32,
8958 });
8959 assert!(
8960 body["error"]["retry_after_seconds"].is_null(),
8961 "retrying a prompt this model cannot tokenize never helps"
8962 );
8963 }
8964
8965 #[test]
8966 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8967 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8968 token: 99,
8969 vocab_size: 32,
8970 });
8971 assert_eq!(status, StatusCode::BAD_REQUEST);
8972 }
8973
8974 #[test]
8975 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8976 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8978 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8979 let config = generate::KvPoolConfig {
8980 pool: pool.clone(),
8981 queue_wait: Duration::ZERO,
8982 };
8983
8984 let (_, finish, _usage) = run_generation(
8985 &model,
8986 &prompt,
8987 &greedy_params(4),
8988 Some(&config),
8989 None,
8990 None,
8991 None,
8992 None,
8993 None,
8994 )
8995 .unwrap();
8996 assert_eq!(finish, FinishReason::Length);
8997 assert_eq!(
8998 pool.lock().unwrap().free_blocks(),
8999 2,
9000 "a completed request must return its blocks to the pool"
9001 );
9002 }
9003
9004 #[tokio::test]
9009 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
9010 let model = Arc::new(test_model());
9011 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9012
9013 let mut handles = Vec::new();
9014 for _ in 0..8 {
9015 let model = Arc::clone(&model);
9016 let prompt = prompt.clone();
9017 handles.push(tokio::task::spawn_blocking(move || {
9018 run_generation(
9019 &model,
9020 &prompt,
9021 &greedy_params(6),
9022 None,
9023 None,
9024 None,
9025 None,
9026 None,
9027 None,
9028 )
9029 .unwrap()
9030 }));
9031 }
9032
9033 let mut results = Vec::new();
9034 for h in handles {
9035 results.push(h.await.unwrap());
9036 }
9037 for r in &results[1..] {
9042 assert_eq!(r.0, results[0].0, "decoded chunks must match");
9043 assert_eq!(r.1, results[0].1, "finish reason must match");
9044 assert_eq!(
9045 r.2.prompt_tokens, results[0].2.prompt_tokens,
9046 "prompt token count must match"
9047 );
9048 assert_eq!(
9049 r.2.completion_tokens, results[0].2.completion_tokens,
9050 "completion token count must match"
9051 );
9052 }
9053 }
9054
9055 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
9062 let mut header_entries = Vec::new();
9063 let mut data = Vec::new();
9064 for (name, shape, values) in tensors {
9065 let start = data.len();
9066 for v in values {
9067 data.extend_from_slice(&v.to_le_bytes());
9068 }
9069 let end = data.len();
9070 let shape_str = shape
9071 .iter()
9072 .map(|d| d.to_string())
9073 .collect::<Vec<_>>()
9074 .join(",");
9075 header_entries.push(format!(
9076 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
9077 ));
9078 }
9079 let header = format!("{{{}}}", header_entries.join(","));
9080 let header_bytes = header.as_bytes();
9081 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
9082 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
9083 out.extend_from_slice(header_bytes);
9084 out.extend_from_slice(&data);
9085 out
9086 }
9087
9088 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9099 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9100 use ferrox_models::kimi_loader::KimiRealHparams;
9101 use ferrox_moe::{GatingFunction, MoeLayerConfig};
9102
9103 let hidden_dim = 8;
9104 let kda_num_heads = 2;
9105 let kda_head_dim = 3;
9106 let kda_proj = kda_num_heads * kda_head_dim;
9107 let conv_kernel = 4;
9108 let dense_intermediate = 5;
9109 let vocab_size = 256;
9113 let mla_num_heads = 1;
9114 let mla_q_lora_rank = 2;
9115 let mla_kv_lora_rank = 2;
9116 let mla_qk_nope_head_dim = 2;
9117 let mla_qk_rope_head_dim = 2;
9118 let mla_v_head_dim = 2;
9119
9120 let model_cfg = ferrox_models::ModelConfig {
9121 name: "synthetic-kimi-server-test",
9122 n_layers: 1,
9123 hidden_dim,
9124 n_heads: 1,
9125 n_kv_heads: 1,
9126 head_dim: 4,
9127 vocab_size,
9128 rope_theta: 10000.0,
9129 rms_norm_eps: 1e-5,
9130 sliding_window: None,
9131 moe: MoeLayerConfig {
9132 expert_weights_scale: 1.0,
9133 n_experts: 1,
9134 n_experts_active: 1,
9135 n_shared_experts: 0,
9136 hidden_dim,
9137 expert_ffn_dim: 4,
9138 gating: GatingFunction::Sigmoid,
9139 norm_topk_prob: true,
9140 expert_group_count: None,
9141 expert_group_used_count: None,
9142 },
9143 n_dense_leading_layers: 1,
9148 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
9149 kda_layers: vec![1],
9150 full_attn_layers: vec![],
9151 mla: MlaConfig {
9152 num_heads: mla_num_heads,
9153 q_lora_rank: mla_q_lora_rank,
9154 kv_lora_rank: mla_kv_lora_rank,
9155 qk_nope_head_dim: mla_qk_nope_head_dim,
9156 qk_rope_head_dim: mla_qk_rope_head_dim,
9157 v_head_dim: mla_v_head_dim,
9158 use_output_gate: true,
9159 rope: None,
9160 },
9161 kda: KdaConfig {
9162 num_heads: kda_num_heads,
9163 head_dim: kda_head_dim,
9164 short_conv_kernel_size: conv_kernel,
9165 gate_lower_bound: -5.0,
9166 use_full_rank_gate: true,
9167 },
9168 }),
9169 rope_freqs: None,
9170 rope_attn_factor: 1.0,
9171 rope_dim: None,
9172 rope_freqs_long: None,
9173 rope_freqs_short: None,
9174 rope_orig_ctx: None,
9175 rope_layout: ferrox_models::config::RopeLayout::Neox,
9176 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
9177 swa_pattern: None,
9178 swa_dense_first: false,
9179 attn_logit_softcap: None,
9180 final_logit_softcap: None,
9181 embedding_scale: None,
9182 attention_scale: None,
9183 rope_theta_swa: None,
9184 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
9185 best_effort_fields: &["synthetic test config, not a real preset"],
9186 };
9187 let hp = KimiRealHparams {
9188 hidden_dim,
9189 kda_num_heads,
9190 kda_head_dim,
9191 mla_num_heads,
9192 mla_q_lora_rank,
9193 mla_kv_lora_rank,
9194 mla_qk_nope_head_dim,
9195 mla_qk_rope_head_dim,
9196 mla_v_head_dim,
9197 dense_intermediate_dim: dense_intermediate,
9198 moe_hidden_dim: hidden_dim,
9199 moe_intermediate_dim: 4,
9200 n_experts: 1,
9201 num_shared_experts: 0,
9202 };
9203
9204 let prefix = "language_model.model.layers.0";
9208 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
9209 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
9210 name: String,
9211 shape: Vec<usize>,
9212 n: usize| {
9213 tensors.push((name, shape, vec![0.01f32; n]));
9214 };
9215 push(
9216 &mut tensors,
9217 format!("{prefix}.input_layernorm.weight"),
9218 vec![hidden_dim],
9219 hidden_dim,
9220 );
9221 push(
9222 &mut tensors,
9223 format!("{prefix}.post_attention_layernorm.weight"),
9224 vec![hidden_dim],
9225 hidden_dim,
9226 );
9227 push(
9228 &mut tensors,
9229 format!("{prefix}.self_attention_res_norm.weight"),
9230 vec![hidden_dim],
9231 hidden_dim,
9232 );
9233 push(
9234 &mut tensors,
9235 format!("{prefix}.self_attention_res_proj.weight"),
9236 vec![1, hidden_dim],
9237 hidden_dim,
9238 );
9239 push(
9240 &mut tensors,
9241 format!("{prefix}.mlp_res_norm.weight"),
9242 vec![hidden_dim],
9243 hidden_dim,
9244 );
9245 push(
9246 &mut tensors,
9247 format!("{prefix}.mlp_res_proj.weight"),
9248 vec![1, hidden_dim],
9249 hidden_dim,
9250 );
9251 push(
9252 &mut tensors,
9253 format!("{prefix}.self_attn.q_proj.weight"),
9254 vec![kda_proj, hidden_dim],
9255 kda_proj * hidden_dim,
9256 );
9257 push(
9258 &mut tensors,
9259 format!("{prefix}.self_attn.k_proj.weight"),
9260 vec![kda_proj, hidden_dim],
9261 kda_proj * hidden_dim,
9262 );
9263 push(
9264 &mut tensors,
9265 format!("{prefix}.self_attn.v_proj.weight"),
9266 vec![kda_proj, hidden_dim],
9267 kda_proj * hidden_dim,
9268 );
9269 push(
9270 &mut tensors,
9271 format!("{prefix}.self_attn.q_conv1d.weight"),
9272 vec![kda_proj, 1, conv_kernel],
9273 kda_proj * conv_kernel,
9274 );
9275 push(
9276 &mut tensors,
9277 format!("{prefix}.self_attn.k_conv1d.weight"),
9278 vec![kda_proj, 1, conv_kernel],
9279 kda_proj * conv_kernel,
9280 );
9281 push(
9282 &mut tensors,
9283 format!("{prefix}.self_attn.v_conv1d.weight"),
9284 vec![kda_proj, 1, conv_kernel],
9285 kda_proj * conv_kernel,
9286 );
9287 push(
9288 &mut tensors,
9289 format!("{prefix}.self_attn.A_log"),
9290 vec![kda_num_heads],
9291 kda_num_heads,
9292 );
9293 push(
9294 &mut tensors,
9295 format!("{prefix}.self_attn.f_a_proj.weight"),
9296 vec![kda_head_dim, hidden_dim],
9297 kda_head_dim * hidden_dim,
9298 );
9299 push(
9300 &mut tensors,
9301 format!("{prefix}.self_attn.f_b_proj.weight"),
9302 vec![kda_proj, kda_head_dim],
9303 kda_proj * kda_head_dim,
9304 );
9305 push(
9306 &mut tensors,
9307 format!("{prefix}.self_attn.dt_bias"),
9308 vec![kda_proj],
9309 kda_proj,
9310 );
9311 push(
9312 &mut tensors,
9313 format!("{prefix}.self_attn.b_proj.weight"),
9314 vec![kda_num_heads, hidden_dim],
9315 kda_num_heads * hidden_dim,
9316 );
9317 push(
9318 &mut tensors,
9319 format!("{prefix}.self_attn.g_proj.weight"),
9320 vec![kda_proj, hidden_dim],
9321 kda_proj * hidden_dim,
9322 );
9323 push(
9324 &mut tensors,
9325 format!("{prefix}.self_attn.o_norm.weight"),
9326 vec![kda_head_dim],
9327 kda_head_dim,
9328 );
9329 push(
9330 &mut tensors,
9331 format!("{prefix}.self_attn.o_proj.weight"),
9332 vec![hidden_dim, kda_proj],
9333 hidden_dim * kda_proj,
9334 );
9335 push(
9336 &mut tensors,
9337 format!("{prefix}.mlp.gate_proj.weight"),
9338 vec![dense_intermediate, hidden_dim],
9339 dense_intermediate * hidden_dim,
9340 );
9341 push(
9342 &mut tensors,
9343 format!("{prefix}.mlp.up_proj.weight"),
9344 vec![dense_intermediate, hidden_dim],
9345 dense_intermediate * hidden_dim,
9346 );
9347 push(
9348 &mut tensors,
9349 format!("{prefix}.mlp.down_proj.weight"),
9350 vec![hidden_dim, dense_intermediate],
9351 hidden_dim * dense_intermediate,
9352 );
9353 push(
9354 &mut tensors,
9355 "language_model.model.embed_tokens.weight".to_string(),
9356 vec![vocab_size, hidden_dim],
9357 vocab_size * hidden_dim,
9358 );
9359 push(
9360 &mut tensors,
9361 "language_model.lm_head.weight".to_string(),
9362 vec![vocab_size, hidden_dim],
9363 vocab_size * hidden_dim,
9364 );
9365 push(
9366 &mut tensors,
9367 "language_model.model.norm.weight".to_string(),
9368 vec![hidden_dim],
9369 hidden_dim,
9370 );
9371 push(
9372 &mut tensors,
9373 "language_model.model.output_attn_res_norm.weight".to_string(),
9374 vec![hidden_dim],
9375 hidden_dim,
9376 );
9377 push(
9378 &mut tensors,
9379 "language_model.model.output_attn_res_proj.weight".to_string(),
9380 vec![1, hidden_dim],
9381 hidden_dim,
9382 );
9383
9384 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9397 let dir = std::env::temp_dir().join(format!(
9398 "ferrox_server_kimi_e2e_test_{}_{}",
9399 std::process::id(),
9400 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9401 ));
9402 std::fs::create_dir_all(&dir).unwrap();
9403 let shard_bytes = write_safetensors_shard(&tensors);
9404 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9405 let map_entries: Vec<String> = tensors
9406 .iter()
9407 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9408 .collect();
9409 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9410 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9411
9412 use base64::Engine;
9416 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9417 .map(|b| {
9418 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9419 format!("{b64} {b}")
9420 })
9421 .collect();
9422 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9423
9424 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9425 .expect("must load the synthetic Kimi checkpoint end to end");
9426 std::fs::remove_dir_all(&dir).ok();
9427 loaded
9428 }
9429
9430 #[test]
9438 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9439 let loaded = build_synthetic_kimi_loaded();
9440 let state = build_app_state(
9441 StartupModels {
9442 loaded: model::LoadedModel::Kimi(loaded),
9443 embedding: None,
9444 },
9445 None,
9446 None,
9447 None,
9448 false,
9449 None,
9450 Arc::new(health::Detection::ready(health::probe_backends())),
9451 );
9452 let active = state.active().expect("a freshly built state has a model");
9453 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9454 assert!(!active.is_synthetic());
9455
9456 let (_chunks, finish, _usage) = run_generation(
9457 active.generative().unwrap(),
9458 "hi",
9459 &greedy_params(5),
9460 None,
9461 None,
9462 None,
9463 None,
9464 None,
9465 None,
9466 )
9467 .expect("a real Kimi checkpoint must generate without error");
9468 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9469 }
9470
9471 #[test]
9484 fn a_grammar_constrains_the_engine_decode_path() {
9485 let loaded = build_synthetic_kimi_loaded();
9486 let state = build_app_state(
9487 StartupModels {
9488 loaded: model::LoadedModel::Kimi(loaded),
9489 embedding: None,
9490 },
9491 None,
9492 None,
9493 None,
9494 false,
9495 None,
9496 Arc::new(health::Detection::ready(health::probe_backends())),
9497 );
9498 let active = state.active().expect("a freshly built state has a model");
9499
9500 let run = |grammar: Option<&str>| {
9501 let mut params = greedy_params(6);
9502 params.grammar = grammar.map(|src| {
9503 Arc::new(
9504 ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
9505 .expect("test grammar parses"),
9506 )
9507 });
9508 run_generation(
9509 active.generative().unwrap(),
9510 "hi",
9511 ¶ms,
9512 None,
9513 None,
9514 None,
9515 None,
9516 None,
9517 None,
9518 )
9519 };
9520
9521 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
9522 let unconstrained = chunks.concat();
9523 assert!(
9524 unconstrained.chars().any(|c| c != 'a'),
9525 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9526 constrained run below would prove nothing"
9527 );
9528
9529 let (chunks, finish, _) =
9530 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9531 let constrained = chunks.concat();
9532 assert!(
9533 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9534 "the engine decode path served text its grammar forbids ({constrained:?}): \
9535 the constraint was dropped between `generate_engine` and the sampler"
9536 );
9537 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9538 }
9539
9540 #[test]
9550 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9551 let loaded = build_synthetic_kimi_loaded();
9552 let state = build_app_state(
9553 StartupModels {
9554 loaded: model::LoadedModel::Kimi(loaded),
9555 embedding: None,
9556 },
9557 None,
9558 None,
9559 None,
9560 false,
9561 None,
9562 Arc::new(health::Detection::ready(health::probe_backends())),
9563 );
9564
9565 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
9566 let kv_pool_config = generate::KvPoolConfig {
9567 pool: pool.clone(),
9568 queue_wait: Duration::ZERO,
9569 };
9570 let pc = Mutex::new(PrefixCache::new(4));
9571
9572 run_generation(
9573 state
9574 .active()
9575 .expect("a freshly built state has a model")
9576 .generative()
9577 .unwrap(),
9578 "hi",
9579 &greedy_params(5),
9580 Some(&kv_pool_config),
9581 None,
9582 Some(&pc),
9583 None,
9584 None,
9585 None,
9586 )
9587 .expect("a real Kimi checkpoint must generate without error");
9588
9589 assert_eq!(
9590 pool.lock().unwrap().free_blocks(),
9591 4,
9592 "the KV pool must be completely untouched by a Kimi request"
9593 );
9594 let stats = pc.lock().unwrap().stats();
9595 assert_eq!(
9596 stats.hits + stats.misses,
9597 0,
9598 "the prefix cache must never be consulted for a Kimi request"
9599 );
9600 }
9601}