1mod admin;
38mod anthropic;
39mod attribution;
40mod budget;
41mod cache_admin;
42mod cancel;
43mod chat_template;
44mod completion;
45mod conversations;
46mod decode_task;
47mod embeddings;
48mod generate;
49mod grammar_request;
50mod health;
51mod journal;
52mod json_mode;
53mod limits;
54mod loaded;
55mod mcp;
56mod model;
57mod openai_extra;
58mod output;
59mod policy;
60mod response_cache;
61pub(crate) mod responses;
62mod resume;
63mod sample_step;
64mod sampling_knobs;
65mod security;
66mod serving;
67mod session;
68mod sse;
69mod stats;
70mod stop;
71mod stream_events;
72mod tasks;
73mod tool_grammar;
74mod unsupported_sampling;
75
76use std::cell::RefCell;
77use std::convert::Infallible;
78use std::fmt;
79use std::net::{IpAddr, Ipv4Addr, SocketAddr};
80use std::path::PathBuf;
81use std::rc::Rc;
82use std::str::FromStr;
83use std::sync::{Arc, Mutex, MutexGuard};
84use std::time::Duration;
85
86use axum::{
87 extract::State,
88 http::StatusCode,
89 response::sse::{Event, Sse},
90 response::{IntoResponse, Response},
91 routing::{get, post},
92 Json, Router,
93};
94use clap::{Parser, ValueEnum};
95use serde::{Deserialize, Serialize};
96
97use ferrox_core::cache::KvBlockPool;
98use ferrox_models::kimi_tokenizer::KimiTokenizer;
99use ferrox_models::sampling::SamplingParams;
100use ferrox_models::tokenizer::StopTokens;
101use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
102use generate::{FinishReason, GenerationParams};
103pub(crate) use loaded::{ActiveModel, Loaded};
104use model::ServerTokenizer;
105use 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(long = "allow-multiple-instances", default_value_t = false)]
185 allow_multiple_instances: bool,
186}
187
188impl ServerArgs {
189 pub fn parse_llama_style<I>(argv: I) -> Self
196 where
197 I: IntoIterator<Item = String>,
198 {
199 Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
200 }
201}
202
203#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
204enum OffloadDevice {
205 Auto,
206 None,
207 Cpu,
208 Metal,
209 Cuda,
210}
211
212#[derive(Debug, Clone, Copy, PartialEq, Eq)]
213enum GpuLayers {
214 Auto,
215 All,
216 Count(u32),
217}
218
219impl GpuLayers {
220 fn offload_enabled(self) -> bool {
221 !matches!(self, Self::Count(0))
222 }
223}
224
225impl FromStr for GpuLayers {
226 type Err = String;
227
228 fn from_str(value: &str) -> Result<Self, Self::Err> {
229 match value {
230 "auto" => Ok(Self::Auto),
231 "all" => Ok(Self::All),
232 _ => value
233 .parse::<u32>()
234 .map(Self::Count)
235 .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
236 }
237 }
238}
239
240impl fmt::Display for GpuLayers {
241 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
242 match self {
243 Self::Auto => f.write_str("auto"),
244 Self::All => f.write_str("all"),
245 Self::Count(value) => value.fmt(f),
246 }
247 }
248}
249
250pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
259
260pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
263
264fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
265 args.into_iter()
266 .map(|arg| match arg.as_str() {
267 "-ngl" => "--n-gpu-layers".into(),
268 "-dev" => "--device".into(),
269 _ => arg,
270 })
271 .collect()
272}
273
274fn print_available_devices() {
275 println!("Available devices:");
276 println!(" CPU");
277
278 let metal = ferrox_metal::MetalProfile::detect();
279 if let Some(name) = metal.device_name {
280 println!(" Metal: {name}");
281 }
282
283 let cuda = ferrox_cuda::HardwareProfile::detect();
284 if cuda.cuda_available {
285 let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
286 println!(" CUDA: {name}");
287 if cuda.cuda_device_count > 1 {
288 println!(" ({} devices detected)", cuda.cuda_device_count);
289 }
290 }
291}
292
293fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
294 if args.host.is_none() && args.port.is_none() {
295 return None;
296 }
297
298 let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
299 let host = args
300 .host
301 .or_else(|| existing.map(|addr| addr.ip()))
302 .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
303 let port = args
304 .port
305 .or_else(|| existing.map(|addr| addr.port()))
306 .unwrap_or(8383);
307 Some(SocketAddr::new(host, port).to_string())
308}
309
310fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
311 if let Some(model) = &args.model {
312 unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
314 }
315
316 if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
317 unsafe { std::env::set_var("FERROX_ADDR", addr) };
319 }
320
321 if let Some(threads) = args.threads {
322 if threads == 0 {
323 anyhow::bail!("--threads must be greater than zero");
324 }
325 unsafe {
327 std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
328 std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
329 }
330 }
331
332 if args.device.is_none() && args.n_gpu_layers.is_none() {
333 } else {
335 let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
336 let device = if layers.offload_enabled() {
337 args.device.unwrap_or(OffloadDevice::Auto)
338 } else {
339 OffloadDevice::None
340 };
341
342 match device {
343 OffloadDevice::None | OffloadDevice::Cpu => unsafe {
344 std::env::set_var("FERROX_METAL", "0");
345 std::env::set_var("FERROX_METAL_ATTN", "0");
346 std::env::set_var("FERROX_CUDA", "0");
347 },
348 OffloadDevice::Auto => unsafe {
349 std::env::set_var("FERROX_METAL", "auto");
350 std::env::set_var("FERROX_CUDA", "auto");
351 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
352 std::env::set_var("FERROX_METAL_ATTN", "1");
353 }
354 },
355 OffloadDevice::Metal => {
356 #[cfg(not(feature = "metal"))]
357 {
358 anyhow::bail!(
359 "Metal requested but this binary was built without --features metal"
360 );
361 }
362 #[cfg(feature = "metal")]
363 {
364 if !ferrox_metal::MetalProfile::detect().available {
365 anyhow::bail!("Metal requested but no Metal device is available");
366 }
367 unsafe {
368 std::env::set_var("FERROX_METAL", "1");
369 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
370 std::env::set_var("FERROX_METAL_ATTN", "1");
371 }
372 std::env::set_var("FERROX_CUDA", "0");
373 }
374 }
375 }
376 OffloadDevice::Cuda => {
377 #[cfg(not(feature = "cuda"))]
378 {
379 anyhow::bail!(
380 "CUDA requested but this binary was built without --features cuda"
381 );
382 }
383 #[cfg(feature = "cuda")]
384 {
385 if !ferrox_cuda::HardwareProfile::detect().cuda_available {
386 anyhow::bail!("CUDA requested but no CUDA device is available");
387 }
388 unsafe {
389 std::env::set_var("FERROX_CUDA", "1");
390 std::env::set_var("FERROX_METAL", "0");
391 std::env::set_var("FERROX_METAL_ATTN", "0");
392 }
393 }
394 }
395 }
396 }
397
398 Ok(())
399}
400
401#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
410 Gguf(GgufModel),
411 Kimi(KimiModel),
412 Mla(MlaModel),
413 Gemma4(Gemma4Model),
414 Glm52(Glm52Model),
415}
416
417pub(crate) struct GgufModel {
418 decoder: Arc<Decoder>,
419 tokenizer: Arc<ServerTokenizer>,
420 stop_tokens: StopTokens,
421 bos_id: Option<usize>,
422 is_synthetic: bool,
423 chat_template: chat_template::PromptTemplate,
424}
425
426pub(crate) struct KimiModel {
427 engine: KimiEngine,
428 tokenizer: KimiTokenizer,
429 stop_tokens: StopTokens,
430 chat_template: chat_template::PromptTemplate,
431}
432
433pub(crate) struct MlaModel {
434 engine: MlaEngine,
435 tokenizer: ServerTokenizer,
436 stop_tokens: StopTokens,
437 bos_id: Option<usize>,
438 name: String,
439 chat_template: chat_template::PromptTemplate,
440}
441
442pub(crate) struct Gemma4Model {
443 engine: Gemma4Engine,
444 tokenizer: ServerTokenizer,
445 stop_tokens: StopTokens,
446 bos_id: Option<usize>,
447 name: String,
448 chat_template: chat_template::PromptTemplate,
449}
450
451pub(crate) struct Glm52Model {
452 engine: ferrox_models::Glm52Engine,
453 tokenizer: ServerTokenizer,
454 stop_tokens: StopTokens,
455 bos_id: Option<usize>,
456 name: String,
457 chat_template: chat_template::PromptTemplate,
458}
459
460impl Model {
461 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
462 match self {
463 Model::Gguf(m) => m.chat_template.clone(),
464 Model::Kimi(m) => m.chat_template.clone(),
465 Model::Mla(m) => m.chat_template.clone(),
466 Model::Gemma4(m) => m.chat_template.clone(),
467 Model::Glm52(m) => m.chat_template.clone(),
468 }
469 }
470
471 fn is_synthetic(&self) -> bool {
476 match self {
477 Model::Gguf(m) => m.is_synthetic,
478 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
479 }
480 }
481
482 fn tokenizer_kind(&self) -> &'static str {
483 match self {
484 Model::Gguf(m) => m.tokenizer.kind(),
485 Model::Kimi(_) => "kimi-tiktoken-bpe",
486 Model::Mla(m) => m.tokenizer.kind(),
487 Model::Gemma4(m) => m.tokenizer.kind(),
488 Model::Glm52(m) => m.tokenizer.kind(),
489 }
490 }
491
492 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
496 match self {
497 Model::Gguf(m) => m.decoder.expert_store_stats(),
498 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
499 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
500 }
501 }
502
503 pub(crate) fn name(&self) -> &str {
504 match self {
505 Model::Gguf(m) => m.decoder.config.name,
506 Model::Kimi(_) => "kimi-k3",
507 Model::Mla(m) => m.name.as_str(),
508 Model::Gemma4(m) => m.name.as_str(),
509 Model::Glm52(m) => m.name.as_str(),
510 }
511 }
512
513 pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
514 match self {
515 Model::Gguf(m) => m.tokenizer.encode(text),
516 Model::Kimi(m) => m
517 .tokenizer
518 .encode(text)
519 .into_iter()
520 .map(|id| id as usize)
521 .collect(),
522 Model::Mla(m) => m.tokenizer.encode(text),
523 Model::Gemma4(m) => m.tokenizer.encode(text),
524 Model::Glm52(m) => m.tokenizer.encode(text),
525 }
526 }
527
528 pub(crate) fn bos_id(&self) -> Option<usize> {
537 match self {
538 Model::Gguf(m) => m.bos_id,
539 Model::Kimi(_) => None,
540 Model::Mla(m) => m.bos_id,
541 Model::Gemma4(m) => m.bos_id,
542 Model::Glm52(m) => m.bos_id,
543 }
544 }
545
546 pub(crate) fn decode(&self, ids: &[usize]) -> String {
547 match self {
548 Model::Gguf(m) => m.tokenizer.decode(ids),
549 Model::Kimi(m) => {
550 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
551 m.tokenizer.decode(&ids32)
552 }
553 Model::Mla(m) => m.tokenizer.decode(ids),
554 Model::Gemma4(m) => m.tokenizer.decode(ids),
555 Model::Glm52(m) => m.tokenizer.decode(ids),
556 }
557 }
558
559 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
562 match self {
563 Model::Gguf(m) => {
564 let mut caches: Vec<_> = (0..m.decoder.layers.len())
565 .map(|_| {
566 ferrox_core::cache::KvCache::new(
567 m.decoder.config.n_kv_heads,
568 m.decoder.config.head_dim,
569 )
570 })
571 .collect();
572 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
573 }
574 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
575 }
576 }
577
578 pub(crate) fn vocab_size(&self) -> Option<usize> {
579 match self {
580 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
581 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
582 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
583 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
584 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
585 }
586 }
587}
588
589pub(crate) struct AppState {
590 pub(crate) embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
604 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
623 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
627 pub(crate) tasks: Arc<tasks::TaskRegistry>,
629 pub(crate) cancels: Arc<cancel::CancelRegistry>,
633 pub(crate) stats: stats::Stats,
636 pub(crate) streams: resume::StreamRegistry,
639 pub(crate) model_dir: Option<PathBuf>,
641 response_cache: Mutex<ResponseCache>,
645 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
657 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
670 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
680 sessions: session::SessionStore,
686 requests_total: std::sync::atomic::AtomicU64,
687 request_errors_total: std::sync::atomic::AtomicU64,
688 started_at: std::time::Instant,
689 last_request_ms: std::sync::atomic::AtomicU64,
694 detection: Arc<health::Detection>,
696 mcp: Option<mcp::LoadedMcpConfig>,
698 pub(crate) continuous_batching_enabled: bool,
702 loading_model: Mutex<Option<String>>,
706 last_load_error: Mutex<Option<(String, String)>>,
710 pub(crate) serving: Mutex<crate::stats::ServingStats>,
715 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
720 pub(crate) footprint:
726 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
727 pub(crate) started_unix: u64,
737}
738
739pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
745
746impl AppState {
747 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
751 self.active
752 .read()
753 .unwrap_or_else(|p| p.into_inner())
754 .clone()
755 }
756
757 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
763 self.active().ok_or_else(|| {
764 (
765 StatusCode::SERVICE_UNAVAILABLE,
766 Json(serde_json::json!({"error": {
767 "message": "no model is loaded; POST /admin/models/load with an id from \
768 GET /admin/models",
769 "type": "model_not_loaded"
770 }})),
771 )
772 })
773 }
774
775 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
783 Ok(Arc::clone(self.require_active()?.generative()?))
784 }
785
786 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
794 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
795 std::mem::replace(&mut *guard, next)
796 }
797
798 fn mark_request_finished(&self) {
802 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
803 self.last_request_ms
804 .store(ms, std::sync::atomic::Ordering::Relaxed);
805 }
806
807 pub(crate) fn uptime(&self) -> Duration {
808 self.started_at.elapsed()
809 }
810
811 pub(crate) fn requests_total(&self) -> u64 {
812 self.requests_total
813 .load(std::sync::atomic::Ordering::Relaxed)
814 }
815
816 pub(crate) fn errors_total(&self) -> u64 {
817 self.request_errors_total
818 .load(std::sync::atomic::Ordering::Relaxed)
819 }
820
821 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
822 lock_cache(&self.response_cache).stats()
823 }
824
825 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
828 let last = self
829 .last_request_ms
830 .load(std::sync::atomic::Ordering::Relaxed);
831 (last > 0)
832 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
833 .map(|age| age.max(0.0))
834 }
835
836 pub(crate) fn loading_model_id(&self) -> Option<String> {
837 self.loading_model
838 .lock()
839 .unwrap_or_else(|p| p.into_inner())
840 .clone()
841 }
842
843 pub(crate) fn set_loading_model(&self, id: Option<String>) {
844 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
845 }
846
847 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
848 self.last_load_error
849 .lock()
850 .unwrap_or_else(|p| p.into_inner())
851 .clone()
852 }
853
854 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
855 *self
856 .last_load_error
857 .lock()
858 .unwrap_or_else(|p| p.into_inner()) = error;
859 }
860
861 pub(crate) fn active_model_name(&self) -> Option<String> {
870 self.active().map(|a| a.name().to_string())
871 }
872
873 pub(crate) fn embedding_model(&self) -> Option<Arc<ferrox_models::EmbeddingModel>> {
883 self.embedding
884 .clone()
885 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
886 }
887
888 pub(crate) fn embedding_model_name(&self) -> Option<String> {
892 match self.embedding_model() {
893 Some(e) => Some(e.name().to_string()),
894 None => self.active_model_name(),
895 }
896 }
897
898 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
899 self.stats.record(stats::entry(record));
900 }
901}
902
903fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
908 cache
909 .lock()
910 .unwrap_or_else(|poisoned| poisoned.into_inner())
911}
912
913#[derive(Debug, Clone, Deserialize)]
914#[serde(untagged)]
915pub(crate) enum MessageContent {
916 Text(String),
917 Parts(Vec<ContentPart>),
918}
919
920#[derive(Debug, Clone, Deserialize)]
921struct ContentPart {
922 #[serde(rename = "type")]
923 kind: String,
924 #[serde(default)]
925 text: Option<String>,
926 #[serde(default)]
927 image_url: Option<serde_json::Value>,
928}
929
930impl MessageContent {
931 fn as_text(&self) -> String {
932 match self {
933 Self::Text(s) => s.clone(),
934 Self::Parts(parts) => parts
935 .iter()
936 .filter_map(|p| p.text.as_deref())
937 .collect::<Vec<_>>()
938 .join(""),
939 }
940 }
941
942 fn has_image(&self) -> bool {
943 match self {
944 Self::Text(_) => false,
945 Self::Parts(parts) => parts
946 .iter()
947 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
948 }
949 }
950}
951
952#[derive(Debug, Clone, Deserialize)]
953pub(crate) struct ChatMessage {
954 pub(crate) role: String,
955 #[serde(default)]
959 pub(crate) content: Option<MessageContent>,
960 #[serde(default)]
964 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
965 #[serde(default)]
970 #[allow(dead_code)]
971 pub(crate) tool_call_id: Option<String>,
972 #[serde(default, alias = "reasoning")]
988 pub(crate) reasoning_content: Option<String>,
989}
990
991impl ChatMessage {
992 fn rendered_content(&self) -> String {
999 let mut out = self
1000 .content
1001 .as_ref()
1002 .map(MessageContent::as_text)
1003 .unwrap_or_default();
1004 if let Some(calls) = &self.tool_calls {
1005 for call in calls {
1006 out.push_str(&format!(
1007 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
1008 call.function.name, call.function.arguments
1009 ));
1010 }
1011 }
1012 out
1013 }
1014}
1015
1016#[derive(Debug, Clone, Deserialize)]
1017pub(crate) struct ToolCallIn {
1018 #[serde(default)]
1019 #[allow(dead_code)]
1020 id: String,
1021 #[serde(rename = "type", default)]
1022 #[allow(dead_code)]
1023 kind: String,
1024 function: ToolCallFunctionIn,
1025}
1026
1027#[derive(Debug, Clone, Deserialize)]
1028struct ToolCallFunctionIn {
1029 name: String,
1030 arguments: String,
1035}
1036
1037#[derive(Debug, Clone, Deserialize)]
1040struct ToolDef {
1041 #[serde(rename = "type", default)]
1042 #[allow(dead_code)]
1043 kind: String,
1044 function: ToolFunctionDef,
1045}
1046
1047#[derive(Debug, Clone, Deserialize)]
1048struct ToolFunctionDef {
1049 name: String,
1050 #[serde(default)]
1051 description: Option<String>,
1052 #[serde(default)]
1053 parameters: Option<serde_json::Value>,
1054}
1055
1056#[derive(Debug, Clone, Deserialize)]
1067#[serde(untagged)]
1068enum ToolChoice {
1069 Mode(String),
1070 Specific(serde_json::Value),
1071}
1072
1073#[derive(Deserialize)]
1076#[serde(untagged)]
1077enum StopParam {
1078 One(String),
1079 Many(Vec<String>),
1080}
1081
1082#[derive(Deserialize)]
1083struct ChatCompletionRequest {
1084 model: String,
1085 messages: Vec<ChatMessage>,
1086 #[serde(default = "default_max_tokens")]
1087 max_tokens: usize,
1088 #[serde(default)]
1089 temperature: Option<f32>,
1090 #[serde(default)]
1091 top_p: Option<f32>,
1092 #[serde(default)]
1097 min_p: Option<f32>,
1098 #[serde(default)]
1099 top_k: Option<usize>,
1100 #[serde(default)]
1101 repetition_penalty: Option<f32>,
1102 #[serde(default)]
1103 seed: Option<u64>,
1104 #[serde(default)]
1105 stop: Option<StopParam>,
1106 #[serde(default)]
1107 stream: Option<bool>,
1108 #[serde(default)]
1121 stream_resumable: Option<bool>,
1122 #[serde(default)]
1132 ignore_eos: Option<bool>,
1133 #[serde(default)]
1134 tools: Vec<ToolDef>,
1135 #[serde(default)]
1136 tool_choice: Option<ToolChoice>,
1137 #[serde(default)]
1145 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
1146 #[serde(default)]
1154 reasoning_effort: Option<String>,
1155 #[serde(default)]
1159 thinking: Option<ThinkingSwitch>,
1160 #[serde(default)]
1165 session_id: Option<String>,
1166 #[serde(default)]
1168 logprobs: Option<bool>,
1169 #[serde(default)]
1170 top_logprobs: Option<u32>,
1171 #[serde(default)]
1172 n: Option<u32>,
1173 #[serde(default)]
1174 presence_penalty: Option<f32>,
1175 #[serde(default)]
1176 frequency_penalty: Option<f32>,
1177 #[serde(default)]
1178 response_format: Option<serde_json::Value>,
1179 #[serde(default)]
1186 logit_bias: Option<serde_json::Value>,
1187 #[serde(default)]
1196 grammar: Option<String>,
1197}
1198
1199const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1211
1212#[derive(Debug, Clone, Deserialize)]
1214pub(crate) struct ThinkingSwitch {
1215 #[serde(rename = "type")]
1216 pub(crate) kind: String,
1217}
1218
1219const THINKING_KWARG_KEYS: [&str; 4] = [
1223 "enable_thinking",
1224 "thinking",
1225 "thinking_mode",
1226 "reasoning_effort",
1227];
1228
1229const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1233
1234fn default_max_tokens() -> usize {
1235 DEFAULT_CHAT_MAX_TOKENS
1236}
1237
1238impl ChatCompletionRequest {
1239 fn sampling_knobs(&self) -> SamplingKnobs {
1244 SamplingKnobs {
1245 temperature: self.temperature,
1246 top_p: self.top_p,
1247 min_p: self.min_p,
1248 top_k: self.top_k,
1249 repetition_penalty: self.repetition_penalty,
1250 presence_penalty: self.presence_penalty,
1251 frequency_penalty: self.frequency_penalty,
1252 penalty_last_n: None,
1256 }
1257 }
1258
1259 fn sampling_params(&self) -> SamplingParams {
1260 self.sampling_knobs().resolve()
1261 }
1262
1263 fn stop_sequences(&self) -> Vec<String> {
1264 self.stop
1265 .as_ref()
1266 .map(|s| match s {
1267 StopParam::One(v) => vec![v.clone()],
1268 StopParam::Many(v) => v.clone(),
1269 })
1270 .unwrap_or_default()
1271 }
1272
1273 fn tools_active(&self) -> bool {
1278 !self.tools.is_empty()
1279 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1280 }
1281
1282 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1294 match &self.tool_choice {
1295 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1296 Some(ToolChoice::Specific(value)) => {
1297 let name = value
1301 .get("function")
1302 .and_then(|f| f.get("name"))
1303 .or_else(|| value.get("name"))
1304 .and_then(|n| n.as_str());
1305 match name {
1306 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1307 None => Err(invalid_request(
1308 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1309 function.name",
1310 "tool_choice",
1311 )),
1312 }
1313 }
1314 _ => Ok(None),
1315 }
1316 }
1317
1318 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1320 self.tools
1321 .iter()
1322 .map(|t| tool_grammar::ToolSpec {
1323 name: &t.function.name,
1324 parameters: t.function.parameters.as_ref(),
1325 })
1326 .collect()
1327 }
1328
1329 fn thinking_direction(&self) -> Option<bool> {
1372 if let Some(switch) = &self.thinking {
1373 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1374 "disabled" => Some(false),
1375 "enabled" => Some(true),
1376 _ => None,
1379 };
1380 }
1381 let effort = self.reasoning_effort.as_ref()?;
1382 DISABLE_EFFORTS
1383 .contains(&effort.trim().to_ascii_lowercase().as_str())
1384 .then_some(false)
1385 }
1386
1387 fn resolve_template_kwargs(
1388 &self,
1389 template: &chat_template::PromptTemplate,
1390 ) -> serde_json::Map<String, serde_json::Value> {
1391 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1392 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1396
1397 if !caller_steered {
1398 match self.thinking_direction() {
1399 Some(false) => {
1400 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1401 kwargs.insert(k, v);
1402 }
1403 return kwargs;
1406 }
1407 Some(true) => {
1408 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1409 kwargs.insert(k, v);
1410 }
1411 }
1412 None => {}
1413 }
1414 if let Some(effort) = &self.reasoning_effort {
1415 kwargs
1416 .entry("reasoning_effort".to_string())
1417 .or_insert_with(|| serde_json::json!(effort));
1418 }
1419 }
1420
1421 let offered: Vec<serde_json::Value> = if self.tools_active() {
1422 self.tools.iter().map(chat_template::tool_json).collect()
1423 } else {
1424 Vec::new()
1425 };
1426 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1427 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1428 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1429 kwargs.entry(k).or_insert(v);
1430 }
1431 }
1432 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1433 crate::policy::effort::EffortMapping::Mapped(to) => {
1434 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1435 }
1436 crate::policy::effort::EffortMapping::Dropped => {
1437 tracing::debug!(
1438 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1439 enough, so its own default applies"
1440 );
1441 }
1442 crate::policy::effort::EffortMapping::Unchanged => {}
1443 }
1444 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1445 kwargs
1446 }
1447
1448 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1451 if self.max_tokens == 0 {
1458 return Err(invalid_request(
1459 "max_tokens must be at least 1",
1460 "max_tokens",
1461 ));
1462 }
1463 if let Some(switch) = &self.thinking {
1468 let kind = switch.kind.trim().to_ascii_lowercase();
1469 if kind != "enabled" && kind != "disabled" {
1470 return Err(invalid_request(
1471 "thinking.type must be \"enabled\" or \"disabled\"",
1472 "thinking.type",
1473 ));
1474 }
1475 }
1476 for msg in &self.messages {
1477 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1478 return Err(unsupported_feature(
1479 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1480 ));
1481 }
1482 }
1483 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1484 return Err(unsupported_feature(
1485 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1486 ));
1487 }
1488 if self.n.is_some_and(|n| n > 1) {
1489 return Err(unsupported_feature(
1490 "n > 1 is not implemented (single completion only)",
1491 ));
1492 }
1493 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1494 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1502 if let Some(fmt) = &self.response_format {
1503 match fmt.get("type").and_then(|v| v.as_str()) {
1504 Some("json_object") => {}
1505 Some(other) => {
1509 return Err((
1510 StatusCode::BAD_REQUEST,
1511 Json(serde_json::json!({
1512 "error": {
1513 "message": format!(
1514 "response_format type {other:?} is not supported (only json_object)"
1515 )
1516 }
1517 })),
1518 ));
1519 }
1520 None => {
1521 return Err((
1522 StatusCode::BAD_REQUEST,
1523 Json(serde_json::json!({
1524 "error": {
1525 "message": "response_format must include \"type\" (only json_object is supported)"
1526 }
1527 })),
1528 ));
1529 }
1530 }
1531 }
1532 if let Some(forced) = self.forced_tool_choice()? {
1540 if self.tools.is_empty() {
1541 return Err(invalid_request(
1542 "tool_choice forces a tool call, but no tools were offered",
1543 "tool_choice",
1544 ));
1545 }
1546 if let tool_grammar::Forced::Named(name) = forced {
1547 if !self.tools.iter().any(|t| t.function.name == name) {
1548 return Err(invalid_request(
1549 &format!(
1550 "tool_choice names {name:?}, which is not one of the tools offered"
1551 ),
1552 "tool_choice",
1553 ));
1554 }
1555 }
1556 if self.grammar.is_some() {
1559 return Err(invalid_request(
1560 "a forced tool_choice and a \"grammar\" are two different constraints on the \
1561 same generation; send one",
1562 "tool_choice",
1563 ));
1564 }
1565 if self.json_object_mode() {
1566 return Err(invalid_request(
1567 "a forced tool_choice cannot be combined with response_format json_object: \
1568 the tool-call markers are not JSON",
1569 "tool_choice",
1570 ));
1571 }
1572 }
1573 Ok(())
1574 }
1575
1576 fn effective_stop_sequences(&self) -> Vec<String> {
1583 let mut stop = self.stop_sequences();
1584 if self.tools_active() {
1585 stop.push("</tool_call>".to_string());
1586 }
1587 stop
1588 }
1589
1590 fn json_object_mode(&self) -> bool {
1591 self.response_format
1592 .as_ref()
1593 .and_then(|v| v.get("type"))
1594 .and_then(|v| v.as_str())
1595 == Some("json_object")
1596 }
1597
1598 fn generation_params(&self) -> Result<GenerationParams, ApiError> {
1603 Ok(GenerationParams {
1604 max_tokens: self.max_tokens,
1605 sampling: self.sampling_params(),
1606 seed: self.resolved_seed(),
1607 stop: self.effective_stop_sequences(),
1608 stop_token_ids: Vec::new(),
1612 json_object: self.json_object_mode(),
1613 grammar: grammar_request::for_request(
1614 self.grammar.as_deref(),
1615 self.response_format.as_ref(),
1616 )?,
1617 cancel: None,
1620 ignore_eos: self.ignore_eos.unwrap_or(false),
1621 })
1622 }
1623
1624 fn generation_params_for_template(
1635 &self,
1636 template: &chat_template::PromptTemplate,
1637 served_model: &str,
1638 ) -> Result<GenerationParams, ApiError> {
1639 let mut params = self.generation_params()?;
1640 if let Some(stop) = template.end_of_turn() {
1641 if !params.stop.iter().any(|s| s == stop) {
1642 params.stop.push(stop.to_string());
1643 }
1644 }
1645 if let Some(forced) = self.forced_tool_choice()? {
1646 params.grammar = Some(tool_grammar::build(
1650 forced,
1651 &self.tool_specs(),
1652 policy::parser::ToolCallFormat::infer(served_model),
1653 )?);
1654 }
1655 Ok(params)
1656 }
1657
1658 fn is_cacheable(&self) -> bool {
1667 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1668 }
1669
1670 fn cache_key(&self, prompt: &str) -> CacheKey {
1671 CacheKey {
1672 model: self.model.clone(),
1673 prompt: prompt.to_string(),
1674 max_tokens: self.max_tokens,
1675 sampling: response_cache::sampling_key(&self.sampling_params()),
1680 seed: self.seed,
1681 stop: self.effective_stop_sequences(),
1682 }
1683 }
1684
1685 fn resolved_seed(&self) -> u64 {
1686 self.seed.unwrap_or_else(|| {
1687 std::time::SystemTime::now()
1688 .duration_since(std::time::UNIX_EPOCH)
1689 .map(|d| d.as_nanos() as u64)
1690 .unwrap_or(0xDEFA017)
1691 })
1692 }
1693}
1694
1695#[derive(Serialize)]
1696struct ChatCompletionChoice {
1697 index: usize,
1698 message: ChatCompletionResponseMessage,
1699 finish_reason: &'static str,
1700}
1701
1702#[derive(Serialize)]
1703struct ChatCompletionResponseMessage {
1704 role: &'static str,
1705 #[serde(skip_serializing_if = "Option::is_none")]
1706 content: Option<String>,
1707 #[serde(skip_serializing_if = "Option::is_none")]
1711 reasoning_content: Option<String>,
1712 #[serde(skip_serializing_if = "Option::is_none")]
1713 tool_calls: Option<Vec<ToolCallOut>>,
1714}
1715
1716#[derive(Serialize, Clone)]
1717struct ToolCallOut {
1718 id: String,
1719 #[serde(rename = "type")]
1720 kind: &'static str,
1721 function: ToolCallFunctionOut,
1722}
1723
1724#[derive(Serialize, Clone)]
1732struct ToolCallDelta {
1733 index: usize,
1734 #[serde(skip_serializing_if = "Option::is_none")]
1735 id: Option<String>,
1736 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1737 kind: Option<&'static str>,
1738 function: ToolCallFunctionDelta,
1739}
1740
1741#[derive(Serialize, Clone, Default)]
1742struct ToolCallFunctionDelta {
1743 #[serde(skip_serializing_if = "Option::is_none")]
1744 name: Option<String>,
1745 #[serde(skip_serializing_if = "Option::is_none")]
1748 arguments: Option<String>,
1749}
1750
1751impl ToolCallDelta {
1752 fn whole(index: usize, name: String, arguments: String) -> Self {
1754 ToolCallDelta {
1755 index,
1756 id: Some(format!("call_{index}")),
1757 kind: Some("function"),
1758 function: ToolCallFunctionDelta {
1759 name: Some(name),
1760 arguments: Some(arguments),
1761 },
1762 }
1763 }
1764
1765 fn opening(index: usize, name: String) -> Self {
1767 ToolCallDelta {
1768 index,
1769 id: Some(format!("call_{index}")),
1770 kind: Some("function"),
1771 function: ToolCallFunctionDelta {
1772 name: Some(name),
1773 arguments: Some(String::new()),
1774 },
1775 }
1776 }
1777
1778 fn arguments(index: usize, fragment: String) -> Self {
1780 ToolCallDelta {
1781 index,
1782 id: None,
1783 kind: None,
1784 function: ToolCallFunctionDelta {
1785 name: None,
1786 arguments: Some(fragment),
1787 },
1788 }
1789 }
1790}
1791
1792#[derive(Serialize, Clone)]
1793struct ToolCallFunctionOut {
1794 name: String,
1795 arguments: String,
1799}
1800
1801#[derive(Serialize)]
1802struct ChatCompletionResponse {
1803 id: String,
1804 request_id: String,
1810 object: &'static str,
1811 model: String,
1812 choices: Vec<ChatCompletionChoice>,
1813 usage: generate::Usage,
1818 ferrox_cache: &'static str,
1826}
1827
1828#[derive(Serialize)]
1829struct ChatCompletionChunkDelta {
1830 #[serde(skip_serializing_if = "Option::is_none")]
1831 role: Option<&'static str>,
1832 #[serde(skip_serializing_if = "Option::is_none")]
1833 content: Option<String>,
1834 #[serde(skip_serializing_if = "Option::is_none")]
1836 reasoning_content: Option<String>,
1837 #[serde(skip_serializing_if = "Option::is_none")]
1838 tool_calls: Option<Vec<ToolCallDelta>>,
1839}
1840
1841#[derive(Serialize)]
1842struct ChatCompletionChunkChoice {
1843 index: usize,
1844 delta: ChatCompletionChunkDelta,
1845 finish_reason: Option<&'static str>,
1846}
1847
1848#[derive(Serialize)]
1849struct ChatCompletionChunk {
1850 id: String,
1851 #[serde(skip_serializing_if = "Option::is_none")]
1858 request_id: Option<String>,
1859 object: &'static str,
1860 model: String,
1861 choices: Vec<ChatCompletionChunkChoice>,
1862 #[serde(skip_serializing_if = "Option::is_none")]
1865 usage: Option<generate::Usage>,
1866}
1867
1868async fn health(State(state): State<Arc<AppState>>) -> Response {
1874 let snapshot = state.detection.snapshot();
1875 let mut capabilities = snapshot.capabilities;
1876 let active = state.active();
1877
1878 capabilities.push(match active.as_deref() {
1881 None => ferrox_api::Capability::unavailable(
1886 ferrox_api::health::capability::REAL_WEIGHTS,
1887 ferrox_api::health::reason::MODEL_NOT_LOADED,
1888 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1889 ),
1890 Some(active) if active.is_synthetic() => ferrox_api::Capability::unavailable(
1891 ferrox_api::health::capability::REAL_WEIGHTS,
1892 ferrox_api::health::reason::MODEL_NOT_LOADED,
1893 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1894 checkpoint. Output from this model is noise.",
1895 ),
1896 Some(active) if active.encoder().is_some() => ferrox_api::Capability::available(
1901 ferrox_api::health::capability::REAL_WEIGHTS,
1902 format!(
1903 "Serving the real embedding checkpoint '{}'. This is an ENCODER: {} serves \
1904 it, generation endpoints refuse it.",
1905 active.name(),
1906 ferrox_api::routes::V1_EMBEDDINGS,
1907 ),
1908 ),
1909 Some(active) => ferrox_api::Capability::available(
1910 ferrox_api::health::capability::REAL_WEIGHTS,
1911 format!("Serving the real checkpoint '{}'.", active.name()),
1912 ),
1913 });
1914 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1915 ferrox_api::Capability::available(
1916 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1917 "Concurrent requests share one batched decode step.",
1918 )
1919 } else {
1920 ferrox_api::Capability::unavailable(
1921 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1922 ferrox_api::health::reason::DISABLED,
1923 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1924 )
1925 });
1926
1927 let last_request_ms = state
1928 .last_request_ms
1929 .load(std::sync::atomic::Ordering::Relaxed);
1930 let uptime = state.started_at.elapsed();
1931 let health_state = if active.is_none() {
1939 ferrox_api::HealthState::Unavailable
1940 } else {
1941 snapshot.state
1942 };
1943 let body = ferrox_api::HealthResponse {
1944 state: health_state,
1945 reason: match health_state {
1946 ferrox_api::HealthState::Ready => None,
1947 ferrox_api::HealthState::Unavailable => {
1948 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
1949 }
1950 ferrox_api::HealthState::Detecting => {
1951 Some(ferrox_api::health::reason::DETECTING.to_string())
1952 }
1953 },
1954 detail: match health_state {
1955 ferrox_api::HealthState::Ready => None,
1956 ferrox_api::HealthState::Unavailable => Some(
1957 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1958 .to_string(),
1959 ),
1960 ferrox_api::HealthState::Detecting => {
1961 Some("Probing available compute backends.".to_string())
1962 }
1963 },
1964 model: active
1965 .as_deref()
1966 .map(|active| ferrox_api::health::ModelSummary {
1967 id: active.name().to_string(),
1968 tokenizer: active.tokenizer_kind().to_string(),
1969 synthetic_weights: active.is_synthetic(),
1970 }),
1971 capabilities,
1972 version: env!("CARGO_PKG_VERSION").to_string(),
1973 pid: std::process::id(),
1974 uptime_seconds: uptime.as_secs_f64(),
1975 server_time_unix_ms: std::time::SystemTime::now()
1976 .duration_since(std::time::UNIX_EPOCH)
1977 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1978 .unwrap_or(0),
1979 last_request_age_seconds: (last_request_ms > 0)
1980 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1981 .map(|age| age.max(0.0)),
1982 };
1983
1984 let status =
1985 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1986 (status, Json(body)).into_response()
1987}
1988
1989async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1990 let Some(active) = state.active() else {
1994 return Json(serde_json::json!({ "object": "list", "data": [] }));
1995 };
1996 let mut model_entry = serde_json::json!({
1997 "id": active.name(),
1998 "object": "model",
1999 "ferrox_synthetic_weights": active.is_synthetic(),
2000 "ferrox_tokenizer": active.tokenizer_kind(),
2001 });
2002 if let Some(encoder) = active.encoder() {
2008 model_entry["ferrox_model_kind"] = serde_json::json!("embedding");
2009 model_entry["ferrox_endpoints"] = serde_json::json!([ferrox_api::routes::V1_EMBEDDINGS]);
2010 model_entry["ferrox_n_embd"] = serde_json::json!(encoder.n_embd());
2011 model_entry["ferrox_pooling"] = serde_json::json!(encoder.pooling_type().name());
2012 model_entry["ferrox_context_length"] = serde_json::json!(encoder.n_ctx_train());
2013 }
2014 if let Some(model) = active.generative_opt() {
2022 let parser_configured =
2023 crate::policy::parser::ReasoningFormat::infer(active.name()).is_some();
2024 let gears = model.chat_template().think_gears(parser_configured);
2025 if !gears.is_empty() {
2026 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
2027 if let Some(default) = &gears.default {
2028 model_entry["default_reasoning_effort"] = serde_json::json!(default);
2029 }
2030 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
2033 }
2034 }
2035 if let Some(mcp) = &state.mcp {
2036 model_entry["ferrox_mcp"] = mcp.models_metadata();
2037 }
2038 Json(serde_json::json!({
2039 "object": "list",
2040 "data": [model_entry]
2041 }))
2042}
2043
2044async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2058 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
2059 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
2060 let active = state.active();
2061 Json(serde_json::json!({
2062 "model": active.as_ref().map(|a| a.name()),
2063 "state": state
2064 .maintenance
2065 .lock()
2066 .unwrap_or_else(|p| p.into_inner())
2067 .state()
2068 .as_str(),
2069 "uptime_s": state.uptime().as_secs(),
2070 "throughput": {
2071 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2072 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2073 },
2074 "requests": {
2075 "active": state.cancels.live_count(),
2076 "completed": state.stats.recorded_total(),
2077 "p95_ms": state.stats.p95_duration_ms(),
2078 "ttft_mean_ms": state.stats.ttft_mean_ms(),
2079 "prompt_tokens_total": state.stats.tokens_prompt_total(),
2080 "completion_tokens_total": state.stats.tokens_generated_total(),
2081 },
2082 "pools": cache_admin::pool_gauges(&state),
2088 "memory": cache_admin::footprint_json(&state),
2091 }))
2092}
2093
2094#[derive(Deserialize)]
2095struct RequestsQuery {
2096 #[serde(default)]
2097 since: u64,
2098 #[serde(default = "default_requests_limit")]
2099 limit: usize,
2100}
2101
2102fn default_requests_limit() -> usize {
2103 stats::MAX_PAGE
2104}
2105
2106async fn recent_requests(
2114 State(state): State<Arc<AppState>>,
2115 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
2116) -> Json<serde_json::Value> {
2117 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
2118 Json(serde_json::json!({
2119 "requests": rows,
2120 "next_cursor": cursor,
2121 "missed": missed,
2122 "total": state.stats.recorded_total(),
2123 }))
2124}
2125
2126#[derive(Serialize)]
2127struct CombinedCacheStats {
2128 response_cache: response_cache::CacheStats,
2129 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
2131}
2132
2133async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
2134 Json(CombinedCacheStats {
2135 response_cache: lock_cache(&state.response_cache).stats(),
2136 prefix_cache: state
2137 .prefix_cache
2138 .as_ref()
2139 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
2140 })
2141}
2142
2143async fn metrics(State(state): State<Arc<AppState>>) -> Response {
2148 use std::sync::atomic::Ordering;
2149
2150 let cache_stats = lock_cache(&state.response_cache).stats();
2151 let active = state.active();
2152 let requests_total = state.requests_total.load(Ordering::Relaxed);
2153 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2154 let uptime = state.started_at.elapsed().as_secs_f64();
2155
2156 let body = format!(
2157 "# HELP ferrox_requests_total Total chat completion requests received.\n\
2158 # TYPE ferrox_requests_total counter\n\
2159 ferrox_requests_total {requests_total}\n\
2160 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
2161 # TYPE ferrox_request_errors_total counter\n\
2162 ferrox_request_errors_total {errors_total}\n\
2163 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
2164 # TYPE ferrox_cache_hits_total counter\n\
2165 ferrox_cache_hits_total {}\n\
2166 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
2167 # TYPE ferrox_cache_misses_total counter\n\
2168 ferrox_cache_misses_total {}\n\
2169 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
2170 # TYPE ferrox_cache_entries gauge\n\
2171 ferrox_cache_entries {}\n\
2172 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2173 # TYPE ferrox_synthetic_weights gauge\n\
2174 ferrox_synthetic_weights {}\n\
2175 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
2176 # TYPE ferrox_uptime_seconds gauge\n\
2177 ferrox_uptime_seconds {uptime}\n",
2178 cache_stats.hits,
2179 cache_stats.misses,
2180 cache_stats.entries,
2181 active
2185 .as_ref()
2186 .map(|a| a.is_synthetic() as u8)
2187 .unwrap_or(0),
2188 );
2189
2190 let body = match active
2194 .as_ref()
2195 .and_then(|a| a.expert_store_stats())
2196 {
2197 Some(es) => format!(
2198 "{body}\
2199 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
2200 # TYPE ferrox_expert_cache_hits_total counter\n\
2201 ferrox_expert_cache_hits_total {}\n\
2202 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2203 # TYPE ferrox_expert_cache_misses_total counter\n\
2204 ferrox_expert_cache_misses_total {}\n\
2205 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
2206 # TYPE ferrox_expert_cache_evictions_total counter\n\
2207 ferrox_expert_cache_evictions_total {}\n\
2208 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2209 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
2210 ferrox_expert_cache_pass_throughs_total {}\n\
2211 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2212 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
2213 ferrox_expert_cache_bytes_read_total {}\n\
2214 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2215 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
2216 ferrox_expert_cache_resident_bytes {}\n",
2217 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2218 ),
2219 None => body,
2220 };
2221
2222 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2227 Some(batcher) => {
2228 let sched = batcher.stats();
2229 format!(
2230 "{body}\
2231 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2232 # TYPE ferrox_prefill_chunks_total counter\n\
2233 ferrox_prefill_chunks_total {}\n\
2234 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2235 # TYPE ferrox_prefill_tokens_total counter\n\
2236 ferrox_prefill_tokens_total {}\n\
2237 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2238 # TYPE ferrox_decode_steps_total counter\n\
2239 ferrox_decode_steps_total {}\n\
2240 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2241 # TYPE ferrox_scheduler_queue_depth gauge\n\
2242 ferrox_scheduler_queue_depth {}\n\
2243 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2244 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2245 ferrox_scheduler_queue_rejected_total {}\n\
2246 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2247 # TYPE ferrox_kv_blocks_total gauge\n\
2248 ferrox_kv_blocks_total {}\n\
2249 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2250 # TYPE ferrox_kv_blocks_free gauge\n\
2251 ferrox_kv_blocks_free {}\n\
2252 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2253 # TYPE ferrox_kv_block_size gauge\n\
2254 ferrox_kv_block_size {}\n\
2255 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2256 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2257 ferrox_kv_rejected_too_large_total {}\n\
2258 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2259 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2260 ferrox_kv_rejected_context_length_total {}\n\
2261 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2262 # TYPE ferrox_scheduler_aborted_total counter\n\
2263 ferrox_scheduler_aborted_total {}\n",
2264 sched.prefill_chunks,
2265 sched.prefill_tokens,
2266 sched.decode_steps,
2267 sched.queue_depth,
2268 sched.queue_rejected,
2269 sched.kv_blocks_total,
2270 sched.kv_blocks_free,
2271 sched.kv_block_size,
2272 sched.kv_rejected_too_large,
2273 sched.kv_rejected_context_length,
2274 sched.aborted,
2275 )
2276 }
2277 None => body,
2278 };
2279
2280 (
2281 [(
2282 axum::http::header::CONTENT_TYPE,
2283 "text/plain; version=0.0.4",
2284 )],
2285 body,
2286 )
2287 .into_response()
2288}
2289
2290pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2291
2292pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2298 (
2299 StatusCode::BAD_REQUEST,
2300 Json(serde_json::json!({"error": {
2301 "message": message,
2302 "type": "invalid_request_error",
2303 "param": param,
2304 "code": null,
2305 }})),
2306 )
2307}
2308
2309pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2310 (
2311 StatusCode::NOT_IMPLEMENTED,
2312 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2313 )
2314}
2315
2316pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2317 let status = match e {
2318 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2319 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2324 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2330 StatusCode::SERVICE_UNAVAILABLE
2331 }
2332 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2336 };
2337 tracing::warn!("decode error: {e}");
2338 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2339 if let generate::DecodeError::KvBudgetExceeded {
2345 binding,
2346 estimated_bytes,
2347 limit_bytes,
2348 positions,
2349 positions_limit,
2350 ..
2351 } = &e
2352 {
2353 body["error"]["type"] = serde_json::json!("invalid_request_error");
2354 body["error"]["code"] = serde_json::json!(binding);
2355 body["error"]["binding"] = serde_json::json!(binding);
2356 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2357 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2358 body["error"]["positions"] = serde_json::json!(positions);
2359 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2360 }
2361 if let Some(secs) = e.retry_after_secs() {
2365 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2366 }
2367 (status, Json(body))
2368}
2369
2370pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2371 tracing::error!("generation task panicked: {e}");
2372 (
2373 StatusCode::INTERNAL_SERVER_ERROR,
2374 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2375 )
2376}
2377
2378#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2388 model: &Model,
2389 prompt: &str,
2390 params: &GenerationParams,
2391 kv_pool: Option<&generate::KvPoolConfig>,
2392 paged_kv: Option<&generate::PagedKvConfig>,
2393 prefix_cache: Option<&Mutex<PrefixCache>>,
2394 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2395 ceiling: Option<&budget::ContextCeiling>,
2396 mut emit: impl FnMut(&str),
2397) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2398 let synthetic = model.is_synthetic();
2399 let mut chunks = Vec::new();
2400 let params = &{
2406 let mut resolved = params.clone();
2407 resolved.stop_token_ids =
2408 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
2409 resolved
2410 };
2411 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2412 let (finish, usage) = match model {
2413 Model::Gguf(m) => {
2414 if let Some(batcher) = continuous_batcher {
2415 let mut tokens = m.tokenizer.encode(prompt);
2416 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2417 let (finish, _generated_ids, text, usage) =
2418 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?;
2419 if !text.is_empty() {
2420 chunks.push(text);
2421 }
2422 (finish, usage)
2423 } else {
2424 generate::generate(
2425 &m.decoder,
2426 m.tokenizer.as_ref(),
2427 &m.stop_tokens,
2428 m.bos_id,
2429 prompt,
2430 params,
2431 kv_pool,
2432 paged_kv,
2433 prefix_cache,
2434 ceiling,
2435 |chunk| {
2436 chunks.push(chunk.to_string());
2437 if !synthetic {
2438 emit(chunk);
2439 }
2440 },
2441 )?
2442 }
2443 }
2444 Model::Kimi(m) => generate::generate_engine(
2445 &m.engine,
2446 &m.tokenizer,
2447 &m.stop_tokens,
2448 None,
2449 prompt,
2450 params,
2451 |chunk| {
2452 chunks.push(chunk.to_string());
2453 if !synthetic {
2454 emit(chunk);
2455 }
2456 },
2457 )?,
2458 Model::Mla(m) => generate::generate_engine(
2459 &m.engine,
2460 &m.tokenizer,
2461 &m.stop_tokens,
2462 m.bos_id,
2463 prompt,
2464 params,
2465 |chunk| {
2466 chunks.push(chunk.to_string());
2467 if !synthetic {
2468 emit(chunk);
2469 }
2470 },
2471 )?,
2472 Model::Gemma4(m) => generate::generate_engine(
2473 &m.engine,
2474 &m.tokenizer,
2475 &m.stop_tokens,
2476 m.bos_id,
2477 prompt,
2478 params,
2479 |chunk| {
2480 chunks.push(chunk.to_string());
2481 if !synthetic {
2482 emit(chunk);
2483 }
2484 },
2485 )?,
2486 Model::Glm52(m) => generate::generate_engine(
2487 &m.engine,
2488 &m.tokenizer,
2489 &m.stop_tokens,
2490 m.bos_id,
2491 prompt,
2492 params,
2493 |chunk| {
2494 chunks.push(chunk.to_string());
2495 if !synthetic {
2496 emit(chunk);
2497 }
2498 },
2499 )?,
2500 };
2501
2502 let mut full = chunks.concat();
2503 if synthetic {
2504 full = format!(
2505 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2506 to serve a real model. Decoded ids -> {full:?}]"
2507 );
2508 emit(&full);
2509 } else if used_batcher && !full.is_empty() {
2510 emit(&full);
2511 }
2512
2513 Ok((finish, usage, full))
2514}
2515
2516#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2521 model: &Model,
2522 prompt: &str,
2523 params: &GenerationParams,
2524 kv_pool: Option<&generate::KvPoolConfig>,
2525 paged_kv: Option<&generate::PagedKvConfig>,
2526 prefix_cache: Option<&Mutex<PrefixCache>>,
2527 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2528 ceiling: Option<&budget::ContextCeiling>,
2529) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2530 let (finish, usage, full) = run_generation_emit(
2531 model,
2532 prompt,
2533 params,
2534 kv_pool,
2535 paged_kv,
2536 prefix_cache,
2537 continuous_batcher,
2538 ceiling,
2539 |_| {},
2540 )?;
2541 Ok((
2542 if full.is_empty() {
2543 Vec::new()
2544 } else {
2545 vec![full]
2546 },
2547 finish,
2548 usage,
2549 ))
2550}
2551
2552pub(crate) fn prompt_from_messages(
2563 messages: &[ChatMessage],
2564 template: &chat_template::PromptTemplate,
2565 tools: &[ToolDef],
2566 extra: serde_json::Map<String, serde_json::Value>,
2567) -> Result<String, ApiError> {
2568 let rendered = if tools.is_empty() || template.handles_tools() {
2569 template.render(messages, tools, extra)
2570 } else {
2571 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2572 with_preamble.push(ChatMessage {
2573 role: "system".to_string(),
2574 content: Some(MessageContent::Text(tool_preamble(tools))),
2575 tool_calls: None,
2576 tool_call_id: None,
2577 reasoning_content: None,
2578 });
2579 with_preamble.extend_from_slice(messages);
2580 template.render(&with_preamble, &[], extra)
2581 };
2582 rendered.map_err(template_error_response)
2583}
2584
2585fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2590 (
2591 StatusCode::BAD_REQUEST,
2592 Json(serde_json::json!({
2593 "error": {
2594 "message": format!("chat template failed to render: {err}"),
2595 "type": "invalid_request_error",
2596 "param": "messages",
2597 "code": null,
2598 }
2599 })),
2600 )
2601}
2602
2603fn tool_preamble(tools: &[ToolDef]) -> String {
2616 let mut out = String::from(
2617 "You can call tools to help answer the user. To call a tool, respond with \
2618 EXACTLY one line in this format and nothing else:\n\
2619 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2620 object matching that tool's parameters>}}</tool_call>\n\n\
2621 Available tools:\n",
2622 );
2623 for t in tools {
2624 out.push_str(&format!(
2625 "- {}: {}\n parameters (JSON schema): {}\n",
2626 t.function.name,
2627 t.function.description.as_deref().unwrap_or(""),
2628 t.function
2629 .parameters
2630 .as_ref()
2631 .map(|v| v.to_string())
2632 .unwrap_or_else(|| "{}".to_string()),
2633 ));
2634 }
2635 out
2636}
2637
2638fn tool_call_deltas(
2648 events: Vec<crate::policy::parser::ToolCallEvent>,
2649 opened: &std::cell::Cell<usize>,
2650) -> (String, Vec<ToolCallDelta>) {
2651 let mut text = String::new();
2652 let mut deltas = Vec::new();
2653 for event in events {
2654 match event {
2655 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2656 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2657 opened.set(opened.get().max(index + 1));
2658 deltas.push(ToolCallDelta::opening(index, name));
2659 }
2660 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2661 if !fragment.is_empty() {
2662 deltas.push(ToolCallDelta::arguments(index, fragment));
2663 }
2664 }
2665 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2666 }
2667 }
2668 (text, deltas)
2669}
2670
2671fn build_response_message(
2682 text: String,
2683 tools: &[ToolDef],
2684 posture: output::OutputPosture,
2685 base_finish: &'static str,
2686) -> (ChatCompletionResponseMessage, &'static str) {
2687 let parsed = output::parse_output(&text, tools, posture);
2688 let calls: Vec<ToolCallOut> = parsed
2689 .calls
2690 .into_iter()
2691 .enumerate()
2692 .map(|(index, call)| ToolCallOut {
2693 id: format!("call_{index}"),
2694 kind: "function",
2695 function: ToolCallFunctionOut {
2696 name: call.name,
2697 arguments: call.arguments,
2698 },
2699 })
2700 .collect();
2701 if !calls.is_empty() {
2702 return (
2703 ChatCompletionResponseMessage {
2704 role: "assistant",
2705 content: None,
2706 reasoning_content: parsed.reasoning,
2707 tool_calls: Some(calls),
2708 },
2709 "tool_calls",
2710 );
2711 }
2712 (
2713 ChatCompletionResponseMessage {
2714 role: "assistant",
2715 content: Some(parsed.content),
2716 reasoning_content: parsed.reasoning,
2717 tool_calls: None,
2718 },
2719 base_finish,
2720 )
2721}
2722
2723fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2728 let mut history = match &req.session_id {
2729 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2730 None => req.messages.clone(),
2731 };
2732 if req.json_object_mode() {
2733 inject_json_object_system_hint(&mut history);
2734 }
2735 history
2736}
2737
2738fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2739 const HINT: &str =
2740 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2741 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2742 match &mut sys.content {
2743 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2744 s.push_str("\n\n");
2745 s.push_str(HINT);
2746 }
2747 None => {
2748 sys.content = Some(MessageContent::Text(HINT.to_string()));
2749 }
2750 _ => {}
2751 }
2752 } else {
2753 messages.insert(
2754 0,
2755 ChatMessage {
2756 role: "system".to_string(),
2757 content: Some(MessageContent::Text(HINT.to_string())),
2758 tool_calls: None,
2759 tool_call_id: None,
2760 reasoning_content: None,
2761 },
2762 );
2763 }
2764}
2765
2766async fn chat_completions(
2767 State(state): State<Arc<AppState>>,
2768 headers: axum::http::HeaderMap,
2769 Json(req): Json<ChatCompletionRequest>,
2770) -> Response {
2771 let attribution = attribution::Attribution::from_headers(&headers);
2772 state
2773 .requests_total
2774 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2775 let started = std::time::Instant::now();
2776
2777 let request_id = ferrox_api::next_request_id();
2781 let stream = req.stream.unwrap_or(false);
2782
2783 let refusal = cache_admin::check_admission(&state)
2789 .err()
2790 .or_else(|| req.validate_supported_fields().err());
2791 if let Some(err) = refusal {
2792 state
2793 .request_errors_total
2794 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2795 let response = err.into_response();
2796 state.record_request(stats::Record {
2797 request_id: &request_id,
2798 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2799 model: state.active_model_name(),
2800 status: response.status().as_u16(),
2801 stream,
2802 duration_ms: started.elapsed().as_millis() as u64,
2803 usage: None,
2804 attribution: &attribution,
2805 });
2806 return response;
2807 }
2808
2809 let response = if stream {
2810 chat_completions_stream(
2811 Arc::clone(&state),
2812 req,
2813 request_id.clone(),
2814 started,
2815 attribution.clone(),
2816 )
2817 .await
2818 .into_response()
2819 } else {
2820 chat_completions_full(
2821 Arc::clone(&state),
2822 req,
2823 request_id.clone(),
2824 started,
2825 attribution.clone(),
2826 )
2827 .await
2828 .into_response()
2829 };
2830
2831 if response.status().is_client_error() || response.status().is_server_error() {
2832 state
2833 .request_errors_total
2834 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2835 state.record_request(stats::Record {
2839 request_id: &request_id,
2840 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2841 model: state.active_model_name(),
2844 status: response.status().as_u16(),
2845 stream,
2846 duration_ms: started.elapsed().as_millis() as u64,
2847 usage: None,
2848 attribution: &attribution,
2849 });
2850 }
2851 state.mark_request_finished();
2852
2853 response
2854}
2855
2856async fn chat_completions_full(
2857 state: Arc<AppState>,
2858 req: ChatCompletionRequest,
2859 request_id: String,
2860 started: std::time::Instant,
2861 attribution: attribution::Attribution,
2862) -> Result<Json<ChatCompletionResponse>, ApiError> {
2863 let tools_active = req.tools_active();
2864 let active = state.require_active()?;
2868 let history = resolve_history(&state, &req);
2869 let template = active.generative()?.chat_template();
2870 let kwargs = req.resolve_template_kwargs(&template);
2871 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2872 let key = req.is_cacheable().then(|| req.cache_key(&prompt));
2873
2874 let (completion, cache_status) = if let Some(cached) = key
2875 .as_ref()
2876 .and_then(|key| lock_cache(&state.response_cache).get(key))
2877 {
2878 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2879 (cached, "hit")
2880 } else {
2881 let params = req.generation_params_for_template(&template, active.name())?;
2882 let (chunks, finish, usage) = decode_task::buffered(
2883 decode_task::DecodeHandles::take(&state, &active)?,
2884 prompt.clone(),
2885 params,
2886 )
2887 .await?;
2888
2889 let completion = response_cache::CachedCompletion {
2890 content: chunks.concat(),
2891 finish,
2892 usage,
2893 };
2894 let cache_status = if let Some(key) = key {
2895 tracing::debug!("cache miss for key {}", key.digest());
2896 lock_cache(&state.response_cache).put(key, completion.clone());
2897 "miss"
2898 } else {
2899 "skip"
2900 };
2901 (completion, cache_status)
2902 };
2903 let content = completion.content;
2904
2905 if req.json_object_mode() {
2906 json_mode::validate_json_object_output(&content)?;
2907 }
2908
2909 if let Some(id) = &req.session_id {
2913 state.sessions.store_reply(
2914 id,
2915 ChatMessage {
2916 role: "assistant".to_string(),
2917 content: Some(MessageContent::Text(content.clone())),
2918 tool_calls: None,
2919 tool_call_id: None,
2920 reasoning_content: None,
2921 },
2922 );
2923 }
2924
2925 let (message, finish_reason) = build_response_message(
2926 content,
2927 if tools_active { &req.tools } else { &[] },
2928 output::OutputPosture::resolve(active.name(), &prompt),
2929 completion.finish.as_str(),
2930 );
2931
2932 state.record_request(stats::Record {
2933 request_id: &request_id,
2934 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2935 model: Some(active.name().to_string()),
2938 status: 200,
2939 stream: false,
2940 duration_ms: started.elapsed().as_millis() as u64,
2941 usage: Some(&completion.usage),
2942 attribution: &attribution,
2943 });
2944
2945 Ok(Json(ChatCompletionResponse {
2946 id: request_id.clone(),
2947 request_id,
2948 object: "chat.completion",
2949 model: req.model,
2950 choices: vec![ChatCompletionChoice {
2951 index: 0,
2952 message,
2953 finish_reason,
2954 }],
2955 usage: completion.usage,
2956 ferrox_cache: cache_status,
2957 }))
2958}
2959
2960async fn chat_completions_stream(
2961 state: Arc<AppState>,
2962 req: ChatCompletionRequest,
2963 request_id: String,
2964 started: std::time::Instant,
2965 attribution: attribution::Attribution,
2966) -> Result<Response, ApiError> {
2967 let tools_active = req.tools_active();
2969 let active = state.require_active()?;
2973 let history = resolve_history(&state, &req);
2974 let template = active.generative()?.chat_template();
2975 let kwargs = req.resolve_template_kwargs(&template);
2976 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2977 let model_name = req.model.clone();
2978 let session_id = req.session_id.clone();
2979 let sessions = state.sessions.clone();
2980
2981 let model = Arc::clone(active.generative()?);
2982 let kv_pool = state.kv_pool.clone();
2983 let paged_kv = state.paged_kv.clone();
2984 let prefix_cache = state.prefix_cache.clone();
2985 let batcher = active.batcher.clone();
2986 let ceiling = active.ceiling.clone();
2987 let mut params = req.generation_params_for_template(&template, active.name())?;
2988 let stats_state = Arc::clone(&state);
2989 let served_model = active.name().to_string();
2992 let posture = output::OutputPosture::resolve(&served_model, &prompt);
2996 let offered_tools: Vec<ToolDef> = if tools_active {
2999 req.tools.clone()
3000 } else {
3001 Vec::new()
3002 };
3003
3004 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3009 params.cancel = Some(cancel_token.clone());
3010
3011 let overlap = batcher.is_none();
3020
3021 let slot = req
3025 .stream_resumable
3026 .unwrap_or(false)
3027 .then(|| state.streams.register(&request_id));
3028 let emitter = resume::Emitter::new(slot);
3029
3030 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3031 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3035 id: request_id.clone(),
3036 request_id: None,
3037 object: "chat.completion.chunk",
3038 model: model_name.clone(),
3039 choices: vec![ChatCompletionChunkChoice {
3040 index: 0,
3041 delta: ChatCompletionChunkDelta {
3042 role: None,
3043 content: None,
3044 reasoning_content: None,
3045 tool_calls: None,
3046 },
3047 finish_reason: None,
3048 }],
3049 usage: None,
3050 });
3051
3052 tokio::task::spawn_blocking(move || {
3053 let _cancel_guard = cancel_guard;
3056 let tx_chunks = tx.clone();
3057 let orphan_timeout = sse::orphan_timeout_from_env();
3062 let mut first = true;
3063 let head_request_id = request_id.clone();
3064 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3074 Rc::new(RefCell::new(posture.reasoning_parser()));
3075 let emit_reasoning = Rc::clone(&stream_reasoning);
3076 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3080 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3081 );
3082 let emit_tools = Rc::clone(&stream_tools);
3083 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3087 let emit_streamed_calls = Rc::clone(&streamed_calls);
3088 let result = run_generation_emit(
3089 &model,
3090 &prompt,
3091 ¶ms,
3092 kv_pool.as_ref(),
3093 paged_kv.as_ref(),
3094 prefix_cache.as_deref(),
3095 batcher.as_ref(),
3096 ceiling.as_deref(),
3097 |chunk| {
3098 if !overlap || chunk.is_empty() {
3099 return;
3100 }
3101 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3102 Some(parser) => {
3103 let delta = parser.push(chunk);
3104 (delta.reasoning, delta.content)
3105 }
3106 None => (String::new(), chunk.to_string()),
3107 };
3108 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3112 Some(parser) => {
3113 let (text, calls) =
3114 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3115 (text, calls)
3116 }
3117 None => (content, Vec::new()),
3118 };
3119 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3122 return;
3123 }
3124 let role = if first { Some("assistant") } else { None };
3125 let request_id = first.then(|| head_request_id.clone());
3126 first = false;
3127 let payload = ChatCompletionChunk {
3128 id: head_request_id.clone(),
3129 request_id,
3130 object: "chat.completion.chunk",
3131 model: model_name.clone(),
3132 choices: vec![ChatCompletionChunkChoice {
3133 index: 0,
3134 delta: ChatCompletionChunkDelta {
3135 role,
3136 content: (!content.is_empty()).then_some(content),
3137 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3138 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3139 },
3140 finish_reason: None,
3141 }],
3142 usage: None,
3143 };
3144 if let Err(why) =
3153 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3154 {
3155 if why == sse::SendFailure::Orphaned {
3156 tracing::warn!(
3157 "SSE stream {head_request_id} accepted nothing for the orphan \
3158 deadline; treating it as abandoned"
3159 );
3160 }
3161 if !emitter.is_resumable() {
3173 cancel_token.cancel();
3174 }
3175 }
3176 },
3177 );
3178
3179 let mut pending_request_id = first.then(|| request_id.clone());
3185
3186 match result {
3187 Ok((finish, usage, full_text)) => {
3188 if let Some(id) = &session_id {
3189 sessions.store_reply(
3190 id,
3191 ChatMessage {
3192 role: "assistant".to_string(),
3193 content: Some(MessageContent::Text(full_text.clone())),
3194 tool_calls: None,
3195 tool_call_id: None,
3196 reasoning_content: None,
3197 },
3198 );
3199 }
3200 let mut streamed_finish: Option<&'static str> = None;
3206 if overlap {
3207 let tail = stream_reasoning
3208 .borrow_mut()
3209 .as_mut()
3210 .map(|parser| parser.flush())
3211 .unwrap_or_default();
3212 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3213 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3214 let mut events = parser.push(&content);
3215 events.extend(parser.finish());
3216 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3217 content = text;
3218 tool_calls = calls;
3219 }
3220 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3221 let payload = ChatCompletionChunk {
3222 id: request_id.clone(),
3223 request_id: pending_request_id.take(),
3224 object: "chat.completion.chunk",
3225 model: model_name.clone(),
3226 choices: vec![ChatCompletionChunkChoice {
3227 index: 0,
3228 delta: ChatCompletionChunkDelta {
3229 role: None,
3230 content: (!content.is_empty()).then_some(content),
3231 reasoning_content: (!tail.reasoning.is_empty())
3232 .then_some(tail.reasoning),
3233 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3234 },
3235 finish_reason: None,
3236 }],
3237 usage: None,
3238 };
3239 let _ =
3240 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3241 }
3242 if streamed_calls.get() > 0 {
3243 streamed_finish = Some("tool_calls");
3244 }
3245 } else {
3246 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3249 let tool_calls: Vec<ToolCallDelta> = parsed
3250 .calls
3251 .iter()
3252 .enumerate()
3253 .map(|(index, call)| {
3254 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3255 })
3256 .collect();
3257 if !tool_calls.is_empty() {
3258 streamed_finish = Some("tool_calls");
3259 }
3260 if !tool_calls.is_empty()
3261 || !parsed.content.is_empty()
3262 || parsed.reasoning.is_some()
3263 {
3264 let payload = ChatCompletionChunk {
3265 id: request_id.clone(),
3266 request_id: pending_request_id.take(),
3267 object: "chat.completion.chunk",
3268 model: model_name.clone(),
3269 choices: vec![ChatCompletionChunkChoice {
3270 index: 0,
3271 delta: ChatCompletionChunkDelta {
3272 role: Some("assistant"),
3273 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3274 .then(|| parsed.content.clone()),
3275 reasoning_content: parsed.reasoning.clone(),
3276 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3277 },
3278 finish_reason: None,
3279 }],
3280 usage: None,
3281 };
3282 let _ =
3283 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3284 }
3285 }
3286 let final_finish_reason = match streamed_finish {
3290 Some(reason) if finish.as_str() != "length" => reason,
3291 _ => finish.as_str(),
3292 };
3293 let final_payload = ChatCompletionChunk {
3294 id: request_id.clone(),
3295 request_id: pending_request_id.take(),
3296 object: "chat.completion.chunk",
3297 model: model_name,
3298 choices: vec![ChatCompletionChunkChoice {
3299 index: 0,
3300 delta: ChatCompletionChunkDelta {
3301 role: None,
3302 content: None,
3303 reasoning_content: None,
3304 tool_calls: None,
3305 },
3306 finish_reason: Some(final_finish_reason),
3307 }],
3308 usage: Some(usage.clone()),
3309 };
3310 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3311 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3312 stats_state.record_request(stats::Record {
3317 request_id: &request_id,
3318 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3319 model: Some(served_model.clone()),
3320 status: 200,
3321 stream: true,
3322 duration_ms: started.elapsed().as_millis() as u64,
3323 usage: Some(&usage),
3324 attribution: &attribution,
3325 });
3326 }
3327 Err(e) => {
3328 tracing::warn!("decode error on streamed request {request_id}: {e}");
3329 stats_state.record_request(stats::Record {
3335 request_id: &request_id,
3336 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3337 model: Some(served_model.clone()),
3338 status: 500,
3339 stream: true,
3340 duration_ms: started.elapsed().as_millis() as u64,
3341 usage: None,
3342 attribution: &attribution,
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,
3349 choices: vec![ChatCompletionChunkChoice {
3350 index: 0,
3351 delta: ChatCompletionChunkDelta {
3352 role: Some("assistant"),
3353 content: Some(format!("[error: {e}]")),
3354 reasoning_content: None,
3355 tool_calls: None,
3356 },
3357 finish_reason: Some("stop"),
3358 }],
3359 usage: None,
3360 };
3361 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3362 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3363 }
3364 }
3365 drop(emitter);
3369 });
3370
3371 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3372 Ok((
3395 [(
3396 axum::http::HeaderName::from_static("x-accel-buffering"),
3397 axum::http::HeaderValue::from_static("no"),
3398 )],
3399 Sse::new(stream),
3400 )
3401 .into_response())
3402}
3403
3404fn axum_path(template: &str) -> String {
3421 let mut out = String::with_capacity(template.len());
3422 let mut rest = template;
3423 while let Some(open) = rest.find('{') {
3424 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3425 break;
3426 };
3427 out.push_str(&rest[..open]);
3428 out.push(':');
3429 out.push_str(&rest[open + 1..close]);
3430 rest = &rest[close + 1..];
3431 }
3432 out.push_str(rest);
3433 out
3434}
3435
3436async fn cancel_generation(
3445 State(state): State<Arc<AppState>>,
3446 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3447) -> Response {
3448 let cancelled = state.cancels.cancel(&req.request_id);
3449 let status = if cancelled {
3450 StatusCode::OK
3451 } else {
3452 StatusCode::NOT_FOUND
3453 };
3454 let detail = if cancelled {
3455 "the generation was asked to stop; it ends at its next token".to_string()
3456 } else {
3457 "no generation with that request_id is running -- it has already \
3458 finished, was never issued, or was served by a path that does \
3459 not register for cancellation"
3460 .to_string()
3461 };
3462 (
3463 status,
3464 Json(ferrox_api::CancelGenerationResponse {
3465 request_id: req.request_id,
3466 cancelled,
3467 detail,
3468 }),
3469 )
3470 .into_response()
3471}
3472
3473type Activated = (
3477 Loaded,
3478 Option<serving::batch::ContinuousBatcher>,
3479 Option<Arc<budget::ContextCeiling>>,
3480);
3481
3482fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3500 let mut batcher = serving::batch::BatcherConfig::from_env();
3501 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3502 return batcher;
3505 }
3506 let Some(path) = path else {
3507 return batcher;
3508 };
3509 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
3515 let Some((priced, gguf_ctx, source)) = priced else {
3516 return batcher;
3517 };
3518 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3519 tracing::warn!(
3523 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3524 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3525 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3526 admit on a number you choose.",
3527 priced.weights_bytes,
3528 priced.device_budget_bytes,
3529 );
3530 return batcher;
3531 };
3532 tracing::info!("{source}");
3533 tracing::info!("{}", derived.fit);
3534 let adopted = budget::apply_derived(&mut batcher, &derived);
3535 if adopted.max_context {
3536 tracing::info!(
3537 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3538 override with FERROX_CB_MAX_CONTEXT",
3539 derived.max_context
3540 );
3541 }
3542 if adopted.kv_blocks {
3543 tracing::info!(
3544 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3545 derived.kv_blocks,
3546 batcher.kv_block_size
3547 );
3548 }
3549 batcher
3550}
3551
3552pub(crate) fn activate_loaded_model(
3561 loaded: model::LoadedModel,
3562 enable_continuous_batching: bool,
3563 path: Option<&str>,
3564 paged_kv: Option<&generate::PagedKvConfig>,
3565) -> Activated {
3566 match loaded {
3567 model::LoadedModel::Gguf(g) => {
3568 let decoder = Arc::new(g.decoder);
3569 let tokenizer = Arc::new(g.tokenizer);
3570 let config = price_batcher_config(path);
3571 let shape =
3577 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
3578 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3579 let batcher = if enable_continuous_batching {
3580 tracing::info!(
3581 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3582 (stop sequences use the same pending-buffer trim as the private generate loop)"
3583 );
3584 let tok = Arc::clone(&tokenizer);
3585 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
3586 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3587 Arc::clone(&decoder),
3588 decode,
3589 config,
3590 Arc::clone(&ceiling),
3591 paged_kv.cloned(),
3592 ))
3593 } else {
3594 None
3595 };
3596 (
3597 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3598 decoder,
3599 tokenizer,
3600 stop_tokens: g.stop_tokens,
3601 bos_id: g.bos_id,
3602 is_synthetic: g.is_synthetic,
3603 chat_template: g.chat_template,
3604 }))),
3605 batcher,
3606 Some(ceiling),
3607 )
3608 }
3609 model::LoadedModel::Kimi(k) => (
3610 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3611 engine: k.engine,
3612 tokenizer: k.tokenizer,
3613 stop_tokens: k.stop_tokens,
3614 chat_template: k.chat_template,
3615 }))),
3616 None,
3617 None,
3618 ),
3619 model::LoadedModel::Mla(m) => (
3620 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3621 engine: m.engine,
3622 tokenizer: m.tokenizer,
3623 stop_tokens: m.stop_tokens,
3624 bos_id: m.bos_id,
3625 name: m.name,
3626 chat_template: m.chat_template,
3627 }))),
3628 None,
3629 None,
3630 ),
3631 model::LoadedModel::Gemma4(m) => (
3632 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3633 engine: m.engine,
3634 tokenizer: m.tokenizer,
3635 stop_tokens: m.stop_tokens,
3636 bos_id: m.bos_id,
3637 name: m.name,
3638 chat_template: m.chat_template,
3639 }))),
3640 None,
3641 None,
3642 ),
3643 model::LoadedModel::Glm52(g) => (
3644 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3645 engine: g.engine,
3646 tokenizer: g.tokenizer,
3647 stop_tokens: g.stop_tokens,
3648 bos_id: g.bos_id,
3649 name: g.name,
3650 chat_template: g.chat_template,
3651 }))),
3652 None,
3653 None,
3654 ),
3655 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3660 }
3661}
3662
3663struct StartupModels {
3670 loaded: model::LoadedModel,
3671 embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
3672}
3673
3674fn build_app_state(
3675 models: StartupModels,
3676 kv_pool: Option<generate::KvPoolConfig>,
3677 paged_kv: Option<generate::PagedKvConfig>,
3678 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3679 enable_continuous_batching: bool,
3680 mcp: Option<mcp::LoadedMcpConfig>,
3681 detection: Arc<health::Detection>,
3682) -> AppState {
3683 let StartupModels { loaded, embedding } = models;
3684 let (loaded, batcher, ceiling) = activate_loaded_model(
3685 loaded,
3686 enable_continuous_batching,
3687 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
3688 paged_kv.as_ref(),
3689 );
3690 let id = startup_model_id();
3696 AppState {
3697 embedding,
3698 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3699 id,
3700 loaded,
3701 batcher,
3702 ceiling,
3703 }))),
3704 paged_kv,
3705 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3706 tasks: Arc::new(tasks::TaskRegistry::new()),
3707 cancels: Arc::new(cancel::CancelRegistry::new()),
3708 stats: stats::Stats::new(),
3709 streams: resume::StreamRegistry::new(),
3710 model_dir: admin::model_dirs().into_iter().next(),
3711 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3712 kv_pool,
3713 prefix_cache,
3714 sessions: session::SessionStore::new(),
3715 requests_total: std::sync::atomic::AtomicU64::new(0),
3716 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3717 started_at: std::time::Instant::now(),
3718 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3719 detection,
3720 mcp,
3721 continuous_batching_enabled: enable_continuous_batching,
3722 loading_model: Mutex::new(None),
3723 last_load_error: Mutex::new(None),
3724 serving: Mutex::new(crate::stats::ServingStats::default()),
3725 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3726 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3727 started_unix: unix_now(),
3728 }
3729}
3730
3731fn load_embedding_model() -> anyhow::Result<Option<Arc<ferrox_models::EmbeddingModel>>> {
3739 let Ok(path) = std::env::var("FERROX_EMBEDDING_MODEL_PATH") else {
3740 return Ok(None);
3741 };
3742 let model = ferrox_models::EmbeddingModel::from_gguf_path(&path)
3743 .map_err(|e| anyhow::anyhow!("FERROX_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3744 tracing::info!(
3745 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3746 model.name(),
3747 model.architecture(),
3748 model.n_embd(),
3749 model.pooling_type().name(),
3750 model.n_ctx_train(),
3751 );
3752 Ok(Some(Arc::new(model)))
3753}
3754
3755fn unix_now() -> u64 {
3760 std::time::SystemTime::now()
3761 .duration_since(std::time::UNIX_EPOCH)
3762 .map(|d| d.as_secs())
3763 .unwrap_or(0)
3764}
3765
3766fn startup_model_id() -> Option<String> {
3771 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
3772 let configured = std::fs::canonicalize(&configured).ok()?;
3773 admin::discover(&admin::model_dirs())
3774 .into_iter()
3775 .find(|d| {
3776 std::fs::canonicalize(&d.path)
3777 .map(|p| p == configured)
3778 .unwrap_or(false)
3779 })
3780 .map(|d| d.id)
3781}
3782
3783fn init_cpu_pool() {
3792 match ferrox_core::threads::init_cpu_pool() {
3793 Some(n) => eprintln!(
3794 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
3795 ferrox_core::threads::perf_core_count()
3796 ),
3797 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
3798 }
3799}
3800
3801fn announce_ready(addr: SocketAddr, scheme: &str) {
3814 use std::io::Write;
3815 let ready =
3816 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3817 let mut stdout = std::io::stdout().lock();
3818 let _ = writeln!(stdout, "{}", ready.to_line());
3819 let _ = stdout.flush();
3820}
3821
3822async fn shutdown_signal(exit_on_stdin_close: bool) {
3834 if !exit_on_stdin_close {
3835 std::future::pending::<()>().await;
3836 return;
3837 }
3838 let _ = tokio::task::spawn_blocking(|| {
3839 use std::io::Read;
3840 let mut sink = [0u8; 256];
3841 let mut stdin = std::io::stdin().lock();
3842 loop {
3843 match stdin.read(&mut sink) {
3844 Ok(0) => break,
3846 Ok(_) => continue,
3848 Err(e) => {
3849 tracing::warn!("stdin read failed ({e}); treating it as closed");
3850 break;
3851 }
3852 }
3853 }
3854 })
3855 .await;
3856 tracing::info!("stdin closed; shutting down");
3857}
3858
3859fn tokio_worker_threads() -> usize {
3864 std::env::var("FERROX_TOKIO_WORKERS")
3865 .ok()
3866 .and_then(|v| v.trim().parse::<usize>().ok())
3867 .filter(|n| *n > 0)
3868 .unwrap_or(2)
3869}
3870
3871fn install_ring_crypto_provider() {
3888 let _ = rustls::crypto::ring::default_provider().install_default();
3889}
3890
3891pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3897 if args.list_devices {
3898 print_available_devices();
3899 return Ok(());
3900 }
3901 apply_cli_overrides(&args)?;
3902
3903 let _instance = {
3907 use ferrox_core::instance::{register, InstancePolicy};
3908 let policy = if args.allow_multiple_instances {
3909 InstancePolicy::Multi
3910 } else {
3911 InstancePolicy::from_env_or(InstancePolicy::Single)
3912 };
3913 let model = std::env::var("FERROX_MODEL_PATH").ok();
3914 register(
3915 "server",
3916 model.as_deref(),
3917 ferrox_core::instance::current_backend(),
3918 policy,
3919 )
3920 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
3921 };
3922
3923 let journal = journal::Journal::from_env();
3924 eprintln!(
3925 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
3926 journal.path()
3927 );
3928 journal.append(&journal::Record::session_start(
3929 env!("CARGO_PKG_VERSION"),
3930 std::process::id(),
3931 ));
3932 journal::install_panic_hook(journal.clone());
3933
3934 let mcp_config_path = args.mcp_config.clone();
3935 let exit_on_stdin_close = args.exit_on_stdin_close
3936 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
3937 .map(|v| v == "1")
3938 .unwrap_or(false);
3939
3940 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
3944 init_cpu_pool();
3945
3946 let runtime = tokio::runtime::Builder::new_multi_thread()
3947 .worker_threads(tokio_worker_threads())
3948 .enable_all()
3949 .build()?;
3950 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
3951
3952 let reason = match &result {
3953 Ok(()) => "normal".to_string(),
3954 Err(e) => e.to_string(),
3955 };
3956 journal.append(&journal::Record::session_exit(reason));
3957
3958 runtime.shutdown_background();
3963
3964 result
3965}
3966
3967async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
3968 let _ = tracing_subscriber::fmt::try_init();
3974
3975 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
3985 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
3986 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
3987 .map(|v| v == "1")
3988 .unwrap_or(false);
3989 if let Err(msg) =
3990 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
3991 {
3992 anyhow::bail!(msg);
3993 }
3994
3995 let embedding_model = load_embedding_model()?;
4002
4003 let mut loaded = model::load()?;
4004 match &loaded {
4005 model::LoadedModel::Gguf(g) => tracing::info!(
4006 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4007 g.decoder.config.name,
4008 g.is_synthetic,
4009 g.tokenizer.kind()
4010 ),
4011 model::LoadedModel::Kimi(k) => tracing::info!(
4012 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4013 k.tokenizer.vocab_size()
4014 ),
4015 model::LoadedModel::Mla(m) => tracing::info!(
4016 "loaded MLA GGUF '{}' (tokenizer={})",
4017 m.name,
4018 m.tokenizer.kind()
4019 ),
4020 model::LoadedModel::Gemma4(m) => tracing::info!(
4021 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4022 m.name,
4023 m.tokenizer.kind()
4024 ),
4025 model::LoadedModel::Glm52(g) => tracing::info!(
4026 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4027 g.name,
4028 g.tokenizer.kind()
4029 ),
4030 model::LoadedModel::Encoder(_) => {}
4033 }
4034 let metal_default_moe_budget = {
4042 #[cfg(feature = "metal")]
4043 {
4044 ferrox_core::metal_dense_enabled()
4045 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
4046 }
4047 #[cfg(not(feature = "metal"))]
4048 {
4049 false
4050 }
4051 };
4052 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
4053 let budget: u64 = budget_str
4054 .parse()
4055 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4056 match &mut loaded {
4057 model::LoadedModel::Gguf(g) => {
4058 tracing::info!(
4059 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4060 (CUDA and/or Metal matvecs when built with the matching feature)"
4061 );
4062 g.decoder.gpu_vram_budget_bytes = Some(budget);
4063 }
4064 model::LoadedModel::Kimi(_) => {
4065 tracing::warn!(
4066 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4067 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4068 );
4069 }
4070 model::LoadedModel::Mla(_) => {
4071 tracing::warn!(
4072 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4073 FFN path only today; ignoring expert VRAM budget"
4074 );
4075 }
4076 model::LoadedModel::Gemma4(_) => {
4077 tracing::warn!(
4078 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4079 ignoring expert VRAM budget"
4080 );
4081 }
4082 model::LoadedModel::Glm52(_) => {
4083 tracing::warn!(
4084 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4085 GPU expert placement not wired yet; ignoring"
4086 );
4087 }
4088 model::LoadedModel::Encoder(_) => {
4089 tracing::warn!(
4090 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4091 it has no routed experts to place; ignoring"
4092 );
4093 }
4094 }
4095 } else if metal_default_moe_budget {
4096 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4100 if let model::LoadedModel::Gguf(g) = &mut loaded {
4101 tracing::info!(
4102 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4103 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4104 );
4105 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4106 }
4107 }
4108 #[cfg(feature = "cuda")]
4109 {
4110 if ferrox_core::cuda_dense_enabled() {
4111 tracing::info!(
4112 "CUDA dense matvec enabled for WeightMatrix::apply \
4113 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4114 );
4115 } else {
4116 tracing::info!(
4117 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
4118 );
4119 }
4120 }
4121 #[cfg(feature = "metal")]
4122 {
4123 if ferrox_core::metal_dense_enabled() {
4124 tracing::info!(
4125 "Metal dense matvec enabled for WeightMatrix::apply \
4126 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4127 );
4128 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
4129 Some("1") | Some("true") | Some("on") | Some("attn") => {
4130 tracing::info!(
4131 "Metal fused attention requested (FERROX_METAL_ATTN): \
4132 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4133 );
4134 }
4135 _ => {}
4136 }
4137 tracing::info!(
4138 "Metal greedy GPU argmax: temperature<=0 folds \
4139 final_norm+lm_head+argmax into the dense stack"
4140 );
4141 } else {
4142 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
4143 }
4144 }
4145 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
4156 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
4157 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
4158 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4159 panic!(
4160 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
4161 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
4162 );
4163 }
4164 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4165 (Ok(blocks), Ok(block_size), Err(_)) => {
4166 let total_blocks: usize = blocks
4167 .parse()
4168 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
4169 let block_size: usize = block_size
4170 .parse()
4171 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4172 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4178 .ok()
4179 .map(|v| {
4180 v.parse()
4181 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4182 })
4183 .unwrap_or(0);
4184 tracing::info!(
4185 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4186 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4187 );
4188 Some(generate::KvPoolConfig {
4189 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4190 queue_wait: Duration::from_millis(queue_wait_ms),
4191 })
4192 }
4193 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4194 let block_size: usize = block_size
4195 .parse()
4196 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4197 let budget: u64 = byte_budget
4198 .parse()
4199 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
4200 let cfg = match &loaded {
4201 model::LoadedModel::Gguf(g) => &g.decoder.config,
4202 model::LoadedModel::Kimi(_)
4203 | model::LoadedModel::Mla(_)
4204 | model::LoadedModel::Gemma4(_)
4205 | model::LoadedModel::Glm52(_)
4206 | model::LoadedModel::Encoder(_) => {
4207 panic!(
4208 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
4209 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4210 );
4211 }
4212 };
4213 let bytes_per_block = block_size
4214 * cfg.n_layers
4215 * cfg.n_kv_heads
4216 * cfg.head_dim
4217 * 2
4218 * std::mem::size_of::<f32>();
4219 assert!(
4220 bytes_per_block > 0,
4221 "derived KV block byte size must be positive (check model config and block size)"
4222 );
4223 let total_blocks = (budget as usize / bytes_per_block).max(1);
4224 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4225 .ok()
4226 .map(|v| {
4227 v.parse()
4228 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4229 })
4230 .unwrap_or(0);
4231 tracing::info!(
4232 "KV cache block pool enabled from byte budget: {budget} bytes / \
4233 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4234 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4235 cfg.n_layers
4236 );
4237 Some(generate::KvPoolConfig {
4238 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4239 queue_wait: Duration::from_millis(queue_wait_ms),
4240 })
4241 }
4242 (Err(_), Err(_), Err(_)) => None,
4243 (Err(_), Ok(_), Err(_)) => panic!(
4244 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4245 (or unset all three to disable KV cache pooling)"
4246 ),
4247 (Ok(_), Ok(_), Ok(_)) => {
4248 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4249 }
4250 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4251 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4252 set together (or neither, to disable KV cache pooling)"
4253 ),
4254 };
4255 let paged_kv = match (
4261 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4262 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4263 ) {
4264 (Ok(blocks), Ok(block_size)) => {
4265 assert!(
4266 kv_pool.is_none(),
4267 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4268 mutually exclusive: both bound the same KV memory, by different means. \
4269 Set one."
4270 );
4271 let blocks_per_layer: usize = blocks
4285 .parse()
4286 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4287 let block_size: usize = block_size
4288 .parse()
4289 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4290 let gguf = match &loaded {
4291 model::LoadedModel::Gguf(g) => g,
4292 _ => panic!(
4293 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4294 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4295 ),
4296 };
4297 let cfg = &gguf.decoder.config;
4298 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4299 .ok()
4300 .map(|v| {
4301 v.parse()
4302 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4303 })
4304 .unwrap_or(0);
4305 tracing::info!(
4306 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4307 layer across {} layers, shared by all concurrent requests, \
4308 {queue_wait_ms}ms admission queue wait",
4309 cfg.n_layers
4310 );
4311 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4316 block_size,
4317 ))));
4318 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4327 crate::policy::parser::ToolCallFormat::infer(
4328 &std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
4329 )
4330 .opener(),
4331 |text| {
4332 gguf.tokenizer
4333 .encode(text)
4334 .into_iter()
4335 .map(|t| t as u32)
4336 .collect()
4337 },
4338 );
4339 if let Some(id) = anchor_token {
4340 tracing::info!(
4341 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4342 window stops short of where its next turn rejoins"
4343 );
4344 }
4345 let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
4346 .ok()
4347 .map(|v| {
4348 v.parse()
4349 .expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4350 })
4351 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4352 if let Some(window) = cfg.uniform_sliding_window() {
4353 tracing::info!(
4354 "Paged KV window slide enabled: every layer slides by {window} every \
4355 {slide_interval} decode steps, so a request holds its prompt and a \
4356 window rather than its whole context"
4357 );
4358 } else if cfg.kv_block_window().is_some() {
4359 tracing::info!(
4360 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4361 and a page group holds one block in every layer"
4362 );
4363 }
4364 Some(generate::PagedKvConfig {
4365 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4366 cfg.n_layers,
4367 block_size,
4368 blocks_per_layer,
4369 cfg.n_kv_heads,
4370 cfg.head_dim,
4371 )),
4372 queue_wait: Duration::from_millis(queue_wait_ms),
4373 radix,
4374 anchor_token,
4375 slide_interval,
4376 })
4377 }
4378 (Err(_), Err(_)) => None,
4379 _ => panic!(
4380 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4381 (or neither, to disable paged KV)"
4382 ),
4383 };
4384 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4391 let max_entries: usize = v
4392 .parse()
4393 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4394 if kv_pool.is_some() {
4395 tracing::warn!(
4396 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4397 caching will never be consulted while a KV pool is configured"
4398 );
4399 }
4400 assert!(
4407 paged_kv.is_none(),
4408 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4409 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4410 produce, so the cache could never hit. Set one."
4411 );
4412 tracing::info!(
4413 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4414 all requests"
4415 );
4416 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4417 });
4418 if matches!(
4419 loaded,
4420 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4421 ) && (kv_pool.is_some() || prefix_cache.is_some())
4422 {
4423 tracing::warn!(
4424 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4425 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4426 ferrox_models::engine's module docs"
4427 );
4428 }
4429 let enable_cb = std::env::var("FERROX_CONTINUOUS_BATCHING")
4430 .map(|v| v == "1")
4431 .unwrap_or(false)
4432 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
4440 && matches!(loaded, model::LoadedModel::Gguf(_));
4441 if std::env::var("FERROX_CONTINUOUS_BATCHING")
4442 .map(|v| v == "1")
4443 .unwrap_or(false)
4444 && (kv_pool.is_some() || prefix_cache.is_some())
4445 {
4446 tracing::warn!(
4447 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4448 (those modes keep the private generate path)"
4449 );
4450 }
4451 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4452 if let Ok(chunk) = n.parse::<usize>() {
4453 if chunk > 0 {
4454 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4455 }
4456 }
4457 }
4458 if matches!(
4459 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4460 Some("1")
4461 ) {
4462 tracing::warn!(
4463 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4464 (minimal spill; full layer offload still planned)"
4465 );
4466 }
4467
4468 let mcp = match mcp_config_path {
4469 Some(path) => {
4470 let loaded = mcp::load_mcp_config(&path)?;
4471 tracing::info!(
4472 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4473 loaded.path,
4474 loaded.servers.len()
4475 );
4476 Some(loaded)
4477 }
4478 None => None,
4479 };
4480
4481 let detection = health::Detection::spawn();
4485
4486 let state = Arc::new(build_app_state(
4487 StartupModels {
4488 loaded,
4489 embedding: embedding_model,
4490 },
4491 kv_pool,
4492 paged_kv,
4493 prefix_cache,
4494 enable_cb,
4495 mcp,
4496 detection,
4497 ));
4498
4499 use ferrox_api::routes;
4503
4504 let public = Router::new().route(routes::HEALTH, get(health));
4509
4510 let mut protected = Router::new()
4511 .route(routes::V1_MODELS, get(list_models))
4512 .route(routes::V1_RESPONSES, post(responses::responses))
4516 .route(
4517 &axum_path(routes::V1_RESPONSE),
4518 get(responses::responses_get),
4519 )
4520 .route(
4521 &axum_path(routes::V1_RESPONSE_CANCEL),
4522 post(responses::responses_cancel),
4523 )
4524 .route(routes::V1_STATS, get(serving_stats))
4525 .route(routes::V1_REQUESTS, get(recent_requests))
4526 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
4527 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
4528 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
4529 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
4530 .route(routes::V1_CANCEL, post(cancel_generation))
4534 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
4539 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
4540 .route(routes::V1_MESSAGES, post(anthropic::messages))
4541 .route(
4542 routes::V1_MESSAGES_COUNT_TOKENS,
4543 post(anthropic::count_tokens),
4544 )
4545 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
4546 .route(routes::COMPLETION, post(completion::completion))
4551 .route(routes::COMPLETIONS, post(completion::completion))
4552 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
4553 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
4554 .route(routes::TOKENIZE, post(openai_extra::tokenize))
4560 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
4561 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
4562 .route(routes::CACHE_STATS, get(cache_stats))
4563 .route(routes::METRICS, get(metrics))
4564 .route(routes::ADMIN_MODELS, get(admin::models))
4569 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
4570 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
4571 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
4572 .route(routes::ADMIN_TASKS, get(admin::tasks))
4573 .route(&admin::cancel_route(), post(admin::cancel_task))
4574 .route(routes::ADMIN_STATS, get(admin::stats))
4575 .merge(conversations::router());
4579
4580 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4585 tracing::info!("API key auth enabled");
4586 let auth = limits::AuthConfig {
4587 api_key: Arc::new(key),
4588 };
4589 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4590 auth,
4591 limits::require_api_key,
4592 ));
4593 }
4594 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4595 let rpm: u32 = rpm
4596 .parse()
4597 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4598 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4599 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4600 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4601 limiter,
4602 limits::rate_limit,
4603 ));
4604 }
4605 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
4613 let origins = security::parse_cors_origins(&spec)
4614 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
4615 tracing::info!(
4616 "CORS enabled: {} allow-listed origin(s) ({})",
4617 origins.len(),
4618 spec
4619 );
4620 let cors = tower_http::cors::CorsLayer::new()
4621 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4622 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4623 .allow_headers([
4624 axum::http::header::CONTENT_TYPE,
4625 axum::http::header::AUTHORIZATION,
4626 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4632 axum::http::HeaderName::from_static("last-event-id"),
4636 ]);
4637 protected = protected.route_layer(cors);
4638 }
4639
4640 let app = public
4644 .merge(protected)
4645 .layer(axum::middleware::from_fn(limits::retry_after))
4646 .with_state(state);
4647
4648 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4654 install_ring_crypto_provider();
4655 match tls_paths {
4660 Some(paths) => {
4661 let config =
4662 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4663 .await
4664 .map_err(|e| {
4665 anyhow::anyhow!(
4666 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4667 paths.cert,
4668 paths.key
4669 )
4670 })?;
4671 let socket_addr: std::net::SocketAddr = addr
4672 .parse()
4673 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
4674 let listener = std::net::TcpListener::bind(socket_addr)?;
4675 listener.set_nonblocking(true)?;
4683 let bound = listener.local_addr()?;
4684 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
4685 announce_ready(bound, "https");
4686
4687 let handle = axum_server::Handle::new();
4688 let shutdown_handle = handle.clone();
4689 tokio::spawn(async move {
4690 shutdown_signal(exit_on_stdin_close).await;
4691 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4692 });
4693 axum_server::from_tcp_rustls(listener, config)?
4694 .handle(handle)
4695 .serve(app.into_make_service())
4696 .await?;
4697 }
4698 None => {
4699 let listener = tokio::net::TcpListener::bind(&addr).await?;
4700 let bound = listener.local_addr()?;
4701 tracing::info!("ferrox-server listening on {bound}");
4702 announce_ready(bound, "http");
4703 axum::serve(listener, app)
4704 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4705 .await?;
4706 }
4707 }
4708 Ok(())
4709}
4710
4711#[cfg(test)]
4712mod tests {
4713 use super::*;
4714 use ferrox_models::config::test_dense_fixture;
4715
4716 #[test]
4717 fn parses_llama_server_style_options() {
4718 let argv = [
4719 "ferrox-server",
4720 "-m",
4721 "model.gguf",
4722 "--host",
4723 "::1",
4724 "--port",
4725 "9000",
4726 "-t",
4727 "4",
4728 "-dev",
4729 "Metal",
4730 "-ngl",
4731 "all",
4732 ]
4733 .into_iter()
4734 .map(String::from)
4735 .collect();
4736 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4737
4738 assert_eq!(args.model.as_deref(), Some("model.gguf"));
4739 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
4740 assert_eq!(args.port, Some(9000));
4741 assert_eq!(args.threads, Some(4));
4742 assert_eq!(args.device, Some(OffloadDevice::Metal));
4743 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
4744 assert_eq!(
4745 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4746 Some("[::1]:9000")
4747 );
4748 }
4749
4750 #[test]
4751 fn port_zero_survives_argument_parsing_as_a_real_request() {
4752 let argv = ["ferrox-server", "--port", "0"]
4757 .into_iter()
4758 .map(String::from)
4759 .collect();
4760 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4761 assert_eq!(args.port, Some(0));
4762 assert_eq!(
4763 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4764 Some("127.0.0.1:0")
4765 );
4766 }
4767
4768 #[test]
4769 fn stdin_close_exit_is_opt_in() {
4770 let args =
4773 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
4774 assert!(!args.exit_on_stdin_close);
4775 let args = ServerArgs::try_parse_from(
4776 ["ferrox-server", "--exit-on-stdin-close"]
4777 .into_iter()
4778 .map(String::from),
4779 )
4780 .unwrap();
4781 assert!(args.exit_on_stdin_close);
4782 }
4783
4784 #[test]
4785 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4786 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4787 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4788 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
4789 assert_eq!(parsed.port, 51999);
4790 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4791 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
4793 }
4794
4795 fn test_model() -> Model {
4796 let cfg = test_dense_fixture();
4800 Model::Gguf(GgufModel {
4801 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4802 tokenizer: Arc::new(ServerTokenizer::Byte),
4803 stop_tokens: StopTokens::default(),
4804 bos_id: None,
4805 is_synthetic: true,
4806 chat_template: chat_template::PromptTemplate::plain(),
4807 })
4808 }
4809
4810 fn greedy_params(max_tokens: usize) -> GenerationParams {
4811 GenerationParams {
4812 max_tokens,
4813 sampling: SamplingParams::default(),
4814 seed: 1,
4815 stop: Vec::new(),
4816 stop_token_ids: Vec::new(),
4817 json_object: false,
4818 grammar: None,
4819 cancel: None,
4820 ignore_eos: false,
4821 }
4822 }
4823
4824 fn test_model_full_byte_vocab() -> Model {
4828 let mut cfg = test_dense_fixture();
4829 cfg.vocab_size = 256;
4830 Model::Gguf(GgufModel {
4831 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4832 tokenizer: Arc::new(ServerTokenizer::Byte),
4833 stop_tokens: StopTokens::default(),
4834 bos_id: None,
4835 is_synthetic: true,
4836 chat_template: chat_template::PromptTemplate::plain(),
4837 })
4838 }
4839
4840 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4844 AppState {
4845 embedding: None,
4846 paged_kv: None,
4847 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4848 id: None,
4849 loaded: Loaded::Generative(Arc::new(model)),
4850 batcher: None,
4851 ceiling: None,
4852 }))),
4853 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4854 tasks: Arc::new(tasks::TaskRegistry::new()),
4855 cancels: Arc::new(cancel::CancelRegistry::new()),
4856 stats: stats::Stats::new(),
4857 streams: resume::StreamRegistry::new(),
4858 model_dir: None,
4859 response_cache: Mutex::new(response_cache),
4860 kv_pool: None,
4861 prefix_cache: None,
4862 sessions: session::SessionStore::new(),
4863 requests_total: std::sync::atomic::AtomicU64::new(0),
4864 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4865 started_at: std::time::Instant::now(),
4866 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4867 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4868 mcp: None,
4869 continuous_batching_enabled: false,
4870 loading_model: Mutex::new(None),
4871 last_load_error: Mutex::new(None),
4872 serving: Mutex::new(crate::stats::ServingStats::default()),
4873 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4874 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4875 started_unix: unix_now(),
4876 }
4877 }
4878
4879 fn test_app() -> Router {
4888 test_app_with_state(Arc::new(test_state(
4889 test_model_full_byte_vocab(),
4890 ResponseCache::new(1000, Duration::from_secs(3600)),
4891 )))
4892 }
4893
4894 fn test_app_with_state(state: Arc<AppState>) -> Router {
4897 Router::new()
4898 .route(ferrox_api::routes::HEALTH, get(health))
4899 .route(ferrox_api::routes::V1_MODELS, get(list_models))
4900 .route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
4901 .route(
4902 &axum_path(ferrox_api::routes::V1_RESPONSE),
4903 get(responses::responses_get),
4904 )
4905 .route(
4906 &axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
4907 post(responses::responses_cancel),
4908 )
4909 .route(ferrox_api::routes::V1_STATS, get(serving_stats))
4910 .route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
4911 .route(
4912 ferrox_api::routes::V1_CACHE_STATUS,
4913 get(cache_admin::cache_status),
4914 )
4915 .route(
4916 ferrox_api::routes::V1_CACHE_REBUILD,
4917 post(cache_admin::cache_rebuild),
4918 )
4919 .route(
4920 ferrox_api::routes::ADMIN_PREPARE_STOP,
4921 post(cache_admin::prepare_stop),
4922 )
4923 .route("/v1/chat/completions", post(chat_completions))
4924 .route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
4925 .route(
4926 ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
4927 post(anthropic::count_tokens),
4928 )
4929 .route("/v1/tokenize", post(openai_extra::tokenize))
4930 .route("/v1/detokenize", post(openai_extra::detokenize))
4931 .route(ferrox_api::routes::TOKENIZE, post(openai_extra::tokenize))
4935 .route(
4936 ferrox_api::routes::DETOKENIZE,
4937 post(openai_extra::detokenize),
4938 )
4939 .route("/v1/embeddings", post(embeddings::embeddings))
4940 .route("/v1/completions", post(openai_extra::completions))
4941 .route(ferrox_api::routes::COMPLETION, post(completion::completion))
4943 .route(
4944 ferrox_api::routes::COMPLETIONS,
4945 post(completion::completion),
4946 )
4947 .route(
4948 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
4949 post(admin::unload_model),
4950 )
4951 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
4952 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
4953 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
4954 .route(
4955 &axum_path(ferrox_api::routes::V1_STREAM),
4956 get(resume::resume),
4957 )
4958 .route(
4959 &axum_path(ferrox_api::routes::V1_STREAM_POLL),
4960 get(resume::poll),
4961 )
4962 .with_state(state)
4963 }
4964
4965 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4966 let mut cfg = test_dense_fixture();
4967 cfg.name = name;
4968 cfg.vocab_size = vocab_size;
4969 Model::Gguf(GgufModel {
4970 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4971 tokenizer: Arc::new(ServerTokenizer::Byte),
4972 stop_tokens: StopTokens::default(),
4973 bos_id: None,
4974 is_synthetic: true,
4975 chat_template: chat_template::PromptTemplate::plain(),
4976 })
4977 }
4978
4979 fn model_with_template(name: &'static str, source: &str) -> Model {
4983 let mut cfg = test_dense_fixture();
4984 cfg.name = name;
4985 cfg.vocab_size = 256;
4986 Model::Gguf(GgufModel {
4987 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4988 tokenizer: Arc::new(ServerTokenizer::Byte),
4989 stop_tokens: StopTokens::default(),
4990 bos_id: None,
4991 is_synthetic: true,
4992 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4993 Some(source),
4994 Some("qwen3"),
4995 false,
4996 true,
4997 None,
4998 None,
4999 ),
5000 })
5001 }
5002
5003 #[tokio::test]
5012 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5013 let strict = "{% if messages | length > 1 %}\
5016 {{ raise_exception('this template takes one turn') }}\
5017 {% endif %}{{ messages[0].content }}";
5018 let state = Arc::new(test_state(
5019 model_with_template("strict", strict),
5020 ResponseCache::new(4, Duration::from_secs(60)),
5021 ));
5022 let app = test_app_with_state(state);
5023
5024 let (status, body) = post_json_uri(
5025 &app,
5026 "/v1/chat/completions",
5027 serde_json::json!({
5028 "model": "strict",
5029 "stream": true,
5030 "messages": [
5031 {"role": "user", "content": "one"},
5032 {"role": "user", "content": "two"},
5033 ],
5034 }),
5035 )
5036 .await;
5037 assert_eq!(status, StatusCode::BAD_REQUEST);
5038 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5039 assert!(
5040 body["error"]["message"]
5041 .as_str()
5042 .unwrap()
5043 .contains("one turn"),
5044 "the template's own message must reach the caller: {body}"
5045 );
5046
5047 let (status, _) = post_json_uri(
5049 &app,
5050 "/v1/chat/completions",
5051 serde_json::json!({
5052 "model": "strict",
5053 "stream": true,
5054 "max_tokens": 1,
5055 "messages": [{"role": "user", "content": "one"}],
5056 }),
5057 )
5058 .await;
5059 assert_eq!(status, StatusCode::OK);
5060 }
5061
5062 #[tokio::test]
5064 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5065 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5066 {% if reasoning_effort %}\
5067 {% if reasoning_effort not in ['low','medium','high'] %}\
5068 {{ raise_exception('bad effort') }}\
5069 {% endif %}[{{ reasoning_effort }}]\
5070 {% endif %}{{ messages[0].content }}";
5071 let state = Arc::new(test_state(
5072 model_with_template("thinker", reasoning),
5073 ResponseCache::new(4, Duration::from_secs(60)),
5074 ));
5075 let app = test_app_with_state(state);
5076 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5077 assert_eq!(status, StatusCode::OK);
5078 let entry = &models["data"][0];
5079 assert_eq!(
5080 entry["supported_reasoning_efforts"],
5081 serde_json::json!(["off", "low", "medium", "high"])
5082 );
5083 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5084 }
5085
5086 #[tokio::test]
5091 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5092 let app = test_app();
5093 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5094 let entry = &models["data"][0];
5095 assert!(entry.get("supported_reasoning_efforts").is_none());
5096 assert!(entry.get("default_reasoning_effort").is_none());
5097 }
5098
5099 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5100 Arc::new(ActiveModel {
5101 id: Some(name.to_string()),
5102 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5103 batcher: None,
5104 ceiling: None,
5105 })
5106 .tap_into(state)
5107 }
5108
5109 trait TapInto {
5111 fn tap_into(self, state: &AppState) -> Self;
5112 }
5113 impl TapInto for Arc<ActiveModel> {
5114 fn tap_into(self, state: &AppState) -> Self {
5115 state.swap_active(Some(Arc::clone(&self)));
5116 self
5117 }
5118 }
5119
5120 #[test]
5126 fn an_in_flight_request_keeps_the_model_it_started_on() {
5127 let state = test_state(
5128 named_test_model("model-a", 256),
5129 ResponseCache::new(4, Duration::from_secs(60)),
5130 );
5131
5132 let in_flight = state.active().expect("a model is loaded");
5135 assert_eq!(in_flight.name(), "model-a");
5136
5137 active_model(&state, "model-b");
5138
5139 assert_eq!(state.active().unwrap().name(), "model-b");
5141 assert_eq!(in_flight.name(), "model-a");
5143 let (_chunks, finish, _usage) = run_generation(
5144 in_flight.generative().unwrap(),
5145 "hi",
5146 &greedy_params(3),
5147 None,
5148 None,
5149 None,
5150 None,
5151 None,
5152 )
5153 .expect("the old model must still decode after being swapped out");
5154 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5155 }
5156
5157 #[test]
5162 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5163 let state = test_state(
5164 named_test_model("model-a", 256),
5165 ResponseCache::new(4, Duration::from_secs(60)),
5166 );
5167 let in_flight = state.active().expect("a model is loaded");
5168 let weights = Arc::clone(in_flight.generative().unwrap());
5169 assert!(Arc::strong_count(&weights) >= 2);
5170
5171 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5172 id: Some("model-b".to_string()),
5173 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5174 batcher: None,
5175 ceiling: None,
5176 })));
5177 drop(previous);
5178 assert!(Arc::strong_count(&weights) >= 2);
5180 drop(in_flight);
5181 assert_eq!(Arc::strong_count(&weights), 1);
5182 }
5183
5184 #[tokio::test]
5188 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5189 let state = Arc::new(test_state(
5190 named_test_model("model-a", 256),
5191 ResponseCache::new(4, Duration::from_secs(60)),
5192 ));
5193 let app = test_app_with_state(Arc::clone(&state));
5194
5195 let (status, body) = post_json_uri(
5196 &app,
5197 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5198 serde_json::json!({}),
5199 )
5200 .await;
5201 assert_eq!(status, StatusCode::OK);
5202 assert_eq!(body["ok"], true);
5203 assert!(body["active"].is_null());
5204 assert!(state.active().is_none());
5205
5206 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5207 assert_eq!(status, StatusCode::OK);
5208 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5209 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5210
5211 let (status, body) = post_json_uri(
5212 &app,
5213 "/v1/chat/completions",
5214 serde_json::json!({
5215 "model": "x",
5216 "messages": [{"role": "user", "content": "hi"}]
5217 }),
5218 )
5219 .await;
5220 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5221 assert_eq!(body["error"]["type"], "model_not_loaded");
5222 }
5223
5224 #[tokio::test]
5228 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5229 let state = Arc::new(test_state(
5230 named_test_model("model-a", 256),
5231 ResponseCache::new(4, Duration::from_secs(60)),
5232 ));
5233 let app = test_app_with_state(Arc::clone(&state));
5234 state.swap_active(None);
5235
5236 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5237 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5240 assert_eq!(body["state"], "unavailable");
5241 assert_eq!(body["reason"], "model_not_loaded");
5242 assert!(body["model"].is_null());
5243 let real_weights = body["capabilities"]
5244 .as_array()
5245 .unwrap()
5246 .iter()
5247 .find(|c| c["id"] == "real_weights")
5248 .cloned()
5249 .expect("real_weights is always reported");
5250 assert_eq!(real_weights["available"], false);
5251 assert_eq!(real_weights["reason"], "model_not_loaded");
5252 }
5253
5254 #[tokio::test]
5258 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5259 let app = test_app();
5260
5261 let (status, completion) = post_json_uri(
5262 &app,
5263 "/v1/chat/completions",
5264 serde_json::json!({
5265 "model": "x",
5266 "messages": [{"role": "user", "content": "hi"}],
5267 "max_tokens": 4
5268 }),
5269 )
5270 .await;
5271 assert_eq!(status, StatusCode::OK);
5272 let request_id = completion["request_id"].as_str().unwrap().to_string();
5273
5274 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5275 assert_eq!(status, StatusCode::OK);
5276 let recent = stats["recent"].as_array().unwrap();
5277 assert_eq!(recent.len(), 1);
5278 let row = &recent[0];
5279 assert_eq!(row["request_id"], request_id);
5280 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
5281 assert_eq!(row["status"], 200);
5282 assert_eq!(row["stream"], false);
5283 assert!(row["duration_ms"].is_number());
5286 assert!(row["decode_ms"].is_number());
5287 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5288 assert_eq!(
5289 stats["tokens_prompt_total"].as_u64().unwrap(),
5290 row["prompt_tokens"].as_u64().unwrap()
5291 );
5292 }
5293
5294 #[tokio::test]
5298 async fn a_rejected_request_is_recorded_too() {
5299 let state = Arc::new(test_state(
5300 named_test_model("model-a", 256),
5301 ResponseCache::new(4, Duration::from_secs(60)),
5302 ));
5303 let app = test_app_with_state(Arc::clone(&state));
5304 state.swap_active(None);
5305
5306 let (status, _) = post_json_uri(
5307 &app,
5308 "/v1/chat/completions",
5309 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5310 )
5311 .await;
5312 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5313
5314 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5315 let recent = stats["recent"].as_array().unwrap();
5316 assert_eq!(recent.len(), 1);
5317 assert_eq!(recent[0]["status"], 503);
5318 assert_eq!(recent[0]["completion_tokens"], 0);
5319 assert!(recent[0]["decode_ms"].is_null());
5320 assert_eq!(stats["errors_total"], 1);
5321 }
5322
5323 async fn post_json_with_headers(
5327 app: &Router,
5328 uri: &str,
5329 body: serde_json::Value,
5330 headers: &[(&str, &str)],
5331 ) -> (StatusCode, serde_json::Value) {
5332 use http_body_util::BodyExt;
5333 use tower::ServiceExt;
5334
5335 let mut builder = axum::http::Request::builder()
5336 .method("POST")
5337 .uri(uri)
5338 .header("content-type", "application/json");
5339 for (name, value) in headers {
5340 builder = builder.header(*name, *value);
5341 }
5342 let response = app
5343 .clone()
5344 .oneshot(
5345 builder
5346 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5347 .unwrap(),
5348 )
5349 .await
5350 .unwrap();
5351 let status = response.status();
5352 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5353 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5354 (status, json)
5355 }
5356
5357 #[tokio::test]
5361 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5362 let app = test_app();
5363
5364 let (status, _) = post_json_uri(
5365 &app,
5366 ferrox_api::routes::V1_TOKENIZE,
5367 serde_json::json!({"prompt": "hello"}),
5368 )
5369 .await;
5370 assert_eq!(status, StatusCode::OK);
5371 let (status, _) = post_json_uri(
5372 &app,
5373 ferrox_api::routes::V1_DETOKENIZE,
5374 serde_json::json!({"tokens": [104, 105]}),
5375 )
5376 .await;
5377 assert_eq!(status, StatusCode::OK);
5378 let (status, _) = post_json_uri(
5379 &app,
5380 ferrox_api::routes::V1_EMBEDDINGS,
5381 serde_json::json!({"input": "hello"}),
5382 )
5383 .await;
5384 assert_eq!(status, StatusCode::OK);
5385
5386 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5387 let routes: Vec<&str> = stats["recent"]
5388 .as_array()
5389 .unwrap()
5390 .iter()
5391 .map(|row| row["route"].as_str().unwrap())
5392 .collect();
5393 for expected in [
5394 ferrox_api::routes::V1_TOKENIZE,
5395 ferrox_api::routes::V1_DETOKENIZE,
5396 ferrox_api::routes::V1_EMBEDDINGS,
5397 ] {
5398 assert!(
5399 routes.contains(&expected),
5400 "{expected} is missing: {routes:?}"
5401 );
5402 }
5403
5404 let row = |route: &str| {
5405 stats["recent"]
5406 .as_array()
5407 .unwrap()
5408 .iter()
5409 .find(|r| r["route"] == route)
5410 .cloned()
5411 .unwrap()
5412 };
5413 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5417 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5418 assert!(embed["decode_ms"].is_null());
5419 assert_eq!(embed["completion_tokens"], 0);
5420 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5424 assert_eq!(
5425 stats["tokens_prompt_total"].as_u64().unwrap(),
5426 embed["prompt_tokens"].as_u64().unwrap(),
5427 "only the forward pass counted"
5428 );
5429 }
5430
5431 fn streaming_test_app() -> Router {
5436 let mut cfg = test_dense_fixture();
5437 cfg.vocab_size = 256;
5438 let model = Model::Gguf(GgufModel {
5439 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5440 tokenizer: Arc::new(ServerTokenizer::Byte),
5441 stop_tokens: StopTokens::default(),
5442 bos_id: None,
5443 is_synthetic: false,
5444 chat_template: chat_template::PromptTemplate::plain(),
5445 });
5446 test_app_with_state(Arc::new(test_state(
5447 model,
5448 ResponseCache::new(1000, Duration::from_secs(3600)),
5449 )))
5450 }
5451
5452 #[tokio::test]
5456 async fn the_native_completion_wire_is_not_the_openai_one() {
5457 let app = test_app();
5458
5459 let (status, native) = post_json_uri(
5460 &app,
5461 ferrox_api::routes::COMPLETION,
5462 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5463 )
5464 .await;
5465 assert_eq!(status, StatusCode::OK, "{native}");
5466 assert!(native["content"].is_string(), "{native}");
5467 assert_eq!(native["stop"], true);
5468 assert_eq!(native["stop_type"], "limit");
5469 assert_eq!(native["stopping_word"], "");
5470 assert_eq!(native["truncated"], false);
5471 assert_eq!(native["id_slot"], -1);
5472 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5473 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5474 assert!(
5475 native.get("choices").is_none(),
5476 "the native shape has no `choices`: {native}"
5477 );
5478
5479 let (status, openai) = post_json_uri(
5480 &app,
5481 ferrox_api::routes::V1_COMPLETIONS,
5482 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5483 )
5484 .await;
5485 assert_eq!(status, StatusCode::OK);
5486 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5487 assert!(
5488 openai.get("content").is_none(),
5489 "the OpenAI shape has no top-level `content`: {openai}"
5490 );
5491 }
5492
5493 #[tokio::test]
5497 async fn both_native_spellings_reach_the_same_handler() {
5498 let app = test_app();
5499 for route in [
5500 ferrox_api::routes::COMPLETION,
5501 ferrox_api::routes::COMPLETIONS,
5502 ] {
5503 let (status, body) = post_json_uri(
5504 &app,
5505 route,
5506 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5507 )
5508 .await;
5509 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5510 assert_eq!(body["stop"], true, "{route}");
5511 assert!(body["content"].is_string(), "{route}");
5512 }
5513
5514 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5517 let routes: Vec<&str> = stats["recent"]
5518 .as_array()
5519 .unwrap()
5520 .iter()
5521 .map(|row| row["route"].as_str().unwrap())
5522 .collect();
5523 assert!(
5524 routes.contains(&ferrox_api::routes::COMPLETION),
5525 "{routes:?}"
5526 );
5527 assert!(
5528 routes.contains(&ferrox_api::routes::COMPLETIONS),
5529 "{routes:?}"
5530 );
5531 }
5532
5533 #[tokio::test]
5539 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5540 let app = streaming_test_app();
5541 let raw = post_sse_raw_uri(
5542 &app,
5543 ferrox_api::routes::COMPLETION,
5544 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5545 )
5546 .await;
5547
5548 assert!(
5549 !raw.contains("[DONE]"),
5550 "llama.cpp's native stream has no sentinel: {raw}"
5551 );
5552 let frames: Vec<serde_json::Value> = raw
5553 .lines()
5554 .filter_map(|line| line.strip_prefix("data: "))
5555 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5556 .collect();
5557 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5558
5559 let (last, partials) = frames.split_last().unwrap();
5560 assert_eq!(last["stop"], true, "the last frame closes the stream");
5561 assert!(last["timings"].is_object(), "{last}");
5562 assert!(last["stop_type"].is_string(), "{last}");
5563 for partial in partials {
5564 assert_eq!(partial["stop"], false, "{partial}");
5565 assert!(partial["content"].is_string(), "{partial}");
5566 assert!(partial.get("timings").is_none(), "{partial}");
5570 assert!(partial.get("generation_settings").is_none(), "{partial}");
5571 }
5572 let streamed: String = partials
5575 .iter()
5576 .filter_map(|p| p["content"].as_str())
5577 .collect();
5578 assert_eq!(last["content"].as_str().unwrap(), streamed);
5579 }
5580
5581 #[tokio::test]
5586 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5587 let app = test_app();
5588 for body in [
5589 serde_json::json!({"prompt": "hi"}),
5590 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5591 ] {
5592 let (status, refusal) =
5593 post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
5594 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5595 assert!(
5596 refusal["error"]["message"]
5597 .as_str()
5598 .unwrap()
5599 .contains("n_predict"),
5600 "{refusal}"
5601 );
5602 }
5603 let (status, _) = post_json_uri(
5606 &app,
5607 ferrox_api::routes::COMPLETION,
5608 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5609 )
5610 .await;
5611 assert_eq!(status, StatusCode::OK);
5612 }
5613
5614 #[tokio::test]
5623 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5624 let app = streaming_test_app();
5625 let ask = |stop: serde_json::Value| {
5626 let app = app.clone();
5627 async move {
5628 post_json_uri(
5629 &app,
5630 ferrox_api::routes::COMPLETION,
5631 serde_json::json!({
5632 "prompt": "hi",
5633 "n_predict": 64,
5634 "ignore_eos": true,
5635 "stop": stop,
5636 }),
5637 )
5638 .await
5639 .1
5640 }
5641 };
5642
5643 let baseline = ask(serde_json::json!([])).await;
5644 assert_eq!(baseline["stop_type"], "limit");
5645 assert_eq!(baseline["stopping_word"], "");
5646 let text = baseline["content"].as_str().unwrap().to_string();
5647 let sentinel: String = text.chars().skip(1).take(2).collect();
5653 assert_eq!(
5654 sentinel.chars().count(),
5655 2,
5656 "the fixture must produce enough output to cut: {text:?}"
5657 );
5658 let cut = text.find(&sentinel).expect("it came out of this text");
5659
5660 let stopped = ask(serde_json::json!([sentinel])).await;
5661 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5662 assert_eq!(stopped["stopping_word"], sentinel);
5663 assert_eq!(
5664 stopped["content"].as_str().unwrap(),
5665 &text[..cut],
5666 "the answer must be cut at the sentinel, not run past it"
5667 );
5668 }
5669
5670 #[tokio::test]
5676 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5677 let app = test_app();
5678
5679 let (v1_status, v1) = post_json_uri(
5680 &app,
5681 ferrox_api::routes::V1_TOKENIZE,
5682 serde_json::json!({"prompt": "hello"}),
5683 )
5684 .await;
5685 let (alias_status, alias) = post_json_uri(
5686 &app,
5687 ferrox_api::routes::TOKENIZE,
5688 serde_json::json!({"content": "hello"}),
5689 )
5690 .await;
5691 assert_eq!(v1_status, StatusCode::OK);
5692 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5693 assert_eq!(v1["tokens"], alias["tokens"]);
5694 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5695
5696 let (status, both_ways) = post_json_uri(
5699 &app,
5700 ferrox_api::routes::TOKENIZE,
5701 serde_json::json!({"prompt": "hello"}),
5702 )
5703 .await;
5704 assert_eq!(status, StatusCode::OK);
5705 assert_eq!(both_ways["tokens"], v1["tokens"]);
5706 }
5707
5708 #[tokio::test]
5713 async fn detokenize_answers_under_both_dialects_keys() {
5714 let app = test_app();
5715 for route in [
5716 ferrox_api::routes::DETOKENIZE,
5717 ferrox_api::routes::V1_DETOKENIZE,
5718 ] {
5719 let (status, body) =
5720 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5721 assert_eq!(status, StatusCode::OK, "{route}");
5722 assert_eq!(body["text"], "hi", "{route}");
5723 assert_eq!(body["content"], body["text"], "{route}");
5724 }
5725 }
5726
5727 #[tokio::test]
5731 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5732 let app = test_app();
5733 let (status, _) = post_json_uri(
5734 &app,
5735 ferrox_api::routes::TOKENIZE,
5736 serde_json::json!({"content": "hello"}),
5737 )
5738 .await;
5739 assert_eq!(status, StatusCode::OK);
5740
5741 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5742 let routes: Vec<&str> = stats["recent"]
5743 .as_array()
5744 .unwrap()
5745 .iter()
5746 .map(|row| row["route"].as_str().unwrap())
5747 .collect();
5748 assert!(
5749 routes.contains(&ferrox_api::routes::TOKENIZE),
5750 "the alias must be its own row: {routes:?}"
5751 );
5752 assert!(
5753 !routes.contains(&ferrox_api::routes::V1_TOKENIZE),
5754 "nothing called /v1/tokenize: {routes:?}"
5755 );
5756 }
5757
5758 #[tokio::test]
5763 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5764 let mut cfg = test_dense_fixture();
5765 cfg.vocab_size = 256;
5766 let model = Model::Gguf(GgufModel {
5767 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5768 tokenizer: Arc::new(ServerTokenizer::Byte),
5769 stop_tokens: StopTokens::default(),
5770 bos_id: Some(7),
5771 is_synthetic: true,
5772 chat_template: chat_template::PromptTemplate::plain(),
5773 });
5774 let app = test_app_with_state(Arc::new(test_state(
5775 model,
5776 ResponseCache::new(1000, Duration::from_secs(3600)),
5777 )));
5778
5779 let (_, plain) = post_json_uri(
5780 &app,
5781 ferrox_api::routes::TOKENIZE,
5782 serde_json::json!({"content": "hi"}),
5783 )
5784 .await;
5785 let (_, special) = post_json_uri(
5786 &app,
5787 ferrox_api::routes::TOKENIZE,
5788 serde_json::json!({"content": "hi", "add_special": true}),
5789 )
5790 .await;
5791
5792 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5793 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5794 assert_eq!(special["count"], 3);
5795 }
5796
5797 #[tokio::test]
5800 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5801 let app = test_app();
5802 let (status, _) = post_json_uri(
5803 &app,
5804 ferrox_api::routes::V1_EMBEDDINGS,
5805 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5806 )
5807 .await;
5808 assert_eq!(status, StatusCode::BAD_REQUEST);
5809
5810 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5811 let recent = stats["recent"].as_array().unwrap();
5812 assert_eq!(recent.len(), 1);
5813 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
5814 assert_eq!(recent[0]["status"], 400);
5815 assert_eq!(
5816 recent[0]["prompt_tokens"], 0,
5817 "a rejected call embedded nothing"
5818 );
5819 }
5820
5821 #[tokio::test]
5824 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5825 let app = test_app();
5826 let key = "sk-monitor-secret";
5827 let (status, _) = post_json_with_headers(
5828 &app,
5829 "/v1/chat/completions",
5830 serde_json::json!({
5831 "model": "x",
5832 "messages": [{"role": "user", "content": "hi"}],
5833 "max_tokens": 2
5834 }),
5835 &[
5836 ("authorization", &format!("Bearer {key}")),
5837 ("x-ferrox-client", "ferrox-studio"),
5838 ],
5839 )
5840 .await;
5841 assert_eq!(status, StatusCode::OK);
5842
5843 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5844 let row = stats["recent"].as_array().unwrap()[0].clone();
5845 let fingerprint = row["via_api_key"]
5846 .as_str()
5847 .expect("the row names the key that served it")
5848 .to_string();
5849 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5850 assert!(!fingerprint.contains(key));
5851 assert!(
5852 !serde_json::to_string(&stats).unwrap().contains(key),
5853 "the stats payload must not carry the key in any form"
5854 );
5855 assert_eq!(row["client"], "ferrox-studio");
5856 }
5857
5858 #[tokio::test]
5861 async fn different_keys_are_different_callers_and_no_key_is_null() {
5862 let app = test_app();
5863 let body = serde_json::json!({
5864 "model": "x",
5865 "messages": [{"role": "user", "content": "hi"}],
5866 "max_tokens": 1
5867 });
5868 for headers in [
5869 vec![("authorization", "Bearer key-one")],
5870 vec![("authorization", "Bearer key-two")],
5871 vec![],
5872 ] {
5873 let (status, _) =
5874 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5875 assert_eq!(status, StatusCode::OK);
5876 }
5877
5878 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5879 let recent = stats["recent"].as_array().unwrap();
5880 assert_eq!(recent.len(), 3);
5881 let one = recent[0]["via_api_key"].as_str().unwrap();
5882 let two = recent[1]["via_api_key"].as_str().unwrap();
5883 assert_ne!(one, two, "two keys must not collapse into one caller");
5884 assert!(
5885 recent[2]["via_api_key"].is_null(),
5886 "an unauthenticated call is null, not a fingerprint of nothing"
5887 );
5888 assert!(recent[2]["client"].is_null());
5889 }
5890
5891 #[tokio::test]
5896 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5897 let state = Arc::new(test_state(
5898 named_test_model("really-loaded", 256),
5899 ResponseCache::new(4, Duration::from_secs(60)),
5900 ));
5901 let app = test_app_with_state(Arc::clone(&state));
5902
5903 let (status, _) = post_json_uri(
5904 &app,
5905 "/v1/chat/completions",
5906 serde_json::json!({
5907 "model": "gpt-4-turbo-that-is-not-here",
5908 "messages": [{"role": "user", "content": "hi"}],
5909 "max_tokens": 2
5910 }),
5911 )
5912 .await;
5913 assert_eq!(status, StatusCode::OK);
5914
5915 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5916 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5917
5918 state.swap_active(None);
5921 let (status, _) = post_json_uri(
5922 &app,
5923 "/v1/chat/completions",
5924 serde_json::json!({
5925 "model": "gpt-4-turbo-that-is-not-here",
5926 "messages": [{"role": "user", "content": "hi"}]
5927 }),
5928 )
5929 .await;
5930 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5931 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5932 let recent = stats["recent"].as_array().unwrap();
5933 assert!(recent[recent.len() - 1]["model"].is_null());
5934 }
5935
5936 #[tokio::test]
5940 async fn a_streamed_row_names_the_model_it_decoded_against() {
5941 let state = Arc::new(test_state(
5942 named_test_model("model-before", 256),
5943 ResponseCache::new(4, Duration::from_secs(60)),
5944 ));
5945 let app = test_app_with_state(Arc::clone(&state));
5946 let _ = post_sse_raw(&app, resumable_request()).await;
5947 active_model(&state, "model-after");
5949
5950 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5951 assert_eq!(stats["recent"][0]["model"], "model-before");
5952 }
5953
5954 #[tokio::test]
5957 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5958 let app = test_app();
5959 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5960 assert_eq!(status, StatusCode::OK);
5961 assert!(
5962 stats["queue_depth"].is_null(),
5963 "without continuous batching nothing queues, so there is nothing to measure"
5964 );
5965 assert!(stats["queue_rejected_total"].is_null());
5966 assert_eq!(
5967 stats["generating_now"], 0,
5968 "work in progress is measured and really is zero here"
5969 );
5970 }
5971
5972 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5977 post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
5978 }
5979
5980 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
5984 use http_body_util::BodyExt;
5985 use tower::ServiceExt;
5986
5987 let response = app
5988 .clone()
5989 .oneshot(
5990 axum::http::Request::builder()
5991 .method("POST")
5992 .uri(uri)
5993 .header("content-type", "application/json")
5994 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5995 .unwrap(),
5996 )
5997 .await
5998 .unwrap();
5999 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6000 String::from_utf8(bytes.to_vec()).unwrap()
6001 }
6002
6003 async fn get_json_with_headers(
6004 app: &Router,
6005 uri: &str,
6006 headers: &[(&str, &str)],
6007 ) -> (StatusCode, serde_json::Value) {
6008 use http_body_util::BodyExt;
6009 use tower::ServiceExt;
6010
6011 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6012 for (name, value) in headers {
6013 builder = builder.header(*name, *value);
6014 }
6015 let response = app
6016 .clone()
6017 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6018 .await
6019 .unwrap();
6020 let status = response.status();
6021 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6022 (
6023 status,
6024 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6025 )
6026 }
6027
6028 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6029 body.lines()
6030 .filter_map(|line| line.strip_prefix(field))
6031 .map(str::trim)
6032 .collect()
6033 }
6034
6035 fn resumable_request() -> serde_json::Value {
6036 serde_json::json!({
6037 "model": "m",
6038 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6039 "max_tokens": 4,
6040 "temperature": 0,
6041 "stream": true,
6042 "stream_resumable": true,
6043 })
6044 }
6045
6046 #[tokio::test]
6051 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6052 let app = test_app();
6053 let body = post_sse_raw(&app, resumable_request()).await;
6054
6055 let request_id = body
6056 .lines()
6057 .find_map(|l| l.strip_prefix("data: "))
6058 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6059 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6060 .expect("the first chunk names the request");
6061
6062 let ids = sse_field(&body, "id:");
6063 let datas = sse_field(&body, "data:");
6064 assert_eq!(
6065 ids.len(),
6066 datas.len(),
6067 "every event carries an id, or a reconnect cannot name where it stopped"
6068 );
6069 for (i, id) in ids.iter().enumerate() {
6070 assert_eq!(*id, format!("{request_id}:{i}"));
6071 }
6072 let retries = sse_field(&body, "retry:");
6073 assert_eq!(
6074 retries.len(),
6075 1,
6076 "the reconnect delay is stated once, not on every event"
6077 );
6078 assert_eq!(retries[0], "1500");
6079 assert!(
6080 body.contains("data: [DONE]"),
6081 "the end of stream is still stated"
6082 );
6083 }
6084
6085 #[tokio::test]
6089 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6090 let app = test_app();
6091 let mut request = resumable_request();
6092 request["stream_resumable"] = serde_json::json!(false);
6093 let body = post_sse_raw(&app, request).await;
6094 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6095 assert!(
6096 sse_field(&body, "id:").is_empty(),
6097 "an id promises a replay this stream cannot serve"
6098 );
6099 assert!(sse_field(&body, "retry:").is_empty());
6100 }
6101
6102 #[tokio::test]
6106 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6107 let app = test_app();
6108 let body = post_sse_raw(&app, resumable_request()).await;
6109 let request_id = sse_field(&body, "id:")[0]
6110 .rsplit_once(':')
6111 .unwrap()
6112 .0
6113 .to_string();
6114 let streamed: Vec<String> = sse_field(&body, "data:")
6115 .iter()
6116 .map(|d| d.to_string())
6117 .collect();
6118
6119 let (status, polled) = get_json(
6120 &app,
6121 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
6122 )
6123 .await;
6124 assert_eq!(status, StatusCode::OK);
6125 let events: Vec<String> = polled["events"]
6126 .as_array()
6127 .unwrap()
6128 .iter()
6129 .map(|e| e["data"].as_str().unwrap().to_string())
6130 .collect();
6131 assert_eq!(
6132 events, streamed,
6133 "the fallback must deliver the same answer, not a re-run of it"
6134 );
6135 assert_eq!(polled["request_id"], request_id);
6136 assert_eq!(
6137 polled["done"], false,
6138 "events were still being handed out, so the client must ask again"
6139 );
6140
6141 let next = polled["next_index"].as_u64().unwrap();
6144 let (_, drained) = get_json(
6145 &app,
6146 &format!(
6147 "{}?from={next}",
6148 ferrox_api::routes::v1_stream_poll(&request_id)
6149 ),
6150 )
6151 .await;
6152 assert_eq!(drained["done"], true);
6153 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6154 }
6155
6156 #[tokio::test]
6160 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6161 let app = test_app();
6162 let body = post_sse_raw(&app, resumable_request()).await;
6163 let ids = sse_field(&body, "id:");
6164 let datas: Vec<String> = sse_field(&body, "data:")
6165 .iter()
6166 .map(|d| d.to_string())
6167 .collect();
6168 assert!(
6169 ids.len() >= 3,
6170 "need a few events to resume into the middle"
6171 );
6172 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6173
6174 let (status, resumed) = get_json_with_headers(
6175 &app,
6176 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
6177 &[],
6178 )
6179 .await;
6180 assert_eq!(status, StatusCode::OK);
6181 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6182
6183 let (_, tail) = get_json(
6185 &app,
6186 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
6187 )
6188 .await;
6189 let tail_events: Vec<String> = tail["events"]
6190 .as_array()
6191 .unwrap()
6192 .iter()
6193 .map(|e| e["data"].as_str().unwrap().to_string())
6194 .collect();
6195 assert_eq!(tail_events, datas[2..].to_vec());
6196 }
6197
6198 #[tokio::test]
6201 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6202 use http_body_util::BodyExt;
6203 use tower::ServiceExt;
6204
6205 let app = test_app();
6206 let body = post_sse_raw(&app, resumable_request()).await;
6207 let ids = sse_field(&body, "id:");
6208 let datas: Vec<String> = sse_field(&body, "data:")
6209 .iter()
6210 .map(|d| d.to_string())
6211 .collect();
6212 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6213
6214 let response = app
6215 .clone()
6216 .oneshot(
6217 axum::http::Request::builder()
6218 .method("GET")
6219 .uri(ferrox_api::routes::v1_stream(&request_id))
6220 .header("last-event-id", format!("{request_id}:0"))
6221 .body(axum::body::Body::empty())
6222 .unwrap(),
6223 )
6224 .await
6225 .unwrap();
6226 assert_eq!(response.status(), StatusCode::OK);
6227 assert_eq!(
6228 response
6229 .headers()
6230 .get("x-accel-buffering")
6231 .and_then(|v| v.to_str().ok()),
6232 Some("no"),
6233 "the reconnect needs the same anti-buffering header as the stream"
6234 );
6235 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6236 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6237 assert_eq!(
6238 sse_field(&resumed, "data:")
6239 .iter()
6240 .map(|d| d.to_string())
6241 .collect::<Vec<_>>(),
6242 datas[1..].to_vec()
6243 );
6244 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6245 }
6246
6247 #[tokio::test]
6251 async fn a_last_event_id_from_another_stream_is_refused() {
6252 let app = test_app();
6253 let body = post_sse_raw(&app, resumable_request()).await;
6254 let request_id = sse_field(&body, "id:")[0]
6255 .rsplit_once(':')
6256 .unwrap()
6257 .0
6258 .to_string();
6259
6260 let (status, err) = get_json_with_headers(
6261 &app,
6262 &ferrox_api::routes::v1_stream(&request_id),
6263 &[("last-event-id", "chatcmpl-someone-else:3")],
6264 )
6265 .await;
6266 assert_eq!(status, StatusCode::BAD_REQUEST);
6267 assert_eq!(err["error"]["code"], "bad_last_event_id");
6268 }
6269
6270 #[tokio::test]
6274 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6275 let app = test_app();
6276 let mut request = resumable_request();
6277 request["stream_resumable"] = serde_json::json!(false);
6278 let body = post_sse_raw(&app, request).await;
6279 let request_id = body
6280 .lines()
6281 .find_map(|l| l.strip_prefix("data: "))
6282 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6283 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6284 .unwrap();
6285
6286 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
6287 assert_eq!(status, StatusCode::NOT_FOUND);
6288 assert_eq!(err["error"]["code"], "stream_not_found");
6289 assert!(err["error"]["message"]
6290 .as_str()
6291 .unwrap()
6292 .contains("stream_resumable"));
6293 }
6294
6295 #[test]
6299 fn the_axum_stream_patterns_match_the_published_templates() {
6300 assert_eq!(
6301 axum_path(ferrox_api::routes::V1_STREAM),
6302 "/v1/stream/:request_id"
6303 );
6304 assert_eq!(
6305 axum_path(ferrox_api::routes::V1_STREAM_POLL),
6306 "/v1/stream/:request_id/poll"
6307 );
6308 assert_eq!(
6309 ferrox_api::routes::v1_stream("abc"),
6310 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
6311 );
6312 }
6313
6314 #[test]
6324 fn no_published_template_reaches_the_router_with_its_braces() {
6325 for template in [
6326 ferrox_api::routes::V1_STREAM,
6327 ferrox_api::routes::V1_STREAM_POLL,
6328 ferrox_api::routes::V1_RESPONSE,
6329 ferrox_api::routes::V1_RESPONSE_CANCEL,
6330 ferrox_api::routes::ADMIN_TASK_CANCEL,
6331 ] {
6332 assert!(
6333 template.contains('{'),
6334 "{template} is in the template list but has no placeholder"
6335 );
6336 let mounted = axum_path(template);
6337 assert!(
6338 !mounted.contains('{') && !mounted.contains('}'),
6339 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6340 );
6341 assert!(
6342 mounted.contains(':'),
6343 "{template} lost its placeholder entirely and would match one path only"
6344 );
6345 }
6346 }
6347
6348 #[tokio::test]
6356 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6357 let app = test_app();
6358 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6359 assert_eq!(status, StatusCode::NOT_FOUND);
6360 assert!(
6361 !body.is_null(),
6362 "empty body means axum never matched the route, so the id was read as a literal segment"
6363 );
6364 }
6365
6366 #[tokio::test]
6369 async fn the_task_list_starts_empty_rather_than_absent() {
6370 let app = test_app();
6371 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
6372 assert_eq!(status, StatusCode::OK);
6373 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6374 }
6375
6376 async fn post_json_uri(
6377 app: &Router,
6378 uri: &str,
6379 body: serde_json::Value,
6380 ) -> (StatusCode, serde_json::Value) {
6381 use http_body_util::BodyExt;
6382 use tower::ServiceExt;
6383
6384 let response = app
6385 .clone()
6386 .oneshot(
6387 axum::http::Request::builder()
6388 .method("POST")
6389 .uri(uri)
6390 .header("content-type", "application/json")
6391 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6392 .unwrap(),
6393 )
6394 .await
6395 .unwrap();
6396 let status = response.status();
6397 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6398 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6399 (status, json)
6400 }
6401
6402 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6403 post_json_uri(app, "/v1/chat/completions", body).await.1
6404 }
6405
6406 #[tokio::test]
6415 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6416 let app = test_app();
6417 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6418 assert_eq!(status, StatusCode::OK);
6419
6420 let memory = &body["memory"];
6421 if memory.is_null() {
6422 return;
6425 }
6426 assert!(
6427 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6428 "a read that produced a zero is a broken read, not an idle \
6429 engine: {memory}"
6430 );
6431 assert!(
6432 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6433 "the quantity must travel with the number: {memory}"
6434 );
6435 }
6436
6437 #[tokio::test]
6443 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6444 let app = test_app();
6445 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6446 assert_eq!(status, StatusCode::OK);
6447 for pool in ["kv_pages", "window_slots", "state_slots"] {
6448 assert!(
6449 body["pools"][pool].is_null(),
6450 "{pool} must be null rather than a zero row: {}",
6451 body["pools"]
6452 );
6453 }
6454 }
6455
6456 #[tokio::test]
6464 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6465 use http_body_util::BodyExt;
6466 use tower::ServiceExt;
6467
6468 let app = test_app();
6469 let send = |body: serde_json::Value| {
6470 let app = app.clone();
6471 async move {
6472 app.oneshot(
6473 axum::http::Request::builder()
6474 .method("POST")
6475 .uri(ferrox_api::routes::V1_MESSAGES)
6476 .header("content-type", "application/json")
6477 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6478 .unwrap(),
6479 )
6480 .await
6481 .unwrap()
6482 }
6483 };
6484
6485 let ok = send(serde_json::json!({
6486 "model": "test",
6487 "max_tokens": 1,
6488 "messages": [{"role": "user", "content": "hi"}],
6489 }))
6490 .await;
6491 assert_eq!(ok.status(), StatusCode::OK);
6492 let id = ok
6493 .headers()
6494 .get("request-id")
6495 .expect("a served message names its id")
6496 .to_str()
6497 .unwrap()
6498 .to_string();
6499 assert!(!id.is_empty());
6500
6501 let bad = send(serde_json::json!({"model": "test"})).await;
6504 assert!(bad.status().is_client_error());
6505 let other = bad.headers().get("request-id").expect("errors too");
6506 assert_ne!(other.to_str().unwrap(), id);
6507 let _ = bad.into_body().collect().await.unwrap();
6508 }
6509
6510 #[tokio::test]
6518 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6519 let state = Arc::new(test_state(
6520 test_model_full_byte_vocab(),
6521 ResponseCache::new(1000, Duration::from_secs(3600)),
6522 ));
6523 let app = test_app_with_state(Arc::clone(&state));
6524 let body = serde_json::json!({
6525 "model": "test",
6526 "messages": [{"role": "user", "content": "hi"}],
6527 "max_tokens": 1,
6528 });
6529
6530 state
6531 .maintenance
6532 .lock()
6533 .unwrap()
6534 .begin_rebuild()
6535 .expect("a fresh server is serving, so the rebuild starts");
6536 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6537 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6538 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6539
6540 state.maintenance.lock().unwrap().finish_rebuild(true);
6541 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6542 assert_eq!(
6543 status,
6544 StatusCode::OK,
6545 "the gate reopens; a rebuild is not a latch"
6546 );
6547 }
6548
6549 #[tokio::test]
6554 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6555 let app = test_app();
6556 let (status, body) = post_json_uri(
6557 &app,
6558 ferrox_api::routes::V1_CANCEL,
6559 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6560 )
6561 .await;
6562 assert_eq!(status, StatusCode::NOT_FOUND);
6563 assert_eq!(body["cancelled"], serde_json::json!(false));
6564 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6565 assert!(
6566 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6567 "the verdict must carry a human reason: {body}"
6568 );
6569 }
6570
6571 #[tokio::test]
6575 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6576 let state = Arc::new(test_state(
6577 test_model_full_byte_vocab(),
6578 ResponseCache::new(1000, Duration::from_secs(3600)),
6579 ));
6580 let app = test_app_with_state(Arc::clone(&state));
6581 let (token, _guard) = state.cancels.register("chatcmpl-live");
6582
6583 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6584 assert_eq!(status, StatusCode::OK);
6585 assert_eq!(before["generating_now"], serde_json::json!(1));
6586
6587 let (status, body) = post_json_uri(
6588 &app,
6589 ferrox_api::routes::V1_CANCEL,
6590 serde_json::json!({ "request_id": "chatcmpl-live" }),
6591 )
6592 .await;
6593 assert_eq!(status, StatusCode::OK);
6594 assert_eq!(body["cancelled"], serde_json::json!(true));
6595 assert!(
6596 token.is_cancelled(),
6597 "the endpoint answered ok without setting the flag the decode loop reads"
6598 );
6599 }
6600
6601 #[tokio::test]
6602 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6603 let app = test_app();
6604 let (status, tok) =
6605 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6606 assert_eq!(status, StatusCode::OK);
6607 let tokens = tok["tokens"].as_array().unwrap();
6608 assert_eq!(tok["count"], tokens.len());
6609 assert!(!tokens.is_empty());
6610
6611 let (status, detok) = post_json_uri(
6612 &app,
6613 "/v1/detokenize",
6614 serde_json::json!({ "tokens": tokens }),
6615 )
6616 .await;
6617 assert_eq!(status, StatusCode::OK);
6618 assert_eq!(detok["text"], "Hi");
6619
6620 let (status, emb) = post_json_uri(
6621 &app,
6622 "/v1/embeddings",
6623 serde_json::json!({
6624 "input": "Hi",
6625 "embedding_type": "mean"
6626 }),
6627 )
6628 .await;
6629 assert_eq!(status, StatusCode::OK);
6630 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6631 assert!(!vec.is_empty());
6632 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6633 }
6634
6635 #[tokio::test]
6641 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6642 let app = test_app();
6643 let (status, body) = post_json_uri(
6644 &app,
6645 "/v1/embeddings",
6646 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6647 )
6648 .await;
6649 assert_eq!(status, StatusCode::BAD_REQUEST);
6650 let msg = body["error"]["message"].as_str().unwrap();
6651 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6652 }
6653
6654 #[tokio::test]
6659 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6660 async fn a_real_embedding_model_serves_v1_embeddings() {
6661 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6662 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6663 if !path.exists() {
6664 eprintln!("SKIP: {} not present", path.display());
6665 return;
6666 }
6667 let encoder = ferrox_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6668 let mut state = test_state(
6669 test_model_full_byte_vocab(),
6670 ResponseCache::new(1000, Duration::from_secs(3600)),
6671 );
6672 state.embedding = Some(Arc::new(encoder));
6673 let app = test_app_with_state(Arc::new(state));
6674
6675 let (status, body) = post_json_uri(
6676 &app,
6677 "/v1/embeddings",
6678 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6679 )
6680 .await;
6681 assert_eq!(status, StatusCode::OK, "{body}");
6682 assert_eq!(body["model"], "bge-small-en-v1.5");
6683 let data = body["data"].as_array().unwrap();
6684 assert_eq!(data.len(), 2);
6685 for (i, row) in data.iter().enumerate() {
6686 assert_eq!(row["index"], i);
6687 let v: Vec<f64> = row["embedding"]
6688 .as_array()
6689 .unwrap()
6690 .iter()
6691 .map(|x| x.as_f64().unwrap())
6692 .collect();
6693 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6694 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6695 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6696 }
6697 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6700
6701 let (status, mean) = post_json_uri(
6705 &app,
6706 "/v1/embeddings",
6707 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6708 )
6709 .await;
6710 assert_eq!(status, StatusCode::OK);
6711 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6712 }
6713
6714 #[tokio::test]
6726 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6727 async fn an_encoder_can_be_the_loaded_model() {
6728 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6729 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6730 if !path.exists() {
6731 eprintln!("SKIP: {} not present", path.display());
6732 return;
6733 }
6734
6735 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6739 assert!(
6740 matches!(loaded, model::LoadedModel::Encoder(_)),
6741 "an encoder-only GGUF reached a decoder loader"
6742 );
6743 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6744 assert!(
6745 matches!(loaded, Loaded::Encoder(_)),
6746 "the encoder did not stay an encoder through activation"
6747 );
6748 assert!(
6749 batcher.is_none() && ceiling.is_none(),
6750 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6751 );
6752
6753 let state = test_state(
6754 test_model_full_byte_vocab(),
6755 ResponseCache::new(1000, Duration::from_secs(3600)),
6756 );
6757 state.swap_active(Some(Arc::new(ActiveModel {
6758 id: None,
6759 loaded,
6760 batcher,
6761 ceiling,
6762 })));
6763 let app = test_app_with_state(Arc::new(state));
6764
6765 let (status, body) = post_json_uri(
6767 &app,
6768 "/v1/embeddings",
6769 serde_json::json!({ "input": "Hello world" }),
6770 )
6771 .await;
6772 assert_eq!(status, StatusCode::OK, "{body}");
6773 assert_eq!(body["model"], "bge-small-en-v1.5");
6774 let v = body["data"][0]["embedding"].as_array().unwrap();
6775 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6776
6777 let (status, body) = post_json_uri(
6779 &app,
6780 "/v1/chat/completions",
6781 serde_json::json!({
6782 "model": "bge-small-en-v1.5",
6783 "messages": [{"role": "user", "content": "hi"}],
6784 }),
6785 )
6786 .await;
6787 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6788 let msg = body["error"]["message"].as_str().unwrap();
6789 for fact in [
6790 "bge-small-en-v1.5",
6791 "bert",
6792 "embedding model",
6793 "/v1/embeddings",
6794 ] {
6795 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6796 }
6797
6798 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
6800 assert_eq!(status, StatusCode::OK);
6801 let entry = &models["data"][0];
6802 assert_eq!(entry["id"], "bge-small-en-v1.5");
6803 assert_eq!(entry["ferrox_model_kind"], "embedding");
6804 assert_eq!(entry["ferrox_tokenizer"], "gguf-wordpiece");
6805 assert_eq!(entry["ferrox_n_embd"], 384);
6806 assert_eq!(entry["ferrox_pooling"], "CLS");
6807 assert_eq!(
6808 entry["ferrox_endpoints"],
6809 serde_json::json!(["/v1/embeddings"])
6810 );
6811 assert!(entry.get("supported_reasoning_efforts").is_none());
6814
6815 let (status, health) = get_json(&app, ferrox_api::routes::HEALTH).await;
6817 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6818 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6819 assert_eq!(health["model"]["synthetic_weights"], false);
6820 let weights = health["capabilities"]
6821 .as_array()
6822 .unwrap()
6823 .iter()
6824 .find(|c| c["id"] == ferrox_api::health::capability::REAL_WEIGHTS)
6825 .expect("a real-weights capability row");
6826 let detail = weights["detail"].as_str().unwrap_or_default();
6827 assert!(detail.contains("ENCODER"), "{detail}");
6828 }
6829
6830 #[tokio::test]
6835 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6836 use http_body_util::BodyExt;
6837 use tower::ServiceExt;
6838
6839 let fixture = concat!(
6840 "../ferrox-models/tests/fixtures/",
6841 "ferrox_real_moe_test.gguf"
6842 );
6843 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6844 let decoder = Decoder::from_gguf_with_expert_cache(
6845 &fixture,
6846 ferrox_models::config::test_moe_fixture(),
6847 Some(1024 * 1024),
6848 )
6849 .expect("MoE fixture must load store-backed");
6850
6851 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
6855 .layers
6856 .iter()
6857 .map(|_| {
6858 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
6859 })
6860 .collect();
6861 decoder.forward_token(1, 0, &mut caches);
6862
6863 let model = Model::Gguf(GgufModel {
6864 decoder: Arc::new(decoder),
6865 tokenizer: Arc::new(ServerTokenizer::Byte),
6866 stop_tokens: StopTokens::default(),
6867 bos_id: None,
6868 is_synthetic: false,
6869 chat_template: chat_template::PromptTemplate::plain(),
6870 });
6871 let state = Arc::new(test_state(
6872 model,
6873 ResponseCache::new(16, Duration::from_secs(60)),
6874 ));
6875 let app = Router::new()
6876 .route("/metrics", axum::routing::get(metrics))
6877 .route("/v1/chat/completions", post(chat_completions))
6878 .with_state(state);
6879
6880 let fetch_metrics = |app: Router| async move {
6881 let resp = app
6882 .oneshot(
6883 axum::http::Request::builder()
6884 .method("GET")
6885 .uri("/metrics")
6886 .body(axum::body::Body::empty())
6887 .unwrap(),
6888 )
6889 .await
6890 .unwrap();
6891 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6892 String::from_utf8(bytes.to_vec()).unwrap()
6893 };
6894
6895 let after = fetch_metrics(app.clone()).await;
6896 assert!(
6897 after.contains("ferrox_expert_cache_misses_total"),
6898 "streaming model must expose expert-cache metrics: {after}"
6899 );
6900 let misses: u64 = after
6901 .lines()
6902 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
6903 .and_then(|l| l.split_whitespace().nth(1))
6904 .and_then(|v| v.parse().ok())
6905 .expect("misses metric line must parse");
6906 assert!(
6907 misses > 0,
6908 "decode must have read experts through the store: {after}"
6909 );
6910 }
6911
6912 fn weather_tool() -> serde_json::Value {
6913 serde_json::json!({
6914 "type": "function",
6915 "function": {
6916 "name": "get_weather",
6917 "description": "Get the current weather for a location.",
6918 "parameters": {
6919 "type": "object",
6920 "properties": {"location": {"type": "string"}},
6921 "required": ["location"]
6922 }
6923 }
6924 })
6925 }
6926
6927 fn weather_tool_def() -> ToolDef {
6928 ToolDef {
6929 kind: "function".to_string(),
6930 function: ToolFunctionDef {
6931 name: "get_weather".to_string(),
6932 description: Some("Get the current weather for a location.".to_string()),
6933 parameters: Some(serde_json::json!({
6934 "type": "object",
6935 "properties": {"location": {"type": "string"}},
6936 "required": ["location"]
6937 })),
6938 },
6939 }
6940 }
6941
6942 #[test]
6943 fn tool_preamble_mentions_every_tool_name_and_description() {
6944 let preamble = tool_preamble(&[weather_tool_def()]);
6945 assert!(preamble.contains("get_weather"));
6946 assert!(preamble.contains("Get the current weather for a location."));
6947 assert!(preamble.contains("<tool_call>"));
6948 assert!(preamble.contains("</tool_call>"));
6949 }
6950
6951 #[test]
6952 fn a_real_marker_becomes_a_structured_tool_call() {
6953 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6954 let (message, finish) = build_response_message(
6955 text.to_string(),
6956 &[weather_tool_def()],
6957 output::OutputPosture::for_model("test-model"),
6958 "stop",
6959 );
6960 assert_eq!(finish, "tool_calls");
6961 let calls = message.tool_calls.expect("must carry a tool call");
6962 assert_eq!(calls[0].function.name, "get_weather");
6963 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
6964 assert_eq!(parsed["location"], "Paris");
6965 }
6966
6967 #[test]
6968 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
6969 let (message, finish) = build_response_message(
6970 "just an answer".to_string(),
6971 &[weather_tool_def()],
6972 output::OutputPosture::for_model("test-model"),
6973 "stop",
6974 );
6975 assert_eq!(finish, "stop");
6976 assert!(message.tool_calls.is_none());
6977 assert_eq!(message.content.as_deref(), Some("just an answer"));
6978 }
6979
6980 #[test]
6983 fn a_malformed_payload_is_not_a_tool_call() {
6984 let (message, finish) = build_response_message(
6985 "<tool_call>not valid json at all</tool_call>".to_string(),
6986 &[weather_tool_def()],
6987 output::OutputPosture::for_model("test-model"),
6988 "stop",
6989 );
6990 assert_eq!(finish, "stop");
6991 assert!(message.tool_calls.is_none());
6992 }
6993
6994 #[test]
6997 fn a_tool_that_was_never_offered_is_not_returned() {
6998 let (message, finish) = build_response_message(
6999 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7000 &[weather_tool_def()],
7001 output::OutputPosture::for_model("test-model"),
7002 "stop",
7003 );
7004 assert_eq!(finish, "stop");
7005 assert!(message.tool_calls.is_none());
7006 }
7007
7008 #[test]
7010 fn marker_text_with_no_tools_offered_stays_content() {
7011 let (message, finish) = build_response_message(
7012 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7013 &[],
7014 output::OutputPosture::for_model("test-model"),
7015 "stop",
7016 );
7017 assert_eq!(finish, "stop");
7018 assert!(message.tool_calls.is_none());
7019 assert!(message.content.is_some());
7020 }
7021
7022 #[test]
7026 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7027 let opened = std::cell::Cell::new(0usize);
7028 let mut parser = crate::policy::parser::ToolCallParser::new(
7029 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7030 vec![
7031 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7032 "write_file",
7033 serde_json::json!({"type": "object", "properties": {
7034 "path": {"type": "string"},
7035 "contents": {"type": "string"}
7036 }}),
7037 ),
7038 ],
7039 );
7040 let wire = "<tool_call><function=write_file>\
7041 <parameter=path>\n/tmp/x\n</parameter>\
7042 <parameter=contents>\nhello world\n</parameter>\
7043 </function></tool_call>";
7044
7045 let mut deltas = Vec::new();
7046 let mut text = String::new();
7047 for piece in wire.as_bytes().chunks(7) {
7048 let chunk = String::from_utf8_lossy(piece).into_owned();
7049 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7050 text.push_str(&more_text);
7051 deltas.extend(more);
7052 }
7053 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7054 text.push_str(&more_text);
7055 deltas.extend(more);
7056
7057 assert_eq!(opened.get(), 1, "one call opened");
7058 assert!(text.is_empty(), "the markers are not content: {text:?}");
7059
7060 let first = &deltas[0];
7061 assert_eq!(first.index, 0);
7062 assert_eq!(first.id.as_deref(), Some("call_0"));
7063 assert_eq!(first.kind, Some("function"));
7064 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7065
7066 let joined: String = deltas
7069 .iter()
7070 .filter_map(|d| d.function.arguments.clone())
7071 .collect();
7072 let parsed: serde_json::Value =
7073 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7074 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7075 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7076 assert!(
7077 deltas.len() >= 3,
7078 "the arguments arrived in pieces, not whole: {}",
7079 deltas.len()
7080 );
7081 assert!(
7082 deltas[1..].iter().all(|d| d.function.name.is_none()),
7083 "only the opening delta carries identity"
7084 );
7085 }
7086
7087 #[test]
7089 fn text_around_a_streamed_call_is_still_content() {
7090 let opened = std::cell::Cell::new(0usize);
7091 let mut parser = crate::policy::parser::ToolCallParser::new(
7092 crate::policy::parser::ToolCallFormat::Qwen25,
7093 vec![crate::policy::parser::tool_call::ToolSchema::new(
7094 "get_weather",
7095 )],
7096 );
7097 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7098 \"arguments\": {}}</tool_call> done";
7099 let mut text = String::new();
7100 for piece in wire.as_bytes().chunks(5) {
7101 let chunk = String::from_utf8_lossy(piece).into_owned();
7102 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7103 text.push_str(&more);
7104 }
7105 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7106 text.push_str(&more);
7107
7108 assert_eq!(opened.get(), 1);
7109 assert!(text.starts_with("let me check. "), "{text:?}");
7110 assert!(text.ends_with(" done"), "{text:?}");
7111 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7112 }
7113
7114 #[test]
7117 fn a_reasoning_block_is_split_out_of_the_answer() {
7118 let (message, finish) = build_response_message(
7119 "<think>weighing it up</think>The answer is 4.".to_string(),
7120 &[],
7121 output::OutputPosture::for_model("Qwen3-8B"),
7122 "stop",
7123 );
7124 assert_eq!(finish, "stop");
7125 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7126 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7127 }
7128
7129 #[test]
7132 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7133 let (message, _) = build_response_message(
7134 "Use the <think> tag like this.".to_string(),
7135 &[],
7136 output::OutputPosture::for_model("llama-3.1-8b"),
7137 "stop",
7138 );
7139 assert_eq!(
7140 message.content.as_deref(),
7141 Some("Use the <think> tag like this.")
7142 );
7143 assert!(message.reasoning_content.is_none());
7144 }
7145
7146 #[tokio::test]
7154 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7155 let app = test_app();
7156 let body = serde_json::json!({
7157 "model": "m",
7158 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7159 "max_tokens": 4,
7160 "temperature": 0,
7161 });
7162 let resp = post_json(&app, body).await;
7163 let message = &resp["choices"][0]["message"];
7164 assert!(message["content"].is_string());
7165 assert!(message.get("tool_calls").is_none());
7166 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7167 assert_eq!(resp["usage"]["completion_tokens"], 4);
7168 assert_eq!(
7169 resp["usage"]["total_tokens"],
7170 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7171 );
7172 }
7173
7174 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7175 use http_body_util::BodyExt;
7176 use tower::ServiceExt;
7177
7178 let response = app
7179 .clone()
7180 .oneshot(
7181 axum::http::Request::builder()
7182 .method("GET")
7183 .uri(uri)
7184 .body(axum::body::Body::empty())
7185 .unwrap(),
7186 )
7187 .await
7188 .unwrap();
7189 let status = response.status();
7190 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7191 (status, serde_json::from_slice(&bytes).unwrap())
7192 }
7193
7194 #[tokio::test]
7195 async fn health_answers_a_capability_handshake_not_a_boolean() {
7196 let app = test_app();
7197 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7198 assert_eq!(status, StatusCode::OK);
7199
7200 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7201 assert_eq!(health.state, ferrox_api::HealthState::Ready);
7202 assert!(health.pid > 0);
7203 assert!(health.server_time_unix_ms > 0);
7204 assert_eq!(health.last_request_age_seconds, None);
7207
7208 for id in [
7211 ferrox_api::health::capability::CPU,
7212 ferrox_api::health::capability::METAL,
7213 ferrox_api::health::capability::CUDA,
7214 ferrox_api::health::capability::REAL_WEIGHTS,
7215 ferrox_api::health::capability::CONTINUOUS_BATCHING,
7216 ] {
7217 let cap = health
7218 .capability(id)
7219 .unwrap_or_else(|| panic!("{id} missing"));
7220 assert!(!cap.reason.is_empty(), "{cap:?}");
7221 assert!(!cap.detail.is_empty(), "{cap:?}");
7222 }
7223 let weights = health
7227 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
7228 .unwrap();
7229 assert!(!weights.available);
7230 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
7231 assert!(health.model.as_ref().unwrap().synthetic_weights);
7232 }
7233
7234 #[tokio::test]
7235 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7236 let app = test_app();
7237 let _ = post_json(
7238 &app,
7239 serde_json::json!({
7240 "model": "m",
7241 "messages": [{"role": "user", "content": "\u{1}"}],
7242 "max_tokens": 1,
7243 "temperature": 0,
7244 }),
7245 )
7246 .await;
7247 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7248 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7249 let age = health
7250 .last_request_age_seconds
7251 .expect("a served request is evidence of liveness");
7252 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7253 }
7254
7255 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7257 use http_body_util::BodyExt;
7258 use tower::ServiceExt;
7259
7260 let response = app
7261 .clone()
7262 .oneshot(
7263 axum::http::Request::builder()
7264 .method("POST")
7265 .uri("/v1/chat/completions")
7266 .header("content-type", "application/json")
7267 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7268 .unwrap(),
7269 )
7270 .await
7271 .unwrap();
7272 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7273 String::from_utf8(bytes.to_vec())
7274 .unwrap()
7275 .lines()
7276 .filter_map(|line| line.strip_prefix("data: "))
7277 .filter(|payload| *payload != "[DONE]")
7278 .map(|payload| serde_json::from_str(payload).unwrap())
7279 .collect()
7280 }
7281
7282 #[tokio::test]
7283 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7284 let app = test_app();
7285 let chunks = post_sse_chunks(
7286 &app,
7287 serde_json::json!({
7288 "model": "m",
7289 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7290 "max_tokens": 4,
7291 "temperature": 0,
7292 "stream": true,
7293 }),
7294 )
7295 .await;
7296
7297 assert!(!chunks.is_empty());
7298 let request_id = chunks[0]["request_id"]
7299 .as_str()
7300 .expect("the first chunk names the request")
7301 .to_string();
7302 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7303 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7306 assert!(
7307 chunk.get("request_id").is_none(),
7308 "chunk {i} repeats request_id"
7309 );
7310 }
7311 for chunk in &chunks {
7314 assert_eq!(chunk["id"], serde_json::json!(request_id));
7315 }
7316
7317 let other = post_sse_chunks(
7318 &app,
7319 serde_json::json!({
7320 "model": "m",
7321 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7322 "max_tokens": 4,
7323 "temperature": 0,
7324 "stream": true,
7325 }),
7326 )
7327 .await;
7328 assert_ne!(
7329 other[0]["request_id"].as_str().unwrap(),
7330 request_id,
7331 "two concurrent chats must not share an id"
7332 );
7333 }
7334
7335 #[tokio::test]
7336 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7337 let app = test_app();
7338 let resp = post_json(
7339 &app,
7340 serde_json::json!({
7341 "model": "m",
7342 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7343 "max_tokens": 2,
7344 "temperature": 0,
7345 }),
7346 )
7347 .await;
7348 assert_eq!(resp["id"], resp["request_id"]);
7349 assert!(resp["request_id"]
7350 .as_str()
7351 .unwrap()
7352 .starts_with("chatcmpl-"));
7353 }
7354
7355 #[tokio::test]
7358 async fn usage_carries_separate_prefill_and_decode_timings() {
7359 let app = test_app();
7360 let resp = post_json(
7361 &app,
7362 serde_json::json!({
7363 "model": "m",
7364 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7365 "max_tokens": 4,
7366 "temperature": 0,
7367 }),
7368 )
7369 .await;
7370 let usage = &resp["usage"];
7371 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7372 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7373 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7374 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7375 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7377 }
7378
7379 #[tokio::test]
7387 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7388 let app = test_app();
7389 let body = serde_json::json!({
7390 "model": "m",
7391 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7392 "max_tokens": 4,
7393 "temperature": 0,
7394 "tools": [weather_tool()],
7395 });
7396 let resp = post_json(&app, body).await;
7397 let message = &resp["choices"][0]["message"];
7398 assert!(
7399 message["content"].is_string(),
7400 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7401 );
7402 assert!(message.get("tool_calls").is_none());
7403 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7407 }
7408
7409 #[tokio::test]
7413 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7414 let app = test_app();
7415 let body = serde_json::json!({
7416 "model": "m",
7417 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7418 "max_tokens": 3,
7419 "temperature": 0,
7420 });
7421 let first = post_json(&app, body.clone()).await;
7422 assert_eq!(first["ferrox_cache"], "miss");
7423 let second = post_json(&app, body).await;
7424 assert_eq!(second["ferrox_cache"], "hit");
7425 assert_eq!(
7426 first["choices"][0]["message"]["content"],
7427 second["choices"][0]["message"]["content"]
7428 );
7429 assert_eq!(
7430 first["choices"][0]["finish_reason"],
7431 second["choices"][0]["finish_reason"]
7432 );
7433 assert_eq!(first["usage"], second["usage"]);
7434 assert_eq!(second["usage"]["completion_tokens"], 3);
7435 }
7436
7437 #[tokio::test]
7443 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7444 let session_app = test_app();
7445 let manual_app = test_app();
7446
7447 let turn1 = post_json(
7449 &session_app,
7450 serde_json::json!({
7451 "model": "m",
7452 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7453 "session_id": "s1",
7454 "max_tokens": 5,
7455 "temperature": 0,
7456 }),
7457 )
7458 .await;
7459 let reply1 = turn1["choices"][0]["message"]["content"]
7460 .as_str()
7461 .unwrap()
7462 .to_string();
7463
7464 let manual_turn1 = post_json(
7469 &manual_app,
7470 serde_json::json!({
7471 "model": "m",
7472 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7473 "max_tokens": 5,
7474 "temperature": 0,
7475 }),
7476 )
7477 .await;
7478 assert_eq!(
7479 manual_turn1["choices"][0]["message"]["content"]
7480 .as_str()
7481 .unwrap(),
7482 reply1
7483 );
7484
7485 let turn2 = post_json(
7487 &session_app,
7488 serde_json::json!({
7489 "model": "m",
7490 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7491 "session_id": "s1",
7492 "max_tokens": 5,
7493 "temperature": 0,
7494 }),
7495 )
7496 .await;
7497 let reply2 = turn2["choices"][0]["message"]["content"]
7498 .as_str()
7499 .unwrap()
7500 .to_string();
7501
7502 let manual_turn2 = post_json(
7506 &manual_app,
7507 serde_json::json!({
7508 "model": "m",
7509 "messages": [
7510 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7511 {"role": "assistant", "content": reply1},
7512 {"role": "user", "content": "\u{4}\u{5}"},
7513 ],
7514 "max_tokens": 5,
7515 "temperature": 0,
7516 }),
7517 )
7518 .await;
7519 assert_eq!(
7520 manual_turn2["choices"][0]["message"]["content"]
7521 .as_str()
7522 .unwrap(),
7523 reply2,
7524 "resuming a session must produce identical output to manually resending the full history"
7525 );
7526 }
7527
7528 #[test]
7531 fn lock_cache_recovers_from_a_poisoned_mutex() {
7532 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7533
7534 let poison_cache = Arc::clone(&cache);
7535 let _ = std::thread::spawn(move || {
7536 let _guard = poison_cache.lock().unwrap();
7537 panic!("simulated panic while holding the lock");
7538 })
7539 .join();
7540
7541 let recovered = lock_cache(&cache);
7543 assert_eq!(recovered.stats().entries, 0);
7544 }
7545
7546 #[test]
7547 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7548 let mut req_body = serde_json::json!({
7549 "model": "m",
7550 "messages": [{"role": "user", "content": "hi"}],
7551 });
7552 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7553 assert!(
7554 req.is_cacheable(),
7555 "default (temperature 0) must be cacheable"
7556 );
7557
7558 req_body["temperature"] = serde_json::json!(0.8);
7559 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7560 assert!(
7561 !req.is_cacheable(),
7562 "unseeded sampling must never be cacheable"
7563 );
7564
7565 req_body["seed"] = serde_json::json!(42);
7566 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7567 assert!(
7568 req.is_cacheable(),
7569 "sampling with an explicit seed is deterministic and must be cacheable"
7570 );
7571 }
7572
7573 const GRADED: &str = "{% if reasoning_effort %}\
7577 {% if reasoning_effort not in ['low','medium','high'] %}\
7578 {{ raise_exception('unsupported effort') }}\
7579 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7580 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7581
7582 fn graded_template() -> chat_template::PromptTemplate {
7583 chat_template::PromptTemplate::from_gguf_metadata(
7584 Some(GRADED),
7585 Some("qwen3"),
7586 false,
7587 true,
7588 None,
7589 None,
7590 )
7591 }
7592
7593 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7594 serde_json::from_value(value).expect("request")
7595 }
7596
7597 #[test]
7604 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7605 let req = chat_request(serde_json::json!({
7606 "model": "m",
7607 "messages": [{"role": "user", "content": "hi"}],
7608 "grammar": "root ::= \"a\"+",
7609 }));
7610 req.validate_supported_fields()
7611 .expect("a valid grammar is a valid request");
7612 let params = req
7613 .generation_params()
7614 .expect("a valid grammar compiles at params time too");
7615 assert!(
7616 params.grammar.is_some(),
7617 "the grammar was dropped between the wire and the sampler"
7618 );
7619 assert!(
7620 params.needs_vocab_logits(),
7621 "a grammar request that may fold lm_head into a GPU argmax is \
7622 a grammar request served unconstrained"
7623 );
7624
7625 let plain = chat_request(serde_json::json!({
7626 "model": "m",
7627 "messages": [{"role": "user", "content": "hi"}],
7628 }));
7629 assert!(plain.generation_params().unwrap().grammar.is_none());
7630 }
7631
7632 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7633 chat_request(serde_json::json!({
7634 "model": "m",
7635 "messages": [{"role": "user", "content": "weather in Rome?"}],
7636 "tools": [weather_tool()],
7637 "tool_choice": tool_choice,
7638 }))
7639 }
7640
7641 #[test]
7645 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7646 for choice in [
7647 serde_json::json!("required"),
7648 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7649 ] {
7650 let req = tool_request(choice.clone());
7651 req.validate_supported_fields()
7652 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7653 let params = req
7654 .generation_params_for_template(&graded_template(), "Qwen3-8B")
7655 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7656 let grammar = params
7657 .grammar
7658 .as_ref()
7659 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7660 assert!(
7661 grammar.is_awaiting_trigger(),
7662 "the model must be free to think before it calls"
7663 );
7664 assert!(
7665 !grammar.allows_eog(),
7666 "{choice} must not be able to end the turn without a call"
7667 );
7668 assert!(
7674 params.needs_vocab_logits(),
7675 "{choice} would let a backend return a token id instead of logits"
7676 );
7677 assert!(
7678 !generate::greedy_gpu_fold_allowed(¶ms),
7679 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7680 );
7681 }
7682 }
7683
7684 #[test]
7686 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7687 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7688 let req = tool_request(choice.clone());
7689 req.validate_supported_fields().expect("still supported");
7690 let params = match req.generation_params_for_template(&graded_template(), "Qwen3-8B") {
7691 Ok(p) => p,
7692 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7693 };
7694 assert!(
7695 params.grammar.is_none(),
7696 "{choice} does not force a call and must not be constrained"
7697 );
7698 }
7699 }
7700
7701 #[test]
7704 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7705 let req = chat_request(serde_json::json!({
7707 "model": "m",
7708 "messages": [{"role": "user", "content": "hi"}],
7709 "tool_choice": "required",
7710 }));
7711 let (status, _) = req
7712 .validate_supported_fields()
7713 .expect_err("nothing to call");
7714 assert_eq!(status, StatusCode::BAD_REQUEST);
7715
7716 let req =
7718 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7719 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7720 assert_eq!(status, StatusCode::BAD_REQUEST);
7721 assert_eq!(body["error"]["param"], "tool_choice");
7722
7723 let req = tool_request(serde_json::json!({"type": "function"}));
7725 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7726 assert_eq!(status, StatusCode::BAD_REQUEST);
7727
7728 let req = chat_request(serde_json::json!({
7730 "model": "m",
7731 "messages": [{"role": "user", "content": "hi"}],
7732 "tools": [weather_tool()],
7733 "tool_choice": "required",
7734 "grammar": "root ::= \"a\"+",
7735 }));
7736 let (status, _) = req
7737 .validate_supported_fields()
7738 .expect_err("a grammar and a forced call are two constraints");
7739 assert_eq!(status, StatusCode::BAD_REQUEST);
7740
7741 let req = tool_request(serde_json::json!("required"));
7744 let (status, Json(body)) =
7745 match req.generation_params_for_template(&graded_template(), "GLM-4.7") {
7746 Err(e) => e,
7747 Ok(_) => panic!("glm calls are not JSON behind a marker"),
7748 };
7749 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7750 assert!(
7751 body["error"]["message"].as_str().unwrap().contains("glm47"),
7752 "{body}"
7753 );
7754 }
7755
7756 #[test]
7759 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
7760 let req = chat_request(serde_json::json!({
7761 "model": "m",
7762 "messages": [{"role": "user", "content": "hi"}],
7763 "grammar": "root ::= \"a",
7764 }));
7765 let (status, Json(body)) = req
7766 .validate_supported_fields()
7767 .expect_err("this does not parse");
7768 assert_eq!(status, StatusCode::BAD_REQUEST);
7769 assert_eq!(body["error"]["param"], "grammar");
7770 assert!(req.generation_params().is_err(), "and again at params time");
7771 }
7772
7773 #[test]
7778 fn response_format_json_schema_names_the_missing_converter() {
7779 let req = chat_request(serde_json::json!({
7780 "model": "m",
7781 "messages": [{"role": "user", "content": "hi"}],
7782 "response_format": {
7783 "type": "json_schema",
7784 "json_schema": {"name": "x", "schema": {"type": "object"}},
7785 },
7786 }));
7787 let (status, Json(body)) = req
7788 .validate_supported_fields()
7789 .expect_err("not implemented yet");
7790 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7791 let message = body["error"]["message"].as_str().expect("a message");
7792 assert!(
7793 message.contains("JSON schema") && message.contains("grammar"),
7794 "the refusal must say what is missing: {message}"
7795 );
7796 }
7797
7798 #[test]
7801 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
7802 let req = chat_request(serde_json::json!({
7803 "model": "m",
7804 "messages": [{"role": "user", "content": "hi"}],
7805 }));
7806 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
7807 }
7808
7809 #[test]
7814 fn min_p_reaches_the_sampler_from_the_chat_wire() {
7815 let asked = chat_request(serde_json::json!({
7816 "model": "m",
7817 "messages": [{"role": "user", "content": "hi"}],
7818 "min_p": 0.07,
7819 }));
7820 assert_eq!(asked.sampling_params().min_p, 0.07);
7821
7822 let silent = chat_request(serde_json::json!({
7823 "model": "m",
7824 "messages": [{"role": "user", "content": "hi"}],
7825 }));
7826 assert_eq!(
7827 silent.sampling_params().min_p,
7828 0.0,
7829 "an unset min_p must be off, not llama.cpp's CLI default"
7830 );
7831 }
7832
7833 #[test]
7842 fn no_sampler_knob_is_missing_from_the_cache_key() {
7843 let base = serde_json::json!({
7844 "model": "m",
7845 "messages": [{"role": "user", "content": "hi"}],
7846 "seed": 1,
7847 });
7848 let baseline = chat_request(base.clone()).cache_key("prompt");
7849 for (knob, value) in [
7850 ("temperature", serde_json::json!(0.5)),
7851 ("top_p", serde_json::json!(0.9)),
7852 ("min_p", serde_json::json!(0.05)),
7853 ("top_k", serde_json::json!(40)),
7854 ("repetition_penalty", serde_json::json!(1.1)),
7855 ("presence_penalty", serde_json::json!(0.3)),
7856 ("frequency_penalty", serde_json::json!(0.3)),
7857 ] {
7858 let mut body = base.clone();
7859 body[knob] = value;
7860 assert_ne!(
7861 chat_request(body).cache_key("prompt"),
7862 baseline,
7863 "`{knob}` is not in the cache key: two requests differing \
7864 only in it would share one cached answer"
7865 );
7866 }
7867 }
7868
7869 #[test]
7873 fn an_explicit_zero_output_budget_is_a_client_error() {
7874 let req = chat_request(serde_json::json!({
7875 "model": "m",
7876 "messages": [{"role": "user", "content": "hi"}],
7877 "max_tokens": 0,
7878 }));
7879 let (status, body) = req.validate_supported_fields().expect_err("rejected");
7880 assert_eq!(status, StatusCode::BAD_REQUEST);
7881 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
7882 }
7883
7884 #[test]
7887 fn a_request_can_turn_thinking_off() {
7888 let template = graded_template();
7889 for body in [
7890 serde_json::json!({
7891 "model": "m",
7892 "messages": [{"role": "user", "content": "hi"}],
7893 "reasoning_effort": "none",
7894 }),
7895 serde_json::json!({
7896 "model": "m",
7897 "messages": [{"role": "user", "content": "hi"}],
7898 "thinking": {"type": "disabled"},
7899 }),
7900 ] {
7901 let kwargs = chat_request(body).resolve_template_kwargs(&template);
7902 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
7903 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
7904 assert!(!kwargs.contains_key("reasoning_effort"));
7907 }
7908 }
7909
7910 #[test]
7913 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
7914 let template = graded_template();
7915 let kwargs = chat_request(serde_json::json!({
7916 "model": "m",
7917 "messages": [{"role": "user", "content": "hi"}],
7918 "reasoning_effort": "high",
7919 "thinking": {"type": "disabled"},
7920 }))
7921 .resolve_template_kwargs(&template);
7922 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
7923 assert!(!kwargs.contains_key("reasoning_effort"));
7924 }
7925
7926 #[test]
7929 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
7930 let req = chat_request(serde_json::json!({
7931 "model": "m",
7932 "messages": [{"role": "user", "content": "hi"}],
7933 "thinking": {"type": "disable"},
7934 }));
7935 let (status, _) = req.validate_supported_fields().expect_err("rejected");
7936 assert_eq!(status, StatusCode::BAD_REQUEST);
7937 }
7938
7939 #[test]
7942 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
7943 let template = graded_template();
7944 let kwargs = chat_request(serde_json::json!({
7945 "model": "m",
7946 "messages": [{"role": "user", "content": "hi"}],
7947 "reasoning_effort": "none",
7948 "chat_template_kwargs": {"enable_thinking": true},
7949 }))
7950 .resolve_template_kwargs(&template);
7951 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
7952 }
7953
7954 #[test]
7958 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
7959 let template = graded_template();
7960 let req = chat_request(serde_json::json!({
7961 "model": "m",
7962 "messages": [{"role": "user", "content": "hi"}],
7963 "reasoning_effort": "minimal",
7964 }));
7965 let kwargs = req.resolve_template_kwargs(&template);
7966 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
7967 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
7968 assert!(prompt.starts_with("E:low|"), "{prompt}");
7969 }
7970
7971 #[test]
7975 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
7976 let template = graded_template();
7977 let req = chat_request(serde_json::json!({
7978 "model": "m",
7979 "messages": [{"role": "user", "content": "hi"}],
7980 "chat_template_kwargs": {"reasoning_effort": "none"},
7981 }));
7982 let kwargs = req.resolve_template_kwargs(&template);
7983 assert!(!kwargs.contains_key("reasoning_effort"));
7984 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
7985 assert_eq!(prompt, "hi");
7986 }
7987
7988 #[test]
7991 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
7992 let template = graded_template();
7993 let req = chat_request(serde_json::json!({
7994 "model": "m",
7995 "messages": [{"role": "user", "content": "hi"}],
7996 "reasoning_effort": "low",
7997 "chat_template_kwargs": {"reasoning_effort": "high"},
7998 }));
7999 assert_eq!(
8000 req.resolve_template_kwargs(&template)["reasoning_effort"],
8001 serde_json::json!("high")
8002 );
8003 }
8004
8005 #[test]
8009 fn offering_tools_turns_thinking_on_by_itself() {
8010 let template = graded_template();
8011 let quiet = chat_request(serde_json::json!({
8012 "model": "m",
8013 "messages": [{"role": "user", "content": "hi"}],
8014 }));
8015 assert!(!quiet
8016 .resolve_template_kwargs(&template)
8017 .contains_key("enable_thinking"));
8018
8019 let with_tools = chat_request(serde_json::json!({
8020 "model": "m",
8021 "messages": [{"role": "user", "content": "hi"}],
8022 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8023 }));
8024 let kwargs = with_tools.resolve_template_kwargs(&template);
8025 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8026 let prompt =
8027 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8028 assert!(prompt.starts_with("THINK|"), "{prompt}");
8029 }
8030
8031 #[test]
8036 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8037 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8038 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8039 Some("qwen3"),
8040 false,
8041 true,
8042 None,
8043 None,
8044 );
8045 let req = chat_request(serde_json::json!({
8046 "model": "m",
8047 "messages": [{"role": "user", "content": "hi"}],
8048 "chat_template_kwargs": {"enable_thinking": true},
8049 }));
8050 let kwargs = req.resolve_template_kwargs(&opener);
8051 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8052 assert!(prompt.ends_with("<think>"), "{prompt}");
8053
8054 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8057 let (message, _) = build_response_message(
8058 "weighing it up</think>Paris.".to_string(),
8059 &[],
8060 posture,
8061 "stop",
8062 );
8063 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8064 assert_eq!(message.content.as_deref(), Some("Paris."));
8065
8066 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8069 let (message, _) = build_response_message(
8070 "weighing it up</think>Paris.".to_string(),
8071 &[],
8072 closed,
8073 "stop",
8074 );
8075 assert_eq!(message.reasoning_content, None);
8076 }
8077
8078 #[test]
8079 fn stop_param_accepts_both_single_string_and_array() {
8080 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8081 "model": "m",
8082 "messages": [{"role": "user", "content": "hi"}],
8083 "stop": "END",
8084 }))
8085 .unwrap();
8086 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8087
8088 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8089 "model": "m",
8090 "messages": [{"role": "user", "content": "hi"}],
8091 "stop": ["A", "B"],
8092 }))
8093 .unwrap();
8094 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8095 }
8096
8097 #[test]
8098 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8099 let model = test_model();
8100 let result = run_generation(
8101 &model,
8102 "hello",
8103 &greedy_params(4),
8104 None,
8105 None,
8106 None,
8107 None,
8108 None,
8109 );
8110 assert!(matches!(
8111 result,
8112 Err(generate::DecodeError::TokenOutOfVocab { .. })
8113 ));
8114 }
8115
8116 #[test]
8120 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8121 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8123 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8124
8125 let holder_pool = Arc::clone(&pool);
8126 let holder = std::thread::spawn(move || {
8127 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8128 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8130 drop(held);
8131 });
8132 std::thread::sleep(Duration::from_millis(15));
8133
8134 let config = generate::KvPoolConfig {
8135 pool,
8136 queue_wait: Duration::ZERO,
8137 };
8138 let result = run_generation(
8139 &model,
8140 &prompt,
8141 &greedy_params(4),
8142 Some(&config),
8143 None,
8144 None,
8145 None,
8146 None,
8147 );
8148 assert!(matches!(
8149 result,
8150 Err(generate::DecodeError::KvPoolExhausted)
8151 ));
8152
8153 let (status, _body) = decode_error_response(result.unwrap_err());
8154 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8155 holder.join().unwrap();
8156 }
8157
8158 #[test]
8166 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8167 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8169 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
8171 let config = generate::KvPoolConfig {
8172 pool,
8173 queue_wait: Duration::ZERO,
8174 };
8175
8176 let result = run_generation(
8177 &model,
8178 &prompt,
8179 &greedy_params(4),
8180 Some(&config),
8181 None,
8182 None,
8183 None,
8184 None,
8185 );
8186 let err = result.expect_err("one block cannot hold two layers' caches");
8187 assert!(
8188 matches!(
8189 &err,
8190 generate::DecodeError::KvBudgetExceeded { binding, .. }
8191 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
8192 ),
8193 "expected an immovable device-memory refusal, got {err:?}"
8194 );
8195 let (status, _body) = decode_error_response(err);
8196 assert_eq!(status, StatusCode::BAD_REQUEST);
8197 }
8198
8199 #[test]
8205 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8206 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8207 queued: 512,
8208 cap: 512,
8209 });
8210 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8211 assert_eq!(body["error"]["retry_after_seconds"], 1);
8212 let message = body["error"]["message"].as_str().expect("message");
8213 assert!(message.contains("512"), "{message}");
8214 }
8215
8216 #[test]
8217 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8218 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8219 token: 99,
8220 vocab_size: 32,
8221 });
8222 assert!(
8223 body["error"]["retry_after_seconds"].is_null(),
8224 "retrying a prompt this model cannot tokenize never helps"
8225 );
8226 }
8227
8228 #[test]
8229 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8230 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8231 token: 99,
8232 vocab_size: 32,
8233 });
8234 assert_eq!(status, StatusCode::BAD_REQUEST);
8235 }
8236
8237 #[test]
8238 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8239 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8241 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8242 let config = generate::KvPoolConfig {
8243 pool: pool.clone(),
8244 queue_wait: Duration::ZERO,
8245 };
8246
8247 let (_, finish, _usage) = run_generation(
8248 &model,
8249 &prompt,
8250 &greedy_params(4),
8251 Some(&config),
8252 None,
8253 None,
8254 None,
8255 None,
8256 )
8257 .unwrap();
8258 assert_eq!(finish, FinishReason::Length);
8259 assert_eq!(
8260 pool.lock().unwrap().free_blocks(),
8261 2,
8262 "a completed request must return its blocks to the pool"
8263 );
8264 }
8265
8266 #[tokio::test]
8271 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8272 let model = Arc::new(test_model());
8273 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8274
8275 let mut handles = Vec::new();
8276 for _ in 0..8 {
8277 let model = Arc::clone(&model);
8278 let prompt = prompt.clone();
8279 handles.push(tokio::task::spawn_blocking(move || {
8280 run_generation(
8281 &model,
8282 &prompt,
8283 &greedy_params(6),
8284 None,
8285 None,
8286 None,
8287 None,
8288 None,
8289 )
8290 .unwrap()
8291 }));
8292 }
8293
8294 let mut results = Vec::new();
8295 for h in handles {
8296 results.push(h.await.unwrap());
8297 }
8298 for r in &results[1..] {
8303 assert_eq!(r.0, results[0].0, "decoded chunks must match");
8304 assert_eq!(r.1, results[0].1, "finish reason must match");
8305 assert_eq!(
8306 r.2.prompt_tokens, results[0].2.prompt_tokens,
8307 "prompt token count must match"
8308 );
8309 assert_eq!(
8310 r.2.completion_tokens, results[0].2.completion_tokens,
8311 "completion token count must match"
8312 );
8313 }
8314 }
8315
8316 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8323 let mut header_entries = Vec::new();
8324 let mut data = Vec::new();
8325 for (name, shape, values) in tensors {
8326 let start = data.len();
8327 for v in values {
8328 data.extend_from_slice(&v.to_le_bytes());
8329 }
8330 let end = data.len();
8331 let shape_str = shape
8332 .iter()
8333 .map(|d| d.to_string())
8334 .collect::<Vec<_>>()
8335 .join(",");
8336 header_entries.push(format!(
8337 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8338 ));
8339 }
8340 let header = format!("{{{}}}", header_entries.join(","));
8341 let header_bytes = header.as_bytes();
8342 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8343 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8344 out.extend_from_slice(header_bytes);
8345 out.extend_from_slice(&data);
8346 out
8347 }
8348
8349 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8360 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8361 use ferrox_models::kimi_loader::KimiRealHparams;
8362 use ferrox_moe::{GatingFunction, MoeLayerConfig};
8363
8364 let hidden_dim = 8;
8365 let kda_num_heads = 2;
8366 let kda_head_dim = 3;
8367 let kda_proj = kda_num_heads * kda_head_dim;
8368 let conv_kernel = 4;
8369 let dense_intermediate = 5;
8370 let vocab_size = 256;
8374 let mla_num_heads = 1;
8375 let mla_q_lora_rank = 2;
8376 let mla_kv_lora_rank = 2;
8377 let mla_qk_nope_head_dim = 2;
8378 let mla_qk_rope_head_dim = 2;
8379 let mla_v_head_dim = 2;
8380
8381 let model_cfg = ferrox_models::ModelConfig {
8382 name: "synthetic-kimi-server-test",
8383 n_layers: 1,
8384 hidden_dim,
8385 n_heads: 1,
8386 n_kv_heads: 1,
8387 head_dim: 4,
8388 vocab_size,
8389 rope_theta: 10000.0,
8390 rms_norm_eps: 1e-5,
8391 sliding_window: None,
8392 moe: MoeLayerConfig {
8393 expert_weights_scale: 1.0,
8394 n_experts: 1,
8395 n_experts_active: 1,
8396 n_shared_experts: 0,
8397 hidden_dim,
8398 expert_ffn_dim: 4,
8399 gating: GatingFunction::Sigmoid,
8400 norm_topk_prob: true,
8401 expert_group_count: None,
8402 expert_group_used_count: None,
8403 },
8404 n_dense_leading_layers: 1,
8409 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8410 kda_layers: vec![1],
8411 full_attn_layers: vec![],
8412 mla: MlaConfig {
8413 num_heads: mla_num_heads,
8414 q_lora_rank: mla_q_lora_rank,
8415 kv_lora_rank: mla_kv_lora_rank,
8416 qk_nope_head_dim: mla_qk_nope_head_dim,
8417 qk_rope_head_dim: mla_qk_rope_head_dim,
8418 v_head_dim: mla_v_head_dim,
8419 use_output_gate: true,
8420 rope: None,
8421 },
8422 kda: KdaConfig {
8423 num_heads: kda_num_heads,
8424 head_dim: kda_head_dim,
8425 short_conv_kernel_size: conv_kernel,
8426 gate_lower_bound: -5.0,
8427 use_full_rank_gate: true,
8428 },
8429 }),
8430 rope_freqs: None,
8431 rope_attn_factor: 1.0,
8432 rope_dim: None,
8433 rope_freqs_long: None,
8434 rope_freqs_short: None,
8435 rope_orig_ctx: None,
8436 rope_layout: ferrox_models::config::RopeLayout::Neox,
8437 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
8438 swa_pattern: None,
8439 swa_dense_first: false,
8440 attn_logit_softcap: None,
8441 final_logit_softcap: None,
8442 embedding_scale: None,
8443 attention_scale: None,
8444 rope_theta_swa: None,
8445 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
8446 best_effort_fields: &["synthetic test config, not a real preset"],
8447 };
8448 let hp = KimiRealHparams {
8449 hidden_dim,
8450 kda_num_heads,
8451 kda_head_dim,
8452 mla_num_heads,
8453 mla_q_lora_rank,
8454 mla_kv_lora_rank,
8455 mla_qk_nope_head_dim,
8456 mla_qk_rope_head_dim,
8457 mla_v_head_dim,
8458 dense_intermediate_dim: dense_intermediate,
8459 moe_hidden_dim: hidden_dim,
8460 moe_intermediate_dim: 4,
8461 n_experts: 1,
8462 num_shared_experts: 0,
8463 };
8464
8465 let prefix = "language_model.model.layers.0";
8469 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8470 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8471 name: String,
8472 shape: Vec<usize>,
8473 n: usize| {
8474 tensors.push((name, shape, vec![0.01f32; n]));
8475 };
8476 push(
8477 &mut tensors,
8478 format!("{prefix}.input_layernorm.weight"),
8479 vec![hidden_dim],
8480 hidden_dim,
8481 );
8482 push(
8483 &mut tensors,
8484 format!("{prefix}.post_attention_layernorm.weight"),
8485 vec![hidden_dim],
8486 hidden_dim,
8487 );
8488 push(
8489 &mut tensors,
8490 format!("{prefix}.self_attention_res_norm.weight"),
8491 vec![hidden_dim],
8492 hidden_dim,
8493 );
8494 push(
8495 &mut tensors,
8496 format!("{prefix}.self_attention_res_proj.weight"),
8497 vec![1, hidden_dim],
8498 hidden_dim,
8499 );
8500 push(
8501 &mut tensors,
8502 format!("{prefix}.mlp_res_norm.weight"),
8503 vec![hidden_dim],
8504 hidden_dim,
8505 );
8506 push(
8507 &mut tensors,
8508 format!("{prefix}.mlp_res_proj.weight"),
8509 vec![1, hidden_dim],
8510 hidden_dim,
8511 );
8512 push(
8513 &mut tensors,
8514 format!("{prefix}.self_attn.q_proj.weight"),
8515 vec![kda_proj, hidden_dim],
8516 kda_proj * hidden_dim,
8517 );
8518 push(
8519 &mut tensors,
8520 format!("{prefix}.self_attn.k_proj.weight"),
8521 vec![kda_proj, hidden_dim],
8522 kda_proj * hidden_dim,
8523 );
8524 push(
8525 &mut tensors,
8526 format!("{prefix}.self_attn.v_proj.weight"),
8527 vec![kda_proj, hidden_dim],
8528 kda_proj * hidden_dim,
8529 );
8530 push(
8531 &mut tensors,
8532 format!("{prefix}.self_attn.q_conv1d.weight"),
8533 vec![kda_proj, 1, conv_kernel],
8534 kda_proj * conv_kernel,
8535 );
8536 push(
8537 &mut tensors,
8538 format!("{prefix}.self_attn.k_conv1d.weight"),
8539 vec![kda_proj, 1, conv_kernel],
8540 kda_proj * conv_kernel,
8541 );
8542 push(
8543 &mut tensors,
8544 format!("{prefix}.self_attn.v_conv1d.weight"),
8545 vec![kda_proj, 1, conv_kernel],
8546 kda_proj * conv_kernel,
8547 );
8548 push(
8549 &mut tensors,
8550 format!("{prefix}.self_attn.A_log"),
8551 vec![kda_num_heads],
8552 kda_num_heads,
8553 );
8554 push(
8555 &mut tensors,
8556 format!("{prefix}.self_attn.f_a_proj.weight"),
8557 vec![kda_head_dim, hidden_dim],
8558 kda_head_dim * hidden_dim,
8559 );
8560 push(
8561 &mut tensors,
8562 format!("{prefix}.self_attn.f_b_proj.weight"),
8563 vec![kda_proj, kda_head_dim],
8564 kda_proj * kda_head_dim,
8565 );
8566 push(
8567 &mut tensors,
8568 format!("{prefix}.self_attn.dt_bias"),
8569 vec![kda_proj],
8570 kda_proj,
8571 );
8572 push(
8573 &mut tensors,
8574 format!("{prefix}.self_attn.b_proj.weight"),
8575 vec![kda_num_heads, hidden_dim],
8576 kda_num_heads * hidden_dim,
8577 );
8578 push(
8579 &mut tensors,
8580 format!("{prefix}.self_attn.g_proj.weight"),
8581 vec![kda_proj, hidden_dim],
8582 kda_proj * hidden_dim,
8583 );
8584 push(
8585 &mut tensors,
8586 format!("{prefix}.self_attn.o_norm.weight"),
8587 vec![kda_head_dim],
8588 kda_head_dim,
8589 );
8590 push(
8591 &mut tensors,
8592 format!("{prefix}.self_attn.o_proj.weight"),
8593 vec![hidden_dim, kda_proj],
8594 hidden_dim * kda_proj,
8595 );
8596 push(
8597 &mut tensors,
8598 format!("{prefix}.mlp.gate_proj.weight"),
8599 vec![dense_intermediate, hidden_dim],
8600 dense_intermediate * hidden_dim,
8601 );
8602 push(
8603 &mut tensors,
8604 format!("{prefix}.mlp.up_proj.weight"),
8605 vec![dense_intermediate, hidden_dim],
8606 dense_intermediate * hidden_dim,
8607 );
8608 push(
8609 &mut tensors,
8610 format!("{prefix}.mlp.down_proj.weight"),
8611 vec![hidden_dim, dense_intermediate],
8612 hidden_dim * dense_intermediate,
8613 );
8614 push(
8615 &mut tensors,
8616 "language_model.model.embed_tokens.weight".to_string(),
8617 vec![vocab_size, hidden_dim],
8618 vocab_size * hidden_dim,
8619 );
8620 push(
8621 &mut tensors,
8622 "language_model.lm_head.weight".to_string(),
8623 vec![vocab_size, hidden_dim],
8624 vocab_size * hidden_dim,
8625 );
8626 push(
8627 &mut tensors,
8628 "language_model.model.norm.weight".to_string(),
8629 vec![hidden_dim],
8630 hidden_dim,
8631 );
8632 push(
8633 &mut tensors,
8634 "language_model.model.output_attn_res_norm.weight".to_string(),
8635 vec![hidden_dim],
8636 hidden_dim,
8637 );
8638 push(
8639 &mut tensors,
8640 "language_model.model.output_attn_res_proj.weight".to_string(),
8641 vec![1, hidden_dim],
8642 hidden_dim,
8643 );
8644
8645 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
8658 let dir = std::env::temp_dir().join(format!(
8659 "ferrox_server_kimi_e2e_test_{}_{}",
8660 std::process::id(),
8661 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8662 ));
8663 std::fs::create_dir_all(&dir).unwrap();
8664 let shard_bytes = write_safetensors_shard(&tensors);
8665 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
8666 let map_entries: Vec<String> = tensors
8667 .iter()
8668 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
8669 .collect();
8670 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
8671 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
8672
8673 use base64::Engine;
8677 let vocab_lines: Vec<String> = (0..vocab_size as u32)
8678 .map(|b| {
8679 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
8680 format!("{b64} {b}")
8681 })
8682 .collect();
8683 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
8684
8685 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
8686 .expect("must load the synthetic Kimi checkpoint end to end");
8687 std::fs::remove_dir_all(&dir).ok();
8688 loaded
8689 }
8690
8691 #[test]
8699 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
8700 let loaded = build_synthetic_kimi_loaded();
8701 let state = build_app_state(
8702 StartupModels {
8703 loaded: model::LoadedModel::Kimi(loaded),
8704 embedding: None,
8705 },
8706 None,
8707 None,
8708 None,
8709 false,
8710 None,
8711 Arc::new(health::Detection::ready(health::probe_backends())),
8712 );
8713 let active = state.active().expect("a freshly built state has a model");
8714 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
8715 assert!(!active.is_synthetic());
8716
8717 let (_chunks, finish, _usage) = run_generation(
8718 active.generative().unwrap(),
8719 "hi",
8720 &greedy_params(5),
8721 None,
8722 None,
8723 None,
8724 None,
8725 None,
8726 )
8727 .expect("a real Kimi checkpoint must generate without error");
8728 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
8729 }
8730
8731 #[test]
8744 fn a_grammar_constrains_the_engine_decode_path() {
8745 let loaded = build_synthetic_kimi_loaded();
8746 let state = build_app_state(
8747 StartupModels {
8748 loaded: model::LoadedModel::Kimi(loaded),
8749 embedding: None,
8750 },
8751 None,
8752 None,
8753 None,
8754 false,
8755 None,
8756 Arc::new(health::Detection::ready(health::probe_backends())),
8757 );
8758 let active = state.active().expect("a freshly built state has a model");
8759
8760 let run = |grammar: Option<&str>| {
8761 let mut params = greedy_params(6);
8762 params.grammar = grammar.map(|src| {
8763 Arc::new(
8764 ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
8765 .expect("test grammar parses"),
8766 )
8767 });
8768 run_generation(
8769 active.generative().unwrap(),
8770 "hi",
8771 ¶ms,
8772 None,
8773 None,
8774 None,
8775 None,
8776 None,
8777 )
8778 };
8779
8780 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
8781 let unconstrained = chunks.concat();
8782 assert!(
8783 unconstrained.chars().any(|c| c != 'a'),
8784 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
8785 constrained run below would prove nothing"
8786 );
8787
8788 let (chunks, finish, _) =
8789 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
8790 let constrained = chunks.concat();
8791 assert!(
8792 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
8793 "the engine decode path served text its grammar forbids ({constrained:?}): \
8794 the constraint was dropped between `generate_engine` and the sampler"
8795 );
8796 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
8797 }
8798
8799 #[test]
8809 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
8810 let loaded = build_synthetic_kimi_loaded();
8811 let state = build_app_state(
8812 StartupModels {
8813 loaded: model::LoadedModel::Kimi(loaded),
8814 embedding: None,
8815 },
8816 None,
8817 None,
8818 None,
8819 false,
8820 None,
8821 Arc::new(health::Detection::ready(health::probe_backends())),
8822 );
8823
8824 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
8825 let kv_pool_config = generate::KvPoolConfig {
8826 pool: pool.clone(),
8827 queue_wait: Duration::ZERO,
8828 };
8829 let pc = Mutex::new(PrefixCache::new(4));
8830
8831 run_generation(
8832 state
8833 .active()
8834 .expect("a freshly built state has a model")
8835 .generative()
8836 .unwrap(),
8837 "hi",
8838 &greedy_params(5),
8839 Some(&kv_pool_config),
8840 None,
8841 Some(&pc),
8842 None,
8843 None,
8844 )
8845 .expect("a real Kimi checkpoint must generate without error");
8846
8847 assert_eq!(
8848 pool.lock().unwrap().free_blocks(),
8849 4,
8850 "the KV pool must be completely untouched by a Kimi request"
8851 );
8852 let stats = pc.lock().unwrap().stats();
8853 assert_eq!(
8854 stats.hits + stats.misses,
8855 0,
8856 "the prefix cache must never be consulted for a Kimi request"
8857 );
8858 }
8859}