1mod admin;
36mod anthropic;
37mod attribution;
38mod budget;
39mod cache_admin;
40mod cancel;
41mod chat_template;
42mod completion;
43mod conversations;
44mod decode_task;
45mod embeddings;
46mod generate;
47mod grammar_request;
48mod health;
49mod journal;
50mod json_mode;
51mod limits;
52mod loaded;
53mod mcp;
54mod model;
55mod openai_extra;
56mod output;
57mod policy;
58mod rerank;
59mod response_cache;
60pub(crate) mod responses;
61mod resume;
62mod sample_step;
63mod sampling_knobs;
64mod security;
65mod serving;
66mod session;
67mod sse;
68mod stats;
69mod stop;
70mod stream_events;
71mod tasks;
72mod tool_grammar;
73mod unsupported_sampling;
74
75use std::cell::RefCell;
76use std::convert::Infallible;
77use std::fmt;
78use std::net::{IpAddr, Ipv4Addr, SocketAddr};
79use std::path::PathBuf;
80use std::rc::Rc;
81use std::str::FromStr;
82use std::sync::{Arc, Mutex, MutexGuard};
83use std::time::Duration;
84
85use axum::{
86 extract::State,
87 http::StatusCode,
88 response::sse::{Event, Sse},
89 response::{IntoResponse, Response},
90 routing::{get, post},
91 Json, Router,
92};
93use clap::{Parser, ValueEnum};
94use serde::{Deserialize, Serialize};
95
96use ferrox_core::cache::KvBlockPool;
97use ferrox_models::kimi_tokenizer::KimiTokenizer;
98use ferrox_models::sampling::SamplingParams;
99use ferrox_models::tokenizer::StopTokens;
100use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
101use generate::{FinishReason, GenerationParams};
102pub(crate) use loaded::{ActiveModel, Loaded};
103use model::ServerTokenizer;
104use rerank::encoder_endpoints;
105use response_cache::{CacheKey, ResponseCache};
106use sampling_knobs::SamplingKnobs;
107
108#[derive(Parser, Debug, PartialEq)]
112#[command(
119 name = "ferrox-server",
120 about = "OpenAI-compatible Ferrox inference server"
121)]
122pub struct ServerArgs {
123 #[arg(short = 'm', long = "model", value_name = "FILE")]
125 model: Option<String>,
126
127 #[arg(long, value_name = "HOST")]
129 host: Option<IpAddr>,
130
131 #[arg(long, value_name = "PORT")]
136 port: Option<u16>,
137
138 #[arg(short = 't', long = "threads", value_name = "N")]
140 threads: Option<usize>,
141
142 #[arg(
144 long = "device",
145 visible_alias = "dev",
146 value_name = "DEVICE",
147 ignore_case = true
148 )]
149 device: Option<OffloadDevice>,
150
151 #[arg(long = "list-devices", default_value_t = false)]
153 list_devices: bool,
154
155 #[arg(
160 long = "n-gpu-layers",
161 visible_aliases = ["gpu-layers", "ngl"],
162 value_name = "N"
163 )]
164 n_gpu_layers: Option<GpuLayers>,
165
166 #[arg(long = "mcp-config", value_name = "PATH")]
168 mcp_config: Option<PathBuf>,
169
170 #[arg(long = "exit-on-stdin-close", default_value_t = false)]
178 exit_on_stdin_close: bool,
179
180 #[arg(
183 long = "cont-batching",
184 visible_aliases = ["continuous-batching", "cb"],
185 default_value_t = false
186 )]
187 cont_batching: bool,
188
189 #[arg(
192 long = "no-cont-batching",
193 default_value_t = false,
194 conflicts_with = "cont_batching"
195 )]
196 no_cont_batching: bool,
197
198 #[arg(long = "parallel", visible_alias = "np", value_name = "N")]
202 parallel: Option<usize>,
203
204 #[arg(long = "allow-multiple-instances", default_value_t = false)]
209 allow_multiple_instances: bool,
210}
211
212impl ServerArgs {
213 pub fn parse_llama_style<I>(argv: I) -> Self
220 where
221 I: IntoIterator<Item = String>,
222 {
223 Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
224 }
225}
226
227#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
228enum OffloadDevice {
229 Auto,
230 None,
231 Cpu,
232 Metal,
233 Cuda,
234}
235
236#[derive(Debug, Clone, Copy, PartialEq, Eq)]
237enum GpuLayers {
238 Auto,
239 All,
240 Count(u32),
241}
242
243impl GpuLayers {
244 fn offload_enabled(self) -> bool {
245 !matches!(self, Self::Count(0))
246 }
247}
248
249impl FromStr for GpuLayers {
250 type Err = String;
251
252 fn from_str(value: &str) -> Result<Self, Self::Err> {
253 match value {
254 "auto" => Ok(Self::Auto),
255 "all" => Ok(Self::All),
256 _ => value
257 .parse::<u32>()
258 .map(Self::Count)
259 .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
260 }
261 }
262}
263
264impl fmt::Display for GpuLayers {
265 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
266 match self {
267 Self::Auto => f.write_str("auto"),
268 Self::All => f.write_str("all"),
269 Self::Count(value) => value.fmt(f),
270 }
271 }
272}
273
274pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
283
284pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
287
288fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
289 args.into_iter()
290 .map(|arg| match arg.as_str() {
291 "-ngl" => "--n-gpu-layers".into(),
292 "-dev" => "--device".into(),
293 "-cb" => "--cont-batching".into(),
294 "-np" => "--parallel".into(),
295 _ => arg,
296 })
297 .collect()
298}
299
300fn print_available_devices() {
301 println!("Available devices:");
302 println!(" CPU");
303
304 let metal = ferrox_metal::MetalProfile::detect();
305 if let Some(name) = metal.device_name {
306 println!(" Metal: {name}");
307 }
308
309 let cuda = ferrox_cuda::HardwareProfile::detect();
310 if cuda.cuda_available {
311 let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
312 println!(" CUDA: {name}");
313 if cuda.cuda_device_count > 1 {
314 println!(" ({} devices detected)", cuda.cuda_device_count);
315 }
316 }
317}
318
319fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
320 if args.host.is_none() && args.port.is_none() {
321 return None;
322 }
323
324 let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
325 let host = args
326 .host
327 .or_else(|| existing.map(|addr| addr.ip()))
328 .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
329 let port = args
330 .port
331 .or_else(|| existing.map(|addr| addr.port()))
332 .unwrap_or(8383);
333 Some(SocketAddr::new(host, port).to_string())
334}
335
336fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
337 if let Some(model) = &args.model {
338 unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
340 }
341
342 if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
343 unsafe { std::env::set_var("FERROX_ADDR", addr) };
345 }
346
347 if let Some(threads) = args.threads {
348 if threads == 0 {
349 anyhow::bail!("--threads must be greater than zero");
350 }
351 unsafe {
353 std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
354 std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
355 }
356 }
357
358 if args.device.is_none() && args.n_gpu_layers.is_none() {
359 } else {
361 let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
362 let device = if layers.offload_enabled() {
363 args.device.unwrap_or(OffloadDevice::Auto)
364 } else {
365 OffloadDevice::None
366 };
367
368 match device {
369 OffloadDevice::None | OffloadDevice::Cpu => unsafe {
370 std::env::set_var("FERROX_METAL", "0");
371 std::env::set_var("FERROX_METAL_ATTN", "0");
372 std::env::set_var("FERROX_CUDA", "0");
373 },
374 OffloadDevice::Auto => unsafe {
375 std::env::set_var("FERROX_METAL", "auto");
376 std::env::set_var("FERROX_CUDA", "auto");
377 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
378 std::env::set_var("FERROX_METAL_ATTN", "1");
379 }
380 },
381 OffloadDevice::Metal => {
382 #[cfg(not(feature = "metal"))]
383 {
384 anyhow::bail!(
385 "Metal requested but this binary was built without --features metal"
386 );
387 }
388 #[cfg(feature = "metal")]
389 {
390 if !ferrox_metal::MetalProfile::detect().available {
391 anyhow::bail!("Metal requested but no Metal device is available");
392 }
393 unsafe {
394 std::env::set_var("FERROX_METAL", "1");
395 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
396 std::env::set_var("FERROX_METAL_ATTN", "1");
397 }
398 std::env::set_var("FERROX_CUDA", "0");
399 }
400 }
401 }
402 OffloadDevice::Cuda => {
403 #[cfg(not(feature = "cuda"))]
404 {
405 anyhow::bail!(
406 "CUDA requested but this binary was built without --features cuda"
407 );
408 }
409 #[cfg(feature = "cuda")]
410 {
411 if !ferrox_cuda::HardwareProfile::detect().cuda_available {
412 anyhow::bail!("CUDA requested but no CUDA device is available");
413 }
414 unsafe {
415 std::env::set_var("FERROX_CUDA", "1");
416 std::env::set_var("FERROX_METAL", "0");
417 std::env::set_var("FERROX_METAL_ATTN", "0");
418 }
419 }
420 }
421 }
422 }
423
424 if let Some(n) = args.parallel {
425 if n == 0 {
426 anyhow::bail!("--parallel must be greater than zero");
427 }
428 unsafe { std::env::set_var("FERROX_CB_MAX_SEQS", n.to_string()) };
430 }
431
432 if args.cont_batching {
433 unsafe { std::env::set_var("FERROX_CONTINUOUS_BATCHING", "1") };
435 } else if args.no_cont_batching {
436 unsafe { std::env::set_var("FERROX_CONTINUOUS_BATCHING", "0") };
438 } else if args.parallel.is_some() {
439 unsafe { std::env::set_var("FERROX_CONTINUOUS_BATCHING", "1") };
442 }
443
444 Ok(())
445}
446
447#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
456 Gguf(GgufModel),
457 Kimi(KimiModel),
458 Mla(MlaModel),
459 Gemma4(Gemma4Model),
460 Glm52(Glm52Model),
461}
462
463pub(crate) struct GgufModel {
464 decoder: Arc<Decoder>,
465 tokenizer: Arc<ServerTokenizer>,
466 stop_tokens: StopTokens,
467 bos_id: Option<usize>,
468 is_synthetic: bool,
469 chat_template: chat_template::PromptTemplate,
470}
471
472pub(crate) struct KimiModel {
473 engine: KimiEngine,
474 tokenizer: KimiTokenizer,
475 stop_tokens: StopTokens,
476 chat_template: chat_template::PromptTemplate,
477}
478
479pub(crate) struct MlaModel {
480 engine: MlaEngine,
481 tokenizer: ServerTokenizer,
482 stop_tokens: StopTokens,
483 bos_id: Option<usize>,
484 name: String,
485 chat_template: chat_template::PromptTemplate,
486}
487
488pub(crate) struct Gemma4Model {
489 engine: Gemma4Engine,
490 tokenizer: ServerTokenizer,
491 stop_tokens: StopTokens,
492 bos_id: Option<usize>,
493 name: String,
494 chat_template: chat_template::PromptTemplate,
495}
496
497pub(crate) struct Glm52Model {
498 engine: ferrox_models::Glm52Engine,
499 tokenizer: ServerTokenizer,
500 stop_tokens: StopTokens,
501 bos_id: Option<usize>,
502 name: String,
503 chat_template: chat_template::PromptTemplate,
504}
505
506impl Model {
507 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
508 match self {
509 Model::Gguf(m) => m.chat_template.clone(),
510 Model::Kimi(m) => m.chat_template.clone(),
511 Model::Mla(m) => m.chat_template.clone(),
512 Model::Gemma4(m) => m.chat_template.clone(),
513 Model::Glm52(m) => m.chat_template.clone(),
514 }
515 }
516
517 fn is_synthetic(&self) -> bool {
522 match self {
523 Model::Gguf(m) => m.is_synthetic,
524 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
525 }
526 }
527
528 fn tokenizer_kind(&self) -> &'static str {
529 match self {
530 Model::Gguf(m) => m.tokenizer.kind(),
531 Model::Kimi(_) => "kimi-tiktoken-bpe",
532 Model::Mla(m) => m.tokenizer.kind(),
533 Model::Gemma4(m) => m.tokenizer.kind(),
534 Model::Glm52(m) => m.tokenizer.kind(),
535 }
536 }
537
538 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
542 match self {
543 Model::Gguf(m) => m.decoder.expert_store_stats(),
544 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
545 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
546 }
547 }
548
549 pub(crate) fn name(&self) -> &str {
550 match self {
551 Model::Gguf(m) => m.decoder.config.name,
552 Model::Kimi(_) => "kimi-k3",
553 Model::Mla(m) => m.name.as_str(),
554 Model::Gemma4(m) => m.name.as_str(),
555 Model::Glm52(m) => m.name.as_str(),
556 }
557 }
558
559 pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
560 match self {
561 Model::Gguf(m) => m.tokenizer.encode(text),
562 Model::Kimi(m) => m
563 .tokenizer
564 .encode(text)
565 .into_iter()
566 .map(|id| id as usize)
567 .collect(),
568 Model::Mla(m) => m.tokenizer.encode(text),
569 Model::Gemma4(m) => m.tokenizer.encode(text),
570 Model::Glm52(m) => m.tokenizer.encode(text),
571 }
572 }
573
574 pub(crate) fn bos_id(&self) -> Option<usize> {
583 match self {
584 Model::Gguf(m) => m.bos_id,
585 Model::Kimi(_) => None,
586 Model::Mla(m) => m.bos_id,
587 Model::Gemma4(m) => m.bos_id,
588 Model::Glm52(m) => m.bos_id,
589 }
590 }
591
592 pub(crate) fn decode(&self, ids: &[usize]) -> String {
593 match self {
594 Model::Gguf(m) => m.tokenizer.decode(ids),
595 Model::Kimi(m) => {
596 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
597 m.tokenizer.decode(&ids32)
598 }
599 Model::Mla(m) => m.tokenizer.decode(ids),
600 Model::Gemma4(m) => m.tokenizer.decode(ids),
601 Model::Glm52(m) => m.tokenizer.decode(ids),
602 }
603 }
604
605 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
608 match self {
609 Model::Gguf(m) => {
610 let mut caches: Vec<_> = (0..m.decoder.layers.len())
611 .map(|_| {
612 ferrox_core::cache::KvCache::new(
613 m.decoder.config.n_kv_heads,
614 m.decoder.config.head_dim,
615 )
616 })
617 .collect();
618 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
619 }
620 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
621 }
622 }
623
624 pub(crate) fn vocab_size(&self) -> Option<usize> {
625 match self {
626 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
627 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
628 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
629 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
630 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
631 }
632 }
633}
634
635pub(crate) struct AppState {
636 pub(crate) embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
650 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
669 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
673 pub(crate) tasks: Arc<tasks::TaskRegistry>,
675 pub(crate) cancels: Arc<cancel::CancelRegistry>,
679 pub(crate) stats: stats::Stats,
682 pub(crate) streams: resume::StreamRegistry,
685 pub(crate) model_dir: Option<PathBuf>,
687 response_cache: Mutex<ResponseCache>,
691 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
703 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
716 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
726 sessions: session::SessionStore,
732 requests_total: std::sync::atomic::AtomicU64,
733 request_errors_total: std::sync::atomic::AtomicU64,
734 started_at: std::time::Instant,
735 last_request_ms: std::sync::atomic::AtomicU64,
740 detection: Arc<health::Detection>,
742 mcp: Option<mcp::LoadedMcpConfig>,
744 pub(crate) continuous_batching_enabled: bool,
748 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
752 loading_model: Mutex<Option<String>>,
756 last_load_error: Mutex<Option<(String, String)>>,
760 pub(crate) serving: Mutex<crate::stats::ServingStats>,
765 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
770 pub(crate) footprint:
776 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
777 pub(crate) started_unix: u64,
787}
788
789pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
795
796impl AppState {
797 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
801 self.active
802 .read()
803 .unwrap_or_else(|p| p.into_inner())
804 .clone()
805 }
806
807 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
813 self.active().ok_or_else(|| {
814 (
815 StatusCode::SERVICE_UNAVAILABLE,
816 Json(serde_json::json!({"error": {
817 "message": "no model is loaded; POST /admin/models/load with an id from \
818 GET /admin/models",
819 "type": "model_not_loaded"
820 }})),
821 )
822 })
823 }
824
825 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
833 Ok(Arc::clone(self.require_active()?.generative()?))
834 }
835
836 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
844 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
845 std::mem::replace(&mut *guard, next)
846 }
847
848 fn mark_request_finished(&self) {
852 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
853 self.last_request_ms
854 .store(ms, std::sync::atomic::Ordering::Relaxed);
855 }
856
857 pub(crate) fn uptime(&self) -> Duration {
858 self.started_at.elapsed()
859 }
860
861 pub(crate) fn requests_total(&self) -> u64 {
862 self.requests_total
863 .load(std::sync::atomic::Ordering::Relaxed)
864 }
865
866 pub(crate) fn errors_total(&self) -> u64 {
867 self.request_errors_total
868 .load(std::sync::atomic::Ordering::Relaxed)
869 }
870
871 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
872 lock_cache(&self.response_cache).stats()
873 }
874
875 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
878 let last = self
879 .last_request_ms
880 .load(std::sync::atomic::Ordering::Relaxed);
881 (last > 0)
882 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
883 .map(|age| age.max(0.0))
884 }
885
886 pub(crate) fn loading_model_id(&self) -> Option<String> {
887 self.loading_model
888 .lock()
889 .unwrap_or_else(|p| p.into_inner())
890 .clone()
891 }
892
893 pub(crate) fn set_loading_model(&self, id: Option<String>) {
894 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
895 }
896
897 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
898 self.last_load_error
899 .lock()
900 .unwrap_or_else(|p| p.into_inner())
901 .clone()
902 }
903
904 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
905 *self
906 .last_load_error
907 .lock()
908 .unwrap_or_else(|p| p.into_inner()) = error;
909 }
910
911 pub(crate) fn active_model_name(&self) -> Option<String> {
920 self.active().map(|a| a.name().to_string())
921 }
922
923 pub(crate) fn embedding_model(&self) -> Option<Arc<ferrox_models::EmbeddingModel>> {
933 self.embedding
934 .clone()
935 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
936 }
937
938 pub(crate) fn embedding_model_name(&self) -> Option<String> {
942 match self.embedding_model() {
943 Some(e) => Some(e.name().to_string()),
944 None => self.active_model_name(),
945 }
946 }
947
948 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
949 self.stats.record(stats::entry(record));
950 }
951}
952
953fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
958 cache
959 .lock()
960 .unwrap_or_else(|poisoned| poisoned.into_inner())
961}
962
963#[derive(Debug, Clone, Deserialize)]
964#[serde(untagged)]
965pub(crate) enum MessageContent {
966 Text(String),
967 Parts(Vec<ContentPart>),
968}
969
970#[derive(Debug, Clone, Deserialize)]
971struct ContentPart {
972 #[serde(rename = "type")]
973 kind: String,
974 #[serde(default)]
975 text: Option<String>,
976 #[serde(default)]
977 image_url: Option<serde_json::Value>,
978}
979
980impl MessageContent {
981 fn as_text(&self) -> String {
982 match self {
983 Self::Text(s) => s.clone(),
984 Self::Parts(parts) => parts
985 .iter()
986 .filter_map(|p| p.text.as_deref())
987 .collect::<Vec<_>>()
988 .join(""),
989 }
990 }
991
992 fn has_image(&self) -> bool {
993 match self {
994 Self::Text(_) => false,
995 Self::Parts(parts) => parts
996 .iter()
997 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
998 }
999 }
1000}
1001
1002#[derive(Debug, Clone, Deserialize)]
1003pub(crate) struct ChatMessage {
1004 pub(crate) role: String,
1005 #[serde(default)]
1009 pub(crate) content: Option<MessageContent>,
1010 #[serde(default)]
1014 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
1015 #[serde(default)]
1020 #[allow(dead_code)]
1021 pub(crate) tool_call_id: Option<String>,
1022 #[serde(default, alias = "reasoning")]
1038 pub(crate) reasoning_content: Option<String>,
1039}
1040
1041impl ChatMessage {
1042 fn rendered_content(&self) -> String {
1049 let mut out = self
1050 .content
1051 .as_ref()
1052 .map(MessageContent::as_text)
1053 .unwrap_or_default();
1054 if let Some(calls) = &self.tool_calls {
1055 for call in calls {
1056 out.push_str(&format!(
1057 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
1058 call.function.name, call.function.arguments
1059 ));
1060 }
1061 }
1062 out
1063 }
1064}
1065
1066#[derive(Debug, Clone, Deserialize)]
1067pub(crate) struct ToolCallIn {
1068 #[serde(default)]
1069 #[allow(dead_code)]
1070 id: String,
1071 #[serde(rename = "type", default)]
1072 #[allow(dead_code)]
1073 kind: String,
1074 function: ToolCallFunctionIn,
1075}
1076
1077#[derive(Debug, Clone, Deserialize)]
1078struct ToolCallFunctionIn {
1079 name: String,
1080 arguments: String,
1085}
1086
1087#[derive(Debug, Clone, Deserialize)]
1090struct ToolDef {
1091 #[serde(rename = "type", default)]
1092 #[allow(dead_code)]
1093 kind: String,
1094 function: ToolFunctionDef,
1095}
1096
1097#[derive(Debug, Clone, Deserialize)]
1098struct ToolFunctionDef {
1099 name: String,
1100 #[serde(default)]
1101 description: Option<String>,
1102 #[serde(default)]
1103 parameters: Option<serde_json::Value>,
1104}
1105
1106#[derive(Debug, Clone, Deserialize)]
1117#[serde(untagged)]
1118enum ToolChoice {
1119 Mode(String),
1120 Specific(serde_json::Value),
1121}
1122
1123#[derive(Deserialize)]
1126#[serde(untagged)]
1127enum StopParam {
1128 One(String),
1129 Many(Vec<String>),
1130}
1131
1132#[derive(Deserialize)]
1133struct ChatCompletionRequest {
1134 model: String,
1135 messages: Vec<ChatMessage>,
1136 #[serde(default = "default_max_tokens")]
1137 max_tokens: usize,
1138 #[serde(default)]
1139 temperature: Option<f32>,
1140 #[serde(default)]
1141 top_p: Option<f32>,
1142 #[serde(default)]
1147 min_p: Option<f32>,
1148 #[serde(default)]
1149 top_k: Option<usize>,
1150 #[serde(default)]
1151 repetition_penalty: Option<f32>,
1152 #[serde(default)]
1153 seed: Option<u64>,
1154 #[serde(default)]
1155 stop: Option<StopParam>,
1156 #[serde(default)]
1157 stream: Option<bool>,
1158 #[serde(default)]
1171 stream_resumable: Option<bool>,
1172 #[serde(default)]
1182 ignore_eos: Option<bool>,
1183 #[serde(default)]
1184 tools: Vec<ToolDef>,
1185 #[serde(default)]
1186 tool_choice: Option<ToolChoice>,
1187 #[serde(default)]
1195 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
1196 #[serde(default)]
1204 reasoning_effort: Option<String>,
1205 #[serde(default)]
1209 thinking: Option<ThinkingSwitch>,
1210 #[serde(default)]
1215 session_id: Option<String>,
1216 #[serde(default)]
1218 logprobs: Option<bool>,
1219 #[serde(default)]
1220 top_logprobs: Option<u32>,
1221 #[serde(default)]
1222 n: Option<u32>,
1223 #[serde(default)]
1224 presence_penalty: Option<f32>,
1225 #[serde(default)]
1226 frequency_penalty: Option<f32>,
1227 #[serde(default)]
1228 response_format: Option<serde_json::Value>,
1229 #[serde(default)]
1236 logit_bias: Option<serde_json::Value>,
1237 #[serde(default)]
1247 grammar: Option<String>,
1248}
1249
1250const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1262
1263#[derive(Debug, Clone, Deserialize)]
1265pub(crate) struct ThinkingSwitch {
1266 #[serde(rename = "type")]
1267 pub(crate) kind: String,
1268}
1269
1270const THINKING_KWARG_KEYS: [&str; 4] = [
1274 "enable_thinking",
1275 "thinking",
1276 "thinking_mode",
1277 "reasoning_effort",
1278];
1279
1280const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1284
1285fn default_max_tokens() -> usize {
1286 DEFAULT_CHAT_MAX_TOKENS
1287}
1288
1289impl ChatCompletionRequest {
1290 fn sampling_knobs(&self) -> SamplingKnobs {
1295 SamplingKnobs {
1296 temperature: self.temperature,
1297 top_p: self.top_p,
1298 min_p: self.min_p,
1299 top_k: self.top_k,
1300 repetition_penalty: self.repetition_penalty,
1301 presence_penalty: self.presence_penalty,
1302 frequency_penalty: self.frequency_penalty,
1303 penalty_last_n: None,
1307 }
1308 }
1309
1310 fn sampling_params(&self) -> SamplingParams {
1311 self.sampling_knobs().resolve()
1312 }
1313
1314 fn stop_sequences(&self) -> Vec<String> {
1315 self.stop
1316 .as_ref()
1317 .map(|s| match s {
1318 StopParam::One(v) => vec![v.clone()],
1319 StopParam::Many(v) => v.clone(),
1320 })
1321 .unwrap_or_default()
1322 }
1323
1324 fn tools_active(&self) -> bool {
1329 !self.tools.is_empty()
1330 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1331 }
1332
1333 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1345 match &self.tool_choice {
1346 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1347 Some(ToolChoice::Specific(value)) => {
1348 let name = value
1352 .get("function")
1353 .and_then(|f| f.get("name"))
1354 .or_else(|| value.get("name"))
1355 .and_then(|n| n.as_str());
1356 match name {
1357 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1358 None => Err(invalid_request(
1359 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1360 function.name",
1361 "tool_choice",
1362 )),
1363 }
1364 }
1365 _ => Ok(None),
1366 }
1367 }
1368
1369 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1371 self.tools
1372 .iter()
1373 .map(|t| tool_grammar::ToolSpec {
1374 name: &t.function.name,
1375 parameters: t.function.parameters.as_ref(),
1376 })
1377 .collect()
1378 }
1379
1380 fn thinking_direction(&self) -> Option<bool> {
1423 if let Some(switch) = &self.thinking {
1424 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1425 "disabled" => Some(false),
1426 "enabled" => Some(true),
1427 _ => None,
1430 };
1431 }
1432 let effort = self.reasoning_effort.as_ref()?;
1433 DISABLE_EFFORTS
1434 .contains(&effort.trim().to_ascii_lowercase().as_str())
1435 .then_some(false)
1436 }
1437
1438 fn resolve_template_kwargs(
1439 &self,
1440 template: &chat_template::PromptTemplate,
1441 ) -> serde_json::Map<String, serde_json::Value> {
1442 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1443 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1447
1448 if !caller_steered {
1449 match self.thinking_direction() {
1450 Some(false) => {
1451 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1452 kwargs.insert(k, v);
1453 }
1454 return kwargs;
1457 }
1458 Some(true) => {
1459 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1460 kwargs.insert(k, v);
1461 }
1462 }
1463 None => {}
1464 }
1465 if let Some(effort) = &self.reasoning_effort {
1466 kwargs
1467 .entry("reasoning_effort".to_string())
1468 .or_insert_with(|| serde_json::json!(effort));
1469 }
1470 }
1471
1472 let offered: Vec<serde_json::Value> = if self.tools_active() {
1473 self.tools.iter().map(chat_template::tool_json).collect()
1474 } else {
1475 Vec::new()
1476 };
1477 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1478 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1479 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1480 kwargs.entry(k).or_insert(v);
1481 }
1482 }
1483 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1484 crate::policy::effort::EffortMapping::Mapped(to) => {
1485 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1486 }
1487 crate::policy::effort::EffortMapping::Dropped => {
1488 tracing::debug!(
1489 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1490 enough, so its own default applies"
1491 );
1492 }
1493 crate::policy::effort::EffortMapping::Unchanged => {}
1494 }
1495 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1496 kwargs
1497 }
1498
1499 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1502 if self.max_tokens == 0 {
1509 return Err(invalid_request(
1510 "max_tokens must be at least 1",
1511 "max_tokens",
1512 ));
1513 }
1514 if let Some(switch) = &self.thinking {
1519 let kind = switch.kind.trim().to_ascii_lowercase();
1520 if kind != "enabled" && kind != "disabled" {
1521 return Err(invalid_request(
1522 "thinking.type must be \"enabled\" or \"disabled\"",
1523 "thinking.type",
1524 ));
1525 }
1526 }
1527 for msg in &self.messages {
1528 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1529 return Err(unsupported_feature(
1530 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1531 ));
1532 }
1533 }
1534 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1535 return Err(unsupported_feature(
1536 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1537 ));
1538 }
1539 if self.n.is_some_and(|n| n > 1) {
1540 return Err(unsupported_feature(
1541 "n > 1 is not implemented (single completion only)",
1542 ));
1543 }
1544 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1545 let stated_grammar =
1560 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1561 if let Some(forced) = self.forced_tool_choice()? {
1569 if self.tools.is_empty() {
1570 return Err(invalid_request(
1571 "tool_choice forces a tool call, but no tools were offered",
1572 "tool_choice",
1573 ));
1574 }
1575 if let tool_grammar::Forced::Named(name) = forced {
1576 if !self.tools.iter().any(|t| t.function.name == name) {
1577 return Err(invalid_request(
1578 &format!(
1579 "tool_choice names {name:?}, which is not one of the tools offered"
1580 ),
1581 "tool_choice",
1582 ));
1583 }
1584 }
1585 if stated_grammar.is_some() {
1595 return Err(invalid_request(
1596 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1597 are two different constraints on the same generation; send one",
1598 "tool_choice",
1599 ));
1600 }
1601 if self.json_object_mode() {
1602 return Err(invalid_request(
1603 "a forced tool_choice cannot be combined with response_format json_object: \
1604 the tool-call markers are not JSON",
1605 "tool_choice",
1606 ));
1607 }
1608 }
1609 Ok(())
1610 }
1611
1612 fn effective_stop_sequences(&self) -> Vec<String> {
1619 let mut stop = self.stop_sequences();
1620 if self.tools_active() {
1621 stop.push("</tool_call>".to_string());
1622 }
1623 stop
1624 }
1625
1626 fn json_object_mode(&self) -> bool {
1627 self.response_format
1628 .as_ref()
1629 .and_then(|v| v.get("type"))
1630 .and_then(|v| v.as_str())
1631 == Some("json_object")
1632 }
1633
1634 fn generation_params(&self) -> Result<GenerationParams, ApiError> {
1639 Ok(GenerationParams {
1640 max_tokens: self.max_tokens,
1641 sampling: self.sampling_params(),
1642 seed: self.resolved_seed(),
1643 stop: self.effective_stop_sequences(),
1644 stop_token_ids: Vec::new(),
1648 json_object: self.json_object_mode(),
1649 grammar: grammar_request::for_request(
1650 self.grammar.as_deref(),
1651 self.response_format.as_ref(),
1652 )?,
1653 cancel: None,
1656 ignore_eos: self.ignore_eos.unwrap_or(false),
1657 })
1658 }
1659
1660 fn generation_params_for_template(
1671 &self,
1672 template: &chat_template::PromptTemplate,
1673 served_model: &str,
1674 ) -> Result<GenerationParams, ApiError> {
1675 let mut params = self.generation_params()?;
1676 if let Some(stop) = template.end_of_turn() {
1677 if !params.stop.iter().any(|s| s == stop) {
1678 params.stop.push(stop.to_string());
1679 }
1680 }
1681 if let Some(forced) = self.forced_tool_choice()? {
1682 params.grammar = Some(tool_grammar::build(
1686 forced,
1687 &self.tool_specs(),
1688 policy::parser::ToolCallFormat::infer(served_model),
1689 )?);
1690 }
1691 Ok(params)
1692 }
1693
1694 fn is_cacheable(&self) -> bool {
1703 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1704 }
1705
1706 fn cache_key(&self, prompt: &str) -> CacheKey {
1707 CacheKey {
1708 model: self.model.clone(),
1709 prompt: prompt.to_string(),
1710 max_tokens: self.max_tokens,
1711 sampling: response_cache::sampling_key(&self.sampling_params()),
1716 seed: self.seed,
1717 stop: self.effective_stop_sequences(),
1718 }
1719 }
1720
1721 fn resolved_seed(&self) -> u64 {
1722 self.seed.unwrap_or_else(|| {
1723 std::time::SystemTime::now()
1724 .duration_since(std::time::UNIX_EPOCH)
1725 .map(|d| d.as_nanos() as u64)
1726 .unwrap_or(0xDEFA017)
1727 })
1728 }
1729}
1730
1731#[derive(Serialize)]
1732struct ChatCompletionChoice {
1733 index: usize,
1734 message: ChatCompletionResponseMessage,
1735 finish_reason: &'static str,
1736}
1737
1738#[derive(Serialize)]
1739struct ChatCompletionResponseMessage {
1740 role: &'static str,
1741 #[serde(skip_serializing_if = "Option::is_none")]
1742 content: Option<String>,
1743 #[serde(skip_serializing_if = "Option::is_none")]
1747 reasoning_content: Option<String>,
1748 #[serde(skip_serializing_if = "Option::is_none")]
1749 tool_calls: Option<Vec<ToolCallOut>>,
1750}
1751
1752#[derive(Serialize, Clone)]
1753struct ToolCallOut {
1754 id: String,
1755 #[serde(rename = "type")]
1756 kind: &'static str,
1757 function: ToolCallFunctionOut,
1758}
1759
1760#[derive(Serialize, Clone)]
1768struct ToolCallDelta {
1769 index: usize,
1770 #[serde(skip_serializing_if = "Option::is_none")]
1771 id: Option<String>,
1772 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1773 kind: Option<&'static str>,
1774 function: ToolCallFunctionDelta,
1775}
1776
1777#[derive(Serialize, Clone, Default)]
1778struct ToolCallFunctionDelta {
1779 #[serde(skip_serializing_if = "Option::is_none")]
1780 name: Option<String>,
1781 #[serde(skip_serializing_if = "Option::is_none")]
1784 arguments: Option<String>,
1785}
1786
1787impl ToolCallDelta {
1788 fn whole(index: usize, name: String, arguments: String) -> Self {
1790 ToolCallDelta {
1791 index,
1792 id: Some(format!("call_{index}")),
1793 kind: Some("function"),
1794 function: ToolCallFunctionDelta {
1795 name: Some(name),
1796 arguments: Some(arguments),
1797 },
1798 }
1799 }
1800
1801 fn opening(index: usize, name: String) -> Self {
1803 ToolCallDelta {
1804 index,
1805 id: Some(format!("call_{index}")),
1806 kind: Some("function"),
1807 function: ToolCallFunctionDelta {
1808 name: Some(name),
1809 arguments: Some(String::new()),
1810 },
1811 }
1812 }
1813
1814 fn arguments(index: usize, fragment: String) -> Self {
1816 ToolCallDelta {
1817 index,
1818 id: None,
1819 kind: None,
1820 function: ToolCallFunctionDelta {
1821 name: None,
1822 arguments: Some(fragment),
1823 },
1824 }
1825 }
1826}
1827
1828#[derive(Serialize, Clone)]
1829struct ToolCallFunctionOut {
1830 name: String,
1831 arguments: String,
1835}
1836
1837#[derive(Serialize)]
1838struct ChatCompletionResponse {
1839 id: String,
1840 request_id: String,
1846 object: &'static str,
1847 model: String,
1848 choices: Vec<ChatCompletionChoice>,
1849 usage: generate::Usage,
1854 ferrox_cache: &'static str,
1862}
1863
1864#[derive(Serialize)]
1865struct ChatCompletionChunkDelta {
1866 #[serde(skip_serializing_if = "Option::is_none")]
1867 role: Option<&'static str>,
1868 #[serde(skip_serializing_if = "Option::is_none")]
1869 content: Option<String>,
1870 #[serde(skip_serializing_if = "Option::is_none")]
1872 reasoning_content: Option<String>,
1873 #[serde(skip_serializing_if = "Option::is_none")]
1874 tool_calls: Option<Vec<ToolCallDelta>>,
1875}
1876
1877#[derive(Serialize)]
1878struct ChatCompletionChunkChoice {
1879 index: usize,
1880 delta: ChatCompletionChunkDelta,
1881 finish_reason: Option<&'static str>,
1882}
1883
1884#[derive(Serialize)]
1885struct ChatCompletionChunk {
1886 id: String,
1887 #[serde(skip_serializing_if = "Option::is_none")]
1894 request_id: Option<String>,
1895 object: &'static str,
1896 model: String,
1897 choices: Vec<ChatCompletionChunkChoice>,
1898 #[serde(skip_serializing_if = "Option::is_none")]
1901 usage: Option<generate::Usage>,
1902}
1903
1904async fn health(State(state): State<Arc<AppState>>) -> Response {
1910 let snapshot = state.detection.snapshot();
1911 let mut capabilities = snapshot.capabilities;
1912 let active = state.active();
1913
1914 capabilities.push(match active.as_deref() {
1917 None => ferrox_api::Capability::unavailable(
1922 ferrox_api::health::capability::REAL_WEIGHTS,
1923 ferrox_api::health::reason::MODEL_NOT_LOADED,
1924 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1925 ),
1926 Some(active) if active.is_synthetic() => ferrox_api::Capability::unavailable(
1927 ferrox_api::health::capability::REAL_WEIGHTS,
1928 ferrox_api::health::reason::MODEL_NOT_LOADED,
1929 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1930 checkpoint. Output from this model is noise.",
1931 ),
1932 Some(active) if active.encoder().is_some() => {
1942 let endpoints = active
1943 .encoder()
1944 .map(|e| encoder_endpoints(e))
1945 .unwrap_or_default();
1946 let served_by = match endpoints.is_empty() {
1947 true => "no endpoint in this build serves it".to_string(),
1948 false => format!("served by {}", endpoints.join(" and ")),
1949 };
1950 ferrox_api::Capability::available(
1951 ferrox_api::health::capability::REAL_WEIGHTS,
1952 format!(
1953 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1954 {served_by}; generation endpoints refuse it.",
1955 active.name(),
1956 ),
1957 )
1958 }
1959 Some(active) => ferrox_api::Capability::available(
1960 ferrox_api::health::capability::REAL_WEIGHTS,
1961 format!("Serving the real checkpoint '{}'.", active.name()),
1962 ),
1963 });
1964 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1965 ferrox_api::Capability::available(
1966 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1967 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1968 "On by default on Metal. Concurrent requests share one batched decode worker."
1969 } else {
1970 "Concurrent requests share one batched decode step."
1971 },
1972 )
1973 } else if state.metal_private_decode_gate.is_some() {
1974 ferrox_api::Capability::unavailable(
1975 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1976 ferrox_api::health::reason::DISABLED,
1977 "Off; private Metal decodes serialize (one at a time). Set FERROX_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1978 )
1979 } else {
1980 ferrox_api::Capability::unavailable(
1981 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1982 ferrox_api::health::reason::DISABLED,
1983 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1984 )
1985 });
1986
1987 let last_request_ms = state
1988 .last_request_ms
1989 .load(std::sync::atomic::Ordering::Relaxed);
1990 let uptime = state.started_at.elapsed();
1991 let health_state = if active.is_none() {
1999 ferrox_api::HealthState::Unavailable
2000 } else {
2001 snapshot.state
2002 };
2003 let body = ferrox_api::HealthResponse {
2004 state: health_state,
2005 reason: match health_state {
2006 ferrox_api::HealthState::Ready => None,
2007 ferrox_api::HealthState::Unavailable => {
2008 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
2009 }
2010 ferrox_api::HealthState::Detecting => {
2011 Some(ferrox_api::health::reason::DETECTING.to_string())
2012 }
2013 },
2014 detail: match health_state {
2015 ferrox_api::HealthState::Ready => None,
2016 ferrox_api::HealthState::Unavailable => Some(
2017 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
2018 .to_string(),
2019 ),
2020 ferrox_api::HealthState::Detecting => {
2021 Some("Probing available compute backends.".to_string())
2022 }
2023 },
2024 model: active
2025 .as_deref()
2026 .map(|active| ferrox_api::health::ModelSummary {
2027 id: active.name().to_string(),
2028 tokenizer: active.tokenizer_kind().to_string(),
2029 synthetic_weights: active.is_synthetic(),
2030 }),
2031 capabilities,
2032 version: env!("CARGO_PKG_VERSION").to_string(),
2033 pid: std::process::id(),
2034 uptime_seconds: uptime.as_secs_f64(),
2035 server_time_unix_ms: std::time::SystemTime::now()
2036 .duration_since(std::time::UNIX_EPOCH)
2037 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
2038 .unwrap_or(0),
2039 last_request_age_seconds: (last_request_ms > 0)
2040 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
2041 .map(|age| age.max(0.0)),
2042 };
2043
2044 let status =
2045 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
2046 (status, Json(body)).into_response()
2047}
2048
2049async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2050 let Some(active) = state.active() else {
2054 return Json(serde_json::json!({ "object": "list", "data": [] }));
2055 };
2056 let mut model_entry = serde_json::json!({
2057 "id": active.name(),
2058 "object": "model",
2059 "ferrox_synthetic_weights": active.is_synthetic(),
2060 "ferrox_tokenizer": active.tokenizer_kind(),
2061 });
2062 if let Some(encoder) = active.encoder() {
2068 model_entry["ferrox_model_kind"] = serde_json::json!("embedding");
2069 model_entry["ferrox_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
2070 model_entry["ferrox_n_embd"] = serde_json::json!(encoder.n_embd());
2071 model_entry["ferrox_pooling"] = serde_json::json!(encoder.pooling_type().name());
2072 model_entry["ferrox_context_length"] = serde_json::json!(encoder.n_ctx_train());
2073 }
2074 if let Some(model) = active.generative_opt() {
2082 let parser_configured =
2083 crate::policy::parser::ReasoningFormat::infer(active.name()).is_some();
2084 let gears = model.chat_template().think_gears(parser_configured);
2085 if !gears.is_empty() {
2086 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
2087 if let Some(default) = &gears.default {
2088 model_entry["default_reasoning_effort"] = serde_json::json!(default);
2089 }
2090 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
2093 }
2094 }
2095 if let Some(mcp) = &state.mcp {
2096 model_entry["ferrox_mcp"] = mcp.models_metadata();
2097 }
2098 Json(serde_json::json!({
2099 "object": "list",
2100 "data": [model_entry]
2101 }))
2102}
2103
2104async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2118 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
2119 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
2120 let active = state.active();
2121 Json(serde_json::json!({
2122 "model": active.as_ref().map(|a| a.name()),
2123 "state": state
2124 .maintenance
2125 .lock()
2126 .unwrap_or_else(|p| p.into_inner())
2127 .state()
2128 .as_str(),
2129 "uptime_s": state.uptime().as_secs(),
2130 "throughput": {
2131 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2132 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2133 },
2134 "requests": {
2135 "active": state.cancels.live_count(),
2136 "completed": state.stats.recorded_total(),
2137 "p95_ms": state.stats.p95_duration_ms(),
2138 "ttft_mean_ms": state.stats.ttft_mean_ms(),
2139 "prompt_tokens_total": state.stats.tokens_prompt_total(),
2140 "completion_tokens_total": state.stats.tokens_generated_total(),
2141 },
2142 "pools": cache_admin::pool_gauges(&state),
2148 "memory": cache_admin::footprint_json(&state),
2151 }))
2152}
2153
2154#[derive(Deserialize)]
2155struct RequestsQuery {
2156 #[serde(default)]
2157 since: u64,
2158 #[serde(default = "default_requests_limit")]
2159 limit: usize,
2160}
2161
2162fn default_requests_limit() -> usize {
2163 stats::MAX_PAGE
2164}
2165
2166async fn recent_requests(
2174 State(state): State<Arc<AppState>>,
2175 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
2176) -> Json<serde_json::Value> {
2177 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
2178 Json(serde_json::json!({
2179 "requests": rows,
2180 "next_cursor": cursor,
2181 "missed": missed,
2182 "total": state.stats.recorded_total(),
2183 }))
2184}
2185
2186#[derive(Serialize)]
2187struct CombinedCacheStats {
2188 response_cache: response_cache::CacheStats,
2189 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
2191}
2192
2193async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
2194 Json(CombinedCacheStats {
2195 response_cache: lock_cache(&state.response_cache).stats(),
2196 prefix_cache: state
2197 .prefix_cache
2198 .as_ref()
2199 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
2200 })
2201}
2202
2203async fn metrics(State(state): State<Arc<AppState>>) -> Response {
2208 use std::sync::atomic::Ordering;
2209
2210 let cache_stats = lock_cache(&state.response_cache).stats();
2211 let active = state.active();
2212 let requests_total = state.requests_total.load(Ordering::Relaxed);
2213 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2214 let uptime = state.started_at.elapsed().as_secs_f64();
2215
2216 let body = format!(
2217 "# HELP ferrox_requests_total Total chat completion requests received.\n\
2218 # TYPE ferrox_requests_total counter\n\
2219 ferrox_requests_total {requests_total}\n\
2220 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
2221 # TYPE ferrox_request_errors_total counter\n\
2222 ferrox_request_errors_total {errors_total}\n\
2223 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
2224 # TYPE ferrox_cache_hits_total counter\n\
2225 ferrox_cache_hits_total {}\n\
2226 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
2227 # TYPE ferrox_cache_misses_total counter\n\
2228 ferrox_cache_misses_total {}\n\
2229 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
2230 # TYPE ferrox_cache_entries gauge\n\
2231 ferrox_cache_entries {}\n\
2232 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2233 # TYPE ferrox_synthetic_weights gauge\n\
2234 ferrox_synthetic_weights {}\n\
2235 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
2236 # TYPE ferrox_uptime_seconds gauge\n\
2237 ferrox_uptime_seconds {uptime}\n",
2238 cache_stats.hits,
2239 cache_stats.misses,
2240 cache_stats.entries,
2241 active
2245 .as_ref()
2246 .map(|a| a.is_synthetic() as u8)
2247 .unwrap_or(0),
2248 );
2249
2250 let body = match active
2254 .as_ref()
2255 .and_then(|a| a.expert_store_stats())
2256 {
2257 Some(es) => format!(
2258 "{body}\
2259 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
2260 # TYPE ferrox_expert_cache_hits_total counter\n\
2261 ferrox_expert_cache_hits_total {}\n\
2262 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2263 # TYPE ferrox_expert_cache_misses_total counter\n\
2264 ferrox_expert_cache_misses_total {}\n\
2265 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
2266 # TYPE ferrox_expert_cache_evictions_total counter\n\
2267 ferrox_expert_cache_evictions_total {}\n\
2268 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2269 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
2270 ferrox_expert_cache_pass_throughs_total {}\n\
2271 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2272 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
2273 ferrox_expert_cache_bytes_read_total {}\n\
2274 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2275 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
2276 ferrox_expert_cache_resident_bytes {}\n",
2277 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2278 ),
2279 None => body,
2280 };
2281
2282 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2287 Some(batcher) => {
2288 let sched = batcher.stats();
2289 format!(
2290 "{body}\
2291 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2292 # TYPE ferrox_prefill_chunks_total counter\n\
2293 ferrox_prefill_chunks_total {}\n\
2294 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2295 # TYPE ferrox_prefill_tokens_total counter\n\
2296 ferrox_prefill_tokens_total {}\n\
2297 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2298 # TYPE ferrox_decode_steps_total counter\n\
2299 ferrox_decode_steps_total {}\n\
2300 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2301 # TYPE ferrox_scheduler_queue_depth gauge\n\
2302 ferrox_scheduler_queue_depth {}\n\
2303 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2304 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2305 ferrox_scheduler_queue_rejected_total {}\n\
2306 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2307 # TYPE ferrox_kv_blocks_total gauge\n\
2308 ferrox_kv_blocks_total {}\n\
2309 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2310 # TYPE ferrox_kv_blocks_free gauge\n\
2311 ferrox_kv_blocks_free {}\n\
2312 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2313 # TYPE ferrox_kv_block_size gauge\n\
2314 ferrox_kv_block_size {}\n\
2315 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2316 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2317 ferrox_kv_rejected_too_large_total {}\n\
2318 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2319 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2320 ferrox_kv_rejected_context_length_total {}\n\
2321 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2322 # TYPE ferrox_scheduler_aborted_total counter\n\
2323 ferrox_scheduler_aborted_total {}\n",
2324 sched.prefill_chunks,
2325 sched.prefill_tokens,
2326 sched.decode_steps,
2327 sched.queue_depth,
2328 sched.queue_rejected,
2329 sched.kv_blocks_total,
2330 sched.kv_blocks_free,
2331 sched.kv_block_size,
2332 sched.kv_rejected_too_large,
2333 sched.kv_rejected_context_length,
2334 sched.aborted,
2335 )
2336 }
2337 None => body,
2338 };
2339
2340 (
2341 [(
2342 axum::http::header::CONTENT_TYPE,
2343 "text/plain; version=0.0.4",
2344 )],
2345 body,
2346 )
2347 .into_response()
2348}
2349
2350pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2351
2352pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2358 (
2359 StatusCode::BAD_REQUEST,
2360 Json(serde_json::json!({"error": {
2361 "message": message,
2362 "type": "invalid_request_error",
2363 "param": param,
2364 "code": null,
2365 }})),
2366 )
2367}
2368
2369pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2370 (
2371 StatusCode::NOT_IMPLEMENTED,
2372 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2373 )
2374}
2375
2376pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2377 let status = match e {
2378 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2379 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2384 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2390 StatusCode::SERVICE_UNAVAILABLE
2391 }
2392 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2396 };
2397 tracing::warn!("decode error: {e}");
2398 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2399 if let generate::DecodeError::KvBudgetExceeded {
2405 binding,
2406 estimated_bytes,
2407 limit_bytes,
2408 positions,
2409 positions_limit,
2410 ..
2411 } = &e
2412 {
2413 body["error"]["type"] = serde_json::json!("invalid_request_error");
2414 body["error"]["code"] = serde_json::json!(binding);
2415 body["error"]["binding"] = serde_json::json!(binding);
2416 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2417 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2418 body["error"]["positions"] = serde_json::json!(positions);
2419 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2420 }
2421 if let Some(secs) = e.retry_after_secs() {
2425 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2426 }
2427 (status, Json(body))
2428}
2429
2430pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2431 tracing::error!("generation task panicked: {e}");
2432 (
2433 StatusCode::INTERNAL_SERVER_ERROR,
2434 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2435 )
2436}
2437
2438#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2448 model: &Model,
2449 prompt: &str,
2450 params: &GenerationParams,
2451 kv_pool: Option<&generate::KvPoolConfig>,
2452 paged_kv: Option<&generate::PagedKvConfig>,
2453 prefix_cache: Option<&Mutex<PrefixCache>>,
2454 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2455 ceiling: Option<&budget::ContextCeiling>,
2456 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2457 mut emit: impl FnMut(&str),
2458) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2459 let synthetic = model.is_synthetic();
2460 let mut chunks = Vec::new();
2461 let params = &{
2467 let mut resolved = params.clone();
2468 resolved.stop_token_ids =
2469 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
2470 resolved
2471 };
2472 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2473 let _metal_private_guard =
2474 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2475 let (finish, usage) = match model {
2476 Model::Gguf(m) => {
2477 if let Some(batcher) = continuous_batcher {
2478 let mut tokens = m.tokenizer.encode(prompt);
2479 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2480 let (finish, _generated_ids, text, usage) = if synthetic {
2481 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2482 } else {
2483 batcher.generate_streaming(
2484 tokens,
2485 params.clone(),
2486 m.stop_tokens.clone(),
2487 Some(|chunk: &str| {
2488 if !chunk.is_empty() {
2489 chunks.push(chunk.to_string());
2490 emit(chunk);
2491 }
2492 }),
2493 )?
2494 };
2495 if !text.is_empty() && chunks.is_empty() {
2496 chunks.push(text);
2497 }
2498 (finish, usage)
2499 } else {
2500 generate::generate(
2501 &m.decoder,
2502 m.tokenizer.as_ref(),
2503 &m.stop_tokens,
2504 m.bos_id,
2505 prompt,
2506 params,
2507 kv_pool,
2508 paged_kv,
2509 prefix_cache,
2510 ceiling,
2511 |chunk| {
2512 chunks.push(chunk.to_string());
2513 if !synthetic {
2514 emit(chunk);
2515 }
2516 },
2517 )?
2518 }
2519 }
2520 Model::Kimi(m) => generate::generate_engine(
2521 &m.engine,
2522 &m.tokenizer,
2523 &m.stop_tokens,
2524 None,
2525 prompt,
2526 params,
2527 |chunk| {
2528 chunks.push(chunk.to_string());
2529 if !synthetic {
2530 emit(chunk);
2531 }
2532 },
2533 )?,
2534 Model::Mla(m) => generate::generate_engine(
2535 &m.engine,
2536 &m.tokenizer,
2537 &m.stop_tokens,
2538 m.bos_id,
2539 prompt,
2540 params,
2541 |chunk| {
2542 chunks.push(chunk.to_string());
2543 if !synthetic {
2544 emit(chunk);
2545 }
2546 },
2547 )?,
2548 Model::Gemma4(m) => generate::generate_engine(
2549 &m.engine,
2550 &m.tokenizer,
2551 &m.stop_tokens,
2552 m.bos_id,
2553 prompt,
2554 params,
2555 |chunk| {
2556 chunks.push(chunk.to_string());
2557 if !synthetic {
2558 emit(chunk);
2559 }
2560 },
2561 )?,
2562 Model::Glm52(m) => generate::generate_engine(
2563 &m.engine,
2564 &m.tokenizer,
2565 &m.stop_tokens,
2566 m.bos_id,
2567 prompt,
2568 params,
2569 |chunk| {
2570 chunks.push(chunk.to_string());
2571 if !synthetic {
2572 emit(chunk);
2573 }
2574 },
2575 )?,
2576 };
2577
2578 let mut full = chunks.concat();
2579 if synthetic {
2580 full = format!(
2581 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2582 to serve a real model. Decoded ids -> {full:?}]"
2583 );
2584 emit(&full);
2585 } else if used_batcher && !full.is_empty() && chunks.is_empty() {
2586 emit(&full);
2587 }
2588
2589 Ok((finish, usage, full))
2590}
2591
2592#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2597 model: &Model,
2598 prompt: &str,
2599 params: &GenerationParams,
2600 kv_pool: Option<&generate::KvPoolConfig>,
2601 paged_kv: Option<&generate::PagedKvConfig>,
2602 prefix_cache: Option<&Mutex<PrefixCache>>,
2603 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2604 ceiling: Option<&budget::ContextCeiling>,
2605 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2606) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2607 let (finish, usage, full) = run_generation_emit(
2608 model,
2609 prompt,
2610 params,
2611 kv_pool,
2612 paged_kv,
2613 prefix_cache,
2614 continuous_batcher,
2615 ceiling,
2616 metal_private_decode_gate,
2617 |_| {},
2618 )?;
2619 Ok((
2620 if full.is_empty() {
2621 Vec::new()
2622 } else {
2623 vec![full]
2624 },
2625 finish,
2626 usage,
2627 ))
2628}
2629
2630pub(crate) fn prompt_from_messages(
2641 messages: &[ChatMessage],
2642 template: &chat_template::PromptTemplate,
2643 tools: &[ToolDef],
2644 extra: serde_json::Map<String, serde_json::Value>,
2645) -> Result<String, ApiError> {
2646 let rendered = if tools.is_empty() || template.handles_tools() {
2647 template.render(messages, tools, extra)
2648 } else {
2649 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2650 with_preamble.push(ChatMessage {
2651 role: "system".to_string(),
2652 content: Some(MessageContent::Text(tool_preamble(tools))),
2653 tool_calls: None,
2654 tool_call_id: None,
2655 reasoning_content: None,
2656 });
2657 with_preamble.extend_from_slice(messages);
2658 template.render(&with_preamble, &[], extra)
2659 };
2660 rendered.map_err(template_error_response)
2661}
2662
2663fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2668 (
2669 StatusCode::BAD_REQUEST,
2670 Json(serde_json::json!({
2671 "error": {
2672 "message": format!("chat template failed to render: {err}"),
2673 "type": "invalid_request_error",
2674 "param": "messages",
2675 "code": null,
2676 }
2677 })),
2678 )
2679}
2680
2681fn tool_preamble(tools: &[ToolDef]) -> String {
2694 let mut out = String::from(
2695 "You can call tools to help answer the user. To call a tool, respond with \
2696 EXACTLY one line in this format and nothing else:\n\
2697 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2698 object matching that tool's parameters>}}</tool_call>\n\n\
2699 Available tools:\n",
2700 );
2701 for t in tools {
2702 out.push_str(&format!(
2703 "- {}: {}\n parameters (JSON schema): {}\n",
2704 t.function.name,
2705 t.function.description.as_deref().unwrap_or(""),
2706 t.function
2707 .parameters
2708 .as_ref()
2709 .map(|v| v.to_string())
2710 .unwrap_or_else(|| "{}".to_string()),
2711 ));
2712 }
2713 out
2714}
2715
2716fn tool_call_deltas(
2726 events: Vec<crate::policy::parser::ToolCallEvent>,
2727 opened: &std::cell::Cell<usize>,
2728) -> (String, Vec<ToolCallDelta>) {
2729 let mut text = String::new();
2730 let mut deltas = Vec::new();
2731 for event in events {
2732 match event {
2733 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2734 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2735 opened.set(opened.get().max(index + 1));
2736 deltas.push(ToolCallDelta::opening(index, name));
2737 }
2738 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2739 if !fragment.is_empty() {
2740 deltas.push(ToolCallDelta::arguments(index, fragment));
2741 }
2742 }
2743 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2744 }
2745 }
2746 (text, deltas)
2747}
2748
2749fn build_response_message(
2760 text: String,
2761 tools: &[ToolDef],
2762 posture: output::OutputPosture,
2763 base_finish: &'static str,
2764) -> (ChatCompletionResponseMessage, &'static str) {
2765 let parsed = output::parse_output(&text, tools, posture);
2766 let calls: Vec<ToolCallOut> = parsed
2767 .calls
2768 .into_iter()
2769 .enumerate()
2770 .map(|(index, call)| ToolCallOut {
2771 id: format!("call_{index}"),
2772 kind: "function",
2773 function: ToolCallFunctionOut {
2774 name: call.name,
2775 arguments: call.arguments,
2776 },
2777 })
2778 .collect();
2779 if !calls.is_empty() {
2780 return (
2781 ChatCompletionResponseMessage {
2782 role: "assistant",
2783 content: None,
2784 reasoning_content: parsed.reasoning,
2785 tool_calls: Some(calls),
2786 },
2787 "tool_calls",
2788 );
2789 }
2790 (
2791 ChatCompletionResponseMessage {
2792 role: "assistant",
2793 content: Some(parsed.content),
2794 reasoning_content: parsed.reasoning,
2795 tool_calls: None,
2796 },
2797 base_finish,
2798 )
2799}
2800
2801fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2806 let mut history = match &req.session_id {
2807 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2808 None => req.messages.clone(),
2809 };
2810 if req.json_object_mode() {
2811 inject_json_object_system_hint(&mut history);
2812 }
2813 history
2814}
2815
2816fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2817 const HINT: &str =
2818 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2819 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2820 match &mut sys.content {
2821 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2822 s.push_str("\n\n");
2823 s.push_str(HINT);
2824 }
2825 None => {
2826 sys.content = Some(MessageContent::Text(HINT.to_string()));
2827 }
2828 _ => {}
2829 }
2830 } else {
2831 messages.insert(
2832 0,
2833 ChatMessage {
2834 role: "system".to_string(),
2835 content: Some(MessageContent::Text(HINT.to_string())),
2836 tool_calls: None,
2837 tool_call_id: None,
2838 reasoning_content: None,
2839 },
2840 );
2841 }
2842}
2843
2844async fn chat_completions(
2845 State(state): State<Arc<AppState>>,
2846 headers: axum::http::HeaderMap,
2847 Json(req): Json<ChatCompletionRequest>,
2848) -> Response {
2849 let attribution = attribution::Attribution::from_headers(&headers);
2850 state
2851 .requests_total
2852 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2853 let started = std::time::Instant::now();
2854
2855 let request_id = ferrox_api::next_request_id();
2859 let stream = req.stream.unwrap_or(false);
2860
2861 let refusal = cache_admin::check_admission(&state)
2867 .err()
2868 .or_else(|| req.validate_supported_fields().err());
2869 if let Some(err) = refusal {
2870 state
2871 .request_errors_total
2872 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2873 let response = err.into_response();
2874 state.record_request(stats::Record {
2875 request_id: &request_id,
2876 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2877 model: state.active_model_name(),
2878 status: response.status().as_u16(),
2879 stream,
2880 duration_ms: started.elapsed().as_millis() as u64,
2881 usage: None,
2882 attribution: &attribution,
2883 });
2884 return response;
2885 }
2886
2887 let response = if stream {
2888 chat_completions_stream(
2889 Arc::clone(&state),
2890 req,
2891 request_id.clone(),
2892 started,
2893 attribution.clone(),
2894 )
2895 .await
2896 .into_response()
2897 } else {
2898 chat_completions_full(
2899 Arc::clone(&state),
2900 req,
2901 request_id.clone(),
2902 started,
2903 attribution.clone(),
2904 )
2905 .await
2906 .into_response()
2907 };
2908
2909 if response.status().is_client_error() || response.status().is_server_error() {
2910 state
2911 .request_errors_total
2912 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2913 state.record_request(stats::Record {
2917 request_id: &request_id,
2918 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2919 model: state.active_model_name(),
2922 status: response.status().as_u16(),
2923 stream,
2924 duration_ms: started.elapsed().as_millis() as u64,
2925 usage: None,
2926 attribution: &attribution,
2927 });
2928 }
2929 state.mark_request_finished();
2930
2931 response
2932}
2933
2934async fn chat_completions_full(
2935 state: Arc<AppState>,
2936 req: ChatCompletionRequest,
2937 request_id: String,
2938 started: std::time::Instant,
2939 attribution: attribution::Attribution,
2940) -> Result<Json<ChatCompletionResponse>, ApiError> {
2941 let tools_active = req.tools_active();
2942 let active = state.require_active()?;
2946 let history = resolve_history(&state, &req);
2947 let template = active.generative()?.chat_template();
2948 let kwargs = req.resolve_template_kwargs(&template);
2949 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2950 let key = req.is_cacheable().then(|| req.cache_key(&prompt));
2951
2952 let (completion, cache_status) = if let Some(cached) = key
2953 .as_ref()
2954 .and_then(|key| lock_cache(&state.response_cache).get(key))
2955 {
2956 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2957 (cached, "hit")
2958 } else {
2959 let params = req.generation_params_for_template(&template, active.name())?;
2960 let (chunks, finish, usage) = decode_task::buffered(
2961 decode_task::DecodeHandles::take(&state, &active)?,
2962 prompt.clone(),
2963 params,
2964 )
2965 .await?;
2966
2967 let completion = response_cache::CachedCompletion {
2968 content: chunks.concat(),
2969 finish,
2970 usage,
2971 };
2972 let cache_status = if let Some(key) = key {
2973 tracing::debug!("cache miss for key {}", key.digest());
2974 lock_cache(&state.response_cache).put(key, completion.clone());
2975 "miss"
2976 } else {
2977 "skip"
2978 };
2979 (completion, cache_status)
2980 };
2981 let content = completion.content;
2982
2983 if req.json_object_mode() {
2984 json_mode::validate_json_object_output(&content)?;
2985 }
2986
2987 if let Some(id) = &req.session_id {
2991 state.sessions.store_reply(
2992 id,
2993 ChatMessage {
2994 role: "assistant".to_string(),
2995 content: Some(MessageContent::Text(content.clone())),
2996 tool_calls: None,
2997 tool_call_id: None,
2998 reasoning_content: None,
2999 },
3000 );
3001 }
3002
3003 let (message, finish_reason) = build_response_message(
3004 content,
3005 if tools_active { &req.tools } else { &[] },
3006 output::OutputPosture::resolve(active.name(), &prompt),
3007 completion.finish.as_str(),
3008 );
3009
3010 state.record_request(stats::Record {
3011 request_id: &request_id,
3012 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3013 model: Some(active.name().to_string()),
3016 status: 200,
3017 stream: false,
3018 duration_ms: started.elapsed().as_millis() as u64,
3019 usage: Some(&completion.usage),
3020 attribution: &attribution,
3021 });
3022
3023 Ok(Json(ChatCompletionResponse {
3024 id: request_id.clone(),
3025 request_id,
3026 object: "chat.completion",
3027 model: req.model,
3028 choices: vec![ChatCompletionChoice {
3029 index: 0,
3030 message,
3031 finish_reason,
3032 }],
3033 usage: completion.usage,
3034 ferrox_cache: cache_status,
3035 }))
3036}
3037
3038async fn chat_completions_stream(
3039 state: Arc<AppState>,
3040 req: ChatCompletionRequest,
3041 request_id: String,
3042 started: std::time::Instant,
3043 attribution: attribution::Attribution,
3044) -> Result<Response, ApiError> {
3045 let tools_active = req.tools_active();
3047 let active = state.require_active()?;
3051 let history = resolve_history(&state, &req);
3052 let template = active.generative()?.chat_template();
3053 let kwargs = req.resolve_template_kwargs(&template);
3054 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
3055 let model_name = req.model.clone();
3056 let session_id = req.session_id.clone();
3057 let sessions = state.sessions.clone();
3058
3059 let model = Arc::clone(active.generative()?);
3060 let kv_pool = state.kv_pool.clone();
3061 let paged_kv = state.paged_kv.clone();
3062 let prefix_cache = state.prefix_cache.clone();
3063 let batcher = active.batcher.clone();
3064 let ceiling = active.ceiling.clone();
3065 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3066 let mut params = req.generation_params_for_template(&template, active.name())?;
3067 let stats_state = Arc::clone(&state);
3068 let served_model = active.name().to_string();
3071 let posture = output::OutputPosture::resolve(&served_model, &prompt);
3075 let offered_tools: Vec<ToolDef> = if tools_active {
3078 req.tools.clone()
3079 } else {
3080 Vec::new()
3081 };
3082
3083 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3088 params.cancel = Some(cancel_token.clone());
3089
3090 let overlap = true;
3099
3100 let slot = req
3104 .stream_resumable
3105 .unwrap_or(false)
3106 .then(|| state.streams.register(&request_id));
3107 let emitter = resume::Emitter::new(slot);
3108
3109 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3110 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3114 id: request_id.clone(),
3115 request_id: None,
3116 object: "chat.completion.chunk",
3117 model: model_name.clone(),
3118 choices: vec![ChatCompletionChunkChoice {
3119 index: 0,
3120 delta: ChatCompletionChunkDelta {
3121 role: None,
3122 content: None,
3123 reasoning_content: None,
3124 tool_calls: None,
3125 },
3126 finish_reason: None,
3127 }],
3128 usage: None,
3129 });
3130
3131 tokio::task::spawn_blocking(move || {
3132 let _cancel_guard = cancel_guard;
3135 let tx_chunks = tx.clone();
3136 let orphan_timeout = sse::orphan_timeout_from_env();
3141 let mut first = true;
3142 let head_request_id = request_id.clone();
3143 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3153 Rc::new(RefCell::new(posture.reasoning_parser()));
3154 let emit_reasoning = Rc::clone(&stream_reasoning);
3155 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3159 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3160 );
3161 let emit_tools = Rc::clone(&stream_tools);
3162 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3166 let emit_streamed_calls = Rc::clone(&streamed_calls);
3167 let result = run_generation_emit(
3168 &model,
3169 &prompt,
3170 ¶ms,
3171 kv_pool.as_ref(),
3172 paged_kv.as_ref(),
3173 prefix_cache.as_deref(),
3174 batcher.as_ref(),
3175 ceiling.as_deref(),
3176 metal_private_decode_gate.as_deref(),
3177 |chunk| {
3178 if !overlap || chunk.is_empty() {
3179 return;
3180 }
3181 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3182 Some(parser) => {
3183 let delta = parser.push(chunk);
3184 (delta.reasoning, delta.content)
3185 }
3186 None => (String::new(), chunk.to_string()),
3187 };
3188 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3192 Some(parser) => {
3193 let (text, calls) =
3194 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3195 (text, calls)
3196 }
3197 None => (content, Vec::new()),
3198 };
3199 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3202 return;
3203 }
3204 let role = if first { Some("assistant") } else { None };
3205 let request_id = first.then(|| head_request_id.clone());
3206 first = false;
3207 let payload = ChatCompletionChunk {
3208 id: head_request_id.clone(),
3209 request_id,
3210 object: "chat.completion.chunk",
3211 model: model_name.clone(),
3212 choices: vec![ChatCompletionChunkChoice {
3213 index: 0,
3214 delta: ChatCompletionChunkDelta {
3215 role,
3216 content: (!content.is_empty()).then_some(content),
3217 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3218 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3219 },
3220 finish_reason: None,
3221 }],
3222 usage: None,
3223 };
3224 if let Err(why) =
3233 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3234 {
3235 if why == sse::SendFailure::Orphaned {
3236 tracing::warn!(
3237 "SSE stream {head_request_id} accepted nothing for the orphan \
3238 deadline; treating it as abandoned"
3239 );
3240 }
3241 if !emitter.is_resumable() {
3253 cancel_token.cancel();
3254 }
3255 }
3256 },
3257 );
3258
3259 let mut pending_request_id = first.then(|| request_id.clone());
3265
3266 match result {
3267 Ok((finish, usage, full_text)) => {
3268 if let Some(id) = &session_id {
3269 sessions.store_reply(
3270 id,
3271 ChatMessage {
3272 role: "assistant".to_string(),
3273 content: Some(MessageContent::Text(full_text.clone())),
3274 tool_calls: None,
3275 tool_call_id: None,
3276 reasoning_content: None,
3277 },
3278 );
3279 }
3280 let mut streamed_finish: Option<&'static str> = None;
3286 if overlap {
3287 let tail = stream_reasoning
3288 .borrow_mut()
3289 .as_mut()
3290 .map(|parser| parser.flush())
3291 .unwrap_or_default();
3292 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3293 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3294 let mut events = parser.push(&content);
3295 events.extend(parser.finish());
3296 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3297 content = text;
3298 tool_calls = calls;
3299 }
3300 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3301 let payload = ChatCompletionChunk {
3302 id: request_id.clone(),
3303 request_id: pending_request_id.take(),
3304 object: "chat.completion.chunk",
3305 model: model_name.clone(),
3306 choices: vec![ChatCompletionChunkChoice {
3307 index: 0,
3308 delta: ChatCompletionChunkDelta {
3309 role: None,
3310 content: (!content.is_empty()).then_some(content),
3311 reasoning_content: (!tail.reasoning.is_empty())
3312 .then_some(tail.reasoning),
3313 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3314 },
3315 finish_reason: None,
3316 }],
3317 usage: None,
3318 };
3319 let _ =
3320 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3321 }
3322 if streamed_calls.get() > 0 {
3323 streamed_finish = Some("tool_calls");
3324 }
3325 } else {
3326 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3329 let tool_calls: Vec<ToolCallDelta> = parsed
3330 .calls
3331 .iter()
3332 .enumerate()
3333 .map(|(index, call)| {
3334 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3335 })
3336 .collect();
3337 if !tool_calls.is_empty() {
3338 streamed_finish = Some("tool_calls");
3339 }
3340 if !tool_calls.is_empty()
3341 || !parsed.content.is_empty()
3342 || parsed.reasoning.is_some()
3343 {
3344 let payload = ChatCompletionChunk {
3345 id: request_id.clone(),
3346 request_id: pending_request_id.take(),
3347 object: "chat.completion.chunk",
3348 model: model_name.clone(),
3349 choices: vec![ChatCompletionChunkChoice {
3350 index: 0,
3351 delta: ChatCompletionChunkDelta {
3352 role: Some("assistant"),
3353 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3354 .then(|| parsed.content.clone()),
3355 reasoning_content: parsed.reasoning.clone(),
3356 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3357 },
3358 finish_reason: None,
3359 }],
3360 usage: None,
3361 };
3362 let _ =
3363 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3364 }
3365 }
3366 let final_finish_reason = match streamed_finish {
3370 Some(reason) if finish.as_str() != "length" => reason,
3371 _ => finish.as_str(),
3372 };
3373 let final_payload = ChatCompletionChunk {
3374 id: request_id.clone(),
3375 request_id: pending_request_id.take(),
3376 object: "chat.completion.chunk",
3377 model: model_name,
3378 choices: vec![ChatCompletionChunkChoice {
3379 index: 0,
3380 delta: ChatCompletionChunkDelta {
3381 role: None,
3382 content: None,
3383 reasoning_content: None,
3384 tool_calls: None,
3385 },
3386 finish_reason: Some(final_finish_reason),
3387 }],
3388 usage: Some(usage.clone()),
3389 };
3390 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3391 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3392 stats_state.record_request(stats::Record {
3397 request_id: &request_id,
3398 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3399 model: Some(served_model.clone()),
3400 status: 200,
3401 stream: true,
3402 duration_ms: started.elapsed().as_millis() as u64,
3403 usage: Some(&usage),
3404 attribution: &attribution,
3405 });
3406 }
3407 Err(e) => {
3408 tracing::warn!("decode error on streamed request {request_id}: {e}");
3409 stats_state.record_request(stats::Record {
3415 request_id: &request_id,
3416 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3417 model: Some(served_model.clone()),
3418 status: 500,
3419 stream: true,
3420 duration_ms: started.elapsed().as_millis() as u64,
3421 usage: None,
3422 attribution: &attribution,
3423 });
3424 let payload = ChatCompletionChunk {
3425 id: request_id.clone(),
3426 request_id: pending_request_id.take(),
3427 object: "chat.completion.chunk",
3428 model: model_name,
3429 choices: vec![ChatCompletionChunkChoice {
3430 index: 0,
3431 delta: ChatCompletionChunkDelta {
3432 role: Some("assistant"),
3433 content: Some(format!("[error: {e}]")),
3434 reasoning_content: None,
3435 tool_calls: None,
3436 },
3437 finish_reason: Some("stop"),
3438 }],
3439 usage: None,
3440 };
3441 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3442 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3443 }
3444 }
3445 drop(emitter);
3449 });
3450
3451 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3452 Ok((
3475 [(
3476 axum::http::HeaderName::from_static("x-accel-buffering"),
3477 axum::http::HeaderValue::from_static("no"),
3478 )],
3479 Sse::new(stream),
3480 )
3481 .into_response())
3482}
3483
3484fn axum_path(template: &str) -> String {
3501 let mut out = String::with_capacity(template.len());
3502 let mut rest = template;
3503 while let Some(open) = rest.find('{') {
3504 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3505 break;
3506 };
3507 out.push_str(&rest[..open]);
3508 out.push(':');
3509 out.push_str(&rest[open + 1..close]);
3510 rest = &rest[close + 1..];
3511 }
3512 out.push_str(rest);
3513 out
3514}
3515
3516async fn cancel_generation(
3525 State(state): State<Arc<AppState>>,
3526 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3527) -> Response {
3528 let cancelled = state.cancels.cancel(&req.request_id);
3529 let status = if cancelled {
3530 StatusCode::OK
3531 } else {
3532 StatusCode::NOT_FOUND
3533 };
3534 let detail = if cancelled {
3535 "the generation was asked to stop; it ends at its next token".to_string()
3536 } else {
3537 "no generation with that request_id is running -- it has already \
3538 finished, was never issued, or was served by a path that does \
3539 not register for cancellation"
3540 .to_string()
3541 };
3542 (
3543 status,
3544 Json(ferrox_api::CancelGenerationResponse {
3545 request_id: req.request_id,
3546 cancelled,
3547 detail,
3548 }),
3549 )
3550 .into_response()
3551}
3552
3553type Activated = (
3557 Loaded,
3558 Option<serving::batch::ContinuousBatcher>,
3559 Option<Arc<budget::ContextCeiling>>,
3560);
3561
3562fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3580 let mut batcher = serving::batch::BatcherConfig::from_env();
3581 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3582 return batcher;
3585 }
3586 let Some(path) = path else {
3587 return batcher;
3588 };
3589 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
3595 let Some((priced, gguf_ctx, source)) = priced else {
3596 return batcher;
3597 };
3598 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3599 tracing::warn!(
3603 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3604 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3605 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3606 admit on a number you choose.",
3607 priced.weights_bytes,
3608 priced.device_budget_bytes,
3609 );
3610 return batcher;
3611 };
3612 tracing::info!("{source}");
3613 tracing::info!("{}", derived.fit);
3614 let adopted = budget::apply_derived(&mut batcher, &derived);
3615 if adopted.max_context {
3616 tracing::info!(
3617 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3618 override with FERROX_CB_MAX_CONTEXT",
3619 derived.max_context
3620 );
3621 }
3622 if adopted.kv_blocks {
3623 tracing::info!(
3624 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3625 derived.kv_blocks,
3626 batcher.kv_block_size
3627 );
3628 }
3629 batcher
3630}
3631
3632pub(crate) fn activate_loaded_model(
3641 loaded: model::LoadedModel,
3642 enable_continuous_batching: bool,
3643 path: Option<&str>,
3644 paged_kv: Option<&generate::PagedKvConfig>,
3645) -> Activated {
3646 match loaded {
3647 model::LoadedModel::Gguf(g) => {
3648 let decoder = Arc::new(g.decoder);
3649 let tokenizer = Arc::new(g.tokenizer);
3650 let config = price_batcher_config(path);
3651 let shape =
3657 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
3658 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3659 let batcher = if enable_continuous_batching {
3660 tracing::info!(
3661 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3662 (stop sequences use the same pending-buffer trim as the private generate loop)"
3663 );
3664 let tok = Arc::clone(&tokenizer);
3665 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
3666 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3667 Arc::clone(&decoder),
3668 decode,
3669 config,
3670 Arc::clone(&ceiling),
3671 paged_kv.cloned(),
3672 ))
3673 } else {
3674 None
3675 };
3676 (
3677 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3678 decoder,
3679 tokenizer,
3680 stop_tokens: g.stop_tokens,
3681 bos_id: g.bos_id,
3682 is_synthetic: g.is_synthetic,
3683 chat_template: g.chat_template,
3684 }))),
3685 batcher,
3686 Some(ceiling),
3687 )
3688 }
3689 model::LoadedModel::Kimi(k) => (
3690 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3691 engine: k.engine,
3692 tokenizer: k.tokenizer,
3693 stop_tokens: k.stop_tokens,
3694 chat_template: k.chat_template,
3695 }))),
3696 None,
3697 None,
3698 ),
3699 model::LoadedModel::Mla(m) => (
3700 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3701 engine: m.engine,
3702 tokenizer: m.tokenizer,
3703 stop_tokens: m.stop_tokens,
3704 bos_id: m.bos_id,
3705 name: m.name,
3706 chat_template: m.chat_template,
3707 }))),
3708 None,
3709 None,
3710 ),
3711 model::LoadedModel::Gemma4(m) => (
3712 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3713 engine: m.engine,
3714 tokenizer: m.tokenizer,
3715 stop_tokens: m.stop_tokens,
3716 bos_id: m.bos_id,
3717 name: m.name,
3718 chat_template: m.chat_template,
3719 }))),
3720 None,
3721 None,
3722 ),
3723 model::LoadedModel::Glm52(g) => (
3724 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3725 engine: g.engine,
3726 tokenizer: g.tokenizer,
3727 stop_tokens: g.stop_tokens,
3728 bos_id: g.bos_id,
3729 name: g.name,
3730 chat_template: g.chat_template,
3731 }))),
3732 None,
3733 None,
3734 ),
3735 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3740 }
3741}
3742
3743struct StartupModels {
3750 loaded: model::LoadedModel,
3751 embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
3752}
3753
3754fn continuous_batching_env() -> Option<bool> {
3755 match std::env::var("FERROX_CONTINUOUS_BATCHING")
3756 .ok()
3757 .map(|v| v.trim().to_ascii_lowercase())
3758 .as_deref()
3759 {
3760 None => None,
3761 Some("1" | "true" | "yes" | "on") => Some(true),
3762 Some("0" | "false" | "no" | "off") => Some(false),
3763 _ => None,
3764 }
3765}
3766
3767fn metal_private_decode_active() -> bool {
3768 #[cfg(feature = "metal")]
3769 {
3770 BUILT_WITH_METAL
3771 && ferrox_metal::attn::metal_attn_enabled()
3772 && std::env::var("FERROX_METAL").ok().as_deref() != Some("0")
3773 }
3774 #[cfg(not(feature = "metal"))]
3775 {
3776 false
3777 }
3778}
3779
3780fn continuous_batching_compatible(
3781 loaded: &model::LoadedModel,
3782 kv_pool: &Option<generate::KvPoolConfig>,
3783 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3784 paged_kv: &Option<generate::PagedKvConfig>,
3785) -> bool {
3786 matches!(loaded, model::LoadedModel::Gguf(_))
3787 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3788}
3789
3790fn resolve_continuous_batching_enabled(
3791 loaded: &model::LoadedModel,
3792 kv_pool: &Option<generate::KvPoolConfig>,
3793 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3794 paged_kv: &Option<generate::PagedKvConfig>,
3795) -> bool {
3796 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3797 return false;
3798 }
3799 match continuous_batching_env() {
3800 Some(true) => true,
3801 Some(false) => false,
3802 None => metal_private_decode_active(),
3803 }
3804}
3805
3806fn acquire_metal_private_decode_gate(
3807 gate: Option<&std::sync::Mutex<()>>,
3808 used_batcher: bool,
3809) -> Option<std::sync::MutexGuard<'_, ()>> {
3810 if used_batcher {
3811 None
3812 } else {
3813 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3814 }
3815}
3816
3817fn build_app_state(
3818 models: StartupModels,
3819 kv_pool: Option<generate::KvPoolConfig>,
3820 paged_kv: Option<generate::PagedKvConfig>,
3821 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3822 enable_continuous_batching: bool,
3823 mcp: Option<mcp::LoadedMcpConfig>,
3824 detection: Arc<health::Detection>,
3825) -> AppState {
3826 let StartupModels { loaded, embedding } = models;
3827 let (loaded, batcher, ceiling) = activate_loaded_model(
3828 loaded,
3829 enable_continuous_batching,
3830 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
3831 paged_kv.as_ref(),
3832 );
3833 let id = startup_model_id();
3839 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3840 {
3841 None
3842 } else {
3843 tracing::info!(
3844 "Metal private-loop decode will serialize concurrent requests until \
3845 continuous batching is enabled (FERROX_CONTINUOUS_BATCHING=1 or --cont-batching)"
3846 );
3847 Some(Arc::new(std::sync::Mutex::new(())))
3848 };
3849 AppState {
3850 embedding,
3851 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3852 id,
3853 loaded,
3854 batcher,
3855 ceiling,
3856 }))),
3857 paged_kv,
3858 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3859 tasks: Arc::new(tasks::TaskRegistry::new()),
3860 cancels: Arc::new(cancel::CancelRegistry::new()),
3861 stats: stats::Stats::new(),
3862 streams: resume::StreamRegistry::new(),
3863 model_dir: admin::model_dirs().into_iter().next(),
3864 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3865 kv_pool,
3866 prefix_cache,
3867 sessions: session::SessionStore::new(),
3868 requests_total: std::sync::atomic::AtomicU64::new(0),
3869 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3870 started_at: std::time::Instant::now(),
3871 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3872 detection,
3873 mcp,
3874 continuous_batching_enabled: enable_continuous_batching,
3875 metal_private_decode_gate,
3876 loading_model: Mutex::new(None),
3877 last_load_error: Mutex::new(None),
3878 serving: Mutex::new(crate::stats::ServingStats::default()),
3879 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3880 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3881 started_unix: unix_now(),
3882 }
3883}
3884
3885fn load_embedding_model() -> anyhow::Result<Option<Arc<ferrox_models::EmbeddingModel>>> {
3893 let Ok(path) = std::env::var("FERROX_EMBEDDING_MODEL_PATH") else {
3894 return Ok(None);
3895 };
3896 let model = ferrox_models::EmbeddingModel::from_gguf_path(&path)
3897 .map_err(|e| anyhow::anyhow!("FERROX_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3898 tracing::info!(
3899 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3900 model.name(),
3901 model.architecture(),
3902 model.n_embd(),
3903 model.pooling_type().name(),
3904 model.n_ctx_train(),
3905 );
3906 Ok(Some(Arc::new(model)))
3907}
3908
3909fn unix_now() -> u64 {
3914 std::time::SystemTime::now()
3915 .duration_since(std::time::UNIX_EPOCH)
3916 .map(|d| d.as_secs())
3917 .unwrap_or(0)
3918}
3919
3920fn startup_model_id() -> Option<String> {
3925 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
3926 let configured = std::fs::canonicalize(&configured).ok()?;
3927 admin::discover(&admin::model_dirs())
3928 .into_iter()
3929 .find(|d| {
3930 std::fs::canonicalize(&d.path)
3931 .map(|p| p == configured)
3932 .unwrap_or(false)
3933 })
3934 .map(|d| d.id)
3935}
3936
3937fn init_cpu_pool() {
3946 match ferrox_core::threads::init_cpu_pool() {
3947 Some(n) => eprintln!(
3948 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
3949 ferrox_core::threads::perf_core_count()
3950 ),
3951 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
3952 }
3953}
3954
3955fn announce_ready(addr: SocketAddr, scheme: &str) {
3968 use std::io::Write;
3969 let ready =
3970 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3971 let mut stdout = std::io::stdout().lock();
3972 let _ = writeln!(stdout, "{}", ready.to_line());
3973 let _ = stdout.flush();
3974}
3975
3976async fn shutdown_signal(exit_on_stdin_close: bool) {
3988 if !exit_on_stdin_close {
3989 std::future::pending::<()>().await;
3990 return;
3991 }
3992 let _ = tokio::task::spawn_blocking(|| {
3993 use std::io::Read;
3994 let mut sink = [0u8; 256];
3995 let mut stdin = std::io::stdin().lock();
3996 loop {
3997 match stdin.read(&mut sink) {
3998 Ok(0) => break,
4000 Ok(_) => continue,
4002 Err(e) => {
4003 tracing::warn!("stdin read failed ({e}); treating it as closed");
4004 break;
4005 }
4006 }
4007 }
4008 })
4009 .await;
4010 tracing::info!("stdin closed; shutting down");
4011}
4012
4013fn tokio_worker_threads() -> usize {
4018 std::env::var("FERROX_TOKIO_WORKERS")
4019 .ok()
4020 .and_then(|v| v.trim().parse::<usize>().ok())
4021 .filter(|n| *n > 0)
4022 .unwrap_or(2)
4023}
4024
4025fn install_ring_crypto_provider() {
4042 let _ = rustls::crypto::ring::default_provider().install_default();
4043}
4044
4045pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4051 if args.list_devices {
4052 print_available_devices();
4053 return Ok(());
4054 }
4055 apply_cli_overrides(&args)?;
4056
4057 let _instance = {
4061 use ferrox_core::instance::{register, InstancePolicy};
4062 let policy = if args.allow_multiple_instances {
4063 InstancePolicy::Multi
4064 } else {
4065 InstancePolicy::from_env_or(InstancePolicy::Single)
4066 };
4067 let model = std::env::var("FERROX_MODEL_PATH").ok();
4068 register(
4069 "server",
4070 model.as_deref(),
4071 ferrox_core::instance::current_backend(),
4072 policy,
4073 )
4074 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4075 };
4076
4077 let journal = journal::Journal::from_env();
4078 eprintln!(
4079 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
4080 journal.path()
4081 );
4082 journal.append(&journal::Record::session_start(
4083 env!("CARGO_PKG_VERSION"),
4084 std::process::id(),
4085 ));
4086 journal::install_panic_hook(journal.clone());
4087
4088 let mcp_config_path = args.mcp_config.clone();
4089 let exit_on_stdin_close = args.exit_on_stdin_close
4090 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
4091 .map(|v| v == "1")
4092 .unwrap_or(false);
4093
4094 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
4098 init_cpu_pool();
4099
4100 let runtime = tokio::runtime::Builder::new_multi_thread()
4101 .worker_threads(tokio_worker_threads())
4102 .enable_all()
4103 .build()?;
4104 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4105
4106 let reason = match &result {
4107 Ok(()) => "normal".to_string(),
4108 Err(e) => e.to_string(),
4109 };
4110 journal.append(&journal::Record::session_exit(reason));
4111
4112 runtime.shutdown_background();
4117
4118 result
4119}
4120
4121async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4122 let _ = tracing_subscriber::fmt::try_init();
4128
4129 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4139 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
4140 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
4141 .map(|v| v == "1")
4142 .unwrap_or(false);
4143 if let Err(msg) =
4144 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4145 {
4146 anyhow::bail!(msg);
4147 }
4148
4149 let embedding_model = load_embedding_model()?;
4156
4157 let mut loaded = model::load()?;
4158 match &loaded {
4159 model::LoadedModel::Gguf(g) => tracing::info!(
4160 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4161 g.decoder.config.name,
4162 g.is_synthetic,
4163 g.tokenizer.kind()
4164 ),
4165 model::LoadedModel::Kimi(k) => tracing::info!(
4166 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4167 k.tokenizer.vocab_size()
4168 ),
4169 model::LoadedModel::Mla(m) => tracing::info!(
4170 "loaded MLA GGUF '{}' (tokenizer={})",
4171 m.name,
4172 m.tokenizer.kind()
4173 ),
4174 model::LoadedModel::Gemma4(m) => tracing::info!(
4175 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4176 m.name,
4177 m.tokenizer.kind()
4178 ),
4179 model::LoadedModel::Glm52(g) => tracing::info!(
4180 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4181 g.name,
4182 g.tokenizer.kind()
4183 ),
4184 model::LoadedModel::Encoder(_) => {}
4187 }
4188 let metal_default_moe_budget = {
4196 #[cfg(feature = "metal")]
4197 {
4198 ferrox_core::metal_dense_enabled()
4199 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
4200 }
4201 #[cfg(not(feature = "metal"))]
4202 {
4203 false
4204 }
4205 };
4206 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
4207 let budget: u64 = budget_str
4208 .parse()
4209 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4210 match &mut loaded {
4211 model::LoadedModel::Gguf(g) => {
4212 tracing::info!(
4213 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4214 (CUDA and/or Metal matvecs when built with the matching feature)"
4215 );
4216 g.decoder.gpu_vram_budget_bytes = Some(budget);
4217 }
4218 model::LoadedModel::Kimi(_) => {
4219 tracing::warn!(
4220 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4221 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4222 );
4223 }
4224 model::LoadedModel::Mla(_) => {
4225 tracing::warn!(
4226 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4227 FFN path only today; ignoring expert VRAM budget"
4228 );
4229 }
4230 model::LoadedModel::Gemma4(_) => {
4231 tracing::warn!(
4232 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4233 ignoring expert VRAM budget"
4234 );
4235 }
4236 model::LoadedModel::Glm52(_) => {
4237 tracing::warn!(
4238 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4239 GPU expert placement not wired yet; ignoring"
4240 );
4241 }
4242 model::LoadedModel::Encoder(_) => {
4243 tracing::warn!(
4244 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4245 it has no routed experts to place; ignoring"
4246 );
4247 }
4248 }
4249 } else if metal_default_moe_budget {
4250 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4254 if let model::LoadedModel::Gguf(g) = &mut loaded {
4255 tracing::info!(
4256 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4257 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4258 );
4259 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4260 }
4261 }
4262 #[cfg(feature = "cuda")]
4263 {
4264 if ferrox_core::cuda_dense_enabled() {
4265 tracing::info!(
4266 "CUDA dense matvec enabled for WeightMatrix::apply \
4267 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4268 );
4269 } else {
4270 tracing::info!(
4271 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
4272 );
4273 }
4274 }
4275 #[cfg(feature = "metal")]
4276 {
4277 if ferrox_core::metal_dense_enabled() {
4278 tracing::info!(
4279 "Metal dense matvec enabled for WeightMatrix::apply \
4280 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4281 );
4282 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
4283 Some("1") | Some("true") | Some("on") | Some("attn") => {
4284 tracing::info!(
4285 "Metal fused attention requested (FERROX_METAL_ATTN): \
4286 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4287 );
4288 }
4289 _ => {}
4290 }
4291 tracing::info!(
4292 "Metal greedy GPU argmax: temperature<=0 folds \
4293 final_norm+lm_head+argmax into the dense stack"
4294 );
4295 } else {
4296 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
4297 }
4298 }
4299 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
4310 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
4311 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
4312 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4313 panic!(
4314 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
4315 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
4316 );
4317 }
4318 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4319 (Ok(blocks), Ok(block_size), Err(_)) => {
4320 let total_blocks: usize = blocks
4321 .parse()
4322 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
4323 let block_size: usize = block_size
4324 .parse()
4325 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4326 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4332 .ok()
4333 .map(|v| {
4334 v.parse()
4335 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4336 })
4337 .unwrap_or(0);
4338 tracing::info!(
4339 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4340 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4341 );
4342 Some(generate::KvPoolConfig {
4343 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4344 queue_wait: Duration::from_millis(queue_wait_ms),
4345 })
4346 }
4347 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4348 let block_size: usize = block_size
4349 .parse()
4350 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4351 let budget: u64 = byte_budget
4352 .parse()
4353 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
4354 let cfg = match &loaded {
4355 model::LoadedModel::Gguf(g) => &g.decoder.config,
4356 model::LoadedModel::Kimi(_)
4357 | model::LoadedModel::Mla(_)
4358 | model::LoadedModel::Gemma4(_)
4359 | model::LoadedModel::Glm52(_)
4360 | model::LoadedModel::Encoder(_) => {
4361 panic!(
4362 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
4363 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4364 );
4365 }
4366 };
4367 let bytes_per_block = block_size
4368 * cfg.n_layers
4369 * cfg.n_kv_heads
4370 * cfg.head_dim
4371 * 2
4372 * std::mem::size_of::<f32>();
4373 assert!(
4374 bytes_per_block > 0,
4375 "derived KV block byte size must be positive (check model config and block size)"
4376 );
4377 let total_blocks = (budget as usize / bytes_per_block).max(1);
4378 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4379 .ok()
4380 .map(|v| {
4381 v.parse()
4382 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4383 })
4384 .unwrap_or(0);
4385 tracing::info!(
4386 "KV cache block pool enabled from byte budget: {budget} bytes / \
4387 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4388 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4389 cfg.n_layers
4390 );
4391 Some(generate::KvPoolConfig {
4392 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4393 queue_wait: Duration::from_millis(queue_wait_ms),
4394 })
4395 }
4396 (Err(_), Err(_), Err(_)) => None,
4397 (Err(_), Ok(_), Err(_)) => panic!(
4398 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4399 (or unset all three to disable KV cache pooling)"
4400 ),
4401 (Ok(_), Ok(_), Ok(_)) => {
4402 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4403 }
4404 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4405 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4406 set together (or neither, to disable KV cache pooling)"
4407 ),
4408 };
4409 let paged_kv = match (
4415 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4416 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4417 ) {
4418 (Ok(blocks), Ok(block_size)) => {
4419 assert!(
4420 kv_pool.is_none(),
4421 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4422 mutually exclusive: both bound the same KV memory, by different means. \
4423 Set one."
4424 );
4425 let blocks_per_layer: usize = blocks
4439 .parse()
4440 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4441 let block_size: usize = block_size
4442 .parse()
4443 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4444 let gguf = match &loaded {
4445 model::LoadedModel::Gguf(g) => g,
4446 _ => panic!(
4447 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4448 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4449 ),
4450 };
4451 let cfg = &gguf.decoder.config;
4452 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4453 .ok()
4454 .map(|v| {
4455 v.parse()
4456 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4457 })
4458 .unwrap_or(0);
4459 tracing::info!(
4460 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4461 layer across {} layers, shared by all concurrent requests, \
4462 {queue_wait_ms}ms admission queue wait",
4463 cfg.n_layers
4464 );
4465 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4470 block_size,
4471 ))));
4472 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4481 crate::policy::parser::ToolCallFormat::infer(
4482 &std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
4483 )
4484 .opener(),
4485 |text| {
4486 gguf.tokenizer
4487 .encode(text)
4488 .into_iter()
4489 .map(|t| t as u32)
4490 .collect()
4491 },
4492 );
4493 if let Some(id) = anchor_token {
4494 tracing::info!(
4495 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4496 window stops short of where its next turn rejoins"
4497 );
4498 }
4499 let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
4500 .ok()
4501 .map(|v| {
4502 v.parse()
4503 .expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4504 })
4505 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4506 if let Some(window) = cfg.uniform_sliding_window() {
4507 tracing::info!(
4508 "Paged KV window slide enabled: every layer slides by {window} every \
4509 {slide_interval} decode steps, so a request holds its prompt and a \
4510 window rather than its whole context"
4511 );
4512 } else if cfg.kv_block_window().is_some() {
4513 tracing::info!(
4514 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4515 and a page group holds one block in every layer"
4516 );
4517 }
4518 Some(generate::PagedKvConfig {
4519 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4520 cfg.n_layers,
4521 block_size,
4522 blocks_per_layer,
4523 cfg.n_kv_heads,
4524 cfg.head_dim,
4525 )),
4526 queue_wait: Duration::from_millis(queue_wait_ms),
4527 radix,
4528 anchor_token,
4529 slide_interval,
4530 })
4531 }
4532 (Err(_), Err(_)) => None,
4533 _ => panic!(
4534 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4535 (or neither, to disable paged KV)"
4536 ),
4537 };
4538 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4545 let max_entries: usize = v
4546 .parse()
4547 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4548 if kv_pool.is_some() {
4549 tracing::warn!(
4550 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4551 caching will never be consulted while a KV pool is configured"
4552 );
4553 }
4554 assert!(
4561 paged_kv.is_none(),
4562 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4563 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4564 produce, so the cache could never hit. Set one."
4565 );
4566 tracing::info!(
4567 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4568 all requests"
4569 );
4570 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4571 });
4572 if matches!(
4573 loaded,
4574 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4575 ) && (kv_pool.is_some() || prefix_cache.is_some())
4576 {
4577 tracing::warn!(
4578 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4579 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4580 ferrox_models::engine's module docs"
4581 );
4582 }
4583 let enable_cb =
4584 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4585 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4586 tracing::info!(
4587 "continuous batching enabled by default on Metal for safe parallel serving \
4588 (set FERROX_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4589 );
4590 }
4591 if continuous_batching_env() == Some(true)
4592 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4593 && (kv_pool.is_some() || prefix_cache.is_some())
4594 {
4595 tracing::warn!(
4596 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4597 (those modes keep the private generate path)"
4598 );
4599 }
4600 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4601 if let Ok(chunk) = n.parse::<usize>() {
4602 if chunk > 0 {
4603 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4604 }
4605 }
4606 }
4607 if matches!(
4608 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4609 Some("1")
4610 ) {
4611 tracing::warn!(
4612 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4613 (minimal spill; full layer offload still planned)"
4614 );
4615 }
4616
4617 let mcp = match mcp_config_path {
4618 Some(path) => {
4619 let loaded = mcp::load_mcp_config(&path)?;
4620 tracing::info!(
4621 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4622 loaded.path,
4623 loaded.servers.len()
4624 );
4625 Some(loaded)
4626 }
4627 None => None,
4628 };
4629
4630 let detection = health::Detection::spawn();
4634
4635 let state = Arc::new(build_app_state(
4636 StartupModels {
4637 loaded,
4638 embedding: embedding_model,
4639 },
4640 kv_pool,
4641 paged_kv,
4642 prefix_cache,
4643 enable_cb,
4644 mcp,
4645 detection,
4646 ));
4647
4648 use ferrox_api::routes;
4652
4653 let public = Router::new().route(routes::HEALTH, get(health));
4658
4659 let mut protected = Router::new()
4660 .route(routes::V1_MODELS, get(list_models))
4661 .route(routes::V1_RESPONSES, post(responses::responses))
4665 .route(
4666 &axum_path(routes::V1_RESPONSE),
4667 get(responses::responses_get),
4668 )
4669 .route(
4670 &axum_path(routes::V1_RESPONSE_CANCEL),
4671 post(responses::responses_cancel),
4672 )
4673 .route(routes::V1_STATS, get(serving_stats))
4674 .route(routes::V1_REQUESTS, get(recent_requests))
4675 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
4676 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
4677 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
4678 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
4679 .route(routes::V1_CANCEL, post(cancel_generation))
4683 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
4688 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
4689 .route(routes::V1_MESSAGES, post(anthropic::messages))
4690 .route(
4691 routes::V1_MESSAGES_COUNT_TOKENS,
4692 post(anthropic::count_tokens),
4693 )
4694 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
4695 .route(routes::COMPLETION, post(completion::completion))
4700 .route(routes::COMPLETIONS, post(completion::completion))
4701 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
4702 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
4703 .route(routes::TOKENIZE, post(openai_extra::tokenize))
4709 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
4710 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
4711 .route(routes::V1_RERANK, post(rerank::rerank))
4715 .route(routes::RERANK, post(rerank::rerank))
4716 .route(routes::CACHE_STATS, get(cache_stats))
4717 .route(routes::METRICS, get(metrics))
4718 .route(routes::ADMIN_MODELS, get(admin::models))
4723 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
4724 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
4725 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
4726 .route(routes::ADMIN_TASKS, get(admin::tasks))
4727 .route(&admin::cancel_route(), post(admin::cancel_task))
4728 .route(routes::ADMIN_STATS, get(admin::stats))
4729 .merge(conversations::router());
4733
4734 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4739 tracing::info!("API key auth enabled");
4740 let auth = limits::AuthConfig {
4741 api_key: Arc::new(key),
4742 };
4743 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4744 auth,
4745 limits::require_api_key,
4746 ));
4747 }
4748 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4749 let rpm: u32 = rpm
4750 .parse()
4751 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4752 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4753 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4754 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4755 limiter,
4756 limits::rate_limit,
4757 ));
4758 }
4759 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
4767 let origins = security::parse_cors_origins(&spec)
4768 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
4769 tracing::info!(
4770 "CORS enabled: {} allow-listed origin(s) ({})",
4771 origins.len(),
4772 spec
4773 );
4774 let cors = tower_http::cors::CorsLayer::new()
4775 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4776 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4777 .allow_headers([
4778 axum::http::header::CONTENT_TYPE,
4779 axum::http::header::AUTHORIZATION,
4780 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4786 axum::http::HeaderName::from_static("last-event-id"),
4790 ]);
4791 protected = protected.route_layer(cors);
4792 }
4793
4794 let app = public
4798 .merge(protected)
4799 .layer(axum::middleware::from_fn(limits::retry_after))
4800 .with_state(state);
4801
4802 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4808 install_ring_crypto_provider();
4809 match tls_paths {
4814 Some(paths) => {
4815 let config =
4816 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4817 .await
4818 .map_err(|e| {
4819 anyhow::anyhow!(
4820 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4821 paths.cert,
4822 paths.key
4823 )
4824 })?;
4825 let socket_addr: std::net::SocketAddr = addr
4826 .parse()
4827 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
4828 let listener = std::net::TcpListener::bind(socket_addr)?;
4829 listener.set_nonblocking(true)?;
4837 let bound = listener.local_addr()?;
4838 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
4839 announce_ready(bound, "https");
4840
4841 let handle = axum_server::Handle::new();
4842 let shutdown_handle = handle.clone();
4843 tokio::spawn(async move {
4844 shutdown_signal(exit_on_stdin_close).await;
4845 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4846 });
4847 axum_server::from_tcp_rustls(listener, config)?
4848 .handle(handle)
4849 .serve(app.into_make_service())
4850 .await?;
4851 }
4852 None => {
4853 let listener = tokio::net::TcpListener::bind(&addr).await?;
4854 let bound = listener.local_addr()?;
4855 tracing::info!("ferrox-server listening on {bound}");
4856 announce_ready(bound, "http");
4857 axum::serve(listener, app)
4858 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4859 .await?;
4860 }
4861 }
4862 Ok(())
4863}
4864
4865#[cfg(test)]
4866mod tests {
4867 use super::*;
4868 use ferrox_models::config::test_dense_fixture;
4869
4870 #[test]
4871 fn parses_llama_server_style_options() {
4872 let argv = [
4873 "ferrox-server",
4874 "-m",
4875 "model.gguf",
4876 "--host",
4877 "::1",
4878 "--port",
4879 "9000",
4880 "-t",
4881 "4",
4882 "-dev",
4883 "Metal",
4884 "-ngl",
4885 "all",
4886 ]
4887 .into_iter()
4888 .map(String::from)
4889 .collect();
4890 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4891
4892 assert_eq!(args.model.as_deref(), Some("model.gguf"));
4893 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
4894 assert_eq!(args.port, Some(9000));
4895 assert_eq!(args.threads, Some(4));
4896 assert_eq!(args.device, Some(OffloadDevice::Metal));
4897 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
4898 assert_eq!(
4899 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4900 Some("[::1]:9000")
4901 );
4902 }
4903
4904 #[test]
4905 fn port_zero_survives_argument_parsing_as_a_real_request() {
4906 let argv = ["ferrox-server", "--port", "0"]
4911 .into_iter()
4912 .map(String::from)
4913 .collect();
4914 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4915 assert_eq!(args.port, Some(0));
4916 assert_eq!(
4917 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4918 Some("127.0.0.1:0")
4919 );
4920 }
4921
4922 #[test]
4923 fn parallel_flag_parses_and_rewrites_np() {
4924 let argv = ["ferrox-server", "-np", "4"]
4925 .into_iter()
4926 .map(String::from)
4927 .collect();
4928 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4929 assert_eq!(args.parallel, Some(4));
4930 }
4931
4932 #[test]
4933 fn stdin_close_exit_is_opt_in() {
4934 let args =
4937 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
4938 assert!(!args.exit_on_stdin_close);
4939 let args = ServerArgs::try_parse_from(
4940 ["ferrox-server", "--exit-on-stdin-close"]
4941 .into_iter()
4942 .map(String::from),
4943 )
4944 .unwrap();
4945 assert!(args.exit_on_stdin_close);
4946 }
4947
4948 #[test]
4949 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4950 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4951 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4952 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
4953 assert_eq!(parsed.port, 51999);
4954 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4955 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
4957 }
4958
4959 fn test_model() -> Model {
4960 let cfg = test_dense_fixture();
4964 Model::Gguf(GgufModel {
4965 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4966 tokenizer: Arc::new(ServerTokenizer::Byte),
4967 stop_tokens: StopTokens::default(),
4968 bos_id: None,
4969 is_synthetic: true,
4970 chat_template: chat_template::PromptTemplate::plain(),
4971 })
4972 }
4973
4974 fn greedy_params(max_tokens: usize) -> GenerationParams {
4975 GenerationParams {
4976 max_tokens,
4977 sampling: SamplingParams::default(),
4978 seed: 1,
4979 stop: Vec::new(),
4980 stop_token_ids: Vec::new(),
4981 json_object: false,
4982 grammar: None,
4983 cancel: None,
4984 ignore_eos: false,
4985 }
4986 }
4987
4988 fn test_model_full_byte_vocab() -> Model {
4992 let mut cfg = test_dense_fixture();
4993 cfg.vocab_size = 256;
4994 Model::Gguf(GgufModel {
4995 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4996 tokenizer: Arc::new(ServerTokenizer::Byte),
4997 stop_tokens: StopTokens::default(),
4998 bos_id: None,
4999 is_synthetic: true,
5000 chat_template: chat_template::PromptTemplate::plain(),
5001 })
5002 }
5003
5004 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
5008 AppState {
5009 embedding: None,
5010 paged_kv: None,
5011 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
5012 id: None,
5013 loaded: Loaded::Generative(Arc::new(model)),
5014 batcher: None,
5015 ceiling: None,
5016 }))),
5017 load_in_progress: std::sync::atomic::AtomicBool::new(false),
5018 tasks: Arc::new(tasks::TaskRegistry::new()),
5019 cancels: Arc::new(cancel::CancelRegistry::new()),
5020 stats: stats::Stats::new(),
5021 streams: resume::StreamRegistry::new(),
5022 model_dir: None,
5023 response_cache: Mutex::new(response_cache),
5024 kv_pool: None,
5025 prefix_cache: None,
5026 sessions: session::SessionStore::new(),
5027 requests_total: std::sync::atomic::AtomicU64::new(0),
5028 request_errors_total: std::sync::atomic::AtomicU64::new(0),
5029 started_at: std::time::Instant::now(),
5030 last_request_ms: std::sync::atomic::AtomicU64::new(0),
5031 detection: Arc::new(health::Detection::ready(health::probe_backends())),
5032 mcp: None,
5033 continuous_batching_enabled: false,
5034 metal_private_decode_gate: None,
5035 loading_model: Mutex::new(None),
5036 last_load_error: Mutex::new(None),
5037 serving: Mutex::new(crate::stats::ServingStats::default()),
5038 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
5039 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
5040 started_unix: unix_now(),
5041 }
5042 }
5043
5044 fn test_app() -> Router {
5053 test_app_with_state(Arc::new(test_state(
5054 test_model_full_byte_vocab(),
5055 ResponseCache::new(1000, Duration::from_secs(3600)),
5056 )))
5057 }
5058
5059 fn test_app_with_state(state: Arc<AppState>) -> Router {
5062 Router::new()
5063 .route(ferrox_api::routes::HEALTH, get(health))
5064 .route(ferrox_api::routes::V1_MODELS, get(list_models))
5065 .route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
5066 .route(
5067 &axum_path(ferrox_api::routes::V1_RESPONSE),
5068 get(responses::responses_get),
5069 )
5070 .route(
5071 &axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
5072 post(responses::responses_cancel),
5073 )
5074 .route(ferrox_api::routes::V1_STATS, get(serving_stats))
5075 .route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
5076 .route(
5077 ferrox_api::routes::V1_CACHE_STATUS,
5078 get(cache_admin::cache_status),
5079 )
5080 .route(
5081 ferrox_api::routes::V1_CACHE_REBUILD,
5082 post(cache_admin::cache_rebuild),
5083 )
5084 .route(
5085 ferrox_api::routes::ADMIN_PREPARE_STOP,
5086 post(cache_admin::prepare_stop),
5087 )
5088 .route("/v1/chat/completions", post(chat_completions))
5089 .route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
5090 .route(
5091 ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
5092 post(anthropic::count_tokens),
5093 )
5094 .route("/v1/tokenize", post(openai_extra::tokenize))
5095 .route("/v1/detokenize", post(openai_extra::detokenize))
5096 .route(ferrox_api::routes::TOKENIZE, post(openai_extra::tokenize))
5100 .route(
5101 ferrox_api::routes::DETOKENIZE,
5102 post(openai_extra::detokenize),
5103 )
5104 .route("/v1/embeddings", post(embeddings::embeddings))
5105 .route("/v1/completions", post(openai_extra::completions))
5106 .route(ferrox_api::routes::COMPLETION, post(completion::completion))
5108 .route(
5109 ferrox_api::routes::COMPLETIONS,
5110 post(completion::completion),
5111 )
5112 .route(
5113 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5114 post(admin::unload_model),
5115 )
5116 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
5117 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
5118 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
5119 .route(
5120 &axum_path(ferrox_api::routes::V1_STREAM),
5121 get(resume::resume),
5122 )
5123 .route(
5124 &axum_path(ferrox_api::routes::V1_STREAM_POLL),
5125 get(resume::poll),
5126 )
5127 .with_state(state)
5128 }
5129
5130 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
5131 let mut cfg = test_dense_fixture();
5132 cfg.name = name;
5133 cfg.vocab_size = vocab_size;
5134 Model::Gguf(GgufModel {
5135 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5136 tokenizer: Arc::new(ServerTokenizer::Byte),
5137 stop_tokens: StopTokens::default(),
5138 bos_id: None,
5139 is_synthetic: true,
5140 chat_template: chat_template::PromptTemplate::plain(),
5141 })
5142 }
5143
5144 fn model_with_template(name: &'static str, source: &str) -> Model {
5148 let mut cfg = test_dense_fixture();
5149 cfg.name = name;
5150 cfg.vocab_size = 256;
5151 Model::Gguf(GgufModel {
5152 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5153 tokenizer: Arc::new(ServerTokenizer::Byte),
5154 stop_tokens: StopTokens::default(),
5155 bos_id: None,
5156 is_synthetic: true,
5157 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5158 Some(source),
5159 Some("qwen3"),
5160 false,
5161 true,
5162 None,
5163 None,
5164 ),
5165 })
5166 }
5167
5168 #[tokio::test]
5177 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5178 let strict = "{% if messages | length > 1 %}\
5181 {{ raise_exception('this template takes one turn') }}\
5182 {% endif %}{{ messages[0].content }}";
5183 let state = Arc::new(test_state(
5184 model_with_template("strict", strict),
5185 ResponseCache::new(4, Duration::from_secs(60)),
5186 ));
5187 let app = test_app_with_state(state);
5188
5189 let (status, body) = post_json_uri(
5190 &app,
5191 "/v1/chat/completions",
5192 serde_json::json!({
5193 "model": "strict",
5194 "stream": true,
5195 "messages": [
5196 {"role": "user", "content": "one"},
5197 {"role": "user", "content": "two"},
5198 ],
5199 }),
5200 )
5201 .await;
5202 assert_eq!(status, StatusCode::BAD_REQUEST);
5203 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5204 assert!(
5205 body["error"]["message"]
5206 .as_str()
5207 .unwrap()
5208 .contains("one turn"),
5209 "the template's own message must reach the caller: {body}"
5210 );
5211
5212 let (status, _) = post_json_uri(
5214 &app,
5215 "/v1/chat/completions",
5216 serde_json::json!({
5217 "model": "strict",
5218 "stream": true,
5219 "max_tokens": 1,
5220 "messages": [{"role": "user", "content": "one"}],
5221 }),
5222 )
5223 .await;
5224 assert_eq!(status, StatusCode::OK);
5225 }
5226
5227 #[tokio::test]
5229 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5230 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5231 {% if reasoning_effort %}\
5232 {% if reasoning_effort not in ['low','medium','high'] %}\
5233 {{ raise_exception('bad effort') }}\
5234 {% endif %}[{{ reasoning_effort }}]\
5235 {% endif %}{{ messages[0].content }}";
5236 let state = Arc::new(test_state(
5237 model_with_template("thinker", reasoning),
5238 ResponseCache::new(4, Duration::from_secs(60)),
5239 ));
5240 let app = test_app_with_state(state);
5241 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5242 assert_eq!(status, StatusCode::OK);
5243 let entry = &models["data"][0];
5244 assert_eq!(
5245 entry["supported_reasoning_efforts"],
5246 serde_json::json!(["off", "low", "medium", "high"])
5247 );
5248 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5249 }
5250
5251 #[tokio::test]
5256 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5257 let app = test_app();
5258 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5259 let entry = &models["data"][0];
5260 assert!(entry.get("supported_reasoning_efforts").is_none());
5261 assert!(entry.get("default_reasoning_effort").is_none());
5262 }
5263
5264 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5265 Arc::new(ActiveModel {
5266 id: Some(name.to_string()),
5267 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5268 batcher: None,
5269 ceiling: None,
5270 })
5271 .tap_into(state)
5272 }
5273
5274 trait TapInto {
5276 fn tap_into(self, state: &AppState) -> Self;
5277 }
5278 impl TapInto for Arc<ActiveModel> {
5279 fn tap_into(self, state: &AppState) -> Self {
5280 state.swap_active(Some(Arc::clone(&self)));
5281 self
5282 }
5283 }
5284
5285 #[test]
5291 fn an_in_flight_request_keeps_the_model_it_started_on() {
5292 let state = test_state(
5293 named_test_model("model-a", 256),
5294 ResponseCache::new(4, Duration::from_secs(60)),
5295 );
5296
5297 let in_flight = state.active().expect("a model is loaded");
5300 assert_eq!(in_flight.name(), "model-a");
5301
5302 active_model(&state, "model-b");
5303
5304 assert_eq!(state.active().unwrap().name(), "model-b");
5306 assert_eq!(in_flight.name(), "model-a");
5308 let (_chunks, finish, _usage) = run_generation(
5309 in_flight.generative().unwrap(),
5310 "hi",
5311 &greedy_params(3),
5312 None,
5313 None,
5314 None,
5315 None,
5316 None,
5317 None,
5318 )
5319 .expect("the old model must still decode after being swapped out");
5320 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5321 }
5322
5323 #[test]
5328 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5329 let state = test_state(
5330 named_test_model("model-a", 256),
5331 ResponseCache::new(4, Duration::from_secs(60)),
5332 );
5333 let in_flight = state.active().expect("a model is loaded");
5334 let weights = Arc::clone(in_flight.generative().unwrap());
5335 assert!(Arc::strong_count(&weights) >= 2);
5336
5337 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5338 id: Some("model-b".to_string()),
5339 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5340 batcher: None,
5341 ceiling: None,
5342 })));
5343 drop(previous);
5344 assert!(Arc::strong_count(&weights) >= 2);
5346 drop(in_flight);
5347 assert_eq!(Arc::strong_count(&weights), 1);
5348 }
5349
5350 #[tokio::test]
5354 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5355 let state = Arc::new(test_state(
5356 named_test_model("model-a", 256),
5357 ResponseCache::new(4, Duration::from_secs(60)),
5358 ));
5359 let app = test_app_with_state(Arc::clone(&state));
5360
5361 let (status, body) = post_json_uri(
5362 &app,
5363 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5364 serde_json::json!({}),
5365 )
5366 .await;
5367 assert_eq!(status, StatusCode::OK);
5368 assert_eq!(body["ok"], true);
5369 assert!(body["active"].is_null());
5370 assert!(state.active().is_none());
5371
5372 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5373 assert_eq!(status, StatusCode::OK);
5374 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5375 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5376
5377 let (status, body) = post_json_uri(
5378 &app,
5379 "/v1/chat/completions",
5380 serde_json::json!({
5381 "model": "x",
5382 "messages": [{"role": "user", "content": "hi"}]
5383 }),
5384 )
5385 .await;
5386 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5387 assert_eq!(body["error"]["type"], "model_not_loaded");
5388 }
5389
5390 #[tokio::test]
5394 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5395 let state = Arc::new(test_state(
5396 named_test_model("model-a", 256),
5397 ResponseCache::new(4, Duration::from_secs(60)),
5398 ));
5399 let app = test_app_with_state(Arc::clone(&state));
5400 state.swap_active(None);
5401
5402 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5403 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5406 assert_eq!(body["state"], "unavailable");
5407 assert_eq!(body["reason"], "model_not_loaded");
5408 assert!(body["model"].is_null());
5409 let real_weights = body["capabilities"]
5410 .as_array()
5411 .unwrap()
5412 .iter()
5413 .find(|c| c["id"] == "real_weights")
5414 .cloned()
5415 .expect("real_weights is always reported");
5416 assert_eq!(real_weights["available"], false);
5417 assert_eq!(real_weights["reason"], "model_not_loaded");
5418 }
5419
5420 #[tokio::test]
5424 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5425 let app = test_app();
5426
5427 let (status, completion) = post_json_uri(
5428 &app,
5429 "/v1/chat/completions",
5430 serde_json::json!({
5431 "model": "x",
5432 "messages": [{"role": "user", "content": "hi"}],
5433 "max_tokens": 4
5434 }),
5435 )
5436 .await;
5437 assert_eq!(status, StatusCode::OK);
5438 let request_id = completion["request_id"].as_str().unwrap().to_string();
5439
5440 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5441 assert_eq!(status, StatusCode::OK);
5442 let recent = stats["recent"].as_array().unwrap();
5443 assert_eq!(recent.len(), 1);
5444 let row = &recent[0];
5445 assert_eq!(row["request_id"], request_id);
5446 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
5447 assert_eq!(row["status"], 200);
5448 assert_eq!(row["stream"], false);
5449 assert!(row["duration_ms"].is_number());
5452 assert!(row["decode_ms"].is_number());
5453 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5454 assert_eq!(
5455 stats["tokens_prompt_total"].as_u64().unwrap(),
5456 row["prompt_tokens"].as_u64().unwrap()
5457 );
5458 }
5459
5460 #[tokio::test]
5464 async fn a_rejected_request_is_recorded_too() {
5465 let state = Arc::new(test_state(
5466 named_test_model("model-a", 256),
5467 ResponseCache::new(4, Duration::from_secs(60)),
5468 ));
5469 let app = test_app_with_state(Arc::clone(&state));
5470 state.swap_active(None);
5471
5472 let (status, _) = post_json_uri(
5473 &app,
5474 "/v1/chat/completions",
5475 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5476 )
5477 .await;
5478 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5479
5480 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5481 let recent = stats["recent"].as_array().unwrap();
5482 assert_eq!(recent.len(), 1);
5483 assert_eq!(recent[0]["status"], 503);
5484 assert_eq!(recent[0]["completion_tokens"], 0);
5485 assert!(recent[0]["decode_ms"].is_null());
5486 assert_eq!(stats["errors_total"], 1);
5487 }
5488
5489 async fn post_json_with_headers(
5493 app: &Router,
5494 uri: &str,
5495 body: serde_json::Value,
5496 headers: &[(&str, &str)],
5497 ) -> (StatusCode, serde_json::Value) {
5498 use http_body_util::BodyExt;
5499 use tower::ServiceExt;
5500
5501 let mut builder = axum::http::Request::builder()
5502 .method("POST")
5503 .uri(uri)
5504 .header("content-type", "application/json");
5505 for (name, value) in headers {
5506 builder = builder.header(*name, *value);
5507 }
5508 let response = app
5509 .clone()
5510 .oneshot(
5511 builder
5512 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5513 .unwrap(),
5514 )
5515 .await
5516 .unwrap();
5517 let status = response.status();
5518 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5519 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5520 (status, json)
5521 }
5522
5523 #[tokio::test]
5527 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5528 let app = test_app();
5529
5530 let (status, _) = post_json_uri(
5531 &app,
5532 ferrox_api::routes::V1_TOKENIZE,
5533 serde_json::json!({"prompt": "hello"}),
5534 )
5535 .await;
5536 assert_eq!(status, StatusCode::OK);
5537 let (status, _) = post_json_uri(
5538 &app,
5539 ferrox_api::routes::V1_DETOKENIZE,
5540 serde_json::json!({"tokens": [104, 105]}),
5541 )
5542 .await;
5543 assert_eq!(status, StatusCode::OK);
5544 let (status, _) = post_json_uri(
5545 &app,
5546 ferrox_api::routes::V1_EMBEDDINGS,
5547 serde_json::json!({"input": "hello"}),
5548 )
5549 .await;
5550 assert_eq!(status, StatusCode::OK);
5551
5552 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5553 let routes: Vec<&str> = stats["recent"]
5554 .as_array()
5555 .unwrap()
5556 .iter()
5557 .map(|row| row["route"].as_str().unwrap())
5558 .collect();
5559 for expected in [
5560 ferrox_api::routes::V1_TOKENIZE,
5561 ferrox_api::routes::V1_DETOKENIZE,
5562 ferrox_api::routes::V1_EMBEDDINGS,
5563 ] {
5564 assert!(
5565 routes.contains(&expected),
5566 "{expected} is missing: {routes:?}"
5567 );
5568 }
5569
5570 let row = |route: &str| {
5571 stats["recent"]
5572 .as_array()
5573 .unwrap()
5574 .iter()
5575 .find(|r| r["route"] == route)
5576 .cloned()
5577 .unwrap()
5578 };
5579 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5583 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5584 assert!(embed["decode_ms"].is_null());
5585 assert_eq!(embed["completion_tokens"], 0);
5586 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5590 assert_eq!(
5591 stats["tokens_prompt_total"].as_u64().unwrap(),
5592 embed["prompt_tokens"].as_u64().unwrap(),
5593 "only the forward pass counted"
5594 );
5595 }
5596
5597 fn streaming_test_app() -> Router {
5602 let mut cfg = test_dense_fixture();
5603 cfg.vocab_size = 256;
5604 let model = Model::Gguf(GgufModel {
5605 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5606 tokenizer: Arc::new(ServerTokenizer::Byte),
5607 stop_tokens: StopTokens::default(),
5608 bos_id: None,
5609 is_synthetic: false,
5610 chat_template: chat_template::PromptTemplate::plain(),
5611 });
5612 test_app_with_state(Arc::new(test_state(
5613 model,
5614 ResponseCache::new(1000, Duration::from_secs(3600)),
5615 )))
5616 }
5617
5618 #[tokio::test]
5622 async fn the_native_completion_wire_is_not_the_openai_one() {
5623 let app = test_app();
5624
5625 let (status, native) = post_json_uri(
5626 &app,
5627 ferrox_api::routes::COMPLETION,
5628 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5629 )
5630 .await;
5631 assert_eq!(status, StatusCode::OK, "{native}");
5632 assert!(native["content"].is_string(), "{native}");
5633 assert_eq!(native["stop"], true);
5634 assert_eq!(native["stop_type"], "limit");
5635 assert_eq!(native["stopping_word"], "");
5636 assert_eq!(native["truncated"], false);
5637 assert_eq!(native["id_slot"], -1);
5638 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5639 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5640 assert!(
5641 native.get("choices").is_none(),
5642 "the native shape has no `choices`: {native}"
5643 );
5644
5645 let (status, openai) = post_json_uri(
5646 &app,
5647 ferrox_api::routes::V1_COMPLETIONS,
5648 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5649 )
5650 .await;
5651 assert_eq!(status, StatusCode::OK);
5652 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5653 assert!(
5654 openai.get("content").is_none(),
5655 "the OpenAI shape has no top-level `content`: {openai}"
5656 );
5657 }
5658
5659 #[tokio::test]
5663 async fn both_native_spellings_reach_the_same_handler() {
5664 let app = test_app();
5665 for route in [
5666 ferrox_api::routes::COMPLETION,
5667 ferrox_api::routes::COMPLETIONS,
5668 ] {
5669 let (status, body) = post_json_uri(
5670 &app,
5671 route,
5672 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5673 )
5674 .await;
5675 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5676 assert_eq!(body["stop"], true, "{route}");
5677 assert!(body["content"].is_string(), "{route}");
5678 }
5679
5680 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5683 let routes: Vec<&str> = stats["recent"]
5684 .as_array()
5685 .unwrap()
5686 .iter()
5687 .map(|row| row["route"].as_str().unwrap())
5688 .collect();
5689 assert!(
5690 routes.contains(&ferrox_api::routes::COMPLETION),
5691 "{routes:?}"
5692 );
5693 assert!(
5694 routes.contains(&ferrox_api::routes::COMPLETIONS),
5695 "{routes:?}"
5696 );
5697 }
5698
5699 #[tokio::test]
5705 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5706 let app = streaming_test_app();
5707 let raw = post_sse_raw_uri(
5708 &app,
5709 ferrox_api::routes::COMPLETION,
5710 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5711 )
5712 .await;
5713
5714 assert!(
5715 !raw.contains("[DONE]"),
5716 "llama.cpp's native stream has no sentinel: {raw}"
5717 );
5718 let frames: Vec<serde_json::Value> = raw
5719 .lines()
5720 .filter_map(|line| line.strip_prefix("data: "))
5721 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5722 .collect();
5723 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5724
5725 let (last, partials) = frames.split_last().unwrap();
5726 assert_eq!(last["stop"], true, "the last frame closes the stream");
5727 assert!(last["timings"].is_object(), "{last}");
5728 assert!(last["stop_type"].is_string(), "{last}");
5729 for partial in partials {
5730 assert_eq!(partial["stop"], false, "{partial}");
5731 assert!(partial["content"].is_string(), "{partial}");
5732 assert!(partial.get("timings").is_none(), "{partial}");
5736 assert!(partial.get("generation_settings").is_none(), "{partial}");
5737 }
5738 let streamed: String = partials
5741 .iter()
5742 .filter_map(|p| p["content"].as_str())
5743 .collect();
5744 assert_eq!(last["content"].as_str().unwrap(), streamed);
5745 }
5746
5747 #[tokio::test]
5752 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5753 let app = test_app();
5754 for body in [
5755 serde_json::json!({"prompt": "hi"}),
5756 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5757 ] {
5758 let (status, refusal) =
5759 post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
5760 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5761 assert!(
5762 refusal["error"]["message"]
5763 .as_str()
5764 .unwrap()
5765 .contains("n_predict"),
5766 "{refusal}"
5767 );
5768 }
5769 let (status, _) = post_json_uri(
5772 &app,
5773 ferrox_api::routes::COMPLETION,
5774 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5775 )
5776 .await;
5777 assert_eq!(status, StatusCode::OK);
5778 }
5779
5780 #[tokio::test]
5789 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5790 let app = streaming_test_app();
5791 let ask = |stop: serde_json::Value| {
5792 let app = app.clone();
5793 async move {
5794 post_json_uri(
5795 &app,
5796 ferrox_api::routes::COMPLETION,
5797 serde_json::json!({
5798 "prompt": "hi",
5799 "n_predict": 64,
5800 "ignore_eos": true,
5801 "stop": stop,
5802 }),
5803 )
5804 .await
5805 .1
5806 }
5807 };
5808
5809 let baseline = ask(serde_json::json!([])).await;
5810 assert_eq!(baseline["stop_type"], "limit");
5811 assert_eq!(baseline["stopping_word"], "");
5812 let text = baseline["content"].as_str().unwrap().to_string();
5813 let sentinel: String = text.chars().skip(1).take(2).collect();
5819 assert_eq!(
5820 sentinel.chars().count(),
5821 2,
5822 "the fixture must produce enough output to cut: {text:?}"
5823 );
5824 let cut = text.find(&sentinel).expect("it came out of this text");
5825
5826 let stopped = ask(serde_json::json!([sentinel])).await;
5827 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5828 assert_eq!(stopped["stopping_word"], sentinel);
5829 assert_eq!(
5830 stopped["content"].as_str().unwrap(),
5831 &text[..cut],
5832 "the answer must be cut at the sentinel, not run past it"
5833 );
5834 }
5835
5836 #[tokio::test]
5842 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5843 let app = test_app();
5844
5845 let (v1_status, v1) = post_json_uri(
5846 &app,
5847 ferrox_api::routes::V1_TOKENIZE,
5848 serde_json::json!({"prompt": "hello"}),
5849 )
5850 .await;
5851 let (alias_status, alias) = post_json_uri(
5852 &app,
5853 ferrox_api::routes::TOKENIZE,
5854 serde_json::json!({"content": "hello"}),
5855 )
5856 .await;
5857 assert_eq!(v1_status, StatusCode::OK);
5858 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5859 assert_eq!(v1["tokens"], alias["tokens"]);
5860 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5861
5862 let (status, both_ways) = post_json_uri(
5865 &app,
5866 ferrox_api::routes::TOKENIZE,
5867 serde_json::json!({"prompt": "hello"}),
5868 )
5869 .await;
5870 assert_eq!(status, StatusCode::OK);
5871 assert_eq!(both_ways["tokens"], v1["tokens"]);
5872 }
5873
5874 #[tokio::test]
5879 async fn detokenize_answers_under_both_dialects_keys() {
5880 let app = test_app();
5881 for route in [
5882 ferrox_api::routes::DETOKENIZE,
5883 ferrox_api::routes::V1_DETOKENIZE,
5884 ] {
5885 let (status, body) =
5886 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5887 assert_eq!(status, StatusCode::OK, "{route}");
5888 assert_eq!(body["text"], "hi", "{route}");
5889 assert_eq!(body["content"], body["text"], "{route}");
5890 }
5891 }
5892
5893 #[tokio::test]
5897 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5898 let app = test_app();
5899 let (status, _) = post_json_uri(
5900 &app,
5901 ferrox_api::routes::TOKENIZE,
5902 serde_json::json!({"content": "hello"}),
5903 )
5904 .await;
5905 assert_eq!(status, StatusCode::OK);
5906
5907 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5908 let routes: Vec<&str> = stats["recent"]
5909 .as_array()
5910 .unwrap()
5911 .iter()
5912 .map(|row| row["route"].as_str().unwrap())
5913 .collect();
5914 assert!(
5915 routes.contains(&ferrox_api::routes::TOKENIZE),
5916 "the alias must be its own row: {routes:?}"
5917 );
5918 assert!(
5919 !routes.contains(&ferrox_api::routes::V1_TOKENIZE),
5920 "nothing called /v1/tokenize: {routes:?}"
5921 );
5922 }
5923
5924 #[tokio::test]
5929 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5930 let mut cfg = test_dense_fixture();
5931 cfg.vocab_size = 256;
5932 let model = Model::Gguf(GgufModel {
5933 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5934 tokenizer: Arc::new(ServerTokenizer::Byte),
5935 stop_tokens: StopTokens::default(),
5936 bos_id: Some(7),
5937 is_synthetic: true,
5938 chat_template: chat_template::PromptTemplate::plain(),
5939 });
5940 let app = test_app_with_state(Arc::new(test_state(
5941 model,
5942 ResponseCache::new(1000, Duration::from_secs(3600)),
5943 )));
5944
5945 let (_, plain) = post_json_uri(
5946 &app,
5947 ferrox_api::routes::TOKENIZE,
5948 serde_json::json!({"content": "hi"}),
5949 )
5950 .await;
5951 let (_, special) = post_json_uri(
5952 &app,
5953 ferrox_api::routes::TOKENIZE,
5954 serde_json::json!({"content": "hi", "add_special": true}),
5955 )
5956 .await;
5957
5958 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5959 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5960 assert_eq!(special["count"], 3);
5961 }
5962
5963 #[tokio::test]
5966 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5967 let app = test_app();
5968 let (status, _) = post_json_uri(
5969 &app,
5970 ferrox_api::routes::V1_EMBEDDINGS,
5971 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5972 )
5973 .await;
5974 assert_eq!(status, StatusCode::BAD_REQUEST);
5975
5976 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5977 let recent = stats["recent"].as_array().unwrap();
5978 assert_eq!(recent.len(), 1);
5979 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
5980 assert_eq!(recent[0]["status"], 400);
5981 assert_eq!(
5982 recent[0]["prompt_tokens"], 0,
5983 "a rejected call embedded nothing"
5984 );
5985 }
5986
5987 #[tokio::test]
5990 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5991 let app = test_app();
5992 let key = "sk-monitor-secret";
5993 let (status, _) = post_json_with_headers(
5994 &app,
5995 "/v1/chat/completions",
5996 serde_json::json!({
5997 "model": "x",
5998 "messages": [{"role": "user", "content": "hi"}],
5999 "max_tokens": 2
6000 }),
6001 &[
6002 ("authorization", &format!("Bearer {key}")),
6003 ("x-ferrox-client", "ferrox-studio"),
6004 ],
6005 )
6006 .await;
6007 assert_eq!(status, StatusCode::OK);
6008
6009 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6010 let row = stats["recent"].as_array().unwrap()[0].clone();
6011 let fingerprint = row["via_api_key"]
6012 .as_str()
6013 .expect("the row names the key that served it")
6014 .to_string();
6015 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6016 assert!(!fingerprint.contains(key));
6017 assert!(
6018 !serde_json::to_string(&stats).unwrap().contains(key),
6019 "the stats payload must not carry the key in any form"
6020 );
6021 assert_eq!(row["client"], "ferrox-studio");
6022 }
6023
6024 #[tokio::test]
6027 async fn different_keys_are_different_callers_and_no_key_is_null() {
6028 let app = test_app();
6029 let body = serde_json::json!({
6030 "model": "x",
6031 "messages": [{"role": "user", "content": "hi"}],
6032 "max_tokens": 1
6033 });
6034 for headers in [
6035 vec![("authorization", "Bearer key-one")],
6036 vec![("authorization", "Bearer key-two")],
6037 vec![],
6038 ] {
6039 let (status, _) =
6040 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6041 assert_eq!(status, StatusCode::OK);
6042 }
6043
6044 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6045 let recent = stats["recent"].as_array().unwrap();
6046 assert_eq!(recent.len(), 3);
6047 let one = recent[0]["via_api_key"].as_str().unwrap();
6048 let two = recent[1]["via_api_key"].as_str().unwrap();
6049 assert_ne!(one, two, "two keys must not collapse into one caller");
6050 assert!(
6051 recent[2]["via_api_key"].is_null(),
6052 "an unauthenticated call is null, not a fingerprint of nothing"
6053 );
6054 assert!(recent[2]["client"].is_null());
6055 }
6056
6057 #[tokio::test]
6062 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6063 let state = Arc::new(test_state(
6064 named_test_model("really-loaded", 256),
6065 ResponseCache::new(4, Duration::from_secs(60)),
6066 ));
6067 let app = test_app_with_state(Arc::clone(&state));
6068
6069 let (status, _) = post_json_uri(
6070 &app,
6071 "/v1/chat/completions",
6072 serde_json::json!({
6073 "model": "gpt-4-turbo-that-is-not-here",
6074 "messages": [{"role": "user", "content": "hi"}],
6075 "max_tokens": 2
6076 }),
6077 )
6078 .await;
6079 assert_eq!(status, StatusCode::OK);
6080
6081 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6082 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6083
6084 state.swap_active(None);
6087 let (status, _) = post_json_uri(
6088 &app,
6089 "/v1/chat/completions",
6090 serde_json::json!({
6091 "model": "gpt-4-turbo-that-is-not-here",
6092 "messages": [{"role": "user", "content": "hi"}]
6093 }),
6094 )
6095 .await;
6096 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6097 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6098 let recent = stats["recent"].as_array().unwrap();
6099 assert!(recent[recent.len() - 1]["model"].is_null());
6100 }
6101
6102 #[tokio::test]
6106 async fn a_streamed_row_names_the_model_it_decoded_against() {
6107 let state = Arc::new(test_state(
6108 named_test_model("model-before", 256),
6109 ResponseCache::new(4, Duration::from_secs(60)),
6110 ));
6111 let app = test_app_with_state(Arc::clone(&state));
6112 let _ = post_sse_raw(&app, resumable_request()).await;
6113 active_model(&state, "model-after");
6115
6116 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6117 assert_eq!(stats["recent"][0]["model"], "model-before");
6118 }
6119
6120 #[tokio::test]
6123 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6124 let app = test_app();
6125 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6126 assert_eq!(status, StatusCode::OK);
6127 assert!(
6128 stats["queue_depth"].is_null(),
6129 "without continuous batching nothing queues, so there is nothing to measure"
6130 );
6131 assert!(stats["queue_rejected_total"].is_null());
6132 assert_eq!(
6133 stats["generating_now"], 0,
6134 "work in progress is measured and really is zero here"
6135 );
6136 }
6137
6138 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6143 post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
6144 }
6145
6146 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6150 use http_body_util::BodyExt;
6151 use tower::ServiceExt;
6152
6153 let response = app
6154 .clone()
6155 .oneshot(
6156 axum::http::Request::builder()
6157 .method("POST")
6158 .uri(uri)
6159 .header("content-type", "application/json")
6160 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6161 .unwrap(),
6162 )
6163 .await
6164 .unwrap();
6165 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6166 String::from_utf8(bytes.to_vec()).unwrap()
6167 }
6168
6169 async fn get_json_with_headers(
6170 app: &Router,
6171 uri: &str,
6172 headers: &[(&str, &str)],
6173 ) -> (StatusCode, serde_json::Value) {
6174 use http_body_util::BodyExt;
6175 use tower::ServiceExt;
6176
6177 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6178 for (name, value) in headers {
6179 builder = builder.header(*name, *value);
6180 }
6181 let response = app
6182 .clone()
6183 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6184 .await
6185 .unwrap();
6186 let status = response.status();
6187 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6188 (
6189 status,
6190 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6191 )
6192 }
6193
6194 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6195 body.lines()
6196 .filter_map(|line| line.strip_prefix(field))
6197 .map(str::trim)
6198 .collect()
6199 }
6200
6201 fn resumable_request() -> serde_json::Value {
6202 serde_json::json!({
6203 "model": "m",
6204 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6205 "max_tokens": 4,
6206 "temperature": 0,
6207 "stream": true,
6208 "stream_resumable": true,
6209 })
6210 }
6211
6212 #[tokio::test]
6217 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6218 let app = test_app();
6219 let body = post_sse_raw(&app, resumable_request()).await;
6220
6221 let request_id = body
6222 .lines()
6223 .find_map(|l| l.strip_prefix("data: "))
6224 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6225 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6226 .expect("the first chunk names the request");
6227
6228 let ids = sse_field(&body, "id:");
6229 let datas = sse_field(&body, "data:");
6230 assert_eq!(
6231 ids.len(),
6232 datas.len(),
6233 "every event carries an id, or a reconnect cannot name where it stopped"
6234 );
6235 for (i, id) in ids.iter().enumerate() {
6236 assert_eq!(*id, format!("{request_id}:{i}"));
6237 }
6238 let retries = sse_field(&body, "retry:");
6239 assert_eq!(
6240 retries.len(),
6241 1,
6242 "the reconnect delay is stated once, not on every event"
6243 );
6244 assert_eq!(retries[0], "1500");
6245 assert!(
6246 body.contains("data: [DONE]"),
6247 "the end of stream is still stated"
6248 );
6249 }
6250
6251 #[tokio::test]
6255 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6256 let app = test_app();
6257 let mut request = resumable_request();
6258 request["stream_resumable"] = serde_json::json!(false);
6259 let body = post_sse_raw(&app, request).await;
6260 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6261 assert!(
6262 sse_field(&body, "id:").is_empty(),
6263 "an id promises a replay this stream cannot serve"
6264 );
6265 assert!(sse_field(&body, "retry:").is_empty());
6266 }
6267
6268 #[tokio::test]
6272 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6273 let app = test_app();
6274 let body = post_sse_raw(&app, resumable_request()).await;
6275 let request_id = sse_field(&body, "id:")[0]
6276 .rsplit_once(':')
6277 .unwrap()
6278 .0
6279 .to_string();
6280 let streamed: Vec<String> = sse_field(&body, "data:")
6281 .iter()
6282 .map(|d| d.to_string())
6283 .collect();
6284
6285 let (status, polled) = get_json(
6286 &app,
6287 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
6288 )
6289 .await;
6290 assert_eq!(status, StatusCode::OK);
6291 let events: Vec<String> = polled["events"]
6292 .as_array()
6293 .unwrap()
6294 .iter()
6295 .map(|e| e["data"].as_str().unwrap().to_string())
6296 .collect();
6297 assert_eq!(
6298 events, streamed,
6299 "the fallback must deliver the same answer, not a re-run of it"
6300 );
6301 assert_eq!(polled["request_id"], request_id);
6302 assert_eq!(
6303 polled["done"], false,
6304 "events were still being handed out, so the client must ask again"
6305 );
6306
6307 let next = polled["next_index"].as_u64().unwrap();
6310 let (_, drained) = get_json(
6311 &app,
6312 &format!(
6313 "{}?from={next}",
6314 ferrox_api::routes::v1_stream_poll(&request_id)
6315 ),
6316 )
6317 .await;
6318 assert_eq!(drained["done"], true);
6319 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6320 }
6321
6322 #[tokio::test]
6326 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6327 let app = test_app();
6328 let body = post_sse_raw(&app, resumable_request()).await;
6329 let ids = sse_field(&body, "id:");
6330 let datas: Vec<String> = sse_field(&body, "data:")
6331 .iter()
6332 .map(|d| d.to_string())
6333 .collect();
6334 assert!(
6335 ids.len() >= 3,
6336 "need a few events to resume into the middle"
6337 );
6338 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6339
6340 let (status, resumed) = get_json_with_headers(
6341 &app,
6342 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
6343 &[],
6344 )
6345 .await;
6346 assert_eq!(status, StatusCode::OK);
6347 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6348
6349 let (_, tail) = get_json(
6351 &app,
6352 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
6353 )
6354 .await;
6355 let tail_events: Vec<String> = tail["events"]
6356 .as_array()
6357 .unwrap()
6358 .iter()
6359 .map(|e| e["data"].as_str().unwrap().to_string())
6360 .collect();
6361 assert_eq!(tail_events, datas[2..].to_vec());
6362 }
6363
6364 #[tokio::test]
6367 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6368 use http_body_util::BodyExt;
6369 use tower::ServiceExt;
6370
6371 let app = test_app();
6372 let body = post_sse_raw(&app, resumable_request()).await;
6373 let ids = sse_field(&body, "id:");
6374 let datas: Vec<String> = sse_field(&body, "data:")
6375 .iter()
6376 .map(|d| d.to_string())
6377 .collect();
6378 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6379
6380 let response = app
6381 .clone()
6382 .oneshot(
6383 axum::http::Request::builder()
6384 .method("GET")
6385 .uri(ferrox_api::routes::v1_stream(&request_id))
6386 .header("last-event-id", format!("{request_id}:0"))
6387 .body(axum::body::Body::empty())
6388 .unwrap(),
6389 )
6390 .await
6391 .unwrap();
6392 assert_eq!(response.status(), StatusCode::OK);
6393 assert_eq!(
6394 response
6395 .headers()
6396 .get("x-accel-buffering")
6397 .and_then(|v| v.to_str().ok()),
6398 Some("no"),
6399 "the reconnect needs the same anti-buffering header as the stream"
6400 );
6401 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6402 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6403 assert_eq!(
6404 sse_field(&resumed, "data:")
6405 .iter()
6406 .map(|d| d.to_string())
6407 .collect::<Vec<_>>(),
6408 datas[1..].to_vec()
6409 );
6410 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6411 }
6412
6413 #[tokio::test]
6417 async fn a_last_event_id_from_another_stream_is_refused() {
6418 let app = test_app();
6419 let body = post_sse_raw(&app, resumable_request()).await;
6420 let request_id = sse_field(&body, "id:")[0]
6421 .rsplit_once(':')
6422 .unwrap()
6423 .0
6424 .to_string();
6425
6426 let (status, err) = get_json_with_headers(
6427 &app,
6428 &ferrox_api::routes::v1_stream(&request_id),
6429 &[("last-event-id", "chatcmpl-someone-else:3")],
6430 )
6431 .await;
6432 assert_eq!(status, StatusCode::BAD_REQUEST);
6433 assert_eq!(err["error"]["code"], "bad_last_event_id");
6434 }
6435
6436 #[tokio::test]
6440 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6441 let app = test_app();
6442 let mut request = resumable_request();
6443 request["stream_resumable"] = serde_json::json!(false);
6444 let body = post_sse_raw(&app, request).await;
6445 let request_id = body
6446 .lines()
6447 .find_map(|l| l.strip_prefix("data: "))
6448 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6449 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6450 .unwrap();
6451
6452 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
6453 assert_eq!(status, StatusCode::NOT_FOUND);
6454 assert_eq!(err["error"]["code"], "stream_not_found");
6455 assert!(err["error"]["message"]
6456 .as_str()
6457 .unwrap()
6458 .contains("stream_resumable"));
6459 }
6460
6461 #[test]
6465 fn the_axum_stream_patterns_match_the_published_templates() {
6466 assert_eq!(
6467 axum_path(ferrox_api::routes::V1_STREAM),
6468 "/v1/stream/:request_id"
6469 );
6470 assert_eq!(
6471 axum_path(ferrox_api::routes::V1_STREAM_POLL),
6472 "/v1/stream/:request_id/poll"
6473 );
6474 assert_eq!(
6475 ferrox_api::routes::v1_stream("abc"),
6476 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
6477 );
6478 }
6479
6480 #[test]
6490 fn no_published_template_reaches_the_router_with_its_braces() {
6491 for template in [
6492 ferrox_api::routes::V1_STREAM,
6493 ferrox_api::routes::V1_STREAM_POLL,
6494 ferrox_api::routes::V1_RESPONSE,
6495 ferrox_api::routes::V1_RESPONSE_CANCEL,
6496 ferrox_api::routes::ADMIN_TASK_CANCEL,
6497 ] {
6498 assert!(
6499 template.contains('{'),
6500 "{template} is in the template list but has no placeholder"
6501 );
6502 let mounted = axum_path(template);
6503 assert!(
6504 !mounted.contains('{') && !mounted.contains('}'),
6505 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6506 );
6507 assert!(
6508 mounted.contains(':'),
6509 "{template} lost its placeholder entirely and would match one path only"
6510 );
6511 }
6512 }
6513
6514 #[tokio::test]
6522 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6523 let app = test_app();
6524 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6525 assert_eq!(status, StatusCode::NOT_FOUND);
6526 assert!(
6527 !body.is_null(),
6528 "empty body means axum never matched the route, so the id was read as a literal segment"
6529 );
6530 }
6531
6532 #[tokio::test]
6535 async fn the_task_list_starts_empty_rather_than_absent() {
6536 let app = test_app();
6537 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
6538 assert_eq!(status, StatusCode::OK);
6539 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6540 }
6541
6542 async fn post_json_uri(
6543 app: &Router,
6544 uri: &str,
6545 body: serde_json::Value,
6546 ) -> (StatusCode, serde_json::Value) {
6547 use http_body_util::BodyExt;
6548 use tower::ServiceExt;
6549
6550 let response = app
6551 .clone()
6552 .oneshot(
6553 axum::http::Request::builder()
6554 .method("POST")
6555 .uri(uri)
6556 .header("content-type", "application/json")
6557 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6558 .unwrap(),
6559 )
6560 .await
6561 .unwrap();
6562 let status = response.status();
6563 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6564 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6565 (status, json)
6566 }
6567
6568 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6569 post_json_uri(app, "/v1/chat/completions", body).await.1
6570 }
6571
6572 #[tokio::test]
6581 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6582 let app = test_app();
6583 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6584 assert_eq!(status, StatusCode::OK);
6585
6586 let memory = &body["memory"];
6587 if memory.is_null() {
6588 return;
6591 }
6592 assert!(
6593 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6594 "a read that produced a zero is a broken read, not an idle \
6595 engine: {memory}"
6596 );
6597 assert!(
6598 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6599 "the quantity must travel with the number: {memory}"
6600 );
6601 }
6602
6603 #[tokio::test]
6609 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6610 let app = test_app();
6611 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6612 assert_eq!(status, StatusCode::OK);
6613 for pool in ["kv_pages", "window_slots", "state_slots"] {
6614 assert!(
6615 body["pools"][pool].is_null(),
6616 "{pool} must be null rather than a zero row: {}",
6617 body["pools"]
6618 );
6619 }
6620 }
6621
6622 #[tokio::test]
6630 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6631 use http_body_util::BodyExt;
6632 use tower::ServiceExt;
6633
6634 let app = test_app();
6635 let send = |body: serde_json::Value| {
6636 let app = app.clone();
6637 async move {
6638 app.oneshot(
6639 axum::http::Request::builder()
6640 .method("POST")
6641 .uri(ferrox_api::routes::V1_MESSAGES)
6642 .header("content-type", "application/json")
6643 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6644 .unwrap(),
6645 )
6646 .await
6647 .unwrap()
6648 }
6649 };
6650
6651 let ok = send(serde_json::json!({
6652 "model": "test",
6653 "max_tokens": 1,
6654 "messages": [{"role": "user", "content": "hi"}],
6655 }))
6656 .await;
6657 assert_eq!(ok.status(), StatusCode::OK);
6658 let id = ok
6659 .headers()
6660 .get("request-id")
6661 .expect("a served message names its id")
6662 .to_str()
6663 .unwrap()
6664 .to_string();
6665 assert!(!id.is_empty());
6666
6667 let bad = send(serde_json::json!({"model": "test"})).await;
6670 assert!(bad.status().is_client_error());
6671 let other = bad.headers().get("request-id").expect("errors too");
6672 assert_ne!(other.to_str().unwrap(), id);
6673 let _ = bad.into_body().collect().await.unwrap();
6674 }
6675
6676 #[tokio::test]
6684 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6685 let state = Arc::new(test_state(
6686 test_model_full_byte_vocab(),
6687 ResponseCache::new(1000, Duration::from_secs(3600)),
6688 ));
6689 let app = test_app_with_state(Arc::clone(&state));
6690 let body = serde_json::json!({
6691 "model": "test",
6692 "messages": [{"role": "user", "content": "hi"}],
6693 "max_tokens": 1,
6694 });
6695
6696 state
6697 .maintenance
6698 .lock()
6699 .unwrap()
6700 .begin_rebuild()
6701 .expect("a fresh server is serving, so the rebuild starts");
6702 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6703 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6704 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6705
6706 state.maintenance.lock().unwrap().finish_rebuild(true);
6707 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6708 assert_eq!(
6709 status,
6710 StatusCode::OK,
6711 "the gate reopens; a rebuild is not a latch"
6712 );
6713 }
6714
6715 #[tokio::test]
6720 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6721 let app = test_app();
6722 let (status, body) = post_json_uri(
6723 &app,
6724 ferrox_api::routes::V1_CANCEL,
6725 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6726 )
6727 .await;
6728 assert_eq!(status, StatusCode::NOT_FOUND);
6729 assert_eq!(body["cancelled"], serde_json::json!(false));
6730 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6731 assert!(
6732 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6733 "the verdict must carry a human reason: {body}"
6734 );
6735 }
6736
6737 #[tokio::test]
6741 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6742 let state = Arc::new(test_state(
6743 test_model_full_byte_vocab(),
6744 ResponseCache::new(1000, Duration::from_secs(3600)),
6745 ));
6746 let app = test_app_with_state(Arc::clone(&state));
6747 let (token, _guard) = state.cancels.register("chatcmpl-live");
6748
6749 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6750 assert_eq!(status, StatusCode::OK);
6751 assert_eq!(before["generating_now"], serde_json::json!(1));
6752
6753 let (status, body) = post_json_uri(
6754 &app,
6755 ferrox_api::routes::V1_CANCEL,
6756 serde_json::json!({ "request_id": "chatcmpl-live" }),
6757 )
6758 .await;
6759 assert_eq!(status, StatusCode::OK);
6760 assert_eq!(body["cancelled"], serde_json::json!(true));
6761 assert!(
6762 token.is_cancelled(),
6763 "the endpoint answered ok without setting the flag the decode loop reads"
6764 );
6765 }
6766
6767 #[tokio::test]
6768 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6769 let app = test_app();
6770 let (status, tok) =
6771 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6772 assert_eq!(status, StatusCode::OK);
6773 let tokens = tok["tokens"].as_array().unwrap();
6774 assert_eq!(tok["count"], tokens.len());
6775 assert!(!tokens.is_empty());
6776
6777 let (status, detok) = post_json_uri(
6778 &app,
6779 "/v1/detokenize",
6780 serde_json::json!({ "tokens": tokens }),
6781 )
6782 .await;
6783 assert_eq!(status, StatusCode::OK);
6784 assert_eq!(detok["text"], "Hi");
6785
6786 let (status, emb) = post_json_uri(
6787 &app,
6788 "/v1/embeddings",
6789 serde_json::json!({
6790 "input": "Hi",
6791 "embedding_type": "mean"
6792 }),
6793 )
6794 .await;
6795 assert_eq!(status, StatusCode::OK);
6796 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6797 assert!(!vec.is_empty());
6798 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6799 }
6800
6801 #[tokio::test]
6807 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6808 let app = test_app();
6809 let (status, body) = post_json_uri(
6810 &app,
6811 "/v1/embeddings",
6812 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6813 )
6814 .await;
6815 assert_eq!(status, StatusCode::BAD_REQUEST);
6816 let msg = body["error"]["message"].as_str().unwrap();
6817 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6818 }
6819
6820 #[tokio::test]
6825 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6826 async fn a_real_embedding_model_serves_v1_embeddings() {
6827 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6828 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6829 if !path.exists() {
6830 eprintln!("SKIP: {} not present", path.display());
6831 return;
6832 }
6833 let encoder = ferrox_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6834 let mut state = test_state(
6835 test_model_full_byte_vocab(),
6836 ResponseCache::new(1000, Duration::from_secs(3600)),
6837 );
6838 state.embedding = Some(Arc::new(encoder));
6839 let app = test_app_with_state(Arc::new(state));
6840
6841 let (status, body) = post_json_uri(
6842 &app,
6843 "/v1/embeddings",
6844 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6845 )
6846 .await;
6847 assert_eq!(status, StatusCode::OK, "{body}");
6848 assert_eq!(body["model"], "bge-small-en-v1.5");
6849 let data = body["data"].as_array().unwrap();
6850 assert_eq!(data.len(), 2);
6851 for (i, row) in data.iter().enumerate() {
6852 assert_eq!(row["index"], i);
6853 let v: Vec<f64> = row["embedding"]
6854 .as_array()
6855 .unwrap()
6856 .iter()
6857 .map(|x| x.as_f64().unwrap())
6858 .collect();
6859 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6860 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6861 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6862 }
6863 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6866
6867 let (status, mean) = post_json_uri(
6871 &app,
6872 "/v1/embeddings",
6873 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6874 )
6875 .await;
6876 assert_eq!(status, StatusCode::OK);
6877 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6878 }
6879
6880 #[tokio::test]
6892 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6893 async fn an_encoder_can_be_the_loaded_model() {
6894 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6895 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6896 if !path.exists() {
6897 eprintln!("SKIP: {} not present", path.display());
6898 return;
6899 }
6900
6901 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6905 assert!(
6906 matches!(loaded, model::LoadedModel::Encoder(_)),
6907 "an encoder-only GGUF reached a decoder loader"
6908 );
6909 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6910 assert!(
6911 matches!(loaded, Loaded::Encoder(_)),
6912 "the encoder did not stay an encoder through activation"
6913 );
6914 assert!(
6915 batcher.is_none() && ceiling.is_none(),
6916 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6917 );
6918
6919 let state = test_state(
6920 test_model_full_byte_vocab(),
6921 ResponseCache::new(1000, Duration::from_secs(3600)),
6922 );
6923 state.swap_active(Some(Arc::new(ActiveModel {
6924 id: None,
6925 loaded,
6926 batcher,
6927 ceiling,
6928 })));
6929 let app = test_app_with_state(Arc::new(state));
6930
6931 let (status, body) = post_json_uri(
6933 &app,
6934 "/v1/embeddings",
6935 serde_json::json!({ "input": "Hello world" }),
6936 )
6937 .await;
6938 assert_eq!(status, StatusCode::OK, "{body}");
6939 assert_eq!(body["model"], "bge-small-en-v1.5");
6940 let v = body["data"][0]["embedding"].as_array().unwrap();
6941 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6942
6943 let (status, body) = post_json_uri(
6945 &app,
6946 "/v1/chat/completions",
6947 serde_json::json!({
6948 "model": "bge-small-en-v1.5",
6949 "messages": [{"role": "user", "content": "hi"}],
6950 }),
6951 )
6952 .await;
6953 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6954 let msg = body["error"]["message"].as_str().unwrap();
6955 for fact in [
6956 "bge-small-en-v1.5",
6957 "bert",
6958 "embedding model",
6959 "/v1/embeddings",
6960 ] {
6961 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6962 }
6963
6964 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
6966 assert_eq!(status, StatusCode::OK);
6967 let entry = &models["data"][0];
6968 assert_eq!(entry["id"], "bge-small-en-v1.5");
6969 assert_eq!(entry["ferrox_model_kind"], "embedding");
6970 assert_eq!(entry["ferrox_tokenizer"], "gguf-wordpiece");
6971 assert_eq!(entry["ferrox_n_embd"], 384);
6972 assert_eq!(entry["ferrox_pooling"], "CLS");
6973 assert_eq!(
6974 entry["ferrox_endpoints"],
6975 serde_json::json!(["/v1/embeddings"])
6976 );
6977 assert!(entry.get("supported_reasoning_efforts").is_none());
6980
6981 let (status, health) = get_json(&app, ferrox_api::routes::HEALTH).await;
6983 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6984 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6985 assert_eq!(health["model"]["synthetic_weights"], false);
6986 let weights = health["capabilities"]
6987 .as_array()
6988 .unwrap()
6989 .iter()
6990 .find(|c| c["id"] == ferrox_api::health::capability::REAL_WEIGHTS)
6991 .expect("a real-weights capability row");
6992 let detail = weights["detail"].as_str().unwrap_or_default();
6993 assert!(detail.contains("ENCODER"), "{detail}");
6994 }
6995
6996 #[tokio::test]
7001 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7002 use http_body_util::BodyExt;
7003 use tower::ServiceExt;
7004
7005 let fixture = concat!(
7006 "../ferrox-models/tests/fixtures/",
7007 "ferrox_real_moe_test.gguf"
7008 );
7009 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7010 let decoder = Decoder::from_gguf_with_expert_cache(
7011 &fixture,
7012 ferrox_models::config::test_moe_fixture(),
7013 Some(1024 * 1024),
7014 )
7015 .expect("MoE fixture must load store-backed");
7016
7017 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
7021 .layers
7022 .iter()
7023 .map(|_| {
7024 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
7025 })
7026 .collect();
7027 decoder.forward_token(1, 0, &mut caches);
7028
7029 let model = Model::Gguf(GgufModel {
7030 decoder: Arc::new(decoder),
7031 tokenizer: Arc::new(ServerTokenizer::Byte),
7032 stop_tokens: StopTokens::default(),
7033 bos_id: None,
7034 is_synthetic: false,
7035 chat_template: chat_template::PromptTemplate::plain(),
7036 });
7037 let state = Arc::new(test_state(
7038 model,
7039 ResponseCache::new(16, Duration::from_secs(60)),
7040 ));
7041 let app = Router::new()
7042 .route("/metrics", axum::routing::get(metrics))
7043 .route("/v1/chat/completions", post(chat_completions))
7044 .with_state(state);
7045
7046 let fetch_metrics = |app: Router| async move {
7047 let resp = app
7048 .oneshot(
7049 axum::http::Request::builder()
7050 .method("GET")
7051 .uri("/metrics")
7052 .body(axum::body::Body::empty())
7053 .unwrap(),
7054 )
7055 .await
7056 .unwrap();
7057 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7058 String::from_utf8(bytes.to_vec()).unwrap()
7059 };
7060
7061 let after = fetch_metrics(app.clone()).await;
7062 assert!(
7063 after.contains("ferrox_expert_cache_misses_total"),
7064 "streaming model must expose expert-cache metrics: {after}"
7065 );
7066 let misses: u64 = after
7067 .lines()
7068 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
7069 .and_then(|l| l.split_whitespace().nth(1))
7070 .and_then(|v| v.parse().ok())
7071 .expect("misses metric line must parse");
7072 assert!(
7073 misses > 0,
7074 "decode must have read experts through the store: {after}"
7075 );
7076 }
7077
7078 fn weather_tool() -> serde_json::Value {
7079 serde_json::json!({
7080 "type": "function",
7081 "function": {
7082 "name": "get_weather",
7083 "description": "Get the current weather for a location.",
7084 "parameters": {
7085 "type": "object",
7086 "properties": {"location": {"type": "string"}},
7087 "required": ["location"]
7088 }
7089 }
7090 })
7091 }
7092
7093 fn weather_tool_def() -> ToolDef {
7094 ToolDef {
7095 kind: "function".to_string(),
7096 function: ToolFunctionDef {
7097 name: "get_weather".to_string(),
7098 description: Some("Get the current weather for a location.".to_string()),
7099 parameters: Some(serde_json::json!({
7100 "type": "object",
7101 "properties": {"location": {"type": "string"}},
7102 "required": ["location"]
7103 })),
7104 },
7105 }
7106 }
7107
7108 #[test]
7109 fn tool_preamble_mentions_every_tool_name_and_description() {
7110 let preamble = tool_preamble(&[weather_tool_def()]);
7111 assert!(preamble.contains("get_weather"));
7112 assert!(preamble.contains("Get the current weather for a location."));
7113 assert!(preamble.contains("<tool_call>"));
7114 assert!(preamble.contains("</tool_call>"));
7115 }
7116
7117 #[test]
7118 fn a_real_marker_becomes_a_structured_tool_call() {
7119 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7120 let (message, finish) = build_response_message(
7121 text.to_string(),
7122 &[weather_tool_def()],
7123 output::OutputPosture::for_model("test-model"),
7124 "stop",
7125 );
7126 assert_eq!(finish, "tool_calls");
7127 let calls = message.tool_calls.expect("must carry a tool call");
7128 assert_eq!(calls[0].function.name, "get_weather");
7129 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7130 assert_eq!(parsed["location"], "Paris");
7131 }
7132
7133 #[test]
7134 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7135 let (message, finish) = build_response_message(
7136 "just an answer".to_string(),
7137 &[weather_tool_def()],
7138 output::OutputPosture::for_model("test-model"),
7139 "stop",
7140 );
7141 assert_eq!(finish, "stop");
7142 assert!(message.tool_calls.is_none());
7143 assert_eq!(message.content.as_deref(), Some("just an answer"));
7144 }
7145
7146 #[test]
7149 fn a_malformed_payload_is_not_a_tool_call() {
7150 let (message, finish) = build_response_message(
7151 "<tool_call>not valid json at all</tool_call>".to_string(),
7152 &[weather_tool_def()],
7153 output::OutputPosture::for_model("test-model"),
7154 "stop",
7155 );
7156 assert_eq!(finish, "stop");
7157 assert!(message.tool_calls.is_none());
7158 }
7159
7160 #[test]
7163 fn a_tool_that_was_never_offered_is_not_returned() {
7164 let (message, finish) = build_response_message(
7165 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7166 &[weather_tool_def()],
7167 output::OutputPosture::for_model("test-model"),
7168 "stop",
7169 );
7170 assert_eq!(finish, "stop");
7171 assert!(message.tool_calls.is_none());
7172 }
7173
7174 #[test]
7176 fn marker_text_with_no_tools_offered_stays_content() {
7177 let (message, finish) = build_response_message(
7178 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7179 &[],
7180 output::OutputPosture::for_model("test-model"),
7181 "stop",
7182 );
7183 assert_eq!(finish, "stop");
7184 assert!(message.tool_calls.is_none());
7185 assert!(message.content.is_some());
7186 }
7187
7188 #[test]
7192 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7193 let opened = std::cell::Cell::new(0usize);
7194 let mut parser = crate::policy::parser::ToolCallParser::new(
7195 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7196 vec![
7197 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7198 "write_file",
7199 serde_json::json!({"type": "object", "properties": {
7200 "path": {"type": "string"},
7201 "contents": {"type": "string"}
7202 }}),
7203 ),
7204 ],
7205 );
7206 let wire = "<tool_call><function=write_file>\
7207 <parameter=path>\n/tmp/x\n</parameter>\
7208 <parameter=contents>\nhello world\n</parameter>\
7209 </function></tool_call>";
7210
7211 let mut deltas = Vec::new();
7212 let mut text = String::new();
7213 for piece in wire.as_bytes().chunks(7) {
7214 let chunk = String::from_utf8_lossy(piece).into_owned();
7215 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7216 text.push_str(&more_text);
7217 deltas.extend(more);
7218 }
7219 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7220 text.push_str(&more_text);
7221 deltas.extend(more);
7222
7223 assert_eq!(opened.get(), 1, "one call opened");
7224 assert!(text.is_empty(), "the markers are not content: {text:?}");
7225
7226 let first = &deltas[0];
7227 assert_eq!(first.index, 0);
7228 assert_eq!(first.id.as_deref(), Some("call_0"));
7229 assert_eq!(first.kind, Some("function"));
7230 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7231
7232 let joined: String = deltas
7235 .iter()
7236 .filter_map(|d| d.function.arguments.clone())
7237 .collect();
7238 let parsed: serde_json::Value =
7239 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7240 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7241 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7242 assert!(
7243 deltas.len() >= 3,
7244 "the arguments arrived in pieces, not whole: {}",
7245 deltas.len()
7246 );
7247 assert!(
7248 deltas[1..].iter().all(|d| d.function.name.is_none()),
7249 "only the opening delta carries identity"
7250 );
7251 }
7252
7253 #[test]
7255 fn text_around_a_streamed_call_is_still_content() {
7256 let opened = std::cell::Cell::new(0usize);
7257 let mut parser = crate::policy::parser::ToolCallParser::new(
7258 crate::policy::parser::ToolCallFormat::Qwen25,
7259 vec![crate::policy::parser::tool_call::ToolSchema::new(
7260 "get_weather",
7261 )],
7262 );
7263 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7264 \"arguments\": {}}</tool_call> done";
7265 let mut text = String::new();
7266 for piece in wire.as_bytes().chunks(5) {
7267 let chunk = String::from_utf8_lossy(piece).into_owned();
7268 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7269 text.push_str(&more);
7270 }
7271 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7272 text.push_str(&more);
7273
7274 assert_eq!(opened.get(), 1);
7275 assert!(text.starts_with("let me check. "), "{text:?}");
7276 assert!(text.ends_with(" done"), "{text:?}");
7277 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7278 }
7279
7280 #[test]
7283 fn a_reasoning_block_is_split_out_of_the_answer() {
7284 let (message, finish) = build_response_message(
7285 "<think>weighing it up</think>The answer is 4.".to_string(),
7286 &[],
7287 output::OutputPosture::for_model("Qwen3-8B"),
7288 "stop",
7289 );
7290 assert_eq!(finish, "stop");
7291 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7292 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7293 }
7294
7295 #[test]
7298 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7299 let (message, _) = build_response_message(
7300 "Use the <think> tag like this.".to_string(),
7301 &[],
7302 output::OutputPosture::for_model("llama-3.1-8b"),
7303 "stop",
7304 );
7305 assert_eq!(
7306 message.content.as_deref(),
7307 Some("Use the <think> tag like this.")
7308 );
7309 assert!(message.reasoning_content.is_none());
7310 }
7311
7312 #[tokio::test]
7320 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7321 let app = test_app();
7322 let body = serde_json::json!({
7323 "model": "m",
7324 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7325 "max_tokens": 4,
7326 "temperature": 0,
7327 });
7328 let resp = post_json(&app, body).await;
7329 let message = &resp["choices"][0]["message"];
7330 assert!(message["content"].is_string());
7331 assert!(message.get("tool_calls").is_none());
7332 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7333 assert_eq!(resp["usage"]["completion_tokens"], 4);
7334 assert_eq!(
7335 resp["usage"]["total_tokens"],
7336 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7337 );
7338 }
7339
7340 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7341 use http_body_util::BodyExt;
7342 use tower::ServiceExt;
7343
7344 let response = app
7345 .clone()
7346 .oneshot(
7347 axum::http::Request::builder()
7348 .method("GET")
7349 .uri(uri)
7350 .body(axum::body::Body::empty())
7351 .unwrap(),
7352 )
7353 .await
7354 .unwrap();
7355 let status = response.status();
7356 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7357 (status, serde_json::from_slice(&bytes).unwrap())
7358 }
7359
7360 #[tokio::test]
7361 async fn health_answers_a_capability_handshake_not_a_boolean() {
7362 let app = test_app();
7363 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7364 assert_eq!(status, StatusCode::OK);
7365
7366 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7367 assert_eq!(health.state, ferrox_api::HealthState::Ready);
7368 assert!(health.pid > 0);
7369 assert!(health.server_time_unix_ms > 0);
7370 assert_eq!(health.last_request_age_seconds, None);
7373
7374 for id in [
7377 ferrox_api::health::capability::CPU,
7378 ferrox_api::health::capability::METAL,
7379 ferrox_api::health::capability::CUDA,
7380 ferrox_api::health::capability::REAL_WEIGHTS,
7381 ferrox_api::health::capability::CONTINUOUS_BATCHING,
7382 ] {
7383 let cap = health
7384 .capability(id)
7385 .unwrap_or_else(|| panic!("{id} missing"));
7386 assert!(!cap.reason.is_empty(), "{cap:?}");
7387 assert!(!cap.detail.is_empty(), "{cap:?}");
7388 }
7389 let weights = health
7393 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
7394 .unwrap();
7395 assert!(!weights.available);
7396 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
7397 assert!(health.model.as_ref().unwrap().synthetic_weights);
7398 }
7399
7400 #[tokio::test]
7401 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7402 let app = test_app();
7403 let _ = post_json(
7404 &app,
7405 serde_json::json!({
7406 "model": "m",
7407 "messages": [{"role": "user", "content": "\u{1}"}],
7408 "max_tokens": 1,
7409 "temperature": 0,
7410 }),
7411 )
7412 .await;
7413 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7414 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7415 let age = health
7416 .last_request_age_seconds
7417 .expect("a served request is evidence of liveness");
7418 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7419 }
7420
7421 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7423 use http_body_util::BodyExt;
7424 use tower::ServiceExt;
7425
7426 let response = app
7427 .clone()
7428 .oneshot(
7429 axum::http::Request::builder()
7430 .method("POST")
7431 .uri("/v1/chat/completions")
7432 .header("content-type", "application/json")
7433 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7434 .unwrap(),
7435 )
7436 .await
7437 .unwrap();
7438 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7439 String::from_utf8(bytes.to_vec())
7440 .unwrap()
7441 .lines()
7442 .filter_map(|line| line.strip_prefix("data: "))
7443 .filter(|payload| *payload != "[DONE]")
7444 .map(|payload| serde_json::from_str(payload).unwrap())
7445 .collect()
7446 }
7447
7448 #[tokio::test]
7449 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7450 let app = test_app();
7451 let chunks = post_sse_chunks(
7452 &app,
7453 serde_json::json!({
7454 "model": "m",
7455 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7456 "max_tokens": 4,
7457 "temperature": 0,
7458 "stream": true,
7459 }),
7460 )
7461 .await;
7462
7463 assert!(!chunks.is_empty());
7464 let request_id = chunks[0]["request_id"]
7465 .as_str()
7466 .expect("the first chunk names the request")
7467 .to_string();
7468 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7469 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7472 assert!(
7473 chunk.get("request_id").is_none(),
7474 "chunk {i} repeats request_id"
7475 );
7476 }
7477 for chunk in &chunks {
7480 assert_eq!(chunk["id"], serde_json::json!(request_id));
7481 }
7482
7483 let other = post_sse_chunks(
7484 &app,
7485 serde_json::json!({
7486 "model": "m",
7487 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7488 "max_tokens": 4,
7489 "temperature": 0,
7490 "stream": true,
7491 }),
7492 )
7493 .await;
7494 assert_ne!(
7495 other[0]["request_id"].as_str().unwrap(),
7496 request_id,
7497 "two concurrent chats must not share an id"
7498 );
7499 }
7500
7501 #[tokio::test]
7502 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7503 let app = test_app();
7504 let resp = post_json(
7505 &app,
7506 serde_json::json!({
7507 "model": "m",
7508 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7509 "max_tokens": 2,
7510 "temperature": 0,
7511 }),
7512 )
7513 .await;
7514 assert_eq!(resp["id"], resp["request_id"]);
7515 assert!(resp["request_id"]
7516 .as_str()
7517 .unwrap()
7518 .starts_with("chatcmpl-"));
7519 }
7520
7521 #[tokio::test]
7524 async fn usage_carries_separate_prefill_and_decode_timings() {
7525 let app = test_app();
7526 let resp = post_json(
7527 &app,
7528 serde_json::json!({
7529 "model": "m",
7530 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7531 "max_tokens": 4,
7532 "temperature": 0,
7533 }),
7534 )
7535 .await;
7536 let usage = &resp["usage"];
7537 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7538 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7539 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7540 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7541 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7543 }
7544
7545 #[tokio::test]
7553 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7554 let app = test_app();
7555 let body = serde_json::json!({
7556 "model": "m",
7557 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7558 "max_tokens": 4,
7559 "temperature": 0,
7560 "tools": [weather_tool()],
7561 });
7562 let resp = post_json(&app, body).await;
7563 let message = &resp["choices"][0]["message"];
7564 assert!(
7565 message["content"].is_string(),
7566 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7567 );
7568 assert!(message.get("tool_calls").is_none());
7569 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7573 }
7574
7575 #[tokio::test]
7579 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7580 let app = test_app();
7581 let body = serde_json::json!({
7582 "model": "m",
7583 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7584 "max_tokens": 3,
7585 "temperature": 0,
7586 });
7587 let first = post_json(&app, body.clone()).await;
7588 assert_eq!(first["ferrox_cache"], "miss");
7589 let second = post_json(&app, body).await;
7590 assert_eq!(second["ferrox_cache"], "hit");
7591 assert_eq!(
7592 first["choices"][0]["message"]["content"],
7593 second["choices"][0]["message"]["content"]
7594 );
7595 assert_eq!(
7596 first["choices"][0]["finish_reason"],
7597 second["choices"][0]["finish_reason"]
7598 );
7599 assert_eq!(first["usage"], second["usage"]);
7600 assert_eq!(second["usage"]["completion_tokens"], 3);
7601 }
7602
7603 #[tokio::test]
7609 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7610 let session_app = test_app();
7611 let manual_app = test_app();
7612
7613 let turn1 = post_json(
7615 &session_app,
7616 serde_json::json!({
7617 "model": "m",
7618 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7619 "session_id": "s1",
7620 "max_tokens": 5,
7621 "temperature": 0,
7622 }),
7623 )
7624 .await;
7625 let reply1 = turn1["choices"][0]["message"]["content"]
7626 .as_str()
7627 .unwrap()
7628 .to_string();
7629
7630 let manual_turn1 = post_json(
7635 &manual_app,
7636 serde_json::json!({
7637 "model": "m",
7638 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7639 "max_tokens": 5,
7640 "temperature": 0,
7641 }),
7642 )
7643 .await;
7644 assert_eq!(
7645 manual_turn1["choices"][0]["message"]["content"]
7646 .as_str()
7647 .unwrap(),
7648 reply1
7649 );
7650
7651 let turn2 = post_json(
7653 &session_app,
7654 serde_json::json!({
7655 "model": "m",
7656 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7657 "session_id": "s1",
7658 "max_tokens": 5,
7659 "temperature": 0,
7660 }),
7661 )
7662 .await;
7663 let reply2 = turn2["choices"][0]["message"]["content"]
7664 .as_str()
7665 .unwrap()
7666 .to_string();
7667
7668 let manual_turn2 = post_json(
7672 &manual_app,
7673 serde_json::json!({
7674 "model": "m",
7675 "messages": [
7676 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7677 {"role": "assistant", "content": reply1},
7678 {"role": "user", "content": "\u{4}\u{5}"},
7679 ],
7680 "max_tokens": 5,
7681 "temperature": 0,
7682 }),
7683 )
7684 .await;
7685 assert_eq!(
7686 manual_turn2["choices"][0]["message"]["content"]
7687 .as_str()
7688 .unwrap(),
7689 reply2,
7690 "resuming a session must produce identical output to manually resending the full history"
7691 );
7692 }
7693
7694 #[test]
7697 fn lock_cache_recovers_from_a_poisoned_mutex() {
7698 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7699
7700 let poison_cache = Arc::clone(&cache);
7701 let _ = std::thread::spawn(move || {
7702 let _guard = poison_cache.lock().unwrap();
7703 panic!("simulated panic while holding the lock");
7704 })
7705 .join();
7706
7707 let recovered = lock_cache(&cache);
7709 assert_eq!(recovered.stats().entries, 0);
7710 }
7711
7712 #[test]
7713 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7714 let mut req_body = serde_json::json!({
7715 "model": "m",
7716 "messages": [{"role": "user", "content": "hi"}],
7717 });
7718 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7719 assert!(
7720 req.is_cacheable(),
7721 "default (temperature 0) must be cacheable"
7722 );
7723
7724 req_body["temperature"] = serde_json::json!(0.8);
7725 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7726 assert!(
7727 !req.is_cacheable(),
7728 "unseeded sampling must never be cacheable"
7729 );
7730
7731 req_body["seed"] = serde_json::json!(42);
7732 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7733 assert!(
7734 req.is_cacheable(),
7735 "sampling with an explicit seed is deterministic and must be cacheable"
7736 );
7737 }
7738
7739 const GRADED: &str = "{% if reasoning_effort %}\
7743 {% if reasoning_effort not in ['low','medium','high'] %}\
7744 {{ raise_exception('unsupported effort') }}\
7745 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7746 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7747
7748 fn graded_template() -> chat_template::PromptTemplate {
7749 chat_template::PromptTemplate::from_gguf_metadata(
7750 Some(GRADED),
7751 Some("qwen3"),
7752 false,
7753 true,
7754 None,
7755 None,
7756 )
7757 }
7758
7759 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7760 serde_json::from_value(value).expect("request")
7761 }
7762
7763 #[test]
7770 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7771 let req = chat_request(serde_json::json!({
7772 "model": "m",
7773 "messages": [{"role": "user", "content": "hi"}],
7774 "grammar": "root ::= \"a\"+",
7775 }));
7776 req.validate_supported_fields()
7777 .expect("a valid grammar is a valid request");
7778 let params = req
7779 .generation_params()
7780 .expect("a valid grammar compiles at params time too");
7781 assert!(
7782 params.grammar.is_some(),
7783 "the grammar was dropped between the wire and the sampler"
7784 );
7785 assert!(
7786 params.needs_vocab_logits(),
7787 "a grammar request that may fold lm_head into a GPU argmax is \
7788 a grammar request served unconstrained"
7789 );
7790
7791 let plain = chat_request(serde_json::json!({
7792 "model": "m",
7793 "messages": [{"role": "user", "content": "hi"}],
7794 }));
7795 assert!(plain.generation_params().unwrap().grammar.is_none());
7796 }
7797
7798 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7799 chat_request(serde_json::json!({
7800 "model": "m",
7801 "messages": [{"role": "user", "content": "weather in Rome?"}],
7802 "tools": [weather_tool()],
7803 "tool_choice": tool_choice,
7804 }))
7805 }
7806
7807 #[test]
7811 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7812 for choice in [
7813 serde_json::json!("required"),
7814 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7815 ] {
7816 let req = tool_request(choice.clone());
7817 req.validate_supported_fields()
7818 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7819 let params = req
7820 .generation_params_for_template(&graded_template(), "Qwen3-8B")
7821 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7822 let grammar = params
7823 .grammar
7824 .as_ref()
7825 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7826 assert!(
7827 grammar.is_awaiting_trigger(),
7828 "the model must be free to think before it calls"
7829 );
7830 assert!(
7831 !grammar.allows_eog(),
7832 "{choice} must not be able to end the turn without a call"
7833 );
7834 assert!(
7840 params.needs_vocab_logits(),
7841 "{choice} would let a backend return a token id instead of logits"
7842 );
7843 assert!(
7844 !generate::greedy_gpu_fold_allowed(¶ms),
7845 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7846 );
7847 }
7848 }
7849
7850 #[test]
7852 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7853 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7854 let req = tool_request(choice.clone());
7855 req.validate_supported_fields().expect("still supported");
7856 let params = match req.generation_params_for_template(&graded_template(), "Qwen3-8B") {
7857 Ok(p) => p,
7858 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7859 };
7860 assert!(
7861 params.grammar.is_none(),
7862 "{choice} does not force a call and must not be constrained"
7863 );
7864 }
7865 }
7866
7867 #[test]
7870 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7871 let req = chat_request(serde_json::json!({
7873 "model": "m",
7874 "messages": [{"role": "user", "content": "hi"}],
7875 "tool_choice": "required",
7876 }));
7877 let (status, _) = req
7878 .validate_supported_fields()
7879 .expect_err("nothing to call");
7880 assert_eq!(status, StatusCode::BAD_REQUEST);
7881
7882 let req =
7884 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7885 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7886 assert_eq!(status, StatusCode::BAD_REQUEST);
7887 assert_eq!(body["error"]["param"], "tool_choice");
7888
7889 let req = tool_request(serde_json::json!({"type": "function"}));
7891 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7892 assert_eq!(status, StatusCode::BAD_REQUEST);
7893
7894 let req = chat_request(serde_json::json!({
7896 "model": "m",
7897 "messages": [{"role": "user", "content": "hi"}],
7898 "tools": [weather_tool()],
7899 "tool_choice": "required",
7900 "grammar": "root ::= \"a\"+",
7901 }));
7902 let (status, _) = req
7903 .validate_supported_fields()
7904 .expect_err("a grammar and a forced call are two constraints");
7905 assert_eq!(status, StatusCode::BAD_REQUEST);
7906
7907 let req = tool_request(serde_json::json!("required"));
7910 let (status, Json(body)) =
7911 match req.generation_params_for_template(&graded_template(), "GLM-4.7") {
7912 Err(e) => e,
7913 Ok(_) => panic!("glm calls are not JSON behind a marker"),
7914 };
7915 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7916 assert!(
7917 body["error"]["message"].as_str().unwrap().contains("glm47"),
7918 "{body}"
7919 );
7920 }
7921
7922 #[test]
7925 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
7926 let req = chat_request(serde_json::json!({
7927 "model": "m",
7928 "messages": [{"role": "user", "content": "hi"}],
7929 "grammar": "root ::= \"a",
7930 }));
7931 let (status, Json(body)) = req
7932 .validate_supported_fields()
7933 .expect_err("this does not parse");
7934 assert_eq!(status, StatusCode::BAD_REQUEST);
7935 assert_eq!(body["error"]["param"], "grammar");
7936 assert!(req.generation_params().is_err(), "and again at params time");
7937 }
7938
7939 #[test]
7946 fn response_format_json_schema_becomes_the_requests_grammar() {
7947 let req = chat_request(serde_json::json!({
7948 "model": "m",
7949 "messages": [{"role": "user", "content": "hi"}],
7950 "response_format": {
7951 "type": "json_schema",
7952 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
7953 },
7954 }));
7955 req.validate_supported_fields()
7956 .expect("a boolean schema converts");
7957 let params = req.generation_params().expect("and compiles");
7958 let grammar = params.grammar.expect("the schema is the grammar");
7959 let mut g = (*grammar).clone();
7960 g.accept_token(0, b"true").expect("a boolean is accepted");
7961 assert!(g.allows_eog(), "and completes the parse");
7962 assert!(
7963 !params.json_object,
7964 "a schema is not the json_object character-class mask"
7965 );
7966 }
7967
7968 #[test]
7972 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
7973 let req = chat_request(serde_json::json!({
7974 "model": "m",
7975 "messages": [{"role": "user", "content": "hi"}],
7976 "response_format": {
7977 "type": "json_schema",
7978 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
7979 },
7980 }));
7981 let (status, Json(body)) = req
7982 .validate_supported_fields()
7983 .expect_err("minimum has no grammar in this port");
7984 assert_eq!(status, StatusCode::BAD_REQUEST);
7985 assert!(
7986 body["error"]["message"]
7987 .as_str()
7988 .expect("a message")
7989 .contains("minimum"),
7990 "the refusal must name the keyword: {body}"
7991 );
7992 assert!(req.generation_params().is_err(), "and again at params time");
7993 }
7994
7995 #[test]
8001 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8002 let req = chat_request(serde_json::json!({
8003 "model": "m",
8004 "messages": [{"role": "user", "content": "hi"}],
8005 "tool_choice": "required",
8006 "tools": [{
8007 "type": "function",
8008 "function": {"name": "f", "parameters": {"type": "object"}},
8009 }],
8010 "response_format": {
8011 "type": "json_schema",
8012 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8013 },
8014 }));
8015 let (status, Json(body)) = req
8016 .validate_supported_fields()
8017 .expect_err("two constraints, one generation");
8018 assert_eq!(status, StatusCode::BAD_REQUEST);
8019 assert_eq!(body["error"]["param"], "tool_choice");
8020 }
8021
8022 #[test]
8025 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8026 let req = chat_request(serde_json::json!({
8027 "model": "m",
8028 "messages": [{"role": "user", "content": "hi"}],
8029 }));
8030 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8031 }
8032
8033 #[test]
8038 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8039 let asked = chat_request(serde_json::json!({
8040 "model": "m",
8041 "messages": [{"role": "user", "content": "hi"}],
8042 "min_p": 0.07,
8043 }));
8044 assert_eq!(asked.sampling_params().min_p, 0.07);
8045
8046 let silent = chat_request(serde_json::json!({
8047 "model": "m",
8048 "messages": [{"role": "user", "content": "hi"}],
8049 }));
8050 assert_eq!(
8051 silent.sampling_params().min_p,
8052 0.0,
8053 "an unset min_p must be off, not llama.cpp's CLI default"
8054 );
8055 }
8056
8057 #[test]
8066 fn no_sampler_knob_is_missing_from_the_cache_key() {
8067 let base = serde_json::json!({
8068 "model": "m",
8069 "messages": [{"role": "user", "content": "hi"}],
8070 "seed": 1,
8071 });
8072 let baseline = chat_request(base.clone()).cache_key("prompt");
8073 for (knob, value) in [
8074 ("temperature", serde_json::json!(0.5)),
8075 ("top_p", serde_json::json!(0.9)),
8076 ("min_p", serde_json::json!(0.05)),
8077 ("top_k", serde_json::json!(40)),
8078 ("repetition_penalty", serde_json::json!(1.1)),
8079 ("presence_penalty", serde_json::json!(0.3)),
8080 ("frequency_penalty", serde_json::json!(0.3)),
8081 ] {
8082 let mut body = base.clone();
8083 body[knob] = value;
8084 assert_ne!(
8085 chat_request(body).cache_key("prompt"),
8086 baseline,
8087 "`{knob}` is not in the cache key: two requests differing \
8088 only in it would share one cached answer"
8089 );
8090 }
8091 }
8092
8093 #[test]
8097 fn an_explicit_zero_output_budget_is_a_client_error() {
8098 let req = chat_request(serde_json::json!({
8099 "model": "m",
8100 "messages": [{"role": "user", "content": "hi"}],
8101 "max_tokens": 0,
8102 }));
8103 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8104 assert_eq!(status, StatusCode::BAD_REQUEST);
8105 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8106 }
8107
8108 #[test]
8111 fn a_request_can_turn_thinking_off() {
8112 let template = graded_template();
8113 for body in [
8114 serde_json::json!({
8115 "model": "m",
8116 "messages": [{"role": "user", "content": "hi"}],
8117 "reasoning_effort": "none",
8118 }),
8119 serde_json::json!({
8120 "model": "m",
8121 "messages": [{"role": "user", "content": "hi"}],
8122 "thinking": {"type": "disabled"},
8123 }),
8124 ] {
8125 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8126 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8127 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8128 assert!(!kwargs.contains_key("reasoning_effort"));
8131 }
8132 }
8133
8134 #[test]
8137 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8138 let template = graded_template();
8139 let kwargs = chat_request(serde_json::json!({
8140 "model": "m",
8141 "messages": [{"role": "user", "content": "hi"}],
8142 "reasoning_effort": "high",
8143 "thinking": {"type": "disabled"},
8144 }))
8145 .resolve_template_kwargs(&template);
8146 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8147 assert!(!kwargs.contains_key("reasoning_effort"));
8148 }
8149
8150 #[test]
8153 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8154 let req = chat_request(serde_json::json!({
8155 "model": "m",
8156 "messages": [{"role": "user", "content": "hi"}],
8157 "thinking": {"type": "disable"},
8158 }));
8159 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8160 assert_eq!(status, StatusCode::BAD_REQUEST);
8161 }
8162
8163 #[test]
8166 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8167 let template = graded_template();
8168 let kwargs = chat_request(serde_json::json!({
8169 "model": "m",
8170 "messages": [{"role": "user", "content": "hi"}],
8171 "reasoning_effort": "none",
8172 "chat_template_kwargs": {"enable_thinking": true},
8173 }))
8174 .resolve_template_kwargs(&template);
8175 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8176 }
8177
8178 #[test]
8182 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8183 let template = graded_template();
8184 let req = chat_request(serde_json::json!({
8185 "model": "m",
8186 "messages": [{"role": "user", "content": "hi"}],
8187 "reasoning_effort": "minimal",
8188 }));
8189 let kwargs = req.resolve_template_kwargs(&template);
8190 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8191 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8192 assert!(prompt.starts_with("E:low|"), "{prompt}");
8193 }
8194
8195 #[test]
8199 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8200 let template = graded_template();
8201 let req = chat_request(serde_json::json!({
8202 "model": "m",
8203 "messages": [{"role": "user", "content": "hi"}],
8204 "chat_template_kwargs": {"reasoning_effort": "none"},
8205 }));
8206 let kwargs = req.resolve_template_kwargs(&template);
8207 assert!(!kwargs.contains_key("reasoning_effort"));
8208 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8209 assert_eq!(prompt, "hi");
8210 }
8211
8212 #[test]
8215 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8216 let template = graded_template();
8217 let req = chat_request(serde_json::json!({
8218 "model": "m",
8219 "messages": [{"role": "user", "content": "hi"}],
8220 "reasoning_effort": "low",
8221 "chat_template_kwargs": {"reasoning_effort": "high"},
8222 }));
8223 assert_eq!(
8224 req.resolve_template_kwargs(&template)["reasoning_effort"],
8225 serde_json::json!("high")
8226 );
8227 }
8228
8229 #[test]
8233 fn offering_tools_turns_thinking_on_by_itself() {
8234 let template = graded_template();
8235 let quiet = chat_request(serde_json::json!({
8236 "model": "m",
8237 "messages": [{"role": "user", "content": "hi"}],
8238 }));
8239 assert!(!quiet
8240 .resolve_template_kwargs(&template)
8241 .contains_key("enable_thinking"));
8242
8243 let with_tools = chat_request(serde_json::json!({
8244 "model": "m",
8245 "messages": [{"role": "user", "content": "hi"}],
8246 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8247 }));
8248 let kwargs = with_tools.resolve_template_kwargs(&template);
8249 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8250 let prompt =
8251 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8252 assert!(prompt.starts_with("THINK|"), "{prompt}");
8253 }
8254
8255 #[test]
8260 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8261 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8262 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8263 Some("qwen3"),
8264 false,
8265 true,
8266 None,
8267 None,
8268 );
8269 let req = chat_request(serde_json::json!({
8270 "model": "m",
8271 "messages": [{"role": "user", "content": "hi"}],
8272 "chat_template_kwargs": {"enable_thinking": true},
8273 }));
8274 let kwargs = req.resolve_template_kwargs(&opener);
8275 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8276 assert!(prompt.ends_with("<think>"), "{prompt}");
8277
8278 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8281 let (message, _) = build_response_message(
8282 "weighing it up</think>Paris.".to_string(),
8283 &[],
8284 posture,
8285 "stop",
8286 );
8287 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8288 assert_eq!(message.content.as_deref(), Some("Paris."));
8289
8290 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8293 let (message, _) = build_response_message(
8294 "weighing it up</think>Paris.".to_string(),
8295 &[],
8296 closed,
8297 "stop",
8298 );
8299 assert_eq!(message.reasoning_content, None);
8300 }
8301
8302 #[test]
8303 fn stop_param_accepts_both_single_string_and_array() {
8304 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8305 "model": "m",
8306 "messages": [{"role": "user", "content": "hi"}],
8307 "stop": "END",
8308 }))
8309 .unwrap();
8310 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8311
8312 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8313 "model": "m",
8314 "messages": [{"role": "user", "content": "hi"}],
8315 "stop": ["A", "B"],
8316 }))
8317 .unwrap();
8318 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8319 }
8320
8321 #[test]
8322 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8323 let model = test_model();
8324 let result = run_generation(
8325 &model,
8326 "hello",
8327 &greedy_params(4),
8328 None,
8329 None,
8330 None,
8331 None,
8332 None,
8333 None,
8334 );
8335 assert!(matches!(
8336 result,
8337 Err(generate::DecodeError::TokenOutOfVocab { .. })
8338 ));
8339 }
8340
8341 #[test]
8345 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8346 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8348 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8349
8350 let holder_pool = Arc::clone(&pool);
8351 let holder = std::thread::spawn(move || {
8352 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8353 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8355 drop(held);
8356 });
8357 std::thread::sleep(Duration::from_millis(15));
8358
8359 let config = generate::KvPoolConfig {
8360 pool,
8361 queue_wait: Duration::ZERO,
8362 };
8363 let result = run_generation(
8364 &model,
8365 &prompt,
8366 &greedy_params(4),
8367 Some(&config),
8368 None,
8369 None,
8370 None,
8371 None,
8372 None,
8373 );
8374 assert!(matches!(
8375 result,
8376 Err(generate::DecodeError::KvPoolExhausted)
8377 ));
8378
8379 let (status, _body) = decode_error_response(result.unwrap_err());
8380 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8381 holder.join().unwrap();
8382 }
8383
8384 #[test]
8392 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8393 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8395 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
8397 let config = generate::KvPoolConfig {
8398 pool,
8399 queue_wait: Duration::ZERO,
8400 };
8401
8402 let result = run_generation(
8403 &model,
8404 &prompt,
8405 &greedy_params(4),
8406 Some(&config),
8407 None,
8408 None,
8409 None,
8410 None,
8411 None,
8412 );
8413 let err = result.expect_err("one block cannot hold two layers' caches");
8414 assert!(
8415 matches!(
8416 &err,
8417 generate::DecodeError::KvBudgetExceeded { binding, .. }
8418 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
8419 ),
8420 "expected an immovable device-memory refusal, got {err:?}"
8421 );
8422 let (status, _body) = decode_error_response(err);
8423 assert_eq!(status, StatusCode::BAD_REQUEST);
8424 }
8425
8426 #[test]
8432 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8433 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8434 queued: 512,
8435 cap: 512,
8436 });
8437 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8438 assert_eq!(body["error"]["retry_after_seconds"], 1);
8439 let message = body["error"]["message"].as_str().expect("message");
8440 assert!(message.contains("512"), "{message}");
8441 }
8442
8443 #[test]
8444 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8445 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8446 token: 99,
8447 vocab_size: 32,
8448 });
8449 assert!(
8450 body["error"]["retry_after_seconds"].is_null(),
8451 "retrying a prompt this model cannot tokenize never helps"
8452 );
8453 }
8454
8455 #[test]
8456 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8457 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8458 token: 99,
8459 vocab_size: 32,
8460 });
8461 assert_eq!(status, StatusCode::BAD_REQUEST);
8462 }
8463
8464 #[test]
8465 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8466 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8468 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8469 let config = generate::KvPoolConfig {
8470 pool: pool.clone(),
8471 queue_wait: Duration::ZERO,
8472 };
8473
8474 let (_, finish, _usage) = run_generation(
8475 &model,
8476 &prompt,
8477 &greedy_params(4),
8478 Some(&config),
8479 None,
8480 None,
8481 None,
8482 None,
8483 None,
8484 )
8485 .unwrap();
8486 assert_eq!(finish, FinishReason::Length);
8487 assert_eq!(
8488 pool.lock().unwrap().free_blocks(),
8489 2,
8490 "a completed request must return its blocks to the pool"
8491 );
8492 }
8493
8494 #[tokio::test]
8499 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8500 let model = Arc::new(test_model());
8501 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8502
8503 let mut handles = Vec::new();
8504 for _ in 0..8 {
8505 let model = Arc::clone(&model);
8506 let prompt = prompt.clone();
8507 handles.push(tokio::task::spawn_blocking(move || {
8508 run_generation(
8509 &model,
8510 &prompt,
8511 &greedy_params(6),
8512 None,
8513 None,
8514 None,
8515 None,
8516 None,
8517 None,
8518 )
8519 .unwrap()
8520 }));
8521 }
8522
8523 let mut results = Vec::new();
8524 for h in handles {
8525 results.push(h.await.unwrap());
8526 }
8527 for r in &results[1..] {
8532 assert_eq!(r.0, results[0].0, "decoded chunks must match");
8533 assert_eq!(r.1, results[0].1, "finish reason must match");
8534 assert_eq!(
8535 r.2.prompt_tokens, results[0].2.prompt_tokens,
8536 "prompt token count must match"
8537 );
8538 assert_eq!(
8539 r.2.completion_tokens, results[0].2.completion_tokens,
8540 "completion token count must match"
8541 );
8542 }
8543 }
8544
8545 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8552 let mut header_entries = Vec::new();
8553 let mut data = Vec::new();
8554 for (name, shape, values) in tensors {
8555 let start = data.len();
8556 for v in values {
8557 data.extend_from_slice(&v.to_le_bytes());
8558 }
8559 let end = data.len();
8560 let shape_str = shape
8561 .iter()
8562 .map(|d| d.to_string())
8563 .collect::<Vec<_>>()
8564 .join(",");
8565 header_entries.push(format!(
8566 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8567 ));
8568 }
8569 let header = format!("{{{}}}", header_entries.join(","));
8570 let header_bytes = header.as_bytes();
8571 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8572 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8573 out.extend_from_slice(header_bytes);
8574 out.extend_from_slice(&data);
8575 out
8576 }
8577
8578 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8589 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8590 use ferrox_models::kimi_loader::KimiRealHparams;
8591 use ferrox_moe::{GatingFunction, MoeLayerConfig};
8592
8593 let hidden_dim = 8;
8594 let kda_num_heads = 2;
8595 let kda_head_dim = 3;
8596 let kda_proj = kda_num_heads * kda_head_dim;
8597 let conv_kernel = 4;
8598 let dense_intermediate = 5;
8599 let vocab_size = 256;
8603 let mla_num_heads = 1;
8604 let mla_q_lora_rank = 2;
8605 let mla_kv_lora_rank = 2;
8606 let mla_qk_nope_head_dim = 2;
8607 let mla_qk_rope_head_dim = 2;
8608 let mla_v_head_dim = 2;
8609
8610 let model_cfg = ferrox_models::ModelConfig {
8611 name: "synthetic-kimi-server-test",
8612 n_layers: 1,
8613 hidden_dim,
8614 n_heads: 1,
8615 n_kv_heads: 1,
8616 head_dim: 4,
8617 vocab_size,
8618 rope_theta: 10000.0,
8619 rms_norm_eps: 1e-5,
8620 sliding_window: None,
8621 moe: MoeLayerConfig {
8622 expert_weights_scale: 1.0,
8623 n_experts: 1,
8624 n_experts_active: 1,
8625 n_shared_experts: 0,
8626 hidden_dim,
8627 expert_ffn_dim: 4,
8628 gating: GatingFunction::Sigmoid,
8629 norm_topk_prob: true,
8630 expert_group_count: None,
8631 expert_group_used_count: None,
8632 },
8633 n_dense_leading_layers: 1,
8638 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8639 kda_layers: vec![1],
8640 full_attn_layers: vec![],
8641 mla: MlaConfig {
8642 num_heads: mla_num_heads,
8643 q_lora_rank: mla_q_lora_rank,
8644 kv_lora_rank: mla_kv_lora_rank,
8645 qk_nope_head_dim: mla_qk_nope_head_dim,
8646 qk_rope_head_dim: mla_qk_rope_head_dim,
8647 v_head_dim: mla_v_head_dim,
8648 use_output_gate: true,
8649 rope: None,
8650 },
8651 kda: KdaConfig {
8652 num_heads: kda_num_heads,
8653 head_dim: kda_head_dim,
8654 short_conv_kernel_size: conv_kernel,
8655 gate_lower_bound: -5.0,
8656 use_full_rank_gate: true,
8657 },
8658 }),
8659 rope_freqs: None,
8660 rope_attn_factor: 1.0,
8661 rope_dim: None,
8662 rope_freqs_long: None,
8663 rope_freqs_short: None,
8664 rope_orig_ctx: None,
8665 rope_layout: ferrox_models::config::RopeLayout::Neox,
8666 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
8667 swa_pattern: None,
8668 swa_dense_first: false,
8669 attn_logit_softcap: None,
8670 final_logit_softcap: None,
8671 embedding_scale: None,
8672 attention_scale: None,
8673 rope_theta_swa: None,
8674 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
8675 best_effort_fields: &["synthetic test config, not a real preset"],
8676 };
8677 let hp = KimiRealHparams {
8678 hidden_dim,
8679 kda_num_heads,
8680 kda_head_dim,
8681 mla_num_heads,
8682 mla_q_lora_rank,
8683 mla_kv_lora_rank,
8684 mla_qk_nope_head_dim,
8685 mla_qk_rope_head_dim,
8686 mla_v_head_dim,
8687 dense_intermediate_dim: dense_intermediate,
8688 moe_hidden_dim: hidden_dim,
8689 moe_intermediate_dim: 4,
8690 n_experts: 1,
8691 num_shared_experts: 0,
8692 };
8693
8694 let prefix = "language_model.model.layers.0";
8698 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8699 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8700 name: String,
8701 shape: Vec<usize>,
8702 n: usize| {
8703 tensors.push((name, shape, vec![0.01f32; n]));
8704 };
8705 push(
8706 &mut tensors,
8707 format!("{prefix}.input_layernorm.weight"),
8708 vec![hidden_dim],
8709 hidden_dim,
8710 );
8711 push(
8712 &mut tensors,
8713 format!("{prefix}.post_attention_layernorm.weight"),
8714 vec![hidden_dim],
8715 hidden_dim,
8716 );
8717 push(
8718 &mut tensors,
8719 format!("{prefix}.self_attention_res_norm.weight"),
8720 vec![hidden_dim],
8721 hidden_dim,
8722 );
8723 push(
8724 &mut tensors,
8725 format!("{prefix}.self_attention_res_proj.weight"),
8726 vec![1, hidden_dim],
8727 hidden_dim,
8728 );
8729 push(
8730 &mut tensors,
8731 format!("{prefix}.mlp_res_norm.weight"),
8732 vec![hidden_dim],
8733 hidden_dim,
8734 );
8735 push(
8736 &mut tensors,
8737 format!("{prefix}.mlp_res_proj.weight"),
8738 vec![1, hidden_dim],
8739 hidden_dim,
8740 );
8741 push(
8742 &mut tensors,
8743 format!("{prefix}.self_attn.q_proj.weight"),
8744 vec![kda_proj, hidden_dim],
8745 kda_proj * hidden_dim,
8746 );
8747 push(
8748 &mut tensors,
8749 format!("{prefix}.self_attn.k_proj.weight"),
8750 vec![kda_proj, hidden_dim],
8751 kda_proj * hidden_dim,
8752 );
8753 push(
8754 &mut tensors,
8755 format!("{prefix}.self_attn.v_proj.weight"),
8756 vec![kda_proj, hidden_dim],
8757 kda_proj * hidden_dim,
8758 );
8759 push(
8760 &mut tensors,
8761 format!("{prefix}.self_attn.q_conv1d.weight"),
8762 vec![kda_proj, 1, conv_kernel],
8763 kda_proj * conv_kernel,
8764 );
8765 push(
8766 &mut tensors,
8767 format!("{prefix}.self_attn.k_conv1d.weight"),
8768 vec![kda_proj, 1, conv_kernel],
8769 kda_proj * conv_kernel,
8770 );
8771 push(
8772 &mut tensors,
8773 format!("{prefix}.self_attn.v_conv1d.weight"),
8774 vec![kda_proj, 1, conv_kernel],
8775 kda_proj * conv_kernel,
8776 );
8777 push(
8778 &mut tensors,
8779 format!("{prefix}.self_attn.A_log"),
8780 vec![kda_num_heads],
8781 kda_num_heads,
8782 );
8783 push(
8784 &mut tensors,
8785 format!("{prefix}.self_attn.f_a_proj.weight"),
8786 vec![kda_head_dim, hidden_dim],
8787 kda_head_dim * hidden_dim,
8788 );
8789 push(
8790 &mut tensors,
8791 format!("{prefix}.self_attn.f_b_proj.weight"),
8792 vec![kda_proj, kda_head_dim],
8793 kda_proj * kda_head_dim,
8794 );
8795 push(
8796 &mut tensors,
8797 format!("{prefix}.self_attn.dt_bias"),
8798 vec![kda_proj],
8799 kda_proj,
8800 );
8801 push(
8802 &mut tensors,
8803 format!("{prefix}.self_attn.b_proj.weight"),
8804 vec![kda_num_heads, hidden_dim],
8805 kda_num_heads * hidden_dim,
8806 );
8807 push(
8808 &mut tensors,
8809 format!("{prefix}.self_attn.g_proj.weight"),
8810 vec![kda_proj, hidden_dim],
8811 kda_proj * hidden_dim,
8812 );
8813 push(
8814 &mut tensors,
8815 format!("{prefix}.self_attn.o_norm.weight"),
8816 vec![kda_head_dim],
8817 kda_head_dim,
8818 );
8819 push(
8820 &mut tensors,
8821 format!("{prefix}.self_attn.o_proj.weight"),
8822 vec![hidden_dim, kda_proj],
8823 hidden_dim * kda_proj,
8824 );
8825 push(
8826 &mut tensors,
8827 format!("{prefix}.mlp.gate_proj.weight"),
8828 vec![dense_intermediate, hidden_dim],
8829 dense_intermediate * hidden_dim,
8830 );
8831 push(
8832 &mut tensors,
8833 format!("{prefix}.mlp.up_proj.weight"),
8834 vec![dense_intermediate, hidden_dim],
8835 dense_intermediate * hidden_dim,
8836 );
8837 push(
8838 &mut tensors,
8839 format!("{prefix}.mlp.down_proj.weight"),
8840 vec![hidden_dim, dense_intermediate],
8841 hidden_dim * dense_intermediate,
8842 );
8843 push(
8844 &mut tensors,
8845 "language_model.model.embed_tokens.weight".to_string(),
8846 vec![vocab_size, hidden_dim],
8847 vocab_size * hidden_dim,
8848 );
8849 push(
8850 &mut tensors,
8851 "language_model.lm_head.weight".to_string(),
8852 vec![vocab_size, hidden_dim],
8853 vocab_size * hidden_dim,
8854 );
8855 push(
8856 &mut tensors,
8857 "language_model.model.norm.weight".to_string(),
8858 vec![hidden_dim],
8859 hidden_dim,
8860 );
8861 push(
8862 &mut tensors,
8863 "language_model.model.output_attn_res_norm.weight".to_string(),
8864 vec![hidden_dim],
8865 hidden_dim,
8866 );
8867 push(
8868 &mut tensors,
8869 "language_model.model.output_attn_res_proj.weight".to_string(),
8870 vec![1, hidden_dim],
8871 hidden_dim,
8872 );
8873
8874 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
8887 let dir = std::env::temp_dir().join(format!(
8888 "ferrox_server_kimi_e2e_test_{}_{}",
8889 std::process::id(),
8890 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8891 ));
8892 std::fs::create_dir_all(&dir).unwrap();
8893 let shard_bytes = write_safetensors_shard(&tensors);
8894 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
8895 let map_entries: Vec<String> = tensors
8896 .iter()
8897 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
8898 .collect();
8899 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
8900 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
8901
8902 use base64::Engine;
8906 let vocab_lines: Vec<String> = (0..vocab_size as u32)
8907 .map(|b| {
8908 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
8909 format!("{b64} {b}")
8910 })
8911 .collect();
8912 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
8913
8914 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
8915 .expect("must load the synthetic Kimi checkpoint end to end");
8916 std::fs::remove_dir_all(&dir).ok();
8917 loaded
8918 }
8919
8920 #[test]
8928 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
8929 let loaded = build_synthetic_kimi_loaded();
8930 let state = build_app_state(
8931 StartupModels {
8932 loaded: model::LoadedModel::Kimi(loaded),
8933 embedding: None,
8934 },
8935 None,
8936 None,
8937 None,
8938 false,
8939 None,
8940 Arc::new(health::Detection::ready(health::probe_backends())),
8941 );
8942 let active = state.active().expect("a freshly built state has a model");
8943 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
8944 assert!(!active.is_synthetic());
8945
8946 let (_chunks, finish, _usage) = run_generation(
8947 active.generative().unwrap(),
8948 "hi",
8949 &greedy_params(5),
8950 None,
8951 None,
8952 None,
8953 None,
8954 None,
8955 None,
8956 )
8957 .expect("a real Kimi checkpoint must generate without error");
8958 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
8959 }
8960
8961 #[test]
8974 fn a_grammar_constrains_the_engine_decode_path() {
8975 let loaded = build_synthetic_kimi_loaded();
8976 let state = build_app_state(
8977 StartupModels {
8978 loaded: model::LoadedModel::Kimi(loaded),
8979 embedding: None,
8980 },
8981 None,
8982 None,
8983 None,
8984 false,
8985 None,
8986 Arc::new(health::Detection::ready(health::probe_backends())),
8987 );
8988 let active = state.active().expect("a freshly built state has a model");
8989
8990 let run = |grammar: Option<&str>| {
8991 let mut params = greedy_params(6);
8992 params.grammar = grammar.map(|src| {
8993 Arc::new(
8994 ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
8995 .expect("test grammar parses"),
8996 )
8997 });
8998 run_generation(
8999 active.generative().unwrap(),
9000 "hi",
9001 ¶ms,
9002 None,
9003 None,
9004 None,
9005 None,
9006 None,
9007 None,
9008 )
9009 };
9010
9011 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
9012 let unconstrained = chunks.concat();
9013 assert!(
9014 unconstrained.chars().any(|c| c != 'a'),
9015 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9016 constrained run below would prove nothing"
9017 );
9018
9019 let (chunks, finish, _) =
9020 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9021 let constrained = chunks.concat();
9022 assert!(
9023 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9024 "the engine decode path served text its grammar forbids ({constrained:?}): \
9025 the constraint was dropped between `generate_engine` and the sampler"
9026 );
9027 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9028 }
9029
9030 #[test]
9040 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9041 let loaded = build_synthetic_kimi_loaded();
9042 let state = build_app_state(
9043 StartupModels {
9044 loaded: model::LoadedModel::Kimi(loaded),
9045 embedding: None,
9046 },
9047 None,
9048 None,
9049 None,
9050 false,
9051 None,
9052 Arc::new(health::Detection::ready(health::probe_backends())),
9053 );
9054
9055 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
9056 let kv_pool_config = generate::KvPoolConfig {
9057 pool: pool.clone(),
9058 queue_wait: Duration::ZERO,
9059 };
9060 let pc = Mutex::new(PrefixCache::new(4));
9061
9062 run_generation(
9063 state
9064 .active()
9065 .expect("a freshly built state has a model")
9066 .generative()
9067 .unwrap(),
9068 "hi",
9069 &greedy_params(5),
9070 Some(&kv_pool_config),
9071 None,
9072 Some(&pc),
9073 None,
9074 None,
9075 None,
9076 )
9077 .expect("a real Kimi checkpoint must generate without error");
9078
9079 assert_eq!(
9080 pool.lock().unwrap().free_blocks(),
9081 4,
9082 "the KV pool must be completely untouched by a Kimi request"
9083 );
9084 let stats = pc.lock().unwrap().stats();
9085 assert_eq!(
9086 stats.hits + stats.misses,
9087 0,
9088 "the prefix cache must never be consulted for a Kimi request"
9089 );
9090 }
9091}