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