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