1mod admin;
36mod anthropic;
37mod attribution;
38mod budget;
39mod cache_admin;
40mod cancel;
41mod chat_template;
42mod cli;
43mod completion;
44mod conversations;
45mod decode_task;
46mod embeddings;
47mod generate;
48mod grammar_request;
49mod health;
50mod journal;
51mod json_mode;
52mod limits;
53mod loaded;
54mod mcp;
55mod model;
56mod openai_extra;
57mod output;
58mod policy;
59mod prefill_batch;
60mod reasoning_tokens;
61mod rerank;
62mod response_cache;
63pub(crate) mod responses;
64mod resume;
65mod sample_step;
66mod sampling_knobs;
67mod security;
68mod serving;
69mod session;
70mod slots;
71mod sse;
72mod stats;
73mod stop;
74mod stream_events;
75mod tasks;
76mod tool_grammar;
77mod unsupported_sampling;
78mod utf8_stream;
79
80use std::cell::RefCell;
81use std::convert::Infallible;
82use std::net::SocketAddr;
83use std::path::PathBuf;
84use std::rc::Rc;
85use std::sync::{Arc, Mutex, MutexGuard};
86use std::time::Duration;
87
88use axum::{
89 extract::State,
90 http::StatusCode,
91 response::sse::{Event, Sse},
92 response::{IntoResponse, Response},
93 routing::{get, post},
94 Json, Router,
95};
96use serde::{Deserialize, Serialize};
97
98use cli::{apply_cli_overrides, print_available_devices};
99pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
100
101use ferrox_core::cache::KvBlockPool;
102use ferrox_models::kimi_tokenizer::KimiTokenizer;
103use ferrox_models::sampling::SamplingParams;
104use ferrox_models::tokenizer::{SpecialTokens, StopTokens};
105use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
106use generate::{FinishReason, GenerationParams};
107pub(crate) use loaded::{ActiveModel, Loaded};
108use model::ServerTokenizer;
109use rerank::encoder_endpoints;
110use response_cache::{CacheKey, ResponseCache};
111use sampling_knobs::SamplingKnobs;
112
113#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
122 Gguf(GgufModel),
123 Kimi(KimiModel),
124 Mla(MlaModel),
125 Gemma4(Gemma4Model),
126 Glm52(Glm52Model),
127}
128
129pub(crate) struct GgufModel {
130 decoder: Arc<Decoder>,
131 tokenizer: Arc<ServerTokenizer>,
132 stop_tokens: StopTokens,
133 bos_id: Option<usize>,
134 is_synthetic: bool,
135 chat_template: chat_template::PromptTemplate,
136}
137
138pub(crate) struct KimiModel {
139 engine: KimiEngine,
140 tokenizer: KimiTokenizer,
141 stop_tokens: StopTokens,
142 chat_template: chat_template::PromptTemplate,
143}
144
145pub(crate) struct MlaModel {
146 engine: MlaEngine,
147 tokenizer: ServerTokenizer,
148 stop_tokens: StopTokens,
149 bos_id: Option<usize>,
150 name: String,
151 chat_template: chat_template::PromptTemplate,
152}
153
154pub(crate) struct Gemma4Model {
155 engine: Gemma4Engine,
156 tokenizer: ServerTokenizer,
157 stop_tokens: StopTokens,
158 bos_id: Option<usize>,
159 name: String,
160 chat_template: chat_template::PromptTemplate,
161}
162
163pub(crate) struct Glm52Model {
164 engine: ferrox_models::Glm52Engine,
165 tokenizer: ServerTokenizer,
166 stop_tokens: StopTokens,
167 bos_id: Option<usize>,
168 name: String,
169 chat_template: chat_template::PromptTemplate,
170}
171
172impl Model {
173 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
174 match self {
175 Model::Gguf(m) => m.chat_template.clone(),
176 Model::Kimi(m) => m.chat_template.clone(),
177 Model::Mla(m) => m.chat_template.clone(),
178 Model::Gemma4(m) => m.chat_template.clone(),
179 Model::Glm52(m) => m.chat_template.clone(),
180 }
181 }
182
183 fn is_synthetic(&self) -> bool {
188 match self {
189 Model::Gguf(m) => m.is_synthetic,
190 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
191 }
192 }
193
194 fn tokenizer_kind(&self) -> &'static str {
195 match self {
196 Model::Gguf(m) => m.tokenizer.kind(),
197 Model::Kimi(_) => "kimi-tiktoken-bpe",
198 Model::Mla(m) => m.tokenizer.kind(),
199 Model::Gemma4(m) => m.tokenizer.kind(),
200 Model::Glm52(m) => m.tokenizer.kind(),
201 }
202 }
203
204 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
208 match self {
209 Model::Gguf(m) => m.decoder.expert_store_stats(),
210 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
211 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
212 }
213 }
214
215 pub(crate) fn name(&self) -> &str {
216 match self {
217 Model::Gguf(m) => m.decoder.config.name,
218 Model::Kimi(_) => "kimi-k3",
219 Model::Mla(m) => m.name.as_str(),
220 Model::Gemma4(m) => m.name.as_str(),
221 Model::Glm52(m) => m.name.as_str(),
222 }
223 }
224
225 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
252 match self {
253 Model::Gguf(m) => m.tokenizer.encode(text, specials),
254 Model::Kimi(m) => m
255 .tokenizer
256 .encode(text, specials)
257 .into_iter()
258 .map(|id| id as usize)
259 .collect(),
260 Model::Mla(m) => m.tokenizer.encode(text, specials),
261 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
262 Model::Glm52(m) => m.tokenizer.encode(text, specials),
263 }
264 }
265
266 pub(crate) fn bos_id(&self) -> Option<usize> {
275 match self {
276 Model::Gguf(m) => m.bos_id,
277 Model::Kimi(_) => None,
278 Model::Mla(m) => m.bos_id,
279 Model::Gemma4(m) => m.bos_id,
280 Model::Glm52(m) => m.bos_id,
281 }
282 }
283
284 pub(crate) fn decode(&self, ids: &[usize]) -> String {
285 match self {
286 Model::Gguf(m) => m.tokenizer.decode(ids),
287 Model::Kimi(m) => {
288 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
289 m.tokenizer.decode(&ids32)
290 }
291 Model::Mla(m) => m.tokenizer.decode(ids),
292 Model::Gemma4(m) => m.tokenizer.decode(ids),
293 Model::Glm52(m) => m.tokenizer.decode(ids),
294 }
295 }
296
297 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
300 match self {
301 Model::Gguf(m) => {
302 let mut caches: Vec<_> = (0..m.decoder.layers.len())
303 .map(|_| {
304 ferrox_core::cache::KvCache::new(
305 m.decoder.config.n_kv_heads,
306 m.decoder.config.head_dim,
307 )
308 })
309 .collect();
310 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
311 }
312 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
313 }
314 }
315
316 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
322 match self {
323 Model::Gguf(m) => Some(&m.decoder),
324 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
325 }
326 }
327
328 pub(crate) fn vocab_size(&self) -> Option<usize> {
329 match self {
330 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
331 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
332 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
333 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
334 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
335 }
336 }
337
338 fn has_real_vocabulary(&self) -> bool {
345 match self {
346 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
347 Model::Kimi(_) => true,
348 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
349 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
350 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
351 }
352 }
353}
354
355impl ferrox_models::dry::DryVocab for Model {
361 fn n_tokens(&self) -> usize {
362 self.vocab_size().unwrap_or(0)
363 }
364
365 fn detokenize(&self, token: usize) -> String {
366 self.decode(&[token])
367 }
368
369 fn tokenize(&self, text: &str) -> Vec<usize> {
370 self.encode(text, SpecialTokens::AsText)
371 }
372}
373
374pub(crate) struct AppState {
375 pub(crate) embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
389 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
408 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
412 pub(crate) tasks: Arc<tasks::TaskRegistry>,
414 pub(crate) cancels: Arc<cancel::CancelRegistry>,
418 pub(crate) stats: stats::Stats,
421 pub(crate) streams: resume::StreamRegistry,
424 pub(crate) model_dir: Option<PathBuf>,
426 response_cache: Mutex<ResponseCache>,
430 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
442 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
455 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
465 sessions: session::SessionStore,
471 requests_total: std::sync::atomic::AtomicU64,
472 request_errors_total: std::sync::atomic::AtomicU64,
473 started_at: std::time::Instant,
474 last_request_ms: std::sync::atomic::AtomicU64,
479 detection: Arc<health::Detection>,
481 mcp: Option<mcp::LoadedMcpConfig>,
483 pub(crate) continuous_batching_enabled: bool,
487 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
491 loading_model: Mutex<Option<String>>,
495 last_load_error: Mutex<Option<(String, String)>>,
499 pub(crate) serving: Mutex<crate::stats::ServingStats>,
504 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
509 pub(crate) footprint:
515 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
516 pub(crate) started_unix: u64,
526}
527
528pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
534
535impl AppState {
536 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
540 self.active
541 .read()
542 .unwrap_or_else(|p| p.into_inner())
543 .clone()
544 }
545
546 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
552 self.active().ok_or_else(|| {
553 (
554 StatusCode::SERVICE_UNAVAILABLE,
555 Json(serde_json::json!({"error": {
556 "message": "no model is loaded; POST /admin/models/load with an id from \
557 GET /admin/models",
558 "type": "model_not_loaded"
559 }})),
560 )
561 })
562 }
563
564 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
572 Ok(Arc::clone(self.require_active()?.generative()?))
573 }
574
575 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
583 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
584 std::mem::replace(&mut *guard, next)
585 }
586
587 fn mark_request_finished(&self) {
591 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
592 self.last_request_ms
593 .store(ms, std::sync::atomic::Ordering::Relaxed);
594 }
595
596 pub(crate) fn uptime(&self) -> Duration {
597 self.started_at.elapsed()
598 }
599
600 pub(crate) fn requests_total(&self) -> u64 {
601 self.requests_total
602 .load(std::sync::atomic::Ordering::Relaxed)
603 }
604
605 pub(crate) fn errors_total(&self) -> u64 {
606 self.request_errors_total
607 .load(std::sync::atomic::Ordering::Relaxed)
608 }
609
610 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
611 lock_cache(&self.response_cache).stats()
612 }
613
614 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
617 let last = self
618 .last_request_ms
619 .load(std::sync::atomic::Ordering::Relaxed);
620 (last > 0)
621 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
622 .map(|age| age.max(0.0))
623 }
624
625 pub(crate) fn loading_model_id(&self) -> Option<String> {
626 self.loading_model
627 .lock()
628 .unwrap_or_else(|p| p.into_inner())
629 .clone()
630 }
631
632 pub(crate) fn set_loading_model(&self, id: Option<String>) {
633 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
634 }
635
636 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
637 self.last_load_error
638 .lock()
639 .unwrap_or_else(|p| p.into_inner())
640 .clone()
641 }
642
643 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
644 *self
645 .last_load_error
646 .lock()
647 .unwrap_or_else(|p| p.into_inner()) = error;
648 }
649
650 pub(crate) fn active_model_name(&self) -> Option<String> {
659 self.active().map(|a| a.name().to_string())
660 }
661
662 pub(crate) fn embedding_model(&self) -> Option<Arc<ferrox_models::EmbeddingModel>> {
672 self.embedding
673 .clone()
674 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
675 }
676
677 pub(crate) fn embedding_model_name(&self) -> Option<String> {
681 match self.embedding_model() {
682 Some(e) => Some(e.name().to_string()),
683 None => self.active_model_name(),
684 }
685 }
686
687 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
688 self.stats.record(stats::entry(record));
689 }
690}
691
692fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
697 cache
698 .lock()
699 .unwrap_or_else(|poisoned| poisoned.into_inner())
700}
701
702#[derive(Debug, Clone, Deserialize)]
703#[serde(untagged)]
704pub(crate) enum MessageContent {
705 Text(String),
706 Parts(Vec<ContentPart>),
707}
708
709#[derive(Debug, Clone, Deserialize)]
710struct ContentPart {
711 #[serde(rename = "type")]
712 kind: String,
713 #[serde(default)]
714 text: Option<String>,
715 #[serde(default)]
716 image_url: Option<serde_json::Value>,
717}
718
719impl MessageContent {
720 fn as_text(&self) -> String {
721 match self {
722 Self::Text(s) => s.clone(),
723 Self::Parts(parts) => parts
724 .iter()
725 .filter_map(|p| p.text.as_deref())
726 .collect::<Vec<_>>()
727 .join(""),
728 }
729 }
730
731 fn has_image(&self) -> bool {
732 match self {
733 Self::Text(_) => false,
734 Self::Parts(parts) => parts
735 .iter()
736 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
737 }
738 }
739}
740
741#[derive(Debug, Clone, Deserialize)]
742pub(crate) struct ChatMessage {
743 pub(crate) role: String,
744 #[serde(default)]
748 pub(crate) content: Option<MessageContent>,
749 #[serde(default)]
753 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
754 #[serde(default)]
759 #[allow(dead_code)]
760 pub(crate) tool_call_id: Option<String>,
761 #[serde(default, alias = "reasoning")]
777 pub(crate) reasoning_content: Option<String>,
778}
779
780impl ChatMessage {
781 fn rendered_content(&self) -> String {
788 let mut out = self
789 .content
790 .as_ref()
791 .map(MessageContent::as_text)
792 .unwrap_or_default();
793 if let Some(calls) = &self.tool_calls {
794 for call in calls {
795 out.push_str(&format!(
796 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
797 call.function.name, call.function.arguments
798 ));
799 }
800 }
801 out
802 }
803}
804
805#[derive(Debug, Clone, Deserialize)]
806pub(crate) struct ToolCallIn {
807 #[serde(default)]
808 #[allow(dead_code)]
809 id: String,
810 #[serde(rename = "type", default)]
811 #[allow(dead_code)]
812 kind: String,
813 function: ToolCallFunctionIn,
814}
815
816#[derive(Debug, Clone, Deserialize)]
817struct ToolCallFunctionIn {
818 name: String,
819 arguments: String,
824}
825
826#[derive(Debug, Clone, Deserialize)]
829struct ToolDef {
830 #[serde(rename = "type", default)]
831 #[allow(dead_code)]
832 kind: String,
833 function: ToolFunctionDef,
834}
835
836#[derive(Debug, Clone, Deserialize)]
837struct ToolFunctionDef {
838 name: String,
839 #[serde(default)]
840 description: Option<String>,
841 #[serde(default)]
842 parameters: Option<serde_json::Value>,
843}
844
845#[derive(Debug, Clone, Deserialize)]
856#[serde(untagged)]
857enum ToolChoice {
858 Mode(String),
859 Specific(serde_json::Value),
860}
861
862#[derive(Deserialize)]
865#[serde(untagged)]
866enum StopParam {
867 One(String),
868 Many(Vec<String>),
869}
870
871#[derive(Deserialize)]
872struct ChatCompletionRequest {
873 model: String,
874 messages: Vec<ChatMessage>,
875 #[serde(default = "default_max_tokens")]
876 max_tokens: usize,
877 #[serde(default)]
878 temperature: Option<f32>,
879 #[serde(default)]
880 top_p: Option<f32>,
881 #[serde(default)]
886 min_p: Option<f32>,
887 #[serde(default)]
888 top_k: Option<usize>,
889 #[serde(default)]
890 repetition_penalty: Option<f32>,
891 #[serde(flatten)]
895 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
896 #[serde(default)]
897 seed: Option<u64>,
898 #[serde(default)]
899 stop: Option<StopParam>,
900 #[serde(default)]
901 stream: Option<bool>,
902 #[serde(default)]
915 stream_resumable: Option<bool>,
916 #[serde(default)]
926 ignore_eos: Option<bool>,
927 #[serde(default)]
928 tools: Vec<ToolDef>,
929 #[serde(default)]
930 tool_choice: Option<ToolChoice>,
931 #[serde(default)]
939 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
940 #[serde(default)]
948 reasoning_effort: Option<String>,
949 #[serde(default)]
953 thinking: Option<ThinkingSwitch>,
954 #[serde(default)]
959 session_id: Option<String>,
960 #[serde(default)]
962 logprobs: Option<bool>,
963 #[serde(default)]
964 top_logprobs: Option<u32>,
965 #[serde(default)]
966 n: Option<u32>,
967 #[serde(default)]
968 presence_penalty: Option<f32>,
969 #[serde(default)]
970 frequency_penalty: Option<f32>,
971 #[serde(default)]
972 response_format: Option<serde_json::Value>,
973 #[serde(default)]
980 logit_bias: Option<serde_json::Value>,
981 #[serde(default)]
991 samplers: Option<serde_json::Value>,
992 #[serde(default)]
1002 grammar: Option<String>,
1003}
1004
1005const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1017
1018#[derive(Debug, Clone, Deserialize)]
1020pub(crate) struct ThinkingSwitch {
1021 #[serde(rename = "type")]
1022 pub(crate) kind: String,
1023}
1024
1025const THINKING_KWARG_KEYS: [&str; 4] = [
1029 "enable_thinking",
1030 "thinking",
1031 "thinking_mode",
1032 "reasoning_effort",
1033];
1034
1035const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1039
1040fn default_max_tokens() -> usize {
1041 DEFAULT_CHAT_MAX_TOKENS
1042}
1043
1044impl ChatCompletionRequest {
1045 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1054 let mut knobs = SamplingKnobs {
1055 temperature: self.temperature,
1056 top_p: self.top_p,
1057 min_p: self.min_p,
1058 top_k: self.top_k,
1059 repetition_penalty: self.repetition_penalty,
1060 presence_penalty: self.presence_penalty,
1061 frequency_penalty: self.frequency_penalty,
1062 penalty_last_n: None,
1066 sampler_order: unsupported_sampling::parse_sampler_order(
1067 self.samplers.as_ref(),
1068 "/v1/chat/completions",
1069 )?,
1070 ..SamplingKnobs::default()
1071 };
1072 self.extra_samplers.apply(&mut knobs);
1073 Ok(knobs)
1074 }
1075
1076 fn sampling_params(
1077 &self,
1078 model: crate::sampling_knobs::SamplerModel<'_>,
1079 ) -> Result<SamplingParams, ApiError> {
1080 self.sampling_knobs()?.resolve(model).map_err(|e| {
1081 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1082 })
1083 }
1084
1085 fn stop_sequences(&self) -> Vec<String> {
1086 self.stop
1087 .as_ref()
1088 .map(|s| match s {
1089 StopParam::One(v) => vec![v.clone()],
1090 StopParam::Many(v) => v.clone(),
1091 })
1092 .unwrap_or_default()
1093 }
1094
1095 fn tools_active(&self) -> bool {
1100 !self.tools.is_empty()
1101 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1102 }
1103
1104 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1116 match &self.tool_choice {
1117 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1118 Some(ToolChoice::Specific(value)) => {
1119 let name = value
1123 .get("function")
1124 .and_then(|f| f.get("name"))
1125 .or_else(|| value.get("name"))
1126 .and_then(|n| n.as_str());
1127 match name {
1128 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1129 None => Err(invalid_request(
1130 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1131 function.name",
1132 "tool_choice",
1133 )),
1134 }
1135 }
1136 _ => Ok(None),
1137 }
1138 }
1139
1140 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1142 self.tools
1143 .iter()
1144 .map(|t| tool_grammar::ToolSpec {
1145 name: &t.function.name,
1146 parameters: t.function.parameters.as_ref(),
1147 })
1148 .collect()
1149 }
1150
1151 fn thinking_direction(&self) -> Option<bool> {
1194 if let Some(switch) = &self.thinking {
1195 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1196 "disabled" => Some(false),
1197 "enabled" => Some(true),
1198 _ => None,
1201 };
1202 }
1203 let effort = self.reasoning_effort.as_ref()?;
1204 DISABLE_EFFORTS
1205 .contains(&effort.trim().to_ascii_lowercase().as_str())
1206 .then_some(false)
1207 }
1208
1209 fn resolve_template_kwargs(
1210 &self,
1211 template: &chat_template::PromptTemplate,
1212 ) -> serde_json::Map<String, serde_json::Value> {
1213 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1214 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1218
1219 if !caller_steered {
1220 match self.thinking_direction() {
1221 Some(false) => {
1222 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1223 kwargs.insert(k, v);
1224 }
1225 return kwargs;
1228 }
1229 Some(true) => {
1230 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1231 kwargs.insert(k, v);
1232 }
1233 }
1234 None => {}
1235 }
1236 if let Some(effort) = &self.reasoning_effort {
1237 kwargs
1238 .entry("reasoning_effort".to_string())
1239 .or_insert_with(|| serde_json::json!(effort));
1240 }
1241 }
1242
1243 let offered: Vec<serde_json::Value> = if self.tools_active() {
1244 self.tools.iter().map(chat_template::tool_json).collect()
1245 } else {
1246 Vec::new()
1247 };
1248 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1249 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1250 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1251 kwargs.entry(k).or_insert(v);
1252 }
1253 }
1254 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1255 crate::policy::effort::EffortMapping::Mapped(to) => {
1256 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1257 }
1258 crate::policy::effort::EffortMapping::Dropped => {
1259 tracing::debug!(
1260 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1261 enough, so its own default applies"
1262 );
1263 }
1264 crate::policy::effort::EffortMapping::Unchanged => {}
1265 }
1266 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1267 kwargs
1268 }
1269
1270 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1273 if self.max_tokens == 0 {
1280 return Err(invalid_request(
1281 "max_tokens must be at least 1",
1282 "max_tokens",
1283 ));
1284 }
1285 if let Some(switch) = &self.thinking {
1290 let kind = switch.kind.trim().to_ascii_lowercase();
1291 if kind != "enabled" && kind != "disabled" {
1292 return Err(invalid_request(
1293 "thinking.type must be \"enabled\" or \"disabled\"",
1294 "thinking.type",
1295 ));
1296 }
1297 }
1298 for msg in &self.messages {
1299 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1300 return Err(unsupported_feature(
1301 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1302 ));
1303 }
1304 }
1305 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1306 return Err(unsupported_feature(
1307 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1308 ));
1309 }
1310 if self.n.is_some_and(|n| n > 1) {
1311 return Err(unsupported_feature(
1312 "n > 1 is not implemented (single completion only)",
1313 ));
1314 }
1315 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1316 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1320 let stated_grammar =
1335 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1336 if let Some(forced) = self.forced_tool_choice()? {
1344 if self.tools.is_empty() {
1345 return Err(invalid_request(
1346 "tool_choice forces a tool call, but no tools were offered",
1347 "tool_choice",
1348 ));
1349 }
1350 if let tool_grammar::Forced::Named(name) = forced {
1351 if !self.tools.iter().any(|t| t.function.name == name) {
1352 return Err(invalid_request(
1353 &format!(
1354 "tool_choice names {name:?}, which is not one of the tools offered"
1355 ),
1356 "tool_choice",
1357 ));
1358 }
1359 }
1360 if stated_grammar.is_some() {
1370 return Err(invalid_request(
1371 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1372 are two different constraints on the same generation; send one",
1373 "tool_choice",
1374 ));
1375 }
1376 if self.json_object_mode() {
1377 return Err(invalid_request(
1378 "a forced tool_choice cannot be combined with response_format json_object: \
1379 the tool-call markers are not JSON",
1380 "tool_choice",
1381 ));
1382 }
1383 }
1384 Ok(())
1385 }
1386
1387 fn effective_stop_sequences(&self) -> Vec<String> {
1394 let mut stop = self.stop_sequences();
1395 if self.tools_active() {
1396 stop.push("</tool_call>".to_string());
1397 }
1398 stop
1399 }
1400
1401 fn json_object_mode(&self) -> bool {
1402 self.response_format
1403 .as_ref()
1404 .and_then(|v| v.get("type"))
1405 .and_then(|v| v.as_str())
1406 == Some("json_object")
1407 }
1408
1409 fn generation_params(
1414 &self,
1415 model: crate::sampling_knobs::SamplerModel<'_>,
1416 ) -> Result<GenerationParams, ApiError> {
1417 Ok(GenerationParams {
1418 reasoning: None,
1423 max_tokens: self.max_tokens,
1424 sampling: self.sampling_params(model)?,
1425 seed: self.resolved_seed(),
1426 stop: self.effective_stop_sequences(),
1427 stop_token_ids: Vec::new(),
1431 json_object: self.json_object_mode(),
1432 grammar: grammar_request::for_request(
1433 self.grammar.as_deref(),
1434 self.response_format.as_ref(),
1435 )?,
1436 cancel: None,
1439 ignore_eos: self.ignore_eos.unwrap_or(false),
1440 })
1441 }
1442
1443 fn generation_params_for_template(
1454 &self,
1455 template: &chat_template::PromptTemplate,
1456 served_model: &str,
1457 model: crate::sampling_knobs::SamplerModel<'_>,
1458 ) -> Result<GenerationParams, ApiError> {
1459 let mut params = self.generation_params(model)?;
1460 params.reasoning = crate::policy::parser::ReasoningFormat::infer(served_model);
1465 if let Some(stop) = template.end_of_turn() {
1466 if !params.stop.iter().any(|s| s == stop) {
1467 params.stop.push(stop.to_string());
1468 }
1469 }
1470 if let Some(forced) = self.forced_tool_choice()? {
1471 params.grammar = Some(tool_grammar::build(
1475 forced,
1476 &self.tool_specs(),
1477 policy::parser::ToolCallFormat::infer(served_model),
1478 )?);
1479 }
1480 Ok(params)
1481 }
1482
1483 fn is_cacheable(&self) -> bool {
1492 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1493 }
1494
1495 fn cache_key(&self, prompt: &str, params: &GenerationParams) -> CacheKey {
1511 CacheKey {
1512 model: self.model.clone(),
1513 prompt: prompt.to_string(),
1514 generation: response_cache::generation_key(params),
1515 seed: self.seed,
1516 }
1517 }
1518
1519 fn resolved_seed(&self) -> u64 {
1520 self.seed.unwrap_or_else(|| {
1521 std::time::SystemTime::now()
1522 .duration_since(std::time::UNIX_EPOCH)
1523 .map(|d| d.as_nanos() as u64)
1524 .unwrap_or(0xDEFA017)
1525 })
1526 }
1527}
1528
1529#[derive(Serialize)]
1530struct ChatCompletionChoice {
1531 index: usize,
1532 message: ChatCompletionResponseMessage,
1533 finish_reason: &'static str,
1534}
1535
1536#[derive(Serialize)]
1537struct ChatCompletionResponseMessage {
1538 role: &'static str,
1539 #[serde(skip_serializing_if = "Option::is_none")]
1540 content: Option<String>,
1541 #[serde(skip_serializing_if = "Option::is_none")]
1545 reasoning_content: Option<String>,
1546 #[serde(skip_serializing_if = "Option::is_none")]
1547 tool_calls: Option<Vec<ToolCallOut>>,
1548}
1549
1550#[derive(Serialize, Clone)]
1551struct ToolCallOut {
1552 id: String,
1553 #[serde(rename = "type")]
1554 kind: &'static str,
1555 function: ToolCallFunctionOut,
1556}
1557
1558#[derive(Serialize, Clone)]
1566struct ToolCallDelta {
1567 index: usize,
1568 #[serde(skip_serializing_if = "Option::is_none")]
1569 id: Option<String>,
1570 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1571 kind: Option<&'static str>,
1572 function: ToolCallFunctionDelta,
1573}
1574
1575#[derive(Serialize, Clone, Default)]
1576struct ToolCallFunctionDelta {
1577 #[serde(skip_serializing_if = "Option::is_none")]
1578 name: Option<String>,
1579 #[serde(skip_serializing_if = "Option::is_none")]
1582 arguments: Option<String>,
1583}
1584
1585impl ToolCallDelta {
1586 fn whole(index: usize, name: String, arguments: String) -> Self {
1588 ToolCallDelta {
1589 index,
1590 id: Some(format!("call_{index}")),
1591 kind: Some("function"),
1592 function: ToolCallFunctionDelta {
1593 name: Some(name),
1594 arguments: Some(arguments),
1595 },
1596 }
1597 }
1598
1599 fn opening(index: usize, name: String) -> Self {
1601 ToolCallDelta {
1602 index,
1603 id: Some(format!("call_{index}")),
1604 kind: Some("function"),
1605 function: ToolCallFunctionDelta {
1606 name: Some(name),
1607 arguments: Some(String::new()),
1608 },
1609 }
1610 }
1611
1612 fn arguments(index: usize, fragment: String) -> Self {
1614 ToolCallDelta {
1615 index,
1616 id: None,
1617 kind: None,
1618 function: ToolCallFunctionDelta {
1619 name: None,
1620 arguments: Some(fragment),
1621 },
1622 }
1623 }
1624}
1625
1626#[derive(Serialize, Clone)]
1627struct ToolCallFunctionOut {
1628 name: String,
1629 arguments: String,
1633}
1634
1635#[derive(Serialize)]
1636struct ChatCompletionResponse {
1637 id: String,
1638 request_id: String,
1644 object: &'static str,
1645 model: String,
1646 choices: Vec<ChatCompletionChoice>,
1647 usage: generate::Usage,
1652 ferrox_cache: &'static str,
1663}
1664
1665#[derive(Serialize)]
1666struct ChatCompletionChunkDelta {
1667 #[serde(skip_serializing_if = "Option::is_none")]
1668 role: Option<&'static str>,
1669 #[serde(skip_serializing_if = "Option::is_none")]
1670 content: Option<String>,
1671 #[serde(skip_serializing_if = "Option::is_none")]
1673 reasoning_content: Option<String>,
1674 #[serde(skip_serializing_if = "Option::is_none")]
1675 tool_calls: Option<Vec<ToolCallDelta>>,
1676}
1677
1678#[derive(Serialize)]
1679struct ChatCompletionChunkChoice {
1680 index: usize,
1681 delta: ChatCompletionChunkDelta,
1682 finish_reason: Option<&'static str>,
1683}
1684
1685#[derive(Serialize)]
1686struct ChatCompletionChunk {
1687 id: String,
1688 #[serde(skip_serializing_if = "Option::is_none")]
1695 request_id: Option<String>,
1696 object: &'static str,
1697 model: String,
1698 choices: Vec<ChatCompletionChunkChoice>,
1699 #[serde(skip_serializing_if = "Option::is_none")]
1702 usage: Option<generate::Usage>,
1703}
1704
1705async fn health(State(state): State<Arc<AppState>>) -> Response {
1711 let snapshot = state.detection.snapshot();
1712 let mut capabilities = snapshot.capabilities;
1713 let active = state.active();
1714
1715 capabilities.push(match active.as_deref() {
1718 None => ferrox_api::Capability::unavailable(
1723 ferrox_api::health::capability::REAL_WEIGHTS,
1724 ferrox_api::health::reason::MODEL_NOT_LOADED,
1725 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1726 ),
1727 Some(active) if active.is_synthetic() => ferrox_api::Capability::unavailable(
1728 ferrox_api::health::capability::REAL_WEIGHTS,
1729 ferrox_api::health::reason::MODEL_NOT_LOADED,
1730 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1731 checkpoint. Output from this model is noise.",
1732 ),
1733 Some(active) if active.encoder().is_some() => {
1743 let endpoints = active
1744 .encoder()
1745 .map(|e| encoder_endpoints(e))
1746 .unwrap_or_default();
1747 let served_by = match endpoints.is_empty() {
1748 true => "no endpoint in this build serves it".to_string(),
1749 false => format!("served by {}", endpoints.join(" and ")),
1750 };
1751 ferrox_api::Capability::available(
1752 ferrox_api::health::capability::REAL_WEIGHTS,
1753 format!(
1754 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1755 {served_by}; generation endpoints refuse it.",
1756 active.name(),
1757 ),
1758 )
1759 }
1760 Some(active) => ferrox_api::Capability::available(
1761 ferrox_api::health::capability::REAL_WEIGHTS,
1762 format!("Serving the real checkpoint '{}'.", active.name()),
1763 ),
1764 });
1765 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1766 ferrox_api::Capability::available(
1767 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1768 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1769 "On by default on Metal. Concurrent requests share one batched decode worker."
1770 } else {
1771 "Concurrent requests share one batched decode step."
1772 },
1773 )
1774 } else if state.metal_private_decode_gate.is_some() {
1775 ferrox_api::Capability::unavailable(
1776 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1777 ferrox_api::health::reason::DISABLED,
1778 "Off; private Metal decodes serialize (one at a time). Set FERROX_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1779 )
1780 } else {
1781 ferrox_api::Capability::unavailable(
1782 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1783 ferrox_api::health::reason::DISABLED,
1784 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1785 )
1786 });
1787
1788 let last_request_ms = state
1789 .last_request_ms
1790 .load(std::sync::atomic::Ordering::Relaxed);
1791 let uptime = state.started_at.elapsed();
1792 let health_state = if active.is_none() {
1800 ferrox_api::HealthState::Unavailable
1801 } else {
1802 snapshot.state
1803 };
1804 let body = ferrox_api::HealthResponse {
1805 state: health_state,
1806 reason: match health_state {
1807 ferrox_api::HealthState::Ready => None,
1808 ferrox_api::HealthState::Unavailable => {
1809 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
1810 }
1811 ferrox_api::HealthState::Detecting => {
1812 Some(ferrox_api::health::reason::DETECTING.to_string())
1813 }
1814 },
1815 detail: match health_state {
1816 ferrox_api::HealthState::Ready => None,
1817 ferrox_api::HealthState::Unavailable => Some(
1818 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1819 .to_string(),
1820 ),
1821 ferrox_api::HealthState::Detecting => {
1822 Some("Probing available compute backends.".to_string())
1823 }
1824 },
1825 model: active
1826 .as_deref()
1827 .map(|active| ferrox_api::health::ModelSummary {
1828 id: active.name().to_string(),
1829 tokenizer: active.tokenizer_kind().to_string(),
1830 synthetic_weights: active.is_synthetic(),
1831 }),
1832 capabilities,
1833 version: env!("CARGO_PKG_VERSION").to_string(),
1834 pid: std::process::id(),
1835 uptime_seconds: uptime.as_secs_f64(),
1836 server_time_unix_ms: std::time::SystemTime::now()
1837 .duration_since(std::time::UNIX_EPOCH)
1838 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1839 .unwrap_or(0),
1840 last_request_age_seconds: (last_request_ms > 0)
1841 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1842 .map(|age| age.max(0.0)),
1843 };
1844
1845 let status =
1846 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1847 (status, Json(body)).into_response()
1848}
1849
1850async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1851 let Some(active) = state.active() else {
1855 return Json(serde_json::json!({ "object": "list", "data": [] }));
1856 };
1857 let mut model_entry = serde_json::json!({
1858 "id": active.name(),
1859 "object": "model",
1860 "ferrox_synthetic_weights": active.is_synthetic(),
1861 "ferrox_tokenizer": active.tokenizer_kind(),
1862 });
1863 if let Some(encoder) = active.encoder() {
1869 model_entry["ferrox_model_kind"] = serde_json::json!("embedding");
1870 model_entry["ferrox_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1871 model_entry["ferrox_n_embd"] = serde_json::json!(encoder.n_embd());
1872 model_entry["ferrox_pooling"] = serde_json::json!(encoder.pooling_type().name());
1873 model_entry["ferrox_context_length"] = serde_json::json!(encoder.n_ctx_train());
1874 }
1875 if let Some(model) = active.generative_opt() {
1883 let parser_configured =
1884 crate::policy::parser::ReasoningFormat::infer(active.name()).is_some();
1885 let gears = model.chat_template().think_gears(parser_configured);
1886 if !gears.is_empty() {
1887 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1888 if let Some(default) = &gears.default {
1889 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1890 }
1891 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1894 }
1895 }
1896 if let Some(mcp) = &state.mcp {
1897 model_entry["ferrox_mcp"] = mcp.models_metadata();
1898 }
1899 Json(serde_json::json!({
1900 "object": "list",
1901 "data": [model_entry]
1902 }))
1903}
1904
1905async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1919 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1920 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1921 let active = state.active();
1922 Json(serde_json::json!({
1923 "model": active.as_ref().map(|a| a.name()),
1924 "state": state
1925 .maintenance
1926 .lock()
1927 .unwrap_or_else(|p| p.into_inner())
1928 .state()
1929 .as_str(),
1930 "uptime_s": state.uptime().as_secs(),
1931 "throughput": {
1932 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1933 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1934 },
1935 "requests": {
1936 "active": state.cancels.live_count(),
1937 "completed": state.stats.recorded_total(),
1938 "p95_ms": state.stats.p95_duration_ms(),
1939 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1940 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1941 "completion_tokens_total": state.stats.tokens_generated_total(),
1942 },
1943 "pools": cache_admin::pool_gauges(&state),
1949 "memory": cache_admin::footprint_json(&state),
1952 }))
1953}
1954
1955#[derive(Deserialize)]
1956struct RequestsQuery {
1957 #[serde(default)]
1958 since: u64,
1959 #[serde(default = "default_requests_limit")]
1960 limit: usize,
1961}
1962
1963fn default_requests_limit() -> usize {
1964 stats::MAX_PAGE
1965}
1966
1967async fn recent_requests(
1975 State(state): State<Arc<AppState>>,
1976 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1977) -> Json<serde_json::Value> {
1978 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1979 Json(serde_json::json!({
1980 "requests": rows,
1981 "next_cursor": cursor,
1982 "missed": missed,
1983 "total": state.stats.recorded_total(),
1984 }))
1985}
1986
1987#[derive(Serialize)]
1988struct CombinedCacheStats {
1989 response_cache: response_cache::CacheStats,
1990 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
1992}
1993
1994async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1995 Json(CombinedCacheStats {
1996 response_cache: lock_cache(&state.response_cache).stats(),
1997 prefix_cache: state
1998 .prefix_cache
1999 .as_ref()
2000 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
2001 })
2002}
2003
2004async fn metrics(State(state): State<Arc<AppState>>) -> Response {
2009 use std::sync::atomic::Ordering;
2010
2011 let cache_stats = lock_cache(&state.response_cache).stats();
2012 let active = state.active();
2013 let requests_total = state.requests_total.load(Ordering::Relaxed);
2014 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2015 let uptime = state.started_at.elapsed().as_secs_f64();
2016
2017 let body = format!(
2018 "# HELP ferrox_requests_total Total chat completion requests received.\n\
2019 # TYPE ferrox_requests_total counter\n\
2020 ferrox_requests_total {requests_total}\n\
2021 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
2022 # TYPE ferrox_request_errors_total counter\n\
2023 ferrox_request_errors_total {errors_total}\n\
2024 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
2025 # TYPE ferrox_cache_hits_total counter\n\
2026 ferrox_cache_hits_total {}\n\
2027 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
2028 # TYPE ferrox_cache_misses_total counter\n\
2029 ferrox_cache_misses_total {}\n\
2030 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
2031 # TYPE ferrox_cache_entries gauge\n\
2032 ferrox_cache_entries {}\n\
2033 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2034 # TYPE ferrox_synthetic_weights gauge\n\
2035 ferrox_synthetic_weights {}\n\
2036 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
2037 # TYPE ferrox_uptime_seconds gauge\n\
2038 ferrox_uptime_seconds {uptime}\n",
2039 cache_stats.hits,
2040 cache_stats.misses,
2041 cache_stats.entries,
2042 active
2046 .as_ref()
2047 .map(|a| a.is_synthetic() as u8)
2048 .unwrap_or(0),
2049 );
2050
2051 let body = match active
2055 .as_ref()
2056 .and_then(|a| a.expert_store_stats())
2057 {
2058 Some(es) => format!(
2059 "{body}\
2060 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
2061 # TYPE ferrox_expert_cache_hits_total counter\n\
2062 ferrox_expert_cache_hits_total {}\n\
2063 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2064 # TYPE ferrox_expert_cache_misses_total counter\n\
2065 ferrox_expert_cache_misses_total {}\n\
2066 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
2067 # TYPE ferrox_expert_cache_evictions_total counter\n\
2068 ferrox_expert_cache_evictions_total {}\n\
2069 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2070 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
2071 ferrox_expert_cache_pass_throughs_total {}\n\
2072 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2073 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
2074 ferrox_expert_cache_bytes_read_total {}\n\
2075 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2076 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
2077 ferrox_expert_cache_resident_bytes {}\n",
2078 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2079 ),
2080 None => body,
2081 };
2082
2083 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2088 Some(batcher) => {
2089 let sched = batcher.stats();
2090 format!(
2091 "{body}\
2092 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2093 # TYPE ferrox_prefill_chunks_total counter\n\
2094 ferrox_prefill_chunks_total {}\n\
2095 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2096 # TYPE ferrox_prefill_tokens_total counter\n\
2097 ferrox_prefill_tokens_total {}\n\
2098 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2099 # TYPE ferrox_decode_steps_total counter\n\
2100 ferrox_decode_steps_total {}\n\
2101 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2102 # TYPE ferrox_scheduler_queue_depth gauge\n\
2103 ferrox_scheduler_queue_depth {}\n\
2104 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2105 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2106 ferrox_scheduler_queue_rejected_total {}\n\
2107 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2108 # TYPE ferrox_kv_blocks_total gauge\n\
2109 ferrox_kv_blocks_total {}\n\
2110 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2111 # TYPE ferrox_kv_blocks_free gauge\n\
2112 ferrox_kv_blocks_free {}\n\
2113 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2114 # TYPE ferrox_kv_block_size gauge\n\
2115 ferrox_kv_block_size {}\n\
2116 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2117 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2118 ferrox_kv_rejected_too_large_total {}\n\
2119 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2120 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2121 ferrox_kv_rejected_context_length_total {}\n\
2122 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2123 # TYPE ferrox_scheduler_aborted_total counter\n\
2124 ferrox_scheduler_aborted_total {}\n\
2125 # HELP ferrox_scheduler_max_seqs Cap on in-flight sequences (-np / FERROX_CB_MAX_SEQS); 0 when unlimited.\n\
2126 # TYPE ferrox_scheduler_max_seqs gauge\n\
2127 ferrox_scheduler_max_seqs {}\n\
2128 # HELP ferrox_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FERROX_CB_PREFILL_CHUNK).\n\
2129 # TYPE ferrox_scheduler_prefill_chunk gauge\n\
2130 ferrox_scheduler_prefill_chunk {}\n",
2131 sched.prefill_chunks,
2132 sched.prefill_tokens,
2133 sched.decode_steps,
2134 sched.queue_depth,
2135 sched.queue_rejected,
2136 sched.kv_blocks_total,
2137 sched.kv_blocks_free,
2138 sched.kv_block_size,
2139 sched.kv_rejected_too_large,
2140 sched.kv_rejected_context_length,
2141 sched.aborted,
2142 sched.max_seqs,
2143 sched.prefill_chunk,
2144 )
2145 }
2146 None => body,
2147 };
2148
2149 (
2150 [(
2151 axum::http::header::CONTENT_TYPE,
2152 "text/plain; version=0.0.4",
2153 )],
2154 body,
2155 )
2156 .into_response()
2157}
2158
2159pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2160
2161pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2167 (
2168 StatusCode::BAD_REQUEST,
2169 Json(serde_json::json!({"error": {
2170 "message": message,
2171 "type": "invalid_request_error",
2172 "param": param,
2173 "code": null,
2174 }})),
2175 )
2176}
2177
2178pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2179 (
2180 StatusCode::NOT_IMPLEMENTED,
2181 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2182 )
2183}
2184
2185pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2186 let status = match e {
2187 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2188 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2193 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2199 StatusCode::SERVICE_UNAVAILABLE
2200 }
2201 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2205 };
2206 tracing::warn!("decode error: {e}");
2207 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2208 if let generate::DecodeError::KvBudgetExceeded {
2214 binding,
2215 estimated_bytes,
2216 limit_bytes,
2217 positions,
2218 positions_limit,
2219 ..
2220 } = &e
2221 {
2222 body["error"]["type"] = serde_json::json!("invalid_request_error");
2223 body["error"]["code"] = serde_json::json!(binding);
2224 body["error"]["binding"] = serde_json::json!(binding);
2225 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2226 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2227 body["error"]["positions"] = serde_json::json!(positions);
2228 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2229 }
2230 if let Some(secs) = e.retry_after_secs() {
2234 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2235 }
2236 (status, Json(body))
2237}
2238
2239pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2240 tracing::error!("generation task panicked: {e}");
2241 (
2242 StatusCode::INTERNAL_SERVER_ERROR,
2243 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2244 )
2245}
2246
2247#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2257 model: &Model,
2258 prompt: &str,
2259 params: &GenerationParams,
2260 kv_pool: Option<&generate::KvPoolConfig>,
2261 paged_kv: Option<&generate::PagedKvConfig>,
2262 prefix_cache: Option<&Mutex<PrefixCache>>,
2263 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2264 ceiling: Option<&budget::ContextCeiling>,
2265 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2266 mut emit: impl FnMut(&str),
2267) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2268 let synthetic = model.is_synthetic();
2269 let mut chunks = Vec::new();
2270 let params = &{
2276 let mut resolved = params.clone();
2277 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2278 model.encode(text, SpecialTokens::Parse)
2279 });
2280 resolved
2281 };
2282 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2283 let _metal_private_guard =
2284 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2285 let (finish, usage) = match model {
2286 Model::Gguf(m) => {
2287 if let Some(batcher) = continuous_batcher {
2288 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2289 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2290 let (finish, _generated_ids, text, usage) = if synthetic {
2291 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2292 } else {
2293 batcher.generate_streaming(
2294 tokens,
2295 params.clone(),
2296 m.stop_tokens.clone(),
2297 Some(|chunk: &str| {
2298 if !chunk.is_empty() {
2299 chunks.push(chunk.to_string());
2300 emit(chunk);
2301 }
2302 }),
2303 )?
2304 };
2305 if !text.is_empty() && chunks.is_empty() {
2306 chunks.push(text);
2307 }
2308 (finish, usage)
2309 } else {
2310 generate::generate(
2311 &m.decoder,
2312 m.tokenizer.as_ref(),
2313 &m.stop_tokens,
2314 m.bos_id,
2315 prompt,
2316 params,
2317 kv_pool,
2318 paged_kv,
2319 prefix_cache,
2320 ceiling,
2321 |chunk| {
2322 chunks.push(chunk.to_string());
2323 if !synthetic {
2324 emit(chunk);
2325 }
2326 },
2327 )?
2328 }
2329 }
2330 Model::Kimi(m) => generate::generate_engine(
2331 &m.engine,
2332 &m.tokenizer,
2333 &m.stop_tokens,
2334 None,
2335 prompt,
2336 params,
2337 |chunk| {
2338 chunks.push(chunk.to_string());
2339 if !synthetic {
2340 emit(chunk);
2341 }
2342 },
2343 )?,
2344 Model::Mla(m) => generate::generate_engine(
2345 &m.engine,
2346 &m.tokenizer,
2347 &m.stop_tokens,
2348 m.bos_id,
2349 prompt,
2350 params,
2351 |chunk| {
2352 chunks.push(chunk.to_string());
2353 if !synthetic {
2354 emit(chunk);
2355 }
2356 },
2357 )?,
2358 Model::Gemma4(m) => generate::generate_engine(
2359 &m.engine,
2360 &m.tokenizer,
2361 &m.stop_tokens,
2362 m.bos_id,
2363 prompt,
2364 params,
2365 |chunk| {
2366 chunks.push(chunk.to_string());
2367 if !synthetic {
2368 emit(chunk);
2369 }
2370 },
2371 )?,
2372 Model::Glm52(m) => generate::generate_engine(
2373 &m.engine,
2374 &m.tokenizer,
2375 &m.stop_tokens,
2376 m.bos_id,
2377 prompt,
2378 params,
2379 |chunk| {
2380 chunks.push(chunk.to_string());
2381 if !synthetic {
2382 emit(chunk);
2383 }
2384 },
2385 )?,
2386 };
2387
2388 let mut full = chunks.concat();
2389 if synthetic {
2390 full = format!(
2391 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2392 to serve a real model. Decoded ids -> {full:?}]"
2393 );
2394 emit(&full);
2395 } else if used_batcher && !full.is_empty() && chunks.is_empty() {
2396 emit(&full);
2397 }
2398
2399 Ok((finish, usage, full))
2400}
2401
2402#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2407 model: &Model,
2408 prompt: &str,
2409 params: &GenerationParams,
2410 kv_pool: Option<&generate::KvPoolConfig>,
2411 paged_kv: Option<&generate::PagedKvConfig>,
2412 prefix_cache: Option<&Mutex<PrefixCache>>,
2413 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2414 ceiling: Option<&budget::ContextCeiling>,
2415 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2416) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2417 let (finish, usage, full) = run_generation_emit(
2418 model,
2419 prompt,
2420 params,
2421 kv_pool,
2422 paged_kv,
2423 prefix_cache,
2424 continuous_batcher,
2425 ceiling,
2426 metal_private_decode_gate,
2427 |_| {},
2428 )?;
2429 Ok((
2430 if full.is_empty() {
2431 Vec::new()
2432 } else {
2433 vec![full]
2434 },
2435 finish,
2436 usage,
2437 ))
2438}
2439
2440pub(crate) fn prompt_from_messages(
2451 messages: &[ChatMessage],
2452 template: &chat_template::PromptTemplate,
2453 tools: &[ToolDef],
2454 extra: serde_json::Map<String, serde_json::Value>,
2455) -> Result<String, ApiError> {
2456 let rendered = if tools.is_empty() || template.handles_tools() {
2457 template.render(messages, tools, extra)
2458 } else {
2459 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2460 with_preamble.push(ChatMessage {
2461 role: "system".to_string(),
2462 content: Some(MessageContent::Text(tool_preamble(tools))),
2463 tool_calls: None,
2464 tool_call_id: None,
2465 reasoning_content: None,
2466 });
2467 with_preamble.extend_from_slice(messages);
2468 template.render(&with_preamble, &[], extra)
2469 };
2470 rendered.map_err(template_error_response)
2471}
2472
2473fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2478 (
2479 StatusCode::BAD_REQUEST,
2480 Json(serde_json::json!({
2481 "error": {
2482 "message": format!("chat template failed to render: {err}"),
2483 "type": "invalid_request_error",
2484 "param": "messages",
2485 "code": null,
2486 }
2487 })),
2488 )
2489}
2490
2491fn tool_preamble(tools: &[ToolDef]) -> String {
2504 let mut out = String::from(
2505 "You can call tools to help answer the user. To call a tool, respond with \
2506 EXACTLY one line in this format and nothing else:\n\
2507 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2508 object matching that tool's parameters>}}</tool_call>\n\n\
2509 Available tools:\n",
2510 );
2511 for t in tools {
2512 out.push_str(&format!(
2513 "- {}: {}\n parameters (JSON schema): {}\n",
2514 t.function.name,
2515 t.function.description.as_deref().unwrap_or(""),
2516 t.function
2517 .parameters
2518 .as_ref()
2519 .map(|v| v.to_string())
2520 .unwrap_or_else(|| "{}".to_string()),
2521 ));
2522 }
2523 out
2524}
2525
2526fn tool_call_deltas(
2536 events: Vec<crate::policy::parser::ToolCallEvent>,
2537 opened: &std::cell::Cell<usize>,
2538) -> (String, Vec<ToolCallDelta>) {
2539 let mut text = String::new();
2540 let mut deltas = Vec::new();
2541 for event in events {
2542 match event {
2543 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2544 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2545 opened.set(opened.get().max(index + 1));
2546 deltas.push(ToolCallDelta::opening(index, name));
2547 }
2548 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2549 if !fragment.is_empty() {
2550 deltas.push(ToolCallDelta::arguments(index, fragment));
2551 }
2552 }
2553 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2554 }
2555 }
2556 (text, deltas)
2557}
2558
2559fn build_response_message(
2570 text: String,
2571 tools: &[ToolDef],
2572 posture: output::OutputPosture,
2573 base_finish: &'static str,
2574) -> (ChatCompletionResponseMessage, &'static str) {
2575 let parsed = output::parse_output(&text, tools, posture);
2576 let calls: Vec<ToolCallOut> = parsed
2577 .calls
2578 .into_iter()
2579 .enumerate()
2580 .map(|(index, call)| ToolCallOut {
2581 id: format!("call_{index}"),
2582 kind: "function",
2583 function: ToolCallFunctionOut {
2584 name: call.name,
2585 arguments: call.arguments,
2586 },
2587 })
2588 .collect();
2589 if !calls.is_empty() {
2590 return (
2591 ChatCompletionResponseMessage {
2592 role: "assistant",
2593 content: None,
2594 reasoning_content: parsed.reasoning,
2595 tool_calls: Some(calls),
2596 },
2597 "tool_calls",
2598 );
2599 }
2600 (
2601 ChatCompletionResponseMessage {
2602 role: "assistant",
2603 content: Some(parsed.content),
2604 reasoning_content: parsed.reasoning,
2605 tool_calls: None,
2606 },
2607 base_finish,
2608 )
2609}
2610
2611fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2616 let mut history = match &req.session_id {
2617 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2618 None => req.messages.clone(),
2619 };
2620 if req.json_object_mode() {
2621 inject_json_object_system_hint(&mut history);
2622 }
2623 history
2624}
2625
2626fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2627 const HINT: &str =
2628 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2629 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2630 match &mut sys.content {
2631 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2632 s.push_str("\n\n");
2633 s.push_str(HINT);
2634 }
2635 None => {
2636 sys.content = Some(MessageContent::Text(HINT.to_string()));
2637 }
2638 _ => {}
2639 }
2640 } else {
2641 messages.insert(
2642 0,
2643 ChatMessage {
2644 role: "system".to_string(),
2645 content: Some(MessageContent::Text(HINT.to_string())),
2646 tool_calls: None,
2647 tool_call_id: None,
2648 reasoning_content: None,
2649 },
2650 );
2651 }
2652}
2653
2654async fn chat_completions(
2655 State(state): State<Arc<AppState>>,
2656 headers: axum::http::HeaderMap,
2657 Json(req): Json<ChatCompletionRequest>,
2658) -> Response {
2659 let attribution = attribution::Attribution::from_headers(&headers);
2660 state
2661 .requests_total
2662 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2663 let started = std::time::Instant::now();
2664
2665 let request_id = ferrox_api::next_request_id();
2669 let stream = req.stream.unwrap_or(false);
2670
2671 let refusal = cache_admin::check_admission(&state)
2677 .err()
2678 .or_else(|| req.validate_supported_fields().err());
2679 if let Some(err) = refusal {
2680 state
2681 .request_errors_total
2682 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2683 let response = err.into_response();
2684 state.record_request(stats::Record {
2685 request_id: &request_id,
2686 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2687 model: state.active_model_name(),
2688 status: response.status().as_u16(),
2689 stream,
2690 duration_ms: started.elapsed().as_millis() as u64,
2691 usage: None,
2692 attribution: &attribution,
2693 });
2694 return response;
2695 }
2696
2697 let response = if stream {
2698 chat_completions_stream(
2699 Arc::clone(&state),
2700 req,
2701 request_id.clone(),
2702 started,
2703 attribution.clone(),
2704 )
2705 .await
2706 .into_response()
2707 } else {
2708 chat_completions_full(
2709 Arc::clone(&state),
2710 req,
2711 request_id.clone(),
2712 started,
2713 attribution.clone(),
2714 )
2715 .await
2716 .into_response()
2717 };
2718
2719 if response.status().is_client_error() || response.status().is_server_error() {
2720 state
2721 .request_errors_total
2722 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2723 state.record_request(stats::Record {
2727 request_id: &request_id,
2728 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2729 model: state.active_model_name(),
2732 status: response.status().as_u16(),
2733 stream,
2734 duration_ms: started.elapsed().as_millis() as u64,
2735 usage: None,
2736 attribution: &attribution,
2737 });
2738 }
2739 state.mark_request_finished();
2740
2741 response
2742}
2743
2744async fn chat_completions_full(
2745 state: Arc<AppState>,
2746 req: ChatCompletionRequest,
2747 request_id: String,
2748 started: std::time::Instant,
2749 attribution: attribution::Attribution,
2750) -> Result<Json<ChatCompletionResponse>, ApiError> {
2751 let tools_active = req.tools_active();
2752 let active = state.require_active()?;
2756 let history = resolve_history(&state, &req);
2757 let template = active.generative()?.chat_template();
2758 let kwargs = req.resolve_template_kwargs(&template);
2759 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2760 let params =
2768 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2769 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2770
2771 let (completion, cache_status) = if let Some(cached) = key
2772 .as_ref()
2773 .and_then(|key| lock_cache(&state.response_cache).get(key))
2774 {
2775 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2776 (cached, "hit")
2777 } else {
2778 let (chunks, finish, usage) = decode_task::buffered(
2779 decode_task::DecodeHandles::take(&state, &active)?,
2780 prompt.clone(),
2781 params,
2782 )
2783 .await?;
2784
2785 let completion = response_cache::CachedCompletion {
2786 content: chunks.concat(),
2787 finish,
2788 usage,
2789 };
2790 let cache_status = match key {
2796 Some(key) => match completion.clone().cacheable() {
2800 Some(cacheable) => {
2801 tracing::debug!("cache miss for key {}", key.digest());
2802 lock_cache(&state.response_cache).put(key, cacheable);
2803 "miss"
2804 }
2805 None => "skip",
2806 },
2807 None => "skip",
2808 };
2809 (completion, cache_status)
2810 };
2811 let content = completion.content;
2812
2813 if req.json_object_mode() {
2814 json_mode::validate_json_object_output(&content)?;
2815 }
2816
2817 if let Some(id) = &req.session_id {
2821 state.sessions.store_reply(
2822 id,
2823 ChatMessage {
2824 role: "assistant".to_string(),
2825 content: Some(MessageContent::Text(content.clone())),
2826 tool_calls: None,
2827 tool_call_id: None,
2828 reasoning_content: None,
2829 },
2830 );
2831 }
2832
2833 let (message, finish_reason) = build_response_message(
2834 content,
2835 if tools_active { &req.tools } else { &[] },
2836 output::OutputPosture::resolve(active.name(), &prompt),
2837 completion.finish.as_str(),
2838 );
2839
2840 state.record_request(stats::Record {
2841 request_id: &request_id,
2842 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2843 model: Some(active.name().to_string()),
2846 status: 200,
2847 stream: false,
2848 duration_ms: started.elapsed().as_millis() as u64,
2849 usage: Some(&completion.usage),
2850 attribution: &attribution,
2851 });
2852
2853 Ok(Json(ChatCompletionResponse {
2854 id: request_id.clone(),
2855 request_id,
2856 object: "chat.completion",
2857 model: req.model,
2858 choices: vec![ChatCompletionChoice {
2859 index: 0,
2860 message,
2861 finish_reason,
2862 }],
2863 usage: completion.usage,
2864 ferrox_cache: cache_status,
2865 }))
2866}
2867
2868async fn chat_completions_stream(
2869 state: Arc<AppState>,
2870 req: ChatCompletionRequest,
2871 request_id: String,
2872 started: std::time::Instant,
2873 attribution: attribution::Attribution,
2874) -> Result<Response, ApiError> {
2875 let tools_active = req.tools_active();
2877 let active = state.require_active()?;
2881 let history = resolve_history(&state, &req);
2882 let template = active.generative()?.chat_template();
2883 let kwargs = req.resolve_template_kwargs(&template);
2884 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2885 let model_name = req.model.clone();
2886 let session_id = req.session_id.clone();
2887 let sessions = state.sessions.clone();
2888
2889 let model = Arc::clone(active.generative()?);
2890 let kv_pool = state.kv_pool.clone();
2891 let paged_kv = state.paged_kv.clone();
2892 let prefix_cache = state.prefix_cache.clone();
2893 let batcher = active.batcher.clone();
2894 let ceiling = active.ceiling.clone();
2895 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2896 let mut params =
2897 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2898 let stats_state = Arc::clone(&state);
2899 let served_model = active.name().to_string();
2902 let posture = output::OutputPosture::resolve(&served_model, &prompt);
2906 let offered_tools: Vec<ToolDef> = if tools_active {
2909 req.tools.clone()
2910 } else {
2911 Vec::new()
2912 };
2913
2914 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2919 params.cancel = Some(cancel_token.clone());
2920
2921 let overlap = true;
2930
2931 let slot = req
2935 .stream_resumable
2936 .unwrap_or(false)
2937 .then(|| state.streams.register(&request_id));
2938 let emitter = resume::Emitter::new(slot);
2939
2940 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2941 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2945 id: request_id.clone(),
2946 request_id: None,
2947 object: "chat.completion.chunk",
2948 model: model_name.clone(),
2949 choices: vec![ChatCompletionChunkChoice {
2950 index: 0,
2951 delta: ChatCompletionChunkDelta {
2952 role: None,
2953 content: None,
2954 reasoning_content: None,
2955 tool_calls: None,
2956 },
2957 finish_reason: None,
2958 }],
2959 usage: None,
2960 });
2961
2962 tokio::task::spawn_blocking(move || {
2963 let _cancel_guard = cancel_guard;
2966 let tx_chunks = tx.clone();
2967 let orphan_timeout = sse::orphan_timeout_from_env();
2972 let mut first = true;
2973 let head_request_id = request_id.clone();
2974 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
2984 Rc::new(RefCell::new(posture.reasoning_parser()));
2985 let emit_reasoning = Rc::clone(&stream_reasoning);
2986 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
2990 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
2991 );
2992 let emit_tools = Rc::clone(&stream_tools);
2993 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
2997 let emit_streamed_calls = Rc::clone(&streamed_calls);
2998 let result = run_generation_emit(
2999 &model,
3000 &prompt,
3001 ¶ms,
3002 kv_pool.as_ref(),
3003 paged_kv.as_ref(),
3004 prefix_cache.as_deref(),
3005 batcher.as_ref(),
3006 ceiling.as_deref(),
3007 metal_private_decode_gate.as_deref(),
3008 |chunk| {
3009 if !overlap || chunk.is_empty() {
3010 return;
3011 }
3012 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3013 Some(parser) => {
3014 let delta = parser.push(chunk);
3015 (delta.reasoning, delta.content)
3016 }
3017 None => (String::new(), chunk.to_string()),
3018 };
3019 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3023 Some(parser) => {
3024 let (text, calls) =
3025 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3026 (text, calls)
3027 }
3028 None => (content, Vec::new()),
3029 };
3030 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3033 return;
3034 }
3035 let role = if first { Some("assistant") } else { None };
3036 let request_id = first.then(|| head_request_id.clone());
3037 first = false;
3038 let payload = ChatCompletionChunk {
3039 id: head_request_id.clone(),
3040 request_id,
3041 object: "chat.completion.chunk",
3042 model: model_name.clone(),
3043 choices: vec![ChatCompletionChunkChoice {
3044 index: 0,
3045 delta: ChatCompletionChunkDelta {
3046 role,
3047 content: (!content.is_empty()).then_some(content),
3048 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3049 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3050 },
3051 finish_reason: None,
3052 }],
3053 usage: None,
3054 };
3055 if let Err(why) =
3064 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3065 {
3066 if why == sse::SendFailure::Orphaned {
3067 tracing::warn!(
3068 "SSE stream {head_request_id} accepted nothing for the orphan \
3069 deadline; treating it as abandoned"
3070 );
3071 }
3072 if !emitter.is_resumable() {
3084 cancel_token.cancel();
3085 }
3086 }
3087 },
3088 );
3089
3090 let mut pending_request_id = first.then(|| request_id.clone());
3096
3097 match result {
3098 Ok((finish, usage, full_text)) => {
3099 if let Some(id) = &session_id {
3100 sessions.store_reply(
3101 id,
3102 ChatMessage {
3103 role: "assistant".to_string(),
3104 content: Some(MessageContent::Text(full_text.clone())),
3105 tool_calls: None,
3106 tool_call_id: None,
3107 reasoning_content: None,
3108 },
3109 );
3110 }
3111 let mut streamed_finish: Option<&'static str> = None;
3117 if overlap {
3118 let tail = stream_reasoning
3119 .borrow_mut()
3120 .as_mut()
3121 .map(|parser| parser.flush())
3122 .unwrap_or_default();
3123 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3124 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3125 let mut events = parser.push(&content);
3126 events.extend(parser.finish());
3127 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3128 content = text;
3129 tool_calls = calls;
3130 }
3131 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3132 let payload = ChatCompletionChunk {
3133 id: request_id.clone(),
3134 request_id: pending_request_id.take(),
3135 object: "chat.completion.chunk",
3136 model: model_name.clone(),
3137 choices: vec![ChatCompletionChunkChoice {
3138 index: 0,
3139 delta: ChatCompletionChunkDelta {
3140 role: None,
3141 content: (!content.is_empty()).then_some(content),
3142 reasoning_content: (!tail.reasoning.is_empty())
3143 .then_some(tail.reasoning),
3144 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3145 },
3146 finish_reason: None,
3147 }],
3148 usage: None,
3149 };
3150 let _ =
3151 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3152 }
3153 if streamed_calls.get() > 0 {
3154 streamed_finish = Some("tool_calls");
3155 }
3156 } else {
3157 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3160 let tool_calls: Vec<ToolCallDelta> = parsed
3161 .calls
3162 .iter()
3163 .enumerate()
3164 .map(|(index, call)| {
3165 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3166 })
3167 .collect();
3168 if !tool_calls.is_empty() {
3169 streamed_finish = Some("tool_calls");
3170 }
3171 if !tool_calls.is_empty()
3172 || !parsed.content.is_empty()
3173 || parsed.reasoning.is_some()
3174 {
3175 let payload = ChatCompletionChunk {
3176 id: request_id.clone(),
3177 request_id: pending_request_id.take(),
3178 object: "chat.completion.chunk",
3179 model: model_name.clone(),
3180 choices: vec![ChatCompletionChunkChoice {
3181 index: 0,
3182 delta: ChatCompletionChunkDelta {
3183 role: Some("assistant"),
3184 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3185 .then(|| parsed.content.clone()),
3186 reasoning_content: parsed.reasoning.clone(),
3187 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3188 },
3189 finish_reason: None,
3190 }],
3191 usage: None,
3192 };
3193 let _ =
3194 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3195 }
3196 }
3197 let final_finish_reason = match streamed_finish {
3201 Some(reason) if finish.as_str() != "length" => reason,
3202 _ => finish.as_str(),
3203 };
3204 let final_payload = ChatCompletionChunk {
3205 id: request_id.clone(),
3206 request_id: pending_request_id.take(),
3207 object: "chat.completion.chunk",
3208 model: model_name,
3209 choices: vec![ChatCompletionChunkChoice {
3210 index: 0,
3211 delta: ChatCompletionChunkDelta {
3212 role: None,
3213 content: None,
3214 reasoning_content: None,
3215 tool_calls: None,
3216 },
3217 finish_reason: Some(final_finish_reason),
3218 }],
3219 usage: Some(usage.clone()),
3220 };
3221 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3222 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3223 stats_state.record_request(stats::Record {
3228 request_id: &request_id,
3229 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3230 model: Some(served_model.clone()),
3231 status: 200,
3232 stream: true,
3233 duration_ms: started.elapsed().as_millis() as u64,
3234 usage: Some(&usage),
3235 attribution: &attribution,
3236 });
3237 }
3238 Err(e) => {
3239 tracing::warn!("decode error on streamed request {request_id}: {e}");
3240 stats_state.record_request(stats::Record {
3246 request_id: &request_id,
3247 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3248 model: Some(served_model.clone()),
3249 status: 500,
3250 stream: true,
3251 duration_ms: started.elapsed().as_millis() as u64,
3252 usage: None,
3253 attribution: &attribution,
3254 });
3255 let payload = ChatCompletionChunk {
3256 id: request_id.clone(),
3257 request_id: pending_request_id.take(),
3258 object: "chat.completion.chunk",
3259 model: model_name,
3260 choices: vec![ChatCompletionChunkChoice {
3261 index: 0,
3262 delta: ChatCompletionChunkDelta {
3263 role: Some("assistant"),
3264 content: Some(format!("[error: {e}]")),
3265 reasoning_content: None,
3266 tool_calls: None,
3267 },
3268 finish_reason: Some("stop"),
3269 }],
3270 usage: None,
3271 };
3272 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3273 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3274 }
3275 }
3276 drop(emitter);
3280 });
3281
3282 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3283 Ok((
3306 [(
3307 axum::http::HeaderName::from_static("x-accel-buffering"),
3308 axum::http::HeaderValue::from_static("no"),
3309 )],
3310 Sse::new(stream),
3311 )
3312 .into_response())
3313}
3314
3315fn protected_routes() -> Router<Arc<AppState>> {
3346 use ferrox_api::routes;
3347
3348 Router::new()
3349 .route(routes::V1_MODELS, get(list_models))
3350 .route(routes::V1_RESPONSES, post(responses::responses))
3354 .route(
3355 &axum_path(routes::V1_RESPONSE),
3356 get(responses::responses_get),
3357 )
3358 .route(
3359 &axum_path(routes::V1_RESPONSE_CANCEL),
3360 post(responses::responses_cancel),
3361 )
3362 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3363 .route(routes::V1_STATS, get(serving_stats))
3364 .route(routes::V1_REQUESTS, get(recent_requests))
3365 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3366 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3367 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3368 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3369 .route(routes::V1_CANCEL, post(cancel_generation))
3373 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3378 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3379 .route(routes::V1_MESSAGES, post(anthropic::messages))
3380 .route(
3381 routes::V1_MESSAGES_COUNT_TOKENS,
3382 post(anthropic::count_tokens),
3383 )
3384 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3385 .route(routes::COMPLETION, post(completion::completion))
3390 .route(routes::COMPLETIONS, post(completion::completion))
3391 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3392 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3393 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3399 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3400 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3401 .route(routes::V1_RERANK, post(rerank::rerank))
3405 .route(routes::RERANK, post(rerank::rerank))
3406 .route(routes::CACHE_STATS, get(cache_stats))
3407 .route(routes::METRICS, get(metrics))
3408 .route(routes::ADMIN_MODELS, get(admin::models))
3413 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3414 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3415 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3416 .route(routes::ADMIN_TASKS, get(admin::tasks))
3417 .route(&admin::cancel_route(), post(admin::cancel_task))
3418 .route(routes::ADMIN_STATS, get(admin::stats))
3419 .merge(conversations::router())
3423}
3424
3425fn axum_path(template: &str) -> String {
3426 let mut out = String::with_capacity(template.len());
3427 let mut rest = template;
3428 while let Some(open) = rest.find('{') {
3429 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3430 break;
3431 };
3432 out.push_str(&rest[..open]);
3433 out.push(':');
3434 out.push_str(&rest[open + 1..close]);
3435 rest = &rest[close + 1..];
3436 }
3437 out.push_str(rest);
3438 out
3439}
3440
3441async fn cancel_generation(
3450 State(state): State<Arc<AppState>>,
3451 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3452) -> Response {
3453 let cancelled = state.cancels.cancel(&req.request_id);
3454 let status = if cancelled {
3455 StatusCode::OK
3456 } else {
3457 StatusCode::NOT_FOUND
3458 };
3459 let detail = if cancelled {
3460 "the generation was asked to stop; it ends at its next token".to_string()
3461 } else {
3462 "no generation with that request_id is running -- it has already \
3463 finished, was never issued, or was served by a path that does \
3464 not register for cancellation"
3465 .to_string()
3466 };
3467 (
3468 status,
3469 Json(ferrox_api::CancelGenerationResponse {
3470 request_id: req.request_id,
3471 cancelled,
3472 detail,
3473 }),
3474 )
3475 .into_response()
3476}
3477
3478type Activated = (
3482 Loaded,
3483 Option<serving::batch::ContinuousBatcher>,
3484 Option<Arc<budget::ContextCeiling>>,
3485);
3486
3487fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3505 let mut batcher = serving::batch::BatcherConfig::from_env();
3506 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3507 return batcher;
3510 }
3511 let Some(path) = path else {
3512 return batcher;
3513 };
3514 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1);
3520 let Some((priced, gguf_ctx, source)) = priced else {
3521 return batcher;
3522 };
3523 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3524 tracing::warn!(
3528 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3529 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3530 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3531 admit on a number you choose.",
3532 priced.weights_bytes,
3533 priced.device_budget_bytes,
3534 );
3535 return batcher;
3536 };
3537 tracing::info!("{source}");
3538 tracing::info!("{}", derived.fit);
3539 let adopted = budget::apply_derived(&mut batcher, &derived);
3540 if adopted.max_context {
3541 tracing::info!(
3542 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3543 override with FERROX_CB_MAX_CONTEXT",
3544 derived.max_context
3545 );
3546 }
3547 if adopted.kv_blocks {
3548 tracing::info!(
3549 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3550 derived.kv_blocks,
3551 batcher.kv_block_size
3552 );
3553 }
3554 batcher
3555}
3556
3557pub(crate) fn activate_loaded_model(
3566 loaded: model::LoadedModel,
3567 enable_continuous_batching: bool,
3568 path: Option<&str>,
3569 paged_kv: Option<&generate::PagedKvConfig>,
3570) -> Activated {
3571 match loaded {
3572 model::LoadedModel::Gguf(g) => {
3573 let decoder = Arc::new(g.decoder);
3574 let tokenizer = Arc::new(g.tokenizer);
3575 let config = price_batcher_config(path);
3576 let shape =
3582 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32);
3583 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3584 let batcher = if enable_continuous_batching {
3585 tracing::info!(
3586 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3587 (stop sequences use the same pending-buffer trim as the private generate loop)"
3588 );
3589 let tok = Arc::clone(&tokenizer);
3590 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3591 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3592 Arc::clone(&decoder),
3593 decode,
3594 config,
3595 Arc::clone(&ceiling),
3596 paged_kv.cloned(),
3597 ))
3598 } else {
3599 None
3600 };
3601 (
3602 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3603 decoder,
3604 tokenizer,
3605 stop_tokens: g.stop_tokens,
3606 bos_id: g.bos_id,
3607 is_synthetic: g.is_synthetic,
3608 chat_template: g.chat_template,
3609 }))),
3610 batcher,
3611 Some(ceiling),
3612 )
3613 }
3614 model::LoadedModel::Kimi(k) => (
3615 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3616 engine: k.engine,
3617 tokenizer: k.tokenizer,
3618 stop_tokens: k.stop_tokens,
3619 chat_template: k.chat_template,
3620 }))),
3621 None,
3622 None,
3623 ),
3624 model::LoadedModel::Mla(m) => (
3625 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3626 engine: m.engine,
3627 tokenizer: m.tokenizer,
3628 stop_tokens: m.stop_tokens,
3629 bos_id: m.bos_id,
3630 name: m.name,
3631 chat_template: m.chat_template,
3632 }))),
3633 None,
3634 None,
3635 ),
3636 model::LoadedModel::Gemma4(m) => (
3637 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3638 engine: m.engine,
3639 tokenizer: m.tokenizer,
3640 stop_tokens: m.stop_tokens,
3641 bos_id: m.bos_id,
3642 name: m.name,
3643 chat_template: m.chat_template,
3644 }))),
3645 None,
3646 None,
3647 ),
3648 model::LoadedModel::Glm52(g) => (
3649 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3650 engine: g.engine,
3651 tokenizer: g.tokenizer,
3652 stop_tokens: g.stop_tokens,
3653 bos_id: g.bos_id,
3654 name: g.name,
3655 chat_template: g.chat_template,
3656 }))),
3657 None,
3658 None,
3659 ),
3660 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3665 }
3666}
3667
3668struct StartupModels {
3675 loaded: model::LoadedModel,
3676 embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
3677}
3678
3679fn continuous_batching_env() -> Option<bool> {
3680 match std::env::var("FERROX_CONTINUOUS_BATCHING")
3681 .ok()
3682 .map(|v| v.trim().to_ascii_lowercase())
3683 .as_deref()
3684 {
3685 None => None,
3686 Some("1" | "true" | "yes" | "on") => Some(true),
3687 Some("0" | "false" | "no" | "off") => Some(false),
3688 _ => None,
3689 }
3690}
3691
3692fn metal_private_decode_active() -> bool {
3693 #[cfg(feature = "metal")]
3694 {
3695 BUILT_WITH_METAL
3696 && ferrox_metal::attn::metal_attn_enabled()
3697 && std::env::var("FERROX_METAL").ok().as_deref() != Some("0")
3698 }
3699 #[cfg(not(feature = "metal"))]
3700 {
3701 false
3702 }
3703}
3704
3705fn continuous_batching_compatible(
3706 loaded: &model::LoadedModel,
3707 kv_pool: &Option<generate::KvPoolConfig>,
3708 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3709 paged_kv: &Option<generate::PagedKvConfig>,
3710) -> bool {
3711 matches!(loaded, model::LoadedModel::Gguf(_))
3712 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3713}
3714
3715fn resolve_continuous_batching_enabled(
3716 loaded: &model::LoadedModel,
3717 kv_pool: &Option<generate::KvPoolConfig>,
3718 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3719 paged_kv: &Option<generate::PagedKvConfig>,
3720) -> bool {
3721 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3722 return false;
3723 }
3724 match continuous_batching_env() {
3725 Some(true) => true,
3726 Some(false) => false,
3727 None => metal_private_decode_active(),
3728 }
3729}
3730
3731fn acquire_metal_private_decode_gate(
3732 gate: Option<&std::sync::Mutex<()>>,
3733 used_batcher: bool,
3734) -> Option<std::sync::MutexGuard<'_, ()>> {
3735 if used_batcher {
3736 None
3737 } else {
3738 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3739 }
3740}
3741
3742fn build_app_state(
3743 models: StartupModels,
3744 kv_pool: Option<generate::KvPoolConfig>,
3745 paged_kv: Option<generate::PagedKvConfig>,
3746 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3747 enable_continuous_batching: bool,
3748 mcp: Option<mcp::LoadedMcpConfig>,
3749 detection: Arc<health::Detection>,
3750) -> AppState {
3751 let StartupModels { loaded, embedding } = models;
3752 let configured_path = std::env::var("FERROX_MODEL_PATH").ok();
3753 let (loaded, batcher, ceiling) = activate_loaded_model(
3754 loaded,
3755 enable_continuous_batching,
3756 configured_path.as_deref(),
3757 paged_kv.as_ref(),
3758 );
3759 let id = startup_model_id();
3765 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3766 {
3767 None
3768 } else {
3769 tracing::info!(
3770 "Metal private-loop decode will serialize concurrent requests until \
3771 continuous batching is enabled (FERROX_CONTINUOUS_BATCHING=1 or --cont-batching)"
3772 );
3773 Some(Arc::new(std::sync::Mutex::new(())))
3774 };
3775 AppState {
3776 embedding,
3777 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3778 id,
3779 loaded,
3780 batcher,
3781 ceiling,
3782 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3783 }))),
3784 paged_kv,
3785 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3786 tasks: Arc::new(tasks::TaskRegistry::new()),
3787 cancels: Arc::new(cancel::CancelRegistry::new()),
3788 stats: stats::Stats::new(),
3789 streams: resume::StreamRegistry::new(),
3790 model_dir: admin::model_dirs().into_iter().next(),
3791 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3792 kv_pool,
3793 prefix_cache,
3794 sessions: session::SessionStore::new(),
3795 requests_total: std::sync::atomic::AtomicU64::new(0),
3796 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3797 started_at: std::time::Instant::now(),
3798 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3799 detection,
3800 mcp,
3801 continuous_batching_enabled: enable_continuous_batching,
3802 metal_private_decode_gate,
3803 loading_model: Mutex::new(None),
3804 last_load_error: Mutex::new(None),
3805 serving: Mutex::new(crate::stats::ServingStats::default()),
3806 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3807 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3808 started_unix: unix_now(),
3809 }
3810}
3811
3812fn load_embedding_model() -> anyhow::Result<Option<Arc<ferrox_models::EmbeddingModel>>> {
3820 let Ok(path) = std::env::var("FERROX_EMBEDDING_MODEL_PATH") else {
3821 return Ok(None);
3822 };
3823 let model = ferrox_models::EmbeddingModel::from_gguf_path(&path)
3824 .map_err(|e| anyhow::anyhow!("FERROX_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3825 tracing::info!(
3826 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3827 model.name(),
3828 model.architecture(),
3829 model.n_embd(),
3830 model.pooling_type().name(),
3831 model.n_ctx_train(),
3832 );
3833 Ok(Some(Arc::new(model)))
3834}
3835
3836fn unix_now() -> u64 {
3841 std::time::SystemTime::now()
3842 .duration_since(std::time::UNIX_EPOCH)
3843 .map(|d| d.as_secs())
3844 .unwrap_or(0)
3845}
3846
3847fn startup_model_id() -> Option<String> {
3852 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
3853 let configured = std::fs::canonicalize(&configured).ok()?;
3854 admin::discover(&admin::model_dirs())
3855 .into_iter()
3856 .find(|d| {
3857 std::fs::canonicalize(&d.path)
3858 .map(|p| p == configured)
3859 .unwrap_or(false)
3860 })
3861 .map(|d| d.id)
3862}
3863
3864fn init_cpu_pool() {
3873 match ferrox_core::threads::init_cpu_pool() {
3874 Some(n) => eprintln!(
3875 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
3876 ferrox_core::threads::perf_core_count()
3877 ),
3878 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
3879 }
3880}
3881
3882fn announce_ready(addr: SocketAddr, scheme: &str) {
3895 use std::io::Write;
3896 let ready =
3897 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3898 let mut stdout = std::io::stdout().lock();
3899 let _ = writeln!(stdout, "{}", ready.to_line());
3900 let _ = stdout.flush();
3901}
3902
3903async fn shutdown_signal(exit_on_stdin_close: bool) {
3915 if !exit_on_stdin_close {
3916 std::future::pending::<()>().await;
3917 return;
3918 }
3919 let _ = tokio::task::spawn_blocking(|| {
3920 use std::io::Read;
3921 let mut sink = [0u8; 256];
3922 let mut stdin = std::io::stdin().lock();
3923 loop {
3924 match stdin.read(&mut sink) {
3925 Ok(0) => break,
3927 Ok(_) => continue,
3929 Err(e) => {
3930 tracing::warn!("stdin read failed ({e}); treating it as closed");
3931 break;
3932 }
3933 }
3934 }
3935 })
3936 .await;
3937 tracing::info!("stdin closed; shutting down");
3938}
3939
3940fn tokio_worker_threads() -> usize {
3945 std::env::var("FERROX_TOKIO_WORKERS")
3946 .ok()
3947 .and_then(|v| v.trim().parse::<usize>().ok())
3948 .filter(|n| *n > 0)
3949 .unwrap_or(2)
3950}
3951
3952fn install_ring_crypto_provider() {
3969 let _ = rustls::crypto::ring::default_provider().install_default();
3970}
3971
3972pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3978 if args.list_devices {
3979 print_available_devices();
3980 return Ok(());
3981 }
3982 apply_cli_overrides(&args)?;
3983
3984 let _instance = {
3988 use ferrox_core::instance::{register, InstancePolicy};
3989 let policy = if args.allow_multiple_instances {
3990 InstancePolicy::Multi
3991 } else {
3992 InstancePolicy::from_env_or(InstancePolicy::Single)
3993 };
3994 let model = std::env::var("FERROX_MODEL_PATH").ok();
3995 register(
3996 "server",
3997 model.as_deref(),
3998 ferrox_core::instance::current_backend(),
3999 policy,
4000 )
4001 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4002 };
4003
4004 let journal = journal::Journal::from_env();
4005 eprintln!(
4006 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
4007 journal.path()
4008 );
4009 journal.append(&journal::Record::session_start(
4010 env!("CARGO_PKG_VERSION"),
4011 std::process::id(),
4012 ));
4013 journal::install_panic_hook(journal.clone());
4014
4015 let mcp_config_path = args.mcp_config.clone();
4016 let exit_on_stdin_close = args.exit_on_stdin_close
4017 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
4018 .map(|v| v == "1")
4019 .unwrap_or(false);
4020
4021 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
4025 init_cpu_pool();
4026
4027 let runtime = tokio::runtime::Builder::new_multi_thread()
4028 .worker_threads(tokio_worker_threads())
4029 .enable_all()
4030 .build()?;
4031 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4032
4033 let reason = match &result {
4034 Ok(()) => "normal".to_string(),
4035 Err(e) => e.to_string(),
4036 };
4037 journal.append(&journal::Record::session_exit(reason));
4038
4039 runtime.shutdown_background();
4044
4045 result
4046}
4047
4048async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4049 let _ = tracing_subscriber::fmt::try_init();
4055
4056 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4066 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
4067 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
4068 .map(|v| v == "1")
4069 .unwrap_or(false);
4070 if let Err(msg) =
4071 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4072 {
4073 anyhow::bail!(msg);
4074 }
4075
4076 let embedding_model = load_embedding_model()?;
4083
4084 let mut loaded = model::load()?;
4085 match &loaded {
4086 model::LoadedModel::Gguf(g) => tracing::info!(
4087 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4088 g.decoder.config.name,
4089 g.is_synthetic,
4090 g.tokenizer.kind()
4091 ),
4092 model::LoadedModel::Kimi(k) => tracing::info!(
4093 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4094 k.tokenizer.vocab_size()
4095 ),
4096 model::LoadedModel::Mla(m) => tracing::info!(
4097 "loaded MLA GGUF '{}' (tokenizer={})",
4098 m.name,
4099 m.tokenizer.kind()
4100 ),
4101 model::LoadedModel::Gemma4(m) => tracing::info!(
4102 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4103 m.name,
4104 m.tokenizer.kind()
4105 ),
4106 model::LoadedModel::Glm52(g) => tracing::info!(
4107 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4108 g.name,
4109 g.tokenizer.kind()
4110 ),
4111 model::LoadedModel::Encoder(_) => {}
4114 }
4115 let metal_default_moe_budget = {
4123 #[cfg(feature = "metal")]
4124 {
4125 ferrox_core::metal_dense_enabled()
4126 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
4127 }
4128 #[cfg(not(feature = "metal"))]
4129 {
4130 false
4131 }
4132 };
4133 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
4134 let budget: u64 = budget_str
4135 .parse()
4136 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4137 match &mut loaded {
4138 model::LoadedModel::Gguf(g) => {
4139 tracing::info!(
4140 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4141 (CUDA and/or Metal matvecs when built with the matching feature)"
4142 );
4143 g.decoder.gpu_vram_budget_bytes = Some(budget);
4144 }
4145 model::LoadedModel::Kimi(_) => {
4146 tracing::warn!(
4147 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4148 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4149 );
4150 }
4151 model::LoadedModel::Mla(_) => {
4152 tracing::warn!(
4153 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4154 FFN path only today; ignoring expert VRAM budget"
4155 );
4156 }
4157 model::LoadedModel::Gemma4(_) => {
4158 tracing::warn!(
4159 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4160 ignoring expert VRAM budget"
4161 );
4162 }
4163 model::LoadedModel::Glm52(_) => {
4164 tracing::warn!(
4165 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4166 GPU expert placement not wired yet; ignoring"
4167 );
4168 }
4169 model::LoadedModel::Encoder(_) => {
4170 tracing::warn!(
4171 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4172 it has no routed experts to place; ignoring"
4173 );
4174 }
4175 }
4176 } else if metal_default_moe_budget {
4177 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4181 if let model::LoadedModel::Gguf(g) = &mut loaded {
4182 tracing::info!(
4183 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4184 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4185 );
4186 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4187 }
4188 }
4189 #[cfg(feature = "cuda")]
4190 {
4191 if ferrox_core::cuda_dense_enabled() {
4192 tracing::info!(
4193 "CUDA dense matvec enabled for WeightMatrix::apply \
4194 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4195 );
4196 } else {
4197 tracing::info!(
4198 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
4199 );
4200 }
4201 }
4202 #[cfg(feature = "metal")]
4203 {
4204 if ferrox_core::metal_dense_enabled() {
4205 tracing::info!(
4206 "Metal dense matvec enabled for WeightMatrix::apply \
4207 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4208 );
4209 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
4210 Some("1") | Some("true") | Some("on") | Some("attn") => {
4211 tracing::info!(
4212 "Metal fused attention requested (FERROX_METAL_ATTN): \
4213 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4214 );
4215 }
4216 _ => {}
4217 }
4218 tracing::info!(
4219 "Metal greedy GPU argmax: temperature<=0 folds \
4220 final_norm+lm_head+argmax into the dense stack"
4221 );
4222 } else {
4223 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
4224 }
4225 }
4226 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
4237 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
4238 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
4239 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4240 panic!(
4241 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
4242 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
4243 );
4244 }
4245 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4246 (Ok(blocks), Ok(block_size), Err(_)) => {
4247 let total_blocks: usize = blocks
4248 .parse()
4249 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
4250 let block_size: usize = block_size
4251 .parse()
4252 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4253 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4259 .ok()
4260 .map(|v| {
4261 v.parse()
4262 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4263 })
4264 .unwrap_or(0);
4265 tracing::info!(
4266 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4267 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4268 );
4269 Some(generate::KvPoolConfig {
4270 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4271 queue_wait: Duration::from_millis(queue_wait_ms),
4272 })
4273 }
4274 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4275 let block_size: usize = block_size
4276 .parse()
4277 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4278 let budget: u64 = byte_budget
4279 .parse()
4280 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
4281 let cfg = match &loaded {
4282 model::LoadedModel::Gguf(g) => &g.decoder.config,
4283 model::LoadedModel::Kimi(_)
4284 | model::LoadedModel::Mla(_)
4285 | model::LoadedModel::Gemma4(_)
4286 | model::LoadedModel::Glm52(_)
4287 | model::LoadedModel::Encoder(_) => {
4288 panic!(
4289 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
4290 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4291 );
4292 }
4293 };
4294 let bytes_per_block = block_size
4295 * cfg.n_layers
4296 * cfg.n_kv_heads
4297 * cfg.head_dim
4298 * 2
4299 * std::mem::size_of::<f32>();
4300 assert!(
4301 bytes_per_block > 0,
4302 "derived KV block byte size must be positive (check model config and block size)"
4303 );
4304 let total_blocks = (budget as usize / bytes_per_block).max(1);
4305 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4306 .ok()
4307 .map(|v| {
4308 v.parse()
4309 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4310 })
4311 .unwrap_or(0);
4312 tracing::info!(
4313 "KV cache block pool enabled from byte budget: {budget} bytes / \
4314 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4315 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4316 cfg.n_layers
4317 );
4318 Some(generate::KvPoolConfig {
4319 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4320 queue_wait: Duration::from_millis(queue_wait_ms),
4321 })
4322 }
4323 (Err(_), Err(_), Err(_)) => None,
4324 (Err(_), Ok(_), Err(_)) => panic!(
4325 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4326 (or unset all three to disable KV cache pooling)"
4327 ),
4328 (Ok(_), Ok(_), Ok(_)) => {
4329 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4330 }
4331 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4332 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4333 set together (or neither, to disable KV cache pooling)"
4334 ),
4335 };
4336 let paged_kv = match (
4342 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4343 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4344 ) {
4345 (Ok(blocks), Ok(block_size)) => {
4346 assert!(
4347 kv_pool.is_none(),
4348 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4349 mutually exclusive: both bound the same KV memory, by different means. \
4350 Set one."
4351 );
4352 let blocks_per_layer: usize = blocks
4366 .parse()
4367 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4368 let block_size: usize = block_size
4369 .parse()
4370 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4371 let gguf = match &loaded {
4372 model::LoadedModel::Gguf(g) => g,
4373 _ => panic!(
4374 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4375 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4376 ),
4377 };
4378 let cfg = &gguf.decoder.config;
4379 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4380 .ok()
4381 .map(|v| {
4382 v.parse()
4383 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4384 })
4385 .unwrap_or(0);
4386 tracing::info!(
4387 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4388 layer across {} layers, shared by all concurrent requests, \
4389 {queue_wait_ms}ms admission queue wait",
4390 cfg.n_layers
4391 );
4392 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4397 block_size,
4398 ))));
4399 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4408 crate::policy::parser::ToolCallFormat::infer(
4409 &std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
4410 )
4411 .opener(),
4412 |text| {
4413 gguf.tokenizer
4414 .encode(text, SpecialTokens::Parse)
4415 .into_iter()
4416 .map(|t| t as u32)
4417 .collect()
4418 },
4419 );
4420 if let Some(id) = anchor_token {
4421 tracing::info!(
4422 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4423 window stops short of where its next turn rejoins"
4424 );
4425 }
4426 let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
4427 .ok()
4428 .map(|v| {
4429 v.parse()
4430 .expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4431 })
4432 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4433 if let Some(window) = cfg.uniform_sliding_window() {
4434 tracing::info!(
4435 "Paged KV window slide enabled: every layer slides by {window} every \
4436 {slide_interval} decode steps, so a request holds its prompt and a \
4437 window rather than its whole context"
4438 );
4439 } else if cfg.kv_block_window().is_some() {
4440 tracing::info!(
4441 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4442 and a page group holds one block in every layer"
4443 );
4444 }
4445 Some(generate::PagedKvConfig {
4446 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4447 cfg.n_layers,
4448 block_size,
4449 blocks_per_layer,
4450 cfg.n_kv_heads,
4451 cfg.head_dim,
4452 )),
4453 queue_wait: Duration::from_millis(queue_wait_ms),
4454 radix,
4455 anchor_token,
4456 slide_interval,
4457 })
4458 }
4459 (Err(_), Err(_)) => None,
4460 _ => panic!(
4461 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4462 (or neither, to disable paged KV)"
4463 ),
4464 };
4465 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4472 let max_entries: usize = v
4473 .parse()
4474 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4475 if kv_pool.is_some() {
4476 tracing::warn!(
4477 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4478 caching will never be consulted while a KV pool is configured"
4479 );
4480 }
4481 assert!(
4488 paged_kv.is_none(),
4489 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4490 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4491 produce, so the cache could never hit. Set one."
4492 );
4493 tracing::info!(
4494 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4495 all requests"
4496 );
4497 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4498 });
4499 if matches!(
4500 loaded,
4501 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4502 ) && (kv_pool.is_some() || prefix_cache.is_some())
4503 {
4504 tracing::warn!(
4505 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4506 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4507 ferrox_models::engine's module docs"
4508 );
4509 }
4510 let enable_cb =
4511 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4512 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4513 tracing::info!(
4514 "continuous batching enabled by default on Metal for safe parallel serving \
4515 (set FERROX_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4516 );
4517 }
4518 if continuous_batching_env() == Some(true)
4519 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4520 && (kv_pool.is_some() || prefix_cache.is_some())
4521 {
4522 tracing::warn!(
4523 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4524 (those modes keep the private generate path)"
4525 );
4526 }
4527 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4528 if let Ok(chunk) = n.parse::<usize>() {
4529 if chunk > 0 {
4530 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4531 }
4532 }
4533 }
4534 if matches!(
4535 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4536 Some("1")
4537 ) {
4538 tracing::warn!(
4539 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4540 (minimal spill; full layer offload still planned)"
4541 );
4542 }
4543
4544 let mcp = match mcp_config_path {
4545 Some(path) => {
4546 let loaded = mcp::load_mcp_config(&path)?;
4547 tracing::info!(
4548 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4549 loaded.path,
4550 loaded.servers.len()
4551 );
4552 Some(loaded)
4553 }
4554 None => None,
4555 };
4556
4557 let detection = health::Detection::spawn();
4561
4562 let state = Arc::new(build_app_state(
4563 StartupModels {
4564 loaded,
4565 embedding: embedding_model,
4566 },
4567 kv_pool,
4568 paged_kv,
4569 prefix_cache,
4570 enable_cb,
4571 mcp,
4572 detection,
4573 ));
4574
4575 use ferrox_api::routes;
4579
4580 let public = Router::new().route(routes::HEALTH, get(health));
4585
4586 let mut protected = protected_routes();
4587
4588 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4593 tracing::info!("API key auth enabled");
4594 let auth = limits::AuthConfig {
4595 api_key: Arc::new(key),
4596 };
4597 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4598 auth,
4599 limits::require_api_key,
4600 ));
4601 }
4602 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4603 let rpm: u32 = rpm
4604 .parse()
4605 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4606 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4607 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4608 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4609 limiter,
4610 limits::rate_limit,
4611 ));
4612 }
4613 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
4621 let origins = security::parse_cors_origins(&spec)
4622 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
4623 tracing::info!(
4624 "CORS enabled: {} allow-listed origin(s) ({})",
4625 origins.len(),
4626 spec
4627 );
4628 let cors = tower_http::cors::CorsLayer::new()
4629 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4630 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4631 .allow_headers([
4632 axum::http::header::CONTENT_TYPE,
4633 axum::http::header::AUTHORIZATION,
4634 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4640 axum::http::HeaderName::from_static("last-event-id"),
4644 ]);
4645 protected = protected.route_layer(cors);
4646 }
4647
4648 let app = public
4652 .merge(protected)
4653 .layer(axum::middleware::from_fn(limits::retry_after))
4654 .with_state(state);
4655
4656 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4662 install_ring_crypto_provider();
4663 match tls_paths {
4668 Some(paths) => {
4669 let config =
4670 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4671 .await
4672 .map_err(|e| {
4673 anyhow::anyhow!(
4674 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4675 paths.cert,
4676 paths.key
4677 )
4678 })?;
4679 let socket_addr: std::net::SocketAddr = addr
4680 .parse()
4681 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
4682 let listener = std::net::TcpListener::bind(socket_addr)?;
4683 listener.set_nonblocking(true)?;
4691 let bound = listener.local_addr()?;
4692 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
4693 announce_ready(bound, "https");
4694
4695 let handle = axum_server::Handle::new();
4696 let shutdown_handle = handle.clone();
4697 tokio::spawn(async move {
4698 shutdown_signal(exit_on_stdin_close).await;
4699 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4700 });
4701 axum_server::from_tcp_rustls(listener, config)?
4702 .handle(handle)
4703 .serve(app.into_make_service())
4704 .await?;
4705 }
4706 None => {
4707 let listener = tokio::net::TcpListener::bind(&addr).await?;
4708 let bound = listener.local_addr()?;
4709 tracing::info!("ferrox-server listening on {bound}");
4710 announce_ready(bound, "http");
4711 axum::serve(listener, app)
4712 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4713 .await?;
4714 }
4715 }
4716 Ok(())
4717}
4718
4719#[cfg(test)]
4720mod tests {
4721 use super::*;
4722 use ferrox_models::config::test_dense_fixture;
4723
4724 #[test]
4725 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4726 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4727 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4728 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
4729 assert_eq!(parsed.port, 51999);
4730 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4731 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
4733 }
4734
4735 fn test_model() -> Model {
4736 let cfg = test_dense_fixture();
4740 Model::Gguf(GgufModel {
4741 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4742 tokenizer: Arc::new(ServerTokenizer::Byte),
4743 stop_tokens: StopTokens::default(),
4744 bos_id: None,
4745 is_synthetic: true,
4746 chat_template: chat_template::PromptTemplate::plain(),
4747 })
4748 }
4749
4750 fn greedy_params(max_tokens: usize) -> GenerationParams {
4751 GenerationParams {
4752 reasoning: None,
4753 max_tokens,
4754 sampling: SamplingParams::default(),
4755 seed: 1,
4756 stop: Vec::new(),
4757 stop_token_ids: Vec::new(),
4758 json_object: false,
4759 grammar: None,
4760 cancel: None,
4761 ignore_eos: false,
4762 }
4763 }
4764
4765 fn test_model_full_byte_vocab() -> Model {
4769 test_model_full_byte_vocab_with_eos(None)
4770 }
4771
4772 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4779 let mut cfg = test_dense_fixture();
4780 cfg.vocab_size = 256;
4781 Model::Gguf(GgufModel {
4782 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4783 tokenizer: Arc::new(ServerTokenizer::Byte),
4784 stop_tokens: StopTokens::from_eos(eos),
4785 bos_id: None,
4786 is_synthetic: true,
4787 chat_template: chat_template::PromptTemplate::plain(),
4788 })
4789 }
4790
4791 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4795 AppState {
4796 embedding: None,
4797 paged_kv: None,
4798 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4799 id: None,
4800 loaded: Loaded::Generative(Arc::new(model)),
4801 batcher: None,
4802 ceiling: None,
4803 checkpoint_path: None,
4804 }))),
4805 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4806 tasks: Arc::new(tasks::TaskRegistry::new()),
4807 cancels: Arc::new(cancel::CancelRegistry::new()),
4808 stats: stats::Stats::new(),
4809 streams: resume::StreamRegistry::new(),
4810 model_dir: None,
4811 response_cache: Mutex::new(response_cache),
4812 kv_pool: None,
4813 prefix_cache: None,
4814 sessions: session::SessionStore::new(),
4815 requests_total: std::sync::atomic::AtomicU64::new(0),
4816 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4817 started_at: std::time::Instant::now(),
4818 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4819 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4820 mcp: None,
4821 continuous_batching_enabled: false,
4822 metal_private_decode_gate: None,
4823 loading_model: Mutex::new(None),
4824 last_load_error: Mutex::new(None),
4825 serving: Mutex::new(crate::stats::ServingStats::default()),
4826 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4827 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4828 started_unix: unix_now(),
4829 }
4830 }
4831
4832 fn test_app() -> Router {
4841 test_app_with_state(Arc::new(test_state(
4842 test_model_full_byte_vocab(),
4843 ResponseCache::new(1000, Duration::from_secs(3600)),
4844 )))
4845 }
4846
4847 fn test_app_with_state(state: Arc<AppState>) -> Router {
4850 Router::new()
4860 .route(ferrox_api::routes::HEALTH, get(health))
4861 .merge(protected_routes())
4862 .with_state(state)
4863 }
4864
4865 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4866 let mut cfg = test_dense_fixture();
4867 cfg.name = name;
4868 cfg.vocab_size = vocab_size;
4869 Model::Gguf(GgufModel {
4870 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4871 tokenizer: Arc::new(ServerTokenizer::Byte),
4872 stop_tokens: StopTokens::default(),
4873 bos_id: None,
4874 is_synthetic: true,
4875 chat_template: chat_template::PromptTemplate::plain(),
4876 })
4877 }
4878
4879 fn model_with_template(name: &'static str, source: &str) -> Model {
4883 let mut cfg = test_dense_fixture();
4884 cfg.name = name;
4885 cfg.vocab_size = 256;
4886 Model::Gguf(GgufModel {
4887 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4888 tokenizer: Arc::new(ServerTokenizer::Byte),
4889 stop_tokens: StopTokens::default(),
4890 bos_id: None,
4891 is_synthetic: true,
4892 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4893 Some(source),
4894 Some("qwen3"),
4895 false,
4896 true,
4897 None,
4898 None,
4899 ),
4900 })
4901 }
4902
4903 #[tokio::test]
4912 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4913 let strict = "{% if messages | length > 1 %}\
4916 {{ raise_exception('this template takes one turn') }}\
4917 {% endif %}{{ messages[0].content }}";
4918 let state = Arc::new(test_state(
4919 model_with_template("strict", strict),
4920 ResponseCache::new(4, Duration::from_secs(60)),
4921 ));
4922 let app = test_app_with_state(state);
4923
4924 let (status, body) = post_json_uri(
4925 &app,
4926 "/v1/chat/completions",
4927 serde_json::json!({
4928 "model": "strict",
4929 "stream": true,
4930 "messages": [
4931 {"role": "user", "content": "one"},
4932 {"role": "user", "content": "two"},
4933 ],
4934 }),
4935 )
4936 .await;
4937 assert_eq!(status, StatusCode::BAD_REQUEST);
4938 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4939 assert!(
4940 body["error"]["message"]
4941 .as_str()
4942 .unwrap()
4943 .contains("one turn"),
4944 "the template's own message must reach the caller: {body}"
4945 );
4946
4947 let (status, _) = post_json_uri(
4949 &app,
4950 "/v1/chat/completions",
4951 serde_json::json!({
4952 "model": "strict",
4953 "stream": true,
4954 "max_tokens": 1,
4955 "messages": [{"role": "user", "content": "one"}],
4956 }),
4957 )
4958 .await;
4959 assert_eq!(status, StatusCode::OK);
4960 }
4961
4962 #[tokio::test]
4964 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4965 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4966 {% if reasoning_effort %}\
4967 {% if reasoning_effort not in ['low','medium','high'] %}\
4968 {{ raise_exception('bad effort') }}\
4969 {% endif %}[{{ reasoning_effort }}]\
4970 {% endif %}{{ messages[0].content }}";
4971 let state = Arc::new(test_state(
4972 model_with_template("thinker", reasoning),
4973 ResponseCache::new(4, Duration::from_secs(60)),
4974 ));
4975 let app = test_app_with_state(state);
4976 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4977 assert_eq!(status, StatusCode::OK);
4978 let entry = &models["data"][0];
4979 assert_eq!(
4980 entry["supported_reasoning_efforts"],
4981 serde_json::json!(["off", "low", "medium", "high"])
4982 );
4983 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4984 }
4985
4986 #[tokio::test]
4991 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4992 let app = test_app();
4993 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4994 let entry = &models["data"][0];
4995 assert!(entry.get("supported_reasoning_efforts").is_none());
4996 assert!(entry.get("default_reasoning_effort").is_none());
4997 }
4998
4999 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5000 Arc::new(ActiveModel {
5001 id: Some(name.to_string()),
5002 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5003 batcher: None,
5004 ceiling: None,
5005 checkpoint_path: None,
5006 })
5007 .tap_into(state)
5008 }
5009
5010 trait TapInto {
5012 fn tap_into(self, state: &AppState) -> Self;
5013 }
5014 impl TapInto for Arc<ActiveModel> {
5015 fn tap_into(self, state: &AppState) -> Self {
5016 state.swap_active(Some(Arc::clone(&self)));
5017 self
5018 }
5019 }
5020
5021 #[test]
5027 fn an_in_flight_request_keeps_the_model_it_started_on() {
5028 let state = test_state(
5029 named_test_model("model-a", 256),
5030 ResponseCache::new(4, Duration::from_secs(60)),
5031 );
5032
5033 let in_flight = state.active().expect("a model is loaded");
5036 assert_eq!(in_flight.name(), "model-a");
5037
5038 active_model(&state, "model-b");
5039
5040 assert_eq!(state.active().unwrap().name(), "model-b");
5042 assert_eq!(in_flight.name(), "model-a");
5044 let (_chunks, finish, _usage) = run_generation(
5045 in_flight.generative().unwrap(),
5046 "hi",
5047 &greedy_params(3),
5048 None,
5049 None,
5050 None,
5051 None,
5052 None,
5053 None,
5054 )
5055 .expect("the old model must still decode after being swapped out");
5056 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5057 }
5058
5059 #[test]
5064 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5065 let state = test_state(
5066 named_test_model("model-a", 256),
5067 ResponseCache::new(4, Duration::from_secs(60)),
5068 );
5069 let in_flight = state.active().expect("a model is loaded");
5070 let weights = Arc::clone(in_flight.generative().unwrap());
5071 assert!(Arc::strong_count(&weights) >= 2);
5072
5073 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5074 id: Some("model-b".to_string()),
5075 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5076 batcher: None,
5077 ceiling: None,
5078 checkpoint_path: None,
5079 })));
5080 drop(previous);
5081 assert!(Arc::strong_count(&weights) >= 2);
5083 drop(in_flight);
5084 assert_eq!(Arc::strong_count(&weights), 1);
5085 }
5086
5087 #[tokio::test]
5091 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5092 let state = Arc::new(test_state(
5093 named_test_model("model-a", 256),
5094 ResponseCache::new(4, Duration::from_secs(60)),
5095 ));
5096 let app = test_app_with_state(Arc::clone(&state));
5097
5098 let (status, body) = post_json_uri(
5099 &app,
5100 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5101 serde_json::json!({}),
5102 )
5103 .await;
5104 assert_eq!(status, StatusCode::OK);
5105 assert_eq!(body["ok"], true);
5106 assert!(body["active"].is_null());
5107 assert!(state.active().is_none());
5108
5109 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5110 assert_eq!(status, StatusCode::OK);
5111 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5112 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5113
5114 let (status, body) = post_json_uri(
5115 &app,
5116 "/v1/chat/completions",
5117 serde_json::json!({
5118 "model": "x",
5119 "messages": [{"role": "user", "content": "hi"}]
5120 }),
5121 )
5122 .await;
5123 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5124 assert_eq!(body["error"]["type"], "model_not_loaded");
5125 }
5126
5127 #[tokio::test]
5131 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5132 let state = Arc::new(test_state(
5133 named_test_model("model-a", 256),
5134 ResponseCache::new(4, Duration::from_secs(60)),
5135 ));
5136 let app = test_app_with_state(Arc::clone(&state));
5137 state.swap_active(None);
5138
5139 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5140 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5143 assert_eq!(body["state"], "unavailable");
5144 assert_eq!(body["reason"], "model_not_loaded");
5145 assert!(body["model"].is_null());
5146 let real_weights = body["capabilities"]
5147 .as_array()
5148 .unwrap()
5149 .iter()
5150 .find(|c| c["id"] == "real_weights")
5151 .cloned()
5152 .expect("real_weights is always reported");
5153 assert_eq!(real_weights["available"], false);
5154 assert_eq!(real_weights["reason"], "model_not_loaded");
5155 }
5156
5157 #[tokio::test]
5161 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5162 let app = test_app();
5163
5164 let (status, completion) = post_json_uri(
5165 &app,
5166 "/v1/chat/completions",
5167 serde_json::json!({
5168 "model": "x",
5169 "messages": [{"role": "user", "content": "hi"}],
5170 "max_tokens": 4
5171 }),
5172 )
5173 .await;
5174 assert_eq!(status, StatusCode::OK);
5175 let request_id = completion["request_id"].as_str().unwrap().to_string();
5176
5177 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5178 assert_eq!(status, StatusCode::OK);
5179 let recent = stats["recent"].as_array().unwrap();
5180 assert_eq!(recent.len(), 1);
5181 let row = &recent[0];
5182 assert_eq!(row["request_id"], request_id);
5183 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
5184 assert_eq!(row["status"], 200);
5185 assert_eq!(row["stream"], false);
5186 assert!(row["duration_ms"].is_number());
5189 assert!(row["decode_ms"].is_number());
5190 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5191 assert_eq!(
5192 stats["tokens_prompt_total"].as_u64().unwrap(),
5193 row["prompt_tokens"].as_u64().unwrap()
5194 );
5195 }
5196
5197 #[tokio::test]
5201 async fn a_rejected_request_is_recorded_too() {
5202 let state = Arc::new(test_state(
5203 named_test_model("model-a", 256),
5204 ResponseCache::new(4, Duration::from_secs(60)),
5205 ));
5206 let app = test_app_with_state(Arc::clone(&state));
5207 state.swap_active(None);
5208
5209 let (status, _) = post_json_uri(
5210 &app,
5211 "/v1/chat/completions",
5212 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5213 )
5214 .await;
5215 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5216
5217 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5218 let recent = stats["recent"].as_array().unwrap();
5219 assert_eq!(recent.len(), 1);
5220 assert_eq!(recent[0]["status"], 503);
5221 assert_eq!(recent[0]["completion_tokens"], 0);
5222 assert!(recent[0]["decode_ms"].is_null());
5223 assert_eq!(stats["errors_total"], 1);
5224 }
5225
5226 async fn post_json_with_headers(
5230 app: &Router,
5231 uri: &str,
5232 body: serde_json::Value,
5233 headers: &[(&str, &str)],
5234 ) -> (StatusCode, serde_json::Value) {
5235 use http_body_util::BodyExt;
5236 use tower::ServiceExt;
5237
5238 let mut builder = axum::http::Request::builder()
5239 .method("POST")
5240 .uri(uri)
5241 .header("content-type", "application/json");
5242 for (name, value) in headers {
5243 builder = builder.header(*name, *value);
5244 }
5245 let response = app
5246 .clone()
5247 .oneshot(
5248 builder
5249 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5250 .unwrap(),
5251 )
5252 .await
5253 .unwrap();
5254 let status = response.status();
5255 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5256 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5257 (status, json)
5258 }
5259
5260 #[tokio::test]
5264 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5265 let app = test_app();
5266
5267 let (status, _) = post_json_uri(
5268 &app,
5269 ferrox_api::routes::V1_TOKENIZE,
5270 serde_json::json!({"prompt": "hello"}),
5271 )
5272 .await;
5273 assert_eq!(status, StatusCode::OK);
5274 let (status, _) = post_json_uri(
5275 &app,
5276 ferrox_api::routes::V1_DETOKENIZE,
5277 serde_json::json!({"tokens": [104, 105]}),
5278 )
5279 .await;
5280 assert_eq!(status, StatusCode::OK);
5281 let (status, _) = post_json_uri(
5282 &app,
5283 ferrox_api::routes::V1_EMBEDDINGS,
5284 serde_json::json!({"input": "hello"}),
5285 )
5286 .await;
5287 assert_eq!(status, StatusCode::OK);
5288
5289 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5290 let routes: Vec<&str> = stats["recent"]
5291 .as_array()
5292 .unwrap()
5293 .iter()
5294 .map(|row| row["route"].as_str().unwrap())
5295 .collect();
5296 for expected in [
5297 ferrox_api::routes::V1_TOKENIZE,
5298 ferrox_api::routes::V1_DETOKENIZE,
5299 ferrox_api::routes::V1_EMBEDDINGS,
5300 ] {
5301 assert!(
5302 routes.contains(&expected),
5303 "{expected} is missing: {routes:?}"
5304 );
5305 }
5306
5307 let row = |route: &str| {
5308 stats["recent"]
5309 .as_array()
5310 .unwrap()
5311 .iter()
5312 .find(|r| r["route"] == route)
5313 .cloned()
5314 .unwrap()
5315 };
5316 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5320 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5321 assert!(embed["decode_ms"].is_null());
5322 assert_eq!(embed["completion_tokens"], 0);
5323 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5327 assert_eq!(
5328 stats["tokens_prompt_total"].as_u64().unwrap(),
5329 embed["prompt_tokens"].as_u64().unwrap(),
5330 "only the forward pass counted"
5331 );
5332 }
5333
5334 fn streaming_test_app() -> Router {
5339 let mut cfg = test_dense_fixture();
5340 cfg.vocab_size = 256;
5341 let model = Model::Gguf(GgufModel {
5342 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5343 tokenizer: Arc::new(ServerTokenizer::Byte),
5344 stop_tokens: StopTokens::default(),
5345 bos_id: None,
5346 is_synthetic: false,
5347 chat_template: chat_template::PromptTemplate::plain(),
5348 });
5349 test_app_with_state(Arc::new(test_state(
5350 model,
5351 ResponseCache::new(1000, Duration::from_secs(3600)),
5352 )))
5353 }
5354
5355 #[tokio::test]
5359 async fn the_native_completion_wire_is_not_the_openai_one() {
5360 let app = test_app();
5361
5362 let (status, native) = post_json_uri(
5363 &app,
5364 ferrox_api::routes::COMPLETION,
5365 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5366 )
5367 .await;
5368 assert_eq!(status, StatusCode::OK, "{native}");
5369 assert!(native["content"].is_string(), "{native}");
5370 assert_eq!(native["stop"], true);
5371 assert_eq!(native["stop_type"], "limit");
5372 assert_eq!(native["stopping_word"], "");
5373 assert_eq!(native["truncated"], false);
5374 assert_eq!(native["id_slot"], -1);
5375 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5376 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5377 assert!(
5378 native.get("choices").is_none(),
5379 "the native shape has no `choices`: {native}"
5380 );
5381
5382 let (status, openai) = post_json_uri(
5383 &app,
5384 ferrox_api::routes::V1_COMPLETIONS,
5385 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5386 )
5387 .await;
5388 assert_eq!(status, StatusCode::OK);
5389 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5390 assert!(
5391 openai.get("content").is_none(),
5392 "the OpenAI shape has no top-level `content`: {openai}"
5393 );
5394 }
5395
5396 #[tokio::test]
5400 async fn both_native_spellings_reach_the_same_handler() {
5401 let app = test_app();
5402 for route in [
5403 ferrox_api::routes::COMPLETION,
5404 ferrox_api::routes::COMPLETIONS,
5405 ] {
5406 let (status, body) = post_json_uri(
5407 &app,
5408 route,
5409 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5410 )
5411 .await;
5412 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5413 assert_eq!(body["stop"], true, "{route}");
5414 assert!(body["content"].is_string(), "{route}");
5415 }
5416
5417 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5420 let routes: Vec<&str> = stats["recent"]
5421 .as_array()
5422 .unwrap()
5423 .iter()
5424 .map(|row| row["route"].as_str().unwrap())
5425 .collect();
5426 assert!(
5427 routes.contains(&ferrox_api::routes::COMPLETION),
5428 "{routes:?}"
5429 );
5430 assert!(
5431 routes.contains(&ferrox_api::routes::COMPLETIONS),
5432 "{routes:?}"
5433 );
5434 }
5435
5436 #[tokio::test]
5442 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5443 let app = streaming_test_app();
5444 let raw = post_sse_raw_uri(
5445 &app,
5446 ferrox_api::routes::COMPLETION,
5447 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5448 )
5449 .await;
5450
5451 assert!(
5452 !raw.contains("[DONE]"),
5453 "llama.cpp's native stream has no sentinel: {raw}"
5454 );
5455 let frames: Vec<serde_json::Value> = raw
5456 .lines()
5457 .filter_map(|line| line.strip_prefix("data: "))
5458 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5459 .collect();
5460 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5461
5462 let (last, partials) = frames.split_last().unwrap();
5463 assert_eq!(last["stop"], true, "the last frame closes the stream");
5464 assert!(last["timings"].is_object(), "{last}");
5465 assert!(last["stop_type"].is_string(), "{last}");
5466 for partial in partials {
5467 assert_eq!(partial["stop"], false, "{partial}");
5468 assert!(partial["content"].is_string(), "{partial}");
5469 assert!(partial.get("timings").is_none(), "{partial}");
5473 assert!(partial.get("generation_settings").is_none(), "{partial}");
5474 }
5475 let streamed: String = partials
5478 .iter()
5479 .filter_map(|p| p["content"].as_str())
5480 .collect();
5481 assert_eq!(last["content"].as_str().unwrap(), streamed);
5482 }
5483
5484 #[tokio::test]
5489 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5490 let app = test_app();
5491 for body in [
5492 serde_json::json!({"prompt": "hi"}),
5493 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5494 ] {
5495 let (status, refusal) =
5496 post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
5497 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5498 assert!(
5499 refusal["error"]["message"]
5500 .as_str()
5501 .unwrap()
5502 .contains("n_predict"),
5503 "{refusal}"
5504 );
5505 }
5506 let (status, _) = post_json_uri(
5509 &app,
5510 ferrox_api::routes::COMPLETION,
5511 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5512 )
5513 .await;
5514 assert_eq!(status, StatusCode::OK);
5515 }
5516
5517 #[tokio::test]
5526 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5527 let app = streaming_test_app();
5528 let ask = |stop: serde_json::Value| {
5529 let app = app.clone();
5530 async move {
5531 post_json_uri(
5532 &app,
5533 ferrox_api::routes::COMPLETION,
5534 serde_json::json!({
5535 "prompt": "hi",
5536 "n_predict": 64,
5537 "ignore_eos": true,
5538 "stop": stop,
5539 }),
5540 )
5541 .await
5542 .1
5543 }
5544 };
5545
5546 let baseline = ask(serde_json::json!([])).await;
5547 assert_eq!(baseline["stop_type"], "limit");
5548 assert_eq!(baseline["stopping_word"], "");
5549 let text = baseline["content"].as_str().unwrap().to_string();
5550 let sentinel: String = text.chars().skip(1).take(2).collect();
5556 assert_eq!(
5557 sentinel.chars().count(),
5558 2,
5559 "the fixture must produce enough output to cut: {text:?}"
5560 );
5561 let cut = text.find(&sentinel).expect("it came out of this text");
5562
5563 let stopped = ask(serde_json::json!([sentinel])).await;
5564 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5565 assert_eq!(stopped["stopping_word"], sentinel);
5566 assert_eq!(
5567 stopped["content"].as_str().unwrap(),
5568 &text[..cut],
5569 "the answer must be cut at the sentinel, not run past it"
5570 );
5571 }
5572
5573 #[tokio::test]
5579 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5580 let app = test_app();
5581
5582 let (v1_status, v1) = post_json_uri(
5583 &app,
5584 ferrox_api::routes::V1_TOKENIZE,
5585 serde_json::json!({"prompt": "hello"}),
5586 )
5587 .await;
5588 let (alias_status, alias) = post_json_uri(
5589 &app,
5590 ferrox_api::routes::TOKENIZE,
5591 serde_json::json!({"content": "hello"}),
5592 )
5593 .await;
5594 assert_eq!(v1_status, StatusCode::OK);
5595 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5596 assert_eq!(v1["tokens"], alias["tokens"]);
5597 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5598
5599 let (status, both_ways) = post_json_uri(
5602 &app,
5603 ferrox_api::routes::TOKENIZE,
5604 serde_json::json!({"prompt": "hello"}),
5605 )
5606 .await;
5607 assert_eq!(status, StatusCode::OK);
5608 assert_eq!(both_ways["tokens"], v1["tokens"]);
5609 }
5610
5611 #[tokio::test]
5616 async fn detokenize_answers_under_both_dialects_keys() {
5617 let app = test_app();
5618 for route in [
5619 ferrox_api::routes::DETOKENIZE,
5620 ferrox_api::routes::V1_DETOKENIZE,
5621 ] {
5622 let (status, body) =
5623 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5624 assert_eq!(status, StatusCode::OK, "{route}");
5625 assert_eq!(body["text"], "hi", "{route}");
5626 assert_eq!(body["content"], body["text"], "{route}");
5627 }
5628 }
5629
5630 #[tokio::test]
5634 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5635 let app = test_app();
5636 let (status, _) = post_json_uri(
5637 &app,
5638 ferrox_api::routes::TOKENIZE,
5639 serde_json::json!({"content": "hello"}),
5640 )
5641 .await;
5642 assert_eq!(status, StatusCode::OK);
5643
5644 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5645 let routes: Vec<&str> = stats["recent"]
5646 .as_array()
5647 .unwrap()
5648 .iter()
5649 .map(|row| row["route"].as_str().unwrap())
5650 .collect();
5651 assert!(
5652 routes.contains(&ferrox_api::routes::TOKENIZE),
5653 "the alias must be its own row: {routes:?}"
5654 );
5655 assert!(
5656 !routes.contains(&ferrox_api::routes::V1_TOKENIZE),
5657 "nothing called /v1/tokenize: {routes:?}"
5658 );
5659 }
5660
5661 #[tokio::test]
5666 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5667 let mut cfg = test_dense_fixture();
5668 cfg.vocab_size = 256;
5669 let model = Model::Gguf(GgufModel {
5670 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5671 tokenizer: Arc::new(ServerTokenizer::Byte),
5672 stop_tokens: StopTokens::default(),
5673 bos_id: Some(7),
5674 is_synthetic: true,
5675 chat_template: chat_template::PromptTemplate::plain(),
5676 });
5677 let app = test_app_with_state(Arc::new(test_state(
5678 model,
5679 ResponseCache::new(1000, Duration::from_secs(3600)),
5680 )));
5681
5682 let (_, plain) = post_json_uri(
5683 &app,
5684 ferrox_api::routes::TOKENIZE,
5685 serde_json::json!({"content": "hi"}),
5686 )
5687 .await;
5688 let (_, special) = post_json_uri(
5689 &app,
5690 ferrox_api::routes::TOKENIZE,
5691 serde_json::json!({"content": "hi", "add_special": true}),
5692 )
5693 .await;
5694
5695 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5696 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5697 assert_eq!(special["count"], 3);
5698 }
5699
5700 #[tokio::test]
5703 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5704 let app = test_app();
5705 let (status, _) = post_json_uri(
5706 &app,
5707 ferrox_api::routes::V1_EMBEDDINGS,
5708 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5709 )
5710 .await;
5711 assert_eq!(status, StatusCode::BAD_REQUEST);
5712
5713 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5714 let recent = stats["recent"].as_array().unwrap();
5715 assert_eq!(recent.len(), 1);
5716 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
5717 assert_eq!(recent[0]["status"], 400);
5718 assert_eq!(
5719 recent[0]["prompt_tokens"], 0,
5720 "a rejected call embedded nothing"
5721 );
5722 }
5723
5724 #[tokio::test]
5727 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5728 let app = test_app();
5729 let key = "sk-monitor-secret";
5730 let (status, _) = post_json_with_headers(
5731 &app,
5732 "/v1/chat/completions",
5733 serde_json::json!({
5734 "model": "x",
5735 "messages": [{"role": "user", "content": "hi"}],
5736 "max_tokens": 2
5737 }),
5738 &[
5739 ("authorization", &format!("Bearer {key}")),
5740 ("x-ferrox-client", "ferrox-studio"),
5741 ],
5742 )
5743 .await;
5744 assert_eq!(status, StatusCode::OK);
5745
5746 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5747 let row = stats["recent"].as_array().unwrap()[0].clone();
5748 let fingerprint = row["via_api_key"]
5749 .as_str()
5750 .expect("the row names the key that served it")
5751 .to_string();
5752 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5753 assert!(!fingerprint.contains(key));
5754 assert!(
5755 !serde_json::to_string(&stats).unwrap().contains(key),
5756 "the stats payload must not carry the key in any form"
5757 );
5758 assert_eq!(row["client"], "ferrox-studio");
5759 }
5760
5761 #[tokio::test]
5764 async fn different_keys_are_different_callers_and_no_key_is_null() {
5765 let app = test_app();
5766 let body = serde_json::json!({
5767 "model": "x",
5768 "messages": [{"role": "user", "content": "hi"}],
5769 "max_tokens": 1
5770 });
5771 for headers in [
5772 vec![("authorization", "Bearer key-one")],
5773 vec![("authorization", "Bearer key-two")],
5774 vec![],
5775 ] {
5776 let (status, _) =
5777 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5778 assert_eq!(status, StatusCode::OK);
5779 }
5780
5781 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5782 let recent = stats["recent"].as_array().unwrap();
5783 assert_eq!(recent.len(), 3);
5784 let one = recent[0]["via_api_key"].as_str().unwrap();
5785 let two = recent[1]["via_api_key"].as_str().unwrap();
5786 assert_ne!(one, two, "two keys must not collapse into one caller");
5787 assert!(
5788 recent[2]["via_api_key"].is_null(),
5789 "an unauthenticated call is null, not a fingerprint of nothing"
5790 );
5791 assert!(recent[2]["client"].is_null());
5792 }
5793
5794 #[tokio::test]
5799 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5800 let state = Arc::new(test_state(
5801 named_test_model("really-loaded", 256),
5802 ResponseCache::new(4, Duration::from_secs(60)),
5803 ));
5804 let app = test_app_with_state(Arc::clone(&state));
5805
5806 let (status, _) = post_json_uri(
5807 &app,
5808 "/v1/chat/completions",
5809 serde_json::json!({
5810 "model": "gpt-4-turbo-that-is-not-here",
5811 "messages": [{"role": "user", "content": "hi"}],
5812 "max_tokens": 2
5813 }),
5814 )
5815 .await;
5816 assert_eq!(status, StatusCode::OK);
5817
5818 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5819 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5820
5821 state.swap_active(None);
5824 let (status, _) = post_json_uri(
5825 &app,
5826 "/v1/chat/completions",
5827 serde_json::json!({
5828 "model": "gpt-4-turbo-that-is-not-here",
5829 "messages": [{"role": "user", "content": "hi"}]
5830 }),
5831 )
5832 .await;
5833 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5834 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5835 let recent = stats["recent"].as_array().unwrap();
5836 assert!(recent[recent.len() - 1]["model"].is_null());
5837 }
5838
5839 #[tokio::test]
5843 async fn a_streamed_row_names_the_model_it_decoded_against() {
5844 let state = Arc::new(test_state(
5845 named_test_model("model-before", 256),
5846 ResponseCache::new(4, Duration::from_secs(60)),
5847 ));
5848 let app = test_app_with_state(Arc::clone(&state));
5849 let _ = post_sse_raw(&app, resumable_request()).await;
5850 active_model(&state, "model-after");
5852
5853 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5854 assert_eq!(stats["recent"][0]["model"], "model-before");
5855 }
5856
5857 #[tokio::test]
5860 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5861 let app = test_app();
5862 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5863 assert_eq!(status, StatusCode::OK);
5864 assert!(
5865 stats["queue_depth"].is_null(),
5866 "without continuous batching nothing queues, so there is nothing to measure"
5867 );
5868 assert!(stats["queue_rejected_total"].is_null());
5869 assert_eq!(
5870 stats["generating_now"], 0,
5871 "work in progress is measured and really is zero here"
5872 );
5873 }
5874
5875 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5880 post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
5881 }
5882
5883 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
5887 use http_body_util::BodyExt;
5888 use tower::ServiceExt;
5889
5890 let response = app
5891 .clone()
5892 .oneshot(
5893 axum::http::Request::builder()
5894 .method("POST")
5895 .uri(uri)
5896 .header("content-type", "application/json")
5897 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5898 .unwrap(),
5899 )
5900 .await
5901 .unwrap();
5902 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5903 String::from_utf8(bytes.to_vec()).unwrap()
5904 }
5905
5906 async fn get_json_with_headers(
5907 app: &Router,
5908 uri: &str,
5909 headers: &[(&str, &str)],
5910 ) -> (StatusCode, serde_json::Value) {
5911 use http_body_util::BodyExt;
5912 use tower::ServiceExt;
5913
5914 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
5915 for (name, value) in headers {
5916 builder = builder.header(*name, *value);
5917 }
5918 let response = app
5919 .clone()
5920 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
5921 .await
5922 .unwrap();
5923 let status = response.status();
5924 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5925 (
5926 status,
5927 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
5928 )
5929 }
5930
5931 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
5932 body.lines()
5933 .filter_map(|line| line.strip_prefix(field))
5934 .map(str::trim)
5935 .collect()
5936 }
5937
5938 fn resumable_request() -> serde_json::Value {
5939 serde_json::json!({
5940 "model": "m",
5941 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5942 "max_tokens": 4,
5943 "temperature": 0,
5944 "stream": true,
5945 "stream_resumable": true,
5946 })
5947 }
5948
5949 #[tokio::test]
5954 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
5955 let app = test_app();
5956 let body = post_sse_raw(&app, resumable_request()).await;
5957
5958 let request_id = body
5959 .lines()
5960 .find_map(|l| l.strip_prefix("data: "))
5961 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
5962 .and_then(|v| v["request_id"].as_str().map(str::to_string))
5963 .expect("the first chunk names the request");
5964
5965 let ids = sse_field(&body, "id:");
5966 let datas = sse_field(&body, "data:");
5967 assert_eq!(
5968 ids.len(),
5969 datas.len(),
5970 "every event carries an id, or a reconnect cannot name where it stopped"
5971 );
5972 for (i, id) in ids.iter().enumerate() {
5973 assert_eq!(*id, format!("{request_id}:{i}"));
5974 }
5975 let retries = sse_field(&body, "retry:");
5976 assert_eq!(
5977 retries.len(),
5978 1,
5979 "the reconnect delay is stated once, not on every event"
5980 );
5981 assert_eq!(retries[0], "1500");
5982 assert!(
5983 body.contains("data: [DONE]"),
5984 "the end of stream is still stated"
5985 );
5986 }
5987
5988 #[tokio::test]
5992 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
5993 let app = test_app();
5994 let mut request = resumable_request();
5995 request["stream_resumable"] = serde_json::json!(false);
5996 let body = post_sse_raw(&app, request).await;
5997 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
5998 assert!(
5999 sse_field(&body, "id:").is_empty(),
6000 "an id promises a replay this stream cannot serve"
6001 );
6002 assert!(sse_field(&body, "retry:").is_empty());
6003 }
6004
6005 #[tokio::test]
6009 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6010 let app = test_app();
6011 let body = post_sse_raw(&app, resumable_request()).await;
6012 let request_id = sse_field(&body, "id:")[0]
6013 .rsplit_once(':')
6014 .unwrap()
6015 .0
6016 .to_string();
6017 let streamed: Vec<String> = sse_field(&body, "data:")
6018 .iter()
6019 .map(|d| d.to_string())
6020 .collect();
6021
6022 let (status, polled) = get_json(
6023 &app,
6024 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
6025 )
6026 .await;
6027 assert_eq!(status, StatusCode::OK);
6028 let events: Vec<String> = polled["events"]
6029 .as_array()
6030 .unwrap()
6031 .iter()
6032 .map(|e| e["data"].as_str().unwrap().to_string())
6033 .collect();
6034 assert_eq!(
6035 events, streamed,
6036 "the fallback must deliver the same answer, not a re-run of it"
6037 );
6038 assert_eq!(polled["request_id"], request_id);
6039 assert_eq!(
6040 polled["done"], false,
6041 "events were still being handed out, so the client must ask again"
6042 );
6043
6044 let next = polled["next_index"].as_u64().unwrap();
6047 let (_, drained) = get_json(
6048 &app,
6049 &format!(
6050 "{}?from={next}",
6051 ferrox_api::routes::v1_stream_poll(&request_id)
6052 ),
6053 )
6054 .await;
6055 assert_eq!(drained["done"], true);
6056 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6057 }
6058
6059 #[tokio::test]
6063 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6064 let app = test_app();
6065 let body = post_sse_raw(&app, resumable_request()).await;
6066 let ids = sse_field(&body, "id:");
6067 let datas: Vec<String> = sse_field(&body, "data:")
6068 .iter()
6069 .map(|d| d.to_string())
6070 .collect();
6071 assert!(
6072 ids.len() >= 3,
6073 "need a few events to resume into the middle"
6074 );
6075 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6076
6077 let (status, resumed) = get_json_with_headers(
6078 &app,
6079 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
6080 &[],
6081 )
6082 .await;
6083 assert_eq!(status, StatusCode::OK);
6084 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6085
6086 let (_, tail) = get_json(
6088 &app,
6089 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
6090 )
6091 .await;
6092 let tail_events: Vec<String> = tail["events"]
6093 .as_array()
6094 .unwrap()
6095 .iter()
6096 .map(|e| e["data"].as_str().unwrap().to_string())
6097 .collect();
6098 assert_eq!(tail_events, datas[2..].to_vec());
6099 }
6100
6101 #[tokio::test]
6104 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6105 use http_body_util::BodyExt;
6106 use tower::ServiceExt;
6107
6108 let app = test_app();
6109 let body = post_sse_raw(&app, resumable_request()).await;
6110 let ids = sse_field(&body, "id:");
6111 let datas: Vec<String> = sse_field(&body, "data:")
6112 .iter()
6113 .map(|d| d.to_string())
6114 .collect();
6115 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6116
6117 let response = app
6118 .clone()
6119 .oneshot(
6120 axum::http::Request::builder()
6121 .method("GET")
6122 .uri(ferrox_api::routes::v1_stream(&request_id))
6123 .header("last-event-id", format!("{request_id}:0"))
6124 .body(axum::body::Body::empty())
6125 .unwrap(),
6126 )
6127 .await
6128 .unwrap();
6129 assert_eq!(response.status(), StatusCode::OK);
6130 assert_eq!(
6131 response
6132 .headers()
6133 .get("x-accel-buffering")
6134 .and_then(|v| v.to_str().ok()),
6135 Some("no"),
6136 "the reconnect needs the same anti-buffering header as the stream"
6137 );
6138 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6139 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6140 assert_eq!(
6141 sse_field(&resumed, "data:")
6142 .iter()
6143 .map(|d| d.to_string())
6144 .collect::<Vec<_>>(),
6145 datas[1..].to_vec()
6146 );
6147 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6148 }
6149
6150 #[tokio::test]
6154 async fn a_last_event_id_from_another_stream_is_refused() {
6155 let app = test_app();
6156 let body = post_sse_raw(&app, resumable_request()).await;
6157 let request_id = sse_field(&body, "id:")[0]
6158 .rsplit_once(':')
6159 .unwrap()
6160 .0
6161 .to_string();
6162
6163 let (status, err) = get_json_with_headers(
6164 &app,
6165 &ferrox_api::routes::v1_stream(&request_id),
6166 &[("last-event-id", "chatcmpl-someone-else:3")],
6167 )
6168 .await;
6169 assert_eq!(status, StatusCode::BAD_REQUEST);
6170 assert_eq!(err["error"]["code"], "bad_last_event_id");
6171 }
6172
6173 #[tokio::test]
6177 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6178 let app = test_app();
6179 let mut request = resumable_request();
6180 request["stream_resumable"] = serde_json::json!(false);
6181 let body = post_sse_raw(&app, request).await;
6182 let request_id = body
6183 .lines()
6184 .find_map(|l| l.strip_prefix("data: "))
6185 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6186 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6187 .unwrap();
6188
6189 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
6190 assert_eq!(status, StatusCode::NOT_FOUND);
6191 assert_eq!(err["error"]["code"], "stream_not_found");
6192 assert!(err["error"]["message"]
6193 .as_str()
6194 .unwrap()
6195 .contains("stream_resumable"));
6196 }
6197
6198 #[test]
6202 fn the_axum_stream_patterns_match_the_published_templates() {
6203 assert_eq!(
6204 axum_path(ferrox_api::routes::V1_STREAM),
6205 "/v1/stream/:request_id"
6206 );
6207 assert_eq!(
6208 axum_path(ferrox_api::routes::V1_STREAM_POLL),
6209 "/v1/stream/:request_id/poll"
6210 );
6211 assert_eq!(
6212 ferrox_api::routes::v1_stream("abc"),
6213 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
6214 );
6215 }
6216
6217 #[test]
6227 fn no_published_template_reaches_the_router_with_its_braces() {
6228 for template in [
6229 ferrox_api::routes::V1_STREAM,
6230 ferrox_api::routes::V1_STREAM_POLL,
6231 ferrox_api::routes::V1_RESPONSE,
6232 ferrox_api::routes::V1_RESPONSE_CANCEL,
6233 ferrox_api::routes::ADMIN_TASK_CANCEL,
6234 ] {
6235 assert!(
6236 template.contains('{'),
6237 "{template} is in the template list but has no placeholder"
6238 );
6239 let mounted = axum_path(template);
6240 assert!(
6241 !mounted.contains('{') && !mounted.contains('}'),
6242 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6243 );
6244 assert!(
6245 mounted.contains(':'),
6246 "{template} lost its placeholder entirely and would match one path only"
6247 );
6248 }
6249 }
6250
6251 #[tokio::test]
6259 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6260 let app = test_app();
6261 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6262 assert_eq!(status, StatusCode::NOT_FOUND);
6263 assert!(
6264 !body.is_null(),
6265 "empty body means axum never matched the route, so the id was read as a literal segment"
6266 );
6267 }
6268
6269 #[tokio::test]
6272 async fn the_task_list_starts_empty_rather_than_absent() {
6273 let app = test_app();
6274 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
6275 assert_eq!(status, StatusCode::OK);
6276 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6277 }
6278
6279 #[tokio::test]
6288 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6289 assert!(
6290 std::env::var("FERROX_SLOT_SAVE_PATH").is_err(),
6291 "this test asserts the unconfigured behaviour"
6292 );
6293 let app = test_app();
6294 let (status, body) = post_json_uri(
6295 &app,
6296 &format!("{}?action=save", ferrox_api::routes::slots_id(0)),
6297 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6298 )
6299 .await;
6300 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6301 assert!(
6302 body["error"]["message"]
6303 .as_str()
6304 .unwrap()
6305 .contains("--slot-save-path"),
6306 "{body}"
6307 );
6308 }
6309
6310 async fn post_json_uri(
6311 app: &Router,
6312 uri: &str,
6313 body: serde_json::Value,
6314 ) -> (StatusCode, serde_json::Value) {
6315 use http_body_util::BodyExt;
6316 use tower::ServiceExt;
6317
6318 let response = app
6319 .clone()
6320 .oneshot(
6321 axum::http::Request::builder()
6322 .method("POST")
6323 .uri(uri)
6324 .header("content-type", "application/json")
6325 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6326 .unwrap(),
6327 )
6328 .await
6329 .unwrap();
6330 let status = response.status();
6331 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6332 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6333 (status, json)
6334 }
6335
6336 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6337 post_json_uri(app, "/v1/chat/completions", body).await.1
6338 }
6339
6340 #[tokio::test]
6349 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6350 let app = test_app();
6351 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6352 assert_eq!(status, StatusCode::OK);
6353
6354 let memory = &body["memory"];
6355 if memory.is_null() {
6356 return;
6359 }
6360 assert!(
6361 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6362 "a read that produced a zero is a broken read, not an idle \
6363 engine: {memory}"
6364 );
6365 assert!(
6366 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6367 "the quantity must travel with the number: {memory}"
6368 );
6369 }
6370
6371 #[tokio::test]
6377 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6378 let app = test_app();
6379 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6380 assert_eq!(status, StatusCode::OK);
6381 for pool in ["kv_pages", "window_slots", "state_slots"] {
6382 assert!(
6383 body["pools"][pool].is_null(),
6384 "{pool} must be null rather than a zero row: {}",
6385 body["pools"]
6386 );
6387 }
6388 }
6389
6390 #[tokio::test]
6398 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6399 use http_body_util::BodyExt;
6400 use tower::ServiceExt;
6401
6402 let app = test_app();
6403 let send = |body: serde_json::Value| {
6404 let app = app.clone();
6405 async move {
6406 app.oneshot(
6407 axum::http::Request::builder()
6408 .method("POST")
6409 .uri(ferrox_api::routes::V1_MESSAGES)
6410 .header("content-type", "application/json")
6411 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6412 .unwrap(),
6413 )
6414 .await
6415 .unwrap()
6416 }
6417 };
6418
6419 let ok = send(serde_json::json!({
6420 "model": "test",
6421 "max_tokens": 1,
6422 "messages": [{"role": "user", "content": "hi"}],
6423 }))
6424 .await;
6425 assert_eq!(ok.status(), StatusCode::OK);
6426 let id = ok
6427 .headers()
6428 .get("request-id")
6429 .expect("a served message names its id")
6430 .to_str()
6431 .unwrap()
6432 .to_string();
6433 assert!(!id.is_empty());
6434
6435 let bad = send(serde_json::json!({"model": "test"})).await;
6438 assert!(bad.status().is_client_error());
6439 let other = bad.headers().get("request-id").expect("errors too");
6440 assert_ne!(other.to_str().unwrap(), id);
6441 let _ = bad.into_body().collect().await.unwrap();
6442 }
6443
6444 #[tokio::test]
6452 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6453 let state = Arc::new(test_state(
6454 test_model_full_byte_vocab(),
6455 ResponseCache::new(1000, Duration::from_secs(3600)),
6456 ));
6457 let app = test_app_with_state(Arc::clone(&state));
6458 let body = serde_json::json!({
6459 "model": "test",
6460 "messages": [{"role": "user", "content": "hi"}],
6461 "max_tokens": 1,
6462 });
6463
6464 state
6465 .maintenance
6466 .lock()
6467 .unwrap()
6468 .begin_rebuild()
6469 .expect("a fresh server is serving, so the rebuild starts");
6470 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6471 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6472 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6473
6474 state.maintenance.lock().unwrap().finish_rebuild(true);
6475 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6476 assert_eq!(
6477 status,
6478 StatusCode::OK,
6479 "the gate reopens; a rebuild is not a latch"
6480 );
6481 }
6482
6483 #[tokio::test]
6488 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6489 let app = test_app();
6490 let (status, body) = post_json_uri(
6491 &app,
6492 ferrox_api::routes::V1_CANCEL,
6493 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6494 )
6495 .await;
6496 assert_eq!(status, StatusCode::NOT_FOUND);
6497 assert_eq!(body["cancelled"], serde_json::json!(false));
6498 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6499 assert!(
6500 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6501 "the verdict must carry a human reason: {body}"
6502 );
6503 }
6504
6505 #[tokio::test]
6509 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6510 let state = Arc::new(test_state(
6511 test_model_full_byte_vocab(),
6512 ResponseCache::new(1000, Duration::from_secs(3600)),
6513 ));
6514 let app = test_app_with_state(Arc::clone(&state));
6515 let (token, _guard) = state.cancels.register("chatcmpl-live");
6516
6517 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6518 assert_eq!(status, StatusCode::OK);
6519 assert_eq!(before["generating_now"], serde_json::json!(1));
6520
6521 let (status, body) = post_json_uri(
6522 &app,
6523 ferrox_api::routes::V1_CANCEL,
6524 serde_json::json!({ "request_id": "chatcmpl-live" }),
6525 )
6526 .await;
6527 assert_eq!(status, StatusCode::OK);
6528 assert_eq!(body["cancelled"], serde_json::json!(true));
6529 assert!(
6530 token.is_cancelled(),
6531 "the endpoint answered ok without setting the flag the decode loop reads"
6532 );
6533 }
6534
6535 #[tokio::test]
6536 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6537 let app = test_app();
6538 let (status, tok) =
6539 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6540 assert_eq!(status, StatusCode::OK);
6541 let tokens = tok["tokens"].as_array().unwrap();
6542 assert_eq!(tok["count"], tokens.len());
6543 assert!(!tokens.is_empty());
6544
6545 let (status, detok) = post_json_uri(
6546 &app,
6547 "/v1/detokenize",
6548 serde_json::json!({ "tokens": tokens }),
6549 )
6550 .await;
6551 assert_eq!(status, StatusCode::OK);
6552 assert_eq!(detok["text"], "Hi");
6553
6554 let (status, emb) = post_json_uri(
6555 &app,
6556 "/v1/embeddings",
6557 serde_json::json!({
6558 "input": "Hi",
6559 "embedding_type": "mean"
6560 }),
6561 )
6562 .await;
6563 assert_eq!(status, StatusCode::OK);
6564 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6565 assert!(!vec.is_empty());
6566 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6567 }
6568
6569 #[tokio::test]
6575 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6576 let app = test_app();
6577 let (status, body) = post_json_uri(
6578 &app,
6579 "/v1/embeddings",
6580 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6581 )
6582 .await;
6583 assert_eq!(status, StatusCode::BAD_REQUEST);
6584 let msg = body["error"]["message"].as_str().unwrap();
6585 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6586 }
6587
6588 #[tokio::test]
6593 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6594 async fn a_real_embedding_model_serves_v1_embeddings() {
6595 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6596 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6597 if !path.exists() {
6598 eprintln!("SKIP: {} not present", path.display());
6599 return;
6600 }
6601 let encoder = ferrox_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6602 let mut state = test_state(
6603 test_model_full_byte_vocab(),
6604 ResponseCache::new(1000, Duration::from_secs(3600)),
6605 );
6606 state.embedding = Some(Arc::new(encoder));
6607 let app = test_app_with_state(Arc::new(state));
6608
6609 let (status, body) = post_json_uri(
6610 &app,
6611 "/v1/embeddings",
6612 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6613 )
6614 .await;
6615 assert_eq!(status, StatusCode::OK, "{body}");
6616 assert_eq!(body["model"], "bge-small-en-v1.5");
6617 let data = body["data"].as_array().unwrap();
6618 assert_eq!(data.len(), 2);
6619 for (i, row) in data.iter().enumerate() {
6620 assert_eq!(row["index"], i);
6621 let v: Vec<f64> = row["embedding"]
6622 .as_array()
6623 .unwrap()
6624 .iter()
6625 .map(|x| x.as_f64().unwrap())
6626 .collect();
6627 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6628 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6629 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6630 }
6631 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6634
6635 let (status, mean) = post_json_uri(
6639 &app,
6640 "/v1/embeddings",
6641 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6642 )
6643 .await;
6644 assert_eq!(status, StatusCode::OK);
6645 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6646 }
6647
6648 #[tokio::test]
6660 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6661 async fn an_encoder_can_be_the_loaded_model() {
6662 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6663 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6664 if !path.exists() {
6665 eprintln!("SKIP: {} not present", path.display());
6666 return;
6667 }
6668
6669 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6673 assert!(
6674 matches!(loaded, model::LoadedModel::Encoder(_)),
6675 "an encoder-only GGUF reached a decoder loader"
6676 );
6677 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6678 assert!(
6679 matches!(loaded, Loaded::Encoder(_)),
6680 "the encoder did not stay an encoder through activation"
6681 );
6682 assert!(
6683 batcher.is_none() && ceiling.is_none(),
6684 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6685 );
6686
6687 let state = test_state(
6688 test_model_full_byte_vocab(),
6689 ResponseCache::new(1000, Duration::from_secs(3600)),
6690 );
6691 state.swap_active(Some(Arc::new(ActiveModel {
6692 id: None,
6693 loaded,
6694 batcher,
6695 ceiling,
6696 checkpoint_path: None,
6697 })));
6698 let app = test_app_with_state(Arc::new(state));
6699
6700 let (status, body) = post_json_uri(
6702 &app,
6703 "/v1/embeddings",
6704 serde_json::json!({ "input": "Hello world" }),
6705 )
6706 .await;
6707 assert_eq!(status, StatusCode::OK, "{body}");
6708 assert_eq!(body["model"], "bge-small-en-v1.5");
6709 let v = body["data"][0]["embedding"].as_array().unwrap();
6710 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6711
6712 let (status, body) = post_json_uri(
6714 &app,
6715 "/v1/chat/completions",
6716 serde_json::json!({
6717 "model": "bge-small-en-v1.5",
6718 "messages": [{"role": "user", "content": "hi"}],
6719 }),
6720 )
6721 .await;
6722 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6723 let msg = body["error"]["message"].as_str().unwrap();
6724 for fact in [
6725 "bge-small-en-v1.5",
6726 "bert",
6727 "embedding model",
6728 "/v1/embeddings",
6729 ] {
6730 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6731 }
6732
6733 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
6735 assert_eq!(status, StatusCode::OK);
6736 let entry = &models["data"][0];
6737 assert_eq!(entry["id"], "bge-small-en-v1.5");
6738 assert_eq!(entry["ferrox_model_kind"], "embedding");
6739 assert_eq!(entry["ferrox_tokenizer"], "gguf-wordpiece");
6740 assert_eq!(entry["ferrox_n_embd"], 384);
6741 assert_eq!(entry["ferrox_pooling"], "CLS");
6742 assert_eq!(
6743 entry["ferrox_endpoints"],
6744 serde_json::json!(["/v1/embeddings"])
6745 );
6746 assert!(entry.get("supported_reasoning_efforts").is_none());
6749
6750 let (status, health) = get_json(&app, ferrox_api::routes::HEALTH).await;
6752 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6753 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6754 assert_eq!(health["model"]["synthetic_weights"], false);
6755 let weights = health["capabilities"]
6756 .as_array()
6757 .unwrap()
6758 .iter()
6759 .find(|c| c["id"] == ferrox_api::health::capability::REAL_WEIGHTS)
6760 .expect("a real-weights capability row");
6761 let detail = weights["detail"].as_str().unwrap_or_default();
6762 assert!(detail.contains("ENCODER"), "{detail}");
6763 let (status, body) = post_json_uri(
6771 &app,
6772 ferrox_api::routes::V1_TOKENIZE,
6773 serde_json::json!({ "content": "hello world" }),
6774 )
6775 .await;
6776 assert_eq!(
6777 status,
6778 StatusCode::OK,
6779 "an encoder has a real tokenizer: {body}"
6780 );
6781 let tokens = body["tokens"].as_array().expect("tokens array").clone();
6782 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
6783
6784 let (status, body) = post_json_uri(
6785 &app,
6786 ferrox_api::routes::V1_DETOKENIZE,
6787 serde_json::json!({ "tokens": tokens }),
6788 )
6789 .await;
6790 assert_eq!(status, StatusCode::OK, "{body}");
6791 let round_tripped = body["content"].as_str().expect("content").to_string();
6792 assert!(
6793 round_tripped.contains("hello") && round_tripped.contains("world"),
6794 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
6795 );
6796
6797 let (status, _) = post_json_uri(
6800 &app,
6801 "/v1/completions",
6802 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
6803 )
6804 .await;
6805 assert_eq!(
6806 status,
6807 StatusCode::NOT_IMPLEMENTED,
6808 "tokenizing an encoder must not have opened a path to generating with one"
6809 );
6810 }
6811
6812 #[tokio::test]
6817 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6818 use http_body_util::BodyExt;
6819 use tower::ServiceExt;
6820
6821 let fixture = concat!(
6822 "../ferrox-models/tests/fixtures/",
6823 "ferrox_real_moe_test.gguf"
6824 );
6825 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6826 let decoder = Decoder::from_gguf_with_expert_cache(
6827 &fixture,
6828 ferrox_models::config::test_moe_fixture(),
6829 Some(1024 * 1024),
6830 )
6831 .expect("MoE fixture must load store-backed");
6832
6833 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
6837 .layers
6838 .iter()
6839 .map(|_| {
6840 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
6841 })
6842 .collect();
6843 decoder.forward_token(1, 0, &mut caches);
6844
6845 let model = Model::Gguf(GgufModel {
6846 decoder: Arc::new(decoder),
6847 tokenizer: Arc::new(ServerTokenizer::Byte),
6848 stop_tokens: StopTokens::default(),
6849 bos_id: None,
6850 is_synthetic: false,
6851 chat_template: chat_template::PromptTemplate::plain(),
6852 });
6853 let state = Arc::new(test_state(
6854 model,
6855 ResponseCache::new(16, Duration::from_secs(60)),
6856 ));
6857 let app = Router::new()
6858 .route("/metrics", axum::routing::get(metrics))
6859 .route("/v1/chat/completions", post(chat_completions))
6860 .with_state(state);
6861
6862 let fetch_metrics = |app: Router| async move {
6863 let resp = app
6864 .oneshot(
6865 axum::http::Request::builder()
6866 .method("GET")
6867 .uri("/metrics")
6868 .body(axum::body::Body::empty())
6869 .unwrap(),
6870 )
6871 .await
6872 .unwrap();
6873 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6874 String::from_utf8(bytes.to_vec()).unwrap()
6875 };
6876
6877 let after = fetch_metrics(app.clone()).await;
6878 assert!(
6879 after.contains("ferrox_expert_cache_misses_total"),
6880 "streaming model must expose expert-cache metrics: {after}"
6881 );
6882 let misses: u64 = after
6883 .lines()
6884 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
6885 .and_then(|l| l.split_whitespace().nth(1))
6886 .and_then(|v| v.parse().ok())
6887 .expect("misses metric line must parse");
6888 assert!(
6889 misses > 0,
6890 "decode must have read experts through the store: {after}"
6891 );
6892 }
6893
6894 fn weather_tool() -> serde_json::Value {
6895 serde_json::json!({
6896 "type": "function",
6897 "function": {
6898 "name": "get_weather",
6899 "description": "Get the current weather for a location.",
6900 "parameters": {
6901 "type": "object",
6902 "properties": {"location": {"type": "string"}},
6903 "required": ["location"]
6904 }
6905 }
6906 })
6907 }
6908
6909 fn weather_tool_def() -> ToolDef {
6910 ToolDef {
6911 kind: "function".to_string(),
6912 function: ToolFunctionDef {
6913 name: "get_weather".to_string(),
6914 description: Some("Get the current weather for a location.".to_string()),
6915 parameters: Some(serde_json::json!({
6916 "type": "object",
6917 "properties": {"location": {"type": "string"}},
6918 "required": ["location"]
6919 })),
6920 },
6921 }
6922 }
6923
6924 #[test]
6925 fn tool_preamble_mentions_every_tool_name_and_description() {
6926 let preamble = tool_preamble(&[weather_tool_def()]);
6927 assert!(preamble.contains("get_weather"));
6928 assert!(preamble.contains("Get the current weather for a location."));
6929 assert!(preamble.contains("<tool_call>"));
6930 assert!(preamble.contains("</tool_call>"));
6931 }
6932
6933 #[test]
6934 fn a_real_marker_becomes_a_structured_tool_call() {
6935 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6936 let (message, finish) = build_response_message(
6937 text.to_string(),
6938 &[weather_tool_def()],
6939 output::OutputPosture::for_model("test-model"),
6940 "stop",
6941 );
6942 assert_eq!(finish, "tool_calls");
6943 let calls = message.tool_calls.expect("must carry a tool call");
6944 assert_eq!(calls[0].function.name, "get_weather");
6945 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
6946 assert_eq!(parsed["location"], "Paris");
6947 }
6948
6949 #[test]
6950 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
6951 let (message, finish) = build_response_message(
6952 "just an answer".to_string(),
6953 &[weather_tool_def()],
6954 output::OutputPosture::for_model("test-model"),
6955 "stop",
6956 );
6957 assert_eq!(finish, "stop");
6958 assert!(message.tool_calls.is_none());
6959 assert_eq!(message.content.as_deref(), Some("just an answer"));
6960 }
6961
6962 #[test]
6965 fn a_malformed_payload_is_not_a_tool_call() {
6966 let (message, finish) = build_response_message(
6967 "<tool_call>not valid json at all</tool_call>".to_string(),
6968 &[weather_tool_def()],
6969 output::OutputPosture::for_model("test-model"),
6970 "stop",
6971 );
6972 assert_eq!(finish, "stop");
6973 assert!(message.tool_calls.is_none());
6974 }
6975
6976 #[test]
6979 fn a_tool_that_was_never_offered_is_not_returned() {
6980 let (message, finish) = build_response_message(
6981 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
6982 &[weather_tool_def()],
6983 output::OutputPosture::for_model("test-model"),
6984 "stop",
6985 );
6986 assert_eq!(finish, "stop");
6987 assert!(message.tool_calls.is_none());
6988 }
6989
6990 #[test]
6992 fn marker_text_with_no_tools_offered_stays_content() {
6993 let (message, finish) = build_response_message(
6994 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
6995 &[],
6996 output::OutputPosture::for_model("test-model"),
6997 "stop",
6998 );
6999 assert_eq!(finish, "stop");
7000 assert!(message.tool_calls.is_none());
7001 assert!(message.content.is_some());
7002 }
7003
7004 #[test]
7008 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7009 let opened = std::cell::Cell::new(0usize);
7010 let mut parser = crate::policy::parser::ToolCallParser::new(
7011 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7012 vec![
7013 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7014 "write_file",
7015 serde_json::json!({"type": "object", "properties": {
7016 "path": {"type": "string"},
7017 "contents": {"type": "string"}
7018 }}),
7019 ),
7020 ],
7021 );
7022 let wire = "<tool_call><function=write_file>\
7023 <parameter=path>\n/tmp/x\n</parameter>\
7024 <parameter=contents>\nhello world\n</parameter>\
7025 </function></tool_call>";
7026
7027 let mut deltas = Vec::new();
7028 let mut text = String::new();
7029 for piece in wire.as_bytes().chunks(7) {
7030 let chunk = String::from_utf8_lossy(piece).into_owned();
7031 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7032 text.push_str(&more_text);
7033 deltas.extend(more);
7034 }
7035 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7036 text.push_str(&more_text);
7037 deltas.extend(more);
7038
7039 assert_eq!(opened.get(), 1, "one call opened");
7040 assert!(text.is_empty(), "the markers are not content: {text:?}");
7041
7042 let first = &deltas[0];
7043 assert_eq!(first.index, 0);
7044 assert_eq!(first.id.as_deref(), Some("call_0"));
7045 assert_eq!(first.kind, Some("function"));
7046 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7047
7048 let joined: String = deltas
7051 .iter()
7052 .filter_map(|d| d.function.arguments.clone())
7053 .collect();
7054 let parsed: serde_json::Value =
7055 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7056 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7057 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7058 assert!(
7059 deltas.len() >= 3,
7060 "the arguments arrived in pieces, not whole: {}",
7061 deltas.len()
7062 );
7063 assert!(
7064 deltas[1..].iter().all(|d| d.function.name.is_none()),
7065 "only the opening delta carries identity"
7066 );
7067 }
7068
7069 #[test]
7071 fn text_around_a_streamed_call_is_still_content() {
7072 let opened = std::cell::Cell::new(0usize);
7073 let mut parser = crate::policy::parser::ToolCallParser::new(
7074 crate::policy::parser::ToolCallFormat::Qwen25,
7075 vec![crate::policy::parser::tool_call::ToolSchema::new(
7076 "get_weather",
7077 )],
7078 );
7079 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7080 \"arguments\": {}}</tool_call> done";
7081 let mut text = String::new();
7082 for piece in wire.as_bytes().chunks(5) {
7083 let chunk = String::from_utf8_lossy(piece).into_owned();
7084 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7085 text.push_str(&more);
7086 }
7087 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7088 text.push_str(&more);
7089
7090 assert_eq!(opened.get(), 1);
7091 assert!(text.starts_with("let me check. "), "{text:?}");
7092 assert!(text.ends_with(" done"), "{text:?}");
7093 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7094 }
7095
7096 #[test]
7099 fn a_reasoning_block_is_split_out_of_the_answer() {
7100 let (message, finish) = build_response_message(
7101 "<think>weighing it up</think>The answer is 4.".to_string(),
7102 &[],
7103 output::OutputPosture::for_model("Qwen3-8B"),
7104 "stop",
7105 );
7106 assert_eq!(finish, "stop");
7107 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7108 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7109 }
7110
7111 #[test]
7114 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7115 let (message, _) = build_response_message(
7116 "Use the <think> tag like this.".to_string(),
7117 &[],
7118 output::OutputPosture::for_model("llama-3.1-8b"),
7119 "stop",
7120 );
7121 assert_eq!(
7122 message.content.as_deref(),
7123 Some("Use the <think> tag like this.")
7124 );
7125 assert!(message.reasoning_content.is_none());
7126 }
7127
7128 #[tokio::test]
7136 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7137 let app = test_app();
7138 let body = serde_json::json!({
7139 "model": "m",
7140 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7141 "max_tokens": 4,
7142 "temperature": 0,
7143 });
7144 let resp = post_json(&app, body).await;
7145 let message = &resp["choices"][0]["message"];
7146 assert!(message["content"].is_string());
7147 assert!(message.get("tool_calls").is_none());
7148 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7149 assert_eq!(resp["usage"]["completion_tokens"], 4);
7150 assert_eq!(
7151 resp["usage"]["total_tokens"],
7152 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7153 );
7154 }
7155
7156 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7157 use http_body_util::BodyExt;
7158 use tower::ServiceExt;
7159
7160 let response = app
7161 .clone()
7162 .oneshot(
7163 axum::http::Request::builder()
7164 .method("GET")
7165 .uri(uri)
7166 .body(axum::body::Body::empty())
7167 .unwrap(),
7168 )
7169 .await
7170 .unwrap();
7171 let status = response.status();
7172 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7173 (status, serde_json::from_slice(&bytes).unwrap())
7174 }
7175
7176 #[tokio::test]
7177 async fn health_answers_a_capability_handshake_not_a_boolean() {
7178 let app = test_app();
7179 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7180 assert_eq!(status, StatusCode::OK);
7181
7182 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7183 assert_eq!(health.state, ferrox_api::HealthState::Ready);
7184 assert!(health.pid > 0);
7185 assert!(health.server_time_unix_ms > 0);
7186 assert_eq!(health.last_request_age_seconds, None);
7189
7190 for id in [
7193 ferrox_api::health::capability::CPU,
7194 ferrox_api::health::capability::METAL,
7195 ferrox_api::health::capability::CUDA,
7196 ferrox_api::health::capability::REAL_WEIGHTS,
7197 ferrox_api::health::capability::CONTINUOUS_BATCHING,
7198 ] {
7199 let cap = health
7200 .capability(id)
7201 .unwrap_or_else(|| panic!("{id} missing"));
7202 assert!(!cap.reason.is_empty(), "{cap:?}");
7203 assert!(!cap.detail.is_empty(), "{cap:?}");
7204 }
7205 let weights = health
7209 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
7210 .unwrap();
7211 assert!(!weights.available);
7212 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
7213 assert!(health.model.as_ref().unwrap().synthetic_weights);
7214 }
7215
7216 #[tokio::test]
7217 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7218 let app = test_app();
7219 let _ = post_json(
7220 &app,
7221 serde_json::json!({
7222 "model": "m",
7223 "messages": [{"role": "user", "content": "\u{1}"}],
7224 "max_tokens": 1,
7225 "temperature": 0,
7226 }),
7227 )
7228 .await;
7229 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
7230 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
7231 let age = health
7232 .last_request_age_seconds
7233 .expect("a served request is evidence of liveness");
7234 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7235 }
7236
7237 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7239 use http_body_util::BodyExt;
7240 use tower::ServiceExt;
7241
7242 let response = app
7243 .clone()
7244 .oneshot(
7245 axum::http::Request::builder()
7246 .method("POST")
7247 .uri("/v1/chat/completions")
7248 .header("content-type", "application/json")
7249 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7250 .unwrap(),
7251 )
7252 .await
7253 .unwrap();
7254 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7255 String::from_utf8(bytes.to_vec())
7256 .unwrap()
7257 .lines()
7258 .filter_map(|line| line.strip_prefix("data: "))
7259 .filter(|payload| *payload != "[DONE]")
7260 .map(|payload| serde_json::from_str(payload).unwrap())
7261 .collect()
7262 }
7263
7264 #[tokio::test]
7265 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7266 let app = test_app();
7267 let chunks = post_sse_chunks(
7268 &app,
7269 serde_json::json!({
7270 "model": "m",
7271 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7272 "max_tokens": 4,
7273 "temperature": 0,
7274 "stream": true,
7275 }),
7276 )
7277 .await;
7278
7279 assert!(!chunks.is_empty());
7280 let request_id = chunks[0]["request_id"]
7281 .as_str()
7282 .expect("the first chunk names the request")
7283 .to_string();
7284 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7285 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7288 assert!(
7289 chunk.get("request_id").is_none(),
7290 "chunk {i} repeats request_id"
7291 );
7292 }
7293 for chunk in &chunks {
7296 assert_eq!(chunk["id"], serde_json::json!(request_id));
7297 }
7298
7299 let other = post_sse_chunks(
7300 &app,
7301 serde_json::json!({
7302 "model": "m",
7303 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7304 "max_tokens": 4,
7305 "temperature": 0,
7306 "stream": true,
7307 }),
7308 )
7309 .await;
7310 assert_ne!(
7311 other[0]["request_id"].as_str().unwrap(),
7312 request_id,
7313 "two concurrent chats must not share an id"
7314 );
7315 }
7316
7317 #[tokio::test]
7318 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7319 let app = test_app();
7320 let resp = post_json(
7321 &app,
7322 serde_json::json!({
7323 "model": "m",
7324 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7325 "max_tokens": 2,
7326 "temperature": 0,
7327 }),
7328 )
7329 .await;
7330 assert_eq!(resp["id"], resp["request_id"]);
7331 assert!(resp["request_id"]
7332 .as_str()
7333 .unwrap()
7334 .starts_with("chatcmpl-"));
7335 }
7336
7337 #[tokio::test]
7340 async fn usage_carries_separate_prefill_and_decode_timings() {
7341 let app = test_app();
7342 let resp = post_json(
7343 &app,
7344 serde_json::json!({
7345 "model": "m",
7346 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7347 "max_tokens": 4,
7348 "temperature": 0,
7349 }),
7350 )
7351 .await;
7352 let usage = &resp["usage"];
7353 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7354 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7355 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7356 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7357 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7359 }
7360
7361 #[tokio::test]
7369 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7370 let app = test_app();
7371 let body = serde_json::json!({
7372 "model": "m",
7373 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7374 "max_tokens": 4,
7375 "temperature": 0,
7376 "tools": [weather_tool()],
7377 });
7378 let resp = post_json(&app, body).await;
7379 let message = &resp["choices"][0]["message"];
7380 assert!(
7381 message["content"].is_string(),
7382 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7383 );
7384 assert!(message.get("tool_calls").is_none());
7385 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7389 }
7390
7391 #[tokio::test]
7395 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7396 let app = test_app();
7397 let body = serde_json::json!({
7398 "model": "m",
7399 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7400 "max_tokens": 3,
7401 "temperature": 0,
7402 });
7403 let first = post_json(&app, body.clone()).await;
7404 assert_eq!(first["ferrox_cache"], "miss");
7405 let second = post_json(&app, body).await;
7406 assert_eq!(second["ferrox_cache"], "hit");
7407 assert_eq!(
7408 first["choices"][0]["message"]["content"],
7409 second["choices"][0]["message"]["content"]
7410 );
7411 assert_eq!(
7412 first["choices"][0]["finish_reason"],
7413 second["choices"][0]["finish_reason"]
7414 );
7415 assert_eq!(first["usage"], second["usage"]);
7416 assert_eq!(second["usage"]["completion_tokens"], 3);
7417 }
7418
7419 #[tokio::test]
7430 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7431 let app = test_app();
7432 let plain = serde_json::json!({
7433 "model": "m",
7434 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7435 "max_tokens": 3,
7436 "temperature": 0,
7437 });
7438
7439 let first = post_json(&app, plain.clone()).await;
7440 assert_eq!(first["ferrox_cache"], "miss");
7441 let unconstrained = first["choices"][0]["message"]["content"]
7442 .as_str()
7443 .expect("content")
7444 .to_string();
7445
7446 let mut constrained = plain.clone();
7447 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7448 let second = post_json(&app, constrained).await;
7449 assert_eq!(
7450 second["ferrox_cache"], "miss",
7451 "a grammar is part of the key, so this body has never been answered"
7452 );
7453 let constrained_answer = second["choices"][0]["message"]["content"]
7457 .as_str()
7458 .expect("content")
7459 .to_string();
7460 assert!(
7461 constrained_answer.contains("-> \"yes\"]"),
7462 "the grammar must have been compiled AND applied, not skipped \
7463 by a cache hit: {constrained_answer}"
7464 );
7465 assert_ne!(
7466 constrained_answer, unconstrained,
7467 "the constrained request was served the unconstrained answer"
7468 );
7469
7470 let third = post_json(&app, plain).await;
7474 assert_eq!(third["ferrox_cache"], "hit");
7475 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7476 }
7477
7478 #[tokio::test]
7500 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7501 let state = Arc::new(test_state(
7502 test_model_full_byte_vocab(),
7503 ResponseCache::new(1000, Duration::from_secs(3600)),
7504 ));
7505 let app = test_app_with_state(state.clone());
7506 let plain = serde_json::json!({
7507 "model": "m",
7508 "messages": [
7509 {"role": "system", "content": "Answer in JSON when it helps."},
7510 {"role": "user", "content": "\u{1}\u{2}"},
7511 ],
7512 "max_tokens": 3,
7513 "temperature": 0,
7514 });
7515
7516 let first = post_json(&app, plain.clone()).await;
7517 assert_eq!(first["ferrox_cache"], "miss");
7518 assert_eq!(state.cache_stats().entries, 1);
7519
7520 let mut as_json = plain.clone();
7521 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7522 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7523 assert_eq!(
7524 status,
7525 StatusCode::BAD_REQUEST,
7526 "the demo banner is not a JSON object, whoever generated it"
7527 );
7528 assert_eq!(
7529 state.cache_stats().hits,
7530 0,
7531 "a json_object request must not be answered from an entry the \
7532 JSON mask never produced"
7533 );
7534 assert_eq!(
7535 state.cache_stats().entries,
7536 2,
7537 "json_object must key its own entry, not reuse the unconstrained \
7538 one it happens to render the same prompt as"
7539 );
7540 }
7541
7542 #[tokio::test]
7554 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7555 let app = test_app_with_state(Arc::new(test_state(
7556 test_model_full_byte_vocab_with_eos(Some(0x77)),
7557 ResponseCache::new(1000, Duration::from_secs(3600)),
7558 )));
7559 let body = serde_json::json!({
7560 "model": "m",
7561 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7562 "max_tokens": 6,
7563 "temperature": 0,
7564 });
7565
7566 let stopped = post_json(&app, body.clone()).await;
7567 assert_eq!(stopped["ferrox_cache"], "miss");
7568 assert_eq!(
7569 stopped["choices"][0]["finish_reason"], "stop",
7570 "the fixture is only meaningful if the model's EOS really fires here"
7571 );
7572 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7573
7574 let mut ignoring = body.clone();
7575 ignoring["ignore_eos"] = serde_json::json!(true);
7576 let ran_on = post_json(&app, ignoring).await;
7577 assert_eq!(
7578 ran_on["ferrox_cache"], "miss",
7579 "ignore_eos is part of the key, so this body has never been answered"
7580 );
7581 assert_eq!(
7582 ran_on["usage"]["completion_tokens"], 6,
7583 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7584 );
7585 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7586 assert_ne!(
7587 ran_on["choices"][0]["message"]["content"],
7588 stopped["choices"][0]["message"]["content"]
7589 );
7590 }
7591
7592 #[tokio::test]
7598 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7599 let session_app = test_app();
7600 let manual_app = test_app();
7601
7602 let turn1 = post_json(
7604 &session_app,
7605 serde_json::json!({
7606 "model": "m",
7607 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7608 "session_id": "s1",
7609 "max_tokens": 5,
7610 "temperature": 0,
7611 }),
7612 )
7613 .await;
7614 let reply1 = turn1["choices"][0]["message"]["content"]
7615 .as_str()
7616 .unwrap()
7617 .to_string();
7618
7619 let manual_turn1 = post_json(
7624 &manual_app,
7625 serde_json::json!({
7626 "model": "m",
7627 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7628 "max_tokens": 5,
7629 "temperature": 0,
7630 }),
7631 )
7632 .await;
7633 assert_eq!(
7634 manual_turn1["choices"][0]["message"]["content"]
7635 .as_str()
7636 .unwrap(),
7637 reply1
7638 );
7639
7640 let turn2 = post_json(
7642 &session_app,
7643 serde_json::json!({
7644 "model": "m",
7645 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7646 "session_id": "s1",
7647 "max_tokens": 5,
7648 "temperature": 0,
7649 }),
7650 )
7651 .await;
7652 let reply2 = turn2["choices"][0]["message"]["content"]
7653 .as_str()
7654 .unwrap()
7655 .to_string();
7656
7657 let manual_turn2 = post_json(
7661 &manual_app,
7662 serde_json::json!({
7663 "model": "m",
7664 "messages": [
7665 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7666 {"role": "assistant", "content": reply1},
7667 {"role": "user", "content": "\u{4}\u{5}"},
7668 ],
7669 "max_tokens": 5,
7670 "temperature": 0,
7671 }),
7672 )
7673 .await;
7674 assert_eq!(
7675 manual_turn2["choices"][0]["message"]["content"]
7676 .as_str()
7677 .unwrap(),
7678 reply2,
7679 "resuming a session must produce identical output to manually resending the full history"
7680 );
7681 }
7682
7683 #[test]
7686 fn lock_cache_recovers_from_a_poisoned_mutex() {
7687 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7688
7689 let poison_cache = Arc::clone(&cache);
7690 let _ = std::thread::spawn(move || {
7691 let _guard = poison_cache.lock().unwrap();
7692 panic!("simulated panic while holding the lock");
7693 })
7694 .join();
7695
7696 let recovered = lock_cache(&cache);
7698 assert_eq!(recovered.stats().entries, 0);
7699 }
7700
7701 #[test]
7702 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7703 let mut req_body = serde_json::json!({
7704 "model": "m",
7705 "messages": [{"role": "user", "content": "hi"}],
7706 });
7707 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7708 assert!(
7709 req.is_cacheable(),
7710 "default (temperature 0) must be cacheable"
7711 );
7712
7713 req_body["temperature"] = serde_json::json!(0.8);
7714 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7715 assert!(
7716 !req.is_cacheable(),
7717 "unseeded sampling must never be cacheable"
7718 );
7719
7720 req_body["seed"] = serde_json::json!(42);
7721 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7722 assert!(
7723 req.is_cacheable(),
7724 "sampling with an explicit seed is deterministic and must be cacheable"
7725 );
7726 }
7727
7728 const GRADED: &str = "{% if reasoning_effort %}\
7732 {% if reasoning_effort not in ['low','medium','high'] %}\
7733 {{ raise_exception('unsupported effort') }}\
7734 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7735 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7736
7737 fn graded_template() -> chat_template::PromptTemplate {
7738 chat_template::PromptTemplate::from_gguf_metadata(
7739 Some(GRADED),
7740 Some("qwen3"),
7741 false,
7742 true,
7743 None,
7744 None,
7745 )
7746 }
7747
7748 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7749 serde_json::from_value(value).expect("request")
7750 }
7751
7752 #[test]
7759 fn the_reasoning_split_is_resolved_from_the_served_model_not_the_request() {
7760 let req = chat_request(serde_json::json!({
7761 "model": "some-alias",
7764 "messages": [{"role": "user", "content": "hi"}],
7765 }));
7766 let template = chat_template::PromptTemplate::plain();
7767
7768 let thinks = req
7769 .generation_params_for_template(
7770 &template,
7771 "Qwen3-8B",
7772 crate::sampling_knobs::SamplerModel::absent(),
7773 )
7774 .expect("params");
7775 assert!(
7776 thinks.reasoning.is_some(),
7777 "a thinking checkpoint must carry its format into generation"
7778 );
7779
7780 let plain = req
7781 .generation_params_for_template(
7782 &template,
7783 "Llama-3.2-1B-Instruct",
7784 crate::sampling_knobs::SamplerModel::absent(),
7785 )
7786 .expect("params");
7787 assert!(
7788 plain.reasoning.is_none(),
7789 "a checkpoint with no reasoning format must carry none, so the \
7790 usage field stays absent rather than becoming a zero"
7791 );
7792 }
7793
7794 #[test]
7801 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7802 let req = chat_request(serde_json::json!({
7803 "model": "m",
7804 "messages": [{"role": "user", "content": "hi"}],
7805 "grammar": "root ::= \"a\"+",
7806 }));
7807 req.validate_supported_fields()
7808 .expect("a valid grammar is a valid request");
7809 let params = req
7810 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7811 .expect("a valid grammar compiles at params time too");
7812 assert!(
7813 params.grammar.is_some(),
7814 "the grammar was dropped between the wire and the sampler"
7815 );
7816 assert!(
7817 params.needs_vocab_logits(),
7818 "a grammar request that may fold lm_head into a GPU argmax is \
7819 a grammar request served unconstrained"
7820 );
7821
7822 let plain = chat_request(serde_json::json!({
7823 "model": "m",
7824 "messages": [{"role": "user", "content": "hi"}],
7825 }));
7826 assert!(plain
7827 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7828 .unwrap()
7829 .grammar
7830 .is_none());
7831 }
7832
7833 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7834 chat_request(serde_json::json!({
7835 "model": "m",
7836 "messages": [{"role": "user", "content": "weather in Rome?"}],
7837 "tools": [weather_tool()],
7838 "tool_choice": tool_choice,
7839 }))
7840 }
7841
7842 #[test]
7846 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7847 for choice in [
7848 serde_json::json!("required"),
7849 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7850 ] {
7851 let req = tool_request(choice.clone());
7852 req.validate_supported_fields()
7853 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7854 let params = req
7855 .generation_params_for_template(
7856 &graded_template(),
7857 "Qwen3-8B",
7858 crate::sampling_knobs::SamplerModel::absent(),
7859 )
7860 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7861 let grammar = params
7862 .grammar
7863 .as_ref()
7864 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7865 assert!(
7866 grammar.is_awaiting_trigger(),
7867 "the model must be free to think before it calls"
7868 );
7869 assert!(
7870 !grammar.allows_eog(),
7871 "{choice} must not be able to end the turn without a call"
7872 );
7873 assert!(
7879 params.needs_vocab_logits(),
7880 "{choice} would let a backend return a token id instead of logits"
7881 );
7882 assert!(
7883 !generate::greedy_gpu_fold_allowed(¶ms),
7884 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7885 );
7886 }
7887 }
7888
7889 #[test]
7891 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7892 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7893 let req = tool_request(choice.clone());
7894 req.validate_supported_fields().expect("still supported");
7895 let params = match req.generation_params_for_template(
7896 &graded_template(),
7897 "Qwen3-8B",
7898 crate::sampling_knobs::SamplerModel::absent(),
7899 ) {
7900 Ok(p) => p,
7901 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7902 };
7903 assert!(
7904 params.grammar.is_none(),
7905 "{choice} does not force a call and must not be constrained"
7906 );
7907 }
7908 }
7909
7910 #[test]
7913 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7914 let req = chat_request(serde_json::json!({
7916 "model": "m",
7917 "messages": [{"role": "user", "content": "hi"}],
7918 "tool_choice": "required",
7919 }));
7920 let (status, _) = req
7921 .validate_supported_fields()
7922 .expect_err("nothing to call");
7923 assert_eq!(status, StatusCode::BAD_REQUEST);
7924
7925 let req =
7927 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7928 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7929 assert_eq!(status, StatusCode::BAD_REQUEST);
7930 assert_eq!(body["error"]["param"], "tool_choice");
7931
7932 let req = tool_request(serde_json::json!({"type": "function"}));
7934 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7935 assert_eq!(status, StatusCode::BAD_REQUEST);
7936
7937 let req = chat_request(serde_json::json!({
7939 "model": "m",
7940 "messages": [{"role": "user", "content": "hi"}],
7941 "tools": [weather_tool()],
7942 "tool_choice": "required",
7943 "grammar": "root ::= \"a\"+",
7944 }));
7945 let (status, _) = req
7946 .validate_supported_fields()
7947 .expect_err("a grammar and a forced call are two constraints");
7948 assert_eq!(status, StatusCode::BAD_REQUEST);
7949
7950 let req = tool_request(serde_json::json!("required"));
7956 let (status, Json(body)) = match req.generation_params_for_template(
7957 &graded_template(),
7958 "muse-glimmer-8b",
7959 crate::sampling_knobs::SamplerModel::absent(),
7960 ) {
7961 Err(e) => e,
7962 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
7963 };
7964 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7965 assert!(
7966 body["error"]["message"]
7967 .as_str()
7968 .unwrap()
7969 .contains("muse_glimmer"),
7970 "{body}"
7971 );
7972
7973 let req = tool_request(serde_json::json!("required"));
7979 let params = req
7980 .generation_params_for_template(
7981 &graded_template(),
7982 "gemma-4-E2B-it",
7983 crate::sampling_knobs::SamplerModel::absent(),
7984 )
7985 .expect("a gemma4 forced tool_choice is served");
7986 assert!(
7987 params.grammar.is_some(),
7988 "a forced tool_choice must arrive as the generation's grammar"
7989 );
7990 }
7991
7992 #[test]
7995 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
7996 let req = chat_request(serde_json::json!({
7997 "model": "m",
7998 "messages": [{"role": "user", "content": "hi"}],
7999 "grammar": "root ::= \"a",
8000 }));
8001 let (status, Json(body)) = req
8002 .validate_supported_fields()
8003 .expect_err("this does not parse");
8004 assert_eq!(status, StatusCode::BAD_REQUEST);
8005 assert_eq!(body["error"]["param"], "grammar");
8006 assert!(
8007 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8008 .is_err(),
8009 "and again at params time"
8010 );
8011 }
8012
8013 #[test]
8020 fn response_format_json_schema_becomes_the_requests_grammar() {
8021 let req = chat_request(serde_json::json!({
8022 "model": "m",
8023 "messages": [{"role": "user", "content": "hi"}],
8024 "response_format": {
8025 "type": "json_schema",
8026 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8027 },
8028 }));
8029 req.validate_supported_fields()
8030 .expect("a boolean schema converts");
8031 let params = req
8032 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8033 .expect("and compiles");
8034 let grammar = params.grammar.expect("the schema is the grammar");
8035 let mut g = (*grammar).clone();
8036 g.accept_token(0, b"true").expect("a boolean is accepted");
8037 assert!(g.allows_eog(), "and completes the parse");
8038 assert!(
8039 !params.json_object,
8040 "a schema is not the json_object character-class mask"
8041 );
8042 }
8043
8044 #[test]
8048 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8049 let req = chat_request(serde_json::json!({
8050 "model": "m",
8051 "messages": [{"role": "user", "content": "hi"}],
8052 "response_format": {
8053 "type": "json_schema",
8054 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8055 },
8056 }));
8057 let (status, Json(body)) = req
8058 .validate_supported_fields()
8059 .expect_err("minimum has no grammar in this port");
8060 assert_eq!(status, StatusCode::BAD_REQUEST);
8061 assert!(
8062 body["error"]["message"]
8063 .as_str()
8064 .expect("a message")
8065 .contains("minimum"),
8066 "the refusal must name the keyword: {body}"
8067 );
8068 assert!(
8069 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8070 .is_err(),
8071 "and again at params time"
8072 );
8073 }
8074
8075 #[test]
8081 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8082 let req = chat_request(serde_json::json!({
8083 "model": "m",
8084 "messages": [{"role": "user", "content": "hi"}],
8085 "tool_choice": "required",
8086 "tools": [{
8087 "type": "function",
8088 "function": {"name": "f", "parameters": {"type": "object"}},
8089 }],
8090 "response_format": {
8091 "type": "json_schema",
8092 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8093 },
8094 }));
8095 let (status, Json(body)) = req
8096 .validate_supported_fields()
8097 .expect_err("two constraints, one generation");
8098 assert_eq!(status, StatusCode::BAD_REQUEST);
8099 assert_eq!(body["error"]["param"], "tool_choice");
8100 }
8101
8102 #[test]
8105 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8106 let req = chat_request(serde_json::json!({
8107 "model": "m",
8108 "messages": [{"role": "user", "content": "hi"}],
8109 }));
8110 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8111 }
8112
8113 #[test]
8118 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8119 let asked = chat_request(serde_json::json!({
8120 "model": "m",
8121 "messages": [{"role": "user", "content": "hi"}],
8122 "min_p": 0.07,
8123 }));
8124 assert_eq!(
8125 asked
8126 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8127 .expect("knobs")
8128 .min_p,
8129 0.07
8130 );
8131
8132 let silent = chat_request(serde_json::json!({
8133 "model": "m",
8134 "messages": [{"role": "user", "content": "hi"}],
8135 }));
8136 assert_eq!(
8137 silent
8138 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8139 .expect("knobs")
8140 .min_p,
8141 0.0,
8142 "an unset min_p must be off, not llama.cpp's CLI default"
8143 );
8144 }
8145
8146 #[test]
8155 fn no_sampler_knob_is_missing_from_the_cache_key() {
8156 let base = serde_json::json!({
8157 "model": "m",
8158 "messages": [{"role": "user", "content": "hi"}],
8159 "seed": 1,
8160 });
8161 let key_for = |body: serde_json::Value| {
8162 let req = chat_request(body);
8163 let params = req
8164 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8165 .expect("params");
8166 req.cache_key("prompt", ¶ms)
8167 };
8168 let baseline = key_for(base.clone());
8169 for (knob, value) in [
8170 ("temperature", serde_json::json!(0.5)),
8171 ("top_p", serde_json::json!(0.9)),
8172 ("min_p", serde_json::json!(0.05)),
8173 ("top_k", serde_json::json!(40)),
8174 ("repetition_penalty", serde_json::json!(1.1)),
8175 ("presence_penalty", serde_json::json!(0.3)),
8176 ("frequency_penalty", serde_json::json!(0.3)),
8177 (
8178 "samplers",
8179 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8180 ),
8181 ] {
8182 let mut body = base.clone();
8183 body[knob] = value;
8184 assert_ne!(
8185 key_for(body),
8186 baseline,
8187 "`{knob}` is not in the cache key: two requests differing \
8188 only in it would share one cached answer"
8189 );
8190 }
8191 }
8192
8193 #[test]
8203 fn no_constraint_is_missing_from_the_cache_key() {
8204 let base = serde_json::json!({
8205 "model": "m",
8206 "messages": [{"role": "user", "content": "pick one"}],
8207 });
8208 let key_for = |body: serde_json::Value| {
8209 let req = chat_request(body);
8210 let params = req
8211 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8212 .expect("params");
8213 req.cache_key("prompt", ¶ms)
8214 };
8215 let baseline = key_for(base.clone());
8216 for (field, value) in [
8217 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8218 (
8219 "response_format",
8220 serde_json::json!({"type": "json_object"}),
8221 ),
8222 (
8223 "response_format",
8224 serde_json::json!({"type": "json_schema", "json_schema": {
8225 "name": "answer",
8226 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8227 }}),
8228 ),
8229 ("ignore_eos", serde_json::json!(true)),
8230 ("stop", serde_json::json!(["\n"])),
8231 ("max_tokens", serde_json::json!(7)),
8232 ] {
8233 let mut body = base.clone();
8234 body[field] = value.clone();
8235 assert_ne!(
8236 key_for(body),
8237 baseline,
8238 "`{field}: {value}` is not in the cache key: two requests \
8239 differing only in it would share one cached answer"
8240 );
8241 }
8242 }
8243
8244 #[test]
8248 fn an_explicit_zero_output_budget_is_a_client_error() {
8249 let req = chat_request(serde_json::json!({
8250 "model": "m",
8251 "messages": [{"role": "user", "content": "hi"}],
8252 "max_tokens": 0,
8253 }));
8254 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8255 assert_eq!(status, StatusCode::BAD_REQUEST);
8256 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8257 }
8258
8259 #[test]
8262 fn a_request_can_turn_thinking_off() {
8263 let template = graded_template();
8264 for body in [
8265 serde_json::json!({
8266 "model": "m",
8267 "messages": [{"role": "user", "content": "hi"}],
8268 "reasoning_effort": "none",
8269 }),
8270 serde_json::json!({
8271 "model": "m",
8272 "messages": [{"role": "user", "content": "hi"}],
8273 "thinking": {"type": "disabled"},
8274 }),
8275 ] {
8276 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8277 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8278 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8279 assert!(!kwargs.contains_key("reasoning_effort"));
8282 }
8283 }
8284
8285 #[test]
8288 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8289 let template = graded_template();
8290 let kwargs = chat_request(serde_json::json!({
8291 "model": "m",
8292 "messages": [{"role": "user", "content": "hi"}],
8293 "reasoning_effort": "high",
8294 "thinking": {"type": "disabled"},
8295 }))
8296 .resolve_template_kwargs(&template);
8297 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8298 assert!(!kwargs.contains_key("reasoning_effort"));
8299 }
8300
8301 #[test]
8304 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8305 let req = chat_request(serde_json::json!({
8306 "model": "m",
8307 "messages": [{"role": "user", "content": "hi"}],
8308 "thinking": {"type": "disable"},
8309 }));
8310 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8311 assert_eq!(status, StatusCode::BAD_REQUEST);
8312 }
8313
8314 #[test]
8317 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8318 let template = graded_template();
8319 let kwargs = chat_request(serde_json::json!({
8320 "model": "m",
8321 "messages": [{"role": "user", "content": "hi"}],
8322 "reasoning_effort": "none",
8323 "chat_template_kwargs": {"enable_thinking": true},
8324 }))
8325 .resolve_template_kwargs(&template);
8326 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8327 }
8328
8329 #[test]
8333 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8334 let template = graded_template();
8335 let req = chat_request(serde_json::json!({
8336 "model": "m",
8337 "messages": [{"role": "user", "content": "hi"}],
8338 "reasoning_effort": "minimal",
8339 }));
8340 let kwargs = req.resolve_template_kwargs(&template);
8341 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8342 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8343 assert!(prompt.starts_with("E:low|"), "{prompt}");
8344 }
8345
8346 #[test]
8350 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8351 let template = graded_template();
8352 let req = chat_request(serde_json::json!({
8353 "model": "m",
8354 "messages": [{"role": "user", "content": "hi"}],
8355 "chat_template_kwargs": {"reasoning_effort": "none"},
8356 }));
8357 let kwargs = req.resolve_template_kwargs(&template);
8358 assert!(!kwargs.contains_key("reasoning_effort"));
8359 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8360 assert_eq!(prompt, "hi");
8361 }
8362
8363 #[test]
8366 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8367 let template = graded_template();
8368 let req = chat_request(serde_json::json!({
8369 "model": "m",
8370 "messages": [{"role": "user", "content": "hi"}],
8371 "reasoning_effort": "low",
8372 "chat_template_kwargs": {"reasoning_effort": "high"},
8373 }));
8374 assert_eq!(
8375 req.resolve_template_kwargs(&template)["reasoning_effort"],
8376 serde_json::json!("high")
8377 );
8378 }
8379
8380 #[test]
8384 fn offering_tools_turns_thinking_on_by_itself() {
8385 let template = graded_template();
8386 let quiet = chat_request(serde_json::json!({
8387 "model": "m",
8388 "messages": [{"role": "user", "content": "hi"}],
8389 }));
8390 assert!(!quiet
8391 .resolve_template_kwargs(&template)
8392 .contains_key("enable_thinking"));
8393
8394 let with_tools = chat_request(serde_json::json!({
8395 "model": "m",
8396 "messages": [{"role": "user", "content": "hi"}],
8397 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8398 }));
8399 let kwargs = with_tools.resolve_template_kwargs(&template);
8400 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8401 let prompt =
8402 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8403 assert!(prompt.starts_with("THINK|"), "{prompt}");
8404 }
8405
8406 #[test]
8411 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8412 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8413 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8414 Some("qwen3"),
8415 false,
8416 true,
8417 None,
8418 None,
8419 );
8420 let req = chat_request(serde_json::json!({
8421 "model": "m",
8422 "messages": [{"role": "user", "content": "hi"}],
8423 "chat_template_kwargs": {"enable_thinking": true},
8424 }));
8425 let kwargs = req.resolve_template_kwargs(&opener);
8426 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8427 assert!(prompt.ends_with("<think>"), "{prompt}");
8428
8429 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8432 let (message, _) = build_response_message(
8433 "weighing it up</think>Paris.".to_string(),
8434 &[],
8435 posture,
8436 "stop",
8437 );
8438 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8439 assert_eq!(message.content.as_deref(), Some("Paris."));
8440
8441 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8444 let (message, _) = build_response_message(
8445 "weighing it up</think>Paris.".to_string(),
8446 &[],
8447 closed,
8448 "stop",
8449 );
8450 assert_eq!(message.reasoning_content, None);
8451 }
8452
8453 #[test]
8454 fn stop_param_accepts_both_single_string_and_array() {
8455 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8456 "model": "m",
8457 "messages": [{"role": "user", "content": "hi"}],
8458 "stop": "END",
8459 }))
8460 .unwrap();
8461 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8462
8463 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8464 "model": "m",
8465 "messages": [{"role": "user", "content": "hi"}],
8466 "stop": ["A", "B"],
8467 }))
8468 .unwrap();
8469 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8470 }
8471
8472 #[test]
8473 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8474 let model = test_model();
8475 let result = run_generation(
8476 &model,
8477 "hello",
8478 &greedy_params(4),
8479 None,
8480 None,
8481 None,
8482 None,
8483 None,
8484 None,
8485 );
8486 assert!(matches!(
8487 result,
8488 Err(generate::DecodeError::TokenOutOfVocab { .. })
8489 ));
8490 }
8491
8492 #[test]
8496 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8497 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8499 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8500
8501 let holder_pool = Arc::clone(&pool);
8502 let holder = std::thread::spawn(move || {
8503 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8504 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8506 drop(held);
8507 });
8508 std::thread::sleep(Duration::from_millis(15));
8509
8510 let config = generate::KvPoolConfig {
8511 pool,
8512 queue_wait: Duration::ZERO,
8513 };
8514 let result = run_generation(
8515 &model,
8516 &prompt,
8517 &greedy_params(4),
8518 Some(&config),
8519 None,
8520 None,
8521 None,
8522 None,
8523 None,
8524 );
8525 assert!(matches!(
8526 result,
8527 Err(generate::DecodeError::KvPoolExhausted)
8528 ));
8529
8530 let (status, _body) = decode_error_response(result.unwrap_err());
8531 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8532 holder.join().unwrap();
8533 }
8534
8535 #[test]
8543 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8544 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8546 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
8548 let config = generate::KvPoolConfig {
8549 pool,
8550 queue_wait: Duration::ZERO,
8551 };
8552
8553 let result = run_generation(
8554 &model,
8555 &prompt,
8556 &greedy_params(4),
8557 Some(&config),
8558 None,
8559 None,
8560 None,
8561 None,
8562 None,
8563 );
8564 let err = result.expect_err("one block cannot hold two layers' caches");
8565 assert!(
8566 matches!(
8567 &err,
8568 generate::DecodeError::KvBudgetExceeded { binding, .. }
8569 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
8570 ),
8571 "expected an immovable device-memory refusal, got {err:?}"
8572 );
8573 let (status, _body) = decode_error_response(err);
8574 assert_eq!(status, StatusCode::BAD_REQUEST);
8575 }
8576
8577 #[test]
8583 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8584 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8585 queued: 512,
8586 cap: 512,
8587 });
8588 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8589 assert_eq!(body["error"]["retry_after_seconds"], 1);
8590 let message = body["error"]["message"].as_str().expect("message");
8591 assert!(message.contains("512"), "{message}");
8592 }
8593
8594 #[test]
8595 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8596 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8597 token: 99,
8598 vocab_size: 32,
8599 });
8600 assert!(
8601 body["error"]["retry_after_seconds"].is_null(),
8602 "retrying a prompt this model cannot tokenize never helps"
8603 );
8604 }
8605
8606 #[test]
8607 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8608 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8609 token: 99,
8610 vocab_size: 32,
8611 });
8612 assert_eq!(status, StatusCode::BAD_REQUEST);
8613 }
8614
8615 #[test]
8616 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8617 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8619 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
8620 let config = generate::KvPoolConfig {
8621 pool: pool.clone(),
8622 queue_wait: Duration::ZERO,
8623 };
8624
8625 let (_, finish, _usage) = run_generation(
8626 &model,
8627 &prompt,
8628 &greedy_params(4),
8629 Some(&config),
8630 None,
8631 None,
8632 None,
8633 None,
8634 None,
8635 )
8636 .unwrap();
8637 assert_eq!(finish, FinishReason::Length);
8638 assert_eq!(
8639 pool.lock().unwrap().free_blocks(),
8640 2,
8641 "a completed request must return its blocks to the pool"
8642 );
8643 }
8644
8645 #[tokio::test]
8650 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8651 let model = Arc::new(test_model());
8652 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8653
8654 let mut handles = Vec::new();
8655 for _ in 0..8 {
8656 let model = Arc::clone(&model);
8657 let prompt = prompt.clone();
8658 handles.push(tokio::task::spawn_blocking(move || {
8659 run_generation(
8660 &model,
8661 &prompt,
8662 &greedy_params(6),
8663 None,
8664 None,
8665 None,
8666 None,
8667 None,
8668 None,
8669 )
8670 .unwrap()
8671 }));
8672 }
8673
8674 let mut results = Vec::new();
8675 for h in handles {
8676 results.push(h.await.unwrap());
8677 }
8678 for r in &results[1..] {
8683 assert_eq!(r.0, results[0].0, "decoded chunks must match");
8684 assert_eq!(r.1, results[0].1, "finish reason must match");
8685 assert_eq!(
8686 r.2.prompt_tokens, results[0].2.prompt_tokens,
8687 "prompt token count must match"
8688 );
8689 assert_eq!(
8690 r.2.completion_tokens, results[0].2.completion_tokens,
8691 "completion token count must match"
8692 );
8693 }
8694 }
8695
8696 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8703 let mut header_entries = Vec::new();
8704 let mut data = Vec::new();
8705 for (name, shape, values) in tensors {
8706 let start = data.len();
8707 for v in values {
8708 data.extend_from_slice(&v.to_le_bytes());
8709 }
8710 let end = data.len();
8711 let shape_str = shape
8712 .iter()
8713 .map(|d| d.to_string())
8714 .collect::<Vec<_>>()
8715 .join(",");
8716 header_entries.push(format!(
8717 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8718 ));
8719 }
8720 let header = format!("{{{}}}", header_entries.join(","));
8721 let header_bytes = header.as_bytes();
8722 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8723 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8724 out.extend_from_slice(header_bytes);
8725 out.extend_from_slice(&data);
8726 out
8727 }
8728
8729 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8740 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8741 use ferrox_models::kimi_loader::KimiRealHparams;
8742 use ferrox_moe::{GatingFunction, MoeLayerConfig};
8743
8744 let hidden_dim = 8;
8745 let kda_num_heads = 2;
8746 let kda_head_dim = 3;
8747 let kda_proj = kda_num_heads * kda_head_dim;
8748 let conv_kernel = 4;
8749 let dense_intermediate = 5;
8750 let vocab_size = 256;
8754 let mla_num_heads = 1;
8755 let mla_q_lora_rank = 2;
8756 let mla_kv_lora_rank = 2;
8757 let mla_qk_nope_head_dim = 2;
8758 let mla_qk_rope_head_dim = 2;
8759 let mla_v_head_dim = 2;
8760
8761 let model_cfg = ferrox_models::ModelConfig {
8762 rope_layers: ferrox_models::rope_layers::RopeLayers::All,
8763 name: "synthetic-kimi-server-test",
8764 n_layers: 1,
8765 hidden_dim,
8766 n_heads: 1,
8767 n_kv_heads: 1,
8768 head_dim: 4,
8769 vocab_size,
8770 rope_theta: 10000.0,
8771 rms_norm_eps: 1e-5,
8772 sliding_window: None,
8773 moe: MoeLayerConfig {
8774 expert_weights_scale: 1.0,
8775 n_experts: 1,
8776 n_experts_active: 1,
8777 n_shared_experts: 0,
8778 hidden_dim,
8779 expert_ffn_dim: 4,
8780 gating: GatingFunction::Sigmoid,
8781 norm_topk_prob: true,
8782 expert_group_count: None,
8783 expert_group_used_count: None,
8784 },
8785 n_dense_leading_layers: 1,
8790 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8791 kda_layers: vec![1],
8792 full_attn_layers: vec![],
8793 mla: MlaConfig {
8794 num_heads: mla_num_heads,
8795 q_lora_rank: mla_q_lora_rank,
8796 kv_lora_rank: mla_kv_lora_rank,
8797 qk_nope_head_dim: mla_qk_nope_head_dim,
8798 qk_rope_head_dim: mla_qk_rope_head_dim,
8799 v_head_dim: mla_v_head_dim,
8800 use_output_gate: true,
8801 rope: None,
8802 },
8803 kda: KdaConfig {
8804 num_heads: kda_num_heads,
8805 head_dim: kda_head_dim,
8806 short_conv_kernel_size: conv_kernel,
8807 gate_lower_bound: -5.0,
8808 use_full_rank_gate: true,
8809 },
8810 }),
8811 rope_freqs: None,
8812 rope_attn_factor: 1.0,
8813 rope_dim: None,
8814 rope_freqs_long: None,
8815 rope_freqs_short: None,
8816 rope_orig_ctx: None,
8817 rope_layout: ferrox_models::config::RopeLayout::Neox,
8818 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
8819 swa_pattern: None,
8820 swa_dense_first: false,
8821 attn_logit_softcap: None,
8822 final_logit_softcap: None,
8823 embedding_scale: None,
8824 residual_scale: None,
8825 logit_multiplier: None,
8826 attention_scale: None,
8827 rope_theta_swa: None,
8828 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
8829 best_effort_fields: &["synthetic test config, not a real preset"],
8830 };
8831 let hp = KimiRealHparams {
8832 hidden_dim,
8833 kda_num_heads,
8834 kda_head_dim,
8835 mla_num_heads,
8836 mla_q_lora_rank,
8837 mla_kv_lora_rank,
8838 mla_qk_nope_head_dim,
8839 mla_qk_rope_head_dim,
8840 mla_v_head_dim,
8841 dense_intermediate_dim: dense_intermediate,
8842 moe_hidden_dim: hidden_dim,
8843 moe_intermediate_dim: 4,
8844 n_experts: 1,
8845 num_shared_experts: 0,
8846 };
8847
8848 let prefix = "language_model.model.layers.0";
8852 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8853 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8854 name: String,
8855 shape: Vec<usize>,
8856 n: usize| {
8857 tensors.push((name, shape, vec![0.01f32; n]));
8858 };
8859 push(
8860 &mut tensors,
8861 format!("{prefix}.input_layernorm.weight"),
8862 vec![hidden_dim],
8863 hidden_dim,
8864 );
8865 push(
8866 &mut tensors,
8867 format!("{prefix}.post_attention_layernorm.weight"),
8868 vec![hidden_dim],
8869 hidden_dim,
8870 );
8871 push(
8872 &mut tensors,
8873 format!("{prefix}.self_attention_res_norm.weight"),
8874 vec![hidden_dim],
8875 hidden_dim,
8876 );
8877 push(
8878 &mut tensors,
8879 format!("{prefix}.self_attention_res_proj.weight"),
8880 vec![1, hidden_dim],
8881 hidden_dim,
8882 );
8883 push(
8884 &mut tensors,
8885 format!("{prefix}.mlp_res_norm.weight"),
8886 vec![hidden_dim],
8887 hidden_dim,
8888 );
8889 push(
8890 &mut tensors,
8891 format!("{prefix}.mlp_res_proj.weight"),
8892 vec![1, hidden_dim],
8893 hidden_dim,
8894 );
8895 push(
8896 &mut tensors,
8897 format!("{prefix}.self_attn.q_proj.weight"),
8898 vec![kda_proj, hidden_dim],
8899 kda_proj * hidden_dim,
8900 );
8901 push(
8902 &mut tensors,
8903 format!("{prefix}.self_attn.k_proj.weight"),
8904 vec![kda_proj, hidden_dim],
8905 kda_proj * hidden_dim,
8906 );
8907 push(
8908 &mut tensors,
8909 format!("{prefix}.self_attn.v_proj.weight"),
8910 vec![kda_proj, hidden_dim],
8911 kda_proj * hidden_dim,
8912 );
8913 push(
8914 &mut tensors,
8915 format!("{prefix}.self_attn.q_conv1d.weight"),
8916 vec![kda_proj, 1, conv_kernel],
8917 kda_proj * conv_kernel,
8918 );
8919 push(
8920 &mut tensors,
8921 format!("{prefix}.self_attn.k_conv1d.weight"),
8922 vec![kda_proj, 1, conv_kernel],
8923 kda_proj * conv_kernel,
8924 );
8925 push(
8926 &mut tensors,
8927 format!("{prefix}.self_attn.v_conv1d.weight"),
8928 vec![kda_proj, 1, conv_kernel],
8929 kda_proj * conv_kernel,
8930 );
8931 push(
8932 &mut tensors,
8933 format!("{prefix}.self_attn.A_log"),
8934 vec![kda_num_heads],
8935 kda_num_heads,
8936 );
8937 push(
8938 &mut tensors,
8939 format!("{prefix}.self_attn.f_a_proj.weight"),
8940 vec![kda_head_dim, hidden_dim],
8941 kda_head_dim * hidden_dim,
8942 );
8943 push(
8944 &mut tensors,
8945 format!("{prefix}.self_attn.f_b_proj.weight"),
8946 vec![kda_proj, kda_head_dim],
8947 kda_proj * kda_head_dim,
8948 );
8949 push(
8950 &mut tensors,
8951 format!("{prefix}.self_attn.dt_bias"),
8952 vec![kda_proj],
8953 kda_proj,
8954 );
8955 push(
8956 &mut tensors,
8957 format!("{prefix}.self_attn.b_proj.weight"),
8958 vec![kda_num_heads, hidden_dim],
8959 kda_num_heads * hidden_dim,
8960 );
8961 push(
8962 &mut tensors,
8963 format!("{prefix}.self_attn.g_proj.weight"),
8964 vec![kda_proj, hidden_dim],
8965 kda_proj * hidden_dim,
8966 );
8967 push(
8968 &mut tensors,
8969 format!("{prefix}.self_attn.o_norm.weight"),
8970 vec![kda_head_dim],
8971 kda_head_dim,
8972 );
8973 push(
8974 &mut tensors,
8975 format!("{prefix}.self_attn.o_proj.weight"),
8976 vec![hidden_dim, kda_proj],
8977 hidden_dim * kda_proj,
8978 );
8979 push(
8980 &mut tensors,
8981 format!("{prefix}.mlp.gate_proj.weight"),
8982 vec![dense_intermediate, hidden_dim],
8983 dense_intermediate * hidden_dim,
8984 );
8985 push(
8986 &mut tensors,
8987 format!("{prefix}.mlp.up_proj.weight"),
8988 vec![dense_intermediate, hidden_dim],
8989 dense_intermediate * hidden_dim,
8990 );
8991 push(
8992 &mut tensors,
8993 format!("{prefix}.mlp.down_proj.weight"),
8994 vec![hidden_dim, dense_intermediate],
8995 hidden_dim * dense_intermediate,
8996 );
8997 push(
8998 &mut tensors,
8999 "language_model.model.embed_tokens.weight".to_string(),
9000 vec![vocab_size, hidden_dim],
9001 vocab_size * hidden_dim,
9002 );
9003 push(
9004 &mut tensors,
9005 "language_model.lm_head.weight".to_string(),
9006 vec![vocab_size, hidden_dim],
9007 vocab_size * hidden_dim,
9008 );
9009 push(
9010 &mut tensors,
9011 "language_model.model.norm.weight".to_string(),
9012 vec![hidden_dim],
9013 hidden_dim,
9014 );
9015 push(
9016 &mut tensors,
9017 "language_model.model.output_attn_res_norm.weight".to_string(),
9018 vec![hidden_dim],
9019 hidden_dim,
9020 );
9021 push(
9022 &mut tensors,
9023 "language_model.model.output_attn_res_proj.weight".to_string(),
9024 vec![1, hidden_dim],
9025 hidden_dim,
9026 );
9027
9028 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9041 let dir = std::env::temp_dir().join(format!(
9042 "ferrox_server_kimi_e2e_test_{}_{}",
9043 std::process::id(),
9044 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9045 ));
9046 std::fs::create_dir_all(&dir).unwrap();
9047 let shard_bytes = write_safetensors_shard(&tensors);
9048 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9049 let map_entries: Vec<String> = tensors
9050 .iter()
9051 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9052 .collect();
9053 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9054 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9055
9056 use base64::Engine;
9060 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9061 .map(|b| {
9062 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9063 format!("{b64} {b}")
9064 })
9065 .collect();
9066 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9067
9068 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9069 .expect("must load the synthetic Kimi checkpoint end to end");
9070 std::fs::remove_dir_all(&dir).ok();
9071 loaded
9072 }
9073
9074 #[test]
9082 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9083 let loaded = build_synthetic_kimi_loaded();
9084 let state = build_app_state(
9085 StartupModels {
9086 loaded: model::LoadedModel::Kimi(loaded),
9087 embedding: None,
9088 },
9089 None,
9090 None,
9091 None,
9092 false,
9093 None,
9094 Arc::new(health::Detection::ready(health::probe_backends())),
9095 );
9096 let active = state.active().expect("a freshly built state has a model");
9097 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9098 assert!(!active.is_synthetic());
9099
9100 let (_chunks, finish, _usage) = run_generation(
9101 active.generative().unwrap(),
9102 "hi",
9103 &greedy_params(5),
9104 None,
9105 None,
9106 None,
9107 None,
9108 None,
9109 None,
9110 )
9111 .expect("a real Kimi checkpoint must generate without error");
9112 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9113 }
9114
9115 #[test]
9128 fn a_grammar_constrains_the_engine_decode_path() {
9129 let loaded = build_synthetic_kimi_loaded();
9130 let state = build_app_state(
9131 StartupModels {
9132 loaded: model::LoadedModel::Kimi(loaded),
9133 embedding: None,
9134 },
9135 None,
9136 None,
9137 None,
9138 false,
9139 None,
9140 Arc::new(health::Detection::ready(health::probe_backends())),
9141 );
9142 let active = state.active().expect("a freshly built state has a model");
9143
9144 let run = |grammar: Option<&str>| {
9145 let mut params = greedy_params(6);
9146 params.grammar = grammar.map(|src| {
9147 Arc::new(
9148 ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
9149 .expect("test grammar parses"),
9150 )
9151 });
9152 run_generation(
9153 active.generative().unwrap(),
9154 "hi",
9155 ¶ms,
9156 None,
9157 None,
9158 None,
9159 None,
9160 None,
9161 None,
9162 )
9163 };
9164
9165 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
9166 let unconstrained = chunks.concat();
9167 assert!(
9168 unconstrained.chars().any(|c| c != 'a'),
9169 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9170 constrained run below would prove nothing"
9171 );
9172
9173 let (chunks, finish, _) =
9174 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9175 let constrained = chunks.concat();
9176 assert!(
9177 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9178 "the engine decode path served text its grammar forbids ({constrained:?}): \
9179 the constraint was dropped between `generate_engine` and the sampler"
9180 );
9181 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9182 }
9183
9184 #[test]
9194 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9195 let loaded = build_synthetic_kimi_loaded();
9196 let state = build_app_state(
9197 StartupModels {
9198 loaded: model::LoadedModel::Kimi(loaded),
9199 embedding: None,
9200 },
9201 None,
9202 None,
9203 None,
9204 false,
9205 None,
9206 Arc::new(health::Detection::ready(health::probe_backends())),
9207 );
9208
9209 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
9210 let kv_pool_config = generate::KvPoolConfig {
9211 pool: pool.clone(),
9212 queue_wait: Duration::ZERO,
9213 };
9214 let pc = Mutex::new(PrefixCache::new(4));
9215
9216 run_generation(
9217 state
9218 .active()
9219 .expect("a freshly built state has a model")
9220 .generative()
9221 .unwrap(),
9222 "hi",
9223 &greedy_params(5),
9224 Some(&kv_pool_config),
9225 None,
9226 Some(&pc),
9227 None,
9228 None,
9229 None,
9230 )
9231 .expect("a real Kimi checkpoint must generate without error");
9232
9233 assert_eq!(
9234 pool.lock().unwrap().free_blocks(),
9235 4,
9236 "the KV pool must be completely untouched by a Kimi request"
9237 );
9238 let stats = pc.lock().unwrap().stats();
9239 assert_eq!(
9240 stats.hits + stats.misses,
9241 0,
9242 "the prefix cache must never be consulted for a Kimi request"
9243 );
9244 }
9245}