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 samplers: Option<serde_json::Value>,
1437 #[serde(default)]
1447 grammar: Option<String>,
1448}
1449
1450const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1462
1463#[derive(Debug, Clone, Deserialize)]
1465pub(crate) struct ThinkingSwitch {
1466 #[serde(rename = "type")]
1467 pub(crate) kind: String,
1468}
1469
1470const THINKING_KWARG_KEYS: [&str; 4] = [
1474 "enable_thinking",
1475 "thinking",
1476 "thinking_mode",
1477 "reasoning_effort",
1478];
1479
1480const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1484
1485fn default_max_tokens() -> usize {
1486 DEFAULT_CHAT_MAX_TOKENS
1487}
1488
1489impl ChatCompletionRequest {
1490 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1499 Ok(SamplingKnobs {
1500 temperature: self.temperature,
1501 top_p: self.top_p,
1502 min_p: self.min_p,
1503 top_k: self.top_k,
1504 repetition_penalty: self.repetition_penalty,
1505 presence_penalty: self.presence_penalty,
1506 frequency_penalty: self.frequency_penalty,
1507 penalty_last_n: None,
1511 sampler_order: unsupported_sampling::parse_sampler_order(
1512 self.samplers.as_ref(),
1513 "/v1/chat/completions",
1514 )?,
1515 })
1516 }
1517
1518 fn sampling_params(&self) -> Result<SamplingParams, ApiError> {
1519 Ok(self.sampling_knobs()?.resolve())
1520 }
1521
1522 fn stop_sequences(&self) -> Vec<String> {
1523 self.stop
1524 .as_ref()
1525 .map(|s| match s {
1526 StopParam::One(v) => vec![v.clone()],
1527 StopParam::Many(v) => v.clone(),
1528 })
1529 .unwrap_or_default()
1530 }
1531
1532 fn tools_active(&self) -> bool {
1537 !self.tools.is_empty()
1538 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1539 }
1540
1541 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1553 match &self.tool_choice {
1554 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1555 Some(ToolChoice::Specific(value)) => {
1556 let name = value
1560 .get("function")
1561 .and_then(|f| f.get("name"))
1562 .or_else(|| value.get("name"))
1563 .and_then(|n| n.as_str());
1564 match name {
1565 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1566 None => Err(invalid_request(
1567 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1568 function.name",
1569 "tool_choice",
1570 )),
1571 }
1572 }
1573 _ => Ok(None),
1574 }
1575 }
1576
1577 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1579 self.tools
1580 .iter()
1581 .map(|t| tool_grammar::ToolSpec {
1582 name: &t.function.name,
1583 parameters: t.function.parameters.as_ref(),
1584 })
1585 .collect()
1586 }
1587
1588 fn thinking_direction(&self) -> Option<bool> {
1631 if let Some(switch) = &self.thinking {
1632 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1633 "disabled" => Some(false),
1634 "enabled" => Some(true),
1635 _ => None,
1638 };
1639 }
1640 let effort = self.reasoning_effort.as_ref()?;
1641 DISABLE_EFFORTS
1642 .contains(&effort.trim().to_ascii_lowercase().as_str())
1643 .then_some(false)
1644 }
1645
1646 fn resolve_template_kwargs(
1647 &self,
1648 template: &chat_template::PromptTemplate,
1649 ) -> serde_json::Map<String, serde_json::Value> {
1650 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1651 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1655
1656 if !caller_steered {
1657 match self.thinking_direction() {
1658 Some(false) => {
1659 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1660 kwargs.insert(k, v);
1661 }
1662 return kwargs;
1665 }
1666 Some(true) => {
1667 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1668 kwargs.insert(k, v);
1669 }
1670 }
1671 None => {}
1672 }
1673 if let Some(effort) = &self.reasoning_effort {
1674 kwargs
1675 .entry("reasoning_effort".to_string())
1676 .or_insert_with(|| serde_json::json!(effort));
1677 }
1678 }
1679
1680 let offered: Vec<serde_json::Value> = if self.tools_active() {
1681 self.tools.iter().map(chat_template::tool_json).collect()
1682 } else {
1683 Vec::new()
1684 };
1685 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1686 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1687 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1688 kwargs.entry(k).or_insert(v);
1689 }
1690 }
1691 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1692 crate::policy::effort::EffortMapping::Mapped(to) => {
1693 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1694 }
1695 crate::policy::effort::EffortMapping::Dropped => {
1696 tracing::debug!(
1697 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1698 enough, so its own default applies"
1699 );
1700 }
1701 crate::policy::effort::EffortMapping::Unchanged => {}
1702 }
1703 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1704 kwargs
1705 }
1706
1707 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1710 if self.max_tokens == 0 {
1717 return Err(invalid_request(
1718 "max_tokens must be at least 1",
1719 "max_tokens",
1720 ));
1721 }
1722 if let Some(switch) = &self.thinking {
1727 let kind = switch.kind.trim().to_ascii_lowercase();
1728 if kind != "enabled" && kind != "disabled" {
1729 return Err(invalid_request(
1730 "thinking.type must be \"enabled\" or \"disabled\"",
1731 "thinking.type",
1732 ));
1733 }
1734 }
1735 for msg in &self.messages {
1736 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1737 return Err(unsupported_feature(
1738 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1739 ));
1740 }
1741 }
1742 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1743 return Err(unsupported_feature(
1744 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1745 ));
1746 }
1747 if self.n.is_some_and(|n| n > 1) {
1748 return Err(unsupported_feature(
1749 "n > 1 is not implemented (single completion only)",
1750 ));
1751 }
1752 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1753 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1757 let stated_grammar =
1772 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1773 if let Some(forced) = self.forced_tool_choice()? {
1781 if self.tools.is_empty() {
1782 return Err(invalid_request(
1783 "tool_choice forces a tool call, but no tools were offered",
1784 "tool_choice",
1785 ));
1786 }
1787 if let tool_grammar::Forced::Named(name) = forced {
1788 if !self.tools.iter().any(|t| t.function.name == name) {
1789 return Err(invalid_request(
1790 &format!(
1791 "tool_choice names {name:?}, which is not one of the tools offered"
1792 ),
1793 "tool_choice",
1794 ));
1795 }
1796 }
1797 if stated_grammar.is_some() {
1807 return Err(invalid_request(
1808 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1809 are two different constraints on the same generation; send one",
1810 "tool_choice",
1811 ));
1812 }
1813 if self.json_object_mode() {
1814 return Err(invalid_request(
1815 "a forced tool_choice cannot be combined with response_format json_object: \
1816 the tool-call markers are not JSON",
1817 "tool_choice",
1818 ));
1819 }
1820 }
1821 Ok(())
1822 }
1823
1824 fn effective_stop_sequences(&self) -> Vec<String> {
1831 let mut stop = self.stop_sequences();
1832 if self.tools_active() {
1833 stop.push("</tool_call>".to_string());
1834 }
1835 stop
1836 }
1837
1838 fn json_object_mode(&self) -> bool {
1839 self.response_format
1840 .as_ref()
1841 .and_then(|v| v.get("type"))
1842 .and_then(|v| v.as_str())
1843 == Some("json_object")
1844 }
1845
1846 fn generation_params(&self) -> Result<GenerationParams, ApiError> {
1851 Ok(GenerationParams {
1852 max_tokens: self.max_tokens,
1853 sampling: self.sampling_params()?,
1854 seed: self.resolved_seed(),
1855 stop: self.effective_stop_sequences(),
1856 stop_token_ids: Vec::new(),
1860 json_object: self.json_object_mode(),
1861 grammar: grammar_request::for_request(
1862 self.grammar.as_deref(),
1863 self.response_format.as_ref(),
1864 )?,
1865 cancel: None,
1868 ignore_eos: self.ignore_eos.unwrap_or(false),
1869 })
1870 }
1871
1872 fn generation_params_for_template(
1883 &self,
1884 template: &chat_template::PromptTemplate,
1885 served_model: &str,
1886 ) -> Result<GenerationParams, ApiError> {
1887 let mut params = self.generation_params()?;
1888 if let Some(stop) = template.end_of_turn() {
1889 if !params.stop.iter().any(|s| s == stop) {
1890 params.stop.push(stop.to_string());
1891 }
1892 }
1893 if let Some(forced) = self.forced_tool_choice()? {
1894 params.grammar = Some(tool_grammar::build(
1898 forced,
1899 &self.tool_specs(),
1900 policy::parser::ToolCallFormat::infer(served_model),
1901 )?);
1902 }
1903 Ok(params)
1904 }
1905
1906 fn is_cacheable(&self) -> bool {
1915 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1916 }
1917
1918 fn cache_key(&self, prompt: &str, params: &GenerationParams) -> CacheKey {
1934 CacheKey {
1935 model: self.model.clone(),
1936 prompt: prompt.to_string(),
1937 generation: response_cache::generation_key(params),
1938 seed: self.seed,
1939 }
1940 }
1941
1942 fn resolved_seed(&self) -> u64 {
1943 self.seed.unwrap_or_else(|| {
1944 std::time::SystemTime::now()
1945 .duration_since(std::time::UNIX_EPOCH)
1946 .map(|d| d.as_nanos() as u64)
1947 .unwrap_or(0xDEFA017)
1948 })
1949 }
1950}
1951
1952#[derive(Serialize)]
1953struct ChatCompletionChoice {
1954 index: usize,
1955 message: ChatCompletionResponseMessage,
1956 finish_reason: &'static str,
1957}
1958
1959#[derive(Serialize)]
1960struct ChatCompletionResponseMessage {
1961 role: &'static str,
1962 #[serde(skip_serializing_if = "Option::is_none")]
1963 content: Option<String>,
1964 #[serde(skip_serializing_if = "Option::is_none")]
1968 reasoning_content: Option<String>,
1969 #[serde(skip_serializing_if = "Option::is_none")]
1970 tool_calls: Option<Vec<ToolCallOut>>,
1971}
1972
1973#[derive(Serialize, Clone)]
1974struct ToolCallOut {
1975 id: String,
1976 #[serde(rename = "type")]
1977 kind: &'static str,
1978 function: ToolCallFunctionOut,
1979}
1980
1981#[derive(Serialize, Clone)]
1989struct ToolCallDelta {
1990 index: usize,
1991 #[serde(skip_serializing_if = "Option::is_none")]
1992 id: Option<String>,
1993 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1994 kind: Option<&'static str>,
1995 function: ToolCallFunctionDelta,
1996}
1997
1998#[derive(Serialize, Clone, Default)]
1999struct ToolCallFunctionDelta {
2000 #[serde(skip_serializing_if = "Option::is_none")]
2001 name: Option<String>,
2002 #[serde(skip_serializing_if = "Option::is_none")]
2005 arguments: Option<String>,
2006}
2007
2008impl ToolCallDelta {
2009 fn whole(index: usize, name: String, arguments: String) -> Self {
2011 ToolCallDelta {
2012 index,
2013 id: Some(format!("call_{index}")),
2014 kind: Some("function"),
2015 function: ToolCallFunctionDelta {
2016 name: Some(name),
2017 arguments: Some(arguments),
2018 },
2019 }
2020 }
2021
2022 fn opening(index: usize, name: String) -> Self {
2024 ToolCallDelta {
2025 index,
2026 id: Some(format!("call_{index}")),
2027 kind: Some("function"),
2028 function: ToolCallFunctionDelta {
2029 name: Some(name),
2030 arguments: Some(String::new()),
2031 },
2032 }
2033 }
2034
2035 fn arguments(index: usize, fragment: String) -> Self {
2037 ToolCallDelta {
2038 index,
2039 id: None,
2040 kind: None,
2041 function: ToolCallFunctionDelta {
2042 name: None,
2043 arguments: Some(fragment),
2044 },
2045 }
2046 }
2047}
2048
2049#[derive(Serialize, Clone)]
2050struct ToolCallFunctionOut {
2051 name: String,
2052 arguments: String,
2056}
2057
2058#[derive(Serialize)]
2059struct ChatCompletionResponse {
2060 id: String,
2061 request_id: String,
2067 object: &'static str,
2068 model: String,
2069 choices: Vec<ChatCompletionChoice>,
2070 usage: generate::Usage,
2075 ferrox_cache: &'static str,
2086}
2087
2088#[derive(Serialize)]
2089struct ChatCompletionChunkDelta {
2090 #[serde(skip_serializing_if = "Option::is_none")]
2091 role: Option<&'static str>,
2092 #[serde(skip_serializing_if = "Option::is_none")]
2093 content: Option<String>,
2094 #[serde(skip_serializing_if = "Option::is_none")]
2096 reasoning_content: Option<String>,
2097 #[serde(skip_serializing_if = "Option::is_none")]
2098 tool_calls: Option<Vec<ToolCallDelta>>,
2099}
2100
2101#[derive(Serialize)]
2102struct ChatCompletionChunkChoice {
2103 index: usize,
2104 delta: ChatCompletionChunkDelta,
2105 finish_reason: Option<&'static str>,
2106}
2107
2108#[derive(Serialize)]
2109struct ChatCompletionChunk {
2110 id: String,
2111 #[serde(skip_serializing_if = "Option::is_none")]
2118 request_id: Option<String>,
2119 object: &'static str,
2120 model: String,
2121 choices: Vec<ChatCompletionChunkChoice>,
2122 #[serde(skip_serializing_if = "Option::is_none")]
2125 usage: Option<generate::Usage>,
2126}
2127
2128async fn health(State(state): State<Arc<AppState>>) -> Response {
2134 let snapshot = state.detection.snapshot();
2135 let mut capabilities = snapshot.capabilities;
2136 let active = state.active();
2137
2138 capabilities.push(match active.as_deref() {
2141 None => ferrox_api::Capability::unavailable(
2146 ferrox_api::health::capability::REAL_WEIGHTS,
2147 ferrox_api::health::reason::MODEL_NOT_LOADED,
2148 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
2149 ),
2150 Some(active) if active.is_synthetic() => ferrox_api::Capability::unavailable(
2151 ferrox_api::health::capability::REAL_WEIGHTS,
2152 ferrox_api::health::reason::MODEL_NOT_LOADED,
2153 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
2154 checkpoint. Output from this model is noise.",
2155 ),
2156 Some(active) if active.encoder().is_some() => {
2166 let endpoints = active
2167 .encoder()
2168 .map(|e| encoder_endpoints(e))
2169 .unwrap_or_default();
2170 let served_by = match endpoints.is_empty() {
2171 true => "no endpoint in this build serves it".to_string(),
2172 false => format!("served by {}", endpoints.join(" and ")),
2173 };
2174 ferrox_api::Capability::available(
2175 ferrox_api::health::capability::REAL_WEIGHTS,
2176 format!(
2177 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
2178 {served_by}; generation endpoints refuse it.",
2179 active.name(),
2180 ),
2181 )
2182 }
2183 Some(active) => ferrox_api::Capability::available(
2184 ferrox_api::health::capability::REAL_WEIGHTS,
2185 format!("Serving the real checkpoint '{}'.", active.name()),
2186 ),
2187 });
2188 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
2189 ferrox_api::Capability::available(
2190 ferrox_api::health::capability::CONTINUOUS_BATCHING,
2191 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
2192 "On by default on Metal. Concurrent requests share one batched decode worker."
2193 } else {
2194 "Concurrent requests share one batched decode step."
2195 },
2196 )
2197 } else if state.metal_private_decode_gate.is_some() {
2198 ferrox_api::Capability::unavailable(
2199 ferrox_api::health::capability::CONTINUOUS_BATCHING,
2200 ferrox_api::health::reason::DISABLED,
2201 "Off; private Metal decodes serialize (one at a time). Set FERROX_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
2202 )
2203 } else {
2204 ferrox_api::Capability::unavailable(
2205 ferrox_api::health::capability::CONTINUOUS_BATCHING,
2206 ferrox_api::health::reason::DISABLED,
2207 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
2208 )
2209 });
2210
2211 let last_request_ms = state
2212 .last_request_ms
2213 .load(std::sync::atomic::Ordering::Relaxed);
2214 let uptime = state.started_at.elapsed();
2215 let health_state = if active.is_none() {
2223 ferrox_api::HealthState::Unavailable
2224 } else {
2225 snapshot.state
2226 };
2227 let body = ferrox_api::HealthResponse {
2228 state: health_state,
2229 reason: match health_state {
2230 ferrox_api::HealthState::Ready => None,
2231 ferrox_api::HealthState::Unavailable => {
2232 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
2233 }
2234 ferrox_api::HealthState::Detecting => {
2235 Some(ferrox_api::health::reason::DETECTING.to_string())
2236 }
2237 },
2238 detail: match health_state {
2239 ferrox_api::HealthState::Ready => None,
2240 ferrox_api::HealthState::Unavailable => Some(
2241 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
2242 .to_string(),
2243 ),
2244 ferrox_api::HealthState::Detecting => {
2245 Some("Probing available compute backends.".to_string())
2246 }
2247 },
2248 model: active
2249 .as_deref()
2250 .map(|active| ferrox_api::health::ModelSummary {
2251 id: active.name().to_string(),
2252 tokenizer: active.tokenizer_kind().to_string(),
2253 synthetic_weights: active.is_synthetic(),
2254 }),
2255 capabilities,
2256 version: env!("CARGO_PKG_VERSION").to_string(),
2257 pid: std::process::id(),
2258 uptime_seconds: uptime.as_secs_f64(),
2259 server_time_unix_ms: std::time::SystemTime::now()
2260 .duration_since(std::time::UNIX_EPOCH)
2261 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
2262 .unwrap_or(0),
2263 last_request_age_seconds: (last_request_ms > 0)
2264 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
2265 .map(|age| age.max(0.0)),
2266 };
2267
2268 let status =
2269 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
2270 (status, Json(body)).into_response()
2271}
2272
2273async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2274 let Some(active) = state.active() else {
2278 return Json(serde_json::json!({ "object": "list", "data": [] }));
2279 };
2280 let mut model_entry = serde_json::json!({
2281 "id": active.name(),
2282 "object": "model",
2283 "ferrox_synthetic_weights": active.is_synthetic(),
2284 "ferrox_tokenizer": active.tokenizer_kind(),
2285 });
2286 if let Some(encoder) = active.encoder() {
2292 model_entry["ferrox_model_kind"] = serde_json::json!("embedding");
2293 model_entry["ferrox_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
2294 model_entry["ferrox_n_embd"] = serde_json::json!(encoder.n_embd());
2295 model_entry["ferrox_pooling"] = serde_json::json!(encoder.pooling_type().name());
2296 model_entry["ferrox_context_length"] = serde_json::json!(encoder.n_ctx_train());
2297 }
2298 if let Some(model) = active.generative_opt() {
2306 let parser_configured =
2307 crate::policy::parser::ReasoningFormat::infer(active.name()).is_some();
2308 let gears = model.chat_template().think_gears(parser_configured);
2309 if !gears.is_empty() {
2310 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
2311 if let Some(default) = &gears.default {
2312 model_entry["default_reasoning_effort"] = serde_json::json!(default);
2313 }
2314 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
2317 }
2318 }
2319 if let Some(mcp) = &state.mcp {
2320 model_entry["ferrox_mcp"] = mcp.models_metadata();
2321 }
2322 Json(serde_json::json!({
2323 "object": "list",
2324 "data": [model_entry]
2325 }))
2326}
2327
2328async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2342 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
2343 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
2344 let active = state.active();
2345 Json(serde_json::json!({
2346 "model": active.as_ref().map(|a| a.name()),
2347 "state": state
2348 .maintenance
2349 .lock()
2350 .unwrap_or_else(|p| p.into_inner())
2351 .state()
2352 .as_str(),
2353 "uptime_s": state.uptime().as_secs(),
2354 "throughput": {
2355 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2356 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2357 },
2358 "requests": {
2359 "active": state.cancels.live_count(),
2360 "completed": state.stats.recorded_total(),
2361 "p95_ms": state.stats.p95_duration_ms(),
2362 "ttft_mean_ms": state.stats.ttft_mean_ms(),
2363 "prompt_tokens_total": state.stats.tokens_prompt_total(),
2364 "completion_tokens_total": state.stats.tokens_generated_total(),
2365 },
2366 "pools": cache_admin::pool_gauges(&state),
2372 "memory": cache_admin::footprint_json(&state),
2375 }))
2376}
2377
2378#[derive(Deserialize)]
2379struct RequestsQuery {
2380 #[serde(default)]
2381 since: u64,
2382 #[serde(default = "default_requests_limit")]
2383 limit: usize,
2384}
2385
2386fn default_requests_limit() -> usize {
2387 stats::MAX_PAGE
2388}
2389
2390async fn recent_requests(
2398 State(state): State<Arc<AppState>>,
2399 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
2400) -> Json<serde_json::Value> {
2401 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
2402 Json(serde_json::json!({
2403 "requests": rows,
2404 "next_cursor": cursor,
2405 "missed": missed,
2406 "total": state.stats.recorded_total(),
2407 }))
2408}
2409
2410#[derive(Serialize)]
2411struct CombinedCacheStats {
2412 response_cache: response_cache::CacheStats,
2413 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
2415}
2416
2417async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
2418 Json(CombinedCacheStats {
2419 response_cache: lock_cache(&state.response_cache).stats(),
2420 prefix_cache: state
2421 .prefix_cache
2422 .as_ref()
2423 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
2424 })
2425}
2426
2427async fn metrics(State(state): State<Arc<AppState>>) -> Response {
2432 use std::sync::atomic::Ordering;
2433
2434 let cache_stats = lock_cache(&state.response_cache).stats();
2435 let active = state.active();
2436 let requests_total = state.requests_total.load(Ordering::Relaxed);
2437 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2438 let uptime = state.started_at.elapsed().as_secs_f64();
2439
2440 let body = format!(
2441 "# HELP ferrox_requests_total Total chat completion requests received.\n\
2442 # TYPE ferrox_requests_total counter\n\
2443 ferrox_requests_total {requests_total}\n\
2444 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
2445 # TYPE ferrox_request_errors_total counter\n\
2446 ferrox_request_errors_total {errors_total}\n\
2447 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
2448 # TYPE ferrox_cache_hits_total counter\n\
2449 ferrox_cache_hits_total {}\n\
2450 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
2451 # TYPE ferrox_cache_misses_total counter\n\
2452 ferrox_cache_misses_total {}\n\
2453 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
2454 # TYPE ferrox_cache_entries gauge\n\
2455 ferrox_cache_entries {}\n\
2456 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2457 # TYPE ferrox_synthetic_weights gauge\n\
2458 ferrox_synthetic_weights {}\n\
2459 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
2460 # TYPE ferrox_uptime_seconds gauge\n\
2461 ferrox_uptime_seconds {uptime}\n",
2462 cache_stats.hits,
2463 cache_stats.misses,
2464 cache_stats.entries,
2465 active
2469 .as_ref()
2470 .map(|a| a.is_synthetic() as u8)
2471 .unwrap_or(0),
2472 );
2473
2474 let body = match active
2478 .as_ref()
2479 .and_then(|a| a.expert_store_stats())
2480 {
2481 Some(es) => format!(
2482 "{body}\
2483 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
2484 # TYPE ferrox_expert_cache_hits_total counter\n\
2485 ferrox_expert_cache_hits_total {}\n\
2486 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2487 # TYPE ferrox_expert_cache_misses_total counter\n\
2488 ferrox_expert_cache_misses_total {}\n\
2489 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
2490 # TYPE ferrox_expert_cache_evictions_total counter\n\
2491 ferrox_expert_cache_evictions_total {}\n\
2492 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2493 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
2494 ferrox_expert_cache_pass_throughs_total {}\n\
2495 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2496 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
2497 ferrox_expert_cache_bytes_read_total {}\n\
2498 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2499 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
2500 ferrox_expert_cache_resident_bytes {}\n",
2501 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2502 ),
2503 None => body,
2504 };
2505
2506 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2511 Some(batcher) => {
2512 let sched = batcher.stats();
2513 format!(
2514 "{body}\
2515 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2516 # TYPE ferrox_prefill_chunks_total counter\n\
2517 ferrox_prefill_chunks_total {}\n\
2518 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2519 # TYPE ferrox_prefill_tokens_total counter\n\
2520 ferrox_prefill_tokens_total {}\n\
2521 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2522 # TYPE ferrox_decode_steps_total counter\n\
2523 ferrox_decode_steps_total {}\n\
2524 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2525 # TYPE ferrox_scheduler_queue_depth gauge\n\
2526 ferrox_scheduler_queue_depth {}\n\
2527 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2528 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2529 ferrox_scheduler_queue_rejected_total {}\n\
2530 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2531 # TYPE ferrox_kv_blocks_total gauge\n\
2532 ferrox_kv_blocks_total {}\n\
2533 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2534 # TYPE ferrox_kv_blocks_free gauge\n\
2535 ferrox_kv_blocks_free {}\n\
2536 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2537 # TYPE ferrox_kv_block_size gauge\n\
2538 ferrox_kv_block_size {}\n\
2539 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2540 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2541 ferrox_kv_rejected_too_large_total {}\n\
2542 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2543 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2544 ferrox_kv_rejected_context_length_total {}\n\
2545 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2546 # TYPE ferrox_scheduler_aborted_total counter\n\
2547 ferrox_scheduler_aborted_total {}\n",
2548 sched.prefill_chunks,
2549 sched.prefill_tokens,
2550 sched.decode_steps,
2551 sched.queue_depth,
2552 sched.queue_rejected,
2553 sched.kv_blocks_total,
2554 sched.kv_blocks_free,
2555 sched.kv_block_size,
2556 sched.kv_rejected_too_large,
2557 sched.kv_rejected_context_length,
2558 sched.aborted,
2559 )
2560 }
2561 None => body,
2562 };
2563
2564 (
2565 [(
2566 axum::http::header::CONTENT_TYPE,
2567 "text/plain; version=0.0.4",
2568 )],
2569 body,
2570 )
2571 .into_response()
2572}
2573
2574pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2575
2576pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2582 (
2583 StatusCode::BAD_REQUEST,
2584 Json(serde_json::json!({"error": {
2585 "message": message,
2586 "type": "invalid_request_error",
2587 "param": param,
2588 "code": null,
2589 }})),
2590 )
2591}
2592
2593pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2594 (
2595 StatusCode::NOT_IMPLEMENTED,
2596 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2597 )
2598}
2599
2600pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2601 let status = match e {
2602 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2603 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2608 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2614 StatusCode::SERVICE_UNAVAILABLE
2615 }
2616 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2620 };
2621 tracing::warn!("decode error: {e}");
2622 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2623 if let generate::DecodeError::KvBudgetExceeded {
2629 binding,
2630 estimated_bytes,
2631 limit_bytes,
2632 positions,
2633 positions_limit,
2634 ..
2635 } = &e
2636 {
2637 body["error"]["type"] = serde_json::json!("invalid_request_error");
2638 body["error"]["code"] = serde_json::json!(binding);
2639 body["error"]["binding"] = serde_json::json!(binding);
2640 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2641 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2642 body["error"]["positions"] = serde_json::json!(positions);
2643 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2644 }
2645 if let Some(secs) = e.retry_after_secs() {
2649 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2650 }
2651 (status, Json(body))
2652}
2653
2654pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2655 tracing::error!("generation task panicked: {e}");
2656 (
2657 StatusCode::INTERNAL_SERVER_ERROR,
2658 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2659 )
2660}
2661
2662#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2672 model: &Model,
2673 prompt: &str,
2674 params: &GenerationParams,
2675 kv_pool: Option<&generate::KvPoolConfig>,
2676 paged_kv: Option<&generate::PagedKvConfig>,
2677 prefix_cache: Option<&Mutex<PrefixCache>>,
2678 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2679 ceiling: Option<&budget::ContextCeiling>,
2680 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2681 mut emit: impl FnMut(&str),
2682) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2683 let synthetic = model.is_synthetic();
2684 let mut chunks = Vec::new();
2685 let params = &{
2691 let mut resolved = params.clone();
2692 resolved.stop_token_ids =
2693 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
2694 resolved
2695 };
2696 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2697 let _metal_private_guard =
2698 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2699 let (finish, usage) = match model {
2700 Model::Gguf(m) => {
2701 if let Some(batcher) = continuous_batcher {
2702 let mut tokens = m.tokenizer.encode(prompt);
2703 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2704 let (finish, _generated_ids, text, usage) = if synthetic {
2705 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2706 } else {
2707 batcher.generate_streaming(
2708 tokens,
2709 params.clone(),
2710 m.stop_tokens.clone(),
2711 Some(|chunk: &str| {
2712 if !chunk.is_empty() {
2713 chunks.push(chunk.to_string());
2714 emit(chunk);
2715 }
2716 }),
2717 )?
2718 };
2719 if !text.is_empty() && chunks.is_empty() {
2720 chunks.push(text);
2721 }
2722 (finish, usage)
2723 } else {
2724 generate::generate(
2725 &m.decoder,
2726 m.tokenizer.as_ref(),
2727 &m.stop_tokens,
2728 m.bos_id,
2729 prompt,
2730 params,
2731 kv_pool,
2732 paged_kv,
2733 prefix_cache,
2734 ceiling,
2735 |chunk| {
2736 chunks.push(chunk.to_string());
2737 if !synthetic {
2738 emit(chunk);
2739 }
2740 },
2741 )?
2742 }
2743 }
2744 Model::Kimi(m) => generate::generate_engine(
2745 &m.engine,
2746 &m.tokenizer,
2747 &m.stop_tokens,
2748 None,
2749 prompt,
2750 params,
2751 |chunk| {
2752 chunks.push(chunk.to_string());
2753 if !synthetic {
2754 emit(chunk);
2755 }
2756 },
2757 )?,
2758 Model::Mla(m) => generate::generate_engine(
2759 &m.engine,
2760 &m.tokenizer,
2761 &m.stop_tokens,
2762 m.bos_id,
2763 prompt,
2764 params,
2765 |chunk| {
2766 chunks.push(chunk.to_string());
2767 if !synthetic {
2768 emit(chunk);
2769 }
2770 },
2771 )?,
2772 Model::Gemma4(m) => generate::generate_engine(
2773 &m.engine,
2774 &m.tokenizer,
2775 &m.stop_tokens,
2776 m.bos_id,
2777 prompt,
2778 params,
2779 |chunk| {
2780 chunks.push(chunk.to_string());
2781 if !synthetic {
2782 emit(chunk);
2783 }
2784 },
2785 )?,
2786 Model::Glm52(m) => generate::generate_engine(
2787 &m.engine,
2788 &m.tokenizer,
2789 &m.stop_tokens,
2790 m.bos_id,
2791 prompt,
2792 params,
2793 |chunk| {
2794 chunks.push(chunk.to_string());
2795 if !synthetic {
2796 emit(chunk);
2797 }
2798 },
2799 )?,
2800 };
2801
2802 let mut full = chunks.concat();
2803 if synthetic {
2804 full = format!(
2805 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2806 to serve a real model. Decoded ids -> {full:?}]"
2807 );
2808 emit(&full);
2809 } else if used_batcher && !full.is_empty() && chunks.is_empty() {
2810 emit(&full);
2811 }
2812
2813 Ok((finish, usage, full))
2814}
2815
2816#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2821 model: &Model,
2822 prompt: &str,
2823 params: &GenerationParams,
2824 kv_pool: Option<&generate::KvPoolConfig>,
2825 paged_kv: Option<&generate::PagedKvConfig>,
2826 prefix_cache: Option<&Mutex<PrefixCache>>,
2827 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2828 ceiling: Option<&budget::ContextCeiling>,
2829 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2830) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2831 let (finish, usage, full) = run_generation_emit(
2832 model,
2833 prompt,
2834 params,
2835 kv_pool,
2836 paged_kv,
2837 prefix_cache,
2838 continuous_batcher,
2839 ceiling,
2840 metal_private_decode_gate,
2841 |_| {},
2842 )?;
2843 Ok((
2844 if full.is_empty() {
2845 Vec::new()
2846 } else {
2847 vec![full]
2848 },
2849 finish,
2850 usage,
2851 ))
2852}
2853
2854pub(crate) fn prompt_from_messages(
2865 messages: &[ChatMessage],
2866 template: &chat_template::PromptTemplate,
2867 tools: &[ToolDef],
2868 extra: serde_json::Map<String, serde_json::Value>,
2869) -> Result<String, ApiError> {
2870 let rendered = if tools.is_empty() || template.handles_tools() {
2871 template.render(messages, tools, extra)
2872 } else {
2873 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2874 with_preamble.push(ChatMessage {
2875 role: "system".to_string(),
2876 content: Some(MessageContent::Text(tool_preamble(tools))),
2877 tool_calls: None,
2878 tool_call_id: None,
2879 reasoning_content: None,
2880 });
2881 with_preamble.extend_from_slice(messages);
2882 template.render(&with_preamble, &[], extra)
2883 };
2884 rendered.map_err(template_error_response)
2885}
2886
2887fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2892 (
2893 StatusCode::BAD_REQUEST,
2894 Json(serde_json::json!({
2895 "error": {
2896 "message": format!("chat template failed to render: {err}"),
2897 "type": "invalid_request_error",
2898 "param": "messages",
2899 "code": null,
2900 }
2901 })),
2902 )
2903}
2904
2905fn tool_preamble(tools: &[ToolDef]) -> String {
2918 let mut out = String::from(
2919 "You can call tools to help answer the user. To call a tool, respond with \
2920 EXACTLY one line in this format and nothing else:\n\
2921 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2922 object matching that tool's parameters>}}</tool_call>\n\n\
2923 Available tools:\n",
2924 );
2925 for t in tools {
2926 out.push_str(&format!(
2927 "- {}: {}\n parameters (JSON schema): {}\n",
2928 t.function.name,
2929 t.function.description.as_deref().unwrap_or(""),
2930 t.function
2931 .parameters
2932 .as_ref()
2933 .map(|v| v.to_string())
2934 .unwrap_or_else(|| "{}".to_string()),
2935 ));
2936 }
2937 out
2938}
2939
2940fn tool_call_deltas(
2950 events: Vec<crate::policy::parser::ToolCallEvent>,
2951 opened: &std::cell::Cell<usize>,
2952) -> (String, Vec<ToolCallDelta>) {
2953 let mut text = String::new();
2954 let mut deltas = Vec::new();
2955 for event in events {
2956 match event {
2957 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2958 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2959 opened.set(opened.get().max(index + 1));
2960 deltas.push(ToolCallDelta::opening(index, name));
2961 }
2962 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2963 if !fragment.is_empty() {
2964 deltas.push(ToolCallDelta::arguments(index, fragment));
2965 }
2966 }
2967 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2968 }
2969 }
2970 (text, deltas)
2971}
2972
2973fn build_response_message(
2984 text: String,
2985 tools: &[ToolDef],
2986 posture: output::OutputPosture,
2987 base_finish: &'static str,
2988) -> (ChatCompletionResponseMessage, &'static str) {
2989 let parsed = output::parse_output(&text, tools, posture);
2990 let calls: Vec<ToolCallOut> = parsed
2991 .calls
2992 .into_iter()
2993 .enumerate()
2994 .map(|(index, call)| ToolCallOut {
2995 id: format!("call_{index}"),
2996 kind: "function",
2997 function: ToolCallFunctionOut {
2998 name: call.name,
2999 arguments: call.arguments,
3000 },
3001 })
3002 .collect();
3003 if !calls.is_empty() {
3004 return (
3005 ChatCompletionResponseMessage {
3006 role: "assistant",
3007 content: None,
3008 reasoning_content: parsed.reasoning,
3009 tool_calls: Some(calls),
3010 },
3011 "tool_calls",
3012 );
3013 }
3014 (
3015 ChatCompletionResponseMessage {
3016 role: "assistant",
3017 content: Some(parsed.content),
3018 reasoning_content: parsed.reasoning,
3019 tool_calls: None,
3020 },
3021 base_finish,
3022 )
3023}
3024
3025fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
3030 let mut history = match &req.session_id {
3031 Some(id) => state.sessions.extend_and_get(id, &req.messages),
3032 None => req.messages.clone(),
3033 };
3034 if req.json_object_mode() {
3035 inject_json_object_system_hint(&mut history);
3036 }
3037 history
3038}
3039
3040fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
3041 const HINT: &str =
3042 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
3043 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
3044 match &mut sys.content {
3045 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
3046 s.push_str("\n\n");
3047 s.push_str(HINT);
3048 }
3049 None => {
3050 sys.content = Some(MessageContent::Text(HINT.to_string()));
3051 }
3052 _ => {}
3053 }
3054 } else {
3055 messages.insert(
3056 0,
3057 ChatMessage {
3058 role: "system".to_string(),
3059 content: Some(MessageContent::Text(HINT.to_string())),
3060 tool_calls: None,
3061 tool_call_id: None,
3062 reasoning_content: None,
3063 },
3064 );
3065 }
3066}
3067
3068async fn chat_completions(
3069 State(state): State<Arc<AppState>>,
3070 headers: axum::http::HeaderMap,
3071 Json(req): Json<ChatCompletionRequest>,
3072) -> Response {
3073 let attribution = attribution::Attribution::from_headers(&headers);
3074 state
3075 .requests_total
3076 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3077 let started = std::time::Instant::now();
3078
3079 let request_id = ferrox_api::next_request_id();
3083 let stream = req.stream.unwrap_or(false);
3084
3085 let refusal = cache_admin::check_admission(&state)
3091 .err()
3092 .or_else(|| req.validate_supported_fields().err());
3093 if let Some(err) = refusal {
3094 state
3095 .request_errors_total
3096 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3097 let response = err.into_response();
3098 state.record_request(stats::Record {
3099 request_id: &request_id,
3100 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3101 model: state.active_model_name(),
3102 status: response.status().as_u16(),
3103 stream,
3104 duration_ms: started.elapsed().as_millis() as u64,
3105 usage: None,
3106 attribution: &attribution,
3107 });
3108 return response;
3109 }
3110
3111 let response = if stream {
3112 chat_completions_stream(
3113 Arc::clone(&state),
3114 req,
3115 request_id.clone(),
3116 started,
3117 attribution.clone(),
3118 )
3119 .await
3120 .into_response()
3121 } else {
3122 chat_completions_full(
3123 Arc::clone(&state),
3124 req,
3125 request_id.clone(),
3126 started,
3127 attribution.clone(),
3128 )
3129 .await
3130 .into_response()
3131 };
3132
3133 if response.status().is_client_error() || response.status().is_server_error() {
3134 state
3135 .request_errors_total
3136 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3137 state.record_request(stats::Record {
3141 request_id: &request_id,
3142 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3143 model: state.active_model_name(),
3146 status: response.status().as_u16(),
3147 stream,
3148 duration_ms: started.elapsed().as_millis() as u64,
3149 usage: None,
3150 attribution: &attribution,
3151 });
3152 }
3153 state.mark_request_finished();
3154
3155 response
3156}
3157
3158async fn chat_completions_full(
3159 state: Arc<AppState>,
3160 req: ChatCompletionRequest,
3161 request_id: String,
3162 started: std::time::Instant,
3163 attribution: attribution::Attribution,
3164) -> Result<Json<ChatCompletionResponse>, ApiError> {
3165 let tools_active = req.tools_active();
3166 let active = state.require_active()?;
3170 let history = resolve_history(&state, &req);
3171 let template = active.generative()?.chat_template();
3172 let kwargs = req.resolve_template_kwargs(&template);
3173 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
3174 let params = req.generation_params_for_template(&template, active.name())?;
3182 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
3183
3184 let (completion, cache_status) = if let Some(cached) = key
3185 .as_ref()
3186 .and_then(|key| lock_cache(&state.response_cache).get(key))
3187 {
3188 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
3189 (cached, "hit")
3190 } else {
3191 let (chunks, finish, usage) = decode_task::buffered(
3192 decode_task::DecodeHandles::take(&state, &active)?,
3193 prompt.clone(),
3194 params,
3195 )
3196 .await?;
3197
3198 let completion = response_cache::CachedCompletion {
3199 content: chunks.concat(),
3200 finish,
3201 usage,
3202 };
3203 let cache_status = match key {
3209 Some(key) => match completion.clone().cacheable() {
3213 Some(cacheable) => {
3214 tracing::debug!("cache miss for key {}", key.digest());
3215 lock_cache(&state.response_cache).put(key, cacheable);
3216 "miss"
3217 }
3218 None => "skip",
3219 },
3220 None => "skip",
3221 };
3222 (completion, cache_status)
3223 };
3224 let content = completion.content;
3225
3226 if req.json_object_mode() {
3227 json_mode::validate_json_object_output(&content)?;
3228 }
3229
3230 if let Some(id) = &req.session_id {
3234 state.sessions.store_reply(
3235 id,
3236 ChatMessage {
3237 role: "assistant".to_string(),
3238 content: Some(MessageContent::Text(content.clone())),
3239 tool_calls: None,
3240 tool_call_id: None,
3241 reasoning_content: None,
3242 },
3243 );
3244 }
3245
3246 let (message, finish_reason) = build_response_message(
3247 content,
3248 if tools_active { &req.tools } else { &[] },
3249 output::OutputPosture::resolve(active.name(), &prompt),
3250 completion.finish.as_str(),
3251 );
3252
3253 state.record_request(stats::Record {
3254 request_id: &request_id,
3255 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3256 model: Some(active.name().to_string()),
3259 status: 200,
3260 stream: false,
3261 duration_ms: started.elapsed().as_millis() as u64,
3262 usage: Some(&completion.usage),
3263 attribution: &attribution,
3264 });
3265
3266 Ok(Json(ChatCompletionResponse {
3267 id: request_id.clone(),
3268 request_id,
3269 object: "chat.completion",
3270 model: req.model,
3271 choices: vec![ChatCompletionChoice {
3272 index: 0,
3273 message,
3274 finish_reason,
3275 }],
3276 usage: completion.usage,
3277 ferrox_cache: cache_status,
3278 }))
3279}
3280
3281async fn chat_completions_stream(
3282 state: Arc<AppState>,
3283 req: ChatCompletionRequest,
3284 request_id: String,
3285 started: std::time::Instant,
3286 attribution: attribution::Attribution,
3287) -> Result<Response, ApiError> {
3288 let tools_active = req.tools_active();
3290 let active = state.require_active()?;
3294 let history = resolve_history(&state, &req);
3295 let template = active.generative()?.chat_template();
3296 let kwargs = req.resolve_template_kwargs(&template);
3297 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
3298 let model_name = req.model.clone();
3299 let session_id = req.session_id.clone();
3300 let sessions = state.sessions.clone();
3301
3302 let model = Arc::clone(active.generative()?);
3303 let kv_pool = state.kv_pool.clone();
3304 let paged_kv = state.paged_kv.clone();
3305 let prefix_cache = state.prefix_cache.clone();
3306 let batcher = active.batcher.clone();
3307 let ceiling = active.ceiling.clone();
3308 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3309 let mut params = req.generation_params_for_template(&template, active.name())?;
3310 let stats_state = Arc::clone(&state);
3311 let served_model = active.name().to_string();
3314 let posture = output::OutputPosture::resolve(&served_model, &prompt);
3318 let offered_tools: Vec<ToolDef> = if tools_active {
3321 req.tools.clone()
3322 } else {
3323 Vec::new()
3324 };
3325
3326 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3331 params.cancel = Some(cancel_token.clone());
3332
3333 let overlap = true;
3342
3343 let slot = req
3347 .stream_resumable
3348 .unwrap_or(false)
3349 .then(|| state.streams.register(&request_id));
3350 let emitter = resume::Emitter::new(slot);
3351
3352 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3353 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3357 id: request_id.clone(),
3358 request_id: None,
3359 object: "chat.completion.chunk",
3360 model: model_name.clone(),
3361 choices: vec![ChatCompletionChunkChoice {
3362 index: 0,
3363 delta: ChatCompletionChunkDelta {
3364 role: None,
3365 content: None,
3366 reasoning_content: None,
3367 tool_calls: None,
3368 },
3369 finish_reason: None,
3370 }],
3371 usage: None,
3372 });
3373
3374 tokio::task::spawn_blocking(move || {
3375 let _cancel_guard = cancel_guard;
3378 let tx_chunks = tx.clone();
3379 let orphan_timeout = sse::orphan_timeout_from_env();
3384 let mut first = true;
3385 let head_request_id = request_id.clone();
3386 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3396 Rc::new(RefCell::new(posture.reasoning_parser()));
3397 let emit_reasoning = Rc::clone(&stream_reasoning);
3398 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3402 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3403 );
3404 let emit_tools = Rc::clone(&stream_tools);
3405 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3409 let emit_streamed_calls = Rc::clone(&streamed_calls);
3410 let result = run_generation_emit(
3411 &model,
3412 &prompt,
3413 ¶ms,
3414 kv_pool.as_ref(),
3415 paged_kv.as_ref(),
3416 prefix_cache.as_deref(),
3417 batcher.as_ref(),
3418 ceiling.as_deref(),
3419 metal_private_decode_gate.as_deref(),
3420 |chunk| {
3421 if !overlap || chunk.is_empty() {
3422 return;
3423 }
3424 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3425 Some(parser) => {
3426 let delta = parser.push(chunk);
3427 (delta.reasoning, delta.content)
3428 }
3429 None => (String::new(), chunk.to_string()),
3430 };
3431 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3435 Some(parser) => {
3436 let (text, calls) =
3437 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3438 (text, calls)
3439 }
3440 None => (content, Vec::new()),
3441 };
3442 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3445 return;
3446 }
3447 let role = if first { Some("assistant") } else { None };
3448 let request_id = first.then(|| head_request_id.clone());
3449 first = false;
3450 let payload = ChatCompletionChunk {
3451 id: head_request_id.clone(),
3452 request_id,
3453 object: "chat.completion.chunk",
3454 model: model_name.clone(),
3455 choices: vec![ChatCompletionChunkChoice {
3456 index: 0,
3457 delta: ChatCompletionChunkDelta {
3458 role,
3459 content: (!content.is_empty()).then_some(content),
3460 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3461 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3462 },
3463 finish_reason: None,
3464 }],
3465 usage: None,
3466 };
3467 if let Err(why) =
3476 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3477 {
3478 if why == sse::SendFailure::Orphaned {
3479 tracing::warn!(
3480 "SSE stream {head_request_id} accepted nothing for the orphan \
3481 deadline; treating it as abandoned"
3482 );
3483 }
3484 if !emitter.is_resumable() {
3496 cancel_token.cancel();
3497 }
3498 }
3499 },
3500 );
3501
3502 let mut pending_request_id = first.then(|| request_id.clone());
3508
3509 match result {
3510 Ok((finish, usage, full_text)) => {
3511 if let Some(id) = &session_id {
3512 sessions.store_reply(
3513 id,
3514 ChatMessage {
3515 role: "assistant".to_string(),
3516 content: Some(MessageContent::Text(full_text.clone())),
3517 tool_calls: None,
3518 tool_call_id: None,
3519 reasoning_content: None,
3520 },
3521 );
3522 }
3523 let mut streamed_finish: Option<&'static str> = None;
3529 if overlap {
3530 let tail = stream_reasoning
3531 .borrow_mut()
3532 .as_mut()
3533 .map(|parser| parser.flush())
3534 .unwrap_or_default();
3535 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3536 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3537 let mut events = parser.push(&content);
3538 events.extend(parser.finish());
3539 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3540 content = text;
3541 tool_calls = calls;
3542 }
3543 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3544 let payload = ChatCompletionChunk {
3545 id: request_id.clone(),
3546 request_id: pending_request_id.take(),
3547 object: "chat.completion.chunk",
3548 model: model_name.clone(),
3549 choices: vec![ChatCompletionChunkChoice {
3550 index: 0,
3551 delta: ChatCompletionChunkDelta {
3552 role: None,
3553 content: (!content.is_empty()).then_some(content),
3554 reasoning_content: (!tail.reasoning.is_empty())
3555 .then_some(tail.reasoning),
3556 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3557 },
3558 finish_reason: None,
3559 }],
3560 usage: None,
3561 };
3562 let _ =
3563 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3564 }
3565 if streamed_calls.get() > 0 {
3566 streamed_finish = Some("tool_calls");
3567 }
3568 } else {
3569 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3572 let tool_calls: Vec<ToolCallDelta> = parsed
3573 .calls
3574 .iter()
3575 .enumerate()
3576 .map(|(index, call)| {
3577 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3578 })
3579 .collect();
3580 if !tool_calls.is_empty() {
3581 streamed_finish = Some("tool_calls");
3582 }
3583 if !tool_calls.is_empty()
3584 || !parsed.content.is_empty()
3585 || parsed.reasoning.is_some()
3586 {
3587 let payload = ChatCompletionChunk {
3588 id: request_id.clone(),
3589 request_id: pending_request_id.take(),
3590 object: "chat.completion.chunk",
3591 model: model_name.clone(),
3592 choices: vec![ChatCompletionChunkChoice {
3593 index: 0,
3594 delta: ChatCompletionChunkDelta {
3595 role: Some("assistant"),
3596 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3597 .then(|| parsed.content.clone()),
3598 reasoning_content: parsed.reasoning.clone(),
3599 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3600 },
3601 finish_reason: None,
3602 }],
3603 usage: None,
3604 };
3605 let _ =
3606 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3607 }
3608 }
3609 let final_finish_reason = match streamed_finish {
3613 Some(reason) if finish.as_str() != "length" => reason,
3614 _ => finish.as_str(),
3615 };
3616 let final_payload = ChatCompletionChunk {
3617 id: request_id.clone(),
3618 request_id: pending_request_id.take(),
3619 object: "chat.completion.chunk",
3620 model: model_name,
3621 choices: vec![ChatCompletionChunkChoice {
3622 index: 0,
3623 delta: ChatCompletionChunkDelta {
3624 role: None,
3625 content: None,
3626 reasoning_content: None,
3627 tool_calls: None,
3628 },
3629 finish_reason: Some(final_finish_reason),
3630 }],
3631 usage: Some(usage.clone()),
3632 };
3633 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3634 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3635 stats_state.record_request(stats::Record {
3640 request_id: &request_id,
3641 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3642 model: Some(served_model.clone()),
3643 status: 200,
3644 stream: true,
3645 duration_ms: started.elapsed().as_millis() as u64,
3646 usage: Some(&usage),
3647 attribution: &attribution,
3648 });
3649 }
3650 Err(e) => {
3651 tracing::warn!("decode error on streamed request {request_id}: {e}");
3652 stats_state.record_request(stats::Record {
3658 request_id: &request_id,
3659 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3660 model: Some(served_model.clone()),
3661 status: 500,
3662 stream: true,
3663 duration_ms: started.elapsed().as_millis() as u64,
3664 usage: None,
3665 attribution: &attribution,
3666 });
3667 let payload = ChatCompletionChunk {
3668 id: request_id.clone(),
3669 request_id: pending_request_id.take(),
3670 object: "chat.completion.chunk",
3671 model: model_name,
3672 choices: vec![ChatCompletionChunkChoice {
3673 index: 0,
3674 delta: ChatCompletionChunkDelta {
3675 role: Some("assistant"),
3676 content: Some(format!("[error: {e}]")),
3677 reasoning_content: None,
3678 tool_calls: None,
3679 },
3680 finish_reason: Some("stop"),
3681 }],
3682 usage: None,
3683 };
3684 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3685 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3686 }
3687 }
3688 drop(emitter);
3692 });
3693
3694 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3695 Ok((
3718 [(
3719 axum::http::HeaderName::from_static("x-accel-buffering"),
3720 axum::http::HeaderValue::from_static("no"),
3721 )],
3722 Sse::new(stream),
3723 )
3724 .into_response())
3725}
3726
3727fn axum_path(template: &str) -> String {
3744 let mut out = String::with_capacity(template.len());
3745 let mut rest = template;
3746 while let Some(open) = rest.find('{') {
3747 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3748 break;
3749 };
3750 out.push_str(&rest[..open]);
3751 out.push(':');
3752 out.push_str(&rest[open + 1..close]);
3753 rest = &rest[close + 1..];
3754 }
3755 out.push_str(rest);
3756 out
3757}
3758
3759async fn cancel_generation(
3768 State(state): State<Arc<AppState>>,
3769 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3770) -> Response {
3771 let cancelled = state.cancels.cancel(&req.request_id);
3772 let status = if cancelled {
3773 StatusCode::OK
3774 } else {
3775 StatusCode::NOT_FOUND
3776 };
3777 let detail = if cancelled {
3778 "the generation was asked to stop; it ends at its next token".to_string()
3779 } else {
3780 "no generation with that request_id is running -- it has already \
3781 finished, was never issued, or was served by a path that does \
3782 not register for cancellation"
3783 .to_string()
3784 };
3785 (
3786 status,
3787 Json(ferrox_api::CancelGenerationResponse {
3788 request_id: req.request_id,
3789 cancelled,
3790 detail,
3791 }),
3792 )
3793 .into_response()
3794}
3795
3796type Activated = (
3800 Loaded,
3801 Option<serving::batch::ContinuousBatcher>,
3802 Option<Arc<budget::ContextCeiling>>,
3803);
3804
3805fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3823 let mut batcher = serving::batch::BatcherConfig::from_env();
3824 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3825 return batcher;
3828 }
3829 let Some(path) = path else {
3830 return batcher;
3831 };
3832 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1);
3838 let Some((priced, gguf_ctx, source)) = priced else {
3839 return batcher;
3840 };
3841 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3842 tracing::warn!(
3846 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3847 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3848 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3849 admit on a number you choose.",
3850 priced.weights_bytes,
3851 priced.device_budget_bytes,
3852 );
3853 return batcher;
3854 };
3855 tracing::info!("{source}");
3856 tracing::info!("{}", derived.fit);
3857 let adopted = budget::apply_derived(&mut batcher, &derived);
3858 if adopted.max_context {
3859 tracing::info!(
3860 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3861 override with FERROX_CB_MAX_CONTEXT",
3862 derived.max_context
3863 );
3864 }
3865 if adopted.kv_blocks {
3866 tracing::info!(
3867 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3868 derived.kv_blocks,
3869 batcher.kv_block_size
3870 );
3871 }
3872 batcher
3873}
3874
3875pub(crate) fn activate_loaded_model(
3884 loaded: model::LoadedModel,
3885 enable_continuous_batching: bool,
3886 path: Option<&str>,
3887 paged_kv: Option<&generate::PagedKvConfig>,
3888) -> Activated {
3889 match loaded {
3890 model::LoadedModel::Gguf(g) => {
3891 let decoder = Arc::new(g.decoder);
3892 let tokenizer = Arc::new(g.tokenizer);
3893 let config = price_batcher_config(path);
3894 let shape =
3900 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32);
3901 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3902 let batcher = if enable_continuous_batching {
3903 tracing::info!(
3904 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3905 (stop sequences use the same pending-buffer trim as the private generate loop)"
3906 );
3907 let tok = Arc::clone(&tokenizer);
3908 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
3909 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3910 Arc::clone(&decoder),
3911 decode,
3912 config,
3913 Arc::clone(&ceiling),
3914 paged_kv.cloned(),
3915 ))
3916 } else {
3917 None
3918 };
3919 (
3920 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3921 decoder,
3922 tokenizer,
3923 stop_tokens: g.stop_tokens,
3924 bos_id: g.bos_id,
3925 is_synthetic: g.is_synthetic,
3926 chat_template: g.chat_template,
3927 }))),
3928 batcher,
3929 Some(ceiling),
3930 )
3931 }
3932 model::LoadedModel::Kimi(k) => (
3933 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3934 engine: k.engine,
3935 tokenizer: k.tokenizer,
3936 stop_tokens: k.stop_tokens,
3937 chat_template: k.chat_template,
3938 }))),
3939 None,
3940 None,
3941 ),
3942 model::LoadedModel::Mla(m) => (
3943 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3944 engine: m.engine,
3945 tokenizer: m.tokenizer,
3946 stop_tokens: m.stop_tokens,
3947 bos_id: m.bos_id,
3948 name: m.name,
3949 chat_template: m.chat_template,
3950 }))),
3951 None,
3952 None,
3953 ),
3954 model::LoadedModel::Gemma4(m) => (
3955 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3956 engine: m.engine,
3957 tokenizer: m.tokenizer,
3958 stop_tokens: m.stop_tokens,
3959 bos_id: m.bos_id,
3960 name: m.name,
3961 chat_template: m.chat_template,
3962 }))),
3963 None,
3964 None,
3965 ),
3966 model::LoadedModel::Glm52(g) => (
3967 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3968 engine: g.engine,
3969 tokenizer: g.tokenizer,
3970 stop_tokens: g.stop_tokens,
3971 bos_id: g.bos_id,
3972 name: g.name,
3973 chat_template: g.chat_template,
3974 }))),
3975 None,
3976 None,
3977 ),
3978 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3983 }
3984}
3985
3986struct StartupModels {
3993 loaded: model::LoadedModel,
3994 embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
3995}
3996
3997fn continuous_batching_env() -> Option<bool> {
3998 match std::env::var("FERROX_CONTINUOUS_BATCHING")
3999 .ok()
4000 .map(|v| v.trim().to_ascii_lowercase())
4001 .as_deref()
4002 {
4003 None => None,
4004 Some("1" | "true" | "yes" | "on") => Some(true),
4005 Some("0" | "false" | "no" | "off") => Some(false),
4006 _ => None,
4007 }
4008}
4009
4010fn metal_private_decode_active() -> bool {
4011 #[cfg(feature = "metal")]
4012 {
4013 BUILT_WITH_METAL
4014 && ferrox_metal::attn::metal_attn_enabled()
4015 && std::env::var("FERROX_METAL").ok().as_deref() != Some("0")
4016 }
4017 #[cfg(not(feature = "metal"))]
4018 {
4019 false
4020 }
4021}
4022
4023fn continuous_batching_compatible(
4024 loaded: &model::LoadedModel,
4025 kv_pool: &Option<generate::KvPoolConfig>,
4026 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
4027 paged_kv: &Option<generate::PagedKvConfig>,
4028) -> bool {
4029 matches!(loaded, model::LoadedModel::Gguf(_))
4030 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
4031}
4032
4033fn resolve_continuous_batching_enabled(
4034 loaded: &model::LoadedModel,
4035 kv_pool: &Option<generate::KvPoolConfig>,
4036 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
4037 paged_kv: &Option<generate::PagedKvConfig>,
4038) -> bool {
4039 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
4040 return false;
4041 }
4042 match continuous_batching_env() {
4043 Some(true) => true,
4044 Some(false) => false,
4045 None => metal_private_decode_active(),
4046 }
4047}
4048
4049fn acquire_metal_private_decode_gate(
4050 gate: Option<&std::sync::Mutex<()>>,
4051 used_batcher: bool,
4052) -> Option<std::sync::MutexGuard<'_, ()>> {
4053 if used_batcher {
4054 None
4055 } else {
4056 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
4057 }
4058}
4059
4060fn build_app_state(
4061 models: StartupModels,
4062 kv_pool: Option<generate::KvPoolConfig>,
4063 paged_kv: Option<generate::PagedKvConfig>,
4064 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
4065 enable_continuous_batching: bool,
4066 mcp: Option<mcp::LoadedMcpConfig>,
4067 detection: Arc<health::Detection>,
4068) -> AppState {
4069 let StartupModels { loaded, embedding } = models;
4070 let (loaded, batcher, ceiling) = activate_loaded_model(
4071 loaded,
4072 enable_continuous_batching,
4073 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
4074 paged_kv.as_ref(),
4075 );
4076 let id = startup_model_id();
4082 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
4083 {
4084 None
4085 } else {
4086 tracing::info!(
4087 "Metal private-loop decode will serialize concurrent requests until \
4088 continuous batching is enabled (FERROX_CONTINUOUS_BATCHING=1 or --cont-batching)"
4089 );
4090 Some(Arc::new(std::sync::Mutex::new(())))
4091 };
4092 AppState {
4093 embedding,
4094 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4095 id,
4096 loaded,
4097 batcher,
4098 ceiling,
4099 }))),
4100 paged_kv,
4101 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4102 tasks: Arc::new(tasks::TaskRegistry::new()),
4103 cancels: Arc::new(cancel::CancelRegistry::new()),
4104 stats: stats::Stats::new(),
4105 streams: resume::StreamRegistry::new(),
4106 model_dir: admin::model_dirs().into_iter().next(),
4107 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
4108 kv_pool,
4109 prefix_cache,
4110 sessions: session::SessionStore::new(),
4111 requests_total: std::sync::atomic::AtomicU64::new(0),
4112 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4113 started_at: std::time::Instant::now(),
4114 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4115 detection,
4116 mcp,
4117 continuous_batching_enabled: enable_continuous_batching,
4118 metal_private_decode_gate,
4119 loading_model: Mutex::new(None),
4120 last_load_error: Mutex::new(None),
4121 serving: Mutex::new(crate::stats::ServingStats::default()),
4122 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4123 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4124 started_unix: unix_now(),
4125 }
4126}
4127
4128fn load_embedding_model() -> anyhow::Result<Option<Arc<ferrox_models::EmbeddingModel>>> {
4136 let Ok(path) = std::env::var("FERROX_EMBEDDING_MODEL_PATH") else {
4137 return Ok(None);
4138 };
4139 let model = ferrox_models::EmbeddingModel::from_gguf_path(&path)
4140 .map_err(|e| anyhow::anyhow!("FERROX_EMBEDDING_MODEL_PATH={path}: {e}"))?;
4141 tracing::info!(
4142 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
4143 model.name(),
4144 model.architecture(),
4145 model.n_embd(),
4146 model.pooling_type().name(),
4147 model.n_ctx_train(),
4148 );
4149 Ok(Some(Arc::new(model)))
4150}
4151
4152fn unix_now() -> u64 {
4157 std::time::SystemTime::now()
4158 .duration_since(std::time::UNIX_EPOCH)
4159 .map(|d| d.as_secs())
4160 .unwrap_or(0)
4161}
4162
4163fn startup_model_id() -> Option<String> {
4168 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
4169 let configured = std::fs::canonicalize(&configured).ok()?;
4170 admin::discover(&admin::model_dirs())
4171 .into_iter()
4172 .find(|d| {
4173 std::fs::canonicalize(&d.path)
4174 .map(|p| p == configured)
4175 .unwrap_or(false)
4176 })
4177 .map(|d| d.id)
4178}
4179
4180fn init_cpu_pool() {
4189 match ferrox_core::threads::init_cpu_pool() {
4190 Some(n) => eprintln!(
4191 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
4192 ferrox_core::threads::perf_core_count()
4193 ),
4194 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
4195 }
4196}
4197
4198fn announce_ready(addr: SocketAddr, scheme: &str) {
4211 use std::io::Write;
4212 let ready =
4213 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
4214 let mut stdout = std::io::stdout().lock();
4215 let _ = writeln!(stdout, "{}", ready.to_line());
4216 let _ = stdout.flush();
4217}
4218
4219async fn shutdown_signal(exit_on_stdin_close: bool) {
4231 if !exit_on_stdin_close {
4232 std::future::pending::<()>().await;
4233 return;
4234 }
4235 let _ = tokio::task::spawn_blocking(|| {
4236 use std::io::Read;
4237 let mut sink = [0u8; 256];
4238 let mut stdin = std::io::stdin().lock();
4239 loop {
4240 match stdin.read(&mut sink) {
4241 Ok(0) => break,
4243 Ok(_) => continue,
4245 Err(e) => {
4246 tracing::warn!("stdin read failed ({e}); treating it as closed");
4247 break;
4248 }
4249 }
4250 }
4251 })
4252 .await;
4253 tracing::info!("stdin closed; shutting down");
4254}
4255
4256fn tokio_worker_threads() -> usize {
4261 std::env::var("FERROX_TOKIO_WORKERS")
4262 .ok()
4263 .and_then(|v| v.trim().parse::<usize>().ok())
4264 .filter(|n| *n > 0)
4265 .unwrap_or(2)
4266}
4267
4268fn install_ring_crypto_provider() {
4285 let _ = rustls::crypto::ring::default_provider().install_default();
4286}
4287
4288pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4294 if args.list_devices {
4295 print_available_devices();
4296 return Ok(());
4297 }
4298 apply_cli_overrides(&args)?;
4299
4300 let _instance = {
4304 use ferrox_core::instance::{register, InstancePolicy};
4305 let policy = if args.allow_multiple_instances {
4306 InstancePolicy::Multi
4307 } else {
4308 InstancePolicy::from_env_or(InstancePolicy::Single)
4309 };
4310 let model = std::env::var("FERROX_MODEL_PATH").ok();
4311 register(
4312 "server",
4313 model.as_deref(),
4314 ferrox_core::instance::current_backend(),
4315 policy,
4316 )
4317 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4318 };
4319
4320 let journal = journal::Journal::from_env();
4321 eprintln!(
4322 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
4323 journal.path()
4324 );
4325 journal.append(&journal::Record::session_start(
4326 env!("CARGO_PKG_VERSION"),
4327 std::process::id(),
4328 ));
4329 journal::install_panic_hook(journal.clone());
4330
4331 let mcp_config_path = args.mcp_config.clone();
4332 let exit_on_stdin_close = args.exit_on_stdin_close
4333 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
4334 .map(|v| v == "1")
4335 .unwrap_or(false);
4336
4337 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
4341 init_cpu_pool();
4342
4343 let runtime = tokio::runtime::Builder::new_multi_thread()
4344 .worker_threads(tokio_worker_threads())
4345 .enable_all()
4346 .build()?;
4347 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4348
4349 let reason = match &result {
4350 Ok(()) => "normal".to_string(),
4351 Err(e) => e.to_string(),
4352 };
4353 journal.append(&journal::Record::session_exit(reason));
4354
4355 runtime.shutdown_background();
4360
4361 result
4362}
4363
4364async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4365 let _ = tracing_subscriber::fmt::try_init();
4371
4372 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4382 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
4383 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
4384 .map(|v| v == "1")
4385 .unwrap_or(false);
4386 if let Err(msg) =
4387 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4388 {
4389 anyhow::bail!(msg);
4390 }
4391
4392 let embedding_model = load_embedding_model()?;
4399
4400 let mut loaded = model::load()?;
4401 match &loaded {
4402 model::LoadedModel::Gguf(g) => tracing::info!(
4403 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4404 g.decoder.config.name,
4405 g.is_synthetic,
4406 g.tokenizer.kind()
4407 ),
4408 model::LoadedModel::Kimi(k) => tracing::info!(
4409 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4410 k.tokenizer.vocab_size()
4411 ),
4412 model::LoadedModel::Mla(m) => tracing::info!(
4413 "loaded MLA GGUF '{}' (tokenizer={})",
4414 m.name,
4415 m.tokenizer.kind()
4416 ),
4417 model::LoadedModel::Gemma4(m) => tracing::info!(
4418 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4419 m.name,
4420 m.tokenizer.kind()
4421 ),
4422 model::LoadedModel::Glm52(g) => tracing::info!(
4423 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4424 g.name,
4425 g.tokenizer.kind()
4426 ),
4427 model::LoadedModel::Encoder(_) => {}
4430 }
4431 let metal_default_moe_budget = {
4439 #[cfg(feature = "metal")]
4440 {
4441 ferrox_core::metal_dense_enabled()
4442 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
4443 }
4444 #[cfg(not(feature = "metal"))]
4445 {
4446 false
4447 }
4448 };
4449 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
4450 let budget: u64 = budget_str
4451 .parse()
4452 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4453 match &mut loaded {
4454 model::LoadedModel::Gguf(g) => {
4455 tracing::info!(
4456 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4457 (CUDA and/or Metal matvecs when built with the matching feature)"
4458 );
4459 g.decoder.gpu_vram_budget_bytes = Some(budget);
4460 }
4461 model::LoadedModel::Kimi(_) => {
4462 tracing::warn!(
4463 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4464 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4465 );
4466 }
4467 model::LoadedModel::Mla(_) => {
4468 tracing::warn!(
4469 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4470 FFN path only today; ignoring expert VRAM budget"
4471 );
4472 }
4473 model::LoadedModel::Gemma4(_) => {
4474 tracing::warn!(
4475 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4476 ignoring expert VRAM budget"
4477 );
4478 }
4479 model::LoadedModel::Glm52(_) => {
4480 tracing::warn!(
4481 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4482 GPU expert placement not wired yet; ignoring"
4483 );
4484 }
4485 model::LoadedModel::Encoder(_) => {
4486 tracing::warn!(
4487 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4488 it has no routed experts to place; ignoring"
4489 );
4490 }
4491 }
4492 } else if metal_default_moe_budget {
4493 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4497 if let model::LoadedModel::Gguf(g) = &mut loaded {
4498 tracing::info!(
4499 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4500 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4501 );
4502 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4503 }
4504 }
4505 #[cfg(feature = "cuda")]
4506 {
4507 if ferrox_core::cuda_dense_enabled() {
4508 tracing::info!(
4509 "CUDA dense matvec enabled for WeightMatrix::apply \
4510 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4511 );
4512 } else {
4513 tracing::info!(
4514 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
4515 );
4516 }
4517 }
4518 #[cfg(feature = "metal")]
4519 {
4520 if ferrox_core::metal_dense_enabled() {
4521 tracing::info!(
4522 "Metal dense matvec enabled for WeightMatrix::apply \
4523 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4524 );
4525 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
4526 Some("1") | Some("true") | Some("on") | Some("attn") => {
4527 tracing::info!(
4528 "Metal fused attention requested (FERROX_METAL_ATTN): \
4529 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4530 );
4531 }
4532 _ => {}
4533 }
4534 tracing::info!(
4535 "Metal greedy GPU argmax: temperature<=0 folds \
4536 final_norm+lm_head+argmax into the dense stack"
4537 );
4538 } else {
4539 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
4540 }
4541 }
4542 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
4553 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
4554 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
4555 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4556 panic!(
4557 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
4558 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
4559 );
4560 }
4561 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4562 (Ok(blocks), Ok(block_size), Err(_)) => {
4563 let total_blocks: usize = blocks
4564 .parse()
4565 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
4566 let block_size: usize = block_size
4567 .parse()
4568 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4569 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4575 .ok()
4576 .map(|v| {
4577 v.parse()
4578 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4579 })
4580 .unwrap_or(0);
4581 tracing::info!(
4582 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4583 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4584 );
4585 Some(generate::KvPoolConfig {
4586 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4587 queue_wait: Duration::from_millis(queue_wait_ms),
4588 })
4589 }
4590 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4591 let block_size: usize = block_size
4592 .parse()
4593 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4594 let budget: u64 = byte_budget
4595 .parse()
4596 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
4597 let cfg = match &loaded {
4598 model::LoadedModel::Gguf(g) => &g.decoder.config,
4599 model::LoadedModel::Kimi(_)
4600 | model::LoadedModel::Mla(_)
4601 | model::LoadedModel::Gemma4(_)
4602 | model::LoadedModel::Glm52(_)
4603 | model::LoadedModel::Encoder(_) => {
4604 panic!(
4605 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
4606 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4607 );
4608 }
4609 };
4610 let bytes_per_block = block_size
4611 * cfg.n_layers
4612 * cfg.n_kv_heads
4613 * cfg.head_dim
4614 * 2
4615 * std::mem::size_of::<f32>();
4616 assert!(
4617 bytes_per_block > 0,
4618 "derived KV block byte size must be positive (check model config and block size)"
4619 );
4620 let total_blocks = (budget as usize / bytes_per_block).max(1);
4621 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4622 .ok()
4623 .map(|v| {
4624 v.parse()
4625 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4626 })
4627 .unwrap_or(0);
4628 tracing::info!(
4629 "KV cache block pool enabled from byte budget: {budget} bytes / \
4630 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4631 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4632 cfg.n_layers
4633 );
4634 Some(generate::KvPoolConfig {
4635 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4636 queue_wait: Duration::from_millis(queue_wait_ms),
4637 })
4638 }
4639 (Err(_), Err(_), Err(_)) => None,
4640 (Err(_), Ok(_), Err(_)) => panic!(
4641 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4642 (or unset all three to disable KV cache pooling)"
4643 ),
4644 (Ok(_), Ok(_), Ok(_)) => {
4645 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4646 }
4647 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4648 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4649 set together (or neither, to disable KV cache pooling)"
4650 ),
4651 };
4652 let paged_kv = match (
4658 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4659 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4660 ) {
4661 (Ok(blocks), Ok(block_size)) => {
4662 assert!(
4663 kv_pool.is_none(),
4664 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4665 mutually exclusive: both bound the same KV memory, by different means. \
4666 Set one."
4667 );
4668 let blocks_per_layer: usize = blocks
4682 .parse()
4683 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4684 let block_size: usize = block_size
4685 .parse()
4686 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4687 let gguf = match &loaded {
4688 model::LoadedModel::Gguf(g) => g,
4689 _ => panic!(
4690 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4691 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4692 ),
4693 };
4694 let cfg = &gguf.decoder.config;
4695 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4696 .ok()
4697 .map(|v| {
4698 v.parse()
4699 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4700 })
4701 .unwrap_or(0);
4702 tracing::info!(
4703 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4704 layer across {} layers, shared by all concurrent requests, \
4705 {queue_wait_ms}ms admission queue wait",
4706 cfg.n_layers
4707 );
4708 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4713 block_size,
4714 ))));
4715 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4724 crate::policy::parser::ToolCallFormat::infer(
4725 &std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
4726 )
4727 .opener(),
4728 |text| {
4729 gguf.tokenizer
4730 .encode(text)
4731 .into_iter()
4732 .map(|t| t as u32)
4733 .collect()
4734 },
4735 );
4736 if let Some(id) = anchor_token {
4737 tracing::info!(
4738 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4739 window stops short of where its next turn rejoins"
4740 );
4741 }
4742 let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
4743 .ok()
4744 .map(|v| {
4745 v.parse()
4746 .expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4747 })
4748 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4749 if let Some(window) = cfg.uniform_sliding_window() {
4750 tracing::info!(
4751 "Paged KV window slide enabled: every layer slides by {window} every \
4752 {slide_interval} decode steps, so a request holds its prompt and a \
4753 window rather than its whole context"
4754 );
4755 } else if cfg.kv_block_window().is_some() {
4756 tracing::info!(
4757 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4758 and a page group holds one block in every layer"
4759 );
4760 }
4761 Some(generate::PagedKvConfig {
4762 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4763 cfg.n_layers,
4764 block_size,
4765 blocks_per_layer,
4766 cfg.n_kv_heads,
4767 cfg.head_dim,
4768 )),
4769 queue_wait: Duration::from_millis(queue_wait_ms),
4770 radix,
4771 anchor_token,
4772 slide_interval,
4773 })
4774 }
4775 (Err(_), Err(_)) => None,
4776 _ => panic!(
4777 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4778 (or neither, to disable paged KV)"
4779 ),
4780 };
4781 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4788 let max_entries: usize = v
4789 .parse()
4790 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4791 if kv_pool.is_some() {
4792 tracing::warn!(
4793 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4794 caching will never be consulted while a KV pool is configured"
4795 );
4796 }
4797 assert!(
4804 paged_kv.is_none(),
4805 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4806 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4807 produce, so the cache could never hit. Set one."
4808 );
4809 tracing::info!(
4810 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4811 all requests"
4812 );
4813 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4814 });
4815 if matches!(
4816 loaded,
4817 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4818 ) && (kv_pool.is_some() || prefix_cache.is_some())
4819 {
4820 tracing::warn!(
4821 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4822 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4823 ferrox_models::engine's module docs"
4824 );
4825 }
4826 let enable_cb =
4827 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4828 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4829 tracing::info!(
4830 "continuous batching enabled by default on Metal for safe parallel serving \
4831 (set FERROX_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4832 );
4833 }
4834 if continuous_batching_env() == Some(true)
4835 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4836 && (kv_pool.is_some() || prefix_cache.is_some())
4837 {
4838 tracing::warn!(
4839 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4840 (those modes keep the private generate path)"
4841 );
4842 }
4843 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4844 if let Ok(chunk) = n.parse::<usize>() {
4845 if chunk > 0 {
4846 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4847 }
4848 }
4849 }
4850 if matches!(
4851 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4852 Some("1")
4853 ) {
4854 tracing::warn!(
4855 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4856 (minimal spill; full layer offload still planned)"
4857 );
4858 }
4859
4860 let mcp = match mcp_config_path {
4861 Some(path) => {
4862 let loaded = mcp::load_mcp_config(&path)?;
4863 tracing::info!(
4864 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4865 loaded.path,
4866 loaded.servers.len()
4867 );
4868 Some(loaded)
4869 }
4870 None => None,
4871 };
4872
4873 let detection = health::Detection::spawn();
4877
4878 let state = Arc::new(build_app_state(
4879 StartupModels {
4880 loaded,
4881 embedding: embedding_model,
4882 },
4883 kv_pool,
4884 paged_kv,
4885 prefix_cache,
4886 enable_cb,
4887 mcp,
4888 detection,
4889 ));
4890
4891 use ferrox_api::routes;
4895
4896 let public = Router::new().route(routes::HEALTH, get(health));
4901
4902 let mut protected = Router::new()
4903 .route(routes::V1_MODELS, get(list_models))
4904 .route(routes::V1_RESPONSES, post(responses::responses))
4908 .route(
4909 &axum_path(routes::V1_RESPONSE),
4910 get(responses::responses_get),
4911 )
4912 .route(
4913 &axum_path(routes::V1_RESPONSE_CANCEL),
4914 post(responses::responses_cancel),
4915 )
4916 .route(routes::V1_STATS, get(serving_stats))
4917 .route(routes::V1_REQUESTS, get(recent_requests))
4918 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
4919 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
4920 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
4921 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
4922 .route(routes::V1_CANCEL, post(cancel_generation))
4926 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
4931 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
4932 .route(routes::V1_MESSAGES, post(anthropic::messages))
4933 .route(
4934 routes::V1_MESSAGES_COUNT_TOKENS,
4935 post(anthropic::count_tokens),
4936 )
4937 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
4938 .route(routes::COMPLETION, post(completion::completion))
4943 .route(routes::COMPLETIONS, post(completion::completion))
4944 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
4945 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
4946 .route(routes::TOKENIZE, post(openai_extra::tokenize))
4952 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
4953 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
4954 .route(routes::V1_RERANK, post(rerank::rerank))
4958 .route(routes::RERANK, post(rerank::rerank))
4959 .route(routes::CACHE_STATS, get(cache_stats))
4960 .route(routes::METRICS, get(metrics))
4961 .route(routes::ADMIN_MODELS, get(admin::models))
4966 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
4967 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
4968 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
4969 .route(routes::ADMIN_TASKS, get(admin::tasks))
4970 .route(&admin::cancel_route(), post(admin::cancel_task))
4971 .route(routes::ADMIN_STATS, get(admin::stats))
4972 .merge(conversations::router());
4976
4977 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4982 tracing::info!("API key auth enabled");
4983 let auth = limits::AuthConfig {
4984 api_key: Arc::new(key),
4985 };
4986 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4987 auth,
4988 limits::require_api_key,
4989 ));
4990 }
4991 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4992 let rpm: u32 = rpm
4993 .parse()
4994 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4995 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4996 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4997 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4998 limiter,
4999 limits::rate_limit,
5000 ));
5001 }
5002 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
5010 let origins = security::parse_cors_origins(&spec)
5011 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
5012 tracing::info!(
5013 "CORS enabled: {} allow-listed origin(s) ({})",
5014 origins.len(),
5015 spec
5016 );
5017 let cors = tower_http::cors::CorsLayer::new()
5018 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
5019 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
5020 .allow_headers([
5021 axum::http::header::CONTENT_TYPE,
5022 axum::http::header::AUTHORIZATION,
5023 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
5029 axum::http::HeaderName::from_static("last-event-id"),
5033 ]);
5034 protected = protected.route_layer(cors);
5035 }
5036
5037 let app = public
5041 .merge(protected)
5042 .layer(axum::middleware::from_fn(limits::retry_after))
5043 .with_state(state);
5044
5045 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
5051 install_ring_crypto_provider();
5052 match tls_paths {
5057 Some(paths) => {
5058 let config =
5059 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
5060 .await
5061 .map_err(|e| {
5062 anyhow::anyhow!(
5063 "failed to load TLS cert/key ({:?}, {:?}): {e}",
5064 paths.cert,
5065 paths.key
5066 )
5067 })?;
5068 let socket_addr: std::net::SocketAddr = addr
5069 .parse()
5070 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
5071 let listener = std::net::TcpListener::bind(socket_addr)?;
5072 listener.set_nonblocking(true)?;
5080 let bound = listener.local_addr()?;
5081 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
5082 announce_ready(bound, "https");
5083
5084 let handle = axum_server::Handle::new();
5085 let shutdown_handle = handle.clone();
5086 tokio::spawn(async move {
5087 shutdown_signal(exit_on_stdin_close).await;
5088 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
5089 });
5090 axum_server::from_tcp_rustls(listener, config)?
5091 .handle(handle)
5092 .serve(app.into_make_service())
5093 .await?;
5094 }
5095 None => {
5096 let listener = tokio::net::TcpListener::bind(&addr).await?;
5097 let bound = listener.local_addr()?;
5098 tracing::info!("ferrox-server listening on {bound}");
5099 announce_ready(bound, "http");
5100 axum::serve(listener, app)
5101 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
5102 .await?;
5103 }
5104 }
5105 Ok(())
5106}
5107
5108#[cfg(test)]
5109mod tests {
5110 use super::*;
5111 use ferrox_models::config::test_dense_fixture;
5112
5113 #[test]
5114 fn parses_llama_server_style_options() {
5115 let argv = [
5116 "ferrox-server",
5117 "-m",
5118 "model.gguf",
5119 "--host",
5120 "::1",
5121 "--port",
5122 "9000",
5123 "-t",
5124 "4",
5125 "-dev",
5126 "Metal",
5127 "-ngl",
5128 "all",
5129 ]
5130 .into_iter()
5131 .map(String::from)
5132 .collect();
5133 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
5134
5135 assert_eq!(args.model.as_deref(), Some("model.gguf"));
5136 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
5137 assert_eq!(args.port, Some(9000));
5138 assert_eq!(args.threads, Some(4));
5139 assert_eq!(args.device, Some(OffloadDevice::Metal));
5140 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
5141 assert_eq!(
5142 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
5143 Some("[::1]:9000")
5144 );
5145 }
5146
5147 #[test]
5148 fn port_zero_survives_argument_parsing_as_a_real_request() {
5149 let argv = ["ferrox-server", "--port", "0"]
5154 .into_iter()
5155 .map(String::from)
5156 .collect();
5157 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
5158 assert_eq!(args.port, Some(0));
5159 assert_eq!(
5160 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
5161 Some("127.0.0.1:0")
5162 );
5163 }
5164
5165 #[test]
5166 fn parallel_flag_parses_and_rewrites_np() {
5167 let argv = ["ferrox-server", "-np", "4"]
5168 .into_iter()
5169 .map(String::from)
5170 .collect();
5171 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
5172 assert_eq!(args.parallel, Some(4));
5173 }
5174
5175 #[test]
5176 fn stdin_close_exit_is_opt_in() {
5177 let args =
5180 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
5181 assert!(!args.exit_on_stdin_close);
5182 let args = ServerArgs::try_parse_from(
5183 ["ferrox-server", "--exit-on-stdin-close"]
5184 .into_iter()
5185 .map(String::from),
5186 )
5187 .unwrap();
5188 assert!(args.exit_on_stdin_close);
5189 }
5190
5191 #[test]
5192 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
5193 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
5194 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
5195 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
5196 assert_eq!(parsed.port, 51999);
5197 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
5198 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
5200 }
5201
5202 fn test_model() -> Model {
5203 let cfg = test_dense_fixture();
5207 Model::Gguf(GgufModel {
5208 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
5209 tokenizer: Arc::new(ServerTokenizer::Byte),
5210 stop_tokens: StopTokens::default(),
5211 bos_id: None,
5212 is_synthetic: true,
5213 chat_template: chat_template::PromptTemplate::plain(),
5214 })
5215 }
5216
5217 fn greedy_params(max_tokens: usize) -> GenerationParams {
5218 GenerationParams {
5219 max_tokens,
5220 sampling: SamplingParams::default(),
5221 seed: 1,
5222 stop: Vec::new(),
5223 stop_token_ids: Vec::new(),
5224 json_object: false,
5225 grammar: None,
5226 cancel: None,
5227 ignore_eos: false,
5228 }
5229 }
5230
5231 fn test_model_full_byte_vocab() -> Model {
5235 test_model_full_byte_vocab_with_eos(None)
5236 }
5237
5238 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
5245 let mut cfg = test_dense_fixture();
5246 cfg.vocab_size = 256;
5247 Model::Gguf(GgufModel {
5248 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5249 tokenizer: Arc::new(ServerTokenizer::Byte),
5250 stop_tokens: StopTokens::from_eos(eos),
5251 bos_id: None,
5252 is_synthetic: true,
5253 chat_template: chat_template::PromptTemplate::plain(),
5254 })
5255 }
5256
5257 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
5261 AppState {
5262 embedding: None,
5263 paged_kv: None,
5264 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
5265 id: None,
5266 loaded: Loaded::Generative(Arc::new(model)),
5267 batcher: None,
5268 ceiling: None,
5269 }))),
5270 load_in_progress: std::sync::atomic::AtomicBool::new(false),
5271 tasks: Arc::new(tasks::TaskRegistry::new()),
5272 cancels: Arc::new(cancel::CancelRegistry::new()),
5273 stats: stats::Stats::new(),
5274 streams: resume::StreamRegistry::new(),
5275 model_dir: None,
5276 response_cache: Mutex::new(response_cache),
5277 kv_pool: None,
5278 prefix_cache: None,
5279 sessions: session::SessionStore::new(),
5280 requests_total: std::sync::atomic::AtomicU64::new(0),
5281 request_errors_total: std::sync::atomic::AtomicU64::new(0),
5282 started_at: std::time::Instant::now(),
5283 last_request_ms: std::sync::atomic::AtomicU64::new(0),
5284 detection: Arc::new(health::Detection::ready(health::probe_backends())),
5285 mcp: None,
5286 continuous_batching_enabled: false,
5287 metal_private_decode_gate: None,
5288 loading_model: Mutex::new(None),
5289 last_load_error: Mutex::new(None),
5290 serving: Mutex::new(crate::stats::ServingStats::default()),
5291 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
5292 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
5293 started_unix: unix_now(),
5294 }
5295 }
5296
5297 fn test_app() -> Router {
5306 test_app_with_state(Arc::new(test_state(
5307 test_model_full_byte_vocab(),
5308 ResponseCache::new(1000, Duration::from_secs(3600)),
5309 )))
5310 }
5311
5312 fn test_app_with_state(state: Arc<AppState>) -> Router {
5315 Router::new()
5316 .route(ferrox_api::routes::HEALTH, get(health))
5317 .route(ferrox_api::routes::V1_MODELS, get(list_models))
5318 .route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
5319 .route(
5320 &axum_path(ferrox_api::routes::V1_RESPONSE),
5321 get(responses::responses_get),
5322 )
5323 .route(
5324 &axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
5325 post(responses::responses_cancel),
5326 )
5327 .route(ferrox_api::routes::V1_STATS, get(serving_stats))
5328 .route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
5329 .route(
5330 ferrox_api::routes::V1_CACHE_STATUS,
5331 get(cache_admin::cache_status),
5332 )
5333 .route(
5334 ferrox_api::routes::V1_CACHE_REBUILD,
5335 post(cache_admin::cache_rebuild),
5336 )
5337 .route(
5338 ferrox_api::routes::ADMIN_PREPARE_STOP,
5339 post(cache_admin::prepare_stop),
5340 )
5341 .route("/v1/chat/completions", post(chat_completions))
5342 .route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
5343 .route(
5344 ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
5345 post(anthropic::count_tokens),
5346 )
5347 .route("/v1/tokenize", post(openai_extra::tokenize))
5348 .route("/v1/detokenize", post(openai_extra::detokenize))
5349 .route(ferrox_api::routes::TOKENIZE, post(openai_extra::tokenize))
5353 .route(
5354 ferrox_api::routes::DETOKENIZE,
5355 post(openai_extra::detokenize),
5356 )
5357 .route("/v1/embeddings", post(embeddings::embeddings))
5358 .route("/v1/completions", post(openai_extra::completions))
5359 .route(ferrox_api::routes::COMPLETION, post(completion::completion))
5361 .route(
5362 ferrox_api::routes::COMPLETIONS,
5363 post(completion::completion),
5364 )
5365 .route(
5366 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5367 post(admin::unload_model),
5368 )
5369 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
5370 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
5371 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
5372 .route(
5373 &axum_path(ferrox_api::routes::V1_STREAM),
5374 get(resume::resume),
5375 )
5376 .route(
5377 &axum_path(ferrox_api::routes::V1_STREAM_POLL),
5378 get(resume::poll),
5379 )
5380 .with_state(state)
5381 }
5382
5383 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
5384 let mut cfg = test_dense_fixture();
5385 cfg.name = name;
5386 cfg.vocab_size = vocab_size;
5387 Model::Gguf(GgufModel {
5388 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5389 tokenizer: Arc::new(ServerTokenizer::Byte),
5390 stop_tokens: StopTokens::default(),
5391 bos_id: None,
5392 is_synthetic: true,
5393 chat_template: chat_template::PromptTemplate::plain(),
5394 })
5395 }
5396
5397 fn model_with_template(name: &'static str, source: &str) -> Model {
5401 let mut cfg = test_dense_fixture();
5402 cfg.name = name;
5403 cfg.vocab_size = 256;
5404 Model::Gguf(GgufModel {
5405 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5406 tokenizer: Arc::new(ServerTokenizer::Byte),
5407 stop_tokens: StopTokens::default(),
5408 bos_id: None,
5409 is_synthetic: true,
5410 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5411 Some(source),
5412 Some("qwen3"),
5413 false,
5414 true,
5415 None,
5416 None,
5417 ),
5418 })
5419 }
5420
5421 #[tokio::test]
5430 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5431 let strict = "{% if messages | length > 1 %}\
5434 {{ raise_exception('this template takes one turn') }}\
5435 {% endif %}{{ messages[0].content }}";
5436 let state = Arc::new(test_state(
5437 model_with_template("strict", strict),
5438 ResponseCache::new(4, Duration::from_secs(60)),
5439 ));
5440 let app = test_app_with_state(state);
5441
5442 let (status, body) = post_json_uri(
5443 &app,
5444 "/v1/chat/completions",
5445 serde_json::json!({
5446 "model": "strict",
5447 "stream": true,
5448 "messages": [
5449 {"role": "user", "content": "one"},
5450 {"role": "user", "content": "two"},
5451 ],
5452 }),
5453 )
5454 .await;
5455 assert_eq!(status, StatusCode::BAD_REQUEST);
5456 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5457 assert!(
5458 body["error"]["message"]
5459 .as_str()
5460 .unwrap()
5461 .contains("one turn"),
5462 "the template's own message must reach the caller: {body}"
5463 );
5464
5465 let (status, _) = post_json_uri(
5467 &app,
5468 "/v1/chat/completions",
5469 serde_json::json!({
5470 "model": "strict",
5471 "stream": true,
5472 "max_tokens": 1,
5473 "messages": [{"role": "user", "content": "one"}],
5474 }),
5475 )
5476 .await;
5477 assert_eq!(status, StatusCode::OK);
5478 }
5479
5480 #[tokio::test]
5482 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5483 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5484 {% if reasoning_effort %}\
5485 {% if reasoning_effort not in ['low','medium','high'] %}\
5486 {{ raise_exception('bad effort') }}\
5487 {% endif %}[{{ reasoning_effort }}]\
5488 {% endif %}{{ messages[0].content }}";
5489 let state = Arc::new(test_state(
5490 model_with_template("thinker", reasoning),
5491 ResponseCache::new(4, Duration::from_secs(60)),
5492 ));
5493 let app = test_app_with_state(state);
5494 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5495 assert_eq!(status, StatusCode::OK);
5496 let entry = &models["data"][0];
5497 assert_eq!(
5498 entry["supported_reasoning_efforts"],
5499 serde_json::json!(["off", "low", "medium", "high"])
5500 );
5501 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5502 }
5503
5504 #[tokio::test]
5509 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5510 let app = test_app();
5511 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5512 let entry = &models["data"][0];
5513 assert!(entry.get("supported_reasoning_efforts").is_none());
5514 assert!(entry.get("default_reasoning_effort").is_none());
5515 }
5516
5517 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5518 Arc::new(ActiveModel {
5519 id: Some(name.to_string()),
5520 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5521 batcher: None,
5522 ceiling: None,
5523 })
5524 .tap_into(state)
5525 }
5526
5527 trait TapInto {
5529 fn tap_into(self, state: &AppState) -> Self;
5530 }
5531 impl TapInto for Arc<ActiveModel> {
5532 fn tap_into(self, state: &AppState) -> Self {
5533 state.swap_active(Some(Arc::clone(&self)));
5534 self
5535 }
5536 }
5537
5538 #[test]
5544 fn an_in_flight_request_keeps_the_model_it_started_on() {
5545 let state = test_state(
5546 named_test_model("model-a", 256),
5547 ResponseCache::new(4, Duration::from_secs(60)),
5548 );
5549
5550 let in_flight = state.active().expect("a model is loaded");
5553 assert_eq!(in_flight.name(), "model-a");
5554
5555 active_model(&state, "model-b");
5556
5557 assert_eq!(state.active().unwrap().name(), "model-b");
5559 assert_eq!(in_flight.name(), "model-a");
5561 let (_chunks, finish, _usage) = run_generation(
5562 in_flight.generative().unwrap(),
5563 "hi",
5564 &greedy_params(3),
5565 None,
5566 None,
5567 None,
5568 None,
5569 None,
5570 None,
5571 )
5572 .expect("the old model must still decode after being swapped out");
5573 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5574 }
5575
5576 #[test]
5581 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5582 let state = test_state(
5583 named_test_model("model-a", 256),
5584 ResponseCache::new(4, Duration::from_secs(60)),
5585 );
5586 let in_flight = state.active().expect("a model is loaded");
5587 let weights = Arc::clone(in_flight.generative().unwrap());
5588 assert!(Arc::strong_count(&weights) >= 2);
5589
5590 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5591 id: Some("model-b".to_string()),
5592 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5593 batcher: None,
5594 ceiling: None,
5595 })));
5596 drop(previous);
5597 assert!(Arc::strong_count(&weights) >= 2);
5599 drop(in_flight);
5600 assert_eq!(Arc::strong_count(&weights), 1);
5601 }
5602
5603 #[tokio::test]
5607 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5608 let state = Arc::new(test_state(
5609 named_test_model("model-a", 256),
5610 ResponseCache::new(4, Duration::from_secs(60)),
5611 ));
5612 let app = test_app_with_state(Arc::clone(&state));
5613
5614 let (status, body) = post_json_uri(
5615 &app,
5616 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5617 serde_json::json!({}),
5618 )
5619 .await;
5620 assert_eq!(status, StatusCode::OK);
5621 assert_eq!(body["ok"], true);
5622 assert!(body["active"].is_null());
5623 assert!(state.active().is_none());
5624
5625 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5626 assert_eq!(status, StatusCode::OK);
5627 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5628 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5629
5630 let (status, body) = post_json_uri(
5631 &app,
5632 "/v1/chat/completions",
5633 serde_json::json!({
5634 "model": "x",
5635 "messages": [{"role": "user", "content": "hi"}]
5636 }),
5637 )
5638 .await;
5639 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5640 assert_eq!(body["error"]["type"], "model_not_loaded");
5641 }
5642
5643 #[tokio::test]
5647 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5648 let state = Arc::new(test_state(
5649 named_test_model("model-a", 256),
5650 ResponseCache::new(4, Duration::from_secs(60)),
5651 ));
5652 let app = test_app_with_state(Arc::clone(&state));
5653 state.swap_active(None);
5654
5655 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5656 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5659 assert_eq!(body["state"], "unavailable");
5660 assert_eq!(body["reason"], "model_not_loaded");
5661 assert!(body["model"].is_null());
5662 let real_weights = body["capabilities"]
5663 .as_array()
5664 .unwrap()
5665 .iter()
5666 .find(|c| c["id"] == "real_weights")
5667 .cloned()
5668 .expect("real_weights is always reported");
5669 assert_eq!(real_weights["available"], false);
5670 assert_eq!(real_weights["reason"], "model_not_loaded");
5671 }
5672
5673 #[tokio::test]
5677 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5678 let app = test_app();
5679
5680 let (status, completion) = post_json_uri(
5681 &app,
5682 "/v1/chat/completions",
5683 serde_json::json!({
5684 "model": "x",
5685 "messages": [{"role": "user", "content": "hi"}],
5686 "max_tokens": 4
5687 }),
5688 )
5689 .await;
5690 assert_eq!(status, StatusCode::OK);
5691 let request_id = completion["request_id"].as_str().unwrap().to_string();
5692
5693 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5694 assert_eq!(status, StatusCode::OK);
5695 let recent = stats["recent"].as_array().unwrap();
5696 assert_eq!(recent.len(), 1);
5697 let row = &recent[0];
5698 assert_eq!(row["request_id"], request_id);
5699 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
5700 assert_eq!(row["status"], 200);
5701 assert_eq!(row["stream"], false);
5702 assert!(row["duration_ms"].is_number());
5705 assert!(row["decode_ms"].is_number());
5706 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5707 assert_eq!(
5708 stats["tokens_prompt_total"].as_u64().unwrap(),
5709 row["prompt_tokens"].as_u64().unwrap()
5710 );
5711 }
5712
5713 #[tokio::test]
5717 async fn a_rejected_request_is_recorded_too() {
5718 let state = Arc::new(test_state(
5719 named_test_model("model-a", 256),
5720 ResponseCache::new(4, Duration::from_secs(60)),
5721 ));
5722 let app = test_app_with_state(Arc::clone(&state));
5723 state.swap_active(None);
5724
5725 let (status, _) = post_json_uri(
5726 &app,
5727 "/v1/chat/completions",
5728 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5729 )
5730 .await;
5731 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5732
5733 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5734 let recent = stats["recent"].as_array().unwrap();
5735 assert_eq!(recent.len(), 1);
5736 assert_eq!(recent[0]["status"], 503);
5737 assert_eq!(recent[0]["completion_tokens"], 0);
5738 assert!(recent[0]["decode_ms"].is_null());
5739 assert_eq!(stats["errors_total"], 1);
5740 }
5741
5742 async fn post_json_with_headers(
5746 app: &Router,
5747 uri: &str,
5748 body: serde_json::Value,
5749 headers: &[(&str, &str)],
5750 ) -> (StatusCode, serde_json::Value) {
5751 use http_body_util::BodyExt;
5752 use tower::ServiceExt;
5753
5754 let mut builder = axum::http::Request::builder()
5755 .method("POST")
5756 .uri(uri)
5757 .header("content-type", "application/json");
5758 for (name, value) in headers {
5759 builder = builder.header(*name, *value);
5760 }
5761 let response = app
5762 .clone()
5763 .oneshot(
5764 builder
5765 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5766 .unwrap(),
5767 )
5768 .await
5769 .unwrap();
5770 let status = response.status();
5771 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5772 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5773 (status, json)
5774 }
5775
5776 #[tokio::test]
5780 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5781 let app = test_app();
5782
5783 let (status, _) = post_json_uri(
5784 &app,
5785 ferrox_api::routes::V1_TOKENIZE,
5786 serde_json::json!({"prompt": "hello"}),
5787 )
5788 .await;
5789 assert_eq!(status, StatusCode::OK);
5790 let (status, _) = post_json_uri(
5791 &app,
5792 ferrox_api::routes::V1_DETOKENIZE,
5793 serde_json::json!({"tokens": [104, 105]}),
5794 )
5795 .await;
5796 assert_eq!(status, StatusCode::OK);
5797 let (status, _) = post_json_uri(
5798 &app,
5799 ferrox_api::routes::V1_EMBEDDINGS,
5800 serde_json::json!({"input": "hello"}),
5801 )
5802 .await;
5803 assert_eq!(status, StatusCode::OK);
5804
5805 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5806 let routes: Vec<&str> = stats["recent"]
5807 .as_array()
5808 .unwrap()
5809 .iter()
5810 .map(|row| row["route"].as_str().unwrap())
5811 .collect();
5812 for expected in [
5813 ferrox_api::routes::V1_TOKENIZE,
5814 ferrox_api::routes::V1_DETOKENIZE,
5815 ferrox_api::routes::V1_EMBEDDINGS,
5816 ] {
5817 assert!(
5818 routes.contains(&expected),
5819 "{expected} is missing: {routes:?}"
5820 );
5821 }
5822
5823 let row = |route: &str| {
5824 stats["recent"]
5825 .as_array()
5826 .unwrap()
5827 .iter()
5828 .find(|r| r["route"] == route)
5829 .cloned()
5830 .unwrap()
5831 };
5832 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5836 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5837 assert!(embed["decode_ms"].is_null());
5838 assert_eq!(embed["completion_tokens"], 0);
5839 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5843 assert_eq!(
5844 stats["tokens_prompt_total"].as_u64().unwrap(),
5845 embed["prompt_tokens"].as_u64().unwrap(),
5846 "only the forward pass counted"
5847 );
5848 }
5849
5850 fn streaming_test_app() -> Router {
5855 let mut cfg = test_dense_fixture();
5856 cfg.vocab_size = 256;
5857 let model = Model::Gguf(GgufModel {
5858 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5859 tokenizer: Arc::new(ServerTokenizer::Byte),
5860 stop_tokens: StopTokens::default(),
5861 bos_id: None,
5862 is_synthetic: false,
5863 chat_template: chat_template::PromptTemplate::plain(),
5864 });
5865 test_app_with_state(Arc::new(test_state(
5866 model,
5867 ResponseCache::new(1000, Duration::from_secs(3600)),
5868 )))
5869 }
5870
5871 #[tokio::test]
5875 async fn the_native_completion_wire_is_not_the_openai_one() {
5876 let app = test_app();
5877
5878 let (status, native) = post_json_uri(
5879 &app,
5880 ferrox_api::routes::COMPLETION,
5881 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5882 )
5883 .await;
5884 assert_eq!(status, StatusCode::OK, "{native}");
5885 assert!(native["content"].is_string(), "{native}");
5886 assert_eq!(native["stop"], true);
5887 assert_eq!(native["stop_type"], "limit");
5888 assert_eq!(native["stopping_word"], "");
5889 assert_eq!(native["truncated"], false);
5890 assert_eq!(native["id_slot"], -1);
5891 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5892 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5893 assert!(
5894 native.get("choices").is_none(),
5895 "the native shape has no `choices`: {native}"
5896 );
5897
5898 let (status, openai) = post_json_uri(
5899 &app,
5900 ferrox_api::routes::V1_COMPLETIONS,
5901 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5902 )
5903 .await;
5904 assert_eq!(status, StatusCode::OK);
5905 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5906 assert!(
5907 openai.get("content").is_none(),
5908 "the OpenAI shape has no top-level `content`: {openai}"
5909 );
5910 }
5911
5912 #[tokio::test]
5916 async fn both_native_spellings_reach_the_same_handler() {
5917 let app = test_app();
5918 for route in [
5919 ferrox_api::routes::COMPLETION,
5920 ferrox_api::routes::COMPLETIONS,
5921 ] {
5922 let (status, body) = post_json_uri(
5923 &app,
5924 route,
5925 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5926 )
5927 .await;
5928 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5929 assert_eq!(body["stop"], true, "{route}");
5930 assert!(body["content"].is_string(), "{route}");
5931 }
5932
5933 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5936 let routes: Vec<&str> = stats["recent"]
5937 .as_array()
5938 .unwrap()
5939 .iter()
5940 .map(|row| row["route"].as_str().unwrap())
5941 .collect();
5942 assert!(
5943 routes.contains(&ferrox_api::routes::COMPLETION),
5944 "{routes:?}"
5945 );
5946 assert!(
5947 routes.contains(&ferrox_api::routes::COMPLETIONS),
5948 "{routes:?}"
5949 );
5950 }
5951
5952 #[tokio::test]
5958 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5959 let app = streaming_test_app();
5960 let raw = post_sse_raw_uri(
5961 &app,
5962 ferrox_api::routes::COMPLETION,
5963 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5964 )
5965 .await;
5966
5967 assert!(
5968 !raw.contains("[DONE]"),
5969 "llama.cpp's native stream has no sentinel: {raw}"
5970 );
5971 let frames: Vec<serde_json::Value> = raw
5972 .lines()
5973 .filter_map(|line| line.strip_prefix("data: "))
5974 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5975 .collect();
5976 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5977
5978 let (last, partials) = frames.split_last().unwrap();
5979 assert_eq!(last["stop"], true, "the last frame closes the stream");
5980 assert!(last["timings"].is_object(), "{last}");
5981 assert!(last["stop_type"].is_string(), "{last}");
5982 for partial in partials {
5983 assert_eq!(partial["stop"], false, "{partial}");
5984 assert!(partial["content"].is_string(), "{partial}");
5985 assert!(partial.get("timings").is_none(), "{partial}");
5989 assert!(partial.get("generation_settings").is_none(), "{partial}");
5990 }
5991 let streamed: String = partials
5994 .iter()
5995 .filter_map(|p| p["content"].as_str())
5996 .collect();
5997 assert_eq!(last["content"].as_str().unwrap(), streamed);
5998 }
5999
6000 #[tokio::test]
6005 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
6006 let app = test_app();
6007 for body in [
6008 serde_json::json!({"prompt": "hi"}),
6009 serde_json::json!({"prompt": "hi", "n_predict": -1}),
6010 ] {
6011 let (status, refusal) =
6012 post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
6013 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
6014 assert!(
6015 refusal["error"]["message"]
6016 .as_str()
6017 .unwrap()
6018 .contains("n_predict"),
6019 "{refusal}"
6020 );
6021 }
6022 let (status, _) = post_json_uri(
6025 &app,
6026 ferrox_api::routes::COMPLETION,
6027 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6028 )
6029 .await;
6030 assert_eq!(status, StatusCode::OK);
6031 }
6032
6033 #[tokio::test]
6042 async fn a_stop_string_halts_the_answer_and_is_named_back() {
6043 let app = streaming_test_app();
6044 let ask = |stop: serde_json::Value| {
6045 let app = app.clone();
6046 async move {
6047 post_json_uri(
6048 &app,
6049 ferrox_api::routes::COMPLETION,
6050 serde_json::json!({
6051 "prompt": "hi",
6052 "n_predict": 64,
6053 "ignore_eos": true,
6054 "stop": stop,
6055 }),
6056 )
6057 .await
6058 .1
6059 }
6060 };
6061
6062 let baseline = ask(serde_json::json!([])).await;
6063 assert_eq!(baseline["stop_type"], "limit");
6064 assert_eq!(baseline["stopping_word"], "");
6065 let text = baseline["content"].as_str().unwrap().to_string();
6066 let sentinel: String = text.chars().skip(1).take(2).collect();
6072 assert_eq!(
6073 sentinel.chars().count(),
6074 2,
6075 "the fixture must produce enough output to cut: {text:?}"
6076 );
6077 let cut = text.find(&sentinel).expect("it came out of this text");
6078
6079 let stopped = ask(serde_json::json!([sentinel])).await;
6080 assert_eq!(stopped["stop_type"], "word", "{stopped}");
6081 assert_eq!(stopped["stopping_word"], sentinel);
6082 assert_eq!(
6083 stopped["content"].as_str().unwrap(),
6084 &text[..cut],
6085 "the answer must be cut at the sentinel, not run past it"
6086 );
6087 }
6088
6089 #[tokio::test]
6095 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
6096 let app = test_app();
6097
6098 let (v1_status, v1) = post_json_uri(
6099 &app,
6100 ferrox_api::routes::V1_TOKENIZE,
6101 serde_json::json!({"prompt": "hello"}),
6102 )
6103 .await;
6104 let (alias_status, alias) = post_json_uri(
6105 &app,
6106 ferrox_api::routes::TOKENIZE,
6107 serde_json::json!({"content": "hello"}),
6108 )
6109 .await;
6110 assert_eq!(v1_status, StatusCode::OK);
6111 assert_eq!(alias_status, StatusCode::OK, "{alias}");
6112 assert_eq!(v1["tokens"], alias["tokens"]);
6113 assert!(!alias["tokens"].as_array().unwrap().is_empty());
6114
6115 let (status, both_ways) = post_json_uri(
6118 &app,
6119 ferrox_api::routes::TOKENIZE,
6120 serde_json::json!({"prompt": "hello"}),
6121 )
6122 .await;
6123 assert_eq!(status, StatusCode::OK);
6124 assert_eq!(both_ways["tokens"], v1["tokens"]);
6125 }
6126
6127 #[tokio::test]
6132 async fn detokenize_answers_under_both_dialects_keys() {
6133 let app = test_app();
6134 for route in [
6135 ferrox_api::routes::DETOKENIZE,
6136 ferrox_api::routes::V1_DETOKENIZE,
6137 ] {
6138 let (status, body) =
6139 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
6140 assert_eq!(status, StatusCode::OK, "{route}");
6141 assert_eq!(body["text"], "hi", "{route}");
6142 assert_eq!(body["content"], body["text"], "{route}");
6143 }
6144 }
6145
6146 #[tokio::test]
6150 async fn the_alias_is_recorded_under_the_path_the_client_called() {
6151 let app = test_app();
6152 let (status, _) = post_json_uri(
6153 &app,
6154 ferrox_api::routes::TOKENIZE,
6155 serde_json::json!({"content": "hello"}),
6156 )
6157 .await;
6158 assert_eq!(status, StatusCode::OK);
6159
6160 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6161 let routes: Vec<&str> = stats["recent"]
6162 .as_array()
6163 .unwrap()
6164 .iter()
6165 .map(|row| row["route"].as_str().unwrap())
6166 .collect();
6167 assert!(
6168 routes.contains(&ferrox_api::routes::TOKENIZE),
6169 "the alias must be its own row: {routes:?}"
6170 );
6171 assert!(
6172 !routes.contains(&ferrox_api::routes::V1_TOKENIZE),
6173 "nothing called /v1/tokenize: {routes:?}"
6174 );
6175 }
6176
6177 #[tokio::test]
6182 async fn add_special_prepends_the_same_bos_the_decoder_would() {
6183 let mut cfg = test_dense_fixture();
6184 cfg.vocab_size = 256;
6185 let model = Model::Gguf(GgufModel {
6186 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
6187 tokenizer: Arc::new(ServerTokenizer::Byte),
6188 stop_tokens: StopTokens::default(),
6189 bos_id: Some(7),
6190 is_synthetic: true,
6191 chat_template: chat_template::PromptTemplate::plain(),
6192 });
6193 let app = test_app_with_state(Arc::new(test_state(
6194 model,
6195 ResponseCache::new(1000, Duration::from_secs(3600)),
6196 )));
6197
6198 let (_, plain) = post_json_uri(
6199 &app,
6200 ferrox_api::routes::TOKENIZE,
6201 serde_json::json!({"content": "hi"}),
6202 )
6203 .await;
6204 let (_, special) = post_json_uri(
6205 &app,
6206 ferrox_api::routes::TOKENIZE,
6207 serde_json::json!({"content": "hi", "add_special": true}),
6208 )
6209 .await;
6210
6211 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6212 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6213 assert_eq!(special["count"], 3);
6214 }
6215
6216 #[tokio::test]
6219 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6220 let app = test_app();
6221 let (status, _) = post_json_uri(
6222 &app,
6223 ferrox_api::routes::V1_EMBEDDINGS,
6224 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6225 )
6226 .await;
6227 assert_eq!(status, StatusCode::BAD_REQUEST);
6228
6229 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6230 let recent = stats["recent"].as_array().unwrap();
6231 assert_eq!(recent.len(), 1);
6232 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
6233 assert_eq!(recent[0]["status"], 400);
6234 assert_eq!(
6235 recent[0]["prompt_tokens"], 0,
6236 "a rejected call embedded nothing"
6237 );
6238 }
6239
6240 #[tokio::test]
6243 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6244 let app = test_app();
6245 let key = "sk-monitor-secret";
6246 let (status, _) = post_json_with_headers(
6247 &app,
6248 "/v1/chat/completions",
6249 serde_json::json!({
6250 "model": "x",
6251 "messages": [{"role": "user", "content": "hi"}],
6252 "max_tokens": 2
6253 }),
6254 &[
6255 ("authorization", &format!("Bearer {key}")),
6256 ("x-ferrox-client", "ferrox-studio"),
6257 ],
6258 )
6259 .await;
6260 assert_eq!(status, StatusCode::OK);
6261
6262 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6263 let row = stats["recent"].as_array().unwrap()[0].clone();
6264 let fingerprint = row["via_api_key"]
6265 .as_str()
6266 .expect("the row names the key that served it")
6267 .to_string();
6268 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6269 assert!(!fingerprint.contains(key));
6270 assert!(
6271 !serde_json::to_string(&stats).unwrap().contains(key),
6272 "the stats payload must not carry the key in any form"
6273 );
6274 assert_eq!(row["client"], "ferrox-studio");
6275 }
6276
6277 #[tokio::test]
6280 async fn different_keys_are_different_callers_and_no_key_is_null() {
6281 let app = test_app();
6282 let body = serde_json::json!({
6283 "model": "x",
6284 "messages": [{"role": "user", "content": "hi"}],
6285 "max_tokens": 1
6286 });
6287 for headers in [
6288 vec![("authorization", "Bearer key-one")],
6289 vec![("authorization", "Bearer key-two")],
6290 vec![],
6291 ] {
6292 let (status, _) =
6293 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6294 assert_eq!(status, StatusCode::OK);
6295 }
6296
6297 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6298 let recent = stats["recent"].as_array().unwrap();
6299 assert_eq!(recent.len(), 3);
6300 let one = recent[0]["via_api_key"].as_str().unwrap();
6301 let two = recent[1]["via_api_key"].as_str().unwrap();
6302 assert_ne!(one, two, "two keys must not collapse into one caller");
6303 assert!(
6304 recent[2]["via_api_key"].is_null(),
6305 "an unauthenticated call is null, not a fingerprint of nothing"
6306 );
6307 assert!(recent[2]["client"].is_null());
6308 }
6309
6310 #[tokio::test]
6315 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6316 let state = Arc::new(test_state(
6317 named_test_model("really-loaded", 256),
6318 ResponseCache::new(4, Duration::from_secs(60)),
6319 ));
6320 let app = test_app_with_state(Arc::clone(&state));
6321
6322 let (status, _) = post_json_uri(
6323 &app,
6324 "/v1/chat/completions",
6325 serde_json::json!({
6326 "model": "gpt-4-turbo-that-is-not-here",
6327 "messages": [{"role": "user", "content": "hi"}],
6328 "max_tokens": 2
6329 }),
6330 )
6331 .await;
6332 assert_eq!(status, StatusCode::OK);
6333
6334 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6335 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6336
6337 state.swap_active(None);
6340 let (status, _) = post_json_uri(
6341 &app,
6342 "/v1/chat/completions",
6343 serde_json::json!({
6344 "model": "gpt-4-turbo-that-is-not-here",
6345 "messages": [{"role": "user", "content": "hi"}]
6346 }),
6347 )
6348 .await;
6349 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6350 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6351 let recent = stats["recent"].as_array().unwrap();
6352 assert!(recent[recent.len() - 1]["model"].is_null());
6353 }
6354
6355 #[tokio::test]
6359 async fn a_streamed_row_names_the_model_it_decoded_against() {
6360 let state = Arc::new(test_state(
6361 named_test_model("model-before", 256),
6362 ResponseCache::new(4, Duration::from_secs(60)),
6363 ));
6364 let app = test_app_with_state(Arc::clone(&state));
6365 let _ = post_sse_raw(&app, resumable_request()).await;
6366 active_model(&state, "model-after");
6368
6369 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6370 assert_eq!(stats["recent"][0]["model"], "model-before");
6371 }
6372
6373 #[tokio::test]
6376 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6377 let app = test_app();
6378 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6379 assert_eq!(status, StatusCode::OK);
6380 assert!(
6381 stats["queue_depth"].is_null(),
6382 "without continuous batching nothing queues, so there is nothing to measure"
6383 );
6384 assert!(stats["queue_rejected_total"].is_null());
6385 assert_eq!(
6386 stats["generating_now"], 0,
6387 "work in progress is measured and really is zero here"
6388 );
6389 }
6390
6391 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6396 post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
6397 }
6398
6399 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6403 use http_body_util::BodyExt;
6404 use tower::ServiceExt;
6405
6406 let response = app
6407 .clone()
6408 .oneshot(
6409 axum::http::Request::builder()
6410 .method("POST")
6411 .uri(uri)
6412 .header("content-type", "application/json")
6413 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6414 .unwrap(),
6415 )
6416 .await
6417 .unwrap();
6418 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6419 String::from_utf8(bytes.to_vec()).unwrap()
6420 }
6421
6422 async fn get_json_with_headers(
6423 app: &Router,
6424 uri: &str,
6425 headers: &[(&str, &str)],
6426 ) -> (StatusCode, serde_json::Value) {
6427 use http_body_util::BodyExt;
6428 use tower::ServiceExt;
6429
6430 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6431 for (name, value) in headers {
6432 builder = builder.header(*name, *value);
6433 }
6434 let response = app
6435 .clone()
6436 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6437 .await
6438 .unwrap();
6439 let status = response.status();
6440 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6441 (
6442 status,
6443 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6444 )
6445 }
6446
6447 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6448 body.lines()
6449 .filter_map(|line| line.strip_prefix(field))
6450 .map(str::trim)
6451 .collect()
6452 }
6453
6454 fn resumable_request() -> serde_json::Value {
6455 serde_json::json!({
6456 "model": "m",
6457 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6458 "max_tokens": 4,
6459 "temperature": 0,
6460 "stream": true,
6461 "stream_resumable": true,
6462 })
6463 }
6464
6465 #[tokio::test]
6470 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6471 let app = test_app();
6472 let body = post_sse_raw(&app, resumable_request()).await;
6473
6474 let request_id = body
6475 .lines()
6476 .find_map(|l| l.strip_prefix("data: "))
6477 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6478 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6479 .expect("the first chunk names the request");
6480
6481 let ids = sse_field(&body, "id:");
6482 let datas = sse_field(&body, "data:");
6483 assert_eq!(
6484 ids.len(),
6485 datas.len(),
6486 "every event carries an id, or a reconnect cannot name where it stopped"
6487 );
6488 for (i, id) in ids.iter().enumerate() {
6489 assert_eq!(*id, format!("{request_id}:{i}"));
6490 }
6491 let retries = sse_field(&body, "retry:");
6492 assert_eq!(
6493 retries.len(),
6494 1,
6495 "the reconnect delay is stated once, not on every event"
6496 );
6497 assert_eq!(retries[0], "1500");
6498 assert!(
6499 body.contains("data: [DONE]"),
6500 "the end of stream is still stated"
6501 );
6502 }
6503
6504 #[tokio::test]
6508 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6509 let app = test_app();
6510 let mut request = resumable_request();
6511 request["stream_resumable"] = serde_json::json!(false);
6512 let body = post_sse_raw(&app, request).await;
6513 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6514 assert!(
6515 sse_field(&body, "id:").is_empty(),
6516 "an id promises a replay this stream cannot serve"
6517 );
6518 assert!(sse_field(&body, "retry:").is_empty());
6519 }
6520
6521 #[tokio::test]
6525 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6526 let app = test_app();
6527 let body = post_sse_raw(&app, resumable_request()).await;
6528 let request_id = sse_field(&body, "id:")[0]
6529 .rsplit_once(':')
6530 .unwrap()
6531 .0
6532 .to_string();
6533 let streamed: Vec<String> = sse_field(&body, "data:")
6534 .iter()
6535 .map(|d| d.to_string())
6536 .collect();
6537
6538 let (status, polled) = get_json(
6539 &app,
6540 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
6541 )
6542 .await;
6543 assert_eq!(status, StatusCode::OK);
6544 let events: Vec<String> = polled["events"]
6545 .as_array()
6546 .unwrap()
6547 .iter()
6548 .map(|e| e["data"].as_str().unwrap().to_string())
6549 .collect();
6550 assert_eq!(
6551 events, streamed,
6552 "the fallback must deliver the same answer, not a re-run of it"
6553 );
6554 assert_eq!(polled["request_id"], request_id);
6555 assert_eq!(
6556 polled["done"], false,
6557 "events were still being handed out, so the client must ask again"
6558 );
6559
6560 let next = polled["next_index"].as_u64().unwrap();
6563 let (_, drained) = get_json(
6564 &app,
6565 &format!(
6566 "{}?from={next}",
6567 ferrox_api::routes::v1_stream_poll(&request_id)
6568 ),
6569 )
6570 .await;
6571 assert_eq!(drained["done"], true);
6572 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6573 }
6574
6575 #[tokio::test]
6579 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6580 let app = test_app();
6581 let body = post_sse_raw(&app, resumable_request()).await;
6582 let ids = sse_field(&body, "id:");
6583 let datas: Vec<String> = sse_field(&body, "data:")
6584 .iter()
6585 .map(|d| d.to_string())
6586 .collect();
6587 assert!(
6588 ids.len() >= 3,
6589 "need a few events to resume into the middle"
6590 );
6591 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6592
6593 let (status, resumed) = get_json_with_headers(
6594 &app,
6595 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
6596 &[],
6597 )
6598 .await;
6599 assert_eq!(status, StatusCode::OK);
6600 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6601
6602 let (_, tail) = get_json(
6604 &app,
6605 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
6606 )
6607 .await;
6608 let tail_events: Vec<String> = tail["events"]
6609 .as_array()
6610 .unwrap()
6611 .iter()
6612 .map(|e| e["data"].as_str().unwrap().to_string())
6613 .collect();
6614 assert_eq!(tail_events, datas[2..].to_vec());
6615 }
6616
6617 #[tokio::test]
6620 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6621 use http_body_util::BodyExt;
6622 use tower::ServiceExt;
6623
6624 let app = test_app();
6625 let body = post_sse_raw(&app, resumable_request()).await;
6626 let ids = sse_field(&body, "id:");
6627 let datas: Vec<String> = sse_field(&body, "data:")
6628 .iter()
6629 .map(|d| d.to_string())
6630 .collect();
6631 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6632
6633 let response = app
6634 .clone()
6635 .oneshot(
6636 axum::http::Request::builder()
6637 .method("GET")
6638 .uri(ferrox_api::routes::v1_stream(&request_id))
6639 .header("last-event-id", format!("{request_id}:0"))
6640 .body(axum::body::Body::empty())
6641 .unwrap(),
6642 )
6643 .await
6644 .unwrap();
6645 assert_eq!(response.status(), StatusCode::OK);
6646 assert_eq!(
6647 response
6648 .headers()
6649 .get("x-accel-buffering")
6650 .and_then(|v| v.to_str().ok()),
6651 Some("no"),
6652 "the reconnect needs the same anti-buffering header as the stream"
6653 );
6654 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6655 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6656 assert_eq!(
6657 sse_field(&resumed, "data:")
6658 .iter()
6659 .map(|d| d.to_string())
6660 .collect::<Vec<_>>(),
6661 datas[1..].to_vec()
6662 );
6663 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6664 }
6665
6666 #[tokio::test]
6670 async fn a_last_event_id_from_another_stream_is_refused() {
6671 let app = test_app();
6672 let body = post_sse_raw(&app, resumable_request()).await;
6673 let request_id = sse_field(&body, "id:")[0]
6674 .rsplit_once(':')
6675 .unwrap()
6676 .0
6677 .to_string();
6678
6679 let (status, err) = get_json_with_headers(
6680 &app,
6681 &ferrox_api::routes::v1_stream(&request_id),
6682 &[("last-event-id", "chatcmpl-someone-else:3")],
6683 )
6684 .await;
6685 assert_eq!(status, StatusCode::BAD_REQUEST);
6686 assert_eq!(err["error"]["code"], "bad_last_event_id");
6687 }
6688
6689 #[tokio::test]
6693 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6694 let app = test_app();
6695 let mut request = resumable_request();
6696 request["stream_resumable"] = serde_json::json!(false);
6697 let body = post_sse_raw(&app, request).await;
6698 let request_id = body
6699 .lines()
6700 .find_map(|l| l.strip_prefix("data: "))
6701 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6702 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6703 .unwrap();
6704
6705 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
6706 assert_eq!(status, StatusCode::NOT_FOUND);
6707 assert_eq!(err["error"]["code"], "stream_not_found");
6708 assert!(err["error"]["message"]
6709 .as_str()
6710 .unwrap()
6711 .contains("stream_resumable"));
6712 }
6713
6714 #[test]
6718 fn the_axum_stream_patterns_match_the_published_templates() {
6719 assert_eq!(
6720 axum_path(ferrox_api::routes::V1_STREAM),
6721 "/v1/stream/:request_id"
6722 );
6723 assert_eq!(
6724 axum_path(ferrox_api::routes::V1_STREAM_POLL),
6725 "/v1/stream/:request_id/poll"
6726 );
6727 assert_eq!(
6728 ferrox_api::routes::v1_stream("abc"),
6729 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
6730 );
6731 }
6732
6733 #[test]
6743 fn no_published_template_reaches_the_router_with_its_braces() {
6744 for template in [
6745 ferrox_api::routes::V1_STREAM,
6746 ferrox_api::routes::V1_STREAM_POLL,
6747 ferrox_api::routes::V1_RESPONSE,
6748 ferrox_api::routes::V1_RESPONSE_CANCEL,
6749 ferrox_api::routes::ADMIN_TASK_CANCEL,
6750 ] {
6751 assert!(
6752 template.contains('{'),
6753 "{template} is in the template list but has no placeholder"
6754 );
6755 let mounted = axum_path(template);
6756 assert!(
6757 !mounted.contains('{') && !mounted.contains('}'),
6758 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6759 );
6760 assert!(
6761 mounted.contains(':'),
6762 "{template} lost its placeholder entirely and would match one path only"
6763 );
6764 }
6765 }
6766
6767 #[tokio::test]
6775 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6776 let app = test_app();
6777 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6778 assert_eq!(status, StatusCode::NOT_FOUND);
6779 assert!(
6780 !body.is_null(),
6781 "empty body means axum never matched the route, so the id was read as a literal segment"
6782 );
6783 }
6784
6785 #[tokio::test]
6788 async fn the_task_list_starts_empty_rather_than_absent() {
6789 let app = test_app();
6790 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
6791 assert_eq!(status, StatusCode::OK);
6792 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6793 }
6794
6795 async fn post_json_uri(
6796 app: &Router,
6797 uri: &str,
6798 body: serde_json::Value,
6799 ) -> (StatusCode, serde_json::Value) {
6800 use http_body_util::BodyExt;
6801 use tower::ServiceExt;
6802
6803 let response = app
6804 .clone()
6805 .oneshot(
6806 axum::http::Request::builder()
6807 .method("POST")
6808 .uri(uri)
6809 .header("content-type", "application/json")
6810 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6811 .unwrap(),
6812 )
6813 .await
6814 .unwrap();
6815 let status = response.status();
6816 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6817 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6818 (status, json)
6819 }
6820
6821 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6822 post_json_uri(app, "/v1/chat/completions", body).await.1
6823 }
6824
6825 #[tokio::test]
6834 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6835 let app = test_app();
6836 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6837 assert_eq!(status, StatusCode::OK);
6838
6839 let memory = &body["memory"];
6840 if memory.is_null() {
6841 return;
6844 }
6845 assert!(
6846 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6847 "a read that produced a zero is a broken read, not an idle \
6848 engine: {memory}"
6849 );
6850 assert!(
6851 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6852 "the quantity must travel with the number: {memory}"
6853 );
6854 }
6855
6856 #[tokio::test]
6862 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6863 let app = test_app();
6864 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6865 assert_eq!(status, StatusCode::OK);
6866 for pool in ["kv_pages", "window_slots", "state_slots"] {
6867 assert!(
6868 body["pools"][pool].is_null(),
6869 "{pool} must be null rather than a zero row: {}",
6870 body["pools"]
6871 );
6872 }
6873 }
6874
6875 #[tokio::test]
6883 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6884 use http_body_util::BodyExt;
6885 use tower::ServiceExt;
6886
6887 let app = test_app();
6888 let send = |body: serde_json::Value| {
6889 let app = app.clone();
6890 async move {
6891 app.oneshot(
6892 axum::http::Request::builder()
6893 .method("POST")
6894 .uri(ferrox_api::routes::V1_MESSAGES)
6895 .header("content-type", "application/json")
6896 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6897 .unwrap(),
6898 )
6899 .await
6900 .unwrap()
6901 }
6902 };
6903
6904 let ok = send(serde_json::json!({
6905 "model": "test",
6906 "max_tokens": 1,
6907 "messages": [{"role": "user", "content": "hi"}],
6908 }))
6909 .await;
6910 assert_eq!(ok.status(), StatusCode::OK);
6911 let id = ok
6912 .headers()
6913 .get("request-id")
6914 .expect("a served message names its id")
6915 .to_str()
6916 .unwrap()
6917 .to_string();
6918 assert!(!id.is_empty());
6919
6920 let bad = send(serde_json::json!({"model": "test"})).await;
6923 assert!(bad.status().is_client_error());
6924 let other = bad.headers().get("request-id").expect("errors too");
6925 assert_ne!(other.to_str().unwrap(), id);
6926 let _ = bad.into_body().collect().await.unwrap();
6927 }
6928
6929 #[tokio::test]
6937 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6938 let state = Arc::new(test_state(
6939 test_model_full_byte_vocab(),
6940 ResponseCache::new(1000, Duration::from_secs(3600)),
6941 ));
6942 let app = test_app_with_state(Arc::clone(&state));
6943 let body = serde_json::json!({
6944 "model": "test",
6945 "messages": [{"role": "user", "content": "hi"}],
6946 "max_tokens": 1,
6947 });
6948
6949 state
6950 .maintenance
6951 .lock()
6952 .unwrap()
6953 .begin_rebuild()
6954 .expect("a fresh server is serving, so the rebuild starts");
6955 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6956 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6957 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6958
6959 state.maintenance.lock().unwrap().finish_rebuild(true);
6960 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6961 assert_eq!(
6962 status,
6963 StatusCode::OK,
6964 "the gate reopens; a rebuild is not a latch"
6965 );
6966 }
6967
6968 #[tokio::test]
6973 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6974 let app = test_app();
6975 let (status, body) = post_json_uri(
6976 &app,
6977 ferrox_api::routes::V1_CANCEL,
6978 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6979 )
6980 .await;
6981 assert_eq!(status, StatusCode::NOT_FOUND);
6982 assert_eq!(body["cancelled"], serde_json::json!(false));
6983 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6984 assert!(
6985 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6986 "the verdict must carry a human reason: {body}"
6987 );
6988 }
6989
6990 #[tokio::test]
6994 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6995 let state = Arc::new(test_state(
6996 test_model_full_byte_vocab(),
6997 ResponseCache::new(1000, Duration::from_secs(3600)),
6998 ));
6999 let app = test_app_with_state(Arc::clone(&state));
7000 let (token, _guard) = state.cancels.register("chatcmpl-live");
7001
7002 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
7003 assert_eq!(status, StatusCode::OK);
7004 assert_eq!(before["generating_now"], serde_json::json!(1));
7005
7006 let (status, body) = post_json_uri(
7007 &app,
7008 ferrox_api::routes::V1_CANCEL,
7009 serde_json::json!({ "request_id": "chatcmpl-live" }),
7010 )
7011 .await;
7012 assert_eq!(status, StatusCode::OK);
7013 assert_eq!(body["cancelled"], serde_json::json!(true));
7014 assert!(
7015 token.is_cancelled(),
7016 "the endpoint answered ok without setting the flag the decode loop reads"
7017 );
7018 }
7019
7020 #[tokio::test]
7021 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
7022 let app = test_app();
7023 let (status, tok) =
7024 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
7025 assert_eq!(status, StatusCode::OK);
7026 let tokens = tok["tokens"].as_array().unwrap();
7027 assert_eq!(tok["count"], tokens.len());
7028 assert!(!tokens.is_empty());
7029
7030 let (status, detok) = post_json_uri(
7031 &app,
7032 "/v1/detokenize",
7033 serde_json::json!({ "tokens": tokens }),
7034 )
7035 .await;
7036 assert_eq!(status, StatusCode::OK);
7037 assert_eq!(detok["text"], "Hi");
7038
7039 let (status, emb) = post_json_uri(
7040 &app,
7041 "/v1/embeddings",
7042 serde_json::json!({
7043 "input": "Hi",
7044 "embedding_type": "mean"
7045 }),
7046 )
7047 .await;
7048 assert_eq!(status, StatusCode::OK);
7049 let vec = emb["data"][0]["embedding"].as_array().unwrap();
7050 assert!(!vec.is_empty());
7051 assert!(vec.iter().all(|v| v.as_f64().is_some()));
7052 }
7053
7054 #[tokio::test]
7060 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
7061 let app = test_app();
7062 let (status, body) = post_json_uri(
7063 &app,
7064 "/v1/embeddings",
7065 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
7066 )
7067 .await;
7068 assert_eq!(status, StatusCode::BAD_REQUEST);
7069 let msg = body["error"]["message"].as_str().unwrap();
7070 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
7071 }
7072
7073 #[tokio::test]
7078 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7079 async fn a_real_embedding_model_serves_v1_embeddings() {
7080 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7081 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7082 if !path.exists() {
7083 eprintln!("SKIP: {} not present", path.display());
7084 return;
7085 }
7086 let encoder = ferrox_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
7087 let mut state = test_state(
7088 test_model_full_byte_vocab(),
7089 ResponseCache::new(1000, Duration::from_secs(3600)),
7090 );
7091 state.embedding = Some(Arc::new(encoder));
7092 let app = test_app_with_state(Arc::new(state));
7093
7094 let (status, body) = post_json_uri(
7095 &app,
7096 "/v1/embeddings",
7097 serde_json::json!({ "input": ["Hello world", "a second input"] }),
7098 )
7099 .await;
7100 assert_eq!(status, StatusCode::OK, "{body}");
7101 assert_eq!(body["model"], "bge-small-en-v1.5");
7102 let data = body["data"].as_array().unwrap();
7103 assert_eq!(data.len(), 2);
7104 for (i, row) in data.iter().enumerate() {
7105 assert_eq!(row["index"], i);
7106 let v: Vec<f64> = row["embedding"]
7107 .as_array()
7108 .unwrap()
7109 .iter()
7110 .map(|x| x.as_f64().unwrap())
7111 .collect();
7112 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
7113 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
7114 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
7115 }
7116 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
7119
7120 let (status, mean) = post_json_uri(
7124 &app,
7125 "/v1/embeddings",
7126 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
7127 )
7128 .await;
7129 assert_eq!(status, StatusCode::OK);
7130 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
7131 }
7132
7133 #[tokio::test]
7145 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7146 async fn an_encoder_can_be_the_loaded_model() {
7147 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7148 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7149 if !path.exists() {
7150 eprintln!("SKIP: {} not present", path.display());
7151 return;
7152 }
7153
7154 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
7158 assert!(
7159 matches!(loaded, model::LoadedModel::Encoder(_)),
7160 "an encoder-only GGUF reached a decoder loader"
7161 );
7162 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
7163 assert!(
7164 matches!(loaded, Loaded::Encoder(_)),
7165 "the encoder did not stay an encoder through activation"
7166 );
7167 assert!(
7168 batcher.is_none() && ceiling.is_none(),
7169 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7170 );
7171
7172 let state = test_state(
7173 test_model_full_byte_vocab(),
7174 ResponseCache::new(1000, Duration::from_secs(3600)),
7175 );
7176 state.swap_active(Some(Arc::new(ActiveModel {
7177 id: None,
7178 loaded,
7179 batcher,
7180 ceiling,
7181 })));
7182 let app = test_app_with_state(Arc::new(state));
7183
7184 let (status, body) = post_json_uri(
7186 &app,
7187 "/v1/embeddings",
7188 serde_json::json!({ "input": "Hello world" }),
7189 )
7190 .await;
7191 assert_eq!(status, StatusCode::OK, "{body}");
7192 assert_eq!(body["model"], "bge-small-en-v1.5");
7193 let v = body["data"][0]["embedding"].as_array().unwrap();
7194 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
7195
7196 let (status, body) = post_json_uri(
7198 &app,
7199 "/v1/chat/completions",
7200 serde_json::json!({
7201 "model": "bge-small-en-v1.5",
7202 "messages": [{"role": "user", "content": "hi"}],
7203 }),
7204 )
7205 .await;
7206 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
7207 let msg = body["error"]["message"].as_str().unwrap();
7208 for fact in [
7209 "bge-small-en-v1.5",
7210 "bert",
7211 "embedding model",
7212 "/v1/embeddings",
7213 ] {
7214 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
7215 }
7216
7217 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
7219 assert_eq!(status, StatusCode::OK);
7220 let entry = &models["data"][0];
7221 assert_eq!(entry["id"], "bge-small-en-v1.5");
7222 assert_eq!(entry["ferrox_model_kind"], "embedding");
7223 assert_eq!(entry["ferrox_tokenizer"], "gguf-wordpiece");
7224 assert_eq!(entry["ferrox_n_embd"], 384);
7225 assert_eq!(entry["ferrox_pooling"], "CLS");
7226 assert_eq!(
7227 entry["ferrox_endpoints"],
7228 serde_json::json!(["/v1/embeddings"])
7229 );
7230 assert!(entry.get("supported_reasoning_efforts").is_none());
7233
7234 let (status, health) = get_json(&app, ferrox_api::routes::HEALTH).await;
7236 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
7237 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
7238 assert_eq!(health["model"]["synthetic_weights"], false);
7239 let weights = health["capabilities"]
7240 .as_array()
7241 .unwrap()
7242 .iter()
7243 .find(|c| c["id"] == ferrox_api::health::capability::REAL_WEIGHTS)
7244 .expect("a real-weights capability row");
7245 let detail = weights["detail"].as_str().unwrap_or_default();
7246 assert!(detail.contains("ENCODER"), "{detail}");
7247 let (status, body) = post_json_uri(
7255 &app,
7256 ferrox_api::routes::V1_TOKENIZE,
7257 serde_json::json!({ "content": "hello world" }),
7258 )
7259 .await;
7260 assert_eq!(
7261 status,
7262 StatusCode::OK,
7263 "an encoder has a real tokenizer: {body}"
7264 );
7265 let tokens = body["tokens"].as_array().expect("tokens array").clone();
7266 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
7267
7268 let (status, body) = post_json_uri(
7269 &app,
7270 ferrox_api::routes::V1_DETOKENIZE,
7271 serde_json::json!({ "tokens": tokens }),
7272 )
7273 .await;
7274 assert_eq!(status, StatusCode::OK, "{body}");
7275 let round_tripped = body["content"].as_str().expect("content").to_string();
7276 assert!(
7277 round_tripped.contains("hello") && round_tripped.contains("world"),
7278 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
7279 );
7280
7281 let (status, _) = post_json_uri(
7284 &app,
7285 "/v1/completions",
7286 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
7287 )
7288 .await;
7289 assert_eq!(
7290 status,
7291 StatusCode::NOT_IMPLEMENTED,
7292 "tokenizing an encoder must not have opened a path to generating with one"
7293 );
7294 }
7295
7296 #[tokio::test]
7301 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7302 use http_body_util::BodyExt;
7303 use tower::ServiceExt;
7304
7305 let fixture = concat!(
7306 "../ferrox-models/tests/fixtures/",
7307 "ferrox_real_moe_test.gguf"
7308 );
7309 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7310 let decoder = Decoder::from_gguf_with_expert_cache(
7311 &fixture,
7312 ferrox_models::config::test_moe_fixture(),
7313 Some(1024 * 1024),
7314 )
7315 .expect("MoE fixture must load store-backed");
7316
7317 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
7321 .layers
7322 .iter()
7323 .map(|_| {
7324 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
7325 })
7326 .collect();
7327 decoder.forward_token(1, 0, &mut caches);
7328
7329 let model = Model::Gguf(GgufModel {
7330 decoder: Arc::new(decoder),
7331 tokenizer: Arc::new(ServerTokenizer::Byte),
7332 stop_tokens: StopTokens::default(),
7333 bos_id: None,
7334 is_synthetic: false,
7335 chat_template: chat_template::PromptTemplate::plain(),
7336 });
7337 let state = Arc::new(test_state(
7338 model,
7339 ResponseCache::new(16, Duration::from_secs(60)),
7340 ));
7341 let app = Router::new()
7342 .route("/metrics", axum::routing::get(metrics))
7343 .route("/v1/chat/completions", post(chat_completions))
7344 .with_state(state);
7345
7346 let fetch_metrics = |app: Router| async move {
7347 let resp = app
7348 .oneshot(
7349 axum::http::Request::builder()
7350 .method("GET")
7351 .uri("/metrics")
7352 .body(axum::body::Body::empty())
7353 .unwrap(),
7354 )
7355 .await
7356 .unwrap();
7357 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7358 String::from_utf8(bytes.to_vec()).unwrap()
7359 };
7360
7361 let after = fetch_metrics(app.clone()).await;
7362 assert!(
7363 after.contains("ferrox_expert_cache_misses_total"),
7364 "streaming model must expose expert-cache metrics: {after}"
7365 );
7366 let misses: u64 = after
7367 .lines()
7368 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
7369 .and_then(|l| l.split_whitespace().nth(1))
7370 .and_then(|v| v.parse().ok())
7371 .expect("misses metric line must parse");
7372 assert!(
7373 misses > 0,
7374 "decode must have read experts through the store: {after}"
7375 );
7376 }
7377
7378 fn weather_tool() -> serde_json::Value {
7379 serde_json::json!({
7380 "type": "function",
7381 "function": {
7382 "name": "get_weather",
7383 "description": "Get the current weather for a location.",
7384 "parameters": {
7385 "type": "object",
7386 "properties": {"location": {"type": "string"}},
7387 "required": ["location"]
7388 }
7389 }
7390 })
7391 }
7392
7393 fn weather_tool_def() -> ToolDef {
7394 ToolDef {
7395 kind: "function".to_string(),
7396 function: ToolFunctionDef {
7397 name: "get_weather".to_string(),
7398 description: Some("Get the current weather for a location.".to_string()),
7399 parameters: Some(serde_json::json!({
7400 "type": "object",
7401 "properties": {"location": {"type": "string"}},
7402 "required": ["location"]
7403 })),
7404 },
7405 }
7406 }
7407
7408 #[test]
7409 fn tool_preamble_mentions_every_tool_name_and_description() {
7410 let preamble = tool_preamble(&[weather_tool_def()]);
7411 assert!(preamble.contains("get_weather"));
7412 assert!(preamble.contains("Get the current weather for a location."));
7413 assert!(preamble.contains("<tool_call>"));
7414 assert!(preamble.contains("</tool_call>"));
7415 }
7416
7417 #[test]
7418 fn a_real_marker_becomes_a_structured_tool_call() {
7419 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7420 let (message, finish) = build_response_message(
7421 text.to_string(),
7422 &[weather_tool_def()],
7423 output::OutputPosture::for_model("test-model"),
7424 "stop",
7425 );
7426 assert_eq!(finish, "tool_calls");
7427 let calls = message.tool_calls.expect("must carry a tool call");
7428 assert_eq!(calls[0].function.name, "get_weather");
7429 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7430 assert_eq!(parsed["location"], "Paris");
7431 }
7432
7433 #[test]
7434 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7435 let (message, finish) = build_response_message(
7436 "just an answer".to_string(),
7437 &[weather_tool_def()],
7438 output::OutputPosture::for_model("test-model"),
7439 "stop",
7440 );
7441 assert_eq!(finish, "stop");
7442 assert!(message.tool_calls.is_none());
7443 assert_eq!(message.content.as_deref(), Some("just an answer"));
7444 }
7445
7446 #[test]
7449 fn a_malformed_payload_is_not_a_tool_call() {
7450 let (message, finish) = build_response_message(
7451 "<tool_call>not valid json at all</tool_call>".to_string(),
7452 &[weather_tool_def()],
7453 output::OutputPosture::for_model("test-model"),
7454 "stop",
7455 );
7456 assert_eq!(finish, "stop");
7457 assert!(message.tool_calls.is_none());
7458 }
7459
7460 #[test]
7463 fn a_tool_that_was_never_offered_is_not_returned() {
7464 let (message, finish) = build_response_message(
7465 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7466 &[weather_tool_def()],
7467 output::OutputPosture::for_model("test-model"),
7468 "stop",
7469 );
7470 assert_eq!(finish, "stop");
7471 assert!(message.tool_calls.is_none());
7472 }
7473
7474 #[test]
7476 fn marker_text_with_no_tools_offered_stays_content() {
7477 let (message, finish) = build_response_message(
7478 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7479 &[],
7480 output::OutputPosture::for_model("test-model"),
7481 "stop",
7482 );
7483 assert_eq!(finish, "stop");
7484 assert!(message.tool_calls.is_none());
7485 assert!(message.content.is_some());
7486 }
7487
7488 #[test]
7492 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7493 let opened = std::cell::Cell::new(0usize);
7494 let mut parser = crate::policy::parser::ToolCallParser::new(
7495 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7496 vec![
7497 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7498 "write_file",
7499 serde_json::json!({"type": "object", "properties": {
7500 "path": {"type": "string"},
7501 "contents": {"type": "string"}
7502 }}),
7503 ),
7504 ],
7505 );
7506 let wire = "<tool_call><function=write_file>\
7507 <parameter=path>\n/tmp/x\n</parameter>\
7508 <parameter=contents>\nhello world\n</parameter>\
7509 </function></tool_call>";
7510
7511 let mut deltas = Vec::new();
7512 let mut text = String::new();
7513 for piece in wire.as_bytes().chunks(7) {
7514 let chunk = String::from_utf8_lossy(piece).into_owned();
7515 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7516 text.push_str(&more_text);
7517 deltas.extend(more);
7518 }
7519 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7520 text.push_str(&more_text);
7521 deltas.extend(more);
7522
7523 assert_eq!(opened.get(), 1, "one call opened");
7524 assert!(text.is_empty(), "the markers are not content: {text:?}");
7525
7526 let first = &deltas[0];
7527 assert_eq!(first.index, 0);
7528 assert_eq!(first.id.as_deref(), Some("call_0"));
7529 assert_eq!(first.kind, Some("function"));
7530 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7531
7532 let joined: String = deltas
7535 .iter()
7536 .filter_map(|d| d.function.arguments.clone())
7537 .collect();
7538 let parsed: serde_json::Value =
7539 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7540 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7541 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7542 assert!(
7543 deltas.len() >= 3,
7544 "the arguments arrived in pieces, not whole: {}",
7545 deltas.len()
7546 );
7547 assert!(
7548 deltas[1..].iter().all(|d| d.function.name.is_none()),
7549 "only the opening delta carries identity"
7550 );
7551 }
7552
7553 #[test]
7555 fn text_around_a_streamed_call_is_still_content() {
7556 let opened = std::cell::Cell::new(0usize);
7557 let mut parser = crate::policy::parser::ToolCallParser::new(
7558 crate::policy::parser::ToolCallFormat::Qwen25,
7559 vec![crate::policy::parser::tool_call::ToolSchema::new(
7560 "get_weather",
7561 )],
7562 );
7563 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7564 \"arguments\": {}}</tool_call> done";
7565 let mut text = String::new();
7566 for piece in wire.as_bytes().chunks(5) {
7567 let chunk = String::from_utf8_lossy(piece).into_owned();
7568 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7569 text.push_str(&more);
7570 }
7571 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7572 text.push_str(&more);
7573
7574 assert_eq!(opened.get(), 1);
7575 assert!(text.starts_with("let me check. "), "{text:?}");
7576 assert!(text.ends_with(" done"), "{text:?}");
7577 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7578 }
7579
7580 #[test]
7583 fn a_reasoning_block_is_split_out_of_the_answer() {
7584 let (message, finish) = build_response_message(
7585 "<think>weighing it up</think>The answer is 4.".to_string(),
7586 &[],
7587 output::OutputPosture::for_model("Qwen3-8B"),
7588 "stop",
7589 );
7590 assert_eq!(finish, "stop");
7591 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7592 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7593 }
7594
7595 #[test]
7598 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7599 let (message, _) = build_response_message(
7600 "Use the <think> tag like this.".to_string(),
7601 &[],
7602 output::OutputPosture::for_model("llama-3.1-8b"),
7603 "stop",
7604 );
7605 assert_eq!(
7606 message.content.as_deref(),
7607 Some("Use the <think> tag like this.")
7608 );
7609 assert!(message.reasoning_content.is_none());
7610 }
7611
7612 #[tokio::test]
7620 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7621 let app = test_app();
7622 let body = serde_json::json!({
7623 "model": "m",
7624 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7625 "max_tokens": 4,
7626 "temperature": 0,
7627 });
7628 let resp = post_json(&app, body).await;
7629 let message = &resp["choices"][0]["message"];
7630 assert!(message["content"].is_string());
7631 assert!(message.get("tool_calls").is_none());
7632 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7633 assert_eq!(resp["usage"]["completion_tokens"], 4);
7634 assert_eq!(
7635 resp["usage"]["total_tokens"],
7636 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7637 );
7638 }
7639
7640 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7641 use http_body_util::BodyExt;
7642 use tower::ServiceExt;
7643
7644 let response = app
7645 .clone()
7646 .oneshot(
7647 axum::http::Request::builder()
7648 .method("GET")
7649 .uri(uri)
7650 .body(axum::body::Body::empty())
7651 .unwrap(),
7652 )
7653 .await
7654 .unwrap();
7655 let status = response.status();
7656 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7657 (status, serde_json::from_slice(&bytes).unwrap())
7658 }
7659
7660 #[tokio::test]
7661 async fn health_answers_a_capability_handshake_not_a_boolean() {
7662 let app = test_app();
7663 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7664 assert_eq!(status, StatusCode::OK);
7665
7666 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7667 assert_eq!(health.state, ferrox_api::HealthState::Ready);
7668 assert!(health.pid > 0);
7669 assert!(health.server_time_unix_ms > 0);
7670 assert_eq!(health.last_request_age_seconds, None);
7673
7674 for id in [
7677 ferrox_api::health::capability::CPU,
7678 ferrox_api::health::capability::METAL,
7679 ferrox_api::health::capability::CUDA,
7680 ferrox_api::health::capability::REAL_WEIGHTS,
7681 ferrox_api::health::capability::CONTINUOUS_BATCHING,
7682 ] {
7683 let cap = health
7684 .capability(id)
7685 .unwrap_or_else(|| panic!("{id} missing"));
7686 assert!(!cap.reason.is_empty(), "{cap:?}");
7687 assert!(!cap.detail.is_empty(), "{cap:?}");
7688 }
7689 let weights = health
7693 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
7694 .unwrap();
7695 assert!(!weights.available);
7696 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
7697 assert!(health.model.as_ref().unwrap().synthetic_weights);
7698 }
7699
7700 #[tokio::test]
7701 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7702 let app = test_app();
7703 let _ = post_json(
7704 &app,
7705 serde_json::json!({
7706 "model": "m",
7707 "messages": [{"role": "user", "content": "\u{1}"}],
7708 "max_tokens": 1,
7709 "temperature": 0,
7710 }),
7711 )
7712 .await;
7713 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7714 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7715 let age = health
7716 .last_request_age_seconds
7717 .expect("a served request is evidence of liveness");
7718 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7719 }
7720
7721 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7723 use http_body_util::BodyExt;
7724 use tower::ServiceExt;
7725
7726 let response = app
7727 .clone()
7728 .oneshot(
7729 axum::http::Request::builder()
7730 .method("POST")
7731 .uri("/v1/chat/completions")
7732 .header("content-type", "application/json")
7733 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7734 .unwrap(),
7735 )
7736 .await
7737 .unwrap();
7738 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7739 String::from_utf8(bytes.to_vec())
7740 .unwrap()
7741 .lines()
7742 .filter_map(|line| line.strip_prefix("data: "))
7743 .filter(|payload| *payload != "[DONE]")
7744 .map(|payload| serde_json::from_str(payload).unwrap())
7745 .collect()
7746 }
7747
7748 #[tokio::test]
7749 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7750 let app = test_app();
7751 let chunks = post_sse_chunks(
7752 &app,
7753 serde_json::json!({
7754 "model": "m",
7755 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7756 "max_tokens": 4,
7757 "temperature": 0,
7758 "stream": true,
7759 }),
7760 )
7761 .await;
7762
7763 assert!(!chunks.is_empty());
7764 let request_id = chunks[0]["request_id"]
7765 .as_str()
7766 .expect("the first chunk names the request")
7767 .to_string();
7768 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7769 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7772 assert!(
7773 chunk.get("request_id").is_none(),
7774 "chunk {i} repeats request_id"
7775 );
7776 }
7777 for chunk in &chunks {
7780 assert_eq!(chunk["id"], serde_json::json!(request_id));
7781 }
7782
7783 let other = post_sse_chunks(
7784 &app,
7785 serde_json::json!({
7786 "model": "m",
7787 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7788 "max_tokens": 4,
7789 "temperature": 0,
7790 "stream": true,
7791 }),
7792 )
7793 .await;
7794 assert_ne!(
7795 other[0]["request_id"].as_str().unwrap(),
7796 request_id,
7797 "two concurrent chats must not share an id"
7798 );
7799 }
7800
7801 #[tokio::test]
7802 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7803 let app = test_app();
7804 let resp = post_json(
7805 &app,
7806 serde_json::json!({
7807 "model": "m",
7808 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7809 "max_tokens": 2,
7810 "temperature": 0,
7811 }),
7812 )
7813 .await;
7814 assert_eq!(resp["id"], resp["request_id"]);
7815 assert!(resp["request_id"]
7816 .as_str()
7817 .unwrap()
7818 .starts_with("chatcmpl-"));
7819 }
7820
7821 #[tokio::test]
7824 async fn usage_carries_separate_prefill_and_decode_timings() {
7825 let app = test_app();
7826 let resp = post_json(
7827 &app,
7828 serde_json::json!({
7829 "model": "m",
7830 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7831 "max_tokens": 4,
7832 "temperature": 0,
7833 }),
7834 )
7835 .await;
7836 let usage = &resp["usage"];
7837 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7838 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7839 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7840 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7841 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7843 }
7844
7845 #[tokio::test]
7853 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7854 let app = test_app();
7855 let body = serde_json::json!({
7856 "model": "m",
7857 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7858 "max_tokens": 4,
7859 "temperature": 0,
7860 "tools": [weather_tool()],
7861 });
7862 let resp = post_json(&app, body).await;
7863 let message = &resp["choices"][0]["message"];
7864 assert!(
7865 message["content"].is_string(),
7866 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7867 );
7868 assert!(message.get("tool_calls").is_none());
7869 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7873 }
7874
7875 #[tokio::test]
7879 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7880 let app = test_app();
7881 let body = serde_json::json!({
7882 "model": "m",
7883 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7884 "max_tokens": 3,
7885 "temperature": 0,
7886 });
7887 let first = post_json(&app, body.clone()).await;
7888 assert_eq!(first["ferrox_cache"], "miss");
7889 let second = post_json(&app, body).await;
7890 assert_eq!(second["ferrox_cache"], "hit");
7891 assert_eq!(
7892 first["choices"][0]["message"]["content"],
7893 second["choices"][0]["message"]["content"]
7894 );
7895 assert_eq!(
7896 first["choices"][0]["finish_reason"],
7897 second["choices"][0]["finish_reason"]
7898 );
7899 assert_eq!(first["usage"], second["usage"]);
7900 assert_eq!(second["usage"]["completion_tokens"], 3);
7901 }
7902
7903 #[tokio::test]
7914 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7915 let app = test_app();
7916 let plain = serde_json::json!({
7917 "model": "m",
7918 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7919 "max_tokens": 3,
7920 "temperature": 0,
7921 });
7922
7923 let first = post_json(&app, plain.clone()).await;
7924 assert_eq!(first["ferrox_cache"], "miss");
7925 let unconstrained = first["choices"][0]["message"]["content"]
7926 .as_str()
7927 .expect("content")
7928 .to_string();
7929
7930 let mut constrained = plain.clone();
7931 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7932 let second = post_json(&app, constrained).await;
7933 assert_eq!(
7934 second["ferrox_cache"], "miss",
7935 "a grammar is part of the key, so this body has never been answered"
7936 );
7937 let constrained_answer = second["choices"][0]["message"]["content"]
7941 .as_str()
7942 .expect("content")
7943 .to_string();
7944 assert!(
7945 constrained_answer.contains("-> \"yes\"]"),
7946 "the grammar must have been compiled AND applied, not skipped \
7947 by a cache hit: {constrained_answer}"
7948 );
7949 assert_ne!(
7950 constrained_answer, unconstrained,
7951 "the constrained request was served the unconstrained answer"
7952 );
7953
7954 let third = post_json(&app, plain).await;
7958 assert_eq!(third["ferrox_cache"], "hit");
7959 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7960 }
7961
7962 #[tokio::test]
7984 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7985 let state = Arc::new(test_state(
7986 test_model_full_byte_vocab(),
7987 ResponseCache::new(1000, Duration::from_secs(3600)),
7988 ));
7989 let app = test_app_with_state(state.clone());
7990 let plain = serde_json::json!({
7991 "model": "m",
7992 "messages": [
7993 {"role": "system", "content": "Answer in JSON when it helps."},
7994 {"role": "user", "content": "\u{1}\u{2}"},
7995 ],
7996 "max_tokens": 3,
7997 "temperature": 0,
7998 });
7999
8000 let first = post_json(&app, plain.clone()).await;
8001 assert_eq!(first["ferrox_cache"], "miss");
8002 assert_eq!(state.cache_stats().entries, 1);
8003
8004 let mut as_json = plain.clone();
8005 as_json["response_format"] = serde_json::json!({"type": "json_object"});
8006 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
8007 assert_eq!(
8008 status,
8009 StatusCode::BAD_REQUEST,
8010 "the demo banner is not a JSON object, whoever generated it"
8011 );
8012 assert_eq!(
8013 state.cache_stats().hits,
8014 0,
8015 "a json_object request must not be answered from an entry the \
8016 JSON mask never produced"
8017 );
8018 assert_eq!(
8019 state.cache_stats().entries,
8020 2,
8021 "json_object must key its own entry, not reuse the unconstrained \
8022 one it happens to render the same prompt as"
8023 );
8024 }
8025
8026 #[tokio::test]
8038 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
8039 let app = test_app_with_state(Arc::new(test_state(
8040 test_model_full_byte_vocab_with_eos(Some(0x77)),
8041 ResponseCache::new(1000, Duration::from_secs(3600)),
8042 )));
8043 let body = serde_json::json!({
8044 "model": "m",
8045 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8046 "max_tokens": 6,
8047 "temperature": 0,
8048 });
8049
8050 let stopped = post_json(&app, body.clone()).await;
8051 assert_eq!(stopped["ferrox_cache"], "miss");
8052 assert_eq!(
8053 stopped["choices"][0]["finish_reason"], "stop",
8054 "the fixture is only meaningful if the model's EOS really fires here"
8055 );
8056 assert_eq!(stopped["usage"]["completion_tokens"], 1);
8057
8058 let mut ignoring = body.clone();
8059 ignoring["ignore_eos"] = serde_json::json!(true);
8060 let ran_on = post_json(&app, ignoring).await;
8061 assert_eq!(
8062 ran_on["ferrox_cache"], "miss",
8063 "ignore_eos is part of the key, so this body has never been answered"
8064 );
8065 assert_eq!(
8066 ran_on["usage"]["completion_tokens"], 6,
8067 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
8068 );
8069 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
8070 assert_ne!(
8071 ran_on["choices"][0]["message"]["content"],
8072 stopped["choices"][0]["message"]["content"]
8073 );
8074 }
8075
8076 #[tokio::test]
8082 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
8083 let session_app = test_app();
8084 let manual_app = test_app();
8085
8086 let turn1 = post_json(
8088 &session_app,
8089 serde_json::json!({
8090 "model": "m",
8091 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8092 "session_id": "s1",
8093 "max_tokens": 5,
8094 "temperature": 0,
8095 }),
8096 )
8097 .await;
8098 let reply1 = turn1["choices"][0]["message"]["content"]
8099 .as_str()
8100 .unwrap()
8101 .to_string();
8102
8103 let manual_turn1 = post_json(
8108 &manual_app,
8109 serde_json::json!({
8110 "model": "m",
8111 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8112 "max_tokens": 5,
8113 "temperature": 0,
8114 }),
8115 )
8116 .await;
8117 assert_eq!(
8118 manual_turn1["choices"][0]["message"]["content"]
8119 .as_str()
8120 .unwrap(),
8121 reply1
8122 );
8123
8124 let turn2 = post_json(
8126 &session_app,
8127 serde_json::json!({
8128 "model": "m",
8129 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
8130 "session_id": "s1",
8131 "max_tokens": 5,
8132 "temperature": 0,
8133 }),
8134 )
8135 .await;
8136 let reply2 = turn2["choices"][0]["message"]["content"]
8137 .as_str()
8138 .unwrap()
8139 .to_string();
8140
8141 let manual_turn2 = post_json(
8145 &manual_app,
8146 serde_json::json!({
8147 "model": "m",
8148 "messages": [
8149 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
8150 {"role": "assistant", "content": reply1},
8151 {"role": "user", "content": "\u{4}\u{5}"},
8152 ],
8153 "max_tokens": 5,
8154 "temperature": 0,
8155 }),
8156 )
8157 .await;
8158 assert_eq!(
8159 manual_turn2["choices"][0]["message"]["content"]
8160 .as_str()
8161 .unwrap(),
8162 reply2,
8163 "resuming a session must produce identical output to manually resending the full history"
8164 );
8165 }
8166
8167 #[test]
8170 fn lock_cache_recovers_from_a_poisoned_mutex() {
8171 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
8172
8173 let poison_cache = Arc::clone(&cache);
8174 let _ = std::thread::spawn(move || {
8175 let _guard = poison_cache.lock().unwrap();
8176 panic!("simulated panic while holding the lock");
8177 })
8178 .join();
8179
8180 let recovered = lock_cache(&cache);
8182 assert_eq!(recovered.stats().entries, 0);
8183 }
8184
8185 #[test]
8186 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8187 let mut req_body = serde_json::json!({
8188 "model": "m",
8189 "messages": [{"role": "user", "content": "hi"}],
8190 });
8191 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8192 assert!(
8193 req.is_cacheable(),
8194 "default (temperature 0) must be cacheable"
8195 );
8196
8197 req_body["temperature"] = serde_json::json!(0.8);
8198 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8199 assert!(
8200 !req.is_cacheable(),
8201 "unseeded sampling must never be cacheable"
8202 );
8203
8204 req_body["seed"] = serde_json::json!(42);
8205 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
8206 assert!(
8207 req.is_cacheable(),
8208 "sampling with an explicit seed is deterministic and must be cacheable"
8209 );
8210 }
8211
8212 const GRADED: &str = "{% if reasoning_effort %}\
8216 {% if reasoning_effort not in ['low','medium','high'] %}\
8217 {{ raise_exception('unsupported effort') }}\
8218 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
8219 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
8220
8221 fn graded_template() -> chat_template::PromptTemplate {
8222 chat_template::PromptTemplate::from_gguf_metadata(
8223 Some(GRADED),
8224 Some("qwen3"),
8225 false,
8226 true,
8227 None,
8228 None,
8229 )
8230 }
8231
8232 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
8233 serde_json::from_value(value).expect("request")
8234 }
8235
8236 #[test]
8243 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
8244 let req = chat_request(serde_json::json!({
8245 "model": "m",
8246 "messages": [{"role": "user", "content": "hi"}],
8247 "grammar": "root ::= \"a\"+",
8248 }));
8249 req.validate_supported_fields()
8250 .expect("a valid grammar is a valid request");
8251 let params = req
8252 .generation_params()
8253 .expect("a valid grammar compiles at params time too");
8254 assert!(
8255 params.grammar.is_some(),
8256 "the grammar was dropped between the wire and the sampler"
8257 );
8258 assert!(
8259 params.needs_vocab_logits(),
8260 "a grammar request that may fold lm_head into a GPU argmax is \
8261 a grammar request served unconstrained"
8262 );
8263
8264 let plain = chat_request(serde_json::json!({
8265 "model": "m",
8266 "messages": [{"role": "user", "content": "hi"}],
8267 }));
8268 assert!(plain.generation_params().unwrap().grammar.is_none());
8269 }
8270
8271 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
8272 chat_request(serde_json::json!({
8273 "model": "m",
8274 "messages": [{"role": "user", "content": "weather in Rome?"}],
8275 "tools": [weather_tool()],
8276 "tool_choice": tool_choice,
8277 }))
8278 }
8279
8280 #[test]
8284 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
8285 for choice in [
8286 serde_json::json!("required"),
8287 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
8288 ] {
8289 let req = tool_request(choice.clone());
8290 req.validate_supported_fields()
8291 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
8292 let params = req
8293 .generation_params_for_template(&graded_template(), "Qwen3-8B")
8294 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
8295 let grammar = params
8296 .grammar
8297 .as_ref()
8298 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
8299 assert!(
8300 grammar.is_awaiting_trigger(),
8301 "the model must be free to think before it calls"
8302 );
8303 assert!(
8304 !grammar.allows_eog(),
8305 "{choice} must not be able to end the turn without a call"
8306 );
8307 assert!(
8313 params.needs_vocab_logits(),
8314 "{choice} would let a backend return a token id instead of logits"
8315 );
8316 assert!(
8317 !generate::greedy_gpu_fold_allowed(¶ms),
8318 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8319 );
8320 }
8321 }
8322
8323 #[test]
8325 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8326 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8327 let req = tool_request(choice.clone());
8328 req.validate_supported_fields().expect("still supported");
8329 let params = match req.generation_params_for_template(&graded_template(), "Qwen3-8B") {
8330 Ok(p) => p,
8331 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
8332 };
8333 assert!(
8334 params.grammar.is_none(),
8335 "{choice} does not force a call and must not be constrained"
8336 );
8337 }
8338 }
8339
8340 #[test]
8343 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8344 let req = chat_request(serde_json::json!({
8346 "model": "m",
8347 "messages": [{"role": "user", "content": "hi"}],
8348 "tool_choice": "required",
8349 }));
8350 let (status, _) = req
8351 .validate_supported_fields()
8352 .expect_err("nothing to call");
8353 assert_eq!(status, StatusCode::BAD_REQUEST);
8354
8355 let req =
8357 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8358 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8359 assert_eq!(status, StatusCode::BAD_REQUEST);
8360 assert_eq!(body["error"]["param"], "tool_choice");
8361
8362 let req = tool_request(serde_json::json!({"type": "function"}));
8364 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8365 assert_eq!(status, StatusCode::BAD_REQUEST);
8366
8367 let req = chat_request(serde_json::json!({
8369 "model": "m",
8370 "messages": [{"role": "user", "content": "hi"}],
8371 "tools": [weather_tool()],
8372 "tool_choice": "required",
8373 "grammar": "root ::= \"a\"+",
8374 }));
8375 let (status, _) = req
8376 .validate_supported_fields()
8377 .expect_err("a grammar and a forced call are two constraints");
8378 assert_eq!(status, StatusCode::BAD_REQUEST);
8379
8380 let req = tool_request(serde_json::json!("required"));
8386 let (status, Json(body)) =
8387 match req.generation_params_for_template(&graded_template(), "Gemma4-27B") {
8388 Err(e) => e,
8389 Ok(_) => panic!("a gemma4 call's arguments are not an object rule"),
8390 };
8391 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8392 assert!(
8393 body["error"]["message"]
8394 .as_str()
8395 .unwrap()
8396 .contains("gemma4"),
8397 "{body}"
8398 );
8399 }
8400
8401 #[test]
8404 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8405 let req = chat_request(serde_json::json!({
8406 "model": "m",
8407 "messages": [{"role": "user", "content": "hi"}],
8408 "grammar": "root ::= \"a",
8409 }));
8410 let (status, Json(body)) = req
8411 .validate_supported_fields()
8412 .expect_err("this does not parse");
8413 assert_eq!(status, StatusCode::BAD_REQUEST);
8414 assert_eq!(body["error"]["param"], "grammar");
8415 assert!(req.generation_params().is_err(), "and again at params time");
8416 }
8417
8418 #[test]
8425 fn response_format_json_schema_becomes_the_requests_grammar() {
8426 let req = chat_request(serde_json::json!({
8427 "model": "m",
8428 "messages": [{"role": "user", "content": "hi"}],
8429 "response_format": {
8430 "type": "json_schema",
8431 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8432 },
8433 }));
8434 req.validate_supported_fields()
8435 .expect("a boolean schema converts");
8436 let params = req.generation_params().expect("and compiles");
8437 let grammar = params.grammar.expect("the schema is the grammar");
8438 let mut g = (*grammar).clone();
8439 g.accept_token(0, b"true").expect("a boolean is accepted");
8440 assert!(g.allows_eog(), "and completes the parse");
8441 assert!(
8442 !params.json_object,
8443 "a schema is not the json_object character-class mask"
8444 );
8445 }
8446
8447 #[test]
8451 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8452 let req = chat_request(serde_json::json!({
8453 "model": "m",
8454 "messages": [{"role": "user", "content": "hi"}],
8455 "response_format": {
8456 "type": "json_schema",
8457 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8458 },
8459 }));
8460 let (status, Json(body)) = req
8461 .validate_supported_fields()
8462 .expect_err("minimum has no grammar in this port");
8463 assert_eq!(status, StatusCode::BAD_REQUEST);
8464 assert!(
8465 body["error"]["message"]
8466 .as_str()
8467 .expect("a message")
8468 .contains("minimum"),
8469 "the refusal must name the keyword: {body}"
8470 );
8471 assert!(req.generation_params().is_err(), "and again at params time");
8472 }
8473
8474 #[test]
8480 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8481 let req = chat_request(serde_json::json!({
8482 "model": "m",
8483 "messages": [{"role": "user", "content": "hi"}],
8484 "tool_choice": "required",
8485 "tools": [{
8486 "type": "function",
8487 "function": {"name": "f", "parameters": {"type": "object"}},
8488 }],
8489 "response_format": {
8490 "type": "json_schema",
8491 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8492 },
8493 }));
8494 let (status, Json(body)) = req
8495 .validate_supported_fields()
8496 .expect_err("two constraints, one generation");
8497 assert_eq!(status, StatusCode::BAD_REQUEST);
8498 assert_eq!(body["error"]["param"], "tool_choice");
8499 }
8500
8501 #[test]
8504 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8505 let req = chat_request(serde_json::json!({
8506 "model": "m",
8507 "messages": [{"role": "user", "content": "hi"}],
8508 }));
8509 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8510 }
8511
8512 #[test]
8517 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8518 let asked = chat_request(serde_json::json!({
8519 "model": "m",
8520 "messages": [{"role": "user", "content": "hi"}],
8521 "min_p": 0.07,
8522 }));
8523 assert_eq!(asked.sampling_params().expect("knobs").min_p, 0.07);
8524
8525 let silent = chat_request(serde_json::json!({
8526 "model": "m",
8527 "messages": [{"role": "user", "content": "hi"}],
8528 }));
8529 assert_eq!(
8530 silent.sampling_params().expect("knobs").min_p,
8531 0.0,
8532 "an unset min_p must be off, not llama.cpp's CLI default"
8533 );
8534 }
8535
8536 #[test]
8545 fn no_sampler_knob_is_missing_from_the_cache_key() {
8546 let base = serde_json::json!({
8547 "model": "m",
8548 "messages": [{"role": "user", "content": "hi"}],
8549 "seed": 1,
8550 });
8551 let key_for = |body: serde_json::Value| {
8552 let req = chat_request(body);
8553 let params = req.generation_params().expect("params");
8554 req.cache_key("prompt", ¶ms)
8555 };
8556 let baseline = key_for(base.clone());
8557 for (knob, value) in [
8558 ("temperature", serde_json::json!(0.5)),
8559 ("top_p", serde_json::json!(0.9)),
8560 ("min_p", serde_json::json!(0.05)),
8561 ("top_k", serde_json::json!(40)),
8562 ("repetition_penalty", serde_json::json!(1.1)),
8563 ("presence_penalty", serde_json::json!(0.3)),
8564 ("frequency_penalty", serde_json::json!(0.3)),
8565 (
8566 "samplers",
8567 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8568 ),
8569 ] {
8570 let mut body = base.clone();
8571 body[knob] = value;
8572 assert_ne!(
8573 key_for(body),
8574 baseline,
8575 "`{knob}` is not in the cache key: two requests differing \
8576 only in it would share one cached answer"
8577 );
8578 }
8579 }
8580
8581 #[test]
8591 fn no_constraint_is_missing_from_the_cache_key() {
8592 let base = serde_json::json!({
8593 "model": "m",
8594 "messages": [{"role": "user", "content": "pick one"}],
8595 });
8596 let key_for = |body: serde_json::Value| {
8597 let req = chat_request(body);
8598 let params = req.generation_params().expect("params");
8599 req.cache_key("prompt", ¶ms)
8600 };
8601 let baseline = key_for(base.clone());
8602 for (field, value) in [
8603 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8604 (
8605 "response_format",
8606 serde_json::json!({"type": "json_object"}),
8607 ),
8608 (
8609 "response_format",
8610 serde_json::json!({"type": "json_schema", "json_schema": {
8611 "name": "answer",
8612 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8613 }}),
8614 ),
8615 ("ignore_eos", serde_json::json!(true)),
8616 ("stop", serde_json::json!(["\n"])),
8617 ("max_tokens", serde_json::json!(7)),
8618 ] {
8619 let mut body = base.clone();
8620 body[field] = value.clone();
8621 assert_ne!(
8622 key_for(body),
8623 baseline,
8624 "`{field}: {value}` is not in the cache key: two requests \
8625 differing only in it would share one cached answer"
8626 );
8627 }
8628 }
8629
8630 #[test]
8634 fn an_explicit_zero_output_budget_is_a_client_error() {
8635 let req = chat_request(serde_json::json!({
8636 "model": "m",
8637 "messages": [{"role": "user", "content": "hi"}],
8638 "max_tokens": 0,
8639 }));
8640 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8641 assert_eq!(status, StatusCode::BAD_REQUEST);
8642 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8643 }
8644
8645 #[test]
8648 fn a_request_can_turn_thinking_off() {
8649 let template = graded_template();
8650 for body in [
8651 serde_json::json!({
8652 "model": "m",
8653 "messages": [{"role": "user", "content": "hi"}],
8654 "reasoning_effort": "none",
8655 }),
8656 serde_json::json!({
8657 "model": "m",
8658 "messages": [{"role": "user", "content": "hi"}],
8659 "thinking": {"type": "disabled"},
8660 }),
8661 ] {
8662 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8663 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8664 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8665 assert!(!kwargs.contains_key("reasoning_effort"));
8668 }
8669 }
8670
8671 #[test]
8674 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8675 let template = graded_template();
8676 let kwargs = chat_request(serde_json::json!({
8677 "model": "m",
8678 "messages": [{"role": "user", "content": "hi"}],
8679 "reasoning_effort": "high",
8680 "thinking": {"type": "disabled"},
8681 }))
8682 .resolve_template_kwargs(&template);
8683 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8684 assert!(!kwargs.contains_key("reasoning_effort"));
8685 }
8686
8687 #[test]
8690 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8691 let req = chat_request(serde_json::json!({
8692 "model": "m",
8693 "messages": [{"role": "user", "content": "hi"}],
8694 "thinking": {"type": "disable"},
8695 }));
8696 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8697 assert_eq!(status, StatusCode::BAD_REQUEST);
8698 }
8699
8700 #[test]
8703 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8704 let template = graded_template();
8705 let kwargs = chat_request(serde_json::json!({
8706 "model": "m",
8707 "messages": [{"role": "user", "content": "hi"}],
8708 "reasoning_effort": "none",
8709 "chat_template_kwargs": {"enable_thinking": true},
8710 }))
8711 .resolve_template_kwargs(&template);
8712 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8713 }
8714
8715 #[test]
8719 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8720 let template = graded_template();
8721 let req = chat_request(serde_json::json!({
8722 "model": "m",
8723 "messages": [{"role": "user", "content": "hi"}],
8724 "reasoning_effort": "minimal",
8725 }));
8726 let kwargs = req.resolve_template_kwargs(&template);
8727 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8728 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8729 assert!(prompt.starts_with("E:low|"), "{prompt}");
8730 }
8731
8732 #[test]
8736 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8737 let template = graded_template();
8738 let req = chat_request(serde_json::json!({
8739 "model": "m",
8740 "messages": [{"role": "user", "content": "hi"}],
8741 "chat_template_kwargs": {"reasoning_effort": "none"},
8742 }));
8743 let kwargs = req.resolve_template_kwargs(&template);
8744 assert!(!kwargs.contains_key("reasoning_effort"));
8745 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8746 assert_eq!(prompt, "hi");
8747 }
8748
8749 #[test]
8752 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8753 let template = graded_template();
8754 let req = chat_request(serde_json::json!({
8755 "model": "m",
8756 "messages": [{"role": "user", "content": "hi"}],
8757 "reasoning_effort": "low",
8758 "chat_template_kwargs": {"reasoning_effort": "high"},
8759 }));
8760 assert_eq!(
8761 req.resolve_template_kwargs(&template)["reasoning_effort"],
8762 serde_json::json!("high")
8763 );
8764 }
8765
8766 #[test]
8770 fn offering_tools_turns_thinking_on_by_itself() {
8771 let template = graded_template();
8772 let quiet = chat_request(serde_json::json!({
8773 "model": "m",
8774 "messages": [{"role": "user", "content": "hi"}],
8775 }));
8776 assert!(!quiet
8777 .resolve_template_kwargs(&template)
8778 .contains_key("enable_thinking"));
8779
8780 let with_tools = chat_request(serde_json::json!({
8781 "model": "m",
8782 "messages": [{"role": "user", "content": "hi"}],
8783 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8784 }));
8785 let kwargs = with_tools.resolve_template_kwargs(&template);
8786 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8787 let prompt =
8788 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8789 assert!(prompt.starts_with("THINK|"), "{prompt}");
8790 }
8791
8792 #[test]
8797 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8798 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8799 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8800 Some("qwen3"),
8801 false,
8802 true,
8803 None,
8804 None,
8805 );
8806 let req = chat_request(serde_json::json!({
8807 "model": "m",
8808 "messages": [{"role": "user", "content": "hi"}],
8809 "chat_template_kwargs": {"enable_thinking": true},
8810 }));
8811 let kwargs = req.resolve_template_kwargs(&opener);
8812 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8813 assert!(prompt.ends_with("<think>"), "{prompt}");
8814
8815 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8818 let (message, _) = build_response_message(
8819 "weighing it up</think>Paris.".to_string(),
8820 &[],
8821 posture,
8822 "stop",
8823 );
8824 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8825 assert_eq!(message.content.as_deref(), Some("Paris."));
8826
8827 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8830 let (message, _) = build_response_message(
8831 "weighing it up</think>Paris.".to_string(),
8832 &[],
8833 closed,
8834 "stop",
8835 );
8836 assert_eq!(message.reasoning_content, None);
8837 }
8838
8839 #[test]
8840 fn stop_param_accepts_both_single_string_and_array() {
8841 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8842 "model": "m",
8843 "messages": [{"role": "user", "content": "hi"}],
8844 "stop": "END",
8845 }))
8846 .unwrap();
8847 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8848
8849 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8850 "model": "m",
8851 "messages": [{"role": "user", "content": "hi"}],
8852 "stop": ["A", "B"],
8853 }))
8854 .unwrap();
8855 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8856 }
8857
8858 #[test]
8859 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8860 let model = test_model();
8861 let result = run_generation(
8862 &model,
8863 "hello",
8864 &greedy_params(4),
8865 None,
8866 None,
8867 None,
8868 None,
8869 None,
8870 None,
8871 );
8872 assert!(matches!(
8873 result,
8874 Err(generate::DecodeError::TokenOutOfVocab { .. })
8875 ));
8876 }
8877
8878 #[test]
8882 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8883 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8885 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8886
8887 let holder_pool = Arc::clone(&pool);
8888 let holder = std::thread::spawn(move || {
8889 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8890 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8892 drop(held);
8893 });
8894 std::thread::sleep(Duration::from_millis(15));
8895
8896 let config = generate::KvPoolConfig {
8897 pool,
8898 queue_wait: Duration::ZERO,
8899 };
8900 let result = run_generation(
8901 &model,
8902 &prompt,
8903 &greedy_params(4),
8904 Some(&config),
8905 None,
8906 None,
8907 None,
8908 None,
8909 None,
8910 );
8911 assert!(matches!(
8912 result,
8913 Err(generate::DecodeError::KvPoolExhausted)
8914 ));
8915
8916 let (status, _body) = decode_error_response(result.unwrap_err());
8917 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8918 holder.join().unwrap();
8919 }
8920
8921 #[test]
8929 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8930 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8932 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
8934 let config = generate::KvPoolConfig {
8935 pool,
8936 queue_wait: Duration::ZERO,
8937 };
8938
8939 let result = run_generation(
8940 &model,
8941 &prompt,
8942 &greedy_params(4),
8943 Some(&config),
8944 None,
8945 None,
8946 None,
8947 None,
8948 None,
8949 );
8950 let err = result.expect_err("one block cannot hold two layers' caches");
8951 assert!(
8952 matches!(
8953 &err,
8954 generate::DecodeError::KvBudgetExceeded { binding, .. }
8955 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
8956 ),
8957 "expected an immovable device-memory refusal, got {err:?}"
8958 );
8959 let (status, _body) = decode_error_response(err);
8960 assert_eq!(status, StatusCode::BAD_REQUEST);
8961 }
8962
8963 #[test]
8969 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8970 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8971 queued: 512,
8972 cap: 512,
8973 });
8974 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8975 assert_eq!(body["error"]["retry_after_seconds"], 1);
8976 let message = body["error"]["message"].as_str().expect("message");
8977 assert!(message.contains("512"), "{message}");
8978 }
8979
8980 #[test]
8981 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8982 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8983 token: 99,
8984 vocab_size: 32,
8985 });
8986 assert!(
8987 body["error"]["retry_after_seconds"].is_null(),
8988 "retrying a prompt this model cannot tokenize never helps"
8989 );
8990 }
8991
8992 #[test]
8993 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8994 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8995 token: 99,
8996 vocab_size: 32,
8997 });
8998 assert_eq!(status, StatusCode::BAD_REQUEST);
8999 }
9000
9001 #[test]
9002 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
9003 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9005 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
9006 let config = generate::KvPoolConfig {
9007 pool: pool.clone(),
9008 queue_wait: Duration::ZERO,
9009 };
9010
9011 let (_, finish, _usage) = run_generation(
9012 &model,
9013 &prompt,
9014 &greedy_params(4),
9015 Some(&config),
9016 None,
9017 None,
9018 None,
9019 None,
9020 None,
9021 )
9022 .unwrap();
9023 assert_eq!(finish, FinishReason::Length);
9024 assert_eq!(
9025 pool.lock().unwrap().free_blocks(),
9026 2,
9027 "a completed request must return its blocks to the pool"
9028 );
9029 }
9030
9031 #[tokio::test]
9036 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
9037 let model = Arc::new(test_model());
9038 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9039
9040 let mut handles = Vec::new();
9041 for _ in 0..8 {
9042 let model = Arc::clone(&model);
9043 let prompt = prompt.clone();
9044 handles.push(tokio::task::spawn_blocking(move || {
9045 run_generation(
9046 &model,
9047 &prompt,
9048 &greedy_params(6),
9049 None,
9050 None,
9051 None,
9052 None,
9053 None,
9054 None,
9055 )
9056 .unwrap()
9057 }));
9058 }
9059
9060 let mut results = Vec::new();
9061 for h in handles {
9062 results.push(h.await.unwrap());
9063 }
9064 for r in &results[1..] {
9069 assert_eq!(r.0, results[0].0, "decoded chunks must match");
9070 assert_eq!(r.1, results[0].1, "finish reason must match");
9071 assert_eq!(
9072 r.2.prompt_tokens, results[0].2.prompt_tokens,
9073 "prompt token count must match"
9074 );
9075 assert_eq!(
9076 r.2.completion_tokens, results[0].2.completion_tokens,
9077 "completion token count must match"
9078 );
9079 }
9080 }
9081
9082 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
9089 let mut header_entries = Vec::new();
9090 let mut data = Vec::new();
9091 for (name, shape, values) in tensors {
9092 let start = data.len();
9093 for v in values {
9094 data.extend_from_slice(&v.to_le_bytes());
9095 }
9096 let end = data.len();
9097 let shape_str = shape
9098 .iter()
9099 .map(|d| d.to_string())
9100 .collect::<Vec<_>>()
9101 .join(",");
9102 header_entries.push(format!(
9103 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
9104 ));
9105 }
9106 let header = format!("{{{}}}", header_entries.join(","));
9107 let header_bytes = header.as_bytes();
9108 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
9109 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
9110 out.extend_from_slice(header_bytes);
9111 out.extend_from_slice(&data);
9112 out
9113 }
9114
9115 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9126 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9127 use ferrox_models::kimi_loader::KimiRealHparams;
9128 use ferrox_moe::{GatingFunction, MoeLayerConfig};
9129
9130 let hidden_dim = 8;
9131 let kda_num_heads = 2;
9132 let kda_head_dim = 3;
9133 let kda_proj = kda_num_heads * kda_head_dim;
9134 let conv_kernel = 4;
9135 let dense_intermediate = 5;
9136 let vocab_size = 256;
9140 let mla_num_heads = 1;
9141 let mla_q_lora_rank = 2;
9142 let mla_kv_lora_rank = 2;
9143 let mla_qk_nope_head_dim = 2;
9144 let mla_qk_rope_head_dim = 2;
9145 let mla_v_head_dim = 2;
9146
9147 let model_cfg = ferrox_models::ModelConfig {
9148 name: "synthetic-kimi-server-test",
9149 n_layers: 1,
9150 hidden_dim,
9151 n_heads: 1,
9152 n_kv_heads: 1,
9153 head_dim: 4,
9154 vocab_size,
9155 rope_theta: 10000.0,
9156 rms_norm_eps: 1e-5,
9157 sliding_window: None,
9158 moe: MoeLayerConfig {
9159 expert_weights_scale: 1.0,
9160 n_experts: 1,
9161 n_experts_active: 1,
9162 n_shared_experts: 0,
9163 hidden_dim,
9164 expert_ffn_dim: 4,
9165 gating: GatingFunction::Sigmoid,
9166 norm_topk_prob: true,
9167 expert_group_count: None,
9168 expert_group_used_count: None,
9169 },
9170 n_dense_leading_layers: 1,
9175 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
9176 kda_layers: vec![1],
9177 full_attn_layers: vec![],
9178 mla: MlaConfig {
9179 num_heads: mla_num_heads,
9180 q_lora_rank: mla_q_lora_rank,
9181 kv_lora_rank: mla_kv_lora_rank,
9182 qk_nope_head_dim: mla_qk_nope_head_dim,
9183 qk_rope_head_dim: mla_qk_rope_head_dim,
9184 v_head_dim: mla_v_head_dim,
9185 use_output_gate: true,
9186 rope: None,
9187 },
9188 kda: KdaConfig {
9189 num_heads: kda_num_heads,
9190 head_dim: kda_head_dim,
9191 short_conv_kernel_size: conv_kernel,
9192 gate_lower_bound: -5.0,
9193 use_full_rank_gate: true,
9194 },
9195 }),
9196 rope_freqs: None,
9197 rope_attn_factor: 1.0,
9198 rope_dim: None,
9199 rope_freqs_long: None,
9200 rope_freqs_short: None,
9201 rope_orig_ctx: None,
9202 rope_layout: ferrox_models::config::RopeLayout::Neox,
9203 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
9204 swa_pattern: None,
9205 swa_dense_first: false,
9206 attn_logit_softcap: None,
9207 final_logit_softcap: None,
9208 embedding_scale: None,
9209 attention_scale: None,
9210 rope_theta_swa: None,
9211 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
9212 best_effort_fields: &["synthetic test config, not a real preset"],
9213 };
9214 let hp = KimiRealHparams {
9215 hidden_dim,
9216 kda_num_heads,
9217 kda_head_dim,
9218 mla_num_heads,
9219 mla_q_lora_rank,
9220 mla_kv_lora_rank,
9221 mla_qk_nope_head_dim,
9222 mla_qk_rope_head_dim,
9223 mla_v_head_dim,
9224 dense_intermediate_dim: dense_intermediate,
9225 moe_hidden_dim: hidden_dim,
9226 moe_intermediate_dim: 4,
9227 n_experts: 1,
9228 num_shared_experts: 0,
9229 };
9230
9231 let prefix = "language_model.model.layers.0";
9235 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
9236 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
9237 name: String,
9238 shape: Vec<usize>,
9239 n: usize| {
9240 tensors.push((name, shape, vec![0.01f32; n]));
9241 };
9242 push(
9243 &mut tensors,
9244 format!("{prefix}.input_layernorm.weight"),
9245 vec![hidden_dim],
9246 hidden_dim,
9247 );
9248 push(
9249 &mut tensors,
9250 format!("{prefix}.post_attention_layernorm.weight"),
9251 vec![hidden_dim],
9252 hidden_dim,
9253 );
9254 push(
9255 &mut tensors,
9256 format!("{prefix}.self_attention_res_norm.weight"),
9257 vec![hidden_dim],
9258 hidden_dim,
9259 );
9260 push(
9261 &mut tensors,
9262 format!("{prefix}.self_attention_res_proj.weight"),
9263 vec![1, hidden_dim],
9264 hidden_dim,
9265 );
9266 push(
9267 &mut tensors,
9268 format!("{prefix}.mlp_res_norm.weight"),
9269 vec![hidden_dim],
9270 hidden_dim,
9271 );
9272 push(
9273 &mut tensors,
9274 format!("{prefix}.mlp_res_proj.weight"),
9275 vec![1, hidden_dim],
9276 hidden_dim,
9277 );
9278 push(
9279 &mut tensors,
9280 format!("{prefix}.self_attn.q_proj.weight"),
9281 vec![kda_proj, hidden_dim],
9282 kda_proj * hidden_dim,
9283 );
9284 push(
9285 &mut tensors,
9286 format!("{prefix}.self_attn.k_proj.weight"),
9287 vec![kda_proj, hidden_dim],
9288 kda_proj * hidden_dim,
9289 );
9290 push(
9291 &mut tensors,
9292 format!("{prefix}.self_attn.v_proj.weight"),
9293 vec![kda_proj, hidden_dim],
9294 kda_proj * hidden_dim,
9295 );
9296 push(
9297 &mut tensors,
9298 format!("{prefix}.self_attn.q_conv1d.weight"),
9299 vec![kda_proj, 1, conv_kernel],
9300 kda_proj * conv_kernel,
9301 );
9302 push(
9303 &mut tensors,
9304 format!("{prefix}.self_attn.k_conv1d.weight"),
9305 vec![kda_proj, 1, conv_kernel],
9306 kda_proj * conv_kernel,
9307 );
9308 push(
9309 &mut tensors,
9310 format!("{prefix}.self_attn.v_conv1d.weight"),
9311 vec![kda_proj, 1, conv_kernel],
9312 kda_proj * conv_kernel,
9313 );
9314 push(
9315 &mut tensors,
9316 format!("{prefix}.self_attn.A_log"),
9317 vec![kda_num_heads],
9318 kda_num_heads,
9319 );
9320 push(
9321 &mut tensors,
9322 format!("{prefix}.self_attn.f_a_proj.weight"),
9323 vec![kda_head_dim, hidden_dim],
9324 kda_head_dim * hidden_dim,
9325 );
9326 push(
9327 &mut tensors,
9328 format!("{prefix}.self_attn.f_b_proj.weight"),
9329 vec![kda_proj, kda_head_dim],
9330 kda_proj * kda_head_dim,
9331 );
9332 push(
9333 &mut tensors,
9334 format!("{prefix}.self_attn.dt_bias"),
9335 vec![kda_proj],
9336 kda_proj,
9337 );
9338 push(
9339 &mut tensors,
9340 format!("{prefix}.self_attn.b_proj.weight"),
9341 vec![kda_num_heads, hidden_dim],
9342 kda_num_heads * hidden_dim,
9343 );
9344 push(
9345 &mut tensors,
9346 format!("{prefix}.self_attn.g_proj.weight"),
9347 vec![kda_proj, hidden_dim],
9348 kda_proj * hidden_dim,
9349 );
9350 push(
9351 &mut tensors,
9352 format!("{prefix}.self_attn.o_norm.weight"),
9353 vec![kda_head_dim],
9354 kda_head_dim,
9355 );
9356 push(
9357 &mut tensors,
9358 format!("{prefix}.self_attn.o_proj.weight"),
9359 vec![hidden_dim, kda_proj],
9360 hidden_dim * kda_proj,
9361 );
9362 push(
9363 &mut tensors,
9364 format!("{prefix}.mlp.gate_proj.weight"),
9365 vec![dense_intermediate, hidden_dim],
9366 dense_intermediate * hidden_dim,
9367 );
9368 push(
9369 &mut tensors,
9370 format!("{prefix}.mlp.up_proj.weight"),
9371 vec![dense_intermediate, hidden_dim],
9372 dense_intermediate * hidden_dim,
9373 );
9374 push(
9375 &mut tensors,
9376 format!("{prefix}.mlp.down_proj.weight"),
9377 vec![hidden_dim, dense_intermediate],
9378 hidden_dim * dense_intermediate,
9379 );
9380 push(
9381 &mut tensors,
9382 "language_model.model.embed_tokens.weight".to_string(),
9383 vec![vocab_size, hidden_dim],
9384 vocab_size * hidden_dim,
9385 );
9386 push(
9387 &mut tensors,
9388 "language_model.lm_head.weight".to_string(),
9389 vec![vocab_size, hidden_dim],
9390 vocab_size * hidden_dim,
9391 );
9392 push(
9393 &mut tensors,
9394 "language_model.model.norm.weight".to_string(),
9395 vec![hidden_dim],
9396 hidden_dim,
9397 );
9398 push(
9399 &mut tensors,
9400 "language_model.model.output_attn_res_norm.weight".to_string(),
9401 vec![hidden_dim],
9402 hidden_dim,
9403 );
9404 push(
9405 &mut tensors,
9406 "language_model.model.output_attn_res_proj.weight".to_string(),
9407 vec![1, hidden_dim],
9408 hidden_dim,
9409 );
9410
9411 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9424 let dir = std::env::temp_dir().join(format!(
9425 "ferrox_server_kimi_e2e_test_{}_{}",
9426 std::process::id(),
9427 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9428 ));
9429 std::fs::create_dir_all(&dir).unwrap();
9430 let shard_bytes = write_safetensors_shard(&tensors);
9431 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9432 let map_entries: Vec<String> = tensors
9433 .iter()
9434 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9435 .collect();
9436 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9437 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9438
9439 use base64::Engine;
9443 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9444 .map(|b| {
9445 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9446 format!("{b64} {b}")
9447 })
9448 .collect();
9449 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9450
9451 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9452 .expect("must load the synthetic Kimi checkpoint end to end");
9453 std::fs::remove_dir_all(&dir).ok();
9454 loaded
9455 }
9456
9457 #[test]
9465 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9466 let loaded = build_synthetic_kimi_loaded();
9467 let state = build_app_state(
9468 StartupModels {
9469 loaded: model::LoadedModel::Kimi(loaded),
9470 embedding: None,
9471 },
9472 None,
9473 None,
9474 None,
9475 false,
9476 None,
9477 Arc::new(health::Detection::ready(health::probe_backends())),
9478 );
9479 let active = state.active().expect("a freshly built state has a model");
9480 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9481 assert!(!active.is_synthetic());
9482
9483 let (_chunks, finish, _usage) = run_generation(
9484 active.generative().unwrap(),
9485 "hi",
9486 &greedy_params(5),
9487 None,
9488 None,
9489 None,
9490 None,
9491 None,
9492 None,
9493 )
9494 .expect("a real Kimi checkpoint must generate without error");
9495 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9496 }
9497
9498 #[test]
9511 fn a_grammar_constrains_the_engine_decode_path() {
9512 let loaded = build_synthetic_kimi_loaded();
9513 let state = build_app_state(
9514 StartupModels {
9515 loaded: model::LoadedModel::Kimi(loaded),
9516 embedding: None,
9517 },
9518 None,
9519 None,
9520 None,
9521 false,
9522 None,
9523 Arc::new(health::Detection::ready(health::probe_backends())),
9524 );
9525 let active = state.active().expect("a freshly built state has a model");
9526
9527 let run = |grammar: Option<&str>| {
9528 let mut params = greedy_params(6);
9529 params.grammar = grammar.map(|src| {
9530 Arc::new(
9531 ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
9532 .expect("test grammar parses"),
9533 )
9534 });
9535 run_generation(
9536 active.generative().unwrap(),
9537 "hi",
9538 ¶ms,
9539 None,
9540 None,
9541 None,
9542 None,
9543 None,
9544 None,
9545 )
9546 };
9547
9548 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
9549 let unconstrained = chunks.concat();
9550 assert!(
9551 unconstrained.chars().any(|c| c != 'a'),
9552 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9553 constrained run below would prove nothing"
9554 );
9555
9556 let (chunks, finish, _) =
9557 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9558 let constrained = chunks.concat();
9559 assert!(
9560 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9561 "the engine decode path served text its grammar forbids ({constrained:?}): \
9562 the constraint was dropped between `generate_engine` and the sampler"
9563 );
9564 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9565 }
9566
9567 #[test]
9577 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9578 let loaded = build_synthetic_kimi_loaded();
9579 let state = build_app_state(
9580 StartupModels {
9581 loaded: model::LoadedModel::Kimi(loaded),
9582 embedding: None,
9583 },
9584 None,
9585 None,
9586 None,
9587 false,
9588 None,
9589 Arc::new(health::Detection::ready(health::probe_backends())),
9590 );
9591
9592 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
9593 let kv_pool_config = generate::KvPoolConfig {
9594 pool: pool.clone(),
9595 queue_wait: Duration::ZERO,
9596 };
9597 let pc = Mutex::new(PrefixCache::new(4));
9598
9599 run_generation(
9600 state
9601 .active()
9602 .expect("a freshly built state has a model")
9603 .generative()
9604 .unwrap(),
9605 "hi",
9606 &greedy_params(5),
9607 Some(&kv_pool_config),
9608 None,
9609 Some(&pc),
9610 None,
9611 None,
9612 None,
9613 )
9614 .expect("a real Kimi checkpoint must generate without error");
9615
9616 assert_eq!(
9617 pool.lock().unwrap().free_blocks(),
9618 4,
9619 "the KV pool must be completely untouched by a Kimi request"
9620 );
9621 let stats = pc.lock().unwrap().stats();
9622 assert_eq!(
9623 stats.hits + stats.misses,
9624 0,
9625 "the prefix cache must never be consulted for a Kimi request"
9626 );
9627 }
9628}