1mod admin;
38mod anthropic;
39mod attribution;
40mod batch_scheduler;
41mod budget;
42mod cache;
43mod cache_admin;
44mod cancel;
45mod chat_template;
46mod generate;
47mod health;
48mod hub;
49mod journal;
50mod json_mode;
51mod limits;
52mod mcp;
53mod model;
54mod openai_extra;
55mod output;
56pub(crate) mod responses;
57mod resume;
58mod security;
59mod session;
60mod sse;
61mod stats;
62mod stop;
63mod stream_events;
64mod tasks;
65
66use std::cell::RefCell;
67use std::convert::Infallible;
68use std::fmt;
69use std::net::{IpAddr, Ipv4Addr, SocketAddr};
70use std::path::PathBuf;
71use std::rc::Rc;
72use std::str::FromStr;
73use std::sync::{Arc, Mutex, MutexGuard};
74use std::time::Duration;
75
76use axum::{
77 extract::State,
78 http::StatusCode,
79 response::sse::{Event, Sse},
80 response::{IntoResponse, Response},
81 routing::{get, post},
82 Json, Router,
83};
84use clap::{Parser, ValueEnum};
85use serde::{Deserialize, Serialize};
86
87use cache::{CacheKey, ResponseCache};
88use ferrox_core::cache::KvBlockPool;
89use ferrox_models::kimi_tokenizer::KimiTokenizer;
90use ferrox_models::sampling::SamplingParams;
91use ferrox_models::tokenizer::StopTokens;
92use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
93use generate::{FinishReason, GenerationParams};
94use model::ServerTokenizer;
95
96#[derive(Parser, Debug, PartialEq)]
100#[command(
107 name = "ferrox-server",
108 about = "OpenAI-compatible Ferrox inference server"
109)]
110pub struct ServerArgs {
111 #[arg(short = 'm', long = "model", value_name = "FILE")]
113 model: Option<String>,
114
115 #[arg(long, value_name = "HOST")]
117 host: Option<IpAddr>,
118
119 #[arg(long, value_name = "PORT")]
124 port: Option<u16>,
125
126 #[arg(short = 't', long = "threads", value_name = "N")]
128 threads: Option<usize>,
129
130 #[arg(
132 long = "device",
133 visible_alias = "dev",
134 value_name = "DEVICE",
135 ignore_case = true
136 )]
137 device: Option<OffloadDevice>,
138
139 #[arg(long = "list-devices", default_value_t = false)]
141 list_devices: bool,
142
143 #[arg(
148 long = "n-gpu-layers",
149 visible_aliases = ["gpu-layers", "ngl"],
150 value_name = "N"
151 )]
152 n_gpu_layers: Option<GpuLayers>,
153
154 #[arg(long = "mcp-config", value_name = "PATH")]
156 mcp_config: Option<PathBuf>,
157
158 #[arg(long = "exit-on-stdin-close", default_value_t = false)]
166 exit_on_stdin_close: bool,
167
168 #[arg(long = "allow-multiple-instances", default_value_t = false)]
173 allow_multiple_instances: bool,
174}
175
176impl ServerArgs {
177 pub fn parse_llama_style<I>(argv: I) -> Self
184 where
185 I: IntoIterator<Item = String>,
186 {
187 Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
188 }
189}
190
191#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
192enum OffloadDevice {
193 Auto,
194 None,
195 Cpu,
196 Metal,
197 Cuda,
198}
199
200#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201enum GpuLayers {
202 Auto,
203 All,
204 Count(u32),
205}
206
207impl GpuLayers {
208 fn offload_enabled(self) -> bool {
209 !matches!(self, Self::Count(0))
210 }
211}
212
213impl FromStr for GpuLayers {
214 type Err = String;
215
216 fn from_str(value: &str) -> Result<Self, Self::Err> {
217 match value {
218 "auto" => Ok(Self::Auto),
219 "all" => Ok(Self::All),
220 _ => value
221 .parse::<u32>()
222 .map(Self::Count)
223 .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
224 }
225 }
226}
227
228impl fmt::Display for GpuLayers {
229 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
230 match self {
231 Self::Auto => f.write_str("auto"),
232 Self::All => f.write_str("all"),
233 Self::Count(value) => value.fmt(f),
234 }
235 }
236}
237
238pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
247
248pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
251
252fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
253 args.into_iter()
254 .map(|arg| match arg.as_str() {
255 "-ngl" => "--n-gpu-layers".into(),
256 "-dev" => "--device".into(),
257 _ => arg,
258 })
259 .collect()
260}
261
262fn print_available_devices() {
263 println!("Available devices:");
264 println!(" CPU");
265
266 let metal = ferrox_metal::MetalProfile::detect();
267 if let Some(name) = metal.device_name {
268 println!(" Metal: {name}");
269 }
270
271 let cuda = ferrox_cuda::HardwareProfile::detect();
272 if cuda.cuda_available {
273 let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
274 println!(" CUDA: {name}");
275 if cuda.cuda_device_count > 1 {
276 println!(" ({} devices detected)", cuda.cuda_device_count);
277 }
278 }
279}
280
281fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
282 if args.host.is_none() && args.port.is_none() {
283 return None;
284 }
285
286 let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
287 let host = args
288 .host
289 .or_else(|| existing.map(|addr| addr.ip()))
290 .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
291 let port = args
292 .port
293 .or_else(|| existing.map(|addr| addr.port()))
294 .unwrap_or(8383);
295 Some(SocketAddr::new(host, port).to_string())
296}
297
298fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
299 if let Some(model) = &args.model {
300 unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
302 }
303
304 if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
305 unsafe { std::env::set_var("FERROX_ADDR", addr) };
307 }
308
309 if let Some(threads) = args.threads {
310 if threads == 0 {
311 anyhow::bail!("--threads must be greater than zero");
312 }
313 unsafe {
315 std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
316 std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
317 }
318 }
319
320 if args.device.is_none() && args.n_gpu_layers.is_none() {
321 } else {
323 let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
324 let device = if layers.offload_enabled() {
325 args.device.unwrap_or(OffloadDevice::Auto)
326 } else {
327 OffloadDevice::None
328 };
329
330 match device {
331 OffloadDevice::None | OffloadDevice::Cpu => unsafe {
332 std::env::set_var("FERROX_METAL", "0");
333 std::env::set_var("FERROX_METAL_ATTN", "0");
334 std::env::set_var("FERROX_CUDA", "0");
335 },
336 OffloadDevice::Auto => unsafe {
337 std::env::set_var("FERROX_METAL", "auto");
338 std::env::set_var("FERROX_CUDA", "auto");
339 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
340 std::env::set_var("FERROX_METAL_ATTN", "1");
341 }
342 },
343 OffloadDevice::Metal => {
344 #[cfg(not(feature = "metal"))]
345 {
346 anyhow::bail!(
347 "Metal requested but this binary was built without --features metal"
348 );
349 }
350 #[cfg(feature = "metal")]
351 {
352 if !ferrox_metal::MetalProfile::detect().available {
353 anyhow::bail!("Metal requested but no Metal device is available");
354 }
355 unsafe {
356 std::env::set_var("FERROX_METAL", "1");
357 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
358 std::env::set_var("FERROX_METAL_ATTN", "1");
359 }
360 std::env::set_var("FERROX_CUDA", "0");
361 }
362 }
363 }
364 OffloadDevice::Cuda => {
365 #[cfg(not(feature = "cuda"))]
366 {
367 anyhow::bail!(
368 "CUDA requested but this binary was built without --features cuda"
369 );
370 }
371 #[cfg(feature = "cuda")]
372 {
373 if !ferrox_cuda::HardwareProfile::detect().cuda_available {
374 anyhow::bail!("CUDA requested but no CUDA device is available");
375 }
376 unsafe {
377 std::env::set_var("FERROX_CUDA", "1");
378 std::env::set_var("FERROX_METAL", "0");
379 std::env::set_var("FERROX_METAL_ATTN", "0");
380 }
381 }
382 }
383 }
384 }
385
386 Ok(())
387}
388
389#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
398 Gguf(GgufModel),
399 Kimi(KimiModel),
400 Mla(MlaModel),
401 Gemma4(Gemma4Model),
402 Glm52(Glm52Model),
403}
404
405pub(crate) struct GgufModel {
406 decoder: Arc<Decoder>,
407 tokenizer: Arc<ServerTokenizer>,
408 stop_tokens: StopTokens,
409 bos_id: Option<usize>,
410 is_synthetic: bool,
411 chat_template: chat_template::PromptTemplate,
412}
413
414pub(crate) struct KimiModel {
415 engine: KimiEngine,
416 tokenizer: KimiTokenizer,
417 stop_tokens: StopTokens,
418 chat_template: chat_template::PromptTemplate,
419}
420
421pub(crate) struct MlaModel {
422 engine: MlaEngine,
423 tokenizer: ServerTokenizer,
424 stop_tokens: StopTokens,
425 bos_id: Option<usize>,
426 name: String,
427 chat_template: chat_template::PromptTemplate,
428}
429
430pub(crate) struct Gemma4Model {
431 engine: Gemma4Engine,
432 tokenizer: ServerTokenizer,
433 stop_tokens: StopTokens,
434 bos_id: Option<usize>,
435 name: String,
436 chat_template: chat_template::PromptTemplate,
437}
438
439pub(crate) struct Glm52Model {
440 engine: ferrox_models::Glm52Engine,
441 tokenizer: ServerTokenizer,
442 stop_tokens: StopTokens,
443 bos_id: Option<usize>,
444 name: String,
445 chat_template: chat_template::PromptTemplate,
446}
447
448impl Model {
449 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
450 match self {
451 Model::Gguf(m) => m.chat_template.clone(),
452 Model::Kimi(m) => m.chat_template.clone(),
453 Model::Mla(m) => m.chat_template.clone(),
454 Model::Gemma4(m) => m.chat_template.clone(),
455 Model::Glm52(m) => m.chat_template.clone(),
456 }
457 }
458
459 fn is_synthetic(&self) -> bool {
464 match self {
465 Model::Gguf(m) => m.is_synthetic,
466 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
467 }
468 }
469
470 fn tokenizer_kind(&self) -> &'static str {
471 match self {
472 Model::Gguf(m) => m.tokenizer.kind(),
473 Model::Kimi(_) => "kimi-tiktoken-bpe",
474 Model::Mla(m) => m.tokenizer.kind(),
475 Model::Gemma4(m) => m.tokenizer.kind(),
476 Model::Glm52(m) => m.tokenizer.kind(),
477 }
478 }
479
480 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
484 match self {
485 Model::Gguf(m) => m.decoder.expert_store_stats(),
486 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
487 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
488 }
489 }
490
491 pub(crate) fn name(&self) -> &str {
492 match self {
493 Model::Gguf(m) => m.decoder.config.name,
494 Model::Kimi(_) => "kimi-k3",
495 Model::Mla(m) => m.name.as_str(),
496 Model::Gemma4(m) => m.name.as_str(),
497 Model::Glm52(m) => m.name.as_str(),
498 }
499 }
500
501 pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
502 match self {
503 Model::Gguf(m) => m.tokenizer.encode(text),
504 Model::Kimi(m) => m
505 .tokenizer
506 .encode(text)
507 .into_iter()
508 .map(|id| id as usize)
509 .collect(),
510 Model::Mla(m) => m.tokenizer.encode(text),
511 Model::Gemma4(m) => m.tokenizer.encode(text),
512 Model::Glm52(m) => m.tokenizer.encode(text),
513 }
514 }
515
516 pub(crate) fn decode(&self, ids: &[usize]) -> String {
517 match self {
518 Model::Gguf(m) => m.tokenizer.decode(ids),
519 Model::Kimi(m) => {
520 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
521 m.tokenizer.decode(&ids32)
522 }
523 Model::Mla(m) => m.tokenizer.decode(ids),
524 Model::Gemma4(m) => m.tokenizer.decode(ids),
525 Model::Glm52(m) => m.tokenizer.decode(ids),
526 }
527 }
528
529 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
532 match self {
533 Model::Gguf(m) => {
534 let mut caches: Vec<_> = (0..m.decoder.layers.len())
535 .map(|_| {
536 ferrox_core::cache::KvCache::new(
537 m.decoder.config.n_kv_heads,
538 m.decoder.config.head_dim,
539 )
540 })
541 .collect();
542 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
543 }
544 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
545 }
546 }
547
548 pub(crate) fn vocab_size(&self) -> Option<usize> {
549 match self {
550 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
551 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
552 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
553 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
554 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
555 }
556 }
557}
558
559pub(crate) struct ActiveModel {
568 pub(crate) id: Option<String>,
572 pub(crate) model: Arc<Model>,
573 pub(crate) batcher: Option<batch_scheduler::ContinuousBatcher>,
578 pub(crate) ceiling: Option<Arc<budget::ContextCeiling>>,
588}
589
590pub(crate) struct AppState {
591 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
610 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
614 pub(crate) tasks: Arc<tasks::TaskRegistry>,
616 pub(crate) cancels: Arc<cancel::CancelRegistry>,
620 pub(crate) stats: stats::Stats,
623 pub(crate) streams: resume::StreamRegistry,
626 pub(crate) model_dir: Option<PathBuf>,
628 response_cache: Mutex<ResponseCache>,
632 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
644 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
657 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
667 sessions: session::SessionStore,
673 requests_total: std::sync::atomic::AtomicU64,
674 request_errors_total: std::sync::atomic::AtomicU64,
675 started_at: std::time::Instant,
676 last_request_ms: std::sync::atomic::AtomicU64,
681 detection: Arc<health::Detection>,
683 mcp: Option<mcp::LoadedMcpConfig>,
685 pub(crate) continuous_batching_enabled: bool,
689 loading_model: Mutex<Option<String>>,
693 last_load_error: Mutex<Option<(String, String)>>,
697 pub(crate) serving: Mutex<ferrox_edge::ServingStats>,
702 pub(crate) maintenance: Mutex<ferrox_edge::MaintenanceGate>,
707 pub(crate) footprint: Mutex<ferrox_edge::ProbeCache<ferrox_edge::Footprint>>,
713 pub(crate) started_unix: u64,
723}
724
725pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
731
732impl AppState {
733 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
737 self.active
738 .read()
739 .unwrap_or_else(|p| p.into_inner())
740 .clone()
741 }
742
743 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
749 self.active().ok_or_else(|| {
750 (
751 StatusCode::SERVICE_UNAVAILABLE,
752 Json(serde_json::json!({"error": {
753 "message": "no model is loaded; POST /admin/models/load with an id from \
754 GET /admin/models",
755 "type": "model_not_loaded"
756 }})),
757 )
758 })
759 }
760
761 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
764 Ok(Arc::clone(&self.require_active()?.model))
765 }
766
767 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
775 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
776 std::mem::replace(&mut *guard, next)
777 }
778
779 fn mark_request_finished(&self) {
783 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
784 self.last_request_ms
785 .store(ms, std::sync::atomic::Ordering::Relaxed);
786 }
787
788 pub(crate) fn uptime(&self) -> Duration {
789 self.started_at.elapsed()
790 }
791
792 pub(crate) fn requests_total(&self) -> u64 {
793 self.requests_total
794 .load(std::sync::atomic::Ordering::Relaxed)
795 }
796
797 pub(crate) fn errors_total(&self) -> u64 {
798 self.request_errors_total
799 .load(std::sync::atomic::Ordering::Relaxed)
800 }
801
802 pub(crate) fn cache_stats(&self) -> cache::CacheStats {
803 lock_cache(&self.response_cache).stats()
804 }
805
806 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
809 let last = self
810 .last_request_ms
811 .load(std::sync::atomic::Ordering::Relaxed);
812 (last > 0)
813 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
814 .map(|age| age.max(0.0))
815 }
816
817 pub(crate) fn loading_model_id(&self) -> Option<String> {
818 self.loading_model
819 .lock()
820 .unwrap_or_else(|p| p.into_inner())
821 .clone()
822 }
823
824 pub(crate) fn set_loading_model(&self, id: Option<String>) {
825 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
826 }
827
828 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
829 self.last_load_error
830 .lock()
831 .unwrap_or_else(|p| p.into_inner())
832 .clone()
833 }
834
835 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
836 *self
837 .last_load_error
838 .lock()
839 .unwrap_or_else(|p| p.into_inner()) = error;
840 }
841
842 pub(crate) fn active_model_name(&self) -> Option<String> {
851 self.active().map(|a| a.model.name().to_string())
852 }
853
854 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
855 self.stats.record(stats::entry(record));
856 }
857}
858
859fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
864 cache
865 .lock()
866 .unwrap_or_else(|poisoned| poisoned.into_inner())
867}
868
869#[derive(Debug, Clone, Deserialize)]
870#[serde(untagged)]
871pub(crate) enum MessageContent {
872 Text(String),
873 Parts(Vec<ContentPart>),
874}
875
876#[derive(Debug, Clone, Deserialize)]
877struct ContentPart {
878 #[serde(rename = "type")]
879 kind: String,
880 #[serde(default)]
881 text: Option<String>,
882 #[serde(default)]
883 image_url: Option<serde_json::Value>,
884}
885
886impl MessageContent {
887 fn as_text(&self) -> String {
888 match self {
889 Self::Text(s) => s.clone(),
890 Self::Parts(parts) => parts
891 .iter()
892 .filter_map(|p| p.text.as_deref())
893 .collect::<Vec<_>>()
894 .join(""),
895 }
896 }
897
898 fn has_image(&self) -> bool {
899 match self {
900 Self::Text(_) => false,
901 Self::Parts(parts) => parts
902 .iter()
903 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
904 }
905 }
906}
907
908#[derive(Debug, Clone, Deserialize)]
909pub(crate) struct ChatMessage {
910 pub(crate) role: String,
911 #[serde(default)]
915 pub(crate) content: Option<MessageContent>,
916 #[serde(default)]
920 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
921 #[serde(default)]
926 #[allow(dead_code)]
927 pub(crate) tool_call_id: Option<String>,
928 #[serde(default, alias = "reasoning")]
944 pub(crate) reasoning_content: Option<String>,
945}
946
947impl ChatMessage {
948 fn rendered_content(&self) -> String {
955 let mut out = self
956 .content
957 .as_ref()
958 .map(MessageContent::as_text)
959 .unwrap_or_default();
960 if let Some(calls) = &self.tool_calls {
961 for call in calls {
962 out.push_str(&format!(
963 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
964 call.function.name, call.function.arguments
965 ));
966 }
967 }
968 out
969 }
970}
971
972#[derive(Debug, Clone, Deserialize)]
973pub(crate) struct ToolCallIn {
974 #[serde(default)]
975 #[allow(dead_code)]
976 id: String,
977 #[serde(rename = "type", default)]
978 #[allow(dead_code)]
979 kind: String,
980 function: ToolCallFunctionIn,
981}
982
983#[derive(Debug, Clone, Deserialize)]
984struct ToolCallFunctionIn {
985 name: String,
986 arguments: String,
991}
992
993#[derive(Debug, Clone, Deserialize)]
996struct ToolDef {
997 #[serde(rename = "type", default)]
998 #[allow(dead_code)]
999 kind: String,
1000 function: ToolFunctionDef,
1001}
1002
1003#[derive(Debug, Clone, Deserialize)]
1004struct ToolFunctionDef {
1005 name: String,
1006 #[serde(default)]
1007 description: Option<String>,
1008 #[serde(default)]
1009 parameters: Option<serde_json::Value>,
1010}
1011
1012#[derive(Debug, Clone, Deserialize)]
1022#[serde(untagged)]
1023enum ToolChoice {
1024 Mode(String),
1025 #[allow(dead_code)]
1026 Specific(serde_json::Value),
1027}
1028
1029#[derive(Deserialize)]
1032#[serde(untagged)]
1033enum StopParam {
1034 One(String),
1035 Many(Vec<String>),
1036}
1037
1038#[derive(Deserialize)]
1039struct ChatCompletionRequest {
1040 model: String,
1041 messages: Vec<ChatMessage>,
1042 #[serde(default = "default_max_tokens")]
1043 max_tokens: usize,
1044 #[serde(default)]
1045 temperature: Option<f32>,
1046 #[serde(default)]
1047 top_p: Option<f32>,
1048 #[serde(default)]
1049 top_k: Option<usize>,
1050 #[serde(default)]
1051 repetition_penalty: Option<f32>,
1052 #[serde(default)]
1053 seed: Option<u64>,
1054 #[serde(default)]
1055 stop: Option<StopParam>,
1056 #[serde(default)]
1057 stream: Option<bool>,
1058 #[serde(default)]
1071 stream_resumable: Option<bool>,
1072 #[serde(default)]
1082 ignore_eos: Option<bool>,
1083 #[serde(default)]
1084 tools: Vec<ToolDef>,
1085 #[serde(default)]
1086 tool_choice: Option<ToolChoice>,
1087 #[serde(default)]
1095 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
1096 #[serde(default)]
1104 reasoning_effort: Option<String>,
1105 #[serde(default)]
1109 thinking: Option<ThinkingSwitch>,
1110 #[serde(default)]
1115 session_id: Option<String>,
1116 #[serde(default)]
1118 logprobs: Option<bool>,
1119 #[serde(default)]
1120 top_logprobs: Option<u32>,
1121 #[serde(default)]
1122 n: Option<u32>,
1123 #[serde(default)]
1124 presence_penalty: Option<f32>,
1125 #[serde(default)]
1126 frequency_penalty: Option<f32>,
1127 #[serde(default)]
1128 response_format: Option<serde_json::Value>,
1129}
1130
1131const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1143
1144#[derive(Debug, Clone, Deserialize)]
1146pub(crate) struct ThinkingSwitch {
1147 #[serde(rename = "type")]
1148 pub(crate) kind: String,
1149}
1150
1151const THINKING_KWARG_KEYS: [&str; 4] = [
1155 "enable_thinking",
1156 "thinking",
1157 "thinking_mode",
1158 "reasoning_effort",
1159];
1160
1161const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1165
1166fn default_max_tokens() -> usize {
1167 DEFAULT_CHAT_MAX_TOKENS
1168}
1169
1170impl ChatCompletionRequest {
1171 fn sampling_params(&self) -> SamplingParams {
1172 SamplingParams {
1173 temperature: self.temperature.unwrap_or(0.0),
1174 top_p: self.top_p.unwrap_or(1.0),
1175 top_k: self.top_k.unwrap_or(0),
1176 repetition_penalty: self.repetition_penalty.unwrap_or(1.0),
1177 presence_penalty: self.presence_penalty.unwrap_or(0.0),
1178 frequency_penalty: self.frequency_penalty.unwrap_or(0.0),
1179 }
1180 }
1181
1182 fn stop_sequences(&self) -> Vec<String> {
1183 self.stop
1184 .as_ref()
1185 .map(|s| match s {
1186 StopParam::One(v) => vec![v.clone()],
1187 StopParam::Many(v) => v.clone(),
1188 })
1189 .unwrap_or_default()
1190 }
1191
1192 fn tools_active(&self) -> bool {
1197 !self.tools.is_empty()
1198 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1199 }
1200
1201 fn thinking_direction(&self) -> Option<bool> {
1244 if let Some(switch) = &self.thinking {
1245 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1246 "disabled" => Some(false),
1247 "enabled" => Some(true),
1248 _ => None,
1251 };
1252 }
1253 let effort = self.reasoning_effort.as_ref()?;
1254 DISABLE_EFFORTS
1255 .contains(&effort.trim().to_ascii_lowercase().as_str())
1256 .then_some(false)
1257 }
1258
1259 fn resolve_template_kwargs(
1260 &self,
1261 template: &chat_template::PromptTemplate,
1262 ) -> serde_json::Map<String, serde_json::Value> {
1263 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1264 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1268
1269 if !caller_steered {
1270 match self.thinking_direction() {
1271 Some(false) => {
1272 for (k, v) in ferrox_edge::effort::thinking_off_kwargs() {
1273 kwargs.insert(k, v);
1274 }
1275 return kwargs;
1278 }
1279 Some(true) => {
1280 for (k, v) in ferrox_edge::effort::thinking_on_kwargs() {
1281 kwargs.insert(k, v);
1282 }
1283 }
1284 None => {}
1285 }
1286 if let Some(effort) = &self.reasoning_effort {
1287 kwargs
1288 .entry("reasoning_effort".to_string())
1289 .or_insert_with(|| serde_json::json!(effort));
1290 }
1291 }
1292
1293 let offered: Vec<serde_json::Value> = if self.tools_active() {
1294 self.tools.iter().map(chat_template::tool_json).collect()
1295 } else {
1296 Vec::new()
1297 };
1298 let thinking = ferrox_edge::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1299 if thinking == ferrox_edge::ThinkingMode::Thinking {
1300 for (k, v) in ferrox_edge::effort::thinking_on_kwargs() {
1301 kwargs.entry(k).or_insert(v);
1302 }
1303 }
1304 match ferrox_edge::sanitize_effort(&mut kwargs, template.efforts()) {
1305 ferrox_edge::EffortMapping::Mapped(to) => {
1306 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1307 }
1308 ferrox_edge::EffortMapping::Dropped => {
1309 tracing::debug!(
1310 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1311 enough, so its own default applies"
1312 );
1313 }
1314 ferrox_edge::EffortMapping::Unchanged => {}
1315 }
1316 ferrox_edge::broadcast_effort_spellings(&mut kwargs);
1317 kwargs
1318 }
1319
1320 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1323 if self.max_tokens == 0 {
1330 return Err(invalid_request(
1331 "max_tokens must be at least 1",
1332 "max_tokens",
1333 ));
1334 }
1335 if let Some(switch) = &self.thinking {
1340 let kind = switch.kind.trim().to_ascii_lowercase();
1341 if kind != "enabled" && kind != "disabled" {
1342 return Err(invalid_request(
1343 "thinking.type must be \"enabled\" or \"disabled\"",
1344 "thinking.type",
1345 ));
1346 }
1347 }
1348 for msg in &self.messages {
1349 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1350 return Err(unsupported_feature(
1351 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1352 ));
1353 }
1354 }
1355 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1356 return Err(unsupported_feature(
1357 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1358 ));
1359 }
1360 if self.n.is_some_and(|n| n > 1) {
1361 return Err(unsupported_feature(
1362 "n > 1 is not implemented (single completion only)",
1363 ));
1364 }
1365 if let Some(fmt) = &self.response_format {
1366 match fmt.get("type").and_then(|v| v.as_str()) {
1367 Some("json_object") => {}
1368 Some(other) => {
1369 return Err((
1370 StatusCode::BAD_REQUEST,
1371 Json(serde_json::json!({
1372 "error": {
1373 "message": format!(
1374 "response_format type {other:?} is not supported (only json_object)"
1375 )
1376 }
1377 })),
1378 ));
1379 }
1380 None => {
1381 return Err((
1382 StatusCode::BAD_REQUEST,
1383 Json(serde_json::json!({
1384 "error": {
1385 "message": "response_format must include \"type\" (only json_object is supported)"
1386 }
1387 })),
1388 ));
1389 }
1390 }
1391 }
1392 match &self.tool_choice {
1393 Some(ToolChoice::Mode(m)) if m == "required" => {
1394 return Err(unsupported_feature(
1395 "tool_choice=required needs constrained decoding (not implemented)",
1396 ));
1397 }
1398 Some(ToolChoice::Specific(_)) => {
1399 return Err(unsupported_feature(
1400 "named tool_choice is not implemented (use auto/none)",
1401 ));
1402 }
1403 _ => {}
1404 }
1405 Ok(())
1406 }
1407
1408 fn effective_stop_sequences(&self) -> Vec<String> {
1415 let mut stop = self.stop_sequences();
1416 if self.tools_active() {
1417 stop.push("</tool_call>".to_string());
1418 }
1419 stop
1420 }
1421
1422 fn json_object_mode(&self) -> bool {
1423 self.response_format
1424 .as_ref()
1425 .and_then(|v| v.get("type"))
1426 .and_then(|v| v.as_str())
1427 == Some("json_object")
1428 }
1429
1430 fn generation_params(&self) -> GenerationParams {
1431 GenerationParams {
1432 max_tokens: self.max_tokens,
1433 sampling: self.sampling_params(),
1434 seed: self.resolved_seed(),
1435 stop: self.effective_stop_sequences(),
1436 stop_token_ids: Vec::new(),
1440 json_object: self.json_object_mode(),
1441 cancel: None,
1444 ignore_eos: self.ignore_eos.unwrap_or(false),
1445 }
1446 }
1447
1448 fn generation_params_for_template(
1451 &self,
1452 template: &chat_template::PromptTemplate,
1453 ) -> GenerationParams {
1454 let mut params = self.generation_params();
1455 if let Some(stop) = template.end_of_turn() {
1456 if !params.stop.iter().any(|s| s == stop) {
1457 params.stop.push(stop.to_string());
1458 }
1459 }
1460 params
1461 }
1462
1463 fn is_cacheable(&self) -> bool {
1472 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1473 }
1474
1475 fn cache_key(&self, prompt: &str) -> CacheKey {
1476 CacheKey {
1477 model: self.model.clone(),
1478 prompt: prompt.to_string(),
1479 max_tokens: self.max_tokens,
1480 temperature_bits: self.temperature.unwrap_or(0.0).to_bits(),
1481 top_p_bits: self.top_p.unwrap_or(1.0).to_bits(),
1482 top_k: self.top_k.unwrap_or(0),
1483 repetition_penalty_bits: self.repetition_penalty.unwrap_or(1.0).to_bits(),
1484 presence_penalty_bits: self.presence_penalty.unwrap_or(0.0).to_bits(),
1485 frequency_penalty_bits: self.frequency_penalty.unwrap_or(0.0).to_bits(),
1486 seed: self.seed,
1487 stop: self.effective_stop_sequences(),
1488 }
1489 }
1490
1491 fn resolved_seed(&self) -> u64 {
1492 self.seed.unwrap_or_else(|| {
1493 std::time::SystemTime::now()
1494 .duration_since(std::time::UNIX_EPOCH)
1495 .map(|d| d.as_nanos() as u64)
1496 .unwrap_or(0xDEFA017)
1497 })
1498 }
1499}
1500
1501#[derive(Serialize)]
1502struct ChatCompletionChoice {
1503 index: usize,
1504 message: ChatCompletionResponseMessage,
1505 finish_reason: &'static str,
1506}
1507
1508#[derive(Serialize)]
1509struct ChatCompletionResponseMessage {
1510 role: &'static str,
1511 #[serde(skip_serializing_if = "Option::is_none")]
1512 content: Option<String>,
1513 #[serde(skip_serializing_if = "Option::is_none")]
1517 reasoning_content: Option<String>,
1518 #[serde(skip_serializing_if = "Option::is_none")]
1519 tool_calls: Option<Vec<ToolCallOut>>,
1520}
1521
1522#[derive(Serialize, Clone)]
1523struct ToolCallOut {
1524 id: String,
1525 #[serde(rename = "type")]
1526 kind: &'static str,
1527 function: ToolCallFunctionOut,
1528}
1529
1530#[derive(Serialize, Clone)]
1538struct ToolCallDelta {
1539 index: usize,
1540 #[serde(skip_serializing_if = "Option::is_none")]
1541 id: Option<String>,
1542 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1543 kind: Option<&'static str>,
1544 function: ToolCallFunctionDelta,
1545}
1546
1547#[derive(Serialize, Clone, Default)]
1548struct ToolCallFunctionDelta {
1549 #[serde(skip_serializing_if = "Option::is_none")]
1550 name: Option<String>,
1551 #[serde(skip_serializing_if = "Option::is_none")]
1554 arguments: Option<String>,
1555}
1556
1557impl ToolCallDelta {
1558 fn whole(index: usize, name: String, arguments: String) -> Self {
1560 ToolCallDelta {
1561 index,
1562 id: Some(format!("call_{index}")),
1563 kind: Some("function"),
1564 function: ToolCallFunctionDelta {
1565 name: Some(name),
1566 arguments: Some(arguments),
1567 },
1568 }
1569 }
1570
1571 fn opening(index: usize, name: String) -> Self {
1573 ToolCallDelta {
1574 index,
1575 id: Some(format!("call_{index}")),
1576 kind: Some("function"),
1577 function: ToolCallFunctionDelta {
1578 name: Some(name),
1579 arguments: Some(String::new()),
1580 },
1581 }
1582 }
1583
1584 fn arguments(index: usize, fragment: String) -> Self {
1586 ToolCallDelta {
1587 index,
1588 id: None,
1589 kind: None,
1590 function: ToolCallFunctionDelta {
1591 name: None,
1592 arguments: Some(fragment),
1593 },
1594 }
1595 }
1596}
1597
1598#[derive(Serialize, Clone)]
1599struct ToolCallFunctionOut {
1600 name: String,
1601 arguments: String,
1605}
1606
1607#[derive(Serialize)]
1608struct ChatCompletionResponse {
1609 id: String,
1610 request_id: String,
1616 object: &'static str,
1617 model: String,
1618 choices: Vec<ChatCompletionChoice>,
1619 usage: generate::Usage,
1624 ferrox_cache: &'static str,
1632}
1633
1634#[derive(Serialize)]
1635struct ChatCompletionChunkDelta {
1636 #[serde(skip_serializing_if = "Option::is_none")]
1637 role: Option<&'static str>,
1638 #[serde(skip_serializing_if = "Option::is_none")]
1639 content: Option<String>,
1640 #[serde(skip_serializing_if = "Option::is_none")]
1642 reasoning_content: Option<String>,
1643 #[serde(skip_serializing_if = "Option::is_none")]
1644 tool_calls: Option<Vec<ToolCallDelta>>,
1645}
1646
1647#[derive(Serialize)]
1648struct ChatCompletionChunkChoice {
1649 index: usize,
1650 delta: ChatCompletionChunkDelta,
1651 finish_reason: Option<&'static str>,
1652}
1653
1654#[derive(Serialize)]
1655struct ChatCompletionChunk {
1656 id: String,
1657 #[serde(skip_serializing_if = "Option::is_none")]
1664 request_id: Option<String>,
1665 object: &'static str,
1666 model: String,
1667 choices: Vec<ChatCompletionChunkChoice>,
1668 #[serde(skip_serializing_if = "Option::is_none")]
1671 usage: Option<generate::Usage>,
1672}
1673
1674async fn health(State(state): State<Arc<AppState>>) -> Response {
1680 let snapshot = state.detection.snapshot();
1681 let mut capabilities = snapshot.capabilities;
1682 let active = state.active();
1683
1684 capabilities.push(match active.as_deref() {
1687 None => ferrox_api::Capability::unavailable(
1692 ferrox_api::health::capability::REAL_WEIGHTS,
1693 ferrox_api::health::reason::MODEL_NOT_LOADED,
1694 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1695 ),
1696 Some(active) if active.model.is_synthetic() => ferrox_api::Capability::unavailable(
1697 ferrox_api::health::capability::REAL_WEIGHTS,
1698 ferrox_api::health::reason::MODEL_NOT_LOADED,
1699 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1700 checkpoint. Output from this model is noise.",
1701 ),
1702 Some(active) => ferrox_api::Capability::available(
1703 ferrox_api::health::capability::REAL_WEIGHTS,
1704 format!("Serving the real checkpoint '{}'.", active.model.name()),
1705 ),
1706 });
1707 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1708 ferrox_api::Capability::available(
1709 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1710 "Concurrent requests share one batched decode step.",
1711 )
1712 } else {
1713 ferrox_api::Capability::unavailable(
1714 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1715 ferrox_api::health::reason::DISABLED,
1716 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1717 )
1718 });
1719
1720 let last_request_ms = state
1721 .last_request_ms
1722 .load(std::sync::atomic::Ordering::Relaxed);
1723 let uptime = state.started_at.elapsed();
1724 let health_state = if active.is_none() {
1732 ferrox_api::HealthState::Unavailable
1733 } else {
1734 snapshot.state
1735 };
1736 let body = ferrox_api::HealthResponse {
1737 state: health_state,
1738 reason: match health_state {
1739 ferrox_api::HealthState::Ready => None,
1740 ferrox_api::HealthState::Unavailable => {
1741 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
1742 }
1743 ferrox_api::HealthState::Detecting => {
1744 Some(ferrox_api::health::reason::DETECTING.to_string())
1745 }
1746 },
1747 detail: match health_state {
1748 ferrox_api::HealthState::Ready => None,
1749 ferrox_api::HealthState::Unavailable => Some(
1750 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1751 .to_string(),
1752 ),
1753 ferrox_api::HealthState::Detecting => {
1754 Some("Probing available compute backends.".to_string())
1755 }
1756 },
1757 model: active
1758 .as_deref()
1759 .map(|active| ferrox_api::health::ModelSummary {
1760 id: active.model.name().to_string(),
1761 tokenizer: active.model.tokenizer_kind().to_string(),
1762 synthetic_weights: active.model.is_synthetic(),
1763 }),
1764 capabilities,
1765 version: env!("CARGO_PKG_VERSION").to_string(),
1766 pid: std::process::id(),
1767 uptime_seconds: uptime.as_secs_f64(),
1768 server_time_unix_ms: std::time::SystemTime::now()
1769 .duration_since(std::time::UNIX_EPOCH)
1770 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1771 .unwrap_or(0),
1772 last_request_age_seconds: (last_request_ms > 0)
1773 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1774 .map(|age| age.max(0.0)),
1775 };
1776
1777 let status =
1778 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1779 (status, Json(body)).into_response()
1780}
1781
1782async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1783 let Some(active) = state.active() else {
1787 return Json(serde_json::json!({ "object": "list", "data": [] }));
1788 };
1789 let mut model_entry = serde_json::json!({
1790 "id": active.model.name(),
1791 "object": "model",
1792 "ferrox_synthetic_weights": active.model.is_synthetic(),
1793 "ferrox_tokenizer": active.model.tokenizer_kind(),
1794 });
1795 let parser_configured = ferrox_edge::ReasoningFormat::infer(active.model.name()).is_some();
1801 let gears = active.model.chat_template().think_gears(parser_configured);
1802 if !gears.is_empty() {
1803 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1804 if let Some(default) = &gears.default {
1805 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1806 }
1807 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1810 }
1811 if let Some(mcp) = &state.mcp {
1812 model_entry["ferrox_mcp"] = mcp.models_metadata();
1813 }
1814 Json(serde_json::json!({
1815 "object": "list",
1816 "data": [model_entry]
1817 }))
1818}
1819
1820async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1834 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1835 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1836 let active = state.active();
1837 Json(serde_json::json!({
1838 "model": active.as_ref().map(|a| a.model.name()),
1839 "state": state
1840 .maintenance
1841 .lock()
1842 .unwrap_or_else(|p| p.into_inner())
1843 .state()
1844 .as_str(),
1845 "uptime_s": state.uptime().as_secs(),
1846 "throughput": {
1847 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1848 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1849 },
1850 "requests": {
1851 "active": state.cancels.live_count(),
1852 "completed": state.stats.recorded_total(),
1853 "p95_ms": state.stats.p95_duration_ms(),
1854 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1855 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1856 "completion_tokens_total": state.stats.tokens_generated_total(),
1857 },
1858 "pools": cache_admin::pool_gauges(&state),
1864 "memory": cache_admin::footprint_json(&state),
1867 }))
1868}
1869
1870#[derive(Deserialize)]
1871struct RequestsQuery {
1872 #[serde(default)]
1873 since: u64,
1874 #[serde(default = "default_requests_limit")]
1875 limit: usize,
1876}
1877
1878fn default_requests_limit() -> usize {
1879 stats::MAX_PAGE
1880}
1881
1882async fn recent_requests(
1890 State(state): State<Arc<AppState>>,
1891 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1892) -> Json<serde_json::Value> {
1893 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1894 Json(serde_json::json!({
1895 "requests": rows,
1896 "next_cursor": cursor,
1897 "missed": missed,
1898 "total": state.stats.recorded_total(),
1899 }))
1900}
1901
1902#[derive(Serialize)]
1903struct CombinedCacheStats {
1904 response_cache: cache::CacheStats,
1905 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
1907}
1908
1909async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1910 Json(CombinedCacheStats {
1911 response_cache: lock_cache(&state.response_cache).stats(),
1912 prefix_cache: state
1913 .prefix_cache
1914 .as_ref()
1915 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1916 })
1917}
1918
1919async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1924 use std::sync::atomic::Ordering;
1925
1926 let cache_stats = lock_cache(&state.response_cache).stats();
1927 let active = state.active();
1928 let requests_total = state.requests_total.load(Ordering::Relaxed);
1929 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1930 let uptime = state.started_at.elapsed().as_secs_f64();
1931
1932 let body = format!(
1933 "# HELP ferrox_requests_total Total chat completion requests received.\n\
1934 # TYPE ferrox_requests_total counter\n\
1935 ferrox_requests_total {requests_total}\n\
1936 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
1937 # TYPE ferrox_request_errors_total counter\n\
1938 ferrox_request_errors_total {errors_total}\n\
1939 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
1940 # TYPE ferrox_cache_hits_total counter\n\
1941 ferrox_cache_hits_total {}\n\
1942 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
1943 # TYPE ferrox_cache_misses_total counter\n\
1944 ferrox_cache_misses_total {}\n\
1945 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
1946 # TYPE ferrox_cache_entries gauge\n\
1947 ferrox_cache_entries {}\n\
1948 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1949 # TYPE ferrox_synthetic_weights gauge\n\
1950 ferrox_synthetic_weights {}\n\
1951 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
1952 # TYPE ferrox_uptime_seconds gauge\n\
1953 ferrox_uptime_seconds {uptime}\n",
1954 cache_stats.hits,
1955 cache_stats.misses,
1956 cache_stats.entries,
1957 active
1961 .as_ref()
1962 .map(|a| a.model.is_synthetic() as u8)
1963 .unwrap_or(0),
1964 );
1965
1966 let body = match active
1970 .as_ref()
1971 .and_then(|a| a.model.expert_store_stats())
1972 {
1973 Some(es) => format!(
1974 "{body}\
1975 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
1976 # TYPE ferrox_expert_cache_hits_total counter\n\
1977 ferrox_expert_cache_hits_total {}\n\
1978 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1979 # TYPE ferrox_expert_cache_misses_total counter\n\
1980 ferrox_expert_cache_misses_total {}\n\
1981 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
1982 # TYPE ferrox_expert_cache_evictions_total counter\n\
1983 ferrox_expert_cache_evictions_total {}\n\
1984 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1985 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
1986 ferrox_expert_cache_pass_throughs_total {}\n\
1987 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1988 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
1989 ferrox_expert_cache_bytes_read_total {}\n\
1990 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1991 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
1992 ferrox_expert_cache_resident_bytes {}\n",
1993 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1994 ),
1995 None => body,
1996 };
1997
1998 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2003 Some(batcher) => {
2004 let sched = batcher.stats();
2005 format!(
2006 "{body}\
2007 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2008 # TYPE ferrox_prefill_chunks_total counter\n\
2009 ferrox_prefill_chunks_total {}\n\
2010 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2011 # TYPE ferrox_prefill_tokens_total counter\n\
2012 ferrox_prefill_tokens_total {}\n\
2013 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2014 # TYPE ferrox_decode_steps_total counter\n\
2015 ferrox_decode_steps_total {}\n\
2016 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2017 # TYPE ferrox_scheduler_queue_depth gauge\n\
2018 ferrox_scheduler_queue_depth {}\n\
2019 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2020 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2021 ferrox_scheduler_queue_rejected_total {}\n\
2022 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2023 # TYPE ferrox_kv_blocks_total gauge\n\
2024 ferrox_kv_blocks_total {}\n\
2025 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2026 # TYPE ferrox_kv_blocks_free gauge\n\
2027 ferrox_kv_blocks_free {}\n\
2028 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2029 # TYPE ferrox_kv_block_size gauge\n\
2030 ferrox_kv_block_size {}\n\
2031 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2032 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2033 ferrox_kv_rejected_too_large_total {}\n\
2034 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2035 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2036 ferrox_kv_rejected_context_length_total {}\n\
2037 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2038 # TYPE ferrox_scheduler_aborted_total counter\n\
2039 ferrox_scheduler_aborted_total {}\n",
2040 sched.prefill_chunks,
2041 sched.prefill_tokens,
2042 sched.decode_steps,
2043 sched.queue_depth,
2044 sched.queue_rejected,
2045 sched.kv_blocks_total,
2046 sched.kv_blocks_free,
2047 sched.kv_block_size,
2048 sched.kv_rejected_too_large,
2049 sched.kv_rejected_context_length,
2050 sched.aborted,
2051 )
2052 }
2053 None => body,
2054 };
2055
2056 (
2057 [(
2058 axum::http::header::CONTENT_TYPE,
2059 "text/plain; version=0.0.4",
2060 )],
2061 body,
2062 )
2063 .into_response()
2064}
2065
2066pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2067
2068pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2074 (
2075 StatusCode::BAD_REQUEST,
2076 Json(serde_json::json!({"error": {
2077 "message": message,
2078 "type": "invalid_request_error",
2079 "param": param,
2080 "code": null,
2081 }})),
2082 )
2083}
2084
2085pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2086 (
2087 StatusCode::NOT_IMPLEMENTED,
2088 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2089 )
2090}
2091
2092pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2093 let status = match e {
2094 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2095 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2100 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2106 StatusCode::SERVICE_UNAVAILABLE
2107 }
2108 };
2109 tracing::warn!("decode error: {e}");
2110 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2111 if let generate::DecodeError::KvBudgetExceeded {
2117 binding,
2118 estimated_bytes,
2119 limit_bytes,
2120 positions,
2121 positions_limit,
2122 ..
2123 } = &e
2124 {
2125 body["error"]["type"] = serde_json::json!("invalid_request_error");
2126 body["error"]["code"] = serde_json::json!(binding);
2127 body["error"]["binding"] = serde_json::json!(binding);
2128 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2129 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2130 body["error"]["positions"] = serde_json::json!(positions);
2131 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2132 }
2133 if let Some(secs) = e.retry_after_secs() {
2137 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2138 }
2139 (status, Json(body))
2140}
2141
2142pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2143 tracing::error!("generation task panicked: {e}");
2144 (
2145 StatusCode::INTERNAL_SERVER_ERROR,
2146 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2147 )
2148}
2149
2150#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2160 model: &Model,
2161 prompt: &str,
2162 params: &GenerationParams,
2163 kv_pool: Option<&generate::KvPoolConfig>,
2164 paged_kv: Option<&generate::PagedKvConfig>,
2165 prefix_cache: Option<&Mutex<PrefixCache>>,
2166 continuous_batcher: Option<&batch_scheduler::ContinuousBatcher>,
2167 ceiling: Option<&budget::ContextCeiling>,
2168 mut emit: impl FnMut(&str),
2169) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2170 let synthetic = model.is_synthetic();
2171 let mut chunks = Vec::new();
2172 let params = &{
2178 let mut resolved = params.clone();
2179 resolved.stop_token_ids =
2180 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
2181 resolved
2182 };
2183 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2184 let (finish, usage) = match model {
2185 Model::Gguf(m) => {
2186 if let Some(batcher) = continuous_batcher {
2187 let mut tokens = m.tokenizer.encode(prompt);
2188 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2189 let (finish, _generated_ids, text, usage) =
2190 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?;
2191 if !text.is_empty() {
2192 chunks.push(text);
2193 }
2194 (finish, usage)
2195 } else {
2196 generate::generate(
2197 &m.decoder,
2198 m.tokenizer.as_ref(),
2199 &m.stop_tokens,
2200 m.bos_id,
2201 prompt,
2202 params,
2203 kv_pool,
2204 paged_kv,
2205 prefix_cache,
2206 ceiling,
2207 |chunk| {
2208 chunks.push(chunk.to_string());
2209 if !synthetic {
2210 emit(chunk);
2211 }
2212 },
2213 )?
2214 }
2215 }
2216 Model::Kimi(m) => generate::generate_engine(
2217 &m.engine,
2218 &m.tokenizer,
2219 &m.stop_tokens,
2220 None,
2221 prompt,
2222 params,
2223 |chunk| {
2224 chunks.push(chunk.to_string());
2225 if !synthetic {
2226 emit(chunk);
2227 }
2228 },
2229 )?,
2230 Model::Mla(m) => generate::generate_engine(
2231 &m.engine,
2232 &m.tokenizer,
2233 &m.stop_tokens,
2234 m.bos_id,
2235 prompt,
2236 params,
2237 |chunk| {
2238 chunks.push(chunk.to_string());
2239 if !synthetic {
2240 emit(chunk);
2241 }
2242 },
2243 )?,
2244 Model::Gemma4(m) => generate::generate_engine(
2245 &m.engine,
2246 &m.tokenizer,
2247 &m.stop_tokens,
2248 m.bos_id,
2249 prompt,
2250 params,
2251 |chunk| {
2252 chunks.push(chunk.to_string());
2253 if !synthetic {
2254 emit(chunk);
2255 }
2256 },
2257 )?,
2258 Model::Glm52(m) => generate::generate_engine(
2259 &m.engine,
2260 &m.tokenizer,
2261 &m.stop_tokens,
2262 m.bos_id,
2263 prompt,
2264 params,
2265 |chunk| {
2266 chunks.push(chunk.to_string());
2267 if !synthetic {
2268 emit(chunk);
2269 }
2270 },
2271 )?,
2272 };
2273
2274 let mut full = chunks.concat();
2275 if synthetic {
2276 full = format!(
2277 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2278 to serve a real model. Decoded ids -> {full:?}]"
2279 );
2280 emit(&full);
2281 } else if used_batcher && !full.is_empty() {
2282 emit(&full);
2283 }
2284
2285 Ok((finish, usage, full))
2286}
2287
2288#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2293 model: &Model,
2294 prompt: &str,
2295 params: &GenerationParams,
2296 kv_pool: Option<&generate::KvPoolConfig>,
2297 paged_kv: Option<&generate::PagedKvConfig>,
2298 prefix_cache: Option<&Mutex<PrefixCache>>,
2299 continuous_batcher: Option<&batch_scheduler::ContinuousBatcher>,
2300 ceiling: Option<&budget::ContextCeiling>,
2301) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2302 let (finish, usage, full) = run_generation_emit(
2303 model,
2304 prompt,
2305 params,
2306 kv_pool,
2307 paged_kv,
2308 prefix_cache,
2309 continuous_batcher,
2310 ceiling,
2311 |_| {},
2312 )?;
2313 Ok((
2314 if full.is_empty() {
2315 Vec::new()
2316 } else {
2317 vec![full]
2318 },
2319 finish,
2320 usage,
2321 ))
2322}
2323
2324pub(crate) fn prompt_from_messages(
2335 messages: &[ChatMessage],
2336 template: &chat_template::PromptTemplate,
2337 tools: &[ToolDef],
2338 extra: serde_json::Map<String, serde_json::Value>,
2339) -> Result<String, ApiError> {
2340 let rendered = if tools.is_empty() || template.handles_tools() {
2341 template.render(messages, tools, extra)
2342 } else {
2343 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2344 with_preamble.push(ChatMessage {
2345 role: "system".to_string(),
2346 content: Some(MessageContent::Text(tool_preamble(tools))),
2347 tool_calls: None,
2348 tool_call_id: None,
2349 reasoning_content: None,
2350 });
2351 with_preamble.extend_from_slice(messages);
2352 template.render(&with_preamble, &[], extra)
2353 };
2354 rendered.map_err(template_error_response)
2355}
2356
2357fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2362 (
2363 StatusCode::BAD_REQUEST,
2364 Json(serde_json::json!({
2365 "error": {
2366 "message": format!("chat template failed to render: {err}"),
2367 "type": "invalid_request_error",
2368 "param": "messages",
2369 "code": null,
2370 }
2371 })),
2372 )
2373}
2374
2375fn tool_preamble(tools: &[ToolDef]) -> String {
2388 let mut out = String::from(
2389 "You can call tools to help answer the user. To call a tool, respond with \
2390 EXACTLY one line in this format and nothing else:\n\
2391 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2392 object matching that tool's parameters>}}</tool_call>\n\n\
2393 Available tools:\n",
2394 );
2395 for t in tools {
2396 out.push_str(&format!(
2397 "- {}: {}\n parameters (JSON schema): {}\n",
2398 t.function.name,
2399 t.function.description.as_deref().unwrap_or(""),
2400 t.function
2401 .parameters
2402 .as_ref()
2403 .map(|v| v.to_string())
2404 .unwrap_or_else(|| "{}".to_string()),
2405 ));
2406 }
2407 out
2408}
2409
2410fn tool_call_deltas(
2420 events: Vec<ferrox_edge::ToolCallEvent>,
2421 opened: &std::cell::Cell<usize>,
2422) -> (String, Vec<ToolCallDelta>) {
2423 let mut text = String::new();
2424 let mut deltas = Vec::new();
2425 for event in events {
2426 match event {
2427 ferrox_edge::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2428 ferrox_edge::ToolCallEvent::CallStart { index, name } => {
2429 opened.set(opened.get().max(index + 1));
2430 deltas.push(ToolCallDelta::opening(index, name));
2431 }
2432 ferrox_edge::ToolCallEvent::CallArguments { index, fragment } => {
2433 if !fragment.is_empty() {
2434 deltas.push(ToolCallDelta::arguments(index, fragment));
2435 }
2436 }
2437 ferrox_edge::ToolCallEvent::CallEnd { .. } => {}
2438 }
2439 }
2440 (text, deltas)
2441}
2442
2443fn build_response_message(
2454 text: String,
2455 tools: &[ToolDef],
2456 posture: output::OutputPosture,
2457 base_finish: &'static str,
2458) -> (ChatCompletionResponseMessage, &'static str) {
2459 let parsed = output::parse_output(&text, tools, posture);
2460 let calls: Vec<ToolCallOut> = parsed
2461 .calls
2462 .into_iter()
2463 .enumerate()
2464 .map(|(index, call)| ToolCallOut {
2465 id: format!("call_{index}"),
2466 kind: "function",
2467 function: ToolCallFunctionOut {
2468 name: call.name,
2469 arguments: call.arguments,
2470 },
2471 })
2472 .collect();
2473 if !calls.is_empty() {
2474 return (
2475 ChatCompletionResponseMessage {
2476 role: "assistant",
2477 content: None,
2478 reasoning_content: parsed.reasoning,
2479 tool_calls: Some(calls),
2480 },
2481 "tool_calls",
2482 );
2483 }
2484 (
2485 ChatCompletionResponseMessage {
2486 role: "assistant",
2487 content: Some(parsed.content),
2488 reasoning_content: parsed.reasoning,
2489 tool_calls: None,
2490 },
2491 base_finish,
2492 )
2493}
2494
2495fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2500 let mut history = match &req.session_id {
2501 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2502 None => req.messages.clone(),
2503 };
2504 if req.json_object_mode() {
2505 inject_json_object_system_hint(&mut history);
2506 }
2507 history
2508}
2509
2510fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2511 const HINT: &str =
2512 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2513 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2514 match &mut sys.content {
2515 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2516 s.push_str("\n\n");
2517 s.push_str(HINT);
2518 }
2519 None => {
2520 sys.content = Some(MessageContent::Text(HINT.to_string()));
2521 }
2522 _ => {}
2523 }
2524 } else {
2525 messages.insert(
2526 0,
2527 ChatMessage {
2528 role: "system".to_string(),
2529 content: Some(MessageContent::Text(HINT.to_string())),
2530 tool_calls: None,
2531 tool_call_id: None,
2532 reasoning_content: None,
2533 },
2534 );
2535 }
2536}
2537
2538async fn chat_completions(
2539 State(state): State<Arc<AppState>>,
2540 headers: axum::http::HeaderMap,
2541 Json(req): Json<ChatCompletionRequest>,
2542) -> Response {
2543 let attribution = attribution::Attribution::from_headers(&headers);
2544 state
2545 .requests_total
2546 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2547 let started = std::time::Instant::now();
2548
2549 let request_id = ferrox_api::next_request_id();
2553 let stream = req.stream.unwrap_or(false);
2554
2555 let refusal = cache_admin::check_admission(&state)
2561 .err()
2562 .or_else(|| req.validate_supported_fields().err());
2563 if let Some(err) = refusal {
2564 state
2565 .request_errors_total
2566 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2567 let response = err.into_response();
2568 state.record_request(stats::Record {
2569 request_id: &request_id,
2570 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2571 model: state.active_model_name(),
2572 status: response.status().as_u16(),
2573 stream,
2574 duration_ms: started.elapsed().as_millis() as u64,
2575 usage: None,
2576 attribution: &attribution,
2577 });
2578 return response;
2579 }
2580
2581 let response = if stream {
2582 chat_completions_stream(
2583 Arc::clone(&state),
2584 req,
2585 request_id.clone(),
2586 started,
2587 attribution.clone(),
2588 )
2589 .await
2590 .into_response()
2591 } else {
2592 chat_completions_full(
2593 Arc::clone(&state),
2594 req,
2595 request_id.clone(),
2596 started,
2597 attribution.clone(),
2598 )
2599 .await
2600 .into_response()
2601 };
2602
2603 if response.status().is_client_error() || response.status().is_server_error() {
2604 state
2605 .request_errors_total
2606 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2607 state.record_request(stats::Record {
2611 request_id: &request_id,
2612 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2613 model: state.active_model_name(),
2616 status: response.status().as_u16(),
2617 stream,
2618 duration_ms: started.elapsed().as_millis() as u64,
2619 usage: None,
2620 attribution: &attribution,
2621 });
2622 }
2623 state.mark_request_finished();
2624
2625 response
2626}
2627
2628async fn chat_completions_full(
2629 state: Arc<AppState>,
2630 req: ChatCompletionRequest,
2631 request_id: String,
2632 started: std::time::Instant,
2633 attribution: attribution::Attribution,
2634) -> Result<Json<ChatCompletionResponse>, ApiError> {
2635 let tools_active = req.tools_active();
2636 let active = state.require_active()?;
2640 let history = resolve_history(&state, &req);
2641 let template = active.model.chat_template();
2642 let kwargs = req.resolve_template_kwargs(&template);
2643 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2644 let key = req.is_cacheable().then(|| req.cache_key(&prompt));
2645
2646 let (completion, cache_status) = if let Some(cached) = key
2647 .as_ref()
2648 .and_then(|key| lock_cache(&state.response_cache).get(key))
2649 {
2650 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2651 (cached, "hit")
2652 } else {
2653 let model = Arc::clone(&active.model);
2654 let kv_pool = state.kv_pool.clone();
2655 let paged_kv = state.paged_kv.clone();
2656 let prefix_cache = state.prefix_cache.clone();
2657 let batcher = active.batcher.clone();
2658 let ceiling = active.ceiling.clone();
2659 let params = req.generation_params_for_template(&template);
2660 let prompt_for_task = prompt.clone();
2661 let (chunks, finish, usage) = tokio::task::spawn_blocking(move || {
2662 run_generation(
2663 &model,
2664 &prompt_for_task,
2665 ¶ms,
2666 kv_pool.as_ref(),
2667 paged_kv.as_ref(),
2668 prefix_cache.as_deref(),
2669 batcher.as_ref(),
2670 ceiling.as_deref(),
2671 )
2672 })
2673 .await
2674 .map_err(join_error_response)?
2675 .map_err(decode_error_response)?;
2676
2677 let completion = cache::CachedCompletion {
2678 content: chunks.concat(),
2679 finish,
2680 usage,
2681 };
2682 let cache_status = if let Some(key) = key {
2683 tracing::debug!("cache miss for key {}", key.digest());
2684 lock_cache(&state.response_cache).put(key, completion.clone());
2685 "miss"
2686 } else {
2687 "skip"
2688 };
2689 (completion, cache_status)
2690 };
2691 let content = completion.content;
2692
2693 if req.json_object_mode() {
2694 json_mode::validate_json_object_output(&content)?;
2695 }
2696
2697 if let Some(id) = &req.session_id {
2701 state.sessions.store_reply(
2702 id,
2703 ChatMessage {
2704 role: "assistant".to_string(),
2705 content: Some(MessageContent::Text(content.clone())),
2706 tool_calls: None,
2707 tool_call_id: None,
2708 reasoning_content: None,
2709 },
2710 );
2711 }
2712
2713 let (message, finish_reason) = build_response_message(
2714 content,
2715 if tools_active { &req.tools } else { &[] },
2716 output::OutputPosture::resolve(active.model.name(), &prompt),
2717 completion.finish.as_str(),
2718 );
2719
2720 state.record_request(stats::Record {
2721 request_id: &request_id,
2722 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2723 model: Some(active.model.name().to_string()),
2726 status: 200,
2727 stream: false,
2728 duration_ms: started.elapsed().as_millis() as u64,
2729 usage: Some(&completion.usage),
2730 attribution: &attribution,
2731 });
2732
2733 Ok(Json(ChatCompletionResponse {
2734 id: request_id.clone(),
2735 request_id,
2736 object: "chat.completion",
2737 model: req.model,
2738 choices: vec![ChatCompletionChoice {
2739 index: 0,
2740 message,
2741 finish_reason,
2742 }],
2743 usage: completion.usage,
2744 ferrox_cache: cache_status,
2745 }))
2746}
2747
2748async fn chat_completions_stream(
2749 state: Arc<AppState>,
2750 req: ChatCompletionRequest,
2751 request_id: String,
2752 started: std::time::Instant,
2753 attribution: attribution::Attribution,
2754) -> Result<Response, ApiError> {
2755 let tools_active = req.tools_active();
2757 let active = state.require_active()?;
2761 let history = resolve_history(&state, &req);
2762 let template = active.model.chat_template();
2763 let kwargs = req.resolve_template_kwargs(&template);
2764 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2765 let model_name = req.model.clone();
2766 let session_id = req.session_id.clone();
2767 let sessions = state.sessions.clone();
2768
2769 let model = Arc::clone(&active.model);
2770 let kv_pool = state.kv_pool.clone();
2771 let paged_kv = state.paged_kv.clone();
2772 let prefix_cache = state.prefix_cache.clone();
2773 let batcher = active.batcher.clone();
2774 let ceiling = active.ceiling.clone();
2775 let mut params = req.generation_params_for_template(&template);
2776 let stats_state = Arc::clone(&state);
2777 let served_model = active.model.name().to_string();
2780 let posture = output::OutputPosture::resolve(&served_model, &prompt);
2784 let offered_tools: Vec<ToolDef> = if tools_active {
2787 req.tools.clone()
2788 } else {
2789 Vec::new()
2790 };
2791
2792 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2797 params.cancel = Some(cancel_token.clone());
2798
2799 let overlap = batcher.is_none();
2808
2809 let slot = req
2813 .stream_resumable
2814 .unwrap_or(false)
2815 .then(|| state.streams.register(&request_id));
2816 let emitter = resume::Emitter::new(slot);
2817
2818 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2819 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2823 id: request_id.clone(),
2824 request_id: None,
2825 object: "chat.completion.chunk",
2826 model: model_name.clone(),
2827 choices: vec![ChatCompletionChunkChoice {
2828 index: 0,
2829 delta: ChatCompletionChunkDelta {
2830 role: None,
2831 content: None,
2832 reasoning_content: None,
2833 tool_calls: None,
2834 },
2835 finish_reason: None,
2836 }],
2837 usage: None,
2838 });
2839
2840 tokio::task::spawn_blocking(move || {
2841 let _cancel_guard = cancel_guard;
2844 let tx_chunks = tx.clone();
2845 let orphan_timeout = sse::orphan_timeout_from_env();
2850 let mut first = true;
2851 let head_request_id = request_id.clone();
2852 let stream_reasoning: Rc<RefCell<Option<ferrox_edge::ReasoningParser>>> =
2862 Rc::new(RefCell::new(posture.reasoning_parser()));
2863 let emit_reasoning = Rc::clone(&stream_reasoning);
2864 let stream_tools: Rc<RefCell<Option<ferrox_edge::ToolCallParser>>> = Rc::new(RefCell::new(
2868 tools_active.then(|| posture.tool_call_parser(&offered_tools)),
2869 ));
2870 let emit_tools = Rc::clone(&stream_tools);
2871 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
2875 let emit_streamed_calls = Rc::clone(&streamed_calls);
2876 let result = run_generation_emit(
2877 &model,
2878 &prompt,
2879 ¶ms,
2880 kv_pool.as_ref(),
2881 paged_kv.as_ref(),
2882 prefix_cache.as_deref(),
2883 batcher.as_ref(),
2884 ceiling.as_deref(),
2885 |chunk| {
2886 if !overlap || chunk.is_empty() {
2887 return;
2888 }
2889 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
2890 Some(parser) => {
2891 let delta = parser.push(chunk);
2892 (delta.reasoning, delta.content)
2893 }
2894 None => (String::new(), chunk.to_string()),
2895 };
2896 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
2900 Some(parser) => {
2901 let (text, calls) =
2902 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
2903 (text, calls)
2904 }
2905 None => (content, Vec::new()),
2906 };
2907 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
2910 return;
2911 }
2912 let role = if first { Some("assistant") } else { None };
2913 let request_id = first.then(|| head_request_id.clone());
2914 first = false;
2915 let payload = ChatCompletionChunk {
2916 id: head_request_id.clone(),
2917 request_id,
2918 object: "chat.completion.chunk",
2919 model: model_name.clone(),
2920 choices: vec![ChatCompletionChunkChoice {
2921 index: 0,
2922 delta: ChatCompletionChunkDelta {
2923 role,
2924 content: (!content.is_empty()).then_some(content),
2925 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
2926 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
2927 },
2928 finish_reason: None,
2929 }],
2930 usage: None,
2931 };
2932 if let Err(why) =
2941 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
2942 {
2943 if why == sse::SendFailure::Orphaned {
2944 tracing::warn!(
2945 "SSE stream {head_request_id} accepted nothing for the orphan \
2946 deadline; treating it as abandoned"
2947 );
2948 }
2949 if !emitter.is_resumable() {
2961 cancel_token.cancel();
2962 }
2963 }
2964 },
2965 );
2966
2967 let mut pending_request_id = first.then(|| request_id.clone());
2973
2974 match result {
2975 Ok((finish, usage, full_text)) => {
2976 if let Some(id) = &session_id {
2977 sessions.store_reply(
2978 id,
2979 ChatMessage {
2980 role: "assistant".to_string(),
2981 content: Some(MessageContent::Text(full_text.clone())),
2982 tool_calls: None,
2983 tool_call_id: None,
2984 reasoning_content: None,
2985 },
2986 );
2987 }
2988 let mut streamed_finish: Option<&'static str> = None;
2994 if overlap {
2995 let tail = stream_reasoning
2996 .borrow_mut()
2997 .as_mut()
2998 .map(|parser| parser.flush())
2999 .unwrap_or_default();
3000 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3001 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3002 let mut events = parser.push(&content);
3003 events.extend(parser.finish());
3004 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3005 content = text;
3006 tool_calls = calls;
3007 }
3008 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3009 let payload = ChatCompletionChunk {
3010 id: request_id.clone(),
3011 request_id: pending_request_id.take(),
3012 object: "chat.completion.chunk",
3013 model: model_name.clone(),
3014 choices: vec![ChatCompletionChunkChoice {
3015 index: 0,
3016 delta: ChatCompletionChunkDelta {
3017 role: None,
3018 content: (!content.is_empty()).then_some(content),
3019 reasoning_content: (!tail.reasoning.is_empty())
3020 .then_some(tail.reasoning),
3021 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3022 },
3023 finish_reason: None,
3024 }],
3025 usage: None,
3026 };
3027 let _ =
3028 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3029 }
3030 if streamed_calls.get() > 0 {
3031 streamed_finish = Some("tool_calls");
3032 }
3033 } else {
3034 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3037 let tool_calls: Vec<ToolCallDelta> = parsed
3038 .calls
3039 .iter()
3040 .enumerate()
3041 .map(|(index, call)| {
3042 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3043 })
3044 .collect();
3045 if !tool_calls.is_empty() {
3046 streamed_finish = Some("tool_calls");
3047 }
3048 if !tool_calls.is_empty()
3049 || !parsed.content.is_empty()
3050 || parsed.reasoning.is_some()
3051 {
3052 let payload = ChatCompletionChunk {
3053 id: request_id.clone(),
3054 request_id: pending_request_id.take(),
3055 object: "chat.completion.chunk",
3056 model: model_name.clone(),
3057 choices: vec![ChatCompletionChunkChoice {
3058 index: 0,
3059 delta: ChatCompletionChunkDelta {
3060 role: Some("assistant"),
3061 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3062 .then(|| parsed.content.clone()),
3063 reasoning_content: parsed.reasoning.clone(),
3064 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3065 },
3066 finish_reason: None,
3067 }],
3068 usage: None,
3069 };
3070 let _ =
3071 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3072 }
3073 }
3074 let final_finish_reason = match streamed_finish {
3078 Some(reason) if finish.as_str() != "length" => reason,
3079 _ => finish.as_str(),
3080 };
3081 let final_payload = ChatCompletionChunk {
3082 id: request_id.clone(),
3083 request_id: pending_request_id.take(),
3084 object: "chat.completion.chunk",
3085 model: model_name,
3086 choices: vec![ChatCompletionChunkChoice {
3087 index: 0,
3088 delta: ChatCompletionChunkDelta {
3089 role: None,
3090 content: None,
3091 reasoning_content: None,
3092 tool_calls: None,
3093 },
3094 finish_reason: Some(final_finish_reason),
3095 }],
3096 usage: Some(usage.clone()),
3097 };
3098 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3099 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3100 stats_state.record_request(stats::Record {
3105 request_id: &request_id,
3106 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3107 model: Some(served_model.clone()),
3108 status: 200,
3109 stream: true,
3110 duration_ms: started.elapsed().as_millis() as u64,
3111 usage: Some(&usage),
3112 attribution: &attribution,
3113 });
3114 }
3115 Err(e) => {
3116 tracing::warn!("decode error on streamed request {request_id}: {e}");
3117 stats_state.record_request(stats::Record {
3123 request_id: &request_id,
3124 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3125 model: Some(served_model.clone()),
3126 status: 500,
3127 stream: true,
3128 duration_ms: started.elapsed().as_millis() as u64,
3129 usage: None,
3130 attribution: &attribution,
3131 });
3132 let payload = ChatCompletionChunk {
3133 id: request_id.clone(),
3134 request_id: pending_request_id.take(),
3135 object: "chat.completion.chunk",
3136 model: model_name,
3137 choices: vec![ChatCompletionChunkChoice {
3138 index: 0,
3139 delta: ChatCompletionChunkDelta {
3140 role: Some("assistant"),
3141 content: Some(format!("[error: {e}]")),
3142 reasoning_content: None,
3143 tool_calls: None,
3144 },
3145 finish_reason: Some("stop"),
3146 }],
3147 usage: None,
3148 };
3149 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3150 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3151 }
3152 }
3153 drop(emitter);
3157 });
3158
3159 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3160 Ok((
3183 [(
3184 axum::http::HeaderName::from_static("x-accel-buffering"),
3185 axum::http::HeaderValue::from_static("no"),
3186 )],
3187 Sse::new(stream),
3188 )
3189 .into_response())
3190}
3191
3192fn gpu_selected(value: Option<&str>, device_present: bool) -> bool {
3199 match value {
3200 Some("1") => true,
3201 Some("auto") => device_present,
3202 _ => false,
3203 }
3204}
3205
3206fn gpu_offload_resolved() -> bool {
3214 let metal_present = {
3215 #[cfg(feature = "metal")]
3216 {
3217 ferrox_metal::MetalProfile::detect().available
3218 }
3219 #[cfg(not(feature = "metal"))]
3220 {
3221 false
3222 }
3223 };
3224 let cuda_present = {
3225 #[cfg(feature = "cuda")]
3226 {
3227 ferrox_cuda::HardwareProfile::detect().cuda_available
3228 }
3229 #[cfg(not(feature = "cuda"))]
3230 {
3231 false
3232 }
3233 };
3234 gpu_selected(std::env::var("FERROX_METAL").ok().as_deref(), metal_present)
3235 || gpu_selected(std::env::var("FERROX_CUDA").ok().as_deref(), cuda_present)
3236}
3237
3238fn axum_path(template: &str) -> String {
3255 let mut out = String::with_capacity(template.len());
3256 let mut rest = template;
3257 while let Some(open) = rest.find('{') {
3258 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3259 break;
3260 };
3261 out.push_str(&rest[..open]);
3262 out.push(':');
3263 out.push_str(&rest[open + 1..close]);
3264 rest = &rest[close + 1..];
3265 }
3266 out.push_str(rest);
3267 out
3268}
3269
3270async fn cancel_generation(
3279 State(state): State<Arc<AppState>>,
3280 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3281) -> Response {
3282 let cancelled = state.cancels.cancel(&req.request_id);
3283 let status = if cancelled {
3284 StatusCode::OK
3285 } else {
3286 StatusCode::NOT_FOUND
3287 };
3288 let detail = if cancelled {
3289 "the generation was asked to stop; it ends at its next token".to_string()
3290 } else {
3291 "no generation with that request_id is running -- it has already \
3292 finished, was never issued, or was served by a path that does \
3293 not register for cancellation"
3294 .to_string()
3295 };
3296 (
3297 status,
3298 Json(ferrox_api::CancelGenerationResponse {
3299 request_id: req.request_id,
3300 cancelled,
3301 detail,
3302 }),
3303 )
3304 .into_response()
3305}
3306
3307type Activated = (
3311 Model,
3312 Option<batch_scheduler::ContinuousBatcher>,
3313 Option<Arc<budget::ContextCeiling>>,
3314);
3315
3316fn price_batcher_config(path: Option<&str>) -> batch_scheduler::BatcherConfig {
3334 let mut batcher = batch_scheduler::BatcherConfig::from_env();
3335 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3336 return batcher;
3339 }
3340 let Some(path) = path else {
3341 return batcher;
3342 };
3343 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
3349 let Some((priced, gguf_ctx, source)) = priced else {
3350 return batcher;
3351 };
3352 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3353 tracing::warn!(
3357 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3358 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3359 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3360 admit on a number you choose.",
3361 priced.weights_bytes,
3362 priced.device_budget_bytes,
3363 );
3364 return batcher;
3365 };
3366 tracing::info!("{source}");
3367 tracing::info!("{}", derived.fit);
3368 let adopted = budget::apply_derived(&mut batcher, &derived);
3369 if adopted.max_context {
3370 tracing::info!(
3371 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3372 override with FERROX_CB_MAX_CONTEXT",
3373 derived.max_context
3374 );
3375 }
3376 if adopted.kv_blocks {
3377 tracing::info!(
3378 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3379 derived.kv_blocks,
3380 batcher.kv_block_size
3381 );
3382 }
3383 batcher
3384}
3385
3386pub(crate) fn activate_loaded_model(
3395 loaded: model::LoadedModel,
3396 enable_continuous_batching: bool,
3397 path: Option<&str>,
3398) -> Activated {
3399 match loaded {
3400 model::LoadedModel::Gguf(g) => {
3401 let decoder = Arc::new(g.decoder);
3402 let tokenizer = Arc::new(g.tokenizer);
3403 let config = price_batcher_config(path);
3404 let shape =
3410 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
3411 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3412 let batcher = if enable_continuous_batching {
3413 tracing::info!(
3414 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3415 (stop sequences use the same pending-buffer trim as the private generate loop)"
3416 );
3417 let tok = Arc::clone(&tokenizer);
3418 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
3419 Some(batch_scheduler::ContinuousBatcher::spawn_with_ceiling(
3420 Arc::clone(&decoder),
3421 decode,
3422 config,
3423 Arc::clone(&ceiling),
3424 ))
3425 } else {
3426 None
3427 };
3428 (
3429 Model::Gguf(GgufModel {
3430 decoder,
3431 tokenizer,
3432 stop_tokens: g.stop_tokens,
3433 bos_id: g.bos_id,
3434 is_synthetic: g.is_synthetic,
3435 chat_template: g.chat_template,
3436 }),
3437 batcher,
3438 Some(ceiling),
3439 )
3440 }
3441 model::LoadedModel::Kimi(k) => (
3442 Model::Kimi(KimiModel {
3443 engine: k.engine,
3444 tokenizer: k.tokenizer,
3445 stop_tokens: k.stop_tokens,
3446 chat_template: k.chat_template,
3447 }),
3448 None,
3449 None,
3450 ),
3451 model::LoadedModel::Mla(m) => (
3452 Model::Mla(MlaModel {
3453 engine: m.engine,
3454 tokenizer: m.tokenizer,
3455 stop_tokens: m.stop_tokens,
3456 bos_id: m.bos_id,
3457 name: m.name,
3458 chat_template: m.chat_template,
3459 }),
3460 None,
3461 None,
3462 ),
3463 model::LoadedModel::Gemma4(m) => (
3464 Model::Gemma4(Gemma4Model {
3465 engine: m.engine,
3466 tokenizer: m.tokenizer,
3467 stop_tokens: m.stop_tokens,
3468 bos_id: m.bos_id,
3469 name: m.name,
3470 chat_template: m.chat_template,
3471 }),
3472 None,
3473 None,
3474 ),
3475 model::LoadedModel::Glm52(g) => (
3476 Model::Glm52(Glm52Model {
3477 engine: g.engine,
3478 tokenizer: g.tokenizer,
3479 stop_tokens: g.stop_tokens,
3480 bos_id: g.bos_id,
3481 name: g.name,
3482 chat_template: g.chat_template,
3483 }),
3484 None,
3485 None,
3486 ),
3487 }
3488}
3489
3490fn build_app_state(
3491 loaded: model::LoadedModel,
3492 kv_pool: Option<generate::KvPoolConfig>,
3493 paged_kv: Option<generate::PagedKvConfig>,
3494 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3495 enable_continuous_batching: bool,
3496 mcp: Option<mcp::LoadedMcpConfig>,
3497 detection: Arc<health::Detection>,
3498) -> AppState {
3499 let (model, batcher, ceiling) = activate_loaded_model(
3500 loaded,
3501 enable_continuous_batching,
3502 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
3503 );
3504 let id = startup_model_id();
3510 AppState {
3511 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3512 id,
3513 model: Arc::new(model),
3514 batcher,
3515 ceiling,
3516 }))),
3517 paged_kv,
3518 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3519 tasks: Arc::new(tasks::TaskRegistry::new()),
3520 cancels: Arc::new(cancel::CancelRegistry::new()),
3521 stats: stats::Stats::new(),
3522 streams: resume::StreamRegistry::new(),
3523 model_dir: admin::model_dirs().into_iter().next(),
3524 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3525 kv_pool,
3526 prefix_cache,
3527 sessions: session::SessionStore::new(),
3528 requests_total: std::sync::atomic::AtomicU64::new(0),
3529 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3530 started_at: std::time::Instant::now(),
3531 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3532 detection,
3533 mcp,
3534 continuous_batching_enabled: enable_continuous_batching,
3535 loading_model: Mutex::new(None),
3536 last_load_error: Mutex::new(None),
3537 serving: Mutex::new(ferrox_edge::ServingStats::default()),
3538 maintenance: Mutex::new(ferrox_edge::MaintenanceGate::serving()),
3539 footprint: Mutex::new(ferrox_edge::ProbeCache::new(FOOTPRINT_TTL_MS)),
3540 started_unix: unix_now(),
3541 }
3542}
3543
3544fn unix_now() -> u64 {
3549 std::time::SystemTime::now()
3550 .duration_since(std::time::UNIX_EPOCH)
3551 .map(|d| d.as_secs())
3552 .unwrap_or(0)
3553}
3554
3555fn startup_model_id() -> Option<String> {
3560 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
3561 let configured = std::fs::canonicalize(&configured).ok()?;
3562 admin::discover(&admin::model_dirs())
3563 .into_iter()
3564 .find(|d| {
3565 std::fs::canonicalize(&d.path)
3566 .map(|p| p == configured)
3567 .unwrap_or(false)
3568 })
3569 .map(|d| d.id)
3570}
3571
3572fn init_cpu_pool() {
3581 match ferrox_core::threads::init_cpu_pool() {
3582 Some(n) => eprintln!(
3583 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
3584 ferrox_core::threads::perf_core_count()
3585 ),
3586 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
3587 }
3588}
3589
3590fn announce_ready(addr: SocketAddr, scheme: &str) {
3603 use std::io::Write;
3604 let ready =
3605 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3606 let mut stdout = std::io::stdout().lock();
3607 let _ = writeln!(stdout, "{}", ready.to_line());
3608 let _ = stdout.flush();
3609}
3610
3611async fn shutdown_signal(exit_on_stdin_close: bool) {
3623 if !exit_on_stdin_close {
3624 std::future::pending::<()>().await;
3625 return;
3626 }
3627 let _ = tokio::task::spawn_blocking(|| {
3628 use std::io::Read;
3629 let mut sink = [0u8; 256];
3630 let mut stdin = std::io::stdin().lock();
3631 loop {
3632 match stdin.read(&mut sink) {
3633 Ok(0) => break,
3635 Ok(_) => continue,
3637 Err(e) => {
3638 tracing::warn!("stdin read failed ({e}); treating it as closed");
3639 break;
3640 }
3641 }
3642 }
3643 })
3644 .await;
3645 tracing::info!("stdin closed; shutting down");
3646}
3647
3648fn tokio_worker_threads() -> usize {
3653 std::env::var("FERROX_TOKIO_WORKERS")
3654 .ok()
3655 .and_then(|v| v.trim().parse::<usize>().ok())
3656 .filter(|n| *n > 0)
3657 .unwrap_or(2)
3658}
3659
3660fn install_ring_crypto_provider() {
3677 let _ = rustls::crypto::ring::default_provider().install_default();
3678}
3679
3680pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3686 if args.list_devices {
3687 print_available_devices();
3688 return Ok(());
3689 }
3690 apply_cli_overrides(&args)?;
3691
3692 let _instance = {
3696 use ferrox_core::instance::{register, InstancePolicy};
3697 let policy = if args.allow_multiple_instances {
3698 InstancePolicy::Multi
3699 } else {
3700 InstancePolicy::from_env_or(InstancePolicy::Single)
3701 };
3702 let model = std::env::var("FERROX_MODEL_PATH").ok();
3703 register(
3704 "server",
3705 model.as_deref(),
3706 ferrox_core::instance::current_backend(),
3707 policy,
3708 )
3709 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
3710 };
3711
3712 let journal = journal::Journal::from_env();
3713 eprintln!(
3714 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
3715 journal.path()
3716 );
3717 journal.append(&journal::Record::session_start(
3718 env!("CARGO_PKG_VERSION"),
3719 std::process::id(),
3720 ));
3721 journal::install_panic_hook(journal.clone());
3722
3723 let mcp_config_path = args.mcp_config.clone();
3724 let exit_on_stdin_close = args.exit_on_stdin_close
3725 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
3726 .map(|v| v == "1")
3727 .unwrap_or(false);
3728
3729 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
3733 init_cpu_pool();
3734
3735 let runtime = tokio::runtime::Builder::new_multi_thread()
3736 .worker_threads(tokio_worker_threads())
3737 .enable_all()
3738 .build()?;
3739 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
3740
3741 let reason = match &result {
3742 Ok(()) => "normal".to_string(),
3743 Err(e) => e.to_string(),
3744 };
3745 journal.append(&journal::Record::session_exit(reason));
3746
3747 runtime.shutdown_background();
3752
3753 result
3754}
3755
3756async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
3757 let _ = tracing_subscriber::fmt::try_init();
3763
3764 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
3774 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
3775 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
3776 .map(|v| v == "1")
3777 .unwrap_or(false);
3778 if let Err(msg) =
3779 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
3780 {
3781 anyhow::bail!(msg);
3782 }
3783
3784 let mut loaded = model::load()?;
3785 match &loaded {
3786 model::LoadedModel::Gguf(g) => tracing::info!(
3787 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
3788 g.decoder.config.name,
3789 g.is_synthetic,
3790 g.tokenizer.kind()
3791 ),
3792 model::LoadedModel::Kimi(k) => tracing::info!(
3793 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
3794 k.tokenizer.vocab_size()
3795 ),
3796 model::LoadedModel::Mla(m) => tracing::info!(
3797 "loaded MLA GGUF '{}' (tokenizer={})",
3798 m.name,
3799 m.tokenizer.kind()
3800 ),
3801 model::LoadedModel::Gemma4(m) => tracing::info!(
3802 "loaded Gemma4 GGUF '{}' (tokenizer={})",
3803 m.name,
3804 m.tokenizer.kind()
3805 ),
3806 model::LoadedModel::Glm52(g) => tracing::info!(
3807 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
3808 g.name,
3809 g.tokenizer.kind()
3810 ),
3811 }
3812 let metal_default_moe_budget = {
3820 #[cfg(feature = "metal")]
3821 {
3822 ferrox_core::metal_dense_enabled()
3823 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
3824 }
3825 #[cfg(not(feature = "metal"))]
3826 {
3827 false
3828 }
3829 };
3830 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
3831 let budget: u64 = budget_str
3832 .parse()
3833 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
3834 match &mut loaded {
3835 model::LoadedModel::Gguf(g) => {
3836 tracing::info!(
3837 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
3838 (CUDA and/or Metal matvecs when built with the matching feature)"
3839 );
3840 g.decoder.gpu_vram_budget_bytes = Some(budget);
3841 }
3842 model::LoadedModel::Kimi(_) => {
3843 tracing::warn!(
3844 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
3845 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
3846 );
3847 }
3848 model::LoadedModel::Mla(_) => {
3849 tracing::warn!(
3850 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
3851 FFN path only today; ignoring expert VRAM budget"
3852 );
3853 }
3854 model::LoadedModel::Gemma4(_) => {
3855 tracing::warn!(
3856 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
3857 ignoring expert VRAM budget"
3858 );
3859 }
3860 model::LoadedModel::Glm52(_) => {
3861 tracing::warn!(
3862 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
3863 GPU expert placement not wired yet; ignoring"
3864 );
3865 }
3866 }
3867 } else if metal_default_moe_budget {
3868 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
3872 if let model::LoadedModel::Gguf(g) = &mut loaded {
3873 tracing::info!(
3874 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
3875 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
3876 );
3877 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
3878 }
3879 }
3880 #[cfg(feature = "cuda")]
3881 {
3882 if ferrox_core::cuda_dense_enabled() {
3883 tracing::info!(
3884 "CUDA dense matvec enabled for WeightMatrix::apply \
3885 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
3886 );
3887 } else {
3888 tracing::info!(
3889 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
3890 );
3891 }
3892 }
3893 #[cfg(feature = "metal")]
3894 {
3895 if ferrox_core::metal_dense_enabled() {
3896 tracing::info!(
3897 "Metal dense matvec enabled for WeightMatrix::apply \
3898 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
3899 );
3900 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
3901 Some("1") | Some("true") | Some("on") | Some("attn") => {
3902 tracing::info!(
3903 "Metal fused attention requested (FERROX_METAL_ATTN): \
3904 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
3905 );
3906 }
3907 _ => {}
3908 }
3909 match std::env::var("FERROX_METAL_LOGITS").ok().as_deref() {
3910 Some("1") | Some("true") | Some("on") | Some("logits") => {
3911 tracing::info!(
3912 "Metal logits-in-stack enabled (FERROX_METAL_LOGITS): \
3913 final_norm+lm_head on-GPU (often slower than host lm_head)"
3914 );
3915 }
3916 _ => {}
3917 }
3918 match std::env::var("FERROX_METAL_GREEDY_GPU").ok().as_deref() {
3919 Some("0") | Some("false") | Some("off") | Some("no") => {
3920 tracing::info!(
3921 "Metal greedy GPU argmax disabled (FERROX_METAL_GREEDY_GPU=0); \
3922 temperature<=0 uses host lm_head"
3923 );
3924 }
3925 _ => {
3926 tracing::info!(
3927 "Metal greedy GPU argmax on by default (opt out FERROX_METAL_GREEDY_GPU=0): \
3928 temperature<=0 folds final_norm+lm_head+argmax into dense stack"
3929 );
3930 }
3931 }
3932 } else {
3933 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
3934 }
3935 }
3936 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
3947 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
3948 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
3949 if blocks_env.is_ok() && byte_budget_env.is_ok() {
3950 panic!(
3951 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
3952 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
3953 );
3954 }
3955 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
3956 (Ok(blocks), Ok(block_size), Err(_)) => {
3957 let total_blocks: usize = blocks
3958 .parse()
3959 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
3960 let block_size: usize = block_size
3961 .parse()
3962 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
3963 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
3969 .ok()
3970 .map(|v| {
3971 v.parse()
3972 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
3973 })
3974 .unwrap_or(0);
3975 tracing::info!(
3976 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
3977 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
3978 );
3979 Some(generate::KvPoolConfig {
3980 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
3981 queue_wait: Duration::from_millis(queue_wait_ms),
3982 })
3983 }
3984 (Err(_), Ok(block_size), Ok(byte_budget)) => {
3985 let block_size: usize = block_size
3986 .parse()
3987 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
3988 let budget: u64 = byte_budget
3989 .parse()
3990 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
3991 let cfg = match &loaded {
3992 model::LoadedModel::Gguf(g) => &g.decoder.config,
3993 model::LoadedModel::Kimi(_)
3994 | model::LoadedModel::Mla(_)
3995 | model::LoadedModel::Gemma4(_)
3996 | model::LoadedModel::Glm52(_) => {
3997 panic!(
3998 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
3999 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4000 );
4001 }
4002 };
4003 let bytes_per_block = block_size
4004 * cfg.n_layers
4005 * cfg.n_kv_heads
4006 * cfg.head_dim
4007 * 2
4008 * std::mem::size_of::<f32>();
4009 assert!(
4010 bytes_per_block > 0,
4011 "derived KV block byte size must be positive (check model config and block size)"
4012 );
4013 let total_blocks = (budget as usize / bytes_per_block).max(1);
4014 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4015 .ok()
4016 .map(|v| {
4017 v.parse()
4018 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4019 })
4020 .unwrap_or(0);
4021 tracing::info!(
4022 "KV cache block pool enabled from byte budget: {budget} bytes / \
4023 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4024 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4025 cfg.n_layers
4026 );
4027 Some(generate::KvPoolConfig {
4028 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4029 queue_wait: Duration::from_millis(queue_wait_ms),
4030 })
4031 }
4032 (Err(_), Err(_), Err(_)) => None,
4033 (Err(_), Ok(_), Err(_)) => panic!(
4034 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4035 (or unset all three to disable KV cache pooling)"
4036 ),
4037 (Ok(_), Ok(_), Ok(_)) => {
4038 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4039 }
4040 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4041 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4042 set together (or neither, to disable KV cache pooling)"
4043 ),
4044 };
4045 let paged_kv = match (
4051 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4052 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4053 ) {
4054 (Ok(blocks), Ok(block_size)) => {
4055 assert!(
4056 kv_pool.is_none(),
4057 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4058 mutually exclusive: both bound the same KV memory, by different means. \
4059 Set one."
4060 );
4061 let gpu_offload = gpu_offload_resolved();
4075 assert!(
4076 !gpu_offload,
4077 "FERROX_PAGED_KV_BLOCKS is CPU-only right now: on a GPU \
4078 backend the paged attention path returns wrong tokens \
4079 rather than failing, so it is refused instead. Run the \
4080 server with `-dev none -ngl 0` to use paged KV, or unset \
4081 FERROX_PAGED_KV_BLOCKS to use the GPU."
4082 );
4083 let blocks_per_layer: usize = blocks
4084 .parse()
4085 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4086 let block_size: usize = block_size
4087 .parse()
4088 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4089 let cfg = match &loaded {
4090 model::LoadedModel::Gguf(g) => &g.decoder.config,
4091 _ => panic!(
4092 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4093 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4094 ),
4095 };
4096 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4097 .ok()
4098 .map(|v| {
4099 v.parse()
4100 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4101 })
4102 .unwrap_or(0);
4103 tracing::info!(
4104 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4105 layer across {} layers, shared by all concurrent requests, \
4106 {queue_wait_ms}ms admission queue wait",
4107 cfg.n_layers
4108 );
4109 let radix = Some(Arc::new(Mutex::new(ferrox_edge::radix::RadixCache::new(
4114 block_size,
4115 ))));
4116 Some(generate::PagedKvConfig {
4117 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4118 cfg.n_layers,
4119 block_size,
4120 blocks_per_layer,
4121 cfg.n_kv_heads,
4122 cfg.head_dim,
4123 )),
4124 queue_wait: Duration::from_millis(queue_wait_ms),
4125 radix,
4126 })
4127 }
4128 (Err(_), Err(_)) => None,
4129 _ => panic!(
4130 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4131 (or neither, to disable paged KV)"
4132 ),
4133 };
4134 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4141 let max_entries: usize = v
4142 .parse()
4143 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4144 if kv_pool.is_some() {
4145 tracing::warn!(
4146 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4147 caching will never be consulted while a KV pool is configured"
4148 );
4149 }
4150 assert!(
4157 paged_kv.is_none(),
4158 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4159 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4160 produce, so the cache could never hit. Set one."
4161 );
4162 tracing::info!(
4163 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4164 all requests"
4165 );
4166 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4167 });
4168 if matches!(
4169 loaded,
4170 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4171 ) && (kv_pool.is_some() || prefix_cache.is_some())
4172 {
4173 tracing::warn!(
4174 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4175 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4176 ferrox_models::engine's module docs"
4177 );
4178 }
4179 let enable_cb = std::env::var("FERROX_CONTINUOUS_BATCHING")
4180 .map(|v| v == "1")
4181 .unwrap_or(false)
4182 && kv_pool.is_none()
4183 && prefix_cache.is_none()
4184 && matches!(loaded, model::LoadedModel::Gguf(_));
4185 if std::env::var("FERROX_CONTINUOUS_BATCHING")
4186 .map(|v| v == "1")
4187 .unwrap_or(false)
4188 && (kv_pool.is_some() || prefix_cache.is_some())
4189 {
4190 tracing::warn!(
4191 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4192 (those modes keep the private generate path)"
4193 );
4194 }
4195 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4196 if let Ok(chunk) = n.parse::<usize>() {
4197 if chunk > 0 {
4198 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4199 }
4200 }
4201 }
4202 if matches!(
4203 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4204 Some("1")
4205 ) {
4206 tracing::warn!(
4207 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4208 (minimal spill; full layer offload still planned)"
4209 );
4210 }
4211
4212 let mcp = match mcp_config_path {
4213 Some(path) => {
4214 let loaded = mcp::load_mcp_config(&path)?;
4215 tracing::info!(
4216 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4217 loaded.path,
4218 loaded.servers.len()
4219 );
4220 Some(loaded)
4221 }
4222 None => None,
4223 };
4224
4225 let detection = health::Detection::spawn();
4229
4230 let state = Arc::new(build_app_state(
4231 loaded,
4232 kv_pool,
4233 paged_kv,
4234 prefix_cache,
4235 enable_cb,
4236 mcp,
4237 detection,
4238 ));
4239
4240 use ferrox_api::routes;
4244
4245 let public = Router::new().route(routes::HEALTH, get(health));
4250
4251 let mut protected = Router::new()
4252 .route(routes::V1_MODELS, get(list_models))
4253 .route(routes::V1_RESPONSES, post(responses::responses))
4257 .route(
4258 &axum_path(routes::V1_RESPONSE),
4259 get(responses::responses_get),
4260 )
4261 .route(
4262 &axum_path(routes::V1_RESPONSE_CANCEL),
4263 post(responses::responses_cancel),
4264 )
4265 .route(routes::V1_STATS, get(serving_stats))
4266 .route(routes::V1_REQUESTS, get(recent_requests))
4267 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
4268 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
4269 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
4270 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
4271 .route(routes::V1_CANCEL, post(cancel_generation))
4275 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
4280 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
4281 .route(routes::V1_MESSAGES, post(anthropic::messages))
4282 .route(
4283 routes::V1_MESSAGES_COUNT_TOKENS,
4284 post(anthropic::count_tokens),
4285 )
4286 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
4287 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
4288 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
4289 .route(routes::V1_EMBEDDINGS, post(openai_extra::embeddings))
4290 .route(routes::CACHE_STATS, get(cache_stats))
4291 .route(routes::METRICS, get(metrics))
4292 .route(routes::ADMIN_MODELS, get(admin::models))
4297 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
4298 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
4299 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
4300 .route(routes::ADMIN_TASKS, get(admin::tasks))
4301 .route(&admin::cancel_route(), post(admin::cancel_task))
4302 .route(routes::ADMIN_STATS, get(admin::stats));
4303
4304 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4309 tracing::info!("API key auth enabled");
4310 let auth = limits::AuthConfig {
4311 api_key: Arc::new(key),
4312 };
4313 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4314 auth,
4315 limits::require_api_key,
4316 ));
4317 }
4318 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4319 let rpm: u32 = rpm
4320 .parse()
4321 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4322 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4323 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4324 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4325 limiter,
4326 limits::rate_limit,
4327 ));
4328 }
4329 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
4337 let origins = security::parse_cors_origins(&spec)
4338 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
4339 tracing::info!(
4340 "CORS enabled: {} allow-listed origin(s) ({})",
4341 origins.len(),
4342 spec
4343 );
4344 let cors = tower_http::cors::CorsLayer::new()
4345 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4346 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4347 .allow_headers([
4348 axum::http::header::CONTENT_TYPE,
4349 axum::http::header::AUTHORIZATION,
4350 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4356 axum::http::HeaderName::from_static("last-event-id"),
4360 ]);
4361 protected = protected.route_layer(cors);
4362 }
4363
4364 let app = public
4368 .merge(protected)
4369 .layer(axum::middleware::from_fn(limits::retry_after))
4370 .with_state(state);
4371
4372 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4378 install_ring_crypto_provider();
4379 match tls_paths {
4384 Some(paths) => {
4385 let config =
4386 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4387 .await
4388 .map_err(|e| {
4389 anyhow::anyhow!(
4390 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4391 paths.cert,
4392 paths.key
4393 )
4394 })?;
4395 let socket_addr: std::net::SocketAddr = addr
4396 .parse()
4397 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
4398 let listener = std::net::TcpListener::bind(socket_addr)?;
4399 listener.set_nonblocking(true)?;
4407 let bound = listener.local_addr()?;
4408 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
4409 announce_ready(bound, "https");
4410
4411 let handle = axum_server::Handle::new();
4412 let shutdown_handle = handle.clone();
4413 tokio::spawn(async move {
4414 shutdown_signal(exit_on_stdin_close).await;
4415 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4416 });
4417 axum_server::from_tcp_rustls(listener, config)?
4418 .handle(handle)
4419 .serve(app.into_make_service())
4420 .await?;
4421 }
4422 None => {
4423 let listener = tokio::net::TcpListener::bind(&addr).await?;
4424 let bound = listener.local_addr()?;
4425 tracing::info!("ferrox-server listening on {bound}");
4426 announce_ready(bound, "http");
4427 axum::serve(listener, app)
4428 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4429 .await?;
4430 }
4431 }
4432 Ok(())
4433}
4434
4435#[cfg(test)]
4436mod tests {
4437 use super::*;
4438 use ferrox_models::config::test_dense_fixture;
4439
4440 #[test]
4441 fn parses_llama_server_style_options() {
4442 let argv = [
4443 "ferrox-server",
4444 "-m",
4445 "model.gguf",
4446 "--host",
4447 "::1",
4448 "--port",
4449 "9000",
4450 "-t",
4451 "4",
4452 "-dev",
4453 "Metal",
4454 "-ngl",
4455 "all",
4456 ]
4457 .into_iter()
4458 .map(String::from)
4459 .collect();
4460 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4461
4462 assert_eq!(args.model.as_deref(), Some("model.gguf"));
4463 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
4464 assert_eq!(args.port, Some(9000));
4465 assert_eq!(args.threads, Some(4));
4466 assert_eq!(args.device, Some(OffloadDevice::Metal));
4467 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
4468 assert_eq!(
4469 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4470 Some("[::1]:9000")
4471 );
4472 }
4473
4474 #[test]
4475 fn port_zero_survives_argument_parsing_as_a_real_request() {
4476 let argv = ["ferrox-server", "--port", "0"]
4481 .into_iter()
4482 .map(String::from)
4483 .collect();
4484 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4485 assert_eq!(args.port, Some(0));
4486 assert_eq!(
4487 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4488 Some("127.0.0.1:0")
4489 );
4490 }
4491
4492 #[test]
4493 fn stdin_close_exit_is_opt_in() {
4494 let args =
4497 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
4498 assert!(!args.exit_on_stdin_close);
4499 let args = ServerArgs::try_parse_from(
4500 ["ferrox-server", "--exit-on-stdin-close"]
4501 .into_iter()
4502 .map(String::from),
4503 )
4504 .unwrap();
4505 assert!(args.exit_on_stdin_close);
4506 }
4507
4508 #[test]
4509 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4510 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4511 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4512 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
4513 assert_eq!(parsed.port, 51999);
4514 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4515 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
4517 }
4518
4519 fn test_model() -> Model {
4520 let cfg = test_dense_fixture();
4524 Model::Gguf(GgufModel {
4525 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4526 tokenizer: Arc::new(ServerTokenizer::Byte),
4527 stop_tokens: StopTokens::default(),
4528 bos_id: None,
4529 is_synthetic: true,
4530 chat_template: chat_template::PromptTemplate::plain(),
4531 })
4532 }
4533
4534 fn greedy_params(max_tokens: usize) -> GenerationParams {
4535 GenerationParams {
4536 max_tokens,
4537 sampling: SamplingParams::default(),
4538 seed: 1,
4539 stop: Vec::new(),
4540 stop_token_ids: Vec::new(),
4541 json_object: false,
4542 cancel: None,
4543 ignore_eos: false,
4544 }
4545 }
4546
4547 fn test_model_full_byte_vocab() -> Model {
4551 let mut cfg = test_dense_fixture();
4552 cfg.vocab_size = 256;
4553 Model::Gguf(GgufModel {
4554 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4555 tokenizer: Arc::new(ServerTokenizer::Byte),
4556 stop_tokens: StopTokens::default(),
4557 bos_id: None,
4558 is_synthetic: true,
4559 chat_template: chat_template::PromptTemplate::plain(),
4560 })
4561 }
4562
4563 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4567 AppState {
4568 paged_kv: None,
4569 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4570 id: None,
4571 model: Arc::new(model),
4572 batcher: None,
4573 ceiling: None,
4574 }))),
4575 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4576 tasks: Arc::new(tasks::TaskRegistry::new()),
4577 cancels: Arc::new(cancel::CancelRegistry::new()),
4578 stats: stats::Stats::new(),
4579 streams: resume::StreamRegistry::new(),
4580 model_dir: None,
4581 response_cache: Mutex::new(response_cache),
4582 kv_pool: None,
4583 prefix_cache: None,
4584 sessions: session::SessionStore::new(),
4585 requests_total: std::sync::atomic::AtomicU64::new(0),
4586 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4587 started_at: std::time::Instant::now(),
4588 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4589 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4590 mcp: None,
4591 continuous_batching_enabled: false,
4592 loading_model: Mutex::new(None),
4593 last_load_error: Mutex::new(None),
4594 serving: Mutex::new(ferrox_edge::ServingStats::default()),
4595 maintenance: Mutex::new(ferrox_edge::MaintenanceGate::serving()),
4596 footprint: Mutex::new(ferrox_edge::ProbeCache::new(FOOTPRINT_TTL_MS)),
4597 started_unix: unix_now(),
4598 }
4599 }
4600
4601 fn test_app() -> Router {
4610 test_app_with_state(Arc::new(test_state(
4611 test_model_full_byte_vocab(),
4612 ResponseCache::new(1000, Duration::from_secs(3600)),
4613 )))
4614 }
4615
4616 fn test_app_with_state(state: Arc<AppState>) -> Router {
4619 Router::new()
4620 .route(ferrox_api::routes::HEALTH, get(health))
4621 .route(ferrox_api::routes::V1_MODELS, get(list_models))
4622 .route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
4623 .route(
4624 &axum_path(ferrox_api::routes::V1_RESPONSE),
4625 get(responses::responses_get),
4626 )
4627 .route(
4628 &axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
4629 post(responses::responses_cancel),
4630 )
4631 .route(ferrox_api::routes::V1_STATS, get(serving_stats))
4632 .route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
4633 .route(
4634 ferrox_api::routes::V1_CACHE_STATUS,
4635 get(cache_admin::cache_status),
4636 )
4637 .route(
4638 ferrox_api::routes::V1_CACHE_REBUILD,
4639 post(cache_admin::cache_rebuild),
4640 )
4641 .route(
4642 ferrox_api::routes::ADMIN_PREPARE_STOP,
4643 post(cache_admin::prepare_stop),
4644 )
4645 .route("/v1/chat/completions", post(chat_completions))
4646 .route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
4647 .route(
4648 ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
4649 post(anthropic::count_tokens),
4650 )
4651 .route("/v1/tokenize", post(openai_extra::tokenize))
4652 .route("/v1/detokenize", post(openai_extra::detokenize))
4653 .route("/v1/embeddings", post(openai_extra::embeddings))
4654 .route("/v1/completions", post(openai_extra::completions))
4655 .route(
4656 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
4657 post(admin::unload_model),
4658 )
4659 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
4660 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
4661 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
4662 .route(
4663 &axum_path(ferrox_api::routes::V1_STREAM),
4664 get(resume::resume),
4665 )
4666 .route(
4667 &axum_path(ferrox_api::routes::V1_STREAM_POLL),
4668 get(resume::poll),
4669 )
4670 .with_state(state)
4671 }
4672
4673 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4674 let mut cfg = test_dense_fixture();
4675 cfg.name = name;
4676 cfg.vocab_size = vocab_size;
4677 Model::Gguf(GgufModel {
4678 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4679 tokenizer: Arc::new(ServerTokenizer::Byte),
4680 stop_tokens: StopTokens::default(),
4681 bos_id: None,
4682 is_synthetic: true,
4683 chat_template: chat_template::PromptTemplate::plain(),
4684 })
4685 }
4686
4687 fn model_with_template(name: &'static str, source: &str) -> Model {
4691 let mut cfg = test_dense_fixture();
4692 cfg.name = name;
4693 cfg.vocab_size = 256;
4694 Model::Gguf(GgufModel {
4695 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4696 tokenizer: Arc::new(ServerTokenizer::Byte),
4697 stop_tokens: StopTokens::default(),
4698 bos_id: None,
4699 is_synthetic: true,
4700 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4701 Some(source),
4702 Some("qwen3"),
4703 false,
4704 None,
4705 None,
4706 ),
4707 })
4708 }
4709
4710 #[tokio::test]
4719 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4720 let strict = "{% if messages | length > 1 %}\
4723 {{ raise_exception('this template takes one turn') }}\
4724 {% endif %}{{ messages[0].content }}";
4725 let state = Arc::new(test_state(
4726 model_with_template("strict", strict),
4727 ResponseCache::new(4, Duration::from_secs(60)),
4728 ));
4729 let app = test_app_with_state(state);
4730
4731 let (status, body) = post_json_uri(
4732 &app,
4733 "/v1/chat/completions",
4734 serde_json::json!({
4735 "model": "strict",
4736 "stream": true,
4737 "messages": [
4738 {"role": "user", "content": "one"},
4739 {"role": "user", "content": "two"},
4740 ],
4741 }),
4742 )
4743 .await;
4744 assert_eq!(status, StatusCode::BAD_REQUEST);
4745 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4746 assert!(
4747 body["error"]["message"]
4748 .as_str()
4749 .unwrap()
4750 .contains("one turn"),
4751 "the template's own message must reach the caller: {body}"
4752 );
4753
4754 let (status, _) = post_json_uri(
4756 &app,
4757 "/v1/chat/completions",
4758 serde_json::json!({
4759 "model": "strict",
4760 "stream": true,
4761 "max_tokens": 1,
4762 "messages": [{"role": "user", "content": "one"}],
4763 }),
4764 )
4765 .await;
4766 assert_eq!(status, StatusCode::OK);
4767 }
4768
4769 #[tokio::test]
4771 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4772 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4773 {% if reasoning_effort %}\
4774 {% if reasoning_effort not in ['low','medium','high'] %}\
4775 {{ raise_exception('bad effort') }}\
4776 {% endif %}[{{ reasoning_effort }}]\
4777 {% endif %}{{ messages[0].content }}";
4778 let state = Arc::new(test_state(
4779 model_with_template("thinker", reasoning),
4780 ResponseCache::new(4, Duration::from_secs(60)),
4781 ));
4782 let app = test_app_with_state(state);
4783 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4784 assert_eq!(status, StatusCode::OK);
4785 let entry = &models["data"][0];
4786 assert_eq!(
4787 entry["supported_reasoning_efforts"],
4788 serde_json::json!(["off", "low", "medium", "high"])
4789 );
4790 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4791 }
4792
4793 #[tokio::test]
4798 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4799 let app = test_app();
4800 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4801 let entry = &models["data"][0];
4802 assert!(entry.get("supported_reasoning_efforts").is_none());
4803 assert!(entry.get("default_reasoning_effort").is_none());
4804 }
4805
4806 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
4807 Arc::new(ActiveModel {
4808 id: Some(name.to_string()),
4809 model: Arc::new(named_test_model(name, 256)),
4810 batcher: None,
4811 ceiling: None,
4812 })
4813 .tap_into(state)
4814 }
4815
4816 trait TapInto {
4818 fn tap_into(self, state: &AppState) -> Self;
4819 }
4820 impl TapInto for Arc<ActiveModel> {
4821 fn tap_into(self, state: &AppState) -> Self {
4822 state.swap_active(Some(Arc::clone(&self)));
4823 self
4824 }
4825 }
4826
4827 #[test]
4833 fn an_in_flight_request_keeps_the_model_it_started_on() {
4834 let state = test_state(
4835 named_test_model("model-a", 256),
4836 ResponseCache::new(4, Duration::from_secs(60)),
4837 );
4838
4839 let in_flight = state.active().expect("a model is loaded");
4842 assert_eq!(in_flight.model.name(), "model-a");
4843
4844 active_model(&state, "model-b");
4845
4846 assert_eq!(state.active().unwrap().model.name(), "model-b");
4848 assert_eq!(in_flight.model.name(), "model-a");
4850 let (_chunks, finish, _usage) = run_generation(
4851 &in_flight.model,
4852 "hi",
4853 &greedy_params(3),
4854 None,
4855 None,
4856 None,
4857 None,
4858 None,
4859 )
4860 .expect("the old model must still decode after being swapped out");
4861 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
4862 }
4863
4864 #[test]
4869 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
4870 let state = test_state(
4871 named_test_model("model-a", 256),
4872 ResponseCache::new(4, Duration::from_secs(60)),
4873 );
4874 let in_flight = state.active().expect("a model is loaded");
4875 let weights = Arc::clone(&in_flight.model);
4876 assert!(Arc::strong_count(&weights) >= 2);
4877
4878 let previous = state.swap_active(Some(Arc::new(ActiveModel {
4879 id: Some("model-b".to_string()),
4880 model: Arc::new(named_test_model("model-b", 256)),
4881 batcher: None,
4882 ceiling: None,
4883 })));
4884 drop(previous);
4885 assert!(Arc::strong_count(&weights) >= 2);
4887 drop(in_flight);
4888 assert_eq!(Arc::strong_count(&weights), 1);
4889 }
4890
4891 #[tokio::test]
4895 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
4896 let state = Arc::new(test_state(
4897 named_test_model("model-a", 256),
4898 ResponseCache::new(4, Duration::from_secs(60)),
4899 ));
4900 let app = test_app_with_state(Arc::clone(&state));
4901
4902 let (status, body) = post_json_uri(
4903 &app,
4904 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
4905 serde_json::json!({}),
4906 )
4907 .await;
4908 assert_eq!(status, StatusCode::OK);
4909 assert_eq!(body["ok"], true);
4910 assert!(body["active"].is_null());
4911 assert!(state.active().is_none());
4912
4913 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4914 assert_eq!(status, StatusCode::OK);
4915 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4916 assert_eq!(models["data"].as_array().unwrap().len(), 0);
4917
4918 let (status, body) = post_json_uri(
4919 &app,
4920 "/v1/chat/completions",
4921 serde_json::json!({
4922 "model": "x",
4923 "messages": [{"role": "user", "content": "hi"}]
4924 }),
4925 )
4926 .await;
4927 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
4928 assert_eq!(body["error"]["type"], "model_not_loaded");
4929 }
4930
4931 #[tokio::test]
4935 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
4936 let state = Arc::new(test_state(
4937 named_test_model("model-a", 256),
4938 ResponseCache::new(4, Duration::from_secs(60)),
4939 ));
4940 let app = test_app_with_state(Arc::clone(&state));
4941 state.swap_active(None);
4942
4943 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
4944 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
4947 assert_eq!(body["state"], "unavailable");
4948 assert_eq!(body["reason"], "model_not_loaded");
4949 assert!(body["model"].is_null());
4950 let real_weights = body["capabilities"]
4951 .as_array()
4952 .unwrap()
4953 .iter()
4954 .find(|c| c["id"] == "real_weights")
4955 .cloned()
4956 .expect("real_weights is always reported");
4957 assert_eq!(real_weights["available"], false);
4958 assert_eq!(real_weights["reason"], "model_not_loaded");
4959 }
4960
4961 #[tokio::test]
4965 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
4966 let app = test_app();
4967
4968 let (status, completion) = post_json_uri(
4969 &app,
4970 "/v1/chat/completions",
4971 serde_json::json!({
4972 "model": "x",
4973 "messages": [{"role": "user", "content": "hi"}],
4974 "max_tokens": 4
4975 }),
4976 )
4977 .await;
4978 assert_eq!(status, StatusCode::OK);
4979 let request_id = completion["request_id"].as_str().unwrap().to_string();
4980
4981 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4982 assert_eq!(status, StatusCode::OK);
4983 let recent = stats["recent"].as_array().unwrap();
4984 assert_eq!(recent.len(), 1);
4985 let row = &recent[0];
4986 assert_eq!(row["request_id"], request_id);
4987 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
4988 assert_eq!(row["status"], 200);
4989 assert_eq!(row["stream"], false);
4990 assert!(row["duration_ms"].is_number());
4993 assert!(row["decode_ms"].is_number());
4994 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
4995 assert_eq!(
4996 stats["tokens_prompt_total"].as_u64().unwrap(),
4997 row["prompt_tokens"].as_u64().unwrap()
4998 );
4999 }
5000
5001 #[tokio::test]
5005 async fn a_rejected_request_is_recorded_too() {
5006 let state = Arc::new(test_state(
5007 named_test_model("model-a", 256),
5008 ResponseCache::new(4, Duration::from_secs(60)),
5009 ));
5010 let app = test_app_with_state(Arc::clone(&state));
5011 state.swap_active(None);
5012
5013 let (status, _) = post_json_uri(
5014 &app,
5015 "/v1/chat/completions",
5016 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5017 )
5018 .await;
5019 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5020
5021 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5022 let recent = stats["recent"].as_array().unwrap();
5023 assert_eq!(recent.len(), 1);
5024 assert_eq!(recent[0]["status"], 503);
5025 assert_eq!(recent[0]["completion_tokens"], 0);
5026 assert!(recent[0]["decode_ms"].is_null());
5027 assert_eq!(stats["errors_total"], 1);
5028 }
5029
5030 async fn post_json_with_headers(
5034 app: &Router,
5035 uri: &str,
5036 body: serde_json::Value,
5037 headers: &[(&str, &str)],
5038 ) -> (StatusCode, serde_json::Value) {
5039 use http_body_util::BodyExt;
5040 use tower::ServiceExt;
5041
5042 let mut builder = axum::http::Request::builder()
5043 .method("POST")
5044 .uri(uri)
5045 .header("content-type", "application/json");
5046 for (name, value) in headers {
5047 builder = builder.header(*name, *value);
5048 }
5049 let response = app
5050 .clone()
5051 .oneshot(
5052 builder
5053 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5054 .unwrap(),
5055 )
5056 .await
5057 .unwrap();
5058 let status = response.status();
5059 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5060 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5061 (status, json)
5062 }
5063
5064 #[tokio::test]
5068 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5069 let app = test_app();
5070
5071 let (status, _) = post_json_uri(
5072 &app,
5073 ferrox_api::routes::V1_TOKENIZE,
5074 serde_json::json!({"prompt": "hello"}),
5075 )
5076 .await;
5077 assert_eq!(status, StatusCode::OK);
5078 let (status, _) = post_json_uri(
5079 &app,
5080 ferrox_api::routes::V1_DETOKENIZE,
5081 serde_json::json!({"tokens": [104, 105]}),
5082 )
5083 .await;
5084 assert_eq!(status, StatusCode::OK);
5085 let (status, _) = post_json_uri(
5086 &app,
5087 ferrox_api::routes::V1_EMBEDDINGS,
5088 serde_json::json!({"input": "hello"}),
5089 )
5090 .await;
5091 assert_eq!(status, StatusCode::OK);
5092
5093 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5094 let routes: Vec<&str> = stats["recent"]
5095 .as_array()
5096 .unwrap()
5097 .iter()
5098 .map(|row| row["route"].as_str().unwrap())
5099 .collect();
5100 for expected in [
5101 ferrox_api::routes::V1_TOKENIZE,
5102 ferrox_api::routes::V1_DETOKENIZE,
5103 ferrox_api::routes::V1_EMBEDDINGS,
5104 ] {
5105 assert!(
5106 routes.contains(&expected),
5107 "{expected} is missing: {routes:?}"
5108 );
5109 }
5110
5111 let row = |route: &str| {
5112 stats["recent"]
5113 .as_array()
5114 .unwrap()
5115 .iter()
5116 .find(|r| r["route"] == route)
5117 .cloned()
5118 .unwrap()
5119 };
5120 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5124 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5125 assert!(embed["decode_ms"].is_null());
5126 assert_eq!(embed["completion_tokens"], 0);
5127 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5131 assert_eq!(
5132 stats["tokens_prompt_total"].as_u64().unwrap(),
5133 embed["prompt_tokens"].as_u64().unwrap(),
5134 "only the forward pass counted"
5135 );
5136 }
5137
5138 #[tokio::test]
5141 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5142 let app = test_app();
5143 let (status, _) = post_json_uri(
5144 &app,
5145 ferrox_api::routes::V1_EMBEDDINGS,
5146 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5147 )
5148 .await;
5149 assert_eq!(status, StatusCode::BAD_REQUEST);
5150
5151 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5152 let recent = stats["recent"].as_array().unwrap();
5153 assert_eq!(recent.len(), 1);
5154 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
5155 assert_eq!(recent[0]["status"], 400);
5156 assert_eq!(
5157 recent[0]["prompt_tokens"], 0,
5158 "a rejected call embedded nothing"
5159 );
5160 }
5161
5162 #[tokio::test]
5165 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5166 let app = test_app();
5167 let key = "sk-monitor-secret";
5168 let (status, _) = post_json_with_headers(
5169 &app,
5170 "/v1/chat/completions",
5171 serde_json::json!({
5172 "model": "x",
5173 "messages": [{"role": "user", "content": "hi"}],
5174 "max_tokens": 2
5175 }),
5176 &[
5177 ("authorization", &format!("Bearer {key}")),
5178 ("x-ferrox-client", "ferrox-studio"),
5179 ],
5180 )
5181 .await;
5182 assert_eq!(status, StatusCode::OK);
5183
5184 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5185 let row = stats["recent"].as_array().unwrap()[0].clone();
5186 let fingerprint = row["via_api_key"]
5187 .as_str()
5188 .expect("the row names the key that served it")
5189 .to_string();
5190 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5191 assert!(!fingerprint.contains(key));
5192 assert!(
5193 !serde_json::to_string(&stats).unwrap().contains(key),
5194 "the stats payload must not carry the key in any form"
5195 );
5196 assert_eq!(row["client"], "ferrox-studio");
5197 }
5198
5199 #[tokio::test]
5202 async fn different_keys_are_different_callers_and_no_key_is_null() {
5203 let app = test_app();
5204 let body = serde_json::json!({
5205 "model": "x",
5206 "messages": [{"role": "user", "content": "hi"}],
5207 "max_tokens": 1
5208 });
5209 for headers in [
5210 vec![("authorization", "Bearer key-one")],
5211 vec![("authorization", "Bearer key-two")],
5212 vec![],
5213 ] {
5214 let (status, _) =
5215 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5216 assert_eq!(status, StatusCode::OK);
5217 }
5218
5219 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5220 let recent = stats["recent"].as_array().unwrap();
5221 assert_eq!(recent.len(), 3);
5222 let one = recent[0]["via_api_key"].as_str().unwrap();
5223 let two = recent[1]["via_api_key"].as_str().unwrap();
5224 assert_ne!(one, two, "two keys must not collapse into one caller");
5225 assert!(
5226 recent[2]["via_api_key"].is_null(),
5227 "an unauthenticated call is null, not a fingerprint of nothing"
5228 );
5229 assert!(recent[2]["client"].is_null());
5230 }
5231
5232 #[tokio::test]
5237 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5238 let state = Arc::new(test_state(
5239 named_test_model("really-loaded", 256),
5240 ResponseCache::new(4, Duration::from_secs(60)),
5241 ));
5242 let app = test_app_with_state(Arc::clone(&state));
5243
5244 let (status, _) = post_json_uri(
5245 &app,
5246 "/v1/chat/completions",
5247 serde_json::json!({
5248 "model": "gpt-4-turbo-that-is-not-here",
5249 "messages": [{"role": "user", "content": "hi"}],
5250 "max_tokens": 2
5251 }),
5252 )
5253 .await;
5254 assert_eq!(status, StatusCode::OK);
5255
5256 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5257 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5258
5259 state.swap_active(None);
5262 let (status, _) = post_json_uri(
5263 &app,
5264 "/v1/chat/completions",
5265 serde_json::json!({
5266 "model": "gpt-4-turbo-that-is-not-here",
5267 "messages": [{"role": "user", "content": "hi"}]
5268 }),
5269 )
5270 .await;
5271 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5272 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5273 let recent = stats["recent"].as_array().unwrap();
5274 assert!(recent[recent.len() - 1]["model"].is_null());
5275 }
5276
5277 #[tokio::test]
5281 async fn a_streamed_row_names_the_model_it_decoded_against() {
5282 let state = Arc::new(test_state(
5283 named_test_model("model-before", 256),
5284 ResponseCache::new(4, Duration::from_secs(60)),
5285 ));
5286 let app = test_app_with_state(Arc::clone(&state));
5287 let _ = post_sse_raw(&app, resumable_request()).await;
5288 active_model(&state, "model-after");
5290
5291 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5292 assert_eq!(stats["recent"][0]["model"], "model-before");
5293 }
5294
5295 #[tokio::test]
5298 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5299 let app = test_app();
5300 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5301 assert_eq!(status, StatusCode::OK);
5302 assert!(
5303 stats["queue_depth"].is_null(),
5304 "without continuous batching nothing queues, so there is nothing to measure"
5305 );
5306 assert!(stats["queue_rejected_total"].is_null());
5307 assert_eq!(
5308 stats["generating_now"], 0,
5309 "work in progress is measured and really is zero here"
5310 );
5311 }
5312
5313 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5318 use http_body_util::BodyExt;
5319 use tower::ServiceExt;
5320
5321 let response = app
5322 .clone()
5323 .oneshot(
5324 axum::http::Request::builder()
5325 .method("POST")
5326 .uri("/v1/chat/completions")
5327 .header("content-type", "application/json")
5328 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5329 .unwrap(),
5330 )
5331 .await
5332 .unwrap();
5333 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5334 String::from_utf8(bytes.to_vec()).unwrap()
5335 }
5336
5337 async fn get_json_with_headers(
5338 app: &Router,
5339 uri: &str,
5340 headers: &[(&str, &str)],
5341 ) -> (StatusCode, serde_json::Value) {
5342 use http_body_util::BodyExt;
5343 use tower::ServiceExt;
5344
5345 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
5346 for (name, value) in headers {
5347 builder = builder.header(*name, *value);
5348 }
5349 let response = app
5350 .clone()
5351 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
5352 .await
5353 .unwrap();
5354 let status = response.status();
5355 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5356 (
5357 status,
5358 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
5359 )
5360 }
5361
5362 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
5363 body.lines()
5364 .filter_map(|line| line.strip_prefix(field))
5365 .map(str::trim)
5366 .collect()
5367 }
5368
5369 fn resumable_request() -> serde_json::Value {
5370 serde_json::json!({
5371 "model": "m",
5372 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5373 "max_tokens": 4,
5374 "temperature": 0,
5375 "stream": true,
5376 "stream_resumable": true,
5377 })
5378 }
5379
5380 #[tokio::test]
5385 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
5386 let app = test_app();
5387 let body = post_sse_raw(&app, resumable_request()).await;
5388
5389 let request_id = body
5390 .lines()
5391 .find_map(|l| l.strip_prefix("data: "))
5392 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
5393 .and_then(|v| v["request_id"].as_str().map(str::to_string))
5394 .expect("the first chunk names the request");
5395
5396 let ids = sse_field(&body, "id:");
5397 let datas = sse_field(&body, "data:");
5398 assert_eq!(
5399 ids.len(),
5400 datas.len(),
5401 "every event carries an id, or a reconnect cannot name where it stopped"
5402 );
5403 for (i, id) in ids.iter().enumerate() {
5404 assert_eq!(*id, format!("{request_id}:{i}"));
5405 }
5406 let retries = sse_field(&body, "retry:");
5407 assert_eq!(
5408 retries.len(),
5409 1,
5410 "the reconnect delay is stated once, not on every event"
5411 );
5412 assert_eq!(retries[0], "1500");
5413 assert!(
5414 body.contains("data: [DONE]"),
5415 "the end of stream is still stated"
5416 );
5417 }
5418
5419 #[tokio::test]
5423 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
5424 let app = test_app();
5425 let mut request = resumable_request();
5426 request["stream_resumable"] = serde_json::json!(false);
5427 let body = post_sse_raw(&app, request).await;
5428 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
5429 assert!(
5430 sse_field(&body, "id:").is_empty(),
5431 "an id promises a replay this stream cannot serve"
5432 );
5433 assert!(sse_field(&body, "retry:").is_empty());
5434 }
5435
5436 #[tokio::test]
5440 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
5441 let app = test_app();
5442 let body = post_sse_raw(&app, resumable_request()).await;
5443 let request_id = sse_field(&body, "id:")[0]
5444 .rsplit_once(':')
5445 .unwrap()
5446 .0
5447 .to_string();
5448 let streamed: Vec<String> = sse_field(&body, "data:")
5449 .iter()
5450 .map(|d| d.to_string())
5451 .collect();
5452
5453 let (status, polled) = get_json(
5454 &app,
5455 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
5456 )
5457 .await;
5458 assert_eq!(status, StatusCode::OK);
5459 let events: Vec<String> = polled["events"]
5460 .as_array()
5461 .unwrap()
5462 .iter()
5463 .map(|e| e["data"].as_str().unwrap().to_string())
5464 .collect();
5465 assert_eq!(
5466 events, streamed,
5467 "the fallback must deliver the same answer, not a re-run of it"
5468 );
5469 assert_eq!(polled["request_id"], request_id);
5470 assert_eq!(
5471 polled["done"], false,
5472 "events were still being handed out, so the client must ask again"
5473 );
5474
5475 let next = polled["next_index"].as_u64().unwrap();
5478 let (_, drained) = get_json(
5479 &app,
5480 &format!(
5481 "{}?from={next}",
5482 ferrox_api::routes::v1_stream_poll(&request_id)
5483 ),
5484 )
5485 .await;
5486 assert_eq!(drained["done"], true);
5487 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
5488 }
5489
5490 #[tokio::test]
5494 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
5495 let app = test_app();
5496 let body = post_sse_raw(&app, resumable_request()).await;
5497 let ids = sse_field(&body, "id:");
5498 let datas: Vec<String> = sse_field(&body, "data:")
5499 .iter()
5500 .map(|d| d.to_string())
5501 .collect();
5502 assert!(
5503 ids.len() >= 3,
5504 "need a few events to resume into the middle"
5505 );
5506 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
5507
5508 let (status, resumed) = get_json_with_headers(
5509 &app,
5510 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
5511 &[],
5512 )
5513 .await;
5514 assert_eq!(status, StatusCode::OK);
5515 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
5516
5517 let (_, tail) = get_json(
5519 &app,
5520 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
5521 )
5522 .await;
5523 let tail_events: Vec<String> = tail["events"]
5524 .as_array()
5525 .unwrap()
5526 .iter()
5527 .map(|e| e["data"].as_str().unwrap().to_string())
5528 .collect();
5529 assert_eq!(tail_events, datas[2..].to_vec());
5530 }
5531
5532 #[tokio::test]
5535 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
5536 use http_body_util::BodyExt;
5537 use tower::ServiceExt;
5538
5539 let app = test_app();
5540 let body = post_sse_raw(&app, resumable_request()).await;
5541 let ids = sse_field(&body, "id:");
5542 let datas: Vec<String> = sse_field(&body, "data:")
5543 .iter()
5544 .map(|d| d.to_string())
5545 .collect();
5546 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
5547
5548 let response = app
5549 .clone()
5550 .oneshot(
5551 axum::http::Request::builder()
5552 .method("GET")
5553 .uri(ferrox_api::routes::v1_stream(&request_id))
5554 .header("last-event-id", format!("{request_id}:0"))
5555 .body(axum::body::Body::empty())
5556 .unwrap(),
5557 )
5558 .await
5559 .unwrap();
5560 assert_eq!(response.status(), StatusCode::OK);
5561 assert_eq!(
5562 response
5563 .headers()
5564 .get("x-accel-buffering")
5565 .and_then(|v| v.to_str().ok()),
5566 Some("no"),
5567 "the reconnect needs the same anti-buffering header as the stream"
5568 );
5569 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5570 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
5571 assert_eq!(
5572 sse_field(&resumed, "data:")
5573 .iter()
5574 .map(|d| d.to_string())
5575 .collect::<Vec<_>>(),
5576 datas[1..].to_vec()
5577 );
5578 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
5579 }
5580
5581 #[tokio::test]
5585 async fn a_last_event_id_from_another_stream_is_refused() {
5586 let app = test_app();
5587 let body = post_sse_raw(&app, resumable_request()).await;
5588 let request_id = sse_field(&body, "id:")[0]
5589 .rsplit_once(':')
5590 .unwrap()
5591 .0
5592 .to_string();
5593
5594 let (status, err) = get_json_with_headers(
5595 &app,
5596 &ferrox_api::routes::v1_stream(&request_id),
5597 &[("last-event-id", "chatcmpl-someone-else:3")],
5598 )
5599 .await;
5600 assert_eq!(status, StatusCode::BAD_REQUEST);
5601 assert_eq!(err["error"]["code"], "bad_last_event_id");
5602 }
5603
5604 #[tokio::test]
5608 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
5609 let app = test_app();
5610 let mut request = resumable_request();
5611 request["stream_resumable"] = serde_json::json!(false);
5612 let body = post_sse_raw(&app, request).await;
5613 let request_id = body
5614 .lines()
5615 .find_map(|l| l.strip_prefix("data: "))
5616 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
5617 .and_then(|v| v["request_id"].as_str().map(str::to_string))
5618 .unwrap();
5619
5620 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
5621 assert_eq!(status, StatusCode::NOT_FOUND);
5622 assert_eq!(err["error"]["code"], "stream_not_found");
5623 assert!(err["error"]["message"]
5624 .as_str()
5625 .unwrap()
5626 .contains("stream_resumable"));
5627 }
5628
5629 #[test]
5633 fn the_axum_stream_patterns_match_the_published_templates() {
5634 assert_eq!(
5635 axum_path(ferrox_api::routes::V1_STREAM),
5636 "/v1/stream/:request_id"
5637 );
5638 assert_eq!(
5639 axum_path(ferrox_api::routes::V1_STREAM_POLL),
5640 "/v1/stream/:request_id/poll"
5641 );
5642 assert_eq!(
5643 ferrox_api::routes::v1_stream("abc"),
5644 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
5645 );
5646 }
5647
5648 #[test]
5656 fn auto_counts_as_the_gpu_whenever_there_is_one() {
5657 assert!(gpu_selected(Some("1"), true));
5658 assert!(
5659 gpu_selected(Some("auto"), true),
5660 "--ngl 99 with a Metal device present is a GPU run"
5661 );
5662 assert!(
5663 !gpu_selected(Some("auto"), false),
5664 "auto on a box with no GPU is a CPU run, and paged KV is fine there"
5665 );
5666 assert!(!gpu_selected(Some("0"), true));
5667 assert!(!gpu_selected(None, true));
5668 }
5669
5670 #[test]
5680 fn no_published_template_reaches_the_router_with_its_braces() {
5681 for template in [
5682 ferrox_api::routes::V1_STREAM,
5683 ferrox_api::routes::V1_STREAM_POLL,
5684 ferrox_api::routes::V1_RESPONSE,
5685 ferrox_api::routes::V1_RESPONSE_CANCEL,
5686 ferrox_api::routes::ADMIN_TASK_CANCEL,
5687 ] {
5688 assert!(
5689 template.contains('{'),
5690 "{template} is in the template list but has no placeholder"
5691 );
5692 let mounted = axum_path(template);
5693 assert!(
5694 !mounted.contains('{') && !mounted.contains('}'),
5695 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
5696 );
5697 assert!(
5698 mounted.contains(':'),
5699 "{template} lost its placeholder entirely and would match one path only"
5700 );
5701 }
5702 }
5703
5704 #[tokio::test]
5712 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
5713 let app = test_app();
5714 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
5715 assert_eq!(status, StatusCode::NOT_FOUND);
5716 assert!(
5717 !body.is_null(),
5718 "empty body means axum never matched the route, so the id was read as a literal segment"
5719 );
5720 }
5721
5722 #[tokio::test]
5725 async fn the_task_list_starts_empty_rather_than_absent() {
5726 let app = test_app();
5727 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
5728 assert_eq!(status, StatusCode::OK);
5729 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
5730 }
5731
5732 async fn post_json_uri(
5733 app: &Router,
5734 uri: &str,
5735 body: serde_json::Value,
5736 ) -> (StatusCode, serde_json::Value) {
5737 use http_body_util::BodyExt;
5738 use tower::ServiceExt;
5739
5740 let response = app
5741 .clone()
5742 .oneshot(
5743 axum::http::Request::builder()
5744 .method("POST")
5745 .uri(uri)
5746 .header("content-type", "application/json")
5747 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5748 .unwrap(),
5749 )
5750 .await
5751 .unwrap();
5752 let status = response.status();
5753 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5754 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5755 (status, json)
5756 }
5757
5758 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
5759 post_json_uri(app, "/v1/chat/completions", body).await.1
5760 }
5761
5762 #[tokio::test]
5771 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
5772 let app = test_app();
5773 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
5774 assert_eq!(status, StatusCode::OK);
5775
5776 let memory = &body["memory"];
5777 if memory.is_null() {
5778 return;
5781 }
5782 assert!(
5783 memory["bytes"].as_u64().is_some_and(|b| b > 0),
5784 "a read that produced a zero is a broken read, not an idle \
5785 engine: {memory}"
5786 );
5787 assert!(
5788 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
5789 "the quantity must travel with the number: {memory}"
5790 );
5791 }
5792
5793 #[tokio::test]
5799 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
5800 let app = test_app();
5801 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
5802 assert_eq!(status, StatusCode::OK);
5803 for pool in ["kv_pages", "window_slots", "state_slots"] {
5804 assert!(
5805 body["pools"][pool].is_null(),
5806 "{pool} must be null rather than a zero row: {}",
5807 body["pools"]
5808 );
5809 }
5810 }
5811
5812 #[tokio::test]
5820 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
5821 use http_body_util::BodyExt;
5822 use tower::ServiceExt;
5823
5824 let app = test_app();
5825 let send = |body: serde_json::Value| {
5826 let app = app.clone();
5827 async move {
5828 app.oneshot(
5829 axum::http::Request::builder()
5830 .method("POST")
5831 .uri(ferrox_api::routes::V1_MESSAGES)
5832 .header("content-type", "application/json")
5833 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5834 .unwrap(),
5835 )
5836 .await
5837 .unwrap()
5838 }
5839 };
5840
5841 let ok = send(serde_json::json!({
5842 "model": "test",
5843 "max_tokens": 1,
5844 "messages": [{"role": "user", "content": "hi"}],
5845 }))
5846 .await;
5847 assert_eq!(ok.status(), StatusCode::OK);
5848 let id = ok
5849 .headers()
5850 .get("request-id")
5851 .expect("a served message names its id")
5852 .to_str()
5853 .unwrap()
5854 .to_string();
5855 assert!(!id.is_empty());
5856
5857 let bad = send(serde_json::json!({"model": "test"})).await;
5860 assert!(bad.status().is_client_error());
5861 let other = bad.headers().get("request-id").expect("errors too");
5862 assert_ne!(other.to_str().unwrap(), id);
5863 let _ = bad.into_body().collect().await.unwrap();
5864 }
5865
5866 #[tokio::test]
5874 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
5875 let state = Arc::new(test_state(
5876 test_model_full_byte_vocab(),
5877 ResponseCache::new(1000, Duration::from_secs(3600)),
5878 ));
5879 let app = test_app_with_state(Arc::clone(&state));
5880 let body = serde_json::json!({
5881 "model": "test",
5882 "messages": [{"role": "user", "content": "hi"}],
5883 "max_tokens": 1,
5884 });
5885
5886 state
5887 .maintenance
5888 .lock()
5889 .unwrap()
5890 .begin_rebuild()
5891 .expect("a fresh server is serving, so the rebuild starts");
5892 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
5893 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5894 assert_eq!(refused["error"]["type"], "cache_rebuilding");
5895
5896 state.maintenance.lock().unwrap().finish_rebuild(true);
5897 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
5898 assert_eq!(
5899 status,
5900 StatusCode::OK,
5901 "the gate reopens; a rebuild is not a latch"
5902 );
5903 }
5904
5905 #[tokio::test]
5910 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
5911 let app = test_app();
5912 let (status, body) = post_json_uri(
5913 &app,
5914 ferrox_api::routes::V1_CANCEL,
5915 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
5916 )
5917 .await;
5918 assert_eq!(status, StatusCode::NOT_FOUND);
5919 assert_eq!(body["cancelled"], serde_json::json!(false));
5920 assert_eq!(body["request_id"], "chatcmpl-never-issued");
5921 assert!(
5922 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
5923 "the verdict must carry a human reason: {body}"
5924 );
5925 }
5926
5927 #[tokio::test]
5931 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
5932 let state = Arc::new(test_state(
5933 test_model_full_byte_vocab(),
5934 ResponseCache::new(1000, Duration::from_secs(3600)),
5935 ));
5936 let app = test_app_with_state(Arc::clone(&state));
5937 let (token, _guard) = state.cancels.register("chatcmpl-live");
5938
5939 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5940 assert_eq!(status, StatusCode::OK);
5941 assert_eq!(before["generating_now"], serde_json::json!(1));
5942
5943 let (status, body) = post_json_uri(
5944 &app,
5945 ferrox_api::routes::V1_CANCEL,
5946 serde_json::json!({ "request_id": "chatcmpl-live" }),
5947 )
5948 .await;
5949 assert_eq!(status, StatusCode::OK);
5950 assert_eq!(body["cancelled"], serde_json::json!(true));
5951 assert!(
5952 token.is_cancelled(),
5953 "the endpoint answered ok without setting the flag the decode loop reads"
5954 );
5955 }
5956
5957 #[tokio::test]
5958 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
5959 let app = test_app();
5960 let (status, tok) =
5961 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
5962 assert_eq!(status, StatusCode::OK);
5963 let tokens = tok["tokens"].as_array().unwrap();
5964 assert_eq!(tok["count"], tokens.len());
5965 assert!(!tokens.is_empty());
5966
5967 let (status, detok) = post_json_uri(
5968 &app,
5969 "/v1/detokenize",
5970 serde_json::json!({ "tokens": tokens }),
5971 )
5972 .await;
5973 assert_eq!(status, StatusCode::OK);
5974 assert_eq!(detok["text"], "Hi");
5975
5976 let (status, emb) = post_json_uri(
5977 &app,
5978 "/v1/embeddings",
5979 serde_json::json!({
5980 "input": "Hi",
5981 "embedding_type": "mean"
5982 }),
5983 )
5984 .await;
5985 assert_eq!(status, StatusCode::OK);
5986 let vec = emb["data"][0]["embedding"].as_array().unwrap();
5987 assert!(!vec.is_empty());
5988 assert!(vec.iter().all(|v| v.as_f64().is_some()));
5989 }
5990
5991 #[tokio::test]
5996 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
5997 use http_body_util::BodyExt;
5998 use tower::ServiceExt;
5999
6000 let fixture = concat!(
6001 "../ferrox-models/tests/fixtures/",
6002 "ferrox_real_moe_test.gguf"
6003 );
6004 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6005 let decoder = Decoder::from_gguf_with_expert_cache(
6006 &fixture,
6007 ferrox_models::config::test_moe_fixture(),
6008 Some(1024 * 1024),
6009 )
6010 .expect("MoE fixture must load store-backed");
6011
6012 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
6016 .layers
6017 .iter()
6018 .map(|_| {
6019 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
6020 })
6021 .collect();
6022 decoder.forward_token(1, 0, &mut caches);
6023
6024 let model = Model::Gguf(GgufModel {
6025 decoder: Arc::new(decoder),
6026 tokenizer: Arc::new(ServerTokenizer::Byte),
6027 stop_tokens: StopTokens::default(),
6028 bos_id: None,
6029 is_synthetic: false,
6030 chat_template: chat_template::PromptTemplate::plain(),
6031 });
6032 let state = Arc::new(test_state(
6033 model,
6034 ResponseCache::new(16, Duration::from_secs(60)),
6035 ));
6036 let app = Router::new()
6037 .route("/metrics", axum::routing::get(metrics))
6038 .route("/v1/chat/completions", post(chat_completions))
6039 .with_state(state);
6040
6041 let fetch_metrics = |app: Router| async move {
6042 let resp = app
6043 .oneshot(
6044 axum::http::Request::builder()
6045 .method("GET")
6046 .uri("/metrics")
6047 .body(axum::body::Body::empty())
6048 .unwrap(),
6049 )
6050 .await
6051 .unwrap();
6052 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6053 String::from_utf8(bytes.to_vec()).unwrap()
6054 };
6055
6056 let after = fetch_metrics(app.clone()).await;
6057 assert!(
6058 after.contains("ferrox_expert_cache_misses_total"),
6059 "streaming model must expose expert-cache metrics: {after}"
6060 );
6061 let misses: u64 = after
6062 .lines()
6063 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
6064 .and_then(|l| l.split_whitespace().nth(1))
6065 .and_then(|v| v.parse().ok())
6066 .expect("misses metric line must parse");
6067 assert!(
6068 misses > 0,
6069 "decode must have read experts through the store: {after}"
6070 );
6071 }
6072
6073 fn weather_tool() -> serde_json::Value {
6074 serde_json::json!({
6075 "type": "function",
6076 "function": {
6077 "name": "get_weather",
6078 "description": "Get the current weather for a location.",
6079 "parameters": {
6080 "type": "object",
6081 "properties": {"location": {"type": "string"}},
6082 "required": ["location"]
6083 }
6084 }
6085 })
6086 }
6087
6088 fn weather_tool_def() -> ToolDef {
6089 ToolDef {
6090 kind: "function".to_string(),
6091 function: ToolFunctionDef {
6092 name: "get_weather".to_string(),
6093 description: Some("Get the current weather for a location.".to_string()),
6094 parameters: Some(serde_json::json!({
6095 "type": "object",
6096 "properties": {"location": {"type": "string"}},
6097 "required": ["location"]
6098 })),
6099 },
6100 }
6101 }
6102
6103 #[test]
6104 fn tool_preamble_mentions_every_tool_name_and_description() {
6105 let preamble = tool_preamble(&[weather_tool_def()]);
6106 assert!(preamble.contains("get_weather"));
6107 assert!(preamble.contains("Get the current weather for a location."));
6108 assert!(preamble.contains("<tool_call>"));
6109 assert!(preamble.contains("</tool_call>"));
6110 }
6111
6112 #[test]
6113 fn a_real_marker_becomes_a_structured_tool_call() {
6114 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6115 let (message, finish) = build_response_message(
6116 text.to_string(),
6117 &[weather_tool_def()],
6118 output::OutputPosture::for_model("test-model"),
6119 "stop",
6120 );
6121 assert_eq!(finish, "tool_calls");
6122 let calls = message.tool_calls.expect("must carry a tool call");
6123 assert_eq!(calls[0].function.name, "get_weather");
6124 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
6125 assert_eq!(parsed["location"], "Paris");
6126 }
6127
6128 #[test]
6129 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
6130 let (message, finish) = build_response_message(
6131 "just an answer".to_string(),
6132 &[weather_tool_def()],
6133 output::OutputPosture::for_model("test-model"),
6134 "stop",
6135 );
6136 assert_eq!(finish, "stop");
6137 assert!(message.tool_calls.is_none());
6138 assert_eq!(message.content.as_deref(), Some("just an answer"));
6139 }
6140
6141 #[test]
6144 fn a_malformed_payload_is_not_a_tool_call() {
6145 let (message, finish) = build_response_message(
6146 "<tool_call>not valid json at all</tool_call>".to_string(),
6147 &[weather_tool_def()],
6148 output::OutputPosture::for_model("test-model"),
6149 "stop",
6150 );
6151 assert_eq!(finish, "stop");
6152 assert!(message.tool_calls.is_none());
6153 }
6154
6155 #[test]
6158 fn a_tool_that_was_never_offered_is_not_returned() {
6159 let (message, finish) = build_response_message(
6160 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
6161 &[weather_tool_def()],
6162 output::OutputPosture::for_model("test-model"),
6163 "stop",
6164 );
6165 assert_eq!(finish, "stop");
6166 assert!(message.tool_calls.is_none());
6167 }
6168
6169 #[test]
6171 fn marker_text_with_no_tools_offered_stays_content() {
6172 let (message, finish) = build_response_message(
6173 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
6174 &[],
6175 output::OutputPosture::for_model("test-model"),
6176 "stop",
6177 );
6178 assert_eq!(finish, "stop");
6179 assert!(message.tool_calls.is_none());
6180 assert!(message.content.is_some());
6181 }
6182
6183 #[test]
6187 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
6188 let opened = std::cell::Cell::new(0usize);
6189 let mut parser = ferrox_edge::ToolCallParser::new(
6190 ferrox_edge::ToolCallFormat::Qwen3Coder,
6191 vec![ferrox_edge::parser::tool_call::ToolSchema::with_parameters(
6192 "write_file",
6193 serde_json::json!({"type": "object", "properties": {
6194 "path": {"type": "string"},
6195 "contents": {"type": "string"}
6196 }}),
6197 )],
6198 );
6199 let wire = "<tool_call><function=write_file>\
6200 <parameter=path>\n/tmp/x\n</parameter>\
6201 <parameter=contents>\nhello world\n</parameter>\
6202 </function></tool_call>";
6203
6204 let mut deltas = Vec::new();
6205 let mut text = String::new();
6206 for piece in wire.as_bytes().chunks(7) {
6207 let chunk = String::from_utf8_lossy(piece).into_owned();
6208 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
6209 text.push_str(&more_text);
6210 deltas.extend(more);
6211 }
6212 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
6213 text.push_str(&more_text);
6214 deltas.extend(more);
6215
6216 assert_eq!(opened.get(), 1, "one call opened");
6217 assert!(text.is_empty(), "the markers are not content: {text:?}");
6218
6219 let first = &deltas[0];
6220 assert_eq!(first.index, 0);
6221 assert_eq!(first.id.as_deref(), Some("call_0"));
6222 assert_eq!(first.kind, Some("function"));
6223 assert_eq!(first.function.name.as_deref(), Some("write_file"));
6224
6225 let joined: String = deltas
6228 .iter()
6229 .filter_map(|d| d.function.arguments.clone())
6230 .collect();
6231 let parsed: serde_json::Value =
6232 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
6233 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
6234 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
6235 assert!(
6236 deltas.len() >= 3,
6237 "the arguments arrived in pieces, not whole: {}",
6238 deltas.len()
6239 );
6240 assert!(
6241 deltas[1..].iter().all(|d| d.function.name.is_none()),
6242 "only the opening delta carries identity"
6243 );
6244 }
6245
6246 #[test]
6248 fn text_around_a_streamed_call_is_still_content() {
6249 let opened = std::cell::Cell::new(0usize);
6250 let mut parser = ferrox_edge::ToolCallParser::new(
6251 ferrox_edge::ToolCallFormat::Qwen25,
6252 vec![ferrox_edge::parser::tool_call::ToolSchema::new(
6253 "get_weather",
6254 )],
6255 );
6256 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
6257 \"arguments\": {}}</tool_call> done";
6258 let mut text = String::new();
6259 for piece in wire.as_bytes().chunks(5) {
6260 let chunk = String::from_utf8_lossy(piece).into_owned();
6261 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
6262 text.push_str(&more);
6263 }
6264 let (more, _) = tool_call_deltas(parser.finish(), &opened);
6265 text.push_str(&more);
6266
6267 assert_eq!(opened.get(), 1);
6268 assert!(text.starts_with("let me check. "), "{text:?}");
6269 assert!(text.ends_with(" done"), "{text:?}");
6270 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
6271 }
6272
6273 #[test]
6276 fn a_reasoning_block_is_split_out_of_the_answer() {
6277 let (message, finish) = build_response_message(
6278 "<think>weighing it up</think>The answer is 4.".to_string(),
6279 &[],
6280 output::OutputPosture::for_model("Qwen3-8B"),
6281 "stop",
6282 );
6283 assert_eq!(finish, "stop");
6284 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
6285 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
6286 }
6287
6288 #[test]
6291 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
6292 let (message, _) = build_response_message(
6293 "Use the <think> tag like this.".to_string(),
6294 &[],
6295 output::OutputPosture::for_model("llama-3.1-8b"),
6296 "stop",
6297 );
6298 assert_eq!(
6299 message.content.as_deref(),
6300 Some("Use the <think> tag like this.")
6301 );
6302 assert!(message.reasoning_content.is_none());
6303 }
6304
6305 #[tokio::test]
6313 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
6314 let app = test_app();
6315 let body = serde_json::json!({
6316 "model": "m",
6317 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6318 "max_tokens": 4,
6319 "temperature": 0,
6320 });
6321 let resp = post_json(&app, body).await;
6322 let message = &resp["choices"][0]["message"];
6323 assert!(message["content"].is_string());
6324 assert!(message.get("tool_calls").is_none());
6325 assert_eq!(resp["choices"][0]["finish_reason"], "length");
6326 assert_eq!(resp["usage"]["completion_tokens"], 4);
6327 assert_eq!(
6328 resp["usage"]["total_tokens"],
6329 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
6330 );
6331 }
6332
6333 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
6334 use http_body_util::BodyExt;
6335 use tower::ServiceExt;
6336
6337 let response = app
6338 .clone()
6339 .oneshot(
6340 axum::http::Request::builder()
6341 .method("GET")
6342 .uri(uri)
6343 .body(axum::body::Body::empty())
6344 .unwrap(),
6345 )
6346 .await
6347 .unwrap();
6348 let status = response.status();
6349 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6350 (status, serde_json::from_slice(&bytes).unwrap())
6351 }
6352
6353 #[tokio::test]
6354 async fn health_answers_a_capability_handshake_not_a_boolean() {
6355 let app = test_app();
6356 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
6357 assert_eq!(status, StatusCode::OK);
6358
6359 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
6360 assert_eq!(health.state, ferrox_api::HealthState::Ready);
6361 assert!(health.pid > 0);
6362 assert!(health.server_time_unix_ms > 0);
6363 assert_eq!(health.last_request_age_seconds, None);
6366
6367 for id in [
6370 ferrox_api::health::capability::CPU,
6371 ferrox_api::health::capability::METAL,
6372 ferrox_api::health::capability::CUDA,
6373 ferrox_api::health::capability::REAL_WEIGHTS,
6374 ferrox_api::health::capability::CONTINUOUS_BATCHING,
6375 ] {
6376 let cap = health
6377 .capability(id)
6378 .unwrap_or_else(|| panic!("{id} missing"));
6379 assert!(!cap.reason.is_empty(), "{cap:?}");
6380 assert!(!cap.detail.is_empty(), "{cap:?}");
6381 }
6382 let weights = health
6386 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
6387 .unwrap();
6388 assert!(!weights.available);
6389 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
6390 assert!(health.model.as_ref().unwrap().synthetic_weights);
6391 }
6392
6393 #[tokio::test]
6394 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
6395 let app = test_app();
6396 let _ = post_json(
6397 &app,
6398 serde_json::json!({
6399 "model": "m",
6400 "messages": [{"role": "user", "content": "\u{1}"}],
6401 "max_tokens": 1,
6402 "temperature": 0,
6403 }),
6404 )
6405 .await;
6406 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
6407 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
6408 let age = health
6409 .last_request_age_seconds
6410 .expect("a served request is evidence of liveness");
6411 assert!((0.0..5.0).contains(&age), "implausible age {age}");
6412 }
6413
6414 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
6416 use http_body_util::BodyExt;
6417 use tower::ServiceExt;
6418
6419 let response = app
6420 .clone()
6421 .oneshot(
6422 axum::http::Request::builder()
6423 .method("POST")
6424 .uri("/v1/chat/completions")
6425 .header("content-type", "application/json")
6426 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6427 .unwrap(),
6428 )
6429 .await
6430 .unwrap();
6431 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6432 String::from_utf8(bytes.to_vec())
6433 .unwrap()
6434 .lines()
6435 .filter_map(|line| line.strip_prefix("data: "))
6436 .filter(|payload| *payload != "[DONE]")
6437 .map(|payload| serde_json::from_str(payload).unwrap())
6438 .collect()
6439 }
6440
6441 #[tokio::test]
6442 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
6443 let app = test_app();
6444 let chunks = post_sse_chunks(
6445 &app,
6446 serde_json::json!({
6447 "model": "m",
6448 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6449 "max_tokens": 4,
6450 "temperature": 0,
6451 "stream": true,
6452 }),
6453 )
6454 .await;
6455
6456 assert!(!chunks.is_empty());
6457 let request_id = chunks[0]["request_id"]
6458 .as_str()
6459 .expect("the first chunk names the request")
6460 .to_string();
6461 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
6462 for (i, chunk) in chunks.iter().enumerate().skip(1) {
6465 assert!(
6466 chunk.get("request_id").is_none(),
6467 "chunk {i} repeats request_id"
6468 );
6469 }
6470 for chunk in &chunks {
6473 assert_eq!(chunk["id"], serde_json::json!(request_id));
6474 }
6475
6476 let other = post_sse_chunks(
6477 &app,
6478 serde_json::json!({
6479 "model": "m",
6480 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6481 "max_tokens": 4,
6482 "temperature": 0,
6483 "stream": true,
6484 }),
6485 )
6486 .await;
6487 assert_ne!(
6488 other[0]["request_id"].as_str().unwrap(),
6489 request_id,
6490 "two concurrent chats must not share an id"
6491 );
6492 }
6493
6494 #[tokio::test]
6495 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
6496 let app = test_app();
6497 let resp = post_json(
6498 &app,
6499 serde_json::json!({
6500 "model": "m",
6501 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6502 "max_tokens": 2,
6503 "temperature": 0,
6504 }),
6505 )
6506 .await;
6507 assert_eq!(resp["id"], resp["request_id"]);
6508 assert!(resp["request_id"]
6509 .as_str()
6510 .unwrap()
6511 .starts_with("chatcmpl-"));
6512 }
6513
6514 #[tokio::test]
6517 async fn usage_carries_separate_prefill_and_decode_timings() {
6518 let app = test_app();
6519 let resp = post_json(
6520 &app,
6521 serde_json::json!({
6522 "model": "m",
6523 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6524 "max_tokens": 4,
6525 "temperature": 0,
6526 }),
6527 )
6528 .await;
6529 let usage = &resp["usage"];
6530 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
6531 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
6532 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
6533 assert!(usage["predicted_per_second"].is_number(), "{usage}");
6534 assert!(usage.get("cached_tokens").is_none(), "{usage}");
6536 }
6537
6538 #[tokio::test]
6546 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
6547 let app = test_app();
6548 let body = serde_json::json!({
6549 "model": "m",
6550 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6551 "max_tokens": 4,
6552 "temperature": 0,
6553 "tools": [weather_tool()],
6554 });
6555 let resp = post_json(&app, body).await;
6556 let message = &resp["choices"][0]["message"];
6557 assert!(
6558 message["content"].is_string(),
6559 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
6560 );
6561 assert!(message.get("tool_calls").is_none());
6562 assert_eq!(resp["choices"][0]["finish_reason"], "length");
6566 }
6567
6568 #[tokio::test]
6572 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
6573 let app = test_app();
6574 let body = serde_json::json!({
6575 "model": "m",
6576 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
6577 "max_tokens": 3,
6578 "temperature": 0,
6579 });
6580 let first = post_json(&app, body.clone()).await;
6581 assert_eq!(first["ferrox_cache"], "miss");
6582 let second = post_json(&app, body).await;
6583 assert_eq!(second["ferrox_cache"], "hit");
6584 assert_eq!(
6585 first["choices"][0]["message"]["content"],
6586 second["choices"][0]["message"]["content"]
6587 );
6588 assert_eq!(
6589 first["choices"][0]["finish_reason"],
6590 second["choices"][0]["finish_reason"]
6591 );
6592 assert_eq!(first["usage"], second["usage"]);
6593 assert_eq!(second["usage"]["completion_tokens"], 3);
6594 }
6595
6596 #[tokio::test]
6602 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
6603 let session_app = test_app();
6604 let manual_app = test_app();
6605
6606 let turn1 = post_json(
6608 &session_app,
6609 serde_json::json!({
6610 "model": "m",
6611 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6612 "session_id": "s1",
6613 "max_tokens": 5,
6614 "temperature": 0,
6615 }),
6616 )
6617 .await;
6618 let reply1 = turn1["choices"][0]["message"]["content"]
6619 .as_str()
6620 .unwrap()
6621 .to_string();
6622
6623 let manual_turn1 = post_json(
6628 &manual_app,
6629 serde_json::json!({
6630 "model": "m",
6631 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6632 "max_tokens": 5,
6633 "temperature": 0,
6634 }),
6635 )
6636 .await;
6637 assert_eq!(
6638 manual_turn1["choices"][0]["message"]["content"]
6639 .as_str()
6640 .unwrap(),
6641 reply1
6642 );
6643
6644 let turn2 = post_json(
6646 &session_app,
6647 serde_json::json!({
6648 "model": "m",
6649 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
6650 "session_id": "s1",
6651 "max_tokens": 5,
6652 "temperature": 0,
6653 }),
6654 )
6655 .await;
6656 let reply2 = turn2["choices"][0]["message"]["content"]
6657 .as_str()
6658 .unwrap()
6659 .to_string();
6660
6661 let manual_turn2 = post_json(
6665 &manual_app,
6666 serde_json::json!({
6667 "model": "m",
6668 "messages": [
6669 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
6670 {"role": "assistant", "content": reply1},
6671 {"role": "user", "content": "\u{4}\u{5}"},
6672 ],
6673 "max_tokens": 5,
6674 "temperature": 0,
6675 }),
6676 )
6677 .await;
6678 assert_eq!(
6679 manual_turn2["choices"][0]["message"]["content"]
6680 .as_str()
6681 .unwrap(),
6682 reply2,
6683 "resuming a session must produce identical output to manually resending the full history"
6684 );
6685 }
6686
6687 #[test]
6690 fn lock_cache_recovers_from_a_poisoned_mutex() {
6691 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
6692
6693 let poison_cache = Arc::clone(&cache);
6694 let _ = std::thread::spawn(move || {
6695 let _guard = poison_cache.lock().unwrap();
6696 panic!("simulated panic while holding the lock");
6697 })
6698 .join();
6699
6700 let recovered = lock_cache(&cache);
6702 assert_eq!(recovered.stats().entries, 0);
6703 }
6704
6705 #[test]
6706 fn is_cacheable_true_for_greedy_or_seeded_requests() {
6707 let mut req_body = serde_json::json!({
6708 "model": "m",
6709 "messages": [{"role": "user", "content": "hi"}],
6710 });
6711 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
6712 assert!(
6713 req.is_cacheable(),
6714 "default (temperature 0) must be cacheable"
6715 );
6716
6717 req_body["temperature"] = serde_json::json!(0.8);
6718 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
6719 assert!(
6720 !req.is_cacheable(),
6721 "unseeded sampling must never be cacheable"
6722 );
6723
6724 req_body["seed"] = serde_json::json!(42);
6725 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
6726 assert!(
6727 req.is_cacheable(),
6728 "sampling with an explicit seed is deterministic and must be cacheable"
6729 );
6730 }
6731
6732 const GRADED: &str = "{% if reasoning_effort %}\
6736 {% if reasoning_effort not in ['low','medium','high'] %}\
6737 {{ raise_exception('unsupported effort') }}\
6738 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
6739 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
6740
6741 fn graded_template() -> chat_template::PromptTemplate {
6742 chat_template::PromptTemplate::from_gguf_metadata(
6743 Some(GRADED),
6744 Some("qwen3"),
6745 false,
6746 None,
6747 None,
6748 )
6749 }
6750
6751 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
6752 serde_json::from_value(value).expect("request")
6753 }
6754
6755 #[test]
6758 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
6759 let req = chat_request(serde_json::json!({
6760 "model": "m",
6761 "messages": [{"role": "user", "content": "hi"}],
6762 }));
6763 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
6764 }
6765
6766 #[test]
6770 fn an_explicit_zero_output_budget_is_a_client_error() {
6771 let req = chat_request(serde_json::json!({
6772 "model": "m",
6773 "messages": [{"role": "user", "content": "hi"}],
6774 "max_tokens": 0,
6775 }));
6776 let (status, body) = req.validate_supported_fields().expect_err("rejected");
6777 assert_eq!(status, StatusCode::BAD_REQUEST);
6778 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
6779 }
6780
6781 #[test]
6784 fn a_request_can_turn_thinking_off() {
6785 let template = graded_template();
6786 for body in [
6787 serde_json::json!({
6788 "model": "m",
6789 "messages": [{"role": "user", "content": "hi"}],
6790 "reasoning_effort": "none",
6791 }),
6792 serde_json::json!({
6793 "model": "m",
6794 "messages": [{"role": "user", "content": "hi"}],
6795 "thinking": {"type": "disabled"},
6796 }),
6797 ] {
6798 let kwargs = chat_request(body).resolve_template_kwargs(&template);
6799 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
6800 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
6801 assert!(!kwargs.contains_key("reasoning_effort"));
6804 }
6805 }
6806
6807 #[test]
6810 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
6811 let template = graded_template();
6812 let kwargs = chat_request(serde_json::json!({
6813 "model": "m",
6814 "messages": [{"role": "user", "content": "hi"}],
6815 "reasoning_effort": "high",
6816 "thinking": {"type": "disabled"},
6817 }))
6818 .resolve_template_kwargs(&template);
6819 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
6820 assert!(!kwargs.contains_key("reasoning_effort"));
6821 }
6822
6823 #[test]
6826 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
6827 let req = chat_request(serde_json::json!({
6828 "model": "m",
6829 "messages": [{"role": "user", "content": "hi"}],
6830 "thinking": {"type": "disable"},
6831 }));
6832 let (status, _) = req.validate_supported_fields().expect_err("rejected");
6833 assert_eq!(status, StatusCode::BAD_REQUEST);
6834 }
6835
6836 #[test]
6839 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
6840 let template = graded_template();
6841 let kwargs = chat_request(serde_json::json!({
6842 "model": "m",
6843 "messages": [{"role": "user", "content": "hi"}],
6844 "reasoning_effort": "none",
6845 "chat_template_kwargs": {"enable_thinking": true},
6846 }))
6847 .resolve_template_kwargs(&template);
6848 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
6849 }
6850
6851 #[test]
6855 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
6856 let template = graded_template();
6857 let req = chat_request(serde_json::json!({
6858 "model": "m",
6859 "messages": [{"role": "user", "content": "hi"}],
6860 "reasoning_effort": "minimal",
6861 }));
6862 let kwargs = req.resolve_template_kwargs(&template);
6863 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
6864 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
6865 assert!(prompt.starts_with("E:low|"), "{prompt}");
6866 }
6867
6868 #[test]
6872 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
6873 let template = graded_template();
6874 let req = chat_request(serde_json::json!({
6875 "model": "m",
6876 "messages": [{"role": "user", "content": "hi"}],
6877 "chat_template_kwargs": {"reasoning_effort": "none"},
6878 }));
6879 let kwargs = req.resolve_template_kwargs(&template);
6880 assert!(!kwargs.contains_key("reasoning_effort"));
6881 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
6882 assert_eq!(prompt, "hi");
6883 }
6884
6885 #[test]
6888 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
6889 let template = graded_template();
6890 let req = chat_request(serde_json::json!({
6891 "model": "m",
6892 "messages": [{"role": "user", "content": "hi"}],
6893 "reasoning_effort": "low",
6894 "chat_template_kwargs": {"reasoning_effort": "high"},
6895 }));
6896 assert_eq!(
6897 req.resolve_template_kwargs(&template)["reasoning_effort"],
6898 serde_json::json!("high")
6899 );
6900 }
6901
6902 #[test]
6906 fn offering_tools_turns_thinking_on_by_itself() {
6907 let template = graded_template();
6908 let quiet = chat_request(serde_json::json!({
6909 "model": "m",
6910 "messages": [{"role": "user", "content": "hi"}],
6911 }));
6912 assert!(!quiet
6913 .resolve_template_kwargs(&template)
6914 .contains_key("enable_thinking"));
6915
6916 let with_tools = chat_request(serde_json::json!({
6917 "model": "m",
6918 "messages": [{"role": "user", "content": "hi"}],
6919 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
6920 }));
6921 let kwargs = with_tools.resolve_template_kwargs(&template);
6922 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
6923 let prompt =
6924 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
6925 assert!(prompt.starts_with("THINK|"), "{prompt}");
6926 }
6927
6928 #[test]
6933 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
6934 let opener = chat_template::PromptTemplate::from_gguf_metadata(
6935 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
6936 Some("qwen3"),
6937 false,
6938 None,
6939 None,
6940 );
6941 let req = chat_request(serde_json::json!({
6942 "model": "m",
6943 "messages": [{"role": "user", "content": "hi"}],
6944 "chat_template_kwargs": {"enable_thinking": true},
6945 }));
6946 let kwargs = req.resolve_template_kwargs(&opener);
6947 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
6948 assert!(prompt.ends_with("<think>"), "{prompt}");
6949
6950 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
6953 let (message, _) = build_response_message(
6954 "weighing it up</think>Paris.".to_string(),
6955 &[],
6956 posture,
6957 "stop",
6958 );
6959 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
6960 assert_eq!(message.content.as_deref(), Some("Paris."));
6961
6962 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
6965 let (message, _) = build_response_message(
6966 "weighing it up</think>Paris.".to_string(),
6967 &[],
6968 closed,
6969 "stop",
6970 );
6971 assert_eq!(message.reasoning_content, None);
6972 }
6973
6974 #[test]
6975 fn stop_param_accepts_both_single_string_and_array() {
6976 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
6977 "model": "m",
6978 "messages": [{"role": "user", "content": "hi"}],
6979 "stop": "END",
6980 }))
6981 .unwrap();
6982 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
6983
6984 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
6985 "model": "m",
6986 "messages": [{"role": "user", "content": "hi"}],
6987 "stop": ["A", "B"],
6988 }))
6989 .unwrap();
6990 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
6991 }
6992
6993 #[test]
6994 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
6995 let model = test_model();
6996 let result = run_generation(
6997 &model,
6998 "hello",
6999 &greedy_params(4),
7000 None,
7001 None,
7002 None,
7003 None,
7004 None,
7005 );
7006 assert!(matches!(
7007 result,
7008 Err(generate::DecodeError::TokenOutOfVocab { .. })
7009 ));
7010 }
7011
7012 #[test]
7016 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
7017 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7019 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
7020
7021 let holder_pool = Arc::clone(&pool);
7022 let holder = std::thread::spawn(move || {
7023 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
7024 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
7026 drop(held);
7027 });
7028 std::thread::sleep(Duration::from_millis(15));
7029
7030 let config = generate::KvPoolConfig {
7031 pool,
7032 queue_wait: Duration::ZERO,
7033 };
7034 let result = run_generation(
7035 &model,
7036 &prompt,
7037 &greedy_params(4),
7038 Some(&config),
7039 None,
7040 None,
7041 None,
7042 None,
7043 );
7044 assert!(matches!(
7045 result,
7046 Err(generate::DecodeError::KvPoolExhausted)
7047 ));
7048
7049 let (status, _body) = decode_error_response(result.unwrap_err());
7050 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7051 holder.join().unwrap();
7052 }
7053
7054 #[test]
7062 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
7063 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7065 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
7067 let config = generate::KvPoolConfig {
7068 pool,
7069 queue_wait: Duration::ZERO,
7070 };
7071
7072 let result = run_generation(
7073 &model,
7074 &prompt,
7075 &greedy_params(4),
7076 Some(&config),
7077 None,
7078 None,
7079 None,
7080 None,
7081 );
7082 let err = result.expect_err("one block cannot hold two layers' caches");
7083 assert!(
7084 matches!(
7085 &err,
7086 generate::DecodeError::KvBudgetExceeded { binding, .. }
7087 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
7088 ),
7089 "expected an immovable device-memory refusal, got {err:?}"
7090 );
7091 let (status, _body) = decode_error_response(err);
7092 assert_eq!(status, StatusCode::BAD_REQUEST);
7093 }
7094
7095 #[test]
7101 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
7102 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
7103 queued: 512,
7104 cap: 512,
7105 });
7106 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7107 assert_eq!(body["error"]["retry_after_seconds"], 1);
7108 let message = body["error"]["message"].as_str().expect("message");
7109 assert!(message.contains("512"), "{message}");
7110 }
7111
7112 #[test]
7113 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
7114 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
7115 token: 99,
7116 vocab_size: 32,
7117 });
7118 assert!(
7119 body["error"]["retry_after_seconds"].is_null(),
7120 "retrying a prompt this model cannot tokenize never helps"
7121 );
7122 }
7123
7124 #[test]
7125 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
7126 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
7127 token: 99,
7128 vocab_size: 32,
7129 });
7130 assert_eq!(status, StatusCode::BAD_REQUEST);
7131 }
7132
7133 #[test]
7134 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
7135 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7137 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
7138 let config = generate::KvPoolConfig {
7139 pool: pool.clone(),
7140 queue_wait: Duration::ZERO,
7141 };
7142
7143 let (_, finish, _usage) = run_generation(
7144 &model,
7145 &prompt,
7146 &greedy_params(4),
7147 Some(&config),
7148 None,
7149 None,
7150 None,
7151 None,
7152 )
7153 .unwrap();
7154 assert_eq!(finish, FinishReason::Length);
7155 assert_eq!(
7156 pool.lock().unwrap().free_blocks(),
7157 2,
7158 "a completed request must return its blocks to the pool"
7159 );
7160 }
7161
7162 #[tokio::test]
7167 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
7168 let model = Arc::new(test_model());
7169 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7170
7171 let mut handles = Vec::new();
7172 for _ in 0..8 {
7173 let model = Arc::clone(&model);
7174 let prompt = prompt.clone();
7175 handles.push(tokio::task::spawn_blocking(move || {
7176 run_generation(
7177 &model,
7178 &prompt,
7179 &greedy_params(6),
7180 None,
7181 None,
7182 None,
7183 None,
7184 None,
7185 )
7186 .unwrap()
7187 }));
7188 }
7189
7190 let mut results = Vec::new();
7191 for h in handles {
7192 results.push(h.await.unwrap());
7193 }
7194 for r in &results[1..] {
7199 assert_eq!(r.0, results[0].0, "decoded chunks must match");
7200 assert_eq!(r.1, results[0].1, "finish reason must match");
7201 assert_eq!(
7202 r.2.prompt_tokens, results[0].2.prompt_tokens,
7203 "prompt token count must match"
7204 );
7205 assert_eq!(
7206 r.2.completion_tokens, results[0].2.completion_tokens,
7207 "completion token count must match"
7208 );
7209 }
7210 }
7211
7212 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
7219 let mut header_entries = Vec::new();
7220 let mut data = Vec::new();
7221 for (name, shape, values) in tensors {
7222 let start = data.len();
7223 for v in values {
7224 data.extend_from_slice(&v.to_le_bytes());
7225 }
7226 let end = data.len();
7227 let shape_str = shape
7228 .iter()
7229 .map(|d| d.to_string())
7230 .collect::<Vec<_>>()
7231 .join(",");
7232 header_entries.push(format!(
7233 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
7234 ));
7235 }
7236 let header = format!("{{{}}}", header_entries.join(","));
7237 let header_bytes = header.as_bytes();
7238 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
7239 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
7240 out.extend_from_slice(header_bytes);
7241 out.extend_from_slice(&data);
7242 out
7243 }
7244
7245 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
7256 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
7257 use ferrox_models::kimi_loader::KimiRealHparams;
7258 use ferrox_moe::{GatingFunction, MoeLayerConfig};
7259
7260 let hidden_dim = 8;
7261 let kda_num_heads = 2;
7262 let kda_head_dim = 3;
7263 let kda_proj = kda_num_heads * kda_head_dim;
7264 let conv_kernel = 4;
7265 let dense_intermediate = 5;
7266 let vocab_size = 256;
7270 let mla_num_heads = 1;
7271 let mla_q_lora_rank = 2;
7272 let mla_kv_lora_rank = 2;
7273 let mla_qk_nope_head_dim = 2;
7274 let mla_qk_rope_head_dim = 2;
7275 let mla_v_head_dim = 2;
7276
7277 let model_cfg = ferrox_models::ModelConfig {
7278 name: "synthetic-kimi-server-test",
7279 n_layers: 1,
7280 hidden_dim,
7281 n_heads: 1,
7282 n_kv_heads: 1,
7283 head_dim: 4,
7284 vocab_size,
7285 rope_theta: 10000.0,
7286 rms_norm_eps: 1e-5,
7287 sliding_window: None,
7288 moe: MoeLayerConfig {
7289 expert_weights_scale: 1.0,
7290 n_experts: 1,
7291 n_experts_active: 1,
7292 n_shared_experts: 0,
7293 hidden_dim,
7294 expert_ffn_dim: 4,
7295 gating: GatingFunction::Sigmoid,
7296 norm_topk_prob: true,
7297 expert_group_count: None,
7298 expert_group_used_count: None,
7299 },
7300 n_dense_leading_layers: 1,
7305 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
7306 kda_layers: vec![1],
7307 full_attn_layers: vec![],
7308 mla: MlaConfig {
7309 num_heads: mla_num_heads,
7310 q_lora_rank: mla_q_lora_rank,
7311 kv_lora_rank: mla_kv_lora_rank,
7312 qk_nope_head_dim: mla_qk_nope_head_dim,
7313 qk_rope_head_dim: mla_qk_rope_head_dim,
7314 v_head_dim: mla_v_head_dim,
7315 use_output_gate: true,
7316 rope: None,
7317 },
7318 kda: KdaConfig {
7319 num_heads: kda_num_heads,
7320 head_dim: kda_head_dim,
7321 short_conv_kernel_size: conv_kernel,
7322 gate_lower_bound: -5.0,
7323 use_full_rank_gate: true,
7324 },
7325 }),
7326 rope_freqs: None,
7327 rope_attn_factor: 1.0,
7328 rope_dim: None,
7329 rope_freqs_long: None,
7330 rope_freqs_short: None,
7331 rope_orig_ctx: None,
7332 rope_layout: ferrox_models::config::RopeLayout::Neox,
7333 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
7334 swa_pattern: None,
7335 attn_logit_softcap: None,
7336 final_logit_softcap: None,
7337 embedding_scale: None,
7338 attention_scale: None,
7339 rope_theta_swa: None,
7340 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
7341 best_effort_fields: &["synthetic test config, not a real preset"],
7342 };
7343 let hp = KimiRealHparams {
7344 hidden_dim,
7345 kda_num_heads,
7346 kda_head_dim,
7347 mla_num_heads,
7348 mla_q_lora_rank,
7349 mla_kv_lora_rank,
7350 mla_qk_nope_head_dim,
7351 mla_qk_rope_head_dim,
7352 mla_v_head_dim,
7353 dense_intermediate_dim: dense_intermediate,
7354 moe_hidden_dim: hidden_dim,
7355 moe_intermediate_dim: 4,
7356 n_experts: 1,
7357 num_shared_experts: 0,
7358 };
7359
7360 let prefix = "language_model.model.layers.0";
7364 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
7365 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
7366 name: String,
7367 shape: Vec<usize>,
7368 n: usize| {
7369 tensors.push((name, shape, vec![0.01f32; n]));
7370 };
7371 push(
7372 &mut tensors,
7373 format!("{prefix}.input_layernorm.weight"),
7374 vec![hidden_dim],
7375 hidden_dim,
7376 );
7377 push(
7378 &mut tensors,
7379 format!("{prefix}.post_attention_layernorm.weight"),
7380 vec![hidden_dim],
7381 hidden_dim,
7382 );
7383 push(
7384 &mut tensors,
7385 format!("{prefix}.self_attention_res_norm.weight"),
7386 vec![hidden_dim],
7387 hidden_dim,
7388 );
7389 push(
7390 &mut tensors,
7391 format!("{prefix}.self_attention_res_proj.weight"),
7392 vec![1, hidden_dim],
7393 hidden_dim,
7394 );
7395 push(
7396 &mut tensors,
7397 format!("{prefix}.mlp_res_norm.weight"),
7398 vec![hidden_dim],
7399 hidden_dim,
7400 );
7401 push(
7402 &mut tensors,
7403 format!("{prefix}.mlp_res_proj.weight"),
7404 vec![1, hidden_dim],
7405 hidden_dim,
7406 );
7407 push(
7408 &mut tensors,
7409 format!("{prefix}.self_attn.q_proj.weight"),
7410 vec![kda_proj, hidden_dim],
7411 kda_proj * hidden_dim,
7412 );
7413 push(
7414 &mut tensors,
7415 format!("{prefix}.self_attn.k_proj.weight"),
7416 vec![kda_proj, hidden_dim],
7417 kda_proj * hidden_dim,
7418 );
7419 push(
7420 &mut tensors,
7421 format!("{prefix}.self_attn.v_proj.weight"),
7422 vec![kda_proj, hidden_dim],
7423 kda_proj * hidden_dim,
7424 );
7425 push(
7426 &mut tensors,
7427 format!("{prefix}.self_attn.q_conv1d.weight"),
7428 vec![kda_proj, 1, conv_kernel],
7429 kda_proj * conv_kernel,
7430 );
7431 push(
7432 &mut tensors,
7433 format!("{prefix}.self_attn.k_conv1d.weight"),
7434 vec![kda_proj, 1, conv_kernel],
7435 kda_proj * conv_kernel,
7436 );
7437 push(
7438 &mut tensors,
7439 format!("{prefix}.self_attn.v_conv1d.weight"),
7440 vec![kda_proj, 1, conv_kernel],
7441 kda_proj * conv_kernel,
7442 );
7443 push(
7444 &mut tensors,
7445 format!("{prefix}.self_attn.A_log"),
7446 vec![kda_num_heads],
7447 kda_num_heads,
7448 );
7449 push(
7450 &mut tensors,
7451 format!("{prefix}.self_attn.f_a_proj.weight"),
7452 vec![kda_head_dim, hidden_dim],
7453 kda_head_dim * hidden_dim,
7454 );
7455 push(
7456 &mut tensors,
7457 format!("{prefix}.self_attn.f_b_proj.weight"),
7458 vec![kda_proj, kda_head_dim],
7459 kda_proj * kda_head_dim,
7460 );
7461 push(
7462 &mut tensors,
7463 format!("{prefix}.self_attn.dt_bias"),
7464 vec![kda_proj],
7465 kda_proj,
7466 );
7467 push(
7468 &mut tensors,
7469 format!("{prefix}.self_attn.b_proj.weight"),
7470 vec![kda_num_heads, hidden_dim],
7471 kda_num_heads * hidden_dim,
7472 );
7473 push(
7474 &mut tensors,
7475 format!("{prefix}.self_attn.g_proj.weight"),
7476 vec![kda_proj, hidden_dim],
7477 kda_proj * hidden_dim,
7478 );
7479 push(
7480 &mut tensors,
7481 format!("{prefix}.self_attn.o_norm.weight"),
7482 vec![kda_head_dim],
7483 kda_head_dim,
7484 );
7485 push(
7486 &mut tensors,
7487 format!("{prefix}.self_attn.o_proj.weight"),
7488 vec![hidden_dim, kda_proj],
7489 hidden_dim * kda_proj,
7490 );
7491 push(
7492 &mut tensors,
7493 format!("{prefix}.mlp.gate_proj.weight"),
7494 vec![dense_intermediate, hidden_dim],
7495 dense_intermediate * hidden_dim,
7496 );
7497 push(
7498 &mut tensors,
7499 format!("{prefix}.mlp.up_proj.weight"),
7500 vec![dense_intermediate, hidden_dim],
7501 dense_intermediate * hidden_dim,
7502 );
7503 push(
7504 &mut tensors,
7505 format!("{prefix}.mlp.down_proj.weight"),
7506 vec![hidden_dim, dense_intermediate],
7507 hidden_dim * dense_intermediate,
7508 );
7509 push(
7510 &mut tensors,
7511 "language_model.model.embed_tokens.weight".to_string(),
7512 vec![vocab_size, hidden_dim],
7513 vocab_size * hidden_dim,
7514 );
7515 push(
7516 &mut tensors,
7517 "language_model.lm_head.weight".to_string(),
7518 vec![vocab_size, hidden_dim],
7519 vocab_size * hidden_dim,
7520 );
7521 push(
7522 &mut tensors,
7523 "language_model.model.norm.weight".to_string(),
7524 vec![hidden_dim],
7525 hidden_dim,
7526 );
7527 push(
7528 &mut tensors,
7529 "language_model.model.output_attn_res_norm.weight".to_string(),
7530 vec![hidden_dim],
7531 hidden_dim,
7532 );
7533 push(
7534 &mut tensors,
7535 "language_model.model.output_attn_res_proj.weight".to_string(),
7536 vec![1, hidden_dim],
7537 hidden_dim,
7538 );
7539
7540 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
7553 let dir = std::env::temp_dir().join(format!(
7554 "ferrox_server_kimi_e2e_test_{}_{}",
7555 std::process::id(),
7556 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
7557 ));
7558 std::fs::create_dir_all(&dir).unwrap();
7559 let shard_bytes = write_safetensors_shard(&tensors);
7560 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
7561 let map_entries: Vec<String> = tensors
7562 .iter()
7563 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
7564 .collect();
7565 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
7566 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
7567
7568 use base64::Engine;
7572 let vocab_lines: Vec<String> = (0..vocab_size as u32)
7573 .map(|b| {
7574 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
7575 format!("{b64} {b}")
7576 })
7577 .collect();
7578 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
7579
7580 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
7581 .expect("must load the synthetic Kimi checkpoint end to end");
7582 std::fs::remove_dir_all(&dir).ok();
7583 loaded
7584 }
7585
7586 #[test]
7594 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
7595 let loaded = build_synthetic_kimi_loaded();
7596 let state = build_app_state(
7597 model::LoadedModel::Kimi(loaded),
7598 None,
7599 None,
7600 None,
7601 false,
7602 None,
7603 Arc::new(health::Detection::ready(health::probe_backends())),
7604 );
7605 let active = state.active().expect("a freshly built state has a model");
7606 assert_eq!(active.model.tokenizer_kind(), "kimi-tiktoken-bpe");
7607 assert!(!active.model.is_synthetic());
7608
7609 let (_chunks, finish, _usage) = run_generation(
7610 &active.model,
7611 "hi",
7612 &greedy_params(5),
7613 None,
7614 None,
7615 None,
7616 None,
7617 None,
7618 )
7619 .expect("a real Kimi checkpoint must generate without error");
7620 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
7621 }
7622
7623 #[test]
7633 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
7634 let loaded = build_synthetic_kimi_loaded();
7635 let state = build_app_state(
7636 model::LoadedModel::Kimi(loaded),
7637 None,
7638 None,
7639 None,
7640 false,
7641 None,
7642 Arc::new(health::Detection::ready(health::probe_backends())),
7643 );
7644
7645 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
7646 let kv_pool_config = generate::KvPoolConfig {
7647 pool: pool.clone(),
7648 queue_wait: Duration::ZERO,
7649 };
7650 let pc = Mutex::new(PrefixCache::new(4));
7651
7652 run_generation(
7653 &state
7654 .active()
7655 .expect("a freshly built state has a model")
7656 .model,
7657 "hi",
7658 &greedy_params(5),
7659 Some(&kv_pool_config),
7660 None,
7661 Some(&pc),
7662 None,
7663 None,
7664 )
7665 .expect("a real Kimi checkpoint must generate without error");
7666
7667 assert_eq!(
7668 pool.lock().unwrap().free_blocks(),
7669 4,
7670 "the KV pool must be completely untouched by a Kimi request"
7671 );
7672 let stats = pc.lock().unwrap().stats();
7673 assert_eq!(
7674 stats.hits + stats.misses,
7675 0,
7676 "the prefix cache must never be consulted for a Kimi request"
7677 );
7678 }
7679}