1mod admin;
38mod anthropic;
39mod attribution;
40mod batch_scheduler;
41mod budget;
42mod cache;
43mod cancel;
44mod chat_template;
45mod generate;
46mod health;
47mod hub;
48mod journal;
49mod json_mode;
50mod limits;
51mod mcp;
52mod model;
53mod openai_extra;
54mod resume;
55mod security;
56mod session;
57mod sse;
58mod stats;
59mod stop;
60mod tasks;
61
62use std::convert::Infallible;
63use std::fmt;
64use std::net::{IpAddr, Ipv4Addr, SocketAddr};
65use std::path::PathBuf;
66use std::str::FromStr;
67use std::sync::{Arc, Mutex, MutexGuard};
68use std::time::Duration;
69
70use axum::{
71 extract::State,
72 http::StatusCode,
73 response::sse::{Event, KeepAlive, Sse},
74 response::{IntoResponse, Response},
75 routing::{get, post},
76 Json, Router,
77};
78use clap::{Parser, ValueEnum};
79use serde::{Deserialize, Serialize};
80
81use cache::{CacheKey, ResponseCache};
82use ferrox_core::cache::KvBlockPool;
83use ferrox_models::kimi_tokenizer::KimiTokenizer;
84use ferrox_models::sampling::SamplingParams;
85use ferrox_models::tokenizer::StopTokens;
86use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
87use generate::{FinishReason, GenerationParams};
88use model::ServerTokenizer;
89
90#[derive(Parser, Debug, PartialEq)]
94#[command(
101 name = "ferrox-server",
102 about = "OpenAI-compatible Ferrox inference server"
103)]
104pub struct ServerArgs {
105 #[arg(short = 'm', long = "model", value_name = "FILE")]
107 model: Option<String>,
108
109 #[arg(long, value_name = "HOST")]
111 host: Option<IpAddr>,
112
113 #[arg(long, value_name = "PORT")]
118 port: Option<u16>,
119
120 #[arg(short = 't', long = "threads", value_name = "N")]
122 threads: Option<usize>,
123
124 #[arg(
126 long = "device",
127 visible_alias = "dev",
128 value_name = "DEVICE",
129 ignore_case = true
130 )]
131 device: Option<OffloadDevice>,
132
133 #[arg(long = "list-devices", default_value_t = false)]
135 list_devices: bool,
136
137 #[arg(
142 long = "n-gpu-layers",
143 visible_aliases = ["gpu-layers", "ngl"],
144 value_name = "N"
145 )]
146 n_gpu_layers: Option<GpuLayers>,
147
148 #[arg(long = "mcp-config", value_name = "PATH")]
150 mcp_config: Option<PathBuf>,
151
152 #[arg(long = "exit-on-stdin-close", default_value_t = false)]
160 exit_on_stdin_close: bool,
161
162 #[arg(long = "allow-multiple-instances", default_value_t = false)]
167 allow_multiple_instances: bool,
168}
169
170impl ServerArgs {
171 pub fn parse_llama_style<I>(argv: I) -> Self
178 where
179 I: IntoIterator<Item = String>,
180 {
181 Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
182 }
183}
184
185#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
186enum OffloadDevice {
187 Auto,
188 None,
189 Cpu,
190 Metal,
191 Cuda,
192}
193
194#[derive(Debug, Clone, Copy, PartialEq, Eq)]
195enum GpuLayers {
196 Auto,
197 All,
198 Count(u32),
199}
200
201impl GpuLayers {
202 fn offload_enabled(self) -> bool {
203 !matches!(self, Self::Count(0))
204 }
205}
206
207impl FromStr for GpuLayers {
208 type Err = String;
209
210 fn from_str(value: &str) -> Result<Self, Self::Err> {
211 match value {
212 "auto" => Ok(Self::Auto),
213 "all" => Ok(Self::All),
214 _ => value
215 .parse::<u32>()
216 .map(Self::Count)
217 .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
218 }
219 }
220}
221
222impl fmt::Display for GpuLayers {
223 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
224 match self {
225 Self::Auto => f.write_str("auto"),
226 Self::All => f.write_str("all"),
227 Self::Count(value) => value.fmt(f),
228 }
229 }
230}
231
232pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
241
242pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
245
246fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
247 args.into_iter()
248 .map(|arg| match arg.as_str() {
249 "-ngl" => "--n-gpu-layers".into(),
250 "-dev" => "--device".into(),
251 _ => arg,
252 })
253 .collect()
254}
255
256fn print_available_devices() {
257 println!("Available devices:");
258 println!(" CPU");
259
260 let metal = ferrox_metal::MetalProfile::detect();
261 if let Some(name) = metal.device_name {
262 println!(" Metal: {name}");
263 }
264
265 let cuda = ferrox_cuda::HardwareProfile::detect();
266 if cuda.cuda_available {
267 let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
268 println!(" CUDA: {name}");
269 if cuda.cuda_device_count > 1 {
270 println!(" ({} devices detected)", cuda.cuda_device_count);
271 }
272 }
273}
274
275fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
276 if args.host.is_none() && args.port.is_none() {
277 return None;
278 }
279
280 let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
281 let host = args
282 .host
283 .or_else(|| existing.map(|addr| addr.ip()))
284 .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
285 let port = args
286 .port
287 .or_else(|| existing.map(|addr| addr.port()))
288 .unwrap_or(8383);
289 Some(SocketAddr::new(host, port).to_string())
290}
291
292fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
293 if let Some(model) = &args.model {
294 unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
296 }
297
298 if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
299 unsafe { std::env::set_var("FERROX_ADDR", addr) };
301 }
302
303 if let Some(threads) = args.threads {
304 if threads == 0 {
305 anyhow::bail!("--threads must be greater than zero");
306 }
307 unsafe {
309 std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
310 std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
311 }
312 }
313
314 if args.device.is_none() && args.n_gpu_layers.is_none() {
315 } else {
317 let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
318 let device = if layers.offload_enabled() {
319 args.device.unwrap_or(OffloadDevice::Auto)
320 } else {
321 OffloadDevice::None
322 };
323
324 match device {
325 OffloadDevice::None | OffloadDevice::Cpu => unsafe {
326 std::env::set_var("FERROX_METAL", "0");
327 std::env::set_var("FERROX_METAL_ATTN", "0");
328 std::env::set_var("FERROX_CUDA", "0");
329 },
330 OffloadDevice::Auto => unsafe {
331 std::env::set_var("FERROX_METAL", "auto");
332 std::env::set_var("FERROX_CUDA", "auto");
333 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
334 std::env::set_var("FERROX_METAL_ATTN", "1");
335 }
336 },
337 OffloadDevice::Metal => {
338 #[cfg(not(feature = "metal"))]
339 {
340 anyhow::bail!(
341 "Metal requested but this binary was built without --features metal"
342 );
343 }
344 #[cfg(feature = "metal")]
345 {
346 if !ferrox_metal::MetalProfile::detect().available {
347 anyhow::bail!("Metal requested but no Metal device is available");
348 }
349 unsafe {
350 std::env::set_var("FERROX_METAL", "1");
351 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
352 std::env::set_var("FERROX_METAL_ATTN", "1");
353 }
354 std::env::set_var("FERROX_CUDA", "0");
355 }
356 }
357 }
358 OffloadDevice::Cuda => {
359 #[cfg(not(feature = "cuda"))]
360 {
361 anyhow::bail!(
362 "CUDA requested but this binary was built without --features cuda"
363 );
364 }
365 #[cfg(feature = "cuda")]
366 {
367 if !ferrox_cuda::HardwareProfile::detect().cuda_available {
368 anyhow::bail!("CUDA requested but no CUDA device is available");
369 }
370 unsafe {
371 std::env::set_var("FERROX_CUDA", "1");
372 std::env::set_var("FERROX_METAL", "0");
373 std::env::set_var("FERROX_METAL_ATTN", "0");
374 }
375 }
376 }
377 }
378 }
379
380 Ok(())
381}
382
383#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
392 Gguf(GgufModel),
393 Kimi(KimiModel),
394 Mla(MlaModel),
395 Gemma4(Gemma4Model),
396 Glm52(Glm52Model),
397}
398
399pub(crate) struct GgufModel {
400 decoder: Arc<Decoder>,
401 tokenizer: Arc<ServerTokenizer>,
402 stop_tokens: StopTokens,
403 bos_id: Option<usize>,
404 is_synthetic: bool,
405 chat_template: chat_template::ChatTemplate,
406}
407
408pub(crate) struct KimiModel {
409 engine: KimiEngine,
410 tokenizer: KimiTokenizer,
411 stop_tokens: StopTokens,
412 chat_template: chat_template::ChatTemplate,
413}
414
415pub(crate) struct MlaModel {
416 engine: MlaEngine,
417 tokenizer: ServerTokenizer,
418 stop_tokens: StopTokens,
419 bos_id: Option<usize>,
420 name: String,
421 chat_template: chat_template::ChatTemplate,
422}
423
424pub(crate) struct Gemma4Model {
425 engine: Gemma4Engine,
426 tokenizer: ServerTokenizer,
427 stop_tokens: StopTokens,
428 bos_id: Option<usize>,
429 name: String,
430 chat_template: chat_template::ChatTemplate,
431}
432
433pub(crate) struct Glm52Model {
434 engine: ferrox_models::Glm52Engine,
435 tokenizer: ServerTokenizer,
436 stop_tokens: StopTokens,
437 bos_id: Option<usize>,
438 name: String,
439 chat_template: chat_template::ChatTemplate,
440}
441
442impl Model {
443 pub(crate) fn chat_template(&self) -> chat_template::ChatTemplate {
444 match self {
445 Model::Gguf(m) => m.chat_template,
446 Model::Kimi(m) => m.chat_template,
447 Model::Mla(m) => m.chat_template,
448 Model::Gemma4(m) => m.chat_template,
449 Model::Glm52(m) => m.chat_template,
450 }
451 }
452
453 fn is_synthetic(&self) -> bool {
458 match self {
459 Model::Gguf(m) => m.is_synthetic,
460 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
461 }
462 }
463
464 fn tokenizer_kind(&self) -> &'static str {
465 match self {
466 Model::Gguf(m) => m.tokenizer.kind(),
467 Model::Kimi(_) => "kimi-tiktoken-bpe",
468 Model::Mla(m) => m.tokenizer.kind(),
469 Model::Gemma4(m) => m.tokenizer.kind(),
470 Model::Glm52(m) => m.tokenizer.kind(),
471 }
472 }
473
474 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
478 match self {
479 Model::Gguf(m) => m.decoder.expert_store_stats(),
480 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
481 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
482 }
483 }
484
485 pub(crate) fn name(&self) -> &str {
486 match self {
487 Model::Gguf(m) => m.decoder.config.name,
488 Model::Kimi(_) => "kimi-k3",
489 Model::Mla(m) => m.name.as_str(),
490 Model::Gemma4(m) => m.name.as_str(),
491 Model::Glm52(m) => m.name.as_str(),
492 }
493 }
494
495 pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
496 match self {
497 Model::Gguf(m) => m.tokenizer.encode(text),
498 Model::Kimi(m) => m
499 .tokenizer
500 .encode(text)
501 .into_iter()
502 .map(|id| id as usize)
503 .collect(),
504 Model::Mla(m) => m.tokenizer.encode(text),
505 Model::Gemma4(m) => m.tokenizer.encode(text),
506 Model::Glm52(m) => m.tokenizer.encode(text),
507 }
508 }
509
510 pub(crate) fn decode(&self, ids: &[usize]) -> String {
511 match self {
512 Model::Gguf(m) => m.tokenizer.decode(ids),
513 Model::Kimi(m) => {
514 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
515 m.tokenizer.decode(&ids32)
516 }
517 Model::Mla(m) => m.tokenizer.decode(ids),
518 Model::Gemma4(m) => m.tokenizer.decode(ids),
519 Model::Glm52(m) => m.tokenizer.decode(ids),
520 }
521 }
522
523 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
526 match self {
527 Model::Gguf(m) => {
528 let mut caches: Vec<_> = (0..m.decoder.layers.len())
529 .map(|_| {
530 ferrox_core::cache::KvCache::new(
531 m.decoder.config.n_kv_heads,
532 m.decoder.config.head_dim,
533 )
534 })
535 .collect();
536 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
537 }
538 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
539 }
540 }
541
542 pub(crate) fn vocab_size(&self) -> Option<usize> {
543 match self {
544 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
545 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
546 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
547 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
548 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
549 }
550 }
551}
552
553pub(crate) struct ActiveModel {
562 pub(crate) id: Option<String>,
566 pub(crate) model: Arc<Model>,
567 pub(crate) batcher: Option<batch_scheduler::ContinuousBatcher>,
572 pub(crate) ceiling: Option<Arc<budget::ContextCeiling>>,
582}
583
584pub(crate) struct AppState {
585 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
604 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
608 pub(crate) tasks: Arc<tasks::TaskRegistry>,
610 pub(crate) cancels: Arc<cancel::CancelRegistry>,
614 pub(crate) stats: stats::Stats,
617 pub(crate) streams: resume::StreamRegistry,
620 pub(crate) model_dir: Option<PathBuf>,
622 response_cache: Mutex<ResponseCache>,
626 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
638 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
648 sessions: session::SessionStore,
654 requests_total: std::sync::atomic::AtomicU64,
655 request_errors_total: std::sync::atomic::AtomicU64,
656 started_at: std::time::Instant,
657 last_request_ms: std::sync::atomic::AtomicU64,
662 detection: Arc<health::Detection>,
664 mcp: Option<mcp::LoadedMcpConfig>,
666 pub(crate) continuous_batching_enabled: bool,
670 loading_model: Mutex<Option<String>>,
674 last_load_error: Mutex<Option<(String, String)>>,
678}
679
680impl AppState {
681 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
685 self.active
686 .read()
687 .unwrap_or_else(|p| p.into_inner())
688 .clone()
689 }
690
691 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
697 self.active().ok_or_else(|| {
698 (
699 StatusCode::SERVICE_UNAVAILABLE,
700 Json(serde_json::json!({"error": {
701 "message": "no model is loaded; POST /admin/models/load with an id from \
702 GET /admin/models",
703 "type": "model_not_loaded"
704 }})),
705 )
706 })
707 }
708
709 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
712 Ok(Arc::clone(&self.require_active()?.model))
713 }
714
715 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
723 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
724 std::mem::replace(&mut *guard, next)
725 }
726
727 fn mark_request_finished(&self) {
731 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
732 self.last_request_ms
733 .store(ms, std::sync::atomic::Ordering::Relaxed);
734 }
735
736 pub(crate) fn uptime(&self) -> Duration {
737 self.started_at.elapsed()
738 }
739
740 pub(crate) fn requests_total(&self) -> u64 {
741 self.requests_total
742 .load(std::sync::atomic::Ordering::Relaxed)
743 }
744
745 pub(crate) fn errors_total(&self) -> u64 {
746 self.request_errors_total
747 .load(std::sync::atomic::Ordering::Relaxed)
748 }
749
750 pub(crate) fn cache_stats(&self) -> cache::CacheStats {
751 lock_cache(&self.response_cache).stats()
752 }
753
754 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
757 let last = self
758 .last_request_ms
759 .load(std::sync::atomic::Ordering::Relaxed);
760 (last > 0)
761 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
762 .map(|age| age.max(0.0))
763 }
764
765 pub(crate) fn loading_model_id(&self) -> Option<String> {
766 self.loading_model
767 .lock()
768 .unwrap_or_else(|p| p.into_inner())
769 .clone()
770 }
771
772 pub(crate) fn set_loading_model(&self, id: Option<String>) {
773 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
774 }
775
776 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
777 self.last_load_error
778 .lock()
779 .unwrap_or_else(|p| p.into_inner())
780 .clone()
781 }
782
783 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
784 *self
785 .last_load_error
786 .lock()
787 .unwrap_or_else(|p| p.into_inner()) = error;
788 }
789
790 pub(crate) fn active_model_name(&self) -> Option<String> {
799 self.active().map(|a| a.model.name().to_string())
800 }
801
802 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
803 self.stats.record(stats::entry(record));
804 }
805}
806
807fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
812 cache
813 .lock()
814 .unwrap_or_else(|poisoned| poisoned.into_inner())
815}
816
817#[derive(Debug, Clone, Deserialize)]
818#[serde(untagged)]
819pub(crate) enum MessageContent {
820 Text(String),
821 Parts(Vec<ContentPart>),
822}
823
824#[derive(Debug, Clone, Deserialize)]
825struct ContentPart {
826 #[serde(rename = "type")]
827 kind: String,
828 #[serde(default)]
829 text: Option<String>,
830 #[serde(default)]
831 image_url: Option<serde_json::Value>,
832}
833
834impl MessageContent {
835 fn as_text(&self) -> String {
836 match self {
837 Self::Text(s) => s.clone(),
838 Self::Parts(parts) => parts
839 .iter()
840 .filter_map(|p| p.text.as_deref())
841 .collect::<Vec<_>>()
842 .join(""),
843 }
844 }
845
846 fn has_image(&self) -> bool {
847 match self {
848 Self::Text(_) => false,
849 Self::Parts(parts) => parts
850 .iter()
851 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
852 }
853 }
854}
855
856#[derive(Debug, Clone, Deserialize)]
857pub(crate) struct ChatMessage {
858 pub(crate) role: String,
859 #[serde(default)]
863 pub(crate) content: Option<MessageContent>,
864 #[serde(default)]
868 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
869 #[serde(default)]
874 #[allow(dead_code)]
875 pub(crate) tool_call_id: Option<String>,
876}
877
878impl ChatMessage {
879 fn rendered_content(&self) -> String {
886 let mut out = self
887 .content
888 .as_ref()
889 .map(MessageContent::as_text)
890 .unwrap_or_default();
891 if let Some(calls) = &self.tool_calls {
892 for call in calls {
893 out.push_str(&format!(
894 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
895 call.function.name, call.function.arguments
896 ));
897 }
898 }
899 out
900 }
901}
902
903#[derive(Debug, Clone, Deserialize)]
904pub(crate) struct ToolCallIn {
905 #[serde(default)]
906 #[allow(dead_code)]
907 id: String,
908 #[serde(rename = "type", default)]
909 #[allow(dead_code)]
910 kind: String,
911 function: ToolCallFunctionIn,
912}
913
914#[derive(Debug, Clone, Deserialize)]
915struct ToolCallFunctionIn {
916 name: String,
917 arguments: String,
922}
923
924#[derive(Debug, Clone, Deserialize)]
927struct ToolDef {
928 #[serde(rename = "type", default)]
929 #[allow(dead_code)]
930 kind: String,
931 function: ToolFunctionDef,
932}
933
934#[derive(Debug, Clone, Deserialize)]
935struct ToolFunctionDef {
936 name: String,
937 #[serde(default)]
938 description: Option<String>,
939 #[serde(default)]
940 parameters: Option<serde_json::Value>,
941}
942
943#[derive(Debug, Clone, Deserialize)]
953#[serde(untagged)]
954enum ToolChoice {
955 Mode(String),
956 #[allow(dead_code)]
957 Specific(serde_json::Value),
958}
959
960#[derive(Deserialize)]
963#[serde(untagged)]
964enum StopParam {
965 One(String),
966 Many(Vec<String>),
967}
968
969#[derive(Deserialize)]
970struct ChatCompletionRequest {
971 model: String,
972 messages: Vec<ChatMessage>,
973 #[serde(default = "default_max_tokens")]
974 max_tokens: usize,
975 #[serde(default)]
976 temperature: Option<f32>,
977 #[serde(default)]
978 top_p: Option<f32>,
979 #[serde(default)]
980 top_k: Option<usize>,
981 #[serde(default)]
982 repetition_penalty: Option<f32>,
983 #[serde(default)]
984 seed: Option<u64>,
985 #[serde(default)]
986 stop: Option<StopParam>,
987 #[serde(default)]
988 stream: Option<bool>,
989 #[serde(default)]
1002 stream_resumable: Option<bool>,
1003 #[serde(default)]
1004 tools: Vec<ToolDef>,
1005 #[serde(default)]
1006 tool_choice: Option<ToolChoice>,
1007 #[serde(default)]
1012 session_id: Option<String>,
1013 #[serde(default)]
1015 logprobs: Option<bool>,
1016 #[serde(default)]
1017 top_logprobs: Option<u32>,
1018 #[serde(default)]
1019 n: Option<u32>,
1020 #[serde(default)]
1021 presence_penalty: Option<f32>,
1022 #[serde(default)]
1023 frequency_penalty: Option<f32>,
1024 #[serde(default)]
1025 response_format: Option<serde_json::Value>,
1026}
1027
1028fn default_max_tokens() -> usize {
1029 16
1030}
1031
1032impl ChatCompletionRequest {
1033 fn sampling_params(&self) -> SamplingParams {
1034 SamplingParams {
1035 temperature: self.temperature.unwrap_or(0.0),
1036 top_p: self.top_p.unwrap_or(1.0),
1037 top_k: self.top_k.unwrap_or(0),
1038 repetition_penalty: self.repetition_penalty.unwrap_or(1.0),
1039 presence_penalty: self.presence_penalty.unwrap_or(0.0),
1040 frequency_penalty: self.frequency_penalty.unwrap_or(0.0),
1041 }
1042 }
1043
1044 fn stop_sequences(&self) -> Vec<String> {
1045 self.stop
1046 .as_ref()
1047 .map(|s| match s {
1048 StopParam::One(v) => vec![v.clone()],
1049 StopParam::Many(v) => v.clone(),
1050 })
1051 .unwrap_or_default()
1052 }
1053
1054 fn tools_active(&self) -> bool {
1059 !self.tools.is_empty()
1060 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1061 }
1062
1063 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1066 for msg in &self.messages {
1067 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1068 return Err(unsupported_feature(
1069 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1070 ));
1071 }
1072 }
1073 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1074 return Err(unsupported_feature(
1075 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1076 ));
1077 }
1078 if self.n.is_some_and(|n| n > 1) {
1079 return Err(unsupported_feature(
1080 "n > 1 is not implemented (single completion only)",
1081 ));
1082 }
1083 if let Some(fmt) = &self.response_format {
1084 match fmt.get("type").and_then(|v| v.as_str()) {
1085 Some("json_object") => {}
1086 Some(other) => {
1087 return Err((
1088 StatusCode::BAD_REQUEST,
1089 Json(serde_json::json!({
1090 "error": {
1091 "message": format!(
1092 "response_format type {other:?} is not supported (only json_object)"
1093 )
1094 }
1095 })),
1096 ));
1097 }
1098 None => {
1099 return Err((
1100 StatusCode::BAD_REQUEST,
1101 Json(serde_json::json!({
1102 "error": {
1103 "message": "response_format must include \"type\" (only json_object is supported)"
1104 }
1105 })),
1106 ));
1107 }
1108 }
1109 }
1110 match &self.tool_choice {
1111 Some(ToolChoice::Mode(m)) if m == "required" => {
1112 return Err(unsupported_feature(
1113 "tool_choice=required needs constrained decoding (not implemented)",
1114 ));
1115 }
1116 Some(ToolChoice::Specific(_)) => {
1117 return Err(unsupported_feature(
1118 "named tool_choice is not implemented (use auto/none)",
1119 ));
1120 }
1121 _ => {}
1122 }
1123 Ok(())
1124 }
1125
1126 fn effective_stop_sequences(&self) -> Vec<String> {
1133 let mut stop = self.stop_sequences();
1134 if self.tools_active() {
1135 stop.push("</tool_call>".to_string());
1136 }
1137 stop
1138 }
1139
1140 fn json_object_mode(&self) -> bool {
1141 self.response_format
1142 .as_ref()
1143 .and_then(|v| v.get("type"))
1144 .and_then(|v| v.as_str())
1145 == Some("json_object")
1146 }
1147
1148 fn generation_params(&self) -> GenerationParams {
1149 GenerationParams {
1150 max_tokens: self.max_tokens,
1151 sampling: self.sampling_params(),
1152 seed: self.resolved_seed(),
1153 stop: self.effective_stop_sequences(),
1154 stop_token_ids: Vec::new(),
1158 json_object: self.json_object_mode(),
1159 cancel: None,
1162 }
1163 }
1164
1165 fn generation_params_for_template(
1168 &self,
1169 template: chat_template::ChatTemplate,
1170 ) -> GenerationParams {
1171 let mut params = self.generation_params();
1172 if matches!(
1173 template,
1174 chat_template::ChatTemplate::Gemma | chat_template::ChatTemplate::Gemma4
1175 ) {
1176 let stop = match template {
1177 chat_template::ChatTemplate::Gemma => "<end_of_turn>",
1178 chat_template::ChatTemplate::Gemma4 => "<turn|>",
1179 _ => unreachable!(),
1180 };
1181 if !params.stop.iter().any(|s| s == stop) {
1182 params.stop.push(stop.to_string());
1183 }
1184 }
1185 params
1186 }
1187
1188 fn is_cacheable(&self) -> bool {
1197 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1198 }
1199
1200 fn cache_key(&self, prompt: &str) -> CacheKey {
1201 CacheKey {
1202 model: self.model.clone(),
1203 prompt: prompt.to_string(),
1204 max_tokens: self.max_tokens,
1205 temperature_bits: self.temperature.unwrap_or(0.0).to_bits(),
1206 top_p_bits: self.top_p.unwrap_or(1.0).to_bits(),
1207 top_k: self.top_k.unwrap_or(0),
1208 repetition_penalty_bits: self.repetition_penalty.unwrap_or(1.0).to_bits(),
1209 presence_penalty_bits: self.presence_penalty.unwrap_or(0.0).to_bits(),
1210 frequency_penalty_bits: self.frequency_penalty.unwrap_or(0.0).to_bits(),
1211 seed: self.seed,
1212 stop: self.effective_stop_sequences(),
1213 }
1214 }
1215
1216 fn resolved_seed(&self) -> u64 {
1217 self.seed.unwrap_or_else(|| {
1218 std::time::SystemTime::now()
1219 .duration_since(std::time::UNIX_EPOCH)
1220 .map(|d| d.as_nanos() as u64)
1221 .unwrap_or(0xDEFA017)
1222 })
1223 }
1224}
1225
1226#[derive(Serialize)]
1227struct ChatCompletionChoice {
1228 index: usize,
1229 message: ChatCompletionResponseMessage,
1230 finish_reason: &'static str,
1231}
1232
1233#[derive(Serialize)]
1234struct ChatCompletionResponseMessage {
1235 role: &'static str,
1236 #[serde(skip_serializing_if = "Option::is_none")]
1237 content: Option<String>,
1238 #[serde(skip_serializing_if = "Option::is_none")]
1239 tool_calls: Option<Vec<ToolCallOut>>,
1240}
1241
1242#[derive(Serialize, Clone)]
1243struct ToolCallOut {
1244 id: String,
1245 #[serde(rename = "type")]
1246 kind: &'static str,
1247 function: ToolCallFunctionOut,
1248}
1249
1250#[derive(Serialize, Clone)]
1251struct ToolCallFunctionOut {
1252 name: String,
1253 arguments: String,
1257}
1258
1259#[derive(Serialize)]
1260struct ChatCompletionResponse {
1261 id: String,
1262 request_id: String,
1268 object: &'static str,
1269 model: String,
1270 choices: Vec<ChatCompletionChoice>,
1271 usage: generate::Usage,
1276 ferrox_cache: &'static str,
1284}
1285
1286#[derive(Serialize)]
1287struct ChatCompletionChunkDelta {
1288 #[serde(skip_serializing_if = "Option::is_none")]
1289 role: Option<&'static str>,
1290 #[serde(skip_serializing_if = "Option::is_none")]
1291 content: Option<String>,
1292 #[serde(skip_serializing_if = "Option::is_none")]
1293 tool_calls: Option<Vec<ToolCallOut>>,
1294}
1295
1296#[derive(Serialize)]
1297struct ChatCompletionChunkChoice {
1298 index: usize,
1299 delta: ChatCompletionChunkDelta,
1300 finish_reason: Option<&'static str>,
1301}
1302
1303#[derive(Serialize)]
1304struct ChatCompletionChunk {
1305 id: String,
1306 #[serde(skip_serializing_if = "Option::is_none")]
1313 request_id: Option<String>,
1314 object: &'static str,
1315 model: String,
1316 choices: Vec<ChatCompletionChunkChoice>,
1317 #[serde(skip_serializing_if = "Option::is_none")]
1320 usage: Option<generate::Usage>,
1321}
1322
1323async fn health(State(state): State<Arc<AppState>>) -> Response {
1329 let snapshot = state.detection.snapshot();
1330 let mut capabilities = snapshot.capabilities;
1331 let active = state.active();
1332
1333 capabilities.push(match active.as_deref() {
1336 None => ferrox_api::Capability::unavailable(
1341 ferrox_api::health::capability::REAL_WEIGHTS,
1342 ferrox_api::health::reason::MODEL_NOT_LOADED,
1343 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1344 ),
1345 Some(active) if active.model.is_synthetic() => ferrox_api::Capability::unavailable(
1346 ferrox_api::health::capability::REAL_WEIGHTS,
1347 ferrox_api::health::reason::MODEL_NOT_LOADED,
1348 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1349 checkpoint. Output from this model is noise.",
1350 ),
1351 Some(active) => ferrox_api::Capability::available(
1352 ferrox_api::health::capability::REAL_WEIGHTS,
1353 format!("Serving the real checkpoint '{}'.", active.model.name()),
1354 ),
1355 });
1356 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1357 ferrox_api::Capability::available(
1358 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1359 "Concurrent requests share one batched decode step.",
1360 )
1361 } else {
1362 ferrox_api::Capability::unavailable(
1363 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1364 ferrox_api::health::reason::DISABLED,
1365 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1366 )
1367 });
1368
1369 let last_request_ms = state
1370 .last_request_ms
1371 .load(std::sync::atomic::Ordering::Relaxed);
1372 let uptime = state.started_at.elapsed();
1373 let health_state = if active.is_none() {
1381 ferrox_api::HealthState::Unavailable
1382 } else {
1383 snapshot.state
1384 };
1385 let body = ferrox_api::HealthResponse {
1386 state: health_state,
1387 reason: match health_state {
1388 ferrox_api::HealthState::Ready => None,
1389 ferrox_api::HealthState::Unavailable => {
1390 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
1391 }
1392 ferrox_api::HealthState::Detecting => {
1393 Some(ferrox_api::health::reason::DETECTING.to_string())
1394 }
1395 },
1396 detail: match health_state {
1397 ferrox_api::HealthState::Ready => None,
1398 ferrox_api::HealthState::Unavailable => Some(
1399 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1400 .to_string(),
1401 ),
1402 ferrox_api::HealthState::Detecting => {
1403 Some("Probing available compute backends.".to_string())
1404 }
1405 },
1406 model: active
1407 .as_deref()
1408 .map(|active| ferrox_api::health::ModelSummary {
1409 id: active.model.name().to_string(),
1410 tokenizer: active.model.tokenizer_kind().to_string(),
1411 synthetic_weights: active.model.is_synthetic(),
1412 }),
1413 capabilities,
1414 version: env!("CARGO_PKG_VERSION").to_string(),
1415 pid: std::process::id(),
1416 uptime_seconds: uptime.as_secs_f64(),
1417 server_time_unix_ms: std::time::SystemTime::now()
1418 .duration_since(std::time::UNIX_EPOCH)
1419 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1420 .unwrap_or(0),
1421 last_request_age_seconds: (last_request_ms > 0)
1422 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1423 .map(|age| age.max(0.0)),
1424 };
1425
1426 let status =
1427 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1428 (status, Json(body)).into_response()
1429}
1430
1431async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1432 let Some(active) = state.active() else {
1436 return Json(serde_json::json!({ "object": "list", "data": [] }));
1437 };
1438 let mut model_entry = serde_json::json!({
1439 "id": active.model.name(),
1440 "object": "model",
1441 "ferrox_synthetic_weights": active.model.is_synthetic(),
1442 "ferrox_tokenizer": active.model.tokenizer_kind(),
1443 });
1444 if let Some(mcp) = &state.mcp {
1445 model_entry["ferrox_mcp"] = mcp.models_metadata();
1446 }
1447 Json(serde_json::json!({
1448 "object": "list",
1449 "data": [model_entry]
1450 }))
1451}
1452
1453#[derive(Serialize)]
1454struct CombinedCacheStats {
1455 response_cache: cache::CacheStats,
1456 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
1458}
1459
1460async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1461 Json(CombinedCacheStats {
1462 response_cache: lock_cache(&state.response_cache).stats(),
1463 prefix_cache: state
1464 .prefix_cache
1465 .as_ref()
1466 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1467 })
1468}
1469
1470async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1475 use std::sync::atomic::Ordering;
1476
1477 let cache_stats = lock_cache(&state.response_cache).stats();
1478 let active = state.active();
1479 let requests_total = state.requests_total.load(Ordering::Relaxed);
1480 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1481 let uptime = state.started_at.elapsed().as_secs_f64();
1482
1483 let body = format!(
1484 "# HELP ferrox_requests_total Total chat completion requests received.\n\
1485 # TYPE ferrox_requests_total counter\n\
1486 ferrox_requests_total {requests_total}\n\
1487 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
1488 # TYPE ferrox_request_errors_total counter\n\
1489 ferrox_request_errors_total {errors_total}\n\
1490 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
1491 # TYPE ferrox_cache_hits_total counter\n\
1492 ferrox_cache_hits_total {}\n\
1493 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
1494 # TYPE ferrox_cache_misses_total counter\n\
1495 ferrox_cache_misses_total {}\n\
1496 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
1497 # TYPE ferrox_cache_entries gauge\n\
1498 ferrox_cache_entries {}\n\
1499 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1500 # TYPE ferrox_synthetic_weights gauge\n\
1501 ferrox_synthetic_weights {}\n\
1502 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
1503 # TYPE ferrox_uptime_seconds gauge\n\
1504 ferrox_uptime_seconds {uptime}\n",
1505 cache_stats.hits,
1506 cache_stats.misses,
1507 cache_stats.entries,
1508 active
1512 .as_ref()
1513 .map(|a| a.model.is_synthetic() as u8)
1514 .unwrap_or(0),
1515 );
1516
1517 let body = match active
1521 .as_ref()
1522 .and_then(|a| a.model.expert_store_stats())
1523 {
1524 Some(es) => format!(
1525 "{body}\
1526 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
1527 # TYPE ferrox_expert_cache_hits_total counter\n\
1528 ferrox_expert_cache_hits_total {}\n\
1529 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1530 # TYPE ferrox_expert_cache_misses_total counter\n\
1531 ferrox_expert_cache_misses_total {}\n\
1532 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
1533 # TYPE ferrox_expert_cache_evictions_total counter\n\
1534 ferrox_expert_cache_evictions_total {}\n\
1535 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1536 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
1537 ferrox_expert_cache_pass_throughs_total {}\n\
1538 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1539 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
1540 ferrox_expert_cache_bytes_read_total {}\n\
1541 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1542 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
1543 ferrox_expert_cache_resident_bytes {}\n",
1544 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1545 ),
1546 None => body,
1547 };
1548
1549 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
1554 Some(batcher) => {
1555 let sched = batcher.stats();
1556 format!(
1557 "{body}\
1558 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
1559 # TYPE ferrox_prefill_chunks_total counter\n\
1560 ferrox_prefill_chunks_total {}\n\
1561 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
1562 # TYPE ferrox_prefill_tokens_total counter\n\
1563 ferrox_prefill_tokens_total {}\n\
1564 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
1565 # TYPE ferrox_decode_steps_total counter\n\
1566 ferrox_decode_steps_total {}\n\
1567 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
1568 # TYPE ferrox_scheduler_queue_depth gauge\n\
1569 ferrox_scheduler_queue_depth {}\n\
1570 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
1571 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
1572 ferrox_scheduler_queue_rejected_total {}\n\
1573 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
1574 # TYPE ferrox_kv_blocks_total gauge\n\
1575 ferrox_kv_blocks_total {}\n\
1576 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
1577 # TYPE ferrox_kv_blocks_free gauge\n\
1578 ferrox_kv_blocks_free {}\n\
1579 # HELP ferrox_kv_block_size Token positions per KV block.\n\
1580 # TYPE ferrox_kv_block_size gauge\n\
1581 ferrox_kv_block_size {}\n\
1582 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
1583 # TYPE ferrox_kv_rejected_too_large_total counter\n\
1584 ferrox_kv_rejected_too_large_total {}\n\
1585 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
1586 # TYPE ferrox_kv_rejected_context_length_total counter\n\
1587 ferrox_kv_rejected_context_length_total {}\n\
1588 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
1589 # TYPE ferrox_scheduler_aborted_total counter\n\
1590 ferrox_scheduler_aborted_total {}\n",
1591 sched.prefill_chunks,
1592 sched.prefill_tokens,
1593 sched.decode_steps,
1594 sched.queue_depth,
1595 sched.queue_rejected,
1596 sched.kv_blocks_total,
1597 sched.kv_blocks_free,
1598 sched.kv_block_size,
1599 sched.kv_rejected_too_large,
1600 sched.kv_rejected_context_length,
1601 sched.aborted,
1602 )
1603 }
1604 None => body,
1605 };
1606
1607 (
1608 [(
1609 axum::http::header::CONTENT_TYPE,
1610 "text/plain; version=0.0.4",
1611 )],
1612 body,
1613 )
1614 .into_response()
1615}
1616
1617pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
1618
1619pub(crate) fn unsupported_feature(message: &str) -> ApiError {
1620 (
1621 StatusCode::NOT_IMPLEMENTED,
1622 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
1623 )
1624}
1625
1626pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
1627 let status = match e {
1628 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
1629 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
1634 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
1640 StatusCode::SERVICE_UNAVAILABLE
1641 }
1642 };
1643 tracing::warn!("decode error: {e}");
1644 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
1645 if let generate::DecodeError::KvBudgetExceeded {
1651 binding,
1652 estimated_bytes,
1653 limit_bytes,
1654 positions,
1655 positions_limit,
1656 ..
1657 } = &e
1658 {
1659 body["error"]["type"] = serde_json::json!("invalid_request_error");
1660 body["error"]["code"] = serde_json::json!(binding);
1661 body["error"]["binding"] = serde_json::json!(binding);
1662 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
1663 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
1664 body["error"]["positions"] = serde_json::json!(positions);
1665 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
1666 }
1667 if let Some(secs) = e.retry_after_secs() {
1671 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
1672 }
1673 (status, Json(body))
1674}
1675
1676pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
1677 tracing::error!("generation task panicked: {e}");
1678 (
1679 StatusCode::INTERNAL_SERVER_ERROR,
1680 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
1681 )
1682}
1683
1684#[allow(clippy::too_many_arguments)] fn run_generation_emit(
1694 model: &Model,
1695 prompt: &str,
1696 params: &GenerationParams,
1697 kv_pool: Option<&generate::KvPoolConfig>,
1698 prefix_cache: Option<&Mutex<PrefixCache>>,
1699 continuous_batcher: Option<&batch_scheduler::ContinuousBatcher>,
1700 ceiling: Option<&budget::ContextCeiling>,
1701 mut emit: impl FnMut(&str),
1702) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
1703 let synthetic = model.is_synthetic();
1704 let mut chunks = Vec::new();
1705 let params = &{
1711 let mut resolved = params.clone();
1712 resolved.stop_token_ids =
1713 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
1714 resolved
1715 };
1716 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
1717 let (finish, usage) = match model {
1718 Model::Gguf(m) => {
1719 if let Some(batcher) = continuous_batcher {
1720 let mut tokens = m.tokenizer.encode(prompt);
1721 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
1722 let (finish, _generated_ids, text, usage) =
1723 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?;
1724 if !text.is_empty() {
1725 chunks.push(text);
1726 }
1727 (finish, usage)
1728 } else {
1729 generate::generate(
1730 &m.decoder,
1731 m.tokenizer.as_ref(),
1732 &m.stop_tokens,
1733 m.bos_id,
1734 prompt,
1735 params,
1736 kv_pool,
1737 prefix_cache,
1738 ceiling,
1739 |chunk| {
1740 chunks.push(chunk.to_string());
1741 if !synthetic {
1742 emit(chunk);
1743 }
1744 },
1745 )?
1746 }
1747 }
1748 Model::Kimi(m) => generate::generate_engine(
1749 &m.engine,
1750 &m.tokenizer,
1751 &m.stop_tokens,
1752 None,
1753 prompt,
1754 params,
1755 |chunk| {
1756 chunks.push(chunk.to_string());
1757 if !synthetic {
1758 emit(chunk);
1759 }
1760 },
1761 )?,
1762 Model::Mla(m) => generate::generate_engine(
1763 &m.engine,
1764 &m.tokenizer,
1765 &m.stop_tokens,
1766 m.bos_id,
1767 prompt,
1768 params,
1769 |chunk| {
1770 chunks.push(chunk.to_string());
1771 if !synthetic {
1772 emit(chunk);
1773 }
1774 },
1775 )?,
1776 Model::Gemma4(m) => generate::generate_engine(
1777 &m.engine,
1778 &m.tokenizer,
1779 &m.stop_tokens,
1780 m.bos_id,
1781 prompt,
1782 params,
1783 |chunk| {
1784 chunks.push(chunk.to_string());
1785 if !synthetic {
1786 emit(chunk);
1787 }
1788 },
1789 )?,
1790 Model::Glm52(m) => generate::generate_engine(
1791 &m.engine,
1792 &m.tokenizer,
1793 &m.stop_tokens,
1794 m.bos_id,
1795 prompt,
1796 params,
1797 |chunk| {
1798 chunks.push(chunk.to_string());
1799 if !synthetic {
1800 emit(chunk);
1801 }
1802 },
1803 )?,
1804 };
1805
1806 let mut full = chunks.concat();
1807 if synthetic {
1808 full = format!(
1809 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
1810 to serve a real model. Decoded ids -> {full:?}]"
1811 );
1812 emit(&full);
1813 } else if used_batcher && !full.is_empty() {
1814 emit(&full);
1815 }
1816
1817 Ok((finish, usage, full))
1818}
1819
1820pub(crate) fn run_generation(
1823 model: &Model,
1824 prompt: &str,
1825 params: &GenerationParams,
1826 kv_pool: Option<&generate::KvPoolConfig>,
1827 prefix_cache: Option<&Mutex<PrefixCache>>,
1828 continuous_batcher: Option<&batch_scheduler::ContinuousBatcher>,
1829 ceiling: Option<&budget::ContextCeiling>,
1830) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
1831 let (finish, usage, full) = run_generation_emit(
1832 model,
1833 prompt,
1834 params,
1835 kv_pool,
1836 prefix_cache,
1837 continuous_batcher,
1838 ceiling,
1839 |_| {},
1840 )?;
1841 Ok((
1842 if full.is_empty() {
1843 Vec::new()
1844 } else {
1845 vec![full]
1846 },
1847 finish,
1848 usage,
1849 ))
1850}
1851
1852pub(crate) fn prompt_from_messages(
1853 messages: &[ChatMessage],
1854 template: chat_template::ChatTemplate,
1855 tools: &[ToolDef],
1856) -> String {
1857 if tools.is_empty() {
1858 template.render(messages)
1859 } else {
1860 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
1861 with_preamble.push(ChatMessage {
1862 role: "system".to_string(),
1863 content: Some(MessageContent::Text(tool_preamble(tools))),
1864 tool_calls: None,
1865 tool_call_id: None,
1866 });
1867 with_preamble.extend_from_slice(messages);
1868 template.render(&with_preamble)
1869 }
1870}
1871
1872fn tool_preamble(tools: &[ToolDef]) -> String {
1883 let mut out = String::from(
1884 "You can call tools to help answer the user. To call a tool, respond with \
1885 EXACTLY one line in this format and nothing else:\n\
1886 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
1887 object matching that tool's parameters>}}</tool_call>\n\n\
1888 Available tools:\n",
1889 );
1890 for t in tools {
1891 out.push_str(&format!(
1892 "- {}: {}\n parameters (JSON schema): {}\n",
1893 t.function.name,
1894 t.function.description.as_deref().unwrap_or(""),
1895 t.function
1896 .parameters
1897 .as_ref()
1898 .map(|v| v.to_string())
1899 .unwrap_or_else(|| "{}".to_string()),
1900 ));
1901 }
1902 out
1903}
1904
1905fn extract_tool_call(text: &str) -> Option<(String, String)> {
1912 const START: &str = "<tool_call>";
1913 const END: &str = "</tool_call>";
1914 let start = text.find(START)? + START.len();
1915 let end = start + text[start..].find(END)?;
1916 let body = text[start..end].trim();
1917 let value: serde_json::Value = serde_json::from_str(body).ok()?;
1918 let name = value.get("name")?.as_str()?.to_string();
1919 let arguments = value
1920 .get("arguments")
1921 .cloned()
1922 .unwrap_or_else(|| serde_json::json!({}));
1923 Some((name, arguments.to_string()))
1924}
1925
1926fn build_response_message(
1934 content: String,
1935 tools_active: bool,
1936 base_finish: &'static str,
1937) -> (ChatCompletionResponseMessage, &'static str) {
1938 if tools_active {
1939 if let Some((name, arguments)) = extract_tool_call(&content) {
1940 return (
1941 ChatCompletionResponseMessage {
1942 role: "assistant",
1943 content: None,
1944 tool_calls: Some(vec![ToolCallOut {
1945 id: "call_0".to_string(),
1946 kind: "function",
1947 function: ToolCallFunctionOut { name, arguments },
1948 }]),
1949 },
1950 "tool_calls",
1951 );
1952 }
1953 }
1954 (
1955 ChatCompletionResponseMessage {
1956 role: "assistant",
1957 content: Some(content),
1958 tool_calls: None,
1959 },
1960 base_finish,
1961 )
1962}
1963
1964fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
1969 let mut history = match &req.session_id {
1970 Some(id) => state.sessions.extend_and_get(id, &req.messages),
1971 None => req.messages.clone(),
1972 };
1973 if req.json_object_mode() {
1974 inject_json_object_system_hint(&mut history);
1975 }
1976 history
1977}
1978
1979fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
1980 const HINT: &str =
1981 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
1982 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
1983 match &mut sys.content {
1984 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
1985 s.push_str("\n\n");
1986 s.push_str(HINT);
1987 }
1988 None => {
1989 sys.content = Some(MessageContent::Text(HINT.to_string()));
1990 }
1991 _ => {}
1992 }
1993 } else {
1994 messages.insert(
1995 0,
1996 ChatMessage {
1997 role: "system".to_string(),
1998 content: Some(MessageContent::Text(HINT.to_string())),
1999 tool_calls: None,
2000 tool_call_id: None,
2001 },
2002 );
2003 }
2004}
2005
2006async fn chat_completions(
2007 State(state): State<Arc<AppState>>,
2008 headers: axum::http::HeaderMap,
2009 Json(req): Json<ChatCompletionRequest>,
2010) -> Response {
2011 let attribution = attribution::Attribution::from_headers(&headers);
2012 state
2013 .requests_total
2014 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2015 let started = std::time::Instant::now();
2016
2017 let request_id = ferrox_api::next_request_id();
2021 let stream = req.stream.unwrap_or(false);
2022
2023 if let Err(err) = req.validate_supported_fields() {
2024 state
2025 .request_errors_total
2026 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2027 let response = err.into_response();
2028 state.record_request(stats::Record {
2029 request_id: &request_id,
2030 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2031 model: state.active_model_name(),
2032 status: response.status().as_u16(),
2033 stream,
2034 duration_ms: started.elapsed().as_millis() as u64,
2035 usage: None,
2036 attribution: &attribution,
2037 });
2038 return response;
2039 }
2040
2041 let response = if stream {
2042 chat_completions_stream(
2043 Arc::clone(&state),
2044 req,
2045 request_id.clone(),
2046 started,
2047 attribution.clone(),
2048 )
2049 .await
2050 .into_response()
2051 } else {
2052 chat_completions_full(
2053 Arc::clone(&state),
2054 req,
2055 request_id.clone(),
2056 started,
2057 attribution.clone(),
2058 )
2059 .await
2060 .into_response()
2061 };
2062
2063 if response.status().is_client_error() || response.status().is_server_error() {
2064 state
2065 .request_errors_total
2066 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2067 state.record_request(stats::Record {
2071 request_id: &request_id,
2072 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2073 model: state.active_model_name(),
2076 status: response.status().as_u16(),
2077 stream,
2078 duration_ms: started.elapsed().as_millis() as u64,
2079 usage: None,
2080 attribution: &attribution,
2081 });
2082 }
2083 state.mark_request_finished();
2084
2085 response
2086}
2087
2088async fn chat_completions_full(
2089 state: Arc<AppState>,
2090 req: ChatCompletionRequest,
2091 request_id: String,
2092 started: std::time::Instant,
2093 attribution: attribution::Attribution,
2094) -> Result<Json<ChatCompletionResponse>, ApiError> {
2095 let tools_active = req.tools_active();
2096 let active = state.require_active()?;
2100 let history = resolve_history(&state, &req);
2101 let prompt = prompt_from_messages(&history, active.model.chat_template(), &req.tools);
2102 let key = req.is_cacheable().then(|| req.cache_key(&prompt));
2103
2104 let (completion, cache_status) = if let Some(cached) = key
2105 .as_ref()
2106 .and_then(|key| lock_cache(&state.response_cache).get(key))
2107 {
2108 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2109 (cached, "hit")
2110 } else {
2111 let model = Arc::clone(&active.model);
2112 let kv_pool = state.kv_pool.clone();
2113 let prefix_cache = state.prefix_cache.clone();
2114 let batcher = active.batcher.clone();
2115 let ceiling = active.ceiling.clone();
2116 let params = req.generation_params_for_template(active.model.chat_template());
2117 let prompt_for_task = prompt.clone();
2118 let (chunks, finish, usage) = tokio::task::spawn_blocking(move || {
2119 run_generation(
2120 &model,
2121 &prompt_for_task,
2122 ¶ms,
2123 kv_pool.as_ref(),
2124 prefix_cache.as_deref(),
2125 batcher.as_ref(),
2126 ceiling.as_deref(),
2127 )
2128 })
2129 .await
2130 .map_err(join_error_response)?
2131 .map_err(decode_error_response)?;
2132
2133 let completion = cache::CachedCompletion {
2134 content: chunks.concat(),
2135 finish,
2136 usage,
2137 };
2138 let cache_status = if let Some(key) = key {
2139 tracing::debug!("cache miss for key {}", key.digest());
2140 lock_cache(&state.response_cache).put(key, completion.clone());
2141 "miss"
2142 } else {
2143 "skip"
2144 };
2145 (completion, cache_status)
2146 };
2147 let content = completion.content;
2148
2149 if req.json_object_mode() {
2150 json_mode::validate_json_object_output(&content)?;
2151 }
2152
2153 if let Some(id) = &req.session_id {
2157 state.sessions.store_reply(
2158 id,
2159 ChatMessage {
2160 role: "assistant".to_string(),
2161 content: Some(MessageContent::Text(content.clone())),
2162 tool_calls: None,
2163 tool_call_id: None,
2164 },
2165 );
2166 }
2167
2168 let (message, finish_reason) =
2169 build_response_message(content, tools_active, completion.finish.as_str());
2170
2171 state.record_request(stats::Record {
2172 request_id: &request_id,
2173 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2174 model: Some(active.model.name().to_string()),
2177 status: 200,
2178 stream: false,
2179 duration_ms: started.elapsed().as_millis() as u64,
2180 usage: Some(&completion.usage),
2181 attribution: &attribution,
2182 });
2183
2184 Ok(Json(ChatCompletionResponse {
2185 id: request_id.clone(),
2186 request_id,
2187 object: "chat.completion",
2188 model: req.model,
2189 choices: vec![ChatCompletionChoice {
2190 index: 0,
2191 message,
2192 finish_reason,
2193 }],
2194 usage: completion.usage,
2195 ferrox_cache: cache_status,
2196 }))
2197}
2198
2199async fn chat_completions_stream(
2200 state: Arc<AppState>,
2201 req: ChatCompletionRequest,
2202 request_id: String,
2203 started: std::time::Instant,
2204 attribution: attribution::Attribution,
2205) -> Result<Response, ApiError> {
2206 let tools_active = req.tools_active();
2208 let active = state.require_active()?;
2212 let history = resolve_history(&state, &req);
2213 let prompt = prompt_from_messages(&history, active.model.chat_template(), &req.tools);
2214 let model_name = req.model.clone();
2215 let session_id = req.session_id.clone();
2216 let sessions = state.sessions.clone();
2217
2218 let model = Arc::clone(&active.model);
2219 let kv_pool = state.kv_pool.clone();
2220 let prefix_cache = state.prefix_cache.clone();
2221 let batcher = active.batcher.clone();
2222 let ceiling = active.ceiling.clone();
2223 let mut params = req.generation_params_for_template(active.model.chat_template());
2224 let stats_state = Arc::clone(&state);
2225 let served_model = active.model.name().to_string();
2228
2229 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2234 params.cancel = Some(cancel_token.clone());
2235
2236 let overlap = !tools_active && batcher.is_none();
2240
2241 let slot = req
2245 .stream_resumable
2246 .unwrap_or(false)
2247 .then(|| state.streams.register(&request_id));
2248 let emitter = resume::Emitter::new(slot);
2249
2250 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2251
2252 tokio::task::spawn_blocking(move || {
2253 let _cancel_guard = cancel_guard;
2256 let tx_chunks = tx.clone();
2257 let orphan_timeout = sse::orphan_timeout_from_env();
2262 let mut first = true;
2263 let head_request_id = request_id.clone();
2264 let result = run_generation_emit(
2265 &model,
2266 &prompt,
2267 ¶ms,
2268 kv_pool.as_ref(),
2269 prefix_cache.as_deref(),
2270 batcher.as_ref(),
2271 ceiling.as_deref(),
2272 |chunk| {
2273 if !overlap || chunk.is_empty() {
2274 return;
2275 }
2276 let role = if first { Some("assistant") } else { None };
2277 let request_id = first.then(|| head_request_id.clone());
2278 first = false;
2279 let payload = ChatCompletionChunk {
2280 id: head_request_id.clone(),
2281 request_id,
2282 object: "chat.completion.chunk",
2283 model: model_name.clone(),
2284 choices: vec![ChatCompletionChunkChoice {
2285 index: 0,
2286 delta: ChatCompletionChunkDelta {
2287 role,
2288 content: Some(chunk.to_string()),
2289 tool_calls: None,
2290 },
2291 finish_reason: None,
2292 }],
2293 usage: None,
2294 };
2295 if let Err(why) =
2304 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
2305 {
2306 if why == sse::SendFailure::Orphaned {
2307 tracing::warn!(
2308 "SSE stream {head_request_id} accepted nothing for the orphan \
2309 deadline; treating it as abandoned"
2310 );
2311 }
2312 if !emitter.is_resumable() {
2324 cancel_token.cancel();
2325 }
2326 }
2327 },
2328 );
2329
2330 let mut pending_request_id = first.then(|| request_id.clone());
2336
2337 match result {
2338 Ok((finish, usage, full_text)) => {
2339 if let Some(id) = &session_id {
2340 sessions.store_reply(
2341 id,
2342 ChatMessage {
2343 role: "assistant".to_string(),
2344 content: Some(MessageContent::Text(full_text.clone())),
2345 tool_calls: None,
2346 tool_call_id: None,
2347 },
2348 );
2349 }
2350 let tool_call = if tools_active {
2351 extract_tool_call(&full_text)
2352 } else {
2353 None
2354 };
2355 if !overlap {
2356 if let Some((name, arguments)) = &tool_call {
2357 let payload = ChatCompletionChunk {
2358 id: request_id.clone(),
2359 request_id: pending_request_id.take(),
2360 object: "chat.completion.chunk",
2361 model: model_name.clone(),
2362 choices: vec![ChatCompletionChunkChoice {
2363 index: 0,
2364 delta: ChatCompletionChunkDelta {
2365 role: Some("assistant"),
2366 content: None,
2367 tool_calls: Some(vec![ToolCallOut {
2368 id: "call_0".to_string(),
2369 kind: "function",
2370 function: ToolCallFunctionOut {
2371 name: name.clone(),
2372 arguments: arguments.clone(),
2373 },
2374 }]),
2375 },
2376 finish_reason: None,
2377 }],
2378 usage: None,
2379 };
2380 let _ =
2381 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
2382 } else if !full_text.is_empty() {
2383 let payload = ChatCompletionChunk {
2384 id: request_id.clone(),
2385 request_id: pending_request_id.take(),
2386 object: "chat.completion.chunk",
2387 model: model_name.clone(),
2388 choices: vec![ChatCompletionChunkChoice {
2389 index: 0,
2390 delta: ChatCompletionChunkDelta {
2391 role: Some("assistant"),
2392 content: Some(full_text),
2393 tool_calls: None,
2394 },
2395 finish_reason: None,
2396 }],
2397 usage: None,
2398 };
2399 let _ =
2400 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
2401 }
2402 }
2403 let final_finish_reason = if tool_call.is_some() {
2404 "tool_calls"
2405 } else {
2406 finish.as_str()
2407 };
2408 let final_payload = ChatCompletionChunk {
2409 id: request_id.clone(),
2410 request_id: pending_request_id.take(),
2411 object: "chat.completion.chunk",
2412 model: model_name,
2413 choices: vec![ChatCompletionChunkChoice {
2414 index: 0,
2415 delta: ChatCompletionChunkDelta {
2416 role: None,
2417 content: None,
2418 tool_calls: None,
2419 },
2420 finish_reason: Some(final_finish_reason),
2421 }],
2422 usage: Some(usage.clone()),
2423 };
2424 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
2425 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
2426 stats_state.record_request(stats::Record {
2431 request_id: &request_id,
2432 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2433 model: Some(served_model.clone()),
2434 status: 200,
2435 stream: true,
2436 duration_ms: started.elapsed().as_millis() as u64,
2437 usage: Some(&usage),
2438 attribution: &attribution,
2439 });
2440 }
2441 Err(e) => {
2442 tracing::warn!("decode error on streamed request {request_id}: {e}");
2443 stats_state.record_request(stats::Record {
2449 request_id: &request_id,
2450 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2451 model: Some(served_model.clone()),
2452 status: 500,
2453 stream: true,
2454 duration_ms: started.elapsed().as_millis() as u64,
2455 usage: None,
2456 attribution: &attribution,
2457 });
2458 let payload = ChatCompletionChunk {
2459 id: request_id.clone(),
2460 request_id: pending_request_id.take(),
2461 object: "chat.completion.chunk",
2462 model: model_name,
2463 choices: vec![ChatCompletionChunkChoice {
2464 index: 0,
2465 delta: ChatCompletionChunkDelta {
2466 role: Some("assistant"),
2467 content: Some(format!("[error: {e}]")),
2468 tool_calls: None,
2469 },
2470 finish_reason: Some("stop"),
2471 }],
2472 usage: None,
2473 };
2474 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
2475 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
2476 }
2477 }
2478 drop(emitter);
2482 });
2483
2484 let stream =
2485 futures_util::stream::unfold(
2486 rx,
2487 |mut rx| async move { rx.recv().await.map(|ev| (ev, rx)) },
2488 );
2489 Ok((
2502 [(
2503 axum::http::HeaderName::from_static("x-accel-buffering"),
2504 axum::http::HeaderValue::from_static("no"),
2505 )],
2506 Sse::new(stream).keep_alive(KeepAlive::default()),
2507 )
2508 .into_response())
2509}
2510
2511fn resume_route(template: &str) -> String {
2519 template.replace("{request_id}", ":request_id")
2520}
2521
2522async fn cancel_generation(
2531 State(state): State<Arc<AppState>>,
2532 Json(req): Json<ferrox_api::CancelGenerationRequest>,
2533) -> Response {
2534 let cancelled = state.cancels.cancel(&req.request_id);
2535 let status = if cancelled {
2536 StatusCode::OK
2537 } else {
2538 StatusCode::NOT_FOUND
2539 };
2540 let detail = if cancelled {
2541 "the generation was asked to stop; it ends at its next token".to_string()
2542 } else {
2543 "no generation with that request_id is running -- it has already \
2544 finished, was never issued, or was served by a path that does \
2545 not register for cancellation"
2546 .to_string()
2547 };
2548 (
2549 status,
2550 Json(ferrox_api::CancelGenerationResponse {
2551 request_id: req.request_id,
2552 cancelled,
2553 detail,
2554 }),
2555 )
2556 .into_response()
2557}
2558
2559type Activated = (
2563 Model,
2564 Option<batch_scheduler::ContinuousBatcher>,
2565 Option<Arc<budget::ContextCeiling>>,
2566);
2567
2568fn price_batcher_config(path: Option<&str>) -> batch_scheduler::BatcherConfig {
2586 let mut batcher = batch_scheduler::BatcherConfig::from_env();
2587 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
2588 return batcher;
2591 }
2592 let Some(path) = path else {
2593 return batcher;
2594 };
2595 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
2601 let Some((priced, gguf_ctx, source)) = priced else {
2602 return batcher;
2603 };
2604 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
2605 tracing::warn!(
2609 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
2610 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
2611 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
2612 admit on a number you choose.",
2613 priced.weights_bytes,
2614 priced.device_budget_bytes,
2615 );
2616 return batcher;
2617 };
2618 tracing::info!("{source}");
2619 tracing::info!("{}", derived.fit);
2620 let adopted = budget::apply_derived(&mut batcher, &derived);
2621 if adopted.max_context {
2622 tracing::info!(
2623 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
2624 override with FERROX_CB_MAX_CONTEXT",
2625 derived.max_context
2626 );
2627 }
2628 if adopted.kv_blocks {
2629 tracing::info!(
2630 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
2631 derived.kv_blocks,
2632 batcher.kv_block_size
2633 );
2634 }
2635 batcher
2636}
2637
2638pub(crate) fn activate_loaded_model(
2647 loaded: model::LoadedModel,
2648 enable_continuous_batching: bool,
2649 path: Option<&str>,
2650) -> Activated {
2651 match loaded {
2652 model::LoadedModel::Gguf(g) => {
2653 let decoder = Arc::new(g.decoder);
2654 let tokenizer = Arc::new(g.tokenizer);
2655 let config = price_batcher_config(path);
2656 let shape =
2662 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
2663 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
2664 let batcher = if enable_continuous_batching {
2665 tracing::info!(
2666 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
2667 (stop sequences use the same pending-buffer trim as the private generate loop)"
2668 );
2669 let tok = Arc::clone(&tokenizer);
2670 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
2671 Some(batch_scheduler::ContinuousBatcher::spawn_with_ceiling(
2672 Arc::clone(&decoder),
2673 decode,
2674 config,
2675 Arc::clone(&ceiling),
2676 ))
2677 } else {
2678 None
2679 };
2680 (
2681 Model::Gguf(GgufModel {
2682 decoder,
2683 tokenizer,
2684 stop_tokens: g.stop_tokens,
2685 bos_id: g.bos_id,
2686 is_synthetic: g.is_synthetic,
2687 chat_template: g.chat_template,
2688 }),
2689 batcher,
2690 Some(ceiling),
2691 )
2692 }
2693 model::LoadedModel::Kimi(k) => (
2694 Model::Kimi(KimiModel {
2695 engine: k.engine,
2696 tokenizer: k.tokenizer,
2697 stop_tokens: k.stop_tokens,
2698 chat_template: k.chat_template,
2699 }),
2700 None,
2701 None,
2702 ),
2703 model::LoadedModel::Mla(m) => (
2704 Model::Mla(MlaModel {
2705 engine: m.engine,
2706 tokenizer: m.tokenizer,
2707 stop_tokens: m.stop_tokens,
2708 bos_id: m.bos_id,
2709 name: m.name,
2710 chat_template: m.chat_template,
2711 }),
2712 None,
2713 None,
2714 ),
2715 model::LoadedModel::Gemma4(m) => (
2716 Model::Gemma4(Gemma4Model {
2717 engine: m.engine,
2718 tokenizer: m.tokenizer,
2719 stop_tokens: m.stop_tokens,
2720 bos_id: m.bos_id,
2721 name: m.name,
2722 chat_template: m.chat_template,
2723 }),
2724 None,
2725 None,
2726 ),
2727 model::LoadedModel::Glm52(g) => (
2728 Model::Glm52(Glm52Model {
2729 engine: g.engine,
2730 tokenizer: g.tokenizer,
2731 stop_tokens: g.stop_tokens,
2732 bos_id: g.bos_id,
2733 name: g.name,
2734 chat_template: g.chat_template,
2735 }),
2736 None,
2737 None,
2738 ),
2739 }
2740}
2741
2742fn build_app_state(
2743 loaded: model::LoadedModel,
2744 kv_pool: Option<generate::KvPoolConfig>,
2745 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
2746 enable_continuous_batching: bool,
2747 mcp: Option<mcp::LoadedMcpConfig>,
2748 detection: Arc<health::Detection>,
2749) -> AppState {
2750 let (model, batcher, ceiling) = activate_loaded_model(
2751 loaded,
2752 enable_continuous_batching,
2753 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
2754 );
2755 let id = startup_model_id();
2761 AppState {
2762 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
2763 id,
2764 model: Arc::new(model),
2765 batcher,
2766 ceiling,
2767 }))),
2768 load_in_progress: std::sync::atomic::AtomicBool::new(false),
2769 tasks: Arc::new(tasks::TaskRegistry::new()),
2770 cancels: Arc::new(cancel::CancelRegistry::new()),
2771 stats: stats::Stats::new(),
2772 streams: resume::StreamRegistry::new(),
2773 model_dir: admin::model_dirs().into_iter().next(),
2774 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
2775 kv_pool,
2776 prefix_cache,
2777 sessions: session::SessionStore::new(),
2778 requests_total: std::sync::atomic::AtomicU64::new(0),
2779 request_errors_total: std::sync::atomic::AtomicU64::new(0),
2780 started_at: std::time::Instant::now(),
2781 last_request_ms: std::sync::atomic::AtomicU64::new(0),
2782 detection,
2783 mcp,
2784 continuous_batching_enabled: enable_continuous_batching,
2785 loading_model: Mutex::new(None),
2786 last_load_error: Mutex::new(None),
2787 }
2788}
2789
2790fn startup_model_id() -> Option<String> {
2795 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
2796 let configured = std::fs::canonicalize(&configured).ok()?;
2797 admin::discover(&admin::model_dirs())
2798 .into_iter()
2799 .find(|d| {
2800 std::fs::canonicalize(&d.path)
2801 .map(|p| p == configured)
2802 .unwrap_or(false)
2803 })
2804 .map(|d| d.id)
2805}
2806
2807fn init_cpu_pool() {
2816 match ferrox_core::threads::init_cpu_pool() {
2817 Some(n) => eprintln!(
2818 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
2819 ferrox_core::threads::perf_core_count()
2820 ),
2821 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
2822 }
2823}
2824
2825fn announce_ready(addr: SocketAddr, scheme: &str) {
2838 use std::io::Write;
2839 let ready =
2840 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
2841 let mut stdout = std::io::stdout().lock();
2842 let _ = writeln!(stdout, "{}", ready.to_line());
2843 let _ = stdout.flush();
2844}
2845
2846async fn shutdown_signal(exit_on_stdin_close: bool) {
2858 if !exit_on_stdin_close {
2859 std::future::pending::<()>().await;
2860 return;
2861 }
2862 let _ = tokio::task::spawn_blocking(|| {
2863 use std::io::Read;
2864 let mut sink = [0u8; 256];
2865 let mut stdin = std::io::stdin().lock();
2866 loop {
2867 match stdin.read(&mut sink) {
2868 Ok(0) => break,
2870 Ok(_) => continue,
2872 Err(e) => {
2873 tracing::warn!("stdin read failed ({e}); treating it as closed");
2874 break;
2875 }
2876 }
2877 }
2878 })
2879 .await;
2880 tracing::info!("stdin closed; shutting down");
2881}
2882
2883fn tokio_worker_threads() -> usize {
2888 std::env::var("FERROX_TOKIO_WORKERS")
2889 .ok()
2890 .and_then(|v| v.trim().parse::<usize>().ok())
2891 .filter(|n| *n > 0)
2892 .unwrap_or(2)
2893}
2894
2895fn install_ring_crypto_provider() {
2912 let _ = rustls::crypto::ring::default_provider().install_default();
2913}
2914
2915pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
2921 if args.list_devices {
2922 print_available_devices();
2923 return Ok(());
2924 }
2925 apply_cli_overrides(&args)?;
2926
2927 let _instance = {
2931 use ferrox_core::instance::{register, InstancePolicy};
2932 let policy = if args.allow_multiple_instances {
2933 InstancePolicy::Multi
2934 } else {
2935 InstancePolicy::from_env_or(InstancePolicy::Single)
2936 };
2937 let model = std::env::var("FERROX_MODEL_PATH").ok();
2938 register(
2939 "server",
2940 model.as_deref(),
2941 ferrox_core::instance::current_backend(),
2942 policy,
2943 )
2944 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
2945 };
2946
2947 let journal = journal::Journal::from_env();
2948 eprintln!(
2949 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
2950 journal.path()
2951 );
2952 journal.append(&journal::Record::session_start(
2953 env!("CARGO_PKG_VERSION"),
2954 std::process::id(),
2955 ));
2956 journal::install_panic_hook(journal.clone());
2957
2958 let mcp_config_path = args.mcp_config.clone();
2959 let exit_on_stdin_close = args.exit_on_stdin_close
2960 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
2961 .map(|v| v == "1")
2962 .unwrap_or(false);
2963
2964 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
2968 init_cpu_pool();
2969
2970 let runtime = tokio::runtime::Builder::new_multi_thread()
2971 .worker_threads(tokio_worker_threads())
2972 .enable_all()
2973 .build()?;
2974 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
2975
2976 let reason = match &result {
2977 Ok(()) => "normal".to_string(),
2978 Err(e) => e.to_string(),
2979 };
2980 journal.append(&journal::Record::session_exit(reason));
2981
2982 runtime.shutdown_background();
2987
2988 result
2989}
2990
2991async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
2992 let _ = tracing_subscriber::fmt::try_init();
2998
2999 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
3009 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
3010 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
3011 .map(|v| v == "1")
3012 .unwrap_or(false);
3013 if let Err(msg) =
3014 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
3015 {
3016 anyhow::bail!(msg);
3017 }
3018
3019 let mut loaded = model::load()?;
3020 match &loaded {
3021 model::LoadedModel::Gguf(g) => tracing::info!(
3022 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
3023 g.decoder.config.name,
3024 g.is_synthetic,
3025 g.tokenizer.kind()
3026 ),
3027 model::LoadedModel::Kimi(k) => tracing::info!(
3028 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
3029 k.tokenizer.vocab_size()
3030 ),
3031 model::LoadedModel::Mla(m) => tracing::info!(
3032 "loaded MLA GGUF '{}' (tokenizer={})",
3033 m.name,
3034 m.tokenizer.kind()
3035 ),
3036 model::LoadedModel::Gemma4(m) => tracing::info!(
3037 "loaded Gemma4 GGUF '{}' (tokenizer={})",
3038 m.name,
3039 m.tokenizer.kind()
3040 ),
3041 model::LoadedModel::Glm52(g) => tracing::info!(
3042 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
3043 g.name,
3044 g.tokenizer.kind()
3045 ),
3046 }
3047 let metal_default_moe_budget = {
3055 #[cfg(feature = "metal")]
3056 {
3057 ferrox_core::metal_dense_enabled()
3058 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
3059 }
3060 #[cfg(not(feature = "metal"))]
3061 {
3062 false
3063 }
3064 };
3065 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
3066 let budget: u64 = budget_str
3067 .parse()
3068 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
3069 match &mut loaded {
3070 model::LoadedModel::Gguf(g) => {
3071 tracing::info!(
3072 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
3073 (CUDA and/or Metal matvecs when built with the matching feature)"
3074 );
3075 g.decoder.gpu_vram_budget_bytes = Some(budget);
3076 }
3077 model::LoadedModel::Kimi(_) => {
3078 tracing::warn!(
3079 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
3080 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
3081 );
3082 }
3083 model::LoadedModel::Mla(_) => {
3084 tracing::warn!(
3085 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
3086 FFN path only today; ignoring expert VRAM budget"
3087 );
3088 }
3089 model::LoadedModel::Gemma4(_) => {
3090 tracing::warn!(
3091 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
3092 ignoring expert VRAM budget"
3093 );
3094 }
3095 model::LoadedModel::Glm52(_) => {
3096 tracing::warn!(
3097 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
3098 GPU expert placement not wired yet; ignoring"
3099 );
3100 }
3101 }
3102 } else if metal_default_moe_budget {
3103 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
3107 if let model::LoadedModel::Gguf(g) = &mut loaded {
3108 tracing::info!(
3109 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
3110 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
3111 );
3112 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
3113 }
3114 }
3115 #[cfg(feature = "cuda")]
3116 {
3117 if ferrox_core::cuda_dense_enabled() {
3118 tracing::info!(
3119 "CUDA dense matvec enabled for WeightMatrix::apply \
3120 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
3121 );
3122 } else {
3123 tracing::info!(
3124 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
3125 );
3126 }
3127 }
3128 #[cfg(feature = "metal")]
3129 {
3130 if ferrox_core::metal_dense_enabled() {
3131 tracing::info!(
3132 "Metal dense matvec enabled for WeightMatrix::apply \
3133 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
3134 );
3135 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
3136 Some("1") | Some("true") | Some("on") | Some("attn") => {
3137 tracing::info!(
3138 "Metal fused attention requested (FERROX_METAL_ATTN): \
3139 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
3140 );
3141 }
3142 _ => {}
3143 }
3144 match std::env::var("FERROX_METAL_LOGITS").ok().as_deref() {
3145 Some("1") | Some("true") | Some("on") | Some("logits") => {
3146 tracing::info!(
3147 "Metal logits-in-stack enabled (FERROX_METAL_LOGITS): \
3148 final_norm+lm_head on-GPU (often slower than host lm_head)"
3149 );
3150 }
3151 _ => {}
3152 }
3153 match std::env::var("FERROX_METAL_GREEDY_GPU").ok().as_deref() {
3154 Some("0") | Some("false") | Some("off") | Some("no") => {
3155 tracing::info!(
3156 "Metal greedy GPU argmax disabled (FERROX_METAL_GREEDY_GPU=0); \
3157 temperature<=0 uses host lm_head"
3158 );
3159 }
3160 _ => {
3161 tracing::info!(
3162 "Metal greedy GPU argmax on by default (opt out FERROX_METAL_GREEDY_GPU=0): \
3163 temperature<=0 folds final_norm+lm_head+argmax into dense stack"
3164 );
3165 }
3166 }
3167 } else {
3168 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
3169 }
3170 }
3171 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
3182 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
3183 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
3184 if blocks_env.is_ok() && byte_budget_env.is_ok() {
3185 panic!(
3186 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
3187 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
3188 );
3189 }
3190 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
3191 (Ok(blocks), Ok(block_size), Err(_)) => {
3192 let total_blocks: usize = blocks
3193 .parse()
3194 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
3195 let block_size: usize = block_size
3196 .parse()
3197 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
3198 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
3204 .ok()
3205 .map(|v| {
3206 v.parse()
3207 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
3208 })
3209 .unwrap_or(0);
3210 tracing::info!(
3211 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
3212 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
3213 );
3214 Some(generate::KvPoolConfig {
3215 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
3216 queue_wait: Duration::from_millis(queue_wait_ms),
3217 })
3218 }
3219 (Err(_), Ok(block_size), Ok(byte_budget)) => {
3220 let block_size: usize = block_size
3221 .parse()
3222 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
3223 let budget: u64 = byte_budget
3224 .parse()
3225 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
3226 let cfg = match &loaded {
3227 model::LoadedModel::Gguf(g) => &g.decoder.config,
3228 model::LoadedModel::Kimi(_)
3229 | model::LoadedModel::Mla(_)
3230 | model::LoadedModel::Gemma4(_)
3231 | model::LoadedModel::Glm52(_) => {
3232 panic!(
3233 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
3234 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
3235 );
3236 }
3237 };
3238 let bytes_per_block = block_size
3239 * cfg.n_layers
3240 * cfg.n_kv_heads
3241 * cfg.head_dim
3242 * 2
3243 * std::mem::size_of::<f32>();
3244 assert!(
3245 bytes_per_block > 0,
3246 "derived KV block byte size must be positive (check model config and block size)"
3247 );
3248 let total_blocks = (budget as usize / bytes_per_block).max(1);
3249 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
3250 .ok()
3251 .map(|v| {
3252 v.parse()
3253 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
3254 })
3255 .unwrap_or(0);
3256 tracing::info!(
3257 "KV cache block pool enabled from byte budget: {budget} bytes / \
3258 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
3259 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
3260 cfg.n_layers
3261 );
3262 Some(generate::KvPoolConfig {
3263 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
3264 queue_wait: Duration::from_millis(queue_wait_ms),
3265 })
3266 }
3267 (Err(_), Err(_), Err(_)) => None,
3268 (Err(_), Ok(_), Err(_)) => panic!(
3269 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
3270 (or unset all three to disable KV cache pooling)"
3271 ),
3272 (Ok(_), Ok(_), Ok(_)) => {
3273 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
3274 }
3275 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
3276 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
3277 set together (or neither, to disable KV cache pooling)"
3278 ),
3279 };
3280 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
3287 let max_entries: usize = v
3288 .parse()
3289 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
3290 if kv_pool.is_some() {
3291 tracing::warn!(
3292 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
3293 caching will never be consulted while a KV pool is configured"
3294 );
3295 }
3296 tracing::info!(
3297 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
3298 all requests"
3299 );
3300 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
3301 });
3302 if matches!(
3303 loaded,
3304 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
3305 ) && (kv_pool.is_some() || prefix_cache.is_some())
3306 {
3307 tracing::warn!(
3308 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
3309 neither is consulted for those engines (state shapes differ from Decoder KV); see \
3310 ferrox_models::engine's module docs"
3311 );
3312 }
3313 let enable_cb = std::env::var("FERROX_CONTINUOUS_BATCHING")
3314 .map(|v| v == "1")
3315 .unwrap_or(false)
3316 && kv_pool.is_none()
3317 && prefix_cache.is_none()
3318 && matches!(loaded, model::LoadedModel::Gguf(_));
3319 if std::env::var("FERROX_CONTINUOUS_BATCHING")
3320 .map(|v| v == "1")
3321 .unwrap_or(false)
3322 && (kv_pool.is_some() || prefix_cache.is_some())
3323 {
3324 tracing::warn!(
3325 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
3326 (those modes keep the private generate path)"
3327 );
3328 }
3329 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
3330 if let Ok(chunk) = n.parse::<usize>() {
3331 if chunk > 0 {
3332 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
3333 }
3334 }
3335 }
3336 if matches!(
3337 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
3338 Some("1")
3339 ) {
3340 tracing::warn!(
3341 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
3342 (minimal spill; full layer offload still planned)"
3343 );
3344 }
3345
3346 let mcp = match mcp_config_path {
3347 Some(path) => {
3348 let loaded = mcp::load_mcp_config(&path)?;
3349 tracing::info!(
3350 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
3351 loaded.path,
3352 loaded.servers.len()
3353 );
3354 Some(loaded)
3355 }
3356 None => None,
3357 };
3358
3359 let detection = health::Detection::spawn();
3363
3364 let state = Arc::new(build_app_state(
3365 loaded,
3366 kv_pool,
3367 prefix_cache,
3368 enable_cb,
3369 mcp,
3370 detection,
3371 ));
3372
3373 use ferrox_api::routes;
3377
3378 let public = Router::new().route(routes::HEALTH, get(health));
3383
3384 let mut protected = Router::new()
3385 .route(routes::V1_MODELS, get(list_models))
3386 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3387 .route(routes::V1_CANCEL, post(cancel_generation))
3391 .route(&resume_route(routes::V1_STREAM), get(resume::resume))
3396 .route(&resume_route(routes::V1_STREAM_POLL), get(resume::poll))
3397 .route(routes::V1_MESSAGES, post(anthropic::messages))
3398 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3399 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3400 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3401 .route(routes::V1_EMBEDDINGS, post(openai_extra::embeddings))
3402 .route(routes::CACHE_STATS, get(cache_stats))
3403 .route(routes::METRICS, get(metrics))
3404 .route(routes::ADMIN_MODELS, get(admin::models))
3409 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3410 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3411 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3412 .route(routes::ADMIN_TASKS, get(admin::tasks))
3413 .route(&admin::cancel_route(), post(admin::cancel_task))
3414 .route(routes::ADMIN_STATS, get(admin::stats));
3415
3416 if let Ok(key) = std::env::var("FERROX_API_KEY") {
3421 tracing::info!("API key auth enabled");
3422 let auth = limits::AuthConfig {
3423 api_key: Arc::new(key),
3424 };
3425 protected = protected.route_layer(axum::middleware::from_fn_with_state(
3426 auth,
3427 limits::require_api_key,
3428 ));
3429 }
3430 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
3431 let rpm: u32 = rpm
3432 .parse()
3433 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
3434 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
3435 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
3436 protected = protected.route_layer(axum::middleware::from_fn_with_state(
3437 limiter,
3438 limits::rate_limit,
3439 ));
3440 }
3441 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
3449 let origins = security::parse_cors_origins(&spec)
3450 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
3451 tracing::info!(
3452 "CORS enabled: {} allow-listed origin(s) ({})",
3453 origins.len(),
3454 spec
3455 );
3456 let cors = tower_http::cors::CorsLayer::new()
3457 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
3458 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
3459 .allow_headers([
3460 axum::http::header::CONTENT_TYPE,
3461 axum::http::header::AUTHORIZATION,
3462 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
3468 axum::http::HeaderName::from_static("last-event-id"),
3472 ]);
3473 protected = protected.route_layer(cors);
3474 }
3475
3476 let app = public
3480 .merge(protected)
3481 .layer(axum::middleware::from_fn(limits::retry_after))
3482 .with_state(state);
3483
3484 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
3490 install_ring_crypto_provider();
3491 match tls_paths {
3496 Some(paths) => {
3497 let config =
3498 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
3499 .await
3500 .map_err(|e| {
3501 anyhow::anyhow!(
3502 "failed to load TLS cert/key ({:?}, {:?}): {e}",
3503 paths.cert,
3504 paths.key
3505 )
3506 })?;
3507 let socket_addr: std::net::SocketAddr = addr
3508 .parse()
3509 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
3510 let listener = std::net::TcpListener::bind(socket_addr)?;
3511 listener.set_nonblocking(true)?;
3519 let bound = listener.local_addr()?;
3520 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
3521 announce_ready(bound, "https");
3522
3523 let handle = axum_server::Handle::new();
3524 let shutdown_handle = handle.clone();
3525 tokio::spawn(async move {
3526 shutdown_signal(exit_on_stdin_close).await;
3527 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
3528 });
3529 axum_server::from_tcp_rustls(listener, config)?
3530 .handle(handle)
3531 .serve(app.into_make_service())
3532 .await?;
3533 }
3534 None => {
3535 let listener = tokio::net::TcpListener::bind(&addr).await?;
3536 let bound = listener.local_addr()?;
3537 tracing::info!("ferrox-server listening on {bound}");
3538 announce_ready(bound, "http");
3539 axum::serve(listener, app)
3540 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
3541 .await?;
3542 }
3543 }
3544 Ok(())
3545}
3546
3547#[cfg(test)]
3548mod tests {
3549 use super::*;
3550 use ferrox_models::config::test_dense_fixture;
3551
3552 #[test]
3553 fn parses_llama_server_style_options() {
3554 let argv = [
3555 "ferrox-server",
3556 "-m",
3557 "model.gguf",
3558 "--host",
3559 "::1",
3560 "--port",
3561 "9000",
3562 "-t",
3563 "4",
3564 "-dev",
3565 "Metal",
3566 "-ngl",
3567 "all",
3568 ]
3569 .into_iter()
3570 .map(String::from)
3571 .collect();
3572 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
3573
3574 assert_eq!(args.model.as_deref(), Some("model.gguf"));
3575 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
3576 assert_eq!(args.port, Some(9000));
3577 assert_eq!(args.threads, Some(4));
3578 assert_eq!(args.device, Some(OffloadDevice::Metal));
3579 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
3580 assert_eq!(
3581 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
3582 Some("[::1]:9000")
3583 );
3584 }
3585
3586 #[test]
3587 fn port_zero_survives_argument_parsing_as_a_real_request() {
3588 let argv = ["ferrox-server", "--port", "0"]
3593 .into_iter()
3594 .map(String::from)
3595 .collect();
3596 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
3597 assert_eq!(args.port, Some(0));
3598 assert_eq!(
3599 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
3600 Some("127.0.0.1:0")
3601 );
3602 }
3603
3604 #[test]
3605 fn stdin_close_exit_is_opt_in() {
3606 let args =
3609 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
3610 assert!(!args.exit_on_stdin_close);
3611 let args = ServerArgs::try_parse_from(
3612 ["ferrox-server", "--exit-on-stdin-close"]
3613 .into_iter()
3614 .map(String::from),
3615 )
3616 .unwrap();
3617 assert!(args.exit_on_stdin_close);
3618 }
3619
3620 #[test]
3621 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
3622 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
3623 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
3624 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
3625 assert_eq!(parsed.port, 51999);
3626 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
3627 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
3629 }
3630
3631 fn test_model() -> Model {
3632 let cfg = test_dense_fixture();
3636 Model::Gguf(GgufModel {
3637 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
3638 tokenizer: Arc::new(ServerTokenizer::Byte),
3639 stop_tokens: StopTokens::default(),
3640 bos_id: None,
3641 is_synthetic: true,
3642 chat_template: chat_template::ChatTemplate::Plain,
3643 })
3644 }
3645
3646 fn greedy_params(max_tokens: usize) -> GenerationParams {
3647 GenerationParams {
3648 max_tokens,
3649 sampling: SamplingParams::default(),
3650 seed: 1,
3651 stop: Vec::new(),
3652 stop_token_ids: Vec::new(),
3653 json_object: false,
3654 cancel: None,
3655 }
3656 }
3657
3658 fn test_model_full_byte_vocab() -> Model {
3662 let mut cfg = test_dense_fixture();
3663 cfg.vocab_size = 256;
3664 Model::Gguf(GgufModel {
3665 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
3666 tokenizer: Arc::new(ServerTokenizer::Byte),
3667 stop_tokens: StopTokens::default(),
3668 bos_id: None,
3669 is_synthetic: true,
3670 chat_template: chat_template::ChatTemplate::Plain,
3671 })
3672 }
3673
3674 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
3678 AppState {
3679 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3680 id: None,
3681 model: Arc::new(model),
3682 batcher: None,
3683 ceiling: None,
3684 }))),
3685 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3686 tasks: Arc::new(tasks::TaskRegistry::new()),
3687 cancels: Arc::new(cancel::CancelRegistry::new()),
3688 stats: stats::Stats::new(),
3689 streams: resume::StreamRegistry::new(),
3690 model_dir: None,
3691 response_cache: Mutex::new(response_cache),
3692 kv_pool: None,
3693 prefix_cache: None,
3694 sessions: session::SessionStore::new(),
3695 requests_total: std::sync::atomic::AtomicU64::new(0),
3696 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3697 started_at: std::time::Instant::now(),
3698 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3699 detection: Arc::new(health::Detection::ready(health::probe_backends())),
3700 mcp: None,
3701 continuous_batching_enabled: false,
3702 loading_model: Mutex::new(None),
3703 last_load_error: Mutex::new(None),
3704 }
3705 }
3706
3707 fn test_app() -> Router {
3716 test_app_with_state(Arc::new(test_state(
3717 test_model_full_byte_vocab(),
3718 ResponseCache::new(1000, Duration::from_secs(3600)),
3719 )))
3720 }
3721
3722 fn test_app_with_state(state: Arc<AppState>) -> Router {
3725 Router::new()
3726 .route(ferrox_api::routes::HEALTH, get(health))
3727 .route(ferrox_api::routes::V1_MODELS, get(list_models))
3728 .route("/v1/chat/completions", post(chat_completions))
3729 .route("/v1/tokenize", post(openai_extra::tokenize))
3730 .route("/v1/detokenize", post(openai_extra::detokenize))
3731 .route("/v1/embeddings", post(openai_extra::embeddings))
3732 .route("/v1/completions", post(openai_extra::completions))
3733 .route(
3734 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
3735 post(admin::unload_model),
3736 )
3737 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
3738 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
3739 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
3740 .route(
3741 &resume_route(ferrox_api::routes::V1_STREAM),
3742 get(resume::resume),
3743 )
3744 .route(
3745 &resume_route(ferrox_api::routes::V1_STREAM_POLL),
3746 get(resume::poll),
3747 )
3748 .with_state(state)
3749 }
3750
3751 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
3752 let mut cfg = test_dense_fixture();
3753 cfg.name = name;
3754 cfg.vocab_size = vocab_size;
3755 Model::Gguf(GgufModel {
3756 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
3757 tokenizer: Arc::new(ServerTokenizer::Byte),
3758 stop_tokens: StopTokens::default(),
3759 bos_id: None,
3760 is_synthetic: true,
3761 chat_template: chat_template::ChatTemplate::Plain,
3762 })
3763 }
3764
3765 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
3766 Arc::new(ActiveModel {
3767 id: Some(name.to_string()),
3768 model: Arc::new(named_test_model(name, 256)),
3769 batcher: None,
3770 ceiling: None,
3771 })
3772 .tap_into(state)
3773 }
3774
3775 trait TapInto {
3777 fn tap_into(self, state: &AppState) -> Self;
3778 }
3779 impl TapInto for Arc<ActiveModel> {
3780 fn tap_into(self, state: &AppState) -> Self {
3781 state.swap_active(Some(Arc::clone(&self)));
3782 self
3783 }
3784 }
3785
3786 #[test]
3792 fn an_in_flight_request_keeps_the_model_it_started_on() {
3793 let state = test_state(
3794 named_test_model("model-a", 256),
3795 ResponseCache::new(4, Duration::from_secs(60)),
3796 );
3797
3798 let in_flight = state.active().expect("a model is loaded");
3801 assert_eq!(in_flight.model.name(), "model-a");
3802
3803 active_model(&state, "model-b");
3804
3805 assert_eq!(state.active().unwrap().model.name(), "model-b");
3807 assert_eq!(in_flight.model.name(), "model-a");
3809 let (_chunks, finish, _usage) = run_generation(
3810 &in_flight.model,
3811 "hi",
3812 &greedy_params(3),
3813 None,
3814 None,
3815 None,
3816 None,
3817 )
3818 .expect("the old model must still decode after being swapped out");
3819 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
3820 }
3821
3822 #[test]
3827 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
3828 let state = test_state(
3829 named_test_model("model-a", 256),
3830 ResponseCache::new(4, Duration::from_secs(60)),
3831 );
3832 let in_flight = state.active().expect("a model is loaded");
3833 let weights = Arc::clone(&in_flight.model);
3834 assert!(Arc::strong_count(&weights) >= 2);
3835
3836 let previous = state.swap_active(Some(Arc::new(ActiveModel {
3837 id: Some("model-b".to_string()),
3838 model: Arc::new(named_test_model("model-b", 256)),
3839 batcher: None,
3840 ceiling: None,
3841 })));
3842 drop(previous);
3843 assert!(Arc::strong_count(&weights) >= 2);
3845 drop(in_flight);
3846 assert_eq!(Arc::strong_count(&weights), 1);
3847 }
3848
3849 #[tokio::test]
3853 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
3854 let state = Arc::new(test_state(
3855 named_test_model("model-a", 256),
3856 ResponseCache::new(4, Duration::from_secs(60)),
3857 ));
3858 let app = test_app_with_state(Arc::clone(&state));
3859
3860 let (status, body) = post_json_uri(
3861 &app,
3862 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
3863 serde_json::json!({}),
3864 )
3865 .await;
3866 assert_eq!(status, StatusCode::OK);
3867 assert_eq!(body["ok"], true);
3868 assert!(body["active"].is_null());
3869 assert!(state.active().is_none());
3870
3871 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
3872 assert_eq!(status, StatusCode::OK);
3873 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
3874 assert_eq!(models["data"].as_array().unwrap().len(), 0);
3875
3876 let (status, body) = post_json_uri(
3877 &app,
3878 "/v1/chat/completions",
3879 serde_json::json!({
3880 "model": "x",
3881 "messages": [{"role": "user", "content": "hi"}]
3882 }),
3883 )
3884 .await;
3885 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
3886 assert_eq!(body["error"]["type"], "model_not_loaded");
3887 }
3888
3889 #[tokio::test]
3893 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
3894 let state = Arc::new(test_state(
3895 named_test_model("model-a", 256),
3896 ResponseCache::new(4, Duration::from_secs(60)),
3897 ));
3898 let app = test_app_with_state(Arc::clone(&state));
3899 state.swap_active(None);
3900
3901 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
3902 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
3905 assert_eq!(body["state"], "unavailable");
3906 assert_eq!(body["reason"], "model_not_loaded");
3907 assert!(body["model"].is_null());
3908 let real_weights = body["capabilities"]
3909 .as_array()
3910 .unwrap()
3911 .iter()
3912 .find(|c| c["id"] == "real_weights")
3913 .cloned()
3914 .expect("real_weights is always reported");
3915 assert_eq!(real_weights["available"], false);
3916 assert_eq!(real_weights["reason"], "model_not_loaded");
3917 }
3918
3919 #[tokio::test]
3923 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
3924 let app = test_app();
3925
3926 let (status, completion) = post_json_uri(
3927 &app,
3928 "/v1/chat/completions",
3929 serde_json::json!({
3930 "model": "x",
3931 "messages": [{"role": "user", "content": "hi"}],
3932 "max_tokens": 4
3933 }),
3934 )
3935 .await;
3936 assert_eq!(status, StatusCode::OK);
3937 let request_id = completion["request_id"].as_str().unwrap().to_string();
3938
3939 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
3940 assert_eq!(status, StatusCode::OK);
3941 let recent = stats["recent"].as_array().unwrap();
3942 assert_eq!(recent.len(), 1);
3943 let row = &recent[0];
3944 assert_eq!(row["request_id"], request_id);
3945 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
3946 assert_eq!(row["status"], 200);
3947 assert_eq!(row["stream"], false);
3948 assert!(row["duration_ms"].is_number());
3951 assert!(row["decode_ms"].is_number());
3952 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
3953 assert_eq!(
3954 stats["tokens_prompt_total"].as_u64().unwrap(),
3955 row["prompt_tokens"].as_u64().unwrap()
3956 );
3957 }
3958
3959 #[tokio::test]
3963 async fn a_rejected_request_is_recorded_too() {
3964 let state = Arc::new(test_state(
3965 named_test_model("model-a", 256),
3966 ResponseCache::new(4, Duration::from_secs(60)),
3967 ));
3968 let app = test_app_with_state(Arc::clone(&state));
3969 state.swap_active(None);
3970
3971 let (status, _) = post_json_uri(
3972 &app,
3973 "/v1/chat/completions",
3974 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
3975 )
3976 .await;
3977 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
3978
3979 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
3980 let recent = stats["recent"].as_array().unwrap();
3981 assert_eq!(recent.len(), 1);
3982 assert_eq!(recent[0]["status"], 503);
3983 assert_eq!(recent[0]["completion_tokens"], 0);
3984 assert!(recent[0]["decode_ms"].is_null());
3985 assert_eq!(stats["errors_total"], 1);
3986 }
3987
3988 async fn post_json_with_headers(
3992 app: &Router,
3993 uri: &str,
3994 body: serde_json::Value,
3995 headers: &[(&str, &str)],
3996 ) -> (StatusCode, serde_json::Value) {
3997 use http_body_util::BodyExt;
3998 use tower::ServiceExt;
3999
4000 let mut builder = axum::http::Request::builder()
4001 .method("POST")
4002 .uri(uri)
4003 .header("content-type", "application/json");
4004 for (name, value) in headers {
4005 builder = builder.header(*name, *value);
4006 }
4007 let response = app
4008 .clone()
4009 .oneshot(
4010 builder
4011 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
4012 .unwrap(),
4013 )
4014 .await
4015 .unwrap();
4016 let status = response.status();
4017 let bytes = response.into_body().collect().await.unwrap().to_bytes();
4018 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
4019 (status, json)
4020 }
4021
4022 #[tokio::test]
4026 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
4027 let app = test_app();
4028
4029 let (status, _) = post_json_uri(
4030 &app,
4031 ferrox_api::routes::V1_TOKENIZE,
4032 serde_json::json!({"prompt": "hello"}),
4033 )
4034 .await;
4035 assert_eq!(status, StatusCode::OK);
4036 let (status, _) = post_json_uri(
4037 &app,
4038 ferrox_api::routes::V1_DETOKENIZE,
4039 serde_json::json!({"tokens": [104, 105]}),
4040 )
4041 .await;
4042 assert_eq!(status, StatusCode::OK);
4043 let (status, _) = post_json_uri(
4044 &app,
4045 ferrox_api::routes::V1_EMBEDDINGS,
4046 serde_json::json!({"input": "hello"}),
4047 )
4048 .await;
4049 assert_eq!(status, StatusCode::OK);
4050
4051 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4052 let routes: Vec<&str> = stats["recent"]
4053 .as_array()
4054 .unwrap()
4055 .iter()
4056 .map(|row| row["route"].as_str().unwrap())
4057 .collect();
4058 for expected in [
4059 ferrox_api::routes::V1_TOKENIZE,
4060 ferrox_api::routes::V1_DETOKENIZE,
4061 ferrox_api::routes::V1_EMBEDDINGS,
4062 ] {
4063 assert!(
4064 routes.contains(&expected),
4065 "{expected} is missing: {routes:?}"
4066 );
4067 }
4068
4069 let row = |route: &str| {
4070 stats["recent"]
4071 .as_array()
4072 .unwrap()
4073 .iter()
4074 .find(|r| r["route"] == route)
4075 .cloned()
4076 .unwrap()
4077 };
4078 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
4082 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
4083 assert!(embed["decode_ms"].is_null());
4084 assert_eq!(embed["completion_tokens"], 0);
4085 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
4089 assert_eq!(
4090 stats["tokens_prompt_total"].as_u64().unwrap(),
4091 embed["prompt_tokens"].as_u64().unwrap(),
4092 "only the forward pass counted"
4093 );
4094 }
4095
4096 #[tokio::test]
4099 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
4100 let app = test_app();
4101 let (status, _) = post_json_uri(
4102 &app,
4103 ferrox_api::routes::V1_EMBEDDINGS,
4104 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
4105 )
4106 .await;
4107 assert_eq!(status, StatusCode::BAD_REQUEST);
4108
4109 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4110 let recent = stats["recent"].as_array().unwrap();
4111 assert_eq!(recent.len(), 1);
4112 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
4113 assert_eq!(recent[0]["status"], 400);
4114 assert_eq!(
4115 recent[0]["prompt_tokens"], 0,
4116 "a rejected call embedded nothing"
4117 );
4118 }
4119
4120 #[tokio::test]
4123 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
4124 let app = test_app();
4125 let key = "sk-monitor-secret";
4126 let (status, _) = post_json_with_headers(
4127 &app,
4128 "/v1/chat/completions",
4129 serde_json::json!({
4130 "model": "x",
4131 "messages": [{"role": "user", "content": "hi"}],
4132 "max_tokens": 2
4133 }),
4134 &[
4135 ("authorization", &format!("Bearer {key}")),
4136 ("x-ferrox-client", "ferrox-studio"),
4137 ],
4138 )
4139 .await;
4140 assert_eq!(status, StatusCode::OK);
4141
4142 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4143 let row = stats["recent"].as_array().unwrap()[0].clone();
4144 let fingerprint = row["via_api_key"]
4145 .as_str()
4146 .expect("the row names the key that served it")
4147 .to_string();
4148 assert_eq!(fingerprint, attribution::key_fingerprint(key));
4149 assert!(!fingerprint.contains(key));
4150 assert!(
4151 !serde_json::to_string(&stats).unwrap().contains(key),
4152 "the stats payload must not carry the key in any form"
4153 );
4154 assert_eq!(row["client"], "ferrox-studio");
4155 }
4156
4157 #[tokio::test]
4160 async fn different_keys_are_different_callers_and_no_key_is_null() {
4161 let app = test_app();
4162 let body = serde_json::json!({
4163 "model": "x",
4164 "messages": [{"role": "user", "content": "hi"}],
4165 "max_tokens": 1
4166 });
4167 for headers in [
4168 vec![("authorization", "Bearer key-one")],
4169 vec![("authorization", "Bearer key-two")],
4170 vec![],
4171 ] {
4172 let (status, _) =
4173 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
4174 assert_eq!(status, StatusCode::OK);
4175 }
4176
4177 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4178 let recent = stats["recent"].as_array().unwrap();
4179 assert_eq!(recent.len(), 3);
4180 let one = recent[0]["via_api_key"].as_str().unwrap();
4181 let two = recent[1]["via_api_key"].as_str().unwrap();
4182 assert_ne!(one, two, "two keys must not collapse into one caller");
4183 assert!(
4184 recent[2]["via_api_key"].is_null(),
4185 "an unauthenticated call is null, not a fingerprint of nothing"
4186 );
4187 assert!(recent[2]["client"].is_null());
4188 }
4189
4190 #[tokio::test]
4195 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
4196 let state = Arc::new(test_state(
4197 named_test_model("really-loaded", 256),
4198 ResponseCache::new(4, Duration::from_secs(60)),
4199 ));
4200 let app = test_app_with_state(Arc::clone(&state));
4201
4202 let (status, _) = post_json_uri(
4203 &app,
4204 "/v1/chat/completions",
4205 serde_json::json!({
4206 "model": "gpt-4-turbo-that-is-not-here",
4207 "messages": [{"role": "user", "content": "hi"}],
4208 "max_tokens": 2
4209 }),
4210 )
4211 .await;
4212 assert_eq!(status, StatusCode::OK);
4213
4214 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4215 assert_eq!(stats["recent"][0]["model"], "really-loaded");
4216
4217 state.swap_active(None);
4220 let (status, _) = post_json_uri(
4221 &app,
4222 "/v1/chat/completions",
4223 serde_json::json!({
4224 "model": "gpt-4-turbo-that-is-not-here",
4225 "messages": [{"role": "user", "content": "hi"}]
4226 }),
4227 )
4228 .await;
4229 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
4230 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4231 let recent = stats["recent"].as_array().unwrap();
4232 assert!(recent[recent.len() - 1]["model"].is_null());
4233 }
4234
4235 #[tokio::test]
4239 async fn a_streamed_row_names_the_model_it_decoded_against() {
4240 let state = Arc::new(test_state(
4241 named_test_model("model-before", 256),
4242 ResponseCache::new(4, Duration::from_secs(60)),
4243 ));
4244 let app = test_app_with_state(Arc::clone(&state));
4245 let _ = post_sse_raw(&app, resumable_request()).await;
4246 active_model(&state, "model-after");
4248
4249 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4250 assert_eq!(stats["recent"][0]["model"], "model-before");
4251 }
4252
4253 #[tokio::test]
4256 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
4257 let app = test_app();
4258 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4259 assert_eq!(status, StatusCode::OK);
4260 assert!(
4261 stats["queue_depth"].is_null(),
4262 "without continuous batching nothing queues, so there is nothing to measure"
4263 );
4264 assert!(stats["queue_rejected_total"].is_null());
4265 assert_eq!(
4266 stats["generating_now"], 0,
4267 "work in progress is measured and really is zero here"
4268 );
4269 }
4270
4271 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
4276 use http_body_util::BodyExt;
4277 use tower::ServiceExt;
4278
4279 let response = app
4280 .clone()
4281 .oneshot(
4282 axum::http::Request::builder()
4283 .method("POST")
4284 .uri("/v1/chat/completions")
4285 .header("content-type", "application/json")
4286 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
4287 .unwrap(),
4288 )
4289 .await
4290 .unwrap();
4291 let bytes = response.into_body().collect().await.unwrap().to_bytes();
4292 String::from_utf8(bytes.to_vec()).unwrap()
4293 }
4294
4295 async fn get_json_with_headers(
4296 app: &Router,
4297 uri: &str,
4298 headers: &[(&str, &str)],
4299 ) -> (StatusCode, serde_json::Value) {
4300 use http_body_util::BodyExt;
4301 use tower::ServiceExt;
4302
4303 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
4304 for (name, value) in headers {
4305 builder = builder.header(*name, *value);
4306 }
4307 let response = app
4308 .clone()
4309 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
4310 .await
4311 .unwrap();
4312 let status = response.status();
4313 let bytes = response.into_body().collect().await.unwrap().to_bytes();
4314 (
4315 status,
4316 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
4317 )
4318 }
4319
4320 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
4321 body.lines()
4322 .filter_map(|line| line.strip_prefix(field))
4323 .map(str::trim)
4324 .collect()
4325 }
4326
4327 fn resumable_request() -> serde_json::Value {
4328 serde_json::json!({
4329 "model": "m",
4330 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
4331 "max_tokens": 4,
4332 "temperature": 0,
4333 "stream": true,
4334 "stream_resumable": true,
4335 })
4336 }
4337
4338 #[tokio::test]
4343 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
4344 let app = test_app();
4345 let body = post_sse_raw(&app, resumable_request()).await;
4346
4347 let request_id = body
4348 .lines()
4349 .find_map(|l| l.strip_prefix("data: "))
4350 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
4351 .and_then(|v| v["request_id"].as_str().map(str::to_string))
4352 .expect("the first chunk names the request");
4353
4354 let ids = sse_field(&body, "id:");
4355 let datas = sse_field(&body, "data:");
4356 assert_eq!(
4357 ids.len(),
4358 datas.len(),
4359 "every event carries an id, or a reconnect cannot name where it stopped"
4360 );
4361 for (i, id) in ids.iter().enumerate() {
4362 assert_eq!(*id, format!("{request_id}:{i}"));
4363 }
4364 let retries = sse_field(&body, "retry:");
4365 assert_eq!(
4366 retries.len(),
4367 1,
4368 "the reconnect delay is stated once, not on every event"
4369 );
4370 assert_eq!(retries[0], "1500");
4371 assert!(
4372 body.contains("data: [DONE]"),
4373 "the end of stream is still stated"
4374 );
4375 }
4376
4377 #[tokio::test]
4381 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
4382 let app = test_app();
4383 let mut request = resumable_request();
4384 request["stream_resumable"] = serde_json::json!(false);
4385 let body = post_sse_raw(&app, request).await;
4386 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
4387 assert!(
4388 sse_field(&body, "id:").is_empty(),
4389 "an id promises a replay this stream cannot serve"
4390 );
4391 assert!(sse_field(&body, "retry:").is_empty());
4392 }
4393
4394 #[tokio::test]
4398 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
4399 let app = test_app();
4400 let body = post_sse_raw(&app, resumable_request()).await;
4401 let request_id = sse_field(&body, "id:")[0]
4402 .rsplit_once(':')
4403 .unwrap()
4404 .0
4405 .to_string();
4406 let streamed: Vec<String> = sse_field(&body, "data:")
4407 .iter()
4408 .map(|d| d.to_string())
4409 .collect();
4410
4411 let (status, polled) = get_json(
4412 &app,
4413 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
4414 )
4415 .await;
4416 assert_eq!(status, StatusCode::OK);
4417 let events: Vec<String> = polled["events"]
4418 .as_array()
4419 .unwrap()
4420 .iter()
4421 .map(|e| e["data"].as_str().unwrap().to_string())
4422 .collect();
4423 assert_eq!(
4424 events, streamed,
4425 "the fallback must deliver the same answer, not a re-run of it"
4426 );
4427 assert_eq!(polled["request_id"], request_id);
4428 assert_eq!(
4429 polled["done"], false,
4430 "events were still being handed out, so the client must ask again"
4431 );
4432
4433 let next = polled["next_index"].as_u64().unwrap();
4436 let (_, drained) = get_json(
4437 &app,
4438 &format!(
4439 "{}?from={next}",
4440 ferrox_api::routes::v1_stream_poll(&request_id)
4441 ),
4442 )
4443 .await;
4444 assert_eq!(drained["done"], true);
4445 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
4446 }
4447
4448 #[tokio::test]
4452 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
4453 let app = test_app();
4454 let body = post_sse_raw(&app, resumable_request()).await;
4455 let ids = sse_field(&body, "id:");
4456 let datas: Vec<String> = sse_field(&body, "data:")
4457 .iter()
4458 .map(|d| d.to_string())
4459 .collect();
4460 assert!(
4461 ids.len() >= 3,
4462 "need a few events to resume into the middle"
4463 );
4464 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
4465
4466 let (status, resumed) = get_json_with_headers(
4467 &app,
4468 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
4469 &[],
4470 )
4471 .await;
4472 assert_eq!(status, StatusCode::OK);
4473 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
4474
4475 let (_, tail) = get_json(
4477 &app,
4478 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
4479 )
4480 .await;
4481 let tail_events: Vec<String> = tail["events"]
4482 .as_array()
4483 .unwrap()
4484 .iter()
4485 .map(|e| e["data"].as_str().unwrap().to_string())
4486 .collect();
4487 assert_eq!(tail_events, datas[2..].to_vec());
4488 }
4489
4490 #[tokio::test]
4493 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
4494 use http_body_util::BodyExt;
4495 use tower::ServiceExt;
4496
4497 let app = test_app();
4498 let body = post_sse_raw(&app, resumable_request()).await;
4499 let ids = sse_field(&body, "id:");
4500 let datas: Vec<String> = sse_field(&body, "data:")
4501 .iter()
4502 .map(|d| d.to_string())
4503 .collect();
4504 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
4505
4506 let response = app
4507 .clone()
4508 .oneshot(
4509 axum::http::Request::builder()
4510 .method("GET")
4511 .uri(ferrox_api::routes::v1_stream(&request_id))
4512 .header("last-event-id", format!("{request_id}:0"))
4513 .body(axum::body::Body::empty())
4514 .unwrap(),
4515 )
4516 .await
4517 .unwrap();
4518 assert_eq!(response.status(), StatusCode::OK);
4519 assert_eq!(
4520 response
4521 .headers()
4522 .get("x-accel-buffering")
4523 .and_then(|v| v.to_str().ok()),
4524 Some("no"),
4525 "the reconnect needs the same anti-buffering header as the stream"
4526 );
4527 let bytes = response.into_body().collect().await.unwrap().to_bytes();
4528 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
4529 assert_eq!(
4530 sse_field(&resumed, "data:")
4531 .iter()
4532 .map(|d| d.to_string())
4533 .collect::<Vec<_>>(),
4534 datas[1..].to_vec()
4535 );
4536 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
4537 }
4538
4539 #[tokio::test]
4543 async fn a_last_event_id_from_another_stream_is_refused() {
4544 let app = test_app();
4545 let body = post_sse_raw(&app, resumable_request()).await;
4546 let request_id = sse_field(&body, "id:")[0]
4547 .rsplit_once(':')
4548 .unwrap()
4549 .0
4550 .to_string();
4551
4552 let (status, err) = get_json_with_headers(
4553 &app,
4554 &ferrox_api::routes::v1_stream(&request_id),
4555 &[("last-event-id", "chatcmpl-someone-else:3")],
4556 )
4557 .await;
4558 assert_eq!(status, StatusCode::BAD_REQUEST);
4559 assert_eq!(err["error"]["code"], "bad_last_event_id");
4560 }
4561
4562 #[tokio::test]
4566 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
4567 let app = test_app();
4568 let mut request = resumable_request();
4569 request["stream_resumable"] = serde_json::json!(false);
4570 let body = post_sse_raw(&app, request).await;
4571 let request_id = body
4572 .lines()
4573 .find_map(|l| l.strip_prefix("data: "))
4574 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
4575 .and_then(|v| v["request_id"].as_str().map(str::to_string))
4576 .unwrap();
4577
4578 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
4579 assert_eq!(status, StatusCode::NOT_FOUND);
4580 assert_eq!(err["error"]["code"], "stream_not_found");
4581 assert!(err["error"]["message"]
4582 .as_str()
4583 .unwrap()
4584 .contains("stream_resumable"));
4585 }
4586
4587 #[test]
4591 fn the_axum_stream_patterns_match_the_published_templates() {
4592 assert_eq!(
4593 resume_route(ferrox_api::routes::V1_STREAM),
4594 "/v1/stream/:request_id"
4595 );
4596 assert_eq!(
4597 resume_route(ferrox_api::routes::V1_STREAM_POLL),
4598 "/v1/stream/:request_id/poll"
4599 );
4600 assert_eq!(
4601 ferrox_api::routes::v1_stream("abc"),
4602 resume_route(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
4603 );
4604 }
4605
4606 #[tokio::test]
4609 async fn the_task_list_starts_empty_rather_than_absent() {
4610 let app = test_app();
4611 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
4612 assert_eq!(status, StatusCode::OK);
4613 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
4614 }
4615
4616 async fn post_json_uri(
4617 app: &Router,
4618 uri: &str,
4619 body: serde_json::Value,
4620 ) -> (StatusCode, serde_json::Value) {
4621 use http_body_util::BodyExt;
4622 use tower::ServiceExt;
4623
4624 let response = app
4625 .clone()
4626 .oneshot(
4627 axum::http::Request::builder()
4628 .method("POST")
4629 .uri(uri)
4630 .header("content-type", "application/json")
4631 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
4632 .unwrap(),
4633 )
4634 .await
4635 .unwrap();
4636 let status = response.status();
4637 let bytes = response.into_body().collect().await.unwrap().to_bytes();
4638 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
4639 (status, json)
4640 }
4641
4642 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
4643 post_json_uri(app, "/v1/chat/completions", body).await.1
4644 }
4645
4646 #[tokio::test]
4651 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
4652 let app = test_app();
4653 let (status, body) = post_json_uri(
4654 &app,
4655 ferrox_api::routes::V1_CANCEL,
4656 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
4657 )
4658 .await;
4659 assert_eq!(status, StatusCode::NOT_FOUND);
4660 assert_eq!(body["cancelled"], serde_json::json!(false));
4661 assert_eq!(body["request_id"], "chatcmpl-never-issued");
4662 assert!(
4663 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
4664 "the verdict must carry a human reason: {body}"
4665 );
4666 }
4667
4668 #[tokio::test]
4672 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
4673 let state = Arc::new(test_state(
4674 test_model_full_byte_vocab(),
4675 ResponseCache::new(1000, Duration::from_secs(3600)),
4676 ));
4677 let app = test_app_with_state(Arc::clone(&state));
4678 let (token, _guard) = state.cancels.register("chatcmpl-live");
4679
4680 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
4681 assert_eq!(status, StatusCode::OK);
4682 assert_eq!(before["generating_now"], serde_json::json!(1));
4683
4684 let (status, body) = post_json_uri(
4685 &app,
4686 ferrox_api::routes::V1_CANCEL,
4687 serde_json::json!({ "request_id": "chatcmpl-live" }),
4688 )
4689 .await;
4690 assert_eq!(status, StatusCode::OK);
4691 assert_eq!(body["cancelled"], serde_json::json!(true));
4692 assert!(
4693 token.is_cancelled(),
4694 "the endpoint answered ok without setting the flag the decode loop reads"
4695 );
4696 }
4697
4698 #[tokio::test]
4699 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
4700 let app = test_app();
4701 let (status, tok) =
4702 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
4703 assert_eq!(status, StatusCode::OK);
4704 let tokens = tok["tokens"].as_array().unwrap();
4705 assert_eq!(tok["count"], tokens.len());
4706 assert!(!tokens.is_empty());
4707
4708 let (status, detok) = post_json_uri(
4709 &app,
4710 "/v1/detokenize",
4711 serde_json::json!({ "tokens": tokens }),
4712 )
4713 .await;
4714 assert_eq!(status, StatusCode::OK);
4715 assert_eq!(detok["text"], "Hi");
4716
4717 let (status, emb) = post_json_uri(
4718 &app,
4719 "/v1/embeddings",
4720 serde_json::json!({
4721 "input": "Hi",
4722 "embedding_type": "mean"
4723 }),
4724 )
4725 .await;
4726 assert_eq!(status, StatusCode::OK);
4727 let vec = emb["data"][0]["embedding"].as_array().unwrap();
4728 assert!(!vec.is_empty());
4729 assert!(vec.iter().all(|v| v.as_f64().is_some()));
4730 }
4731
4732 #[tokio::test]
4737 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
4738 use http_body_util::BodyExt;
4739 use tower::ServiceExt;
4740
4741 let fixture = concat!(
4742 "../ferrox-models/tests/fixtures/",
4743 "ferrox_real_moe_test.gguf"
4744 );
4745 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
4746 let decoder = Decoder::from_gguf_with_expert_cache(
4747 &fixture,
4748 ferrox_models::config::test_moe_fixture(),
4749 Some(1024 * 1024),
4750 )
4751 .expect("MoE fixture must load store-backed");
4752
4753 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
4757 .layers
4758 .iter()
4759 .map(|_| {
4760 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
4761 })
4762 .collect();
4763 decoder.forward_token(1, 0, &mut caches);
4764
4765 let model = Model::Gguf(GgufModel {
4766 decoder: Arc::new(decoder),
4767 tokenizer: Arc::new(ServerTokenizer::Byte),
4768 stop_tokens: StopTokens::default(),
4769 bos_id: None,
4770 is_synthetic: false,
4771 chat_template: chat_template::ChatTemplate::Plain,
4772 });
4773 let state = Arc::new(test_state(
4774 model,
4775 ResponseCache::new(16, Duration::from_secs(60)),
4776 ));
4777 let app = Router::new()
4778 .route("/metrics", axum::routing::get(metrics))
4779 .route("/v1/chat/completions", post(chat_completions))
4780 .with_state(state);
4781
4782 let fetch_metrics = |app: Router| async move {
4783 let resp = app
4784 .oneshot(
4785 axum::http::Request::builder()
4786 .method("GET")
4787 .uri("/metrics")
4788 .body(axum::body::Body::empty())
4789 .unwrap(),
4790 )
4791 .await
4792 .unwrap();
4793 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
4794 String::from_utf8(bytes.to_vec()).unwrap()
4795 };
4796
4797 let after = fetch_metrics(app.clone()).await;
4798 assert!(
4799 after.contains("ferrox_expert_cache_misses_total"),
4800 "streaming model must expose expert-cache metrics: {after}"
4801 );
4802 let misses: u64 = after
4803 .lines()
4804 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
4805 .and_then(|l| l.split_whitespace().nth(1))
4806 .and_then(|v| v.parse().ok())
4807 .expect("misses metric line must parse");
4808 assert!(
4809 misses > 0,
4810 "decode must have read experts through the store: {after}"
4811 );
4812 }
4813
4814 fn weather_tool() -> serde_json::Value {
4815 serde_json::json!({
4816 "type": "function",
4817 "function": {
4818 "name": "get_weather",
4819 "description": "Get the current weather for a location.",
4820 "parameters": {
4821 "type": "object",
4822 "properties": {"location": {"type": "string"}},
4823 "required": ["location"]
4824 }
4825 }
4826 })
4827 }
4828
4829 fn weather_tool_def() -> ToolDef {
4830 ToolDef {
4831 kind: "function".to_string(),
4832 function: ToolFunctionDef {
4833 name: "get_weather".to_string(),
4834 description: Some("Get the current weather for a location.".to_string()),
4835 parameters: Some(serde_json::json!({
4836 "type": "object",
4837 "properties": {"location": {"type": "string"}},
4838 "required": ["location"]
4839 })),
4840 },
4841 }
4842 }
4843
4844 #[test]
4845 fn tool_preamble_mentions_every_tool_name_and_description() {
4846 let preamble = tool_preamble(&[weather_tool_def()]);
4847 assert!(preamble.contains("get_weather"));
4848 assert!(preamble.contains("Get the current weather for a location."));
4849 assert!(preamble.contains("<tool_call>"));
4850 assert!(preamble.contains("</tool_call>"));
4851 }
4852
4853 #[test]
4854 fn extract_tool_call_parses_a_real_marker() {
4855 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
4856 let (name, arguments) = extract_tool_call(text).expect("must find the marker");
4857 assert_eq!(name, "get_weather");
4858 let parsed: serde_json::Value = serde_json::from_str(&arguments).unwrap();
4859 assert_eq!(parsed["location"], "Paris");
4860 }
4861
4862 #[test]
4863 fn extract_tool_call_returns_none_when_no_marker_present() {
4864 assert_eq!(
4865 extract_tool_call("just a plain answer, no markers here"),
4866 None
4867 );
4868 }
4869
4870 #[test]
4871 fn extract_tool_call_returns_none_on_malformed_json_inside_the_marker() {
4872 assert_eq!(
4873 extract_tool_call("<tool_call>not valid json at all</tool_call>"),
4874 None
4875 );
4876 }
4877
4878 #[test]
4879 fn extract_tool_call_defaults_to_empty_arguments_when_the_field_is_absent() {
4880 let (name, arguments) =
4881 extract_tool_call("<tool_call>{\"name\": \"ping\"}</tool_call>").unwrap();
4882 assert_eq!(name, "ping");
4883 assert_eq!(arguments, "{}");
4884 }
4885
4886 #[test]
4887 fn build_response_message_promotes_a_real_tool_call_to_tool_calls() {
4888 let (message, finish) = build_response_message(
4889 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Rome\"}}</tool_call>"
4890 .to_string(),
4891 true,
4892 "stop",
4893 );
4894 assert_eq!(finish, "tool_calls");
4895 assert!(message.content.is_none());
4896 let calls = message.tool_calls.expect("must carry a tool call");
4897 assert_eq!(calls.len(), 1);
4898 assert_eq!(calls[0].function.name, "get_weather");
4899 }
4900
4901 #[test]
4902 fn build_response_message_falls_back_to_plain_text_when_tools_are_inactive() {
4903 let (message, finish) = build_response_message(
4904 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
4905 false,
4906 "stop",
4907 );
4908 assert_eq!(finish, "stop");
4909 assert!(message.tool_calls.is_none());
4910 assert!(message.content.is_some());
4911 }
4912
4913 #[test]
4914 fn build_response_message_falls_back_to_plain_text_when_no_marker_is_present() {
4915 let (message, finish) = build_response_message("just an answer".to_string(), true, "stop");
4916 assert_eq!(finish, "stop");
4917 assert!(message.tool_calls.is_none());
4918 assert_eq!(message.content.as_deref(), Some("just an answer"));
4919 }
4920
4921 #[tokio::test]
4929 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
4930 let app = test_app();
4931 let body = serde_json::json!({
4932 "model": "m",
4933 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
4934 "max_tokens": 4,
4935 "temperature": 0,
4936 });
4937 let resp = post_json(&app, body).await;
4938 let message = &resp["choices"][0]["message"];
4939 assert!(message["content"].is_string());
4940 assert!(message.get("tool_calls").is_none());
4941 assert_eq!(resp["choices"][0]["finish_reason"], "length");
4942 assert_eq!(resp["usage"]["completion_tokens"], 4);
4943 assert_eq!(
4944 resp["usage"]["total_tokens"],
4945 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
4946 );
4947 }
4948
4949 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
4950 use http_body_util::BodyExt;
4951 use tower::ServiceExt;
4952
4953 let response = app
4954 .clone()
4955 .oneshot(
4956 axum::http::Request::builder()
4957 .method("GET")
4958 .uri(uri)
4959 .body(axum::body::Body::empty())
4960 .unwrap(),
4961 )
4962 .await
4963 .unwrap();
4964 let status = response.status();
4965 let bytes = response.into_body().collect().await.unwrap().to_bytes();
4966 (status, serde_json::from_slice(&bytes).unwrap())
4967 }
4968
4969 #[tokio::test]
4970 async fn health_answers_a_capability_handshake_not_a_boolean() {
4971 let app = test_app();
4972 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
4973 assert_eq!(status, StatusCode::OK);
4974
4975 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
4976 assert_eq!(health.state, ferrox_api::HealthState::Ready);
4977 assert!(health.pid > 0);
4978 assert!(health.server_time_unix_ms > 0);
4979 assert_eq!(health.last_request_age_seconds, None);
4982
4983 for id in [
4986 ferrox_api::health::capability::CPU,
4987 ferrox_api::health::capability::METAL,
4988 ferrox_api::health::capability::CUDA,
4989 ferrox_api::health::capability::REAL_WEIGHTS,
4990 ferrox_api::health::capability::CONTINUOUS_BATCHING,
4991 ] {
4992 let cap = health
4993 .capability(id)
4994 .unwrap_or_else(|| panic!("{id} missing"));
4995 assert!(!cap.reason.is_empty(), "{cap:?}");
4996 assert!(!cap.detail.is_empty(), "{cap:?}");
4997 }
4998 let weights = health
5002 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
5003 .unwrap();
5004 assert!(!weights.available);
5005 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
5006 assert!(health.model.as_ref().unwrap().synthetic_weights);
5007 }
5008
5009 #[tokio::test]
5010 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
5011 let app = test_app();
5012 let _ = post_json(
5013 &app,
5014 serde_json::json!({
5015 "model": "m",
5016 "messages": [{"role": "user", "content": "\u{1}"}],
5017 "max_tokens": 1,
5018 "temperature": 0,
5019 }),
5020 )
5021 .await;
5022 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5023 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
5024 let age = health
5025 .last_request_age_seconds
5026 .expect("a served request is evidence of liveness");
5027 assert!((0.0..5.0).contains(&age), "implausible age {age}");
5028 }
5029
5030 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
5032 use http_body_util::BodyExt;
5033 use tower::ServiceExt;
5034
5035 let response = app
5036 .clone()
5037 .oneshot(
5038 axum::http::Request::builder()
5039 .method("POST")
5040 .uri("/v1/chat/completions")
5041 .header("content-type", "application/json")
5042 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5043 .unwrap(),
5044 )
5045 .await
5046 .unwrap();
5047 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5048 String::from_utf8(bytes.to_vec())
5049 .unwrap()
5050 .lines()
5051 .filter_map(|line| line.strip_prefix("data: "))
5052 .filter(|payload| *payload != "[DONE]")
5053 .map(|payload| serde_json::from_str(payload).unwrap())
5054 .collect()
5055 }
5056
5057 #[tokio::test]
5058 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
5059 let app = test_app();
5060 let chunks = post_sse_chunks(
5061 &app,
5062 serde_json::json!({
5063 "model": "m",
5064 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5065 "max_tokens": 4,
5066 "temperature": 0,
5067 "stream": true,
5068 }),
5069 )
5070 .await;
5071
5072 assert!(!chunks.is_empty());
5073 let request_id = chunks[0]["request_id"]
5074 .as_str()
5075 .expect("the first chunk names the request")
5076 .to_string();
5077 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
5078 for (i, chunk) in chunks.iter().enumerate().skip(1) {
5081 assert!(
5082 chunk.get("request_id").is_none(),
5083 "chunk {i} repeats request_id"
5084 );
5085 }
5086 for chunk in &chunks {
5089 assert_eq!(chunk["id"], serde_json::json!(request_id));
5090 }
5091
5092 let other = post_sse_chunks(
5093 &app,
5094 serde_json::json!({
5095 "model": "m",
5096 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5097 "max_tokens": 4,
5098 "temperature": 0,
5099 "stream": true,
5100 }),
5101 )
5102 .await;
5103 assert_ne!(
5104 other[0]["request_id"].as_str().unwrap(),
5105 request_id,
5106 "two concurrent chats must not share an id"
5107 );
5108 }
5109
5110 #[tokio::test]
5111 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
5112 let app = test_app();
5113 let resp = post_json(
5114 &app,
5115 serde_json::json!({
5116 "model": "m",
5117 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5118 "max_tokens": 2,
5119 "temperature": 0,
5120 }),
5121 )
5122 .await;
5123 assert_eq!(resp["id"], resp["request_id"]);
5124 assert!(resp["request_id"]
5125 .as_str()
5126 .unwrap()
5127 .starts_with("chatcmpl-"));
5128 }
5129
5130 #[tokio::test]
5133 async fn usage_carries_separate_prefill_and_decode_timings() {
5134 let app = test_app();
5135 let resp = post_json(
5136 &app,
5137 serde_json::json!({
5138 "model": "m",
5139 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5140 "max_tokens": 4,
5141 "temperature": 0,
5142 }),
5143 )
5144 .await;
5145 let usage = &resp["usage"];
5146 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
5147 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
5148 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
5149 assert!(usage["predicted_per_second"].is_number(), "{usage}");
5150 assert!(usage.get("cached_tokens").is_none(), "{usage}");
5152 }
5153
5154 #[tokio::test]
5162 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
5163 let app = test_app();
5164 let body = serde_json::json!({
5165 "model": "m",
5166 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5167 "max_tokens": 4,
5168 "temperature": 0,
5169 "tools": [weather_tool()],
5170 });
5171 let resp = post_json(&app, body).await;
5172 let message = &resp["choices"][0]["message"];
5173 assert!(
5174 message["content"].is_string(),
5175 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
5176 );
5177 assert!(message.get("tool_calls").is_none());
5178 assert_eq!(resp["choices"][0]["finish_reason"], "length");
5182 }
5183
5184 #[tokio::test]
5188 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
5189 let app = test_app();
5190 let body = serde_json::json!({
5191 "model": "m",
5192 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
5193 "max_tokens": 3,
5194 "temperature": 0,
5195 });
5196 let first = post_json(&app, body.clone()).await;
5197 assert_eq!(first["ferrox_cache"], "miss");
5198 let second = post_json(&app, body).await;
5199 assert_eq!(second["ferrox_cache"], "hit");
5200 assert_eq!(
5201 first["choices"][0]["message"]["content"],
5202 second["choices"][0]["message"]["content"]
5203 );
5204 assert_eq!(
5205 first["choices"][0]["finish_reason"],
5206 second["choices"][0]["finish_reason"]
5207 );
5208 assert_eq!(first["usage"], second["usage"]);
5209 assert_eq!(second["usage"]["completion_tokens"], 3);
5210 }
5211
5212 #[tokio::test]
5218 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
5219 let session_app = test_app();
5220 let manual_app = test_app();
5221
5222 let turn1 = post_json(
5224 &session_app,
5225 serde_json::json!({
5226 "model": "m",
5227 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5228 "session_id": "s1",
5229 "max_tokens": 5,
5230 "temperature": 0,
5231 }),
5232 )
5233 .await;
5234 let reply1 = turn1["choices"][0]["message"]["content"]
5235 .as_str()
5236 .unwrap()
5237 .to_string();
5238
5239 let manual_turn1 = post_json(
5244 &manual_app,
5245 serde_json::json!({
5246 "model": "m",
5247 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5248 "max_tokens": 5,
5249 "temperature": 0,
5250 }),
5251 )
5252 .await;
5253 assert_eq!(
5254 manual_turn1["choices"][0]["message"]["content"]
5255 .as_str()
5256 .unwrap(),
5257 reply1
5258 );
5259
5260 let turn2 = post_json(
5262 &session_app,
5263 serde_json::json!({
5264 "model": "m",
5265 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
5266 "session_id": "s1",
5267 "max_tokens": 5,
5268 "temperature": 0,
5269 }),
5270 )
5271 .await;
5272 let reply2 = turn2["choices"][0]["message"]["content"]
5273 .as_str()
5274 .unwrap()
5275 .to_string();
5276
5277 let manual_turn2 = post_json(
5281 &manual_app,
5282 serde_json::json!({
5283 "model": "m",
5284 "messages": [
5285 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
5286 {"role": "assistant", "content": reply1},
5287 {"role": "user", "content": "\u{4}\u{5}"},
5288 ],
5289 "max_tokens": 5,
5290 "temperature": 0,
5291 }),
5292 )
5293 .await;
5294 assert_eq!(
5295 manual_turn2["choices"][0]["message"]["content"]
5296 .as_str()
5297 .unwrap(),
5298 reply2,
5299 "resuming a session must produce identical output to manually resending the full history"
5300 );
5301 }
5302
5303 #[test]
5306 fn lock_cache_recovers_from_a_poisoned_mutex() {
5307 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
5308
5309 let poison_cache = Arc::clone(&cache);
5310 let _ = std::thread::spawn(move || {
5311 let _guard = poison_cache.lock().unwrap();
5312 panic!("simulated panic while holding the lock");
5313 })
5314 .join();
5315
5316 let recovered = lock_cache(&cache);
5318 assert_eq!(recovered.stats().entries, 0);
5319 }
5320
5321 #[test]
5322 fn is_cacheable_true_for_greedy_or_seeded_requests() {
5323 let mut req_body = serde_json::json!({
5324 "model": "m",
5325 "messages": [{"role": "user", "content": "hi"}],
5326 });
5327 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
5328 assert!(
5329 req.is_cacheable(),
5330 "default (temperature 0) must be cacheable"
5331 );
5332
5333 req_body["temperature"] = serde_json::json!(0.8);
5334 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
5335 assert!(
5336 !req.is_cacheable(),
5337 "unseeded sampling must never be cacheable"
5338 );
5339
5340 req_body["seed"] = serde_json::json!(42);
5341 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
5342 assert!(
5343 req.is_cacheable(),
5344 "sampling with an explicit seed is deterministic and must be cacheable"
5345 );
5346 }
5347
5348 #[test]
5349 fn stop_param_accepts_both_single_string_and_array() {
5350 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
5351 "model": "m",
5352 "messages": [{"role": "user", "content": "hi"}],
5353 "stop": "END",
5354 }))
5355 .unwrap();
5356 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
5357
5358 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
5359 "model": "m",
5360 "messages": [{"role": "user", "content": "hi"}],
5361 "stop": ["A", "B"],
5362 }))
5363 .unwrap();
5364 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
5365 }
5366
5367 #[test]
5368 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
5369 let model = test_model();
5370 let result = run_generation(&model, "hello", &greedy_params(4), None, None, None, None);
5371 assert!(matches!(
5372 result,
5373 Err(generate::DecodeError::TokenOutOfVocab { .. })
5374 ));
5375 }
5376
5377 #[test]
5381 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
5382 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5384 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
5385
5386 let holder_pool = Arc::clone(&pool);
5387 let holder = std::thread::spawn(move || {
5388 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
5389 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
5391 drop(held);
5392 });
5393 std::thread::sleep(Duration::from_millis(15));
5394
5395 let config = generate::KvPoolConfig {
5396 pool,
5397 queue_wait: Duration::ZERO,
5398 };
5399 let result = run_generation(
5400 &model,
5401 &prompt,
5402 &greedy_params(4),
5403 Some(&config),
5404 None,
5405 None,
5406 None,
5407 );
5408 assert!(matches!(
5409 result,
5410 Err(generate::DecodeError::KvPoolExhausted)
5411 ));
5412
5413 let (status, _body) = decode_error_response(result.unwrap_err());
5414 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5415 holder.join().unwrap();
5416 }
5417
5418 #[test]
5426 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
5427 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5429 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
5431 let config = generate::KvPoolConfig {
5432 pool,
5433 queue_wait: Duration::ZERO,
5434 };
5435
5436 let result = run_generation(
5437 &model,
5438 &prompt,
5439 &greedy_params(4),
5440 Some(&config),
5441 None,
5442 None,
5443 None,
5444 );
5445 let err = result.expect_err("one block cannot hold two layers' caches");
5446 assert!(
5447 matches!(
5448 &err,
5449 generate::DecodeError::KvBudgetExceeded { binding, .. }
5450 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
5451 ),
5452 "expected an immovable device-memory refusal, got {err:?}"
5453 );
5454 let (status, _body) = decode_error_response(err);
5455 assert_eq!(status, StatusCode::BAD_REQUEST);
5456 }
5457
5458 #[test]
5464 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
5465 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
5466 queued: 512,
5467 cap: 512,
5468 });
5469 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5470 assert_eq!(body["error"]["retry_after_seconds"], 1);
5471 let message = body["error"]["message"].as_str().expect("message");
5472 assert!(message.contains("512"), "{message}");
5473 }
5474
5475 #[test]
5476 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
5477 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
5478 token: 99,
5479 vocab_size: 32,
5480 });
5481 assert!(
5482 body["error"]["retry_after_seconds"].is_null(),
5483 "retrying a prompt this model cannot tokenize never helps"
5484 );
5485 }
5486
5487 #[test]
5488 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
5489 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
5490 token: 99,
5491 vocab_size: 32,
5492 });
5493 assert_eq!(status, StatusCode::BAD_REQUEST);
5494 }
5495
5496 #[test]
5497 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
5498 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5500 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
5501 let config = generate::KvPoolConfig {
5502 pool: pool.clone(),
5503 queue_wait: Duration::ZERO,
5504 };
5505
5506 let (_, finish, _usage) = run_generation(
5507 &model,
5508 &prompt,
5509 &greedy_params(4),
5510 Some(&config),
5511 None,
5512 None,
5513 None,
5514 )
5515 .unwrap();
5516 assert_eq!(finish, FinishReason::Length);
5517 assert_eq!(
5518 pool.lock().unwrap().free_blocks(),
5519 2,
5520 "a completed request must return its blocks to the pool"
5521 );
5522 }
5523
5524 #[tokio::test]
5529 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
5530 let model = Arc::new(test_model());
5531 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
5532
5533 let mut handles = Vec::new();
5534 for _ in 0..8 {
5535 let model = Arc::clone(&model);
5536 let prompt = prompt.clone();
5537 handles.push(tokio::task::spawn_blocking(move || {
5538 run_generation(&model, &prompt, &greedy_params(6), None, None, None, None).unwrap()
5539 }));
5540 }
5541
5542 let mut results = Vec::new();
5543 for h in handles {
5544 results.push(h.await.unwrap());
5545 }
5546 for r in &results[1..] {
5551 assert_eq!(r.0, results[0].0, "decoded chunks must match");
5552 assert_eq!(r.1, results[0].1, "finish reason must match");
5553 assert_eq!(
5554 r.2.prompt_tokens, results[0].2.prompt_tokens,
5555 "prompt token count must match"
5556 );
5557 assert_eq!(
5558 r.2.completion_tokens, results[0].2.completion_tokens,
5559 "completion token count must match"
5560 );
5561 }
5562 }
5563
5564 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
5571 let mut header_entries = Vec::new();
5572 let mut data = Vec::new();
5573 for (name, shape, values) in tensors {
5574 let start = data.len();
5575 for v in values {
5576 data.extend_from_slice(&v.to_le_bytes());
5577 }
5578 let end = data.len();
5579 let shape_str = shape
5580 .iter()
5581 .map(|d| d.to_string())
5582 .collect::<Vec<_>>()
5583 .join(",");
5584 header_entries.push(format!(
5585 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
5586 ));
5587 }
5588 let header = format!("{{{}}}", header_entries.join(","));
5589 let header_bytes = header.as_bytes();
5590 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
5591 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
5592 out.extend_from_slice(header_bytes);
5593 out.extend_from_slice(&data);
5594 out
5595 }
5596
5597 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
5608 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
5609 use ferrox_models::kimi_loader::KimiRealHparams;
5610 use ferrox_moe::{GatingFunction, MoeLayerConfig};
5611
5612 let hidden_dim = 8;
5613 let kda_num_heads = 2;
5614 let kda_head_dim = 3;
5615 let kda_proj = kda_num_heads * kda_head_dim;
5616 let conv_kernel = 4;
5617 let dense_intermediate = 5;
5618 let vocab_size = 256;
5622 let mla_num_heads = 1;
5623 let mla_q_lora_rank = 2;
5624 let mla_kv_lora_rank = 2;
5625 let mla_qk_nope_head_dim = 2;
5626 let mla_qk_rope_head_dim = 2;
5627 let mla_v_head_dim = 2;
5628
5629 let model_cfg = ferrox_models::ModelConfig {
5630 name: "synthetic-kimi-server-test",
5631 n_layers: 1,
5632 hidden_dim,
5633 n_heads: 1,
5634 n_kv_heads: 1,
5635 head_dim: 4,
5636 vocab_size,
5637 rope_theta: 10000.0,
5638 rms_norm_eps: 1e-5,
5639 sliding_window: None,
5640 moe: MoeLayerConfig {
5641 expert_weights_scale: 1.0,
5642 n_experts: 1,
5643 n_experts_active: 1,
5644 n_shared_experts: 0,
5645 hidden_dim,
5646 expert_ffn_dim: 4,
5647 gating: GatingFunction::Sigmoid,
5648 norm_topk_prob: true,
5649 expert_group_count: None,
5650 expert_group_used_count: None,
5651 },
5652 n_dense_leading_layers: 1,
5657 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
5658 kda_layers: vec![1],
5659 full_attn_layers: vec![],
5660 mla: MlaConfig {
5661 num_heads: mla_num_heads,
5662 q_lora_rank: mla_q_lora_rank,
5663 kv_lora_rank: mla_kv_lora_rank,
5664 qk_nope_head_dim: mla_qk_nope_head_dim,
5665 qk_rope_head_dim: mla_qk_rope_head_dim,
5666 v_head_dim: mla_v_head_dim,
5667 use_output_gate: true,
5668 rope: None,
5669 },
5670 kda: KdaConfig {
5671 num_heads: kda_num_heads,
5672 head_dim: kda_head_dim,
5673 short_conv_kernel_size: conv_kernel,
5674 gate_lower_bound: -5.0,
5675 use_full_rank_gate: true,
5676 },
5677 }),
5678 rope_freqs: None,
5679 rope_attn_factor: 1.0,
5680 rope_dim: None,
5681 rope_freqs_long: None,
5682 rope_freqs_short: None,
5683 rope_orig_ctx: None,
5684 rope_layout: ferrox_models::config::RopeLayout::Neox,
5685 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
5686 swa_pattern: None,
5687 attn_logit_softcap: None,
5688 final_logit_softcap: None,
5689 embedding_scale: None,
5690 attention_scale: None,
5691 rope_theta_swa: None,
5692 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
5693 best_effort_fields: &["synthetic test config, not a real preset"],
5694 };
5695 let hp = KimiRealHparams {
5696 hidden_dim,
5697 kda_num_heads,
5698 kda_head_dim,
5699 mla_num_heads,
5700 mla_q_lora_rank,
5701 mla_kv_lora_rank,
5702 mla_qk_nope_head_dim,
5703 mla_qk_rope_head_dim,
5704 mla_v_head_dim,
5705 dense_intermediate_dim: dense_intermediate,
5706 moe_hidden_dim: hidden_dim,
5707 moe_intermediate_dim: 4,
5708 n_experts: 1,
5709 num_shared_experts: 0,
5710 };
5711
5712 let prefix = "language_model.model.layers.0";
5716 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
5717 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
5718 name: String,
5719 shape: Vec<usize>,
5720 n: usize| {
5721 tensors.push((name, shape, vec![0.01f32; n]));
5722 };
5723 push(
5724 &mut tensors,
5725 format!("{prefix}.input_layernorm.weight"),
5726 vec![hidden_dim],
5727 hidden_dim,
5728 );
5729 push(
5730 &mut tensors,
5731 format!("{prefix}.post_attention_layernorm.weight"),
5732 vec![hidden_dim],
5733 hidden_dim,
5734 );
5735 push(
5736 &mut tensors,
5737 format!("{prefix}.self_attention_res_norm.weight"),
5738 vec![hidden_dim],
5739 hidden_dim,
5740 );
5741 push(
5742 &mut tensors,
5743 format!("{prefix}.self_attention_res_proj.weight"),
5744 vec![1, hidden_dim],
5745 hidden_dim,
5746 );
5747 push(
5748 &mut tensors,
5749 format!("{prefix}.mlp_res_norm.weight"),
5750 vec![hidden_dim],
5751 hidden_dim,
5752 );
5753 push(
5754 &mut tensors,
5755 format!("{prefix}.mlp_res_proj.weight"),
5756 vec![1, hidden_dim],
5757 hidden_dim,
5758 );
5759 push(
5760 &mut tensors,
5761 format!("{prefix}.self_attn.q_proj.weight"),
5762 vec![kda_proj, hidden_dim],
5763 kda_proj * hidden_dim,
5764 );
5765 push(
5766 &mut tensors,
5767 format!("{prefix}.self_attn.k_proj.weight"),
5768 vec![kda_proj, hidden_dim],
5769 kda_proj * hidden_dim,
5770 );
5771 push(
5772 &mut tensors,
5773 format!("{prefix}.self_attn.v_proj.weight"),
5774 vec![kda_proj, hidden_dim],
5775 kda_proj * hidden_dim,
5776 );
5777 push(
5778 &mut tensors,
5779 format!("{prefix}.self_attn.q_conv1d.weight"),
5780 vec![kda_proj, 1, conv_kernel],
5781 kda_proj * conv_kernel,
5782 );
5783 push(
5784 &mut tensors,
5785 format!("{prefix}.self_attn.k_conv1d.weight"),
5786 vec![kda_proj, 1, conv_kernel],
5787 kda_proj * conv_kernel,
5788 );
5789 push(
5790 &mut tensors,
5791 format!("{prefix}.self_attn.v_conv1d.weight"),
5792 vec![kda_proj, 1, conv_kernel],
5793 kda_proj * conv_kernel,
5794 );
5795 push(
5796 &mut tensors,
5797 format!("{prefix}.self_attn.A_log"),
5798 vec![kda_num_heads],
5799 kda_num_heads,
5800 );
5801 push(
5802 &mut tensors,
5803 format!("{prefix}.self_attn.f_a_proj.weight"),
5804 vec![kda_head_dim, hidden_dim],
5805 kda_head_dim * hidden_dim,
5806 );
5807 push(
5808 &mut tensors,
5809 format!("{prefix}.self_attn.f_b_proj.weight"),
5810 vec![kda_proj, kda_head_dim],
5811 kda_proj * kda_head_dim,
5812 );
5813 push(
5814 &mut tensors,
5815 format!("{prefix}.self_attn.dt_bias"),
5816 vec![kda_proj],
5817 kda_proj,
5818 );
5819 push(
5820 &mut tensors,
5821 format!("{prefix}.self_attn.b_proj.weight"),
5822 vec![kda_num_heads, hidden_dim],
5823 kda_num_heads * hidden_dim,
5824 );
5825 push(
5826 &mut tensors,
5827 format!("{prefix}.self_attn.g_proj.weight"),
5828 vec![kda_proj, hidden_dim],
5829 kda_proj * hidden_dim,
5830 );
5831 push(
5832 &mut tensors,
5833 format!("{prefix}.self_attn.o_norm.weight"),
5834 vec![kda_head_dim],
5835 kda_head_dim,
5836 );
5837 push(
5838 &mut tensors,
5839 format!("{prefix}.self_attn.o_proj.weight"),
5840 vec![hidden_dim, kda_proj],
5841 hidden_dim * kda_proj,
5842 );
5843 push(
5844 &mut tensors,
5845 format!("{prefix}.mlp.gate_proj.weight"),
5846 vec![dense_intermediate, hidden_dim],
5847 dense_intermediate * hidden_dim,
5848 );
5849 push(
5850 &mut tensors,
5851 format!("{prefix}.mlp.up_proj.weight"),
5852 vec![dense_intermediate, hidden_dim],
5853 dense_intermediate * hidden_dim,
5854 );
5855 push(
5856 &mut tensors,
5857 format!("{prefix}.mlp.down_proj.weight"),
5858 vec![hidden_dim, dense_intermediate],
5859 hidden_dim * dense_intermediate,
5860 );
5861 push(
5862 &mut tensors,
5863 "language_model.model.embed_tokens.weight".to_string(),
5864 vec![vocab_size, hidden_dim],
5865 vocab_size * hidden_dim,
5866 );
5867 push(
5868 &mut tensors,
5869 "language_model.lm_head.weight".to_string(),
5870 vec![vocab_size, hidden_dim],
5871 vocab_size * hidden_dim,
5872 );
5873 push(
5874 &mut tensors,
5875 "language_model.model.norm.weight".to_string(),
5876 vec![hidden_dim],
5877 hidden_dim,
5878 );
5879 push(
5880 &mut tensors,
5881 "language_model.model.output_attn_res_norm.weight".to_string(),
5882 vec![hidden_dim],
5883 hidden_dim,
5884 );
5885 push(
5886 &mut tensors,
5887 "language_model.model.output_attn_res_proj.weight".to_string(),
5888 vec![1, hidden_dim],
5889 hidden_dim,
5890 );
5891
5892 let dir = std::env::temp_dir().join(format!(
5893 "ferrox_server_kimi_e2e_test_{}_{}",
5894 std::process::id(),
5895 vocab_size
5896 ));
5897 std::fs::create_dir_all(&dir).unwrap();
5898 let shard_bytes = write_safetensors_shard(&tensors);
5899 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
5900 let map_entries: Vec<String> = tensors
5901 .iter()
5902 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
5903 .collect();
5904 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
5905 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
5906
5907 use base64::Engine;
5911 let vocab_lines: Vec<String> = (0..vocab_size as u32)
5912 .map(|b| {
5913 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
5914 format!("{b64} {b}")
5915 })
5916 .collect();
5917 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
5918
5919 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
5920 .expect("must load the synthetic Kimi checkpoint end to end");
5921 std::fs::remove_dir_all(&dir).ok();
5922 loaded
5923 }
5924
5925 #[test]
5933 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
5934 let loaded = build_synthetic_kimi_loaded();
5935 let state = build_app_state(
5936 model::LoadedModel::Kimi(loaded),
5937 None,
5938 None,
5939 false,
5940 None,
5941 Arc::new(health::Detection::ready(health::probe_backends())),
5942 );
5943 let active = state.active().expect("a freshly built state has a model");
5944 assert_eq!(active.model.tokenizer_kind(), "kimi-tiktoken-bpe");
5945 assert!(!active.model.is_synthetic());
5946
5947 let (_chunks, finish, _usage) = run_generation(
5948 &active.model,
5949 "hi",
5950 &greedy_params(5),
5951 None,
5952 None,
5953 None,
5954 None,
5955 )
5956 .expect("a real Kimi checkpoint must generate without error");
5957 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5958 }
5959
5960 #[test]
5970 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
5971 let loaded = build_synthetic_kimi_loaded();
5972 let state = build_app_state(
5973 model::LoadedModel::Kimi(loaded),
5974 None,
5975 None,
5976 false,
5977 None,
5978 Arc::new(health::Detection::ready(health::probe_backends())),
5979 );
5980
5981 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
5982 let kv_pool_config = generate::KvPoolConfig {
5983 pool: pool.clone(),
5984 queue_wait: Duration::ZERO,
5985 };
5986 let pc = Mutex::new(PrefixCache::new(4));
5987
5988 run_generation(
5989 &state
5990 .active()
5991 .expect("a freshly built state has a model")
5992 .model,
5993 "hi",
5994 &greedy_params(5),
5995 Some(&kv_pool_config),
5996 Some(&pc),
5997 None,
5998 None,
5999 )
6000 .expect("a real Kimi checkpoint must generate without error");
6001
6002 assert_eq!(
6003 pool.lock().unwrap().free_blocks(),
6004 4,
6005 "the KV pool must be completely untouched by a Kimi request"
6006 );
6007 let stats = pc.lock().unwrap().stats();
6008 assert_eq!(
6009 stats.hits + stats.misses,
6010 0,
6011 "the prefix cache must never be consulted for a Kimi request"
6012 );
6013 }
6014}