1mod admin;
38mod anthropic;
39mod attribution;
40mod budget;
41mod cache_admin;
42mod cancel;
43mod chat_template;
44mod completion;
45mod conversations;
46mod decode_task;
47mod generate;
48mod grammar_request;
49mod health;
50mod journal;
51mod json_mode;
52mod limits;
53mod mcp;
54mod model;
55mod openai_extra;
56mod output;
57mod policy;
58mod response_cache;
59pub(crate) mod responses;
60mod resume;
61mod sample_step;
62mod sampling_knobs;
63mod security;
64mod serving;
65mod session;
66mod sse;
67mod stats;
68mod stop;
69mod stream_events;
70mod tasks;
71mod tool_grammar;
72mod unsupported_sampling;
73
74use std::cell::RefCell;
75use std::convert::Infallible;
76use std::fmt;
77use std::net::{IpAddr, Ipv4Addr, SocketAddr};
78use std::path::PathBuf;
79use std::rc::Rc;
80use std::str::FromStr;
81use std::sync::{Arc, Mutex, MutexGuard};
82use std::time::Duration;
83
84use axum::{
85 extract::State,
86 http::StatusCode,
87 response::sse::{Event, Sse},
88 response::{IntoResponse, Response},
89 routing::{get, post},
90 Json, Router,
91};
92use clap::{Parser, ValueEnum};
93use serde::{Deserialize, Serialize};
94
95use ferrox_core::cache::KvBlockPool;
96use ferrox_models::kimi_tokenizer::KimiTokenizer;
97use ferrox_models::sampling::SamplingParams;
98use ferrox_models::tokenizer::StopTokens;
99use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
100use generate::{FinishReason, GenerationParams};
101use model::ServerTokenizer;
102use response_cache::{CacheKey, ResponseCache};
103use sampling_knobs::SamplingKnobs;
104
105#[derive(Parser, Debug, PartialEq)]
109#[command(
116 name = "ferrox-server",
117 about = "OpenAI-compatible Ferrox inference server"
118)]
119pub struct ServerArgs {
120 #[arg(short = 'm', long = "model", value_name = "FILE")]
122 model: Option<String>,
123
124 #[arg(long, value_name = "HOST")]
126 host: Option<IpAddr>,
127
128 #[arg(long, value_name = "PORT")]
133 port: Option<u16>,
134
135 #[arg(short = 't', long = "threads", value_name = "N")]
137 threads: Option<usize>,
138
139 #[arg(
141 long = "device",
142 visible_alias = "dev",
143 value_name = "DEVICE",
144 ignore_case = true
145 )]
146 device: Option<OffloadDevice>,
147
148 #[arg(long = "list-devices", default_value_t = false)]
150 list_devices: bool,
151
152 #[arg(
157 long = "n-gpu-layers",
158 visible_aliases = ["gpu-layers", "ngl"],
159 value_name = "N"
160 )]
161 n_gpu_layers: Option<GpuLayers>,
162
163 #[arg(long = "mcp-config", value_name = "PATH")]
165 mcp_config: Option<PathBuf>,
166
167 #[arg(long = "exit-on-stdin-close", default_value_t = false)]
175 exit_on_stdin_close: bool,
176
177 #[arg(long = "allow-multiple-instances", default_value_t = false)]
182 allow_multiple_instances: bool,
183}
184
185impl ServerArgs {
186 pub fn parse_llama_style<I>(argv: I) -> Self
193 where
194 I: IntoIterator<Item = String>,
195 {
196 Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
197 }
198}
199
200#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
201enum OffloadDevice {
202 Auto,
203 None,
204 Cpu,
205 Metal,
206 Cuda,
207}
208
209#[derive(Debug, Clone, Copy, PartialEq, Eq)]
210enum GpuLayers {
211 Auto,
212 All,
213 Count(u32),
214}
215
216impl GpuLayers {
217 fn offload_enabled(self) -> bool {
218 !matches!(self, Self::Count(0))
219 }
220}
221
222impl FromStr for GpuLayers {
223 type Err = String;
224
225 fn from_str(value: &str) -> Result<Self, Self::Err> {
226 match value {
227 "auto" => Ok(Self::Auto),
228 "all" => Ok(Self::All),
229 _ => value
230 .parse::<u32>()
231 .map(Self::Count)
232 .map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
233 }
234 }
235}
236
237impl fmt::Display for GpuLayers {
238 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
239 match self {
240 Self::Auto => f.write_str("auto"),
241 Self::All => f.write_str("all"),
242 Self::Count(value) => value.fmt(f),
243 }
244 }
245}
246
247pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
256
257pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
260
261fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
262 args.into_iter()
263 .map(|arg| match arg.as_str() {
264 "-ngl" => "--n-gpu-layers".into(),
265 "-dev" => "--device".into(),
266 _ => arg,
267 })
268 .collect()
269}
270
271fn print_available_devices() {
272 println!("Available devices:");
273 println!(" CPU");
274
275 let metal = ferrox_metal::MetalProfile::detect();
276 if let Some(name) = metal.device_name {
277 println!(" Metal: {name}");
278 }
279
280 let cuda = ferrox_cuda::HardwareProfile::detect();
281 if cuda.cuda_available {
282 let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
283 println!(" CUDA: {name}");
284 if cuda.cuda_device_count > 1 {
285 println!(" ({} devices detected)", cuda.cuda_device_count);
286 }
287 }
288}
289
290fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
291 if args.host.is_none() && args.port.is_none() {
292 return None;
293 }
294
295 let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
296 let host = args
297 .host
298 .or_else(|| existing.map(|addr| addr.ip()))
299 .unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
300 let port = args
301 .port
302 .or_else(|| existing.map(|addr| addr.port()))
303 .unwrap_or(8383);
304 Some(SocketAddr::new(host, port).to_string())
305}
306
307fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
308 if let Some(model) = &args.model {
309 unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
311 }
312
313 if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
314 unsafe { std::env::set_var("FERROX_ADDR", addr) };
316 }
317
318 if let Some(threads) = args.threads {
319 if threads == 0 {
320 anyhow::bail!("--threads must be greater than zero");
321 }
322 unsafe {
324 std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
325 std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
326 }
327 }
328
329 if args.device.is_none() && args.n_gpu_layers.is_none() {
330 } else {
332 let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
333 let device = if layers.offload_enabled() {
334 args.device.unwrap_or(OffloadDevice::Auto)
335 } else {
336 OffloadDevice::None
337 };
338
339 match device {
340 OffloadDevice::None | OffloadDevice::Cpu => unsafe {
341 std::env::set_var("FERROX_METAL", "0");
342 std::env::set_var("FERROX_METAL_ATTN", "0");
343 std::env::set_var("FERROX_CUDA", "0");
344 },
345 OffloadDevice::Auto => unsafe {
346 std::env::set_var("FERROX_METAL", "auto");
347 std::env::set_var("FERROX_CUDA", "auto");
348 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
349 std::env::set_var("FERROX_METAL_ATTN", "1");
350 }
351 },
352 OffloadDevice::Metal => {
353 #[cfg(not(feature = "metal"))]
354 {
355 anyhow::bail!(
356 "Metal requested but this binary was built without --features metal"
357 );
358 }
359 #[cfg(feature = "metal")]
360 {
361 if !ferrox_metal::MetalProfile::detect().available {
362 anyhow::bail!("Metal requested but no Metal device is available");
363 }
364 unsafe {
365 std::env::set_var("FERROX_METAL", "1");
366 if std::env::var_os("FERROX_METAL_ATTN").is_none() {
367 std::env::set_var("FERROX_METAL_ATTN", "1");
368 }
369 std::env::set_var("FERROX_CUDA", "0");
370 }
371 }
372 }
373 OffloadDevice::Cuda => {
374 #[cfg(not(feature = "cuda"))]
375 {
376 anyhow::bail!(
377 "CUDA requested but this binary was built without --features cuda"
378 );
379 }
380 #[cfg(feature = "cuda")]
381 {
382 if !ferrox_cuda::HardwareProfile::detect().cuda_available {
383 anyhow::bail!("CUDA requested but no CUDA device is available");
384 }
385 unsafe {
386 std::env::set_var("FERROX_CUDA", "1");
387 std::env::set_var("FERROX_METAL", "0");
388 std::env::set_var("FERROX_METAL_ATTN", "0");
389 }
390 }
391 }
392 }
393 }
394
395 Ok(())
396}
397
398#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
407 Gguf(GgufModel),
408 Kimi(KimiModel),
409 Mla(MlaModel),
410 Gemma4(Gemma4Model),
411 Glm52(Glm52Model),
412}
413
414pub(crate) struct GgufModel {
415 decoder: Arc<Decoder>,
416 tokenizer: Arc<ServerTokenizer>,
417 stop_tokens: StopTokens,
418 bos_id: Option<usize>,
419 is_synthetic: bool,
420 chat_template: chat_template::PromptTemplate,
421}
422
423pub(crate) struct KimiModel {
424 engine: KimiEngine,
425 tokenizer: KimiTokenizer,
426 stop_tokens: StopTokens,
427 chat_template: chat_template::PromptTemplate,
428}
429
430pub(crate) struct MlaModel {
431 engine: MlaEngine,
432 tokenizer: ServerTokenizer,
433 stop_tokens: StopTokens,
434 bos_id: Option<usize>,
435 name: String,
436 chat_template: chat_template::PromptTemplate,
437}
438
439pub(crate) struct Gemma4Model {
440 engine: Gemma4Engine,
441 tokenizer: ServerTokenizer,
442 stop_tokens: StopTokens,
443 bos_id: Option<usize>,
444 name: String,
445 chat_template: chat_template::PromptTemplate,
446}
447
448pub(crate) struct Glm52Model {
449 engine: ferrox_models::Glm52Engine,
450 tokenizer: ServerTokenizer,
451 stop_tokens: StopTokens,
452 bos_id: Option<usize>,
453 name: String,
454 chat_template: chat_template::PromptTemplate,
455}
456
457impl Model {
458 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
459 match self {
460 Model::Gguf(m) => m.chat_template.clone(),
461 Model::Kimi(m) => m.chat_template.clone(),
462 Model::Mla(m) => m.chat_template.clone(),
463 Model::Gemma4(m) => m.chat_template.clone(),
464 Model::Glm52(m) => m.chat_template.clone(),
465 }
466 }
467
468 fn is_synthetic(&self) -> bool {
473 match self {
474 Model::Gguf(m) => m.is_synthetic,
475 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
476 }
477 }
478
479 fn tokenizer_kind(&self) -> &'static str {
480 match self {
481 Model::Gguf(m) => m.tokenizer.kind(),
482 Model::Kimi(_) => "kimi-tiktoken-bpe",
483 Model::Mla(m) => m.tokenizer.kind(),
484 Model::Gemma4(m) => m.tokenizer.kind(),
485 Model::Glm52(m) => m.tokenizer.kind(),
486 }
487 }
488
489 fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
493 match self {
494 Model::Gguf(m) => m.decoder.expert_store_stats(),
495 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
496 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
497 }
498 }
499
500 pub(crate) fn name(&self) -> &str {
501 match self {
502 Model::Gguf(m) => m.decoder.config.name,
503 Model::Kimi(_) => "kimi-k3",
504 Model::Mla(m) => m.name.as_str(),
505 Model::Gemma4(m) => m.name.as_str(),
506 Model::Glm52(m) => m.name.as_str(),
507 }
508 }
509
510 pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
511 match self {
512 Model::Gguf(m) => m.tokenizer.encode(text),
513 Model::Kimi(m) => m
514 .tokenizer
515 .encode(text)
516 .into_iter()
517 .map(|id| id as usize)
518 .collect(),
519 Model::Mla(m) => m.tokenizer.encode(text),
520 Model::Gemma4(m) => m.tokenizer.encode(text),
521 Model::Glm52(m) => m.tokenizer.encode(text),
522 }
523 }
524
525 pub(crate) fn bos_id(&self) -> Option<usize> {
534 match self {
535 Model::Gguf(m) => m.bos_id,
536 Model::Kimi(_) => None,
537 Model::Mla(m) => m.bos_id,
538 Model::Gemma4(m) => m.bos_id,
539 Model::Glm52(m) => m.bos_id,
540 }
541 }
542
543 pub(crate) fn decode(&self, ids: &[usize]) -> String {
544 match self {
545 Model::Gguf(m) => m.tokenizer.decode(ids),
546 Model::Kimi(m) => {
547 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
548 m.tokenizer.decode(&ids32)
549 }
550 Model::Mla(m) => m.tokenizer.decode(ids),
551 Model::Gemma4(m) => m.tokenizer.decode(ids),
552 Model::Glm52(m) => m.tokenizer.decode(ids),
553 }
554 }
555
556 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
559 match self {
560 Model::Gguf(m) => {
561 let mut caches: Vec<_> = (0..m.decoder.layers.len())
562 .map(|_| {
563 ferrox_core::cache::KvCache::new(
564 m.decoder.config.n_kv_heads,
565 m.decoder.config.head_dim,
566 )
567 })
568 .collect();
569 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
570 }
571 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
572 }
573 }
574
575 pub(crate) fn vocab_size(&self) -> Option<usize> {
576 match self {
577 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
578 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
579 Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
580 Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
581 Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
582 }
583 }
584}
585
586pub(crate) struct ActiveModel {
595 pub(crate) id: Option<String>,
599 pub(crate) model: Arc<Model>,
600 pub(crate) batcher: Option<serving::batch::ContinuousBatcher>,
605 pub(crate) ceiling: Option<Arc<budget::ContextCeiling>>,
615}
616
617pub(crate) struct AppState {
618 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
637 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
641 pub(crate) tasks: Arc<tasks::TaskRegistry>,
643 pub(crate) cancels: Arc<cancel::CancelRegistry>,
647 pub(crate) stats: stats::Stats,
650 pub(crate) streams: resume::StreamRegistry,
653 pub(crate) model_dir: Option<PathBuf>,
655 response_cache: Mutex<ResponseCache>,
659 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
671 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
684 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
694 sessions: session::SessionStore,
700 requests_total: std::sync::atomic::AtomicU64,
701 request_errors_total: std::sync::atomic::AtomicU64,
702 started_at: std::time::Instant,
703 last_request_ms: std::sync::atomic::AtomicU64,
708 detection: Arc<health::Detection>,
710 mcp: Option<mcp::LoadedMcpConfig>,
712 pub(crate) continuous_batching_enabled: bool,
716 loading_model: Mutex<Option<String>>,
720 last_load_error: Mutex<Option<(String, String)>>,
724 pub(crate) serving: Mutex<crate::stats::ServingStats>,
729 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
734 pub(crate) footprint:
740 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
741 pub(crate) started_unix: u64,
751}
752
753pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
759
760impl AppState {
761 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
765 self.active
766 .read()
767 .unwrap_or_else(|p| p.into_inner())
768 .clone()
769 }
770
771 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
777 self.active().ok_or_else(|| {
778 (
779 StatusCode::SERVICE_UNAVAILABLE,
780 Json(serde_json::json!({"error": {
781 "message": "no model is loaded; POST /admin/models/load with an id from \
782 GET /admin/models",
783 "type": "model_not_loaded"
784 }})),
785 )
786 })
787 }
788
789 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
792 Ok(Arc::clone(&self.require_active()?.model))
793 }
794
795 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
803 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
804 std::mem::replace(&mut *guard, next)
805 }
806
807 fn mark_request_finished(&self) {
811 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
812 self.last_request_ms
813 .store(ms, std::sync::atomic::Ordering::Relaxed);
814 }
815
816 pub(crate) fn uptime(&self) -> Duration {
817 self.started_at.elapsed()
818 }
819
820 pub(crate) fn requests_total(&self) -> u64 {
821 self.requests_total
822 .load(std::sync::atomic::Ordering::Relaxed)
823 }
824
825 pub(crate) fn errors_total(&self) -> u64 {
826 self.request_errors_total
827 .load(std::sync::atomic::Ordering::Relaxed)
828 }
829
830 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
831 lock_cache(&self.response_cache).stats()
832 }
833
834 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
837 let last = self
838 .last_request_ms
839 .load(std::sync::atomic::Ordering::Relaxed);
840 (last > 0)
841 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
842 .map(|age| age.max(0.0))
843 }
844
845 pub(crate) fn loading_model_id(&self) -> Option<String> {
846 self.loading_model
847 .lock()
848 .unwrap_or_else(|p| p.into_inner())
849 .clone()
850 }
851
852 pub(crate) fn set_loading_model(&self, id: Option<String>) {
853 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
854 }
855
856 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
857 self.last_load_error
858 .lock()
859 .unwrap_or_else(|p| p.into_inner())
860 .clone()
861 }
862
863 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
864 *self
865 .last_load_error
866 .lock()
867 .unwrap_or_else(|p| p.into_inner()) = error;
868 }
869
870 pub(crate) fn active_model_name(&self) -> Option<String> {
879 self.active().map(|a| a.model.name().to_string())
880 }
881
882 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
883 self.stats.record(stats::entry(record));
884 }
885}
886
887fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
892 cache
893 .lock()
894 .unwrap_or_else(|poisoned| poisoned.into_inner())
895}
896
897#[derive(Debug, Clone, Deserialize)]
898#[serde(untagged)]
899pub(crate) enum MessageContent {
900 Text(String),
901 Parts(Vec<ContentPart>),
902}
903
904#[derive(Debug, Clone, Deserialize)]
905struct ContentPart {
906 #[serde(rename = "type")]
907 kind: String,
908 #[serde(default)]
909 text: Option<String>,
910 #[serde(default)]
911 image_url: Option<serde_json::Value>,
912}
913
914impl MessageContent {
915 fn as_text(&self) -> String {
916 match self {
917 Self::Text(s) => s.clone(),
918 Self::Parts(parts) => parts
919 .iter()
920 .filter_map(|p| p.text.as_deref())
921 .collect::<Vec<_>>()
922 .join(""),
923 }
924 }
925
926 fn has_image(&self) -> bool {
927 match self {
928 Self::Text(_) => false,
929 Self::Parts(parts) => parts
930 .iter()
931 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
932 }
933 }
934}
935
936#[derive(Debug, Clone, Deserialize)]
937pub(crate) struct ChatMessage {
938 pub(crate) role: String,
939 #[serde(default)]
943 pub(crate) content: Option<MessageContent>,
944 #[serde(default)]
948 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
949 #[serde(default)]
954 #[allow(dead_code)]
955 pub(crate) tool_call_id: Option<String>,
956 #[serde(default, alias = "reasoning")]
972 pub(crate) reasoning_content: Option<String>,
973}
974
975impl ChatMessage {
976 fn rendered_content(&self) -> String {
983 let mut out = self
984 .content
985 .as_ref()
986 .map(MessageContent::as_text)
987 .unwrap_or_default();
988 if let Some(calls) = &self.tool_calls {
989 for call in calls {
990 out.push_str(&format!(
991 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
992 call.function.name, call.function.arguments
993 ));
994 }
995 }
996 out
997 }
998}
999
1000#[derive(Debug, Clone, Deserialize)]
1001pub(crate) struct ToolCallIn {
1002 #[serde(default)]
1003 #[allow(dead_code)]
1004 id: String,
1005 #[serde(rename = "type", default)]
1006 #[allow(dead_code)]
1007 kind: String,
1008 function: ToolCallFunctionIn,
1009}
1010
1011#[derive(Debug, Clone, Deserialize)]
1012struct ToolCallFunctionIn {
1013 name: String,
1014 arguments: String,
1019}
1020
1021#[derive(Debug, Clone, Deserialize)]
1024struct ToolDef {
1025 #[serde(rename = "type", default)]
1026 #[allow(dead_code)]
1027 kind: String,
1028 function: ToolFunctionDef,
1029}
1030
1031#[derive(Debug, Clone, Deserialize)]
1032struct ToolFunctionDef {
1033 name: String,
1034 #[serde(default)]
1035 description: Option<String>,
1036 #[serde(default)]
1037 parameters: Option<serde_json::Value>,
1038}
1039
1040#[derive(Debug, Clone, Deserialize)]
1051#[serde(untagged)]
1052enum ToolChoice {
1053 Mode(String),
1054 Specific(serde_json::Value),
1055}
1056
1057#[derive(Deserialize)]
1060#[serde(untagged)]
1061enum StopParam {
1062 One(String),
1063 Many(Vec<String>),
1064}
1065
1066#[derive(Deserialize)]
1067struct ChatCompletionRequest {
1068 model: String,
1069 messages: Vec<ChatMessage>,
1070 #[serde(default = "default_max_tokens")]
1071 max_tokens: usize,
1072 #[serde(default)]
1073 temperature: Option<f32>,
1074 #[serde(default)]
1075 top_p: Option<f32>,
1076 #[serde(default)]
1081 min_p: Option<f32>,
1082 #[serde(default)]
1083 top_k: Option<usize>,
1084 #[serde(default)]
1085 repetition_penalty: Option<f32>,
1086 #[serde(default)]
1087 seed: Option<u64>,
1088 #[serde(default)]
1089 stop: Option<StopParam>,
1090 #[serde(default)]
1091 stream: Option<bool>,
1092 #[serde(default)]
1105 stream_resumable: Option<bool>,
1106 #[serde(default)]
1116 ignore_eos: Option<bool>,
1117 #[serde(default)]
1118 tools: Vec<ToolDef>,
1119 #[serde(default)]
1120 tool_choice: Option<ToolChoice>,
1121 #[serde(default)]
1129 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
1130 #[serde(default)]
1138 reasoning_effort: Option<String>,
1139 #[serde(default)]
1143 thinking: Option<ThinkingSwitch>,
1144 #[serde(default)]
1149 session_id: Option<String>,
1150 #[serde(default)]
1152 logprobs: Option<bool>,
1153 #[serde(default)]
1154 top_logprobs: Option<u32>,
1155 #[serde(default)]
1156 n: Option<u32>,
1157 #[serde(default)]
1158 presence_penalty: Option<f32>,
1159 #[serde(default)]
1160 frequency_penalty: Option<f32>,
1161 #[serde(default)]
1162 response_format: Option<serde_json::Value>,
1163 #[serde(default)]
1170 logit_bias: Option<serde_json::Value>,
1171 #[serde(default)]
1180 grammar: Option<String>,
1181}
1182
1183const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1195
1196#[derive(Debug, Clone, Deserialize)]
1198pub(crate) struct ThinkingSwitch {
1199 #[serde(rename = "type")]
1200 pub(crate) kind: String,
1201}
1202
1203const THINKING_KWARG_KEYS: [&str; 4] = [
1207 "enable_thinking",
1208 "thinking",
1209 "thinking_mode",
1210 "reasoning_effort",
1211];
1212
1213const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1217
1218fn default_max_tokens() -> usize {
1219 DEFAULT_CHAT_MAX_TOKENS
1220}
1221
1222impl ChatCompletionRequest {
1223 fn sampling_knobs(&self) -> SamplingKnobs {
1228 SamplingKnobs {
1229 temperature: self.temperature,
1230 top_p: self.top_p,
1231 min_p: self.min_p,
1232 top_k: self.top_k,
1233 repetition_penalty: self.repetition_penalty,
1234 presence_penalty: self.presence_penalty,
1235 frequency_penalty: self.frequency_penalty,
1236 penalty_last_n: None,
1240 }
1241 }
1242
1243 fn sampling_params(&self) -> SamplingParams {
1244 self.sampling_knobs().resolve()
1245 }
1246
1247 fn stop_sequences(&self) -> Vec<String> {
1248 self.stop
1249 .as_ref()
1250 .map(|s| match s {
1251 StopParam::One(v) => vec![v.clone()],
1252 StopParam::Many(v) => v.clone(),
1253 })
1254 .unwrap_or_default()
1255 }
1256
1257 fn tools_active(&self) -> bool {
1262 !self.tools.is_empty()
1263 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1264 }
1265
1266 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1278 match &self.tool_choice {
1279 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1280 Some(ToolChoice::Specific(value)) => {
1281 let name = value
1285 .get("function")
1286 .and_then(|f| f.get("name"))
1287 .or_else(|| value.get("name"))
1288 .and_then(|n| n.as_str());
1289 match name {
1290 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1291 None => Err(invalid_request(
1292 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1293 function.name",
1294 "tool_choice",
1295 )),
1296 }
1297 }
1298 _ => Ok(None),
1299 }
1300 }
1301
1302 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1304 self.tools
1305 .iter()
1306 .map(|t| tool_grammar::ToolSpec {
1307 name: &t.function.name,
1308 parameters: t.function.parameters.as_ref(),
1309 })
1310 .collect()
1311 }
1312
1313 fn thinking_direction(&self) -> Option<bool> {
1356 if let Some(switch) = &self.thinking {
1357 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1358 "disabled" => Some(false),
1359 "enabled" => Some(true),
1360 _ => None,
1363 };
1364 }
1365 let effort = self.reasoning_effort.as_ref()?;
1366 DISABLE_EFFORTS
1367 .contains(&effort.trim().to_ascii_lowercase().as_str())
1368 .then_some(false)
1369 }
1370
1371 fn resolve_template_kwargs(
1372 &self,
1373 template: &chat_template::PromptTemplate,
1374 ) -> serde_json::Map<String, serde_json::Value> {
1375 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1376 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1380
1381 if !caller_steered {
1382 match self.thinking_direction() {
1383 Some(false) => {
1384 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1385 kwargs.insert(k, v);
1386 }
1387 return kwargs;
1390 }
1391 Some(true) => {
1392 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1393 kwargs.insert(k, v);
1394 }
1395 }
1396 None => {}
1397 }
1398 if let Some(effort) = &self.reasoning_effort {
1399 kwargs
1400 .entry("reasoning_effort".to_string())
1401 .or_insert_with(|| serde_json::json!(effort));
1402 }
1403 }
1404
1405 let offered: Vec<serde_json::Value> = if self.tools_active() {
1406 self.tools.iter().map(chat_template::tool_json).collect()
1407 } else {
1408 Vec::new()
1409 };
1410 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1411 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1412 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1413 kwargs.entry(k).or_insert(v);
1414 }
1415 }
1416 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1417 crate::policy::effort::EffortMapping::Mapped(to) => {
1418 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1419 }
1420 crate::policy::effort::EffortMapping::Dropped => {
1421 tracing::debug!(
1422 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1423 enough, so its own default applies"
1424 );
1425 }
1426 crate::policy::effort::EffortMapping::Unchanged => {}
1427 }
1428 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1429 kwargs
1430 }
1431
1432 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1435 if self.max_tokens == 0 {
1442 return Err(invalid_request(
1443 "max_tokens must be at least 1",
1444 "max_tokens",
1445 ));
1446 }
1447 if let Some(switch) = &self.thinking {
1452 let kind = switch.kind.trim().to_ascii_lowercase();
1453 if kind != "enabled" && kind != "disabled" {
1454 return Err(invalid_request(
1455 "thinking.type must be \"enabled\" or \"disabled\"",
1456 "thinking.type",
1457 ));
1458 }
1459 }
1460 for msg in &self.messages {
1461 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1462 return Err(unsupported_feature(
1463 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1464 ));
1465 }
1466 }
1467 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1468 return Err(unsupported_feature(
1469 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1470 ));
1471 }
1472 if self.n.is_some_and(|n| n > 1) {
1473 return Err(unsupported_feature(
1474 "n > 1 is not implemented (single completion only)",
1475 ));
1476 }
1477 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1478 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1486 if let Some(fmt) = &self.response_format {
1487 match fmt.get("type").and_then(|v| v.as_str()) {
1488 Some("json_object") => {}
1489 Some(other) => {
1493 return Err((
1494 StatusCode::BAD_REQUEST,
1495 Json(serde_json::json!({
1496 "error": {
1497 "message": format!(
1498 "response_format type {other:?} is not supported (only json_object)"
1499 )
1500 }
1501 })),
1502 ));
1503 }
1504 None => {
1505 return Err((
1506 StatusCode::BAD_REQUEST,
1507 Json(serde_json::json!({
1508 "error": {
1509 "message": "response_format must include \"type\" (only json_object is supported)"
1510 }
1511 })),
1512 ));
1513 }
1514 }
1515 }
1516 if let Some(forced) = self.forced_tool_choice()? {
1524 if self.tools.is_empty() {
1525 return Err(invalid_request(
1526 "tool_choice forces a tool call, but no tools were offered",
1527 "tool_choice",
1528 ));
1529 }
1530 if let tool_grammar::Forced::Named(name) = forced {
1531 if !self.tools.iter().any(|t| t.function.name == name) {
1532 return Err(invalid_request(
1533 &format!(
1534 "tool_choice names {name:?}, which is not one of the tools offered"
1535 ),
1536 "tool_choice",
1537 ));
1538 }
1539 }
1540 if self.grammar.is_some() {
1543 return Err(invalid_request(
1544 "a forced tool_choice and a \"grammar\" are two different constraints on the \
1545 same generation; send one",
1546 "tool_choice",
1547 ));
1548 }
1549 if self.json_object_mode() {
1550 return Err(invalid_request(
1551 "a forced tool_choice cannot be combined with response_format json_object: \
1552 the tool-call markers are not JSON",
1553 "tool_choice",
1554 ));
1555 }
1556 }
1557 Ok(())
1558 }
1559
1560 fn effective_stop_sequences(&self) -> Vec<String> {
1567 let mut stop = self.stop_sequences();
1568 if self.tools_active() {
1569 stop.push("</tool_call>".to_string());
1570 }
1571 stop
1572 }
1573
1574 fn json_object_mode(&self) -> bool {
1575 self.response_format
1576 .as_ref()
1577 .and_then(|v| v.get("type"))
1578 .and_then(|v| v.as_str())
1579 == Some("json_object")
1580 }
1581
1582 fn generation_params(&self) -> Result<GenerationParams, ApiError> {
1587 Ok(GenerationParams {
1588 max_tokens: self.max_tokens,
1589 sampling: self.sampling_params(),
1590 seed: self.resolved_seed(),
1591 stop: self.effective_stop_sequences(),
1592 stop_token_ids: Vec::new(),
1596 json_object: self.json_object_mode(),
1597 grammar: grammar_request::for_request(
1598 self.grammar.as_deref(),
1599 self.response_format.as_ref(),
1600 )?,
1601 cancel: None,
1604 ignore_eos: self.ignore_eos.unwrap_or(false),
1605 })
1606 }
1607
1608 fn generation_params_for_template(
1619 &self,
1620 template: &chat_template::PromptTemplate,
1621 served_model: &str,
1622 ) -> Result<GenerationParams, ApiError> {
1623 let mut params = self.generation_params()?;
1624 if let Some(stop) = template.end_of_turn() {
1625 if !params.stop.iter().any(|s| s == stop) {
1626 params.stop.push(stop.to_string());
1627 }
1628 }
1629 if let Some(forced) = self.forced_tool_choice()? {
1630 params.grammar = Some(tool_grammar::build(
1634 forced,
1635 &self.tool_specs(),
1636 policy::parser::ToolCallFormat::infer(served_model),
1637 )?);
1638 }
1639 Ok(params)
1640 }
1641
1642 fn is_cacheable(&self) -> bool {
1651 self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
1652 }
1653
1654 fn cache_key(&self, prompt: &str) -> CacheKey {
1655 CacheKey {
1656 model: self.model.clone(),
1657 prompt: prompt.to_string(),
1658 max_tokens: self.max_tokens,
1659 sampling: response_cache::sampling_key(&self.sampling_params()),
1664 seed: self.seed,
1665 stop: self.effective_stop_sequences(),
1666 }
1667 }
1668
1669 fn resolved_seed(&self) -> u64 {
1670 self.seed.unwrap_or_else(|| {
1671 std::time::SystemTime::now()
1672 .duration_since(std::time::UNIX_EPOCH)
1673 .map(|d| d.as_nanos() as u64)
1674 .unwrap_or(0xDEFA017)
1675 })
1676 }
1677}
1678
1679#[derive(Serialize)]
1680struct ChatCompletionChoice {
1681 index: usize,
1682 message: ChatCompletionResponseMessage,
1683 finish_reason: &'static str,
1684}
1685
1686#[derive(Serialize)]
1687struct ChatCompletionResponseMessage {
1688 role: &'static str,
1689 #[serde(skip_serializing_if = "Option::is_none")]
1690 content: Option<String>,
1691 #[serde(skip_serializing_if = "Option::is_none")]
1695 reasoning_content: Option<String>,
1696 #[serde(skip_serializing_if = "Option::is_none")]
1697 tool_calls: Option<Vec<ToolCallOut>>,
1698}
1699
1700#[derive(Serialize, Clone)]
1701struct ToolCallOut {
1702 id: String,
1703 #[serde(rename = "type")]
1704 kind: &'static str,
1705 function: ToolCallFunctionOut,
1706}
1707
1708#[derive(Serialize, Clone)]
1716struct ToolCallDelta {
1717 index: usize,
1718 #[serde(skip_serializing_if = "Option::is_none")]
1719 id: Option<String>,
1720 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1721 kind: Option<&'static str>,
1722 function: ToolCallFunctionDelta,
1723}
1724
1725#[derive(Serialize, Clone, Default)]
1726struct ToolCallFunctionDelta {
1727 #[serde(skip_serializing_if = "Option::is_none")]
1728 name: Option<String>,
1729 #[serde(skip_serializing_if = "Option::is_none")]
1732 arguments: Option<String>,
1733}
1734
1735impl ToolCallDelta {
1736 fn whole(index: usize, name: String, arguments: String) -> Self {
1738 ToolCallDelta {
1739 index,
1740 id: Some(format!("call_{index}")),
1741 kind: Some("function"),
1742 function: ToolCallFunctionDelta {
1743 name: Some(name),
1744 arguments: Some(arguments),
1745 },
1746 }
1747 }
1748
1749 fn opening(index: usize, name: String) -> Self {
1751 ToolCallDelta {
1752 index,
1753 id: Some(format!("call_{index}")),
1754 kind: Some("function"),
1755 function: ToolCallFunctionDelta {
1756 name: Some(name),
1757 arguments: Some(String::new()),
1758 },
1759 }
1760 }
1761
1762 fn arguments(index: usize, fragment: String) -> Self {
1764 ToolCallDelta {
1765 index,
1766 id: None,
1767 kind: None,
1768 function: ToolCallFunctionDelta {
1769 name: None,
1770 arguments: Some(fragment),
1771 },
1772 }
1773 }
1774}
1775
1776#[derive(Serialize, Clone)]
1777struct ToolCallFunctionOut {
1778 name: String,
1779 arguments: String,
1783}
1784
1785#[derive(Serialize)]
1786struct ChatCompletionResponse {
1787 id: String,
1788 request_id: String,
1794 object: &'static str,
1795 model: String,
1796 choices: Vec<ChatCompletionChoice>,
1797 usage: generate::Usage,
1802 ferrox_cache: &'static str,
1810}
1811
1812#[derive(Serialize)]
1813struct ChatCompletionChunkDelta {
1814 #[serde(skip_serializing_if = "Option::is_none")]
1815 role: Option<&'static str>,
1816 #[serde(skip_serializing_if = "Option::is_none")]
1817 content: Option<String>,
1818 #[serde(skip_serializing_if = "Option::is_none")]
1820 reasoning_content: Option<String>,
1821 #[serde(skip_serializing_if = "Option::is_none")]
1822 tool_calls: Option<Vec<ToolCallDelta>>,
1823}
1824
1825#[derive(Serialize)]
1826struct ChatCompletionChunkChoice {
1827 index: usize,
1828 delta: ChatCompletionChunkDelta,
1829 finish_reason: Option<&'static str>,
1830}
1831
1832#[derive(Serialize)]
1833struct ChatCompletionChunk {
1834 id: String,
1835 #[serde(skip_serializing_if = "Option::is_none")]
1842 request_id: Option<String>,
1843 object: &'static str,
1844 model: String,
1845 choices: Vec<ChatCompletionChunkChoice>,
1846 #[serde(skip_serializing_if = "Option::is_none")]
1849 usage: Option<generate::Usage>,
1850}
1851
1852async fn health(State(state): State<Arc<AppState>>) -> Response {
1858 let snapshot = state.detection.snapshot();
1859 let mut capabilities = snapshot.capabilities;
1860 let active = state.active();
1861
1862 capabilities.push(match active.as_deref() {
1865 None => ferrox_api::Capability::unavailable(
1870 ferrox_api::health::capability::REAL_WEIGHTS,
1871 ferrox_api::health::reason::MODEL_NOT_LOADED,
1872 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1873 ),
1874 Some(active) if active.model.is_synthetic() => ferrox_api::Capability::unavailable(
1875 ferrox_api::health::capability::REAL_WEIGHTS,
1876 ferrox_api::health::reason::MODEL_NOT_LOADED,
1877 "Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
1878 checkpoint. Output from this model is noise.",
1879 ),
1880 Some(active) => ferrox_api::Capability::available(
1881 ferrox_api::health::capability::REAL_WEIGHTS,
1882 format!("Serving the real checkpoint '{}'.", active.model.name()),
1883 ),
1884 });
1885 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1886 ferrox_api::Capability::available(
1887 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1888 "Concurrent requests share one batched decode step.",
1889 )
1890 } else {
1891 ferrox_api::Capability::unavailable(
1892 ferrox_api::health::capability::CONTINUOUS_BATCHING,
1893 ferrox_api::health::reason::DISABLED,
1894 "Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1895 )
1896 });
1897
1898 let last_request_ms = state
1899 .last_request_ms
1900 .load(std::sync::atomic::Ordering::Relaxed);
1901 let uptime = state.started_at.elapsed();
1902 let health_state = if active.is_none() {
1910 ferrox_api::HealthState::Unavailable
1911 } else {
1912 snapshot.state
1913 };
1914 let body = ferrox_api::HealthResponse {
1915 state: health_state,
1916 reason: match health_state {
1917 ferrox_api::HealthState::Ready => None,
1918 ferrox_api::HealthState::Unavailable => {
1919 Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
1920 }
1921 ferrox_api::HealthState::Detecting => {
1922 Some(ferrox_api::health::reason::DETECTING.to_string())
1923 }
1924 },
1925 detail: match health_state {
1926 ferrox_api::HealthState::Ready => None,
1927 ferrox_api::HealthState::Unavailable => Some(
1928 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1929 .to_string(),
1930 ),
1931 ferrox_api::HealthState::Detecting => {
1932 Some("Probing available compute backends.".to_string())
1933 }
1934 },
1935 model: active
1936 .as_deref()
1937 .map(|active| ferrox_api::health::ModelSummary {
1938 id: active.model.name().to_string(),
1939 tokenizer: active.model.tokenizer_kind().to_string(),
1940 synthetic_weights: active.model.is_synthetic(),
1941 }),
1942 capabilities,
1943 version: env!("CARGO_PKG_VERSION").to_string(),
1944 pid: std::process::id(),
1945 uptime_seconds: uptime.as_secs_f64(),
1946 server_time_unix_ms: std::time::SystemTime::now()
1947 .duration_since(std::time::UNIX_EPOCH)
1948 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1949 .unwrap_or(0),
1950 last_request_age_seconds: (last_request_ms > 0)
1951 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1952 .map(|age| age.max(0.0)),
1953 };
1954
1955 let status =
1956 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1957 (status, Json(body)).into_response()
1958}
1959
1960async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1961 let Some(active) = state.active() else {
1965 return Json(serde_json::json!({ "object": "list", "data": [] }));
1966 };
1967 let mut model_entry = serde_json::json!({
1968 "id": active.model.name(),
1969 "object": "model",
1970 "ferrox_synthetic_weights": active.model.is_synthetic(),
1971 "ferrox_tokenizer": active.model.tokenizer_kind(),
1972 });
1973 let parser_configured =
1979 crate::policy::parser::ReasoningFormat::infer(active.model.name()).is_some();
1980 let gears = active.model.chat_template().think_gears(parser_configured);
1981 if !gears.is_empty() {
1982 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1983 if let Some(default) = &gears.default {
1984 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1985 }
1986 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1989 }
1990 if let Some(mcp) = &state.mcp {
1991 model_entry["ferrox_mcp"] = mcp.models_metadata();
1992 }
1993 Json(serde_json::json!({
1994 "object": "list",
1995 "data": [model_entry]
1996 }))
1997}
1998
1999async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
2013 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
2014 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
2015 let active = state.active();
2016 Json(serde_json::json!({
2017 "model": active.as_ref().map(|a| a.model.name()),
2018 "state": state
2019 .maintenance
2020 .lock()
2021 .unwrap_or_else(|p| p.into_inner())
2022 .state()
2023 .as_str(),
2024 "uptime_s": state.uptime().as_secs(),
2025 "throughput": {
2026 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2027 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
2028 },
2029 "requests": {
2030 "active": state.cancels.live_count(),
2031 "completed": state.stats.recorded_total(),
2032 "p95_ms": state.stats.p95_duration_ms(),
2033 "ttft_mean_ms": state.stats.ttft_mean_ms(),
2034 "prompt_tokens_total": state.stats.tokens_prompt_total(),
2035 "completion_tokens_total": state.stats.tokens_generated_total(),
2036 },
2037 "pools": cache_admin::pool_gauges(&state),
2043 "memory": cache_admin::footprint_json(&state),
2046 }))
2047}
2048
2049#[derive(Deserialize)]
2050struct RequestsQuery {
2051 #[serde(default)]
2052 since: u64,
2053 #[serde(default = "default_requests_limit")]
2054 limit: usize,
2055}
2056
2057fn default_requests_limit() -> usize {
2058 stats::MAX_PAGE
2059}
2060
2061async fn recent_requests(
2069 State(state): State<Arc<AppState>>,
2070 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
2071) -> Json<serde_json::Value> {
2072 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
2073 Json(serde_json::json!({
2074 "requests": rows,
2075 "next_cursor": cursor,
2076 "missed": missed,
2077 "total": state.stats.recorded_total(),
2078 }))
2079}
2080
2081#[derive(Serialize)]
2082struct CombinedCacheStats {
2083 response_cache: response_cache::CacheStats,
2084 prefix_cache: Option<ferrox_models::PrefixCacheStats>,
2086}
2087
2088async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
2089 Json(CombinedCacheStats {
2090 response_cache: lock_cache(&state.response_cache).stats(),
2091 prefix_cache: state
2092 .prefix_cache
2093 .as_ref()
2094 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
2095 })
2096}
2097
2098async fn metrics(State(state): State<Arc<AppState>>) -> Response {
2103 use std::sync::atomic::Ordering;
2104
2105 let cache_stats = lock_cache(&state.response_cache).stats();
2106 let active = state.active();
2107 let requests_total = state.requests_total.load(Ordering::Relaxed);
2108 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2109 let uptime = state.started_at.elapsed().as_secs_f64();
2110
2111 let body = format!(
2112 "# HELP ferrox_requests_total Total chat completion requests received.\n\
2113 # TYPE ferrox_requests_total counter\n\
2114 ferrox_requests_total {requests_total}\n\
2115 # HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
2116 # TYPE ferrox_request_errors_total counter\n\
2117 ferrox_request_errors_total {errors_total}\n\
2118 # HELP ferrox_cache_hits_total Whole-response cache hits.\n\
2119 # TYPE ferrox_cache_hits_total counter\n\
2120 ferrox_cache_hits_total {}\n\
2121 # HELP ferrox_cache_misses_total Whole-response cache misses.\n\
2122 # TYPE ferrox_cache_misses_total counter\n\
2123 ferrox_cache_misses_total {}\n\
2124 # HELP ferrox_cache_entries Current whole-response cache entry count.\n\
2125 # TYPE ferrox_cache_entries gauge\n\
2126 ferrox_cache_entries {}\n\
2127 # HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2128 # TYPE ferrox_synthetic_weights gauge\n\
2129 ferrox_synthetic_weights {}\n\
2130 # HELP ferrox_uptime_seconds Seconds since this server process started.\n\
2131 # TYPE ferrox_uptime_seconds gauge\n\
2132 ferrox_uptime_seconds {uptime}\n",
2133 cache_stats.hits,
2134 cache_stats.misses,
2135 cache_stats.entries,
2136 active
2140 .as_ref()
2141 .map(|a| a.model.is_synthetic() as u8)
2142 .unwrap_or(0),
2143 );
2144
2145 let body = match active
2149 .as_ref()
2150 .and_then(|a| a.model.expert_store_stats())
2151 {
2152 Some(es) => format!(
2153 "{body}\
2154 # HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
2155 # TYPE ferrox_expert_cache_hits_total counter\n\
2156 ferrox_expert_cache_hits_total {}\n\
2157 # HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2158 # TYPE ferrox_expert_cache_misses_total counter\n\
2159 ferrox_expert_cache_misses_total {}\n\
2160 # HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
2161 # TYPE ferrox_expert_cache_evictions_total counter\n\
2162 ferrox_expert_cache_evictions_total {}\n\
2163 # HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2164 # TYPE ferrox_expert_cache_pass_throughs_total counter\n\
2165 ferrox_expert_cache_pass_throughs_total {}\n\
2166 # HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2167 # TYPE ferrox_expert_cache_bytes_read_total counter\n\
2168 ferrox_expert_cache_bytes_read_total {}\n\
2169 # HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2170 # TYPE ferrox_expert_cache_resident_bytes gauge\n\
2171 ferrox_expert_cache_resident_bytes {}\n",
2172 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2173 ),
2174 None => body,
2175 };
2176
2177 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2182 Some(batcher) => {
2183 let sched = batcher.stats();
2184 format!(
2185 "{body}\
2186 # HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2187 # TYPE ferrox_prefill_chunks_total counter\n\
2188 ferrox_prefill_chunks_total {}\n\
2189 # HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2190 # TYPE ferrox_prefill_tokens_total counter\n\
2191 ferrox_prefill_tokens_total {}\n\
2192 # HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2193 # TYPE ferrox_decode_steps_total counter\n\
2194 ferrox_decode_steps_total {}\n\
2195 # HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2196 # TYPE ferrox_scheduler_queue_depth gauge\n\
2197 ferrox_scheduler_queue_depth {}\n\
2198 # HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2199 # TYPE ferrox_scheduler_queue_rejected_total counter\n\
2200 ferrox_scheduler_queue_rejected_total {}\n\
2201 # HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2202 # TYPE ferrox_kv_blocks_total gauge\n\
2203 ferrox_kv_blocks_total {}\n\
2204 # HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2205 # TYPE ferrox_kv_blocks_free gauge\n\
2206 ferrox_kv_blocks_free {}\n\
2207 # HELP ferrox_kv_block_size Token positions per KV block.\n\
2208 # TYPE ferrox_kv_block_size gauge\n\
2209 ferrox_kv_block_size {}\n\
2210 # HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2211 # TYPE ferrox_kv_rejected_too_large_total counter\n\
2212 ferrox_kv_rejected_too_large_total {}\n\
2213 # HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2214 # TYPE ferrox_kv_rejected_context_length_total counter\n\
2215 ferrox_kv_rejected_context_length_total {}\n\
2216 # HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2217 # TYPE ferrox_scheduler_aborted_total counter\n\
2218 ferrox_scheduler_aborted_total {}\n",
2219 sched.prefill_chunks,
2220 sched.prefill_tokens,
2221 sched.decode_steps,
2222 sched.queue_depth,
2223 sched.queue_rejected,
2224 sched.kv_blocks_total,
2225 sched.kv_blocks_free,
2226 sched.kv_block_size,
2227 sched.kv_rejected_too_large,
2228 sched.kv_rejected_context_length,
2229 sched.aborted,
2230 )
2231 }
2232 None => body,
2233 };
2234
2235 (
2236 [(
2237 axum::http::header::CONTENT_TYPE,
2238 "text/plain; version=0.0.4",
2239 )],
2240 body,
2241 )
2242 .into_response()
2243}
2244
2245pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2246
2247pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2253 (
2254 StatusCode::BAD_REQUEST,
2255 Json(serde_json::json!({"error": {
2256 "message": message,
2257 "type": "invalid_request_error",
2258 "param": param,
2259 "code": null,
2260 }})),
2261 )
2262}
2263
2264pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2265 (
2266 StatusCode::NOT_IMPLEMENTED,
2267 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2268 )
2269}
2270
2271pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2272 let status = match e {
2273 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2274 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2279 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2285 StatusCode::SERVICE_UNAVAILABLE
2286 }
2287 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2291 };
2292 tracing::warn!("decode error: {e}");
2293 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2294 if let generate::DecodeError::KvBudgetExceeded {
2300 binding,
2301 estimated_bytes,
2302 limit_bytes,
2303 positions,
2304 positions_limit,
2305 ..
2306 } = &e
2307 {
2308 body["error"]["type"] = serde_json::json!("invalid_request_error");
2309 body["error"]["code"] = serde_json::json!(binding);
2310 body["error"]["binding"] = serde_json::json!(binding);
2311 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2312 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2313 body["error"]["positions"] = serde_json::json!(positions);
2314 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2315 }
2316 if let Some(secs) = e.retry_after_secs() {
2320 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2321 }
2322 (status, Json(body))
2323}
2324
2325pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2326 tracing::error!("generation task panicked: {e}");
2327 (
2328 StatusCode::INTERNAL_SERVER_ERROR,
2329 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2330 )
2331}
2332
2333#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2343 model: &Model,
2344 prompt: &str,
2345 params: &GenerationParams,
2346 kv_pool: Option<&generate::KvPoolConfig>,
2347 paged_kv: Option<&generate::PagedKvConfig>,
2348 prefix_cache: Option<&Mutex<PrefixCache>>,
2349 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2350 ceiling: Option<&budget::ContextCeiling>,
2351 mut emit: impl FnMut(&str),
2352) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2353 let synthetic = model.is_synthetic();
2354 let mut chunks = Vec::new();
2355 let params = &{
2361 let mut resolved = params.clone();
2362 resolved.stop_token_ids =
2363 crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
2364 resolved
2365 };
2366 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2367 let (finish, usage) = match model {
2368 Model::Gguf(m) => {
2369 if let Some(batcher) = continuous_batcher {
2370 let mut tokens = m.tokenizer.encode(prompt);
2371 ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2372 let (finish, _generated_ids, text, usage) =
2373 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?;
2374 if !text.is_empty() {
2375 chunks.push(text);
2376 }
2377 (finish, usage)
2378 } else {
2379 generate::generate(
2380 &m.decoder,
2381 m.tokenizer.as_ref(),
2382 &m.stop_tokens,
2383 m.bos_id,
2384 prompt,
2385 params,
2386 kv_pool,
2387 paged_kv,
2388 prefix_cache,
2389 ceiling,
2390 |chunk| {
2391 chunks.push(chunk.to_string());
2392 if !synthetic {
2393 emit(chunk);
2394 }
2395 },
2396 )?
2397 }
2398 }
2399 Model::Kimi(m) => generate::generate_engine(
2400 &m.engine,
2401 &m.tokenizer,
2402 &m.stop_tokens,
2403 None,
2404 prompt,
2405 params,
2406 |chunk| {
2407 chunks.push(chunk.to_string());
2408 if !synthetic {
2409 emit(chunk);
2410 }
2411 },
2412 )?,
2413 Model::Mla(m) => generate::generate_engine(
2414 &m.engine,
2415 &m.tokenizer,
2416 &m.stop_tokens,
2417 m.bos_id,
2418 prompt,
2419 params,
2420 |chunk| {
2421 chunks.push(chunk.to_string());
2422 if !synthetic {
2423 emit(chunk);
2424 }
2425 },
2426 )?,
2427 Model::Gemma4(m) => generate::generate_engine(
2428 &m.engine,
2429 &m.tokenizer,
2430 &m.stop_tokens,
2431 m.bos_id,
2432 prompt,
2433 params,
2434 |chunk| {
2435 chunks.push(chunk.to_string());
2436 if !synthetic {
2437 emit(chunk);
2438 }
2439 },
2440 )?,
2441 Model::Glm52(m) => generate::generate_engine(
2442 &m.engine,
2443 &m.tokenizer,
2444 &m.stop_tokens,
2445 m.bos_id,
2446 prompt,
2447 params,
2448 |chunk| {
2449 chunks.push(chunk.to_string());
2450 if !synthetic {
2451 emit(chunk);
2452 }
2453 },
2454 )?,
2455 };
2456
2457 let mut full = chunks.concat();
2458 if synthetic {
2459 full = format!(
2460 "[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
2461 to serve a real model. Decoded ids -> {full:?}]"
2462 );
2463 emit(&full);
2464 } else if used_batcher && !full.is_empty() {
2465 emit(&full);
2466 }
2467
2468 Ok((finish, usage, full))
2469}
2470
2471#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2476 model: &Model,
2477 prompt: &str,
2478 params: &GenerationParams,
2479 kv_pool: Option<&generate::KvPoolConfig>,
2480 paged_kv: Option<&generate::PagedKvConfig>,
2481 prefix_cache: Option<&Mutex<PrefixCache>>,
2482 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2483 ceiling: Option<&budget::ContextCeiling>,
2484) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2485 let (finish, usage, full) = run_generation_emit(
2486 model,
2487 prompt,
2488 params,
2489 kv_pool,
2490 paged_kv,
2491 prefix_cache,
2492 continuous_batcher,
2493 ceiling,
2494 |_| {},
2495 )?;
2496 Ok((
2497 if full.is_empty() {
2498 Vec::new()
2499 } else {
2500 vec![full]
2501 },
2502 finish,
2503 usage,
2504 ))
2505}
2506
2507pub(crate) fn prompt_from_messages(
2518 messages: &[ChatMessage],
2519 template: &chat_template::PromptTemplate,
2520 tools: &[ToolDef],
2521 extra: serde_json::Map<String, serde_json::Value>,
2522) -> Result<String, ApiError> {
2523 let rendered = if tools.is_empty() || template.handles_tools() {
2524 template.render(messages, tools, extra)
2525 } else {
2526 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2527 with_preamble.push(ChatMessage {
2528 role: "system".to_string(),
2529 content: Some(MessageContent::Text(tool_preamble(tools))),
2530 tool_calls: None,
2531 tool_call_id: None,
2532 reasoning_content: None,
2533 });
2534 with_preamble.extend_from_slice(messages);
2535 template.render(&with_preamble, &[], extra)
2536 };
2537 rendered.map_err(template_error_response)
2538}
2539
2540fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
2545 (
2546 StatusCode::BAD_REQUEST,
2547 Json(serde_json::json!({
2548 "error": {
2549 "message": format!("chat template failed to render: {err}"),
2550 "type": "invalid_request_error",
2551 "param": "messages",
2552 "code": null,
2553 }
2554 })),
2555 )
2556}
2557
2558fn tool_preamble(tools: &[ToolDef]) -> String {
2571 let mut out = String::from(
2572 "You can call tools to help answer the user. To call a tool, respond with \
2573 EXACTLY one line in this format and nothing else:\n\
2574 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2575 object matching that tool's parameters>}}</tool_call>\n\n\
2576 Available tools:\n",
2577 );
2578 for t in tools {
2579 out.push_str(&format!(
2580 "- {}: {}\n parameters (JSON schema): {}\n",
2581 t.function.name,
2582 t.function.description.as_deref().unwrap_or(""),
2583 t.function
2584 .parameters
2585 .as_ref()
2586 .map(|v| v.to_string())
2587 .unwrap_or_else(|| "{}".to_string()),
2588 ));
2589 }
2590 out
2591}
2592
2593fn tool_call_deltas(
2603 events: Vec<crate::policy::parser::ToolCallEvent>,
2604 opened: &std::cell::Cell<usize>,
2605) -> (String, Vec<ToolCallDelta>) {
2606 let mut text = String::new();
2607 let mut deltas = Vec::new();
2608 for event in events {
2609 match event {
2610 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2611 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2612 opened.set(opened.get().max(index + 1));
2613 deltas.push(ToolCallDelta::opening(index, name));
2614 }
2615 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2616 if !fragment.is_empty() {
2617 deltas.push(ToolCallDelta::arguments(index, fragment));
2618 }
2619 }
2620 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2621 }
2622 }
2623 (text, deltas)
2624}
2625
2626fn build_response_message(
2637 text: String,
2638 tools: &[ToolDef],
2639 posture: output::OutputPosture,
2640 base_finish: &'static str,
2641) -> (ChatCompletionResponseMessage, &'static str) {
2642 let parsed = output::parse_output(&text, tools, posture);
2643 let calls: Vec<ToolCallOut> = parsed
2644 .calls
2645 .into_iter()
2646 .enumerate()
2647 .map(|(index, call)| ToolCallOut {
2648 id: format!("call_{index}"),
2649 kind: "function",
2650 function: ToolCallFunctionOut {
2651 name: call.name,
2652 arguments: call.arguments,
2653 },
2654 })
2655 .collect();
2656 if !calls.is_empty() {
2657 return (
2658 ChatCompletionResponseMessage {
2659 role: "assistant",
2660 content: None,
2661 reasoning_content: parsed.reasoning,
2662 tool_calls: Some(calls),
2663 },
2664 "tool_calls",
2665 );
2666 }
2667 (
2668 ChatCompletionResponseMessage {
2669 role: "assistant",
2670 content: Some(parsed.content),
2671 reasoning_content: parsed.reasoning,
2672 tool_calls: None,
2673 },
2674 base_finish,
2675 )
2676}
2677
2678fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2683 let mut history = match &req.session_id {
2684 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2685 None => req.messages.clone(),
2686 };
2687 if req.json_object_mode() {
2688 inject_json_object_system_hint(&mut history);
2689 }
2690 history
2691}
2692
2693fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2694 const HINT: &str =
2695 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2696 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2697 match &mut sys.content {
2698 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2699 s.push_str("\n\n");
2700 s.push_str(HINT);
2701 }
2702 None => {
2703 sys.content = Some(MessageContent::Text(HINT.to_string()));
2704 }
2705 _ => {}
2706 }
2707 } else {
2708 messages.insert(
2709 0,
2710 ChatMessage {
2711 role: "system".to_string(),
2712 content: Some(MessageContent::Text(HINT.to_string())),
2713 tool_calls: None,
2714 tool_call_id: None,
2715 reasoning_content: None,
2716 },
2717 );
2718 }
2719}
2720
2721async fn chat_completions(
2722 State(state): State<Arc<AppState>>,
2723 headers: axum::http::HeaderMap,
2724 Json(req): Json<ChatCompletionRequest>,
2725) -> Response {
2726 let attribution = attribution::Attribution::from_headers(&headers);
2727 state
2728 .requests_total
2729 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2730 let started = std::time::Instant::now();
2731
2732 let request_id = ferrox_api::next_request_id();
2736 let stream = req.stream.unwrap_or(false);
2737
2738 let refusal = cache_admin::check_admission(&state)
2744 .err()
2745 .or_else(|| req.validate_supported_fields().err());
2746 if let Some(err) = refusal {
2747 state
2748 .request_errors_total
2749 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2750 let response = err.into_response();
2751 state.record_request(stats::Record {
2752 request_id: &request_id,
2753 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2754 model: state.active_model_name(),
2755 status: response.status().as_u16(),
2756 stream,
2757 duration_ms: started.elapsed().as_millis() as u64,
2758 usage: None,
2759 attribution: &attribution,
2760 });
2761 return response;
2762 }
2763
2764 let response = if stream {
2765 chat_completions_stream(
2766 Arc::clone(&state),
2767 req,
2768 request_id.clone(),
2769 started,
2770 attribution.clone(),
2771 )
2772 .await
2773 .into_response()
2774 } else {
2775 chat_completions_full(
2776 Arc::clone(&state),
2777 req,
2778 request_id.clone(),
2779 started,
2780 attribution.clone(),
2781 )
2782 .await
2783 .into_response()
2784 };
2785
2786 if response.status().is_client_error() || response.status().is_server_error() {
2787 state
2788 .request_errors_total
2789 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2790 state.record_request(stats::Record {
2794 request_id: &request_id,
2795 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2796 model: state.active_model_name(),
2799 status: response.status().as_u16(),
2800 stream,
2801 duration_ms: started.elapsed().as_millis() as u64,
2802 usage: None,
2803 attribution: &attribution,
2804 });
2805 }
2806 state.mark_request_finished();
2807
2808 response
2809}
2810
2811async fn chat_completions_full(
2812 state: Arc<AppState>,
2813 req: ChatCompletionRequest,
2814 request_id: String,
2815 started: std::time::Instant,
2816 attribution: attribution::Attribution,
2817) -> Result<Json<ChatCompletionResponse>, ApiError> {
2818 let tools_active = req.tools_active();
2819 let active = state.require_active()?;
2823 let history = resolve_history(&state, &req);
2824 let template = active.model.chat_template();
2825 let kwargs = req.resolve_template_kwargs(&template);
2826 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2827 let key = req.is_cacheable().then(|| req.cache_key(&prompt));
2828
2829 let (completion, cache_status) = if let Some(cached) = key
2830 .as_ref()
2831 .and_then(|key| lock_cache(&state.response_cache).get(key))
2832 {
2833 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2834 (cached, "hit")
2835 } else {
2836 let params = req.generation_params_for_template(&template, active.model.name())?;
2837 let (chunks, finish, usage) = decode_task::buffered(
2838 decode_task::DecodeHandles::take(&state, &active),
2839 prompt.clone(),
2840 params,
2841 )
2842 .await?;
2843
2844 let completion = response_cache::CachedCompletion {
2845 content: chunks.concat(),
2846 finish,
2847 usage,
2848 };
2849 let cache_status = if let Some(key) = key {
2850 tracing::debug!("cache miss for key {}", key.digest());
2851 lock_cache(&state.response_cache).put(key, completion.clone());
2852 "miss"
2853 } else {
2854 "skip"
2855 };
2856 (completion, cache_status)
2857 };
2858 let content = completion.content;
2859
2860 if req.json_object_mode() {
2861 json_mode::validate_json_object_output(&content)?;
2862 }
2863
2864 if let Some(id) = &req.session_id {
2868 state.sessions.store_reply(
2869 id,
2870 ChatMessage {
2871 role: "assistant".to_string(),
2872 content: Some(MessageContent::Text(content.clone())),
2873 tool_calls: None,
2874 tool_call_id: None,
2875 reasoning_content: None,
2876 },
2877 );
2878 }
2879
2880 let (message, finish_reason) = build_response_message(
2881 content,
2882 if tools_active { &req.tools } else { &[] },
2883 output::OutputPosture::resolve(active.model.name(), &prompt),
2884 completion.finish.as_str(),
2885 );
2886
2887 state.record_request(stats::Record {
2888 request_id: &request_id,
2889 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
2890 model: Some(active.model.name().to_string()),
2893 status: 200,
2894 stream: false,
2895 duration_ms: started.elapsed().as_millis() as u64,
2896 usage: Some(&completion.usage),
2897 attribution: &attribution,
2898 });
2899
2900 Ok(Json(ChatCompletionResponse {
2901 id: request_id.clone(),
2902 request_id,
2903 object: "chat.completion",
2904 model: req.model,
2905 choices: vec![ChatCompletionChoice {
2906 index: 0,
2907 message,
2908 finish_reason,
2909 }],
2910 usage: completion.usage,
2911 ferrox_cache: cache_status,
2912 }))
2913}
2914
2915async fn chat_completions_stream(
2916 state: Arc<AppState>,
2917 req: ChatCompletionRequest,
2918 request_id: String,
2919 started: std::time::Instant,
2920 attribution: attribution::Attribution,
2921) -> Result<Response, ApiError> {
2922 let tools_active = req.tools_active();
2924 let active = state.require_active()?;
2928 let history = resolve_history(&state, &req);
2929 let template = active.model.chat_template();
2930 let kwargs = req.resolve_template_kwargs(&template);
2931 let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
2932 let model_name = req.model.clone();
2933 let session_id = req.session_id.clone();
2934 let sessions = state.sessions.clone();
2935
2936 let model = Arc::clone(&active.model);
2937 let kv_pool = state.kv_pool.clone();
2938 let paged_kv = state.paged_kv.clone();
2939 let prefix_cache = state.prefix_cache.clone();
2940 let batcher = active.batcher.clone();
2941 let ceiling = active.ceiling.clone();
2942 let mut params = req.generation_params_for_template(&template, active.model.name())?;
2943 let stats_state = Arc::clone(&state);
2944 let served_model = active.model.name().to_string();
2947 let posture = output::OutputPosture::resolve(&served_model, &prompt);
2951 let offered_tools: Vec<ToolDef> = if tools_active {
2954 req.tools.clone()
2955 } else {
2956 Vec::new()
2957 };
2958
2959 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2964 params.cancel = Some(cancel_token.clone());
2965
2966 let overlap = batcher.is_none();
2975
2976 let slot = req
2980 .stream_resumable
2981 .unwrap_or(false)
2982 .then(|| state.streams.register(&request_id));
2983 let emitter = resume::Emitter::new(slot);
2984
2985 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2986 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2990 id: request_id.clone(),
2991 request_id: None,
2992 object: "chat.completion.chunk",
2993 model: model_name.clone(),
2994 choices: vec![ChatCompletionChunkChoice {
2995 index: 0,
2996 delta: ChatCompletionChunkDelta {
2997 role: None,
2998 content: None,
2999 reasoning_content: None,
3000 tool_calls: None,
3001 },
3002 finish_reason: None,
3003 }],
3004 usage: None,
3005 });
3006
3007 tokio::task::spawn_blocking(move || {
3008 let _cancel_guard = cancel_guard;
3011 let tx_chunks = tx.clone();
3012 let orphan_timeout = sse::orphan_timeout_from_env();
3017 let mut first = true;
3018 let head_request_id = request_id.clone();
3019 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3029 Rc::new(RefCell::new(posture.reasoning_parser()));
3030 let emit_reasoning = Rc::clone(&stream_reasoning);
3031 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3035 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3036 );
3037 let emit_tools = Rc::clone(&stream_tools);
3038 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3042 let emit_streamed_calls = Rc::clone(&streamed_calls);
3043 let result = run_generation_emit(
3044 &model,
3045 &prompt,
3046 ¶ms,
3047 kv_pool.as_ref(),
3048 paged_kv.as_ref(),
3049 prefix_cache.as_deref(),
3050 batcher.as_ref(),
3051 ceiling.as_deref(),
3052 |chunk| {
3053 if !overlap || chunk.is_empty() {
3054 return;
3055 }
3056 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3057 Some(parser) => {
3058 let delta = parser.push(chunk);
3059 (delta.reasoning, delta.content)
3060 }
3061 None => (String::new(), chunk.to_string()),
3062 };
3063 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3067 Some(parser) => {
3068 let (text, calls) =
3069 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3070 (text, calls)
3071 }
3072 None => (content, Vec::new()),
3073 };
3074 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3077 return;
3078 }
3079 let role = if first { Some("assistant") } else { None };
3080 let request_id = first.then(|| head_request_id.clone());
3081 first = false;
3082 let payload = ChatCompletionChunk {
3083 id: head_request_id.clone(),
3084 request_id,
3085 object: "chat.completion.chunk",
3086 model: model_name.clone(),
3087 choices: vec![ChatCompletionChunkChoice {
3088 index: 0,
3089 delta: ChatCompletionChunkDelta {
3090 role,
3091 content: (!content.is_empty()).then_some(content),
3092 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3093 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3094 },
3095 finish_reason: None,
3096 }],
3097 usage: None,
3098 };
3099 if let Err(why) =
3108 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3109 {
3110 if why == sse::SendFailure::Orphaned {
3111 tracing::warn!(
3112 "SSE stream {head_request_id} accepted nothing for the orphan \
3113 deadline; treating it as abandoned"
3114 );
3115 }
3116 if !emitter.is_resumable() {
3128 cancel_token.cancel();
3129 }
3130 }
3131 },
3132 );
3133
3134 let mut pending_request_id = first.then(|| request_id.clone());
3140
3141 match result {
3142 Ok((finish, usage, full_text)) => {
3143 if let Some(id) = &session_id {
3144 sessions.store_reply(
3145 id,
3146 ChatMessage {
3147 role: "assistant".to_string(),
3148 content: Some(MessageContent::Text(full_text.clone())),
3149 tool_calls: None,
3150 tool_call_id: None,
3151 reasoning_content: None,
3152 },
3153 );
3154 }
3155 let mut streamed_finish: Option<&'static str> = None;
3161 if overlap {
3162 let tail = stream_reasoning
3163 .borrow_mut()
3164 .as_mut()
3165 .map(|parser| parser.flush())
3166 .unwrap_or_default();
3167 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3168 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3169 let mut events = parser.push(&content);
3170 events.extend(parser.finish());
3171 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3172 content = text;
3173 tool_calls = calls;
3174 }
3175 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3176 let payload = ChatCompletionChunk {
3177 id: request_id.clone(),
3178 request_id: pending_request_id.take(),
3179 object: "chat.completion.chunk",
3180 model: model_name.clone(),
3181 choices: vec![ChatCompletionChunkChoice {
3182 index: 0,
3183 delta: ChatCompletionChunkDelta {
3184 role: None,
3185 content: (!content.is_empty()).then_some(content),
3186 reasoning_content: (!tail.reasoning.is_empty())
3187 .then_some(tail.reasoning),
3188 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3189 },
3190 finish_reason: None,
3191 }],
3192 usage: None,
3193 };
3194 let _ =
3195 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3196 }
3197 if streamed_calls.get() > 0 {
3198 streamed_finish = Some("tool_calls");
3199 }
3200 } else {
3201 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3204 let tool_calls: Vec<ToolCallDelta> = parsed
3205 .calls
3206 .iter()
3207 .enumerate()
3208 .map(|(index, call)| {
3209 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3210 })
3211 .collect();
3212 if !tool_calls.is_empty() {
3213 streamed_finish = Some("tool_calls");
3214 }
3215 if !tool_calls.is_empty()
3216 || !parsed.content.is_empty()
3217 || parsed.reasoning.is_some()
3218 {
3219 let payload = ChatCompletionChunk {
3220 id: request_id.clone(),
3221 request_id: pending_request_id.take(),
3222 object: "chat.completion.chunk",
3223 model: model_name.clone(),
3224 choices: vec![ChatCompletionChunkChoice {
3225 index: 0,
3226 delta: ChatCompletionChunkDelta {
3227 role: Some("assistant"),
3228 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3229 .then(|| parsed.content.clone()),
3230 reasoning_content: parsed.reasoning.clone(),
3231 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3232 },
3233 finish_reason: None,
3234 }],
3235 usage: None,
3236 };
3237 let _ =
3238 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3239 }
3240 }
3241 let final_finish_reason = match streamed_finish {
3245 Some(reason) if finish.as_str() != "length" => reason,
3246 _ => finish.as_str(),
3247 };
3248 let final_payload = ChatCompletionChunk {
3249 id: request_id.clone(),
3250 request_id: pending_request_id.take(),
3251 object: "chat.completion.chunk",
3252 model: model_name,
3253 choices: vec![ChatCompletionChunkChoice {
3254 index: 0,
3255 delta: ChatCompletionChunkDelta {
3256 role: None,
3257 content: None,
3258 reasoning_content: None,
3259 tool_calls: None,
3260 },
3261 finish_reason: Some(final_finish_reason),
3262 }],
3263 usage: Some(usage.clone()),
3264 };
3265 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3266 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3267 stats_state.record_request(stats::Record {
3272 request_id: &request_id,
3273 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3274 model: Some(served_model.clone()),
3275 status: 200,
3276 stream: true,
3277 duration_ms: started.elapsed().as_millis() as u64,
3278 usage: Some(&usage),
3279 attribution: &attribution,
3280 });
3281 }
3282 Err(e) => {
3283 tracing::warn!("decode error on streamed request {request_id}: {e}");
3284 stats_state.record_request(stats::Record {
3290 request_id: &request_id,
3291 route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
3292 model: Some(served_model.clone()),
3293 status: 500,
3294 stream: true,
3295 duration_ms: started.elapsed().as_millis() as u64,
3296 usage: None,
3297 attribution: &attribution,
3298 });
3299 let payload = ChatCompletionChunk {
3300 id: request_id.clone(),
3301 request_id: pending_request_id.take(),
3302 object: "chat.completion.chunk",
3303 model: model_name,
3304 choices: vec![ChatCompletionChunkChoice {
3305 index: 0,
3306 delta: ChatCompletionChunkDelta {
3307 role: Some("assistant"),
3308 content: Some(format!("[error: {e}]")),
3309 reasoning_content: None,
3310 tool_calls: None,
3311 },
3312 finish_reason: Some("stop"),
3313 }],
3314 usage: None,
3315 };
3316 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3317 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3318 }
3319 }
3320 drop(emitter);
3324 });
3325
3326 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3327 Ok((
3350 [(
3351 axum::http::HeaderName::from_static("x-accel-buffering"),
3352 axum::http::HeaderValue::from_static("no"),
3353 )],
3354 Sse::new(stream),
3355 )
3356 .into_response())
3357}
3358
3359fn axum_path(template: &str) -> String {
3376 let mut out = String::with_capacity(template.len());
3377 let mut rest = template;
3378 while let Some(open) = rest.find('{') {
3379 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3380 break;
3381 };
3382 out.push_str(&rest[..open]);
3383 out.push(':');
3384 out.push_str(&rest[open + 1..close]);
3385 rest = &rest[close + 1..];
3386 }
3387 out.push_str(rest);
3388 out
3389}
3390
3391async fn cancel_generation(
3400 State(state): State<Arc<AppState>>,
3401 Json(req): Json<ferrox_api::CancelGenerationRequest>,
3402) -> Response {
3403 let cancelled = state.cancels.cancel(&req.request_id);
3404 let status = if cancelled {
3405 StatusCode::OK
3406 } else {
3407 StatusCode::NOT_FOUND
3408 };
3409 let detail = if cancelled {
3410 "the generation was asked to stop; it ends at its next token".to_string()
3411 } else {
3412 "no generation with that request_id is running -- it has already \
3413 finished, was never issued, or was served by a path that does \
3414 not register for cancellation"
3415 .to_string()
3416 };
3417 (
3418 status,
3419 Json(ferrox_api::CancelGenerationResponse {
3420 request_id: req.request_id,
3421 cancelled,
3422 detail,
3423 }),
3424 )
3425 .into_response()
3426}
3427
3428type Activated = (
3432 Model,
3433 Option<serving::batch::ContinuousBatcher>,
3434 Option<Arc<budget::ContextCeiling>>,
3435);
3436
3437fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3455 let mut batcher = serving::batch::BatcherConfig::from_env();
3456 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3457 return batcher;
3460 }
3461 let Some(path) = path else {
3462 return batcher;
3463 };
3464 let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
3470 let Some((priced, gguf_ctx, source)) = priced else {
3471 return batcher;
3472 };
3473 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3474 tracing::warn!(
3478 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3479 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3480 FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
3481 admit on a number you choose.",
3482 priced.weights_bytes,
3483 priced.device_budget_bytes,
3484 );
3485 return batcher;
3486 };
3487 tracing::info!("{source}");
3488 tracing::info!("{}", derived.fit);
3489 let adopted = budget::apply_derived(&mut batcher, &derived);
3490 if adopted.max_context {
3491 tracing::info!(
3492 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3493 override with FERROX_CB_MAX_CONTEXT",
3494 derived.max_context
3495 );
3496 }
3497 if adopted.kv_blocks {
3498 tracing::info!(
3499 "derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
3500 derived.kv_blocks,
3501 batcher.kv_block_size
3502 );
3503 }
3504 batcher
3505}
3506
3507pub(crate) fn activate_loaded_model(
3516 loaded: model::LoadedModel,
3517 enable_continuous_batching: bool,
3518 path: Option<&str>,
3519 paged_kv: Option<&generate::PagedKvConfig>,
3520) -> Activated {
3521 match loaded {
3522 model::LoadedModel::Gguf(g) => {
3523 let decoder = Arc::new(g.decoder);
3524 let tokenizer = Arc::new(g.tokenizer);
3525 let config = price_batcher_config(path);
3526 let shape =
3532 ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
3533 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3534 let batcher = if enable_continuous_batching {
3535 tracing::info!(
3536 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3537 (stop sequences use the same pending-buffer trim as the private generate loop)"
3538 );
3539 let tok = Arc::clone(&tokenizer);
3540 let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
3541 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3542 Arc::clone(&decoder),
3543 decode,
3544 config,
3545 Arc::clone(&ceiling),
3546 paged_kv.cloned(),
3547 ))
3548 } else {
3549 None
3550 };
3551 (
3552 Model::Gguf(GgufModel {
3553 decoder,
3554 tokenizer,
3555 stop_tokens: g.stop_tokens,
3556 bos_id: g.bos_id,
3557 is_synthetic: g.is_synthetic,
3558 chat_template: g.chat_template,
3559 }),
3560 batcher,
3561 Some(ceiling),
3562 )
3563 }
3564 model::LoadedModel::Kimi(k) => (
3565 Model::Kimi(KimiModel {
3566 engine: k.engine,
3567 tokenizer: k.tokenizer,
3568 stop_tokens: k.stop_tokens,
3569 chat_template: k.chat_template,
3570 }),
3571 None,
3572 None,
3573 ),
3574 model::LoadedModel::Mla(m) => (
3575 Model::Mla(MlaModel {
3576 engine: m.engine,
3577 tokenizer: m.tokenizer,
3578 stop_tokens: m.stop_tokens,
3579 bos_id: m.bos_id,
3580 name: m.name,
3581 chat_template: m.chat_template,
3582 }),
3583 None,
3584 None,
3585 ),
3586 model::LoadedModel::Gemma4(m) => (
3587 Model::Gemma4(Gemma4Model {
3588 engine: m.engine,
3589 tokenizer: m.tokenizer,
3590 stop_tokens: m.stop_tokens,
3591 bos_id: m.bos_id,
3592 name: m.name,
3593 chat_template: m.chat_template,
3594 }),
3595 None,
3596 None,
3597 ),
3598 model::LoadedModel::Glm52(g) => (
3599 Model::Glm52(Glm52Model {
3600 engine: g.engine,
3601 tokenizer: g.tokenizer,
3602 stop_tokens: g.stop_tokens,
3603 bos_id: g.bos_id,
3604 name: g.name,
3605 chat_template: g.chat_template,
3606 }),
3607 None,
3608 None,
3609 ),
3610 }
3611}
3612
3613fn build_app_state(
3614 loaded: model::LoadedModel,
3615 kv_pool: Option<generate::KvPoolConfig>,
3616 paged_kv: Option<generate::PagedKvConfig>,
3617 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3618 enable_continuous_batching: bool,
3619 mcp: Option<mcp::LoadedMcpConfig>,
3620 detection: Arc<health::Detection>,
3621) -> AppState {
3622 let (model, batcher, ceiling) = activate_loaded_model(
3623 loaded,
3624 enable_continuous_batching,
3625 std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
3626 paged_kv.as_ref(),
3627 );
3628 let id = startup_model_id();
3634 AppState {
3635 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3636 id,
3637 model: Arc::new(model),
3638 batcher,
3639 ceiling,
3640 }))),
3641 paged_kv,
3642 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3643 tasks: Arc::new(tasks::TaskRegistry::new()),
3644 cancels: Arc::new(cancel::CancelRegistry::new()),
3645 stats: stats::Stats::new(),
3646 streams: resume::StreamRegistry::new(),
3647 model_dir: admin::model_dirs().into_iter().next(),
3648 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3649 kv_pool,
3650 prefix_cache,
3651 sessions: session::SessionStore::new(),
3652 requests_total: std::sync::atomic::AtomicU64::new(0),
3653 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3654 started_at: std::time::Instant::now(),
3655 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3656 detection,
3657 mcp,
3658 continuous_batching_enabled: enable_continuous_batching,
3659 loading_model: Mutex::new(None),
3660 last_load_error: Mutex::new(None),
3661 serving: Mutex::new(crate::stats::ServingStats::default()),
3662 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3663 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3664 started_unix: unix_now(),
3665 }
3666}
3667
3668fn unix_now() -> u64 {
3673 std::time::SystemTime::now()
3674 .duration_since(std::time::UNIX_EPOCH)
3675 .map(|d| d.as_secs())
3676 .unwrap_or(0)
3677}
3678
3679fn startup_model_id() -> Option<String> {
3684 let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
3685 let configured = std::fs::canonicalize(&configured).ok()?;
3686 admin::discover(&admin::model_dirs())
3687 .into_iter()
3688 .find(|d| {
3689 std::fs::canonicalize(&d.path)
3690 .map(|p| p == configured)
3691 .unwrap_or(false)
3692 })
3693 .map(|d| d.id)
3694}
3695
3696fn init_cpu_pool() {
3705 match ferrox_core::threads::init_cpu_pool() {
3706 Some(n) => eprintln!(
3707 "ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
3708 ferrox_core::threads::perf_core_count()
3709 ),
3710 None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
3711 }
3712}
3713
3714fn announce_ready(addr: SocketAddr, scheme: &str) {
3727 use std::io::Write;
3728 let ready =
3729 ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3730 let mut stdout = std::io::stdout().lock();
3731 let _ = writeln!(stdout, "{}", ready.to_line());
3732 let _ = stdout.flush();
3733}
3734
3735async fn shutdown_signal(exit_on_stdin_close: bool) {
3747 if !exit_on_stdin_close {
3748 std::future::pending::<()>().await;
3749 return;
3750 }
3751 let _ = tokio::task::spawn_blocking(|| {
3752 use std::io::Read;
3753 let mut sink = [0u8; 256];
3754 let mut stdin = std::io::stdin().lock();
3755 loop {
3756 match stdin.read(&mut sink) {
3757 Ok(0) => break,
3759 Ok(_) => continue,
3761 Err(e) => {
3762 tracing::warn!("stdin read failed ({e}); treating it as closed");
3763 break;
3764 }
3765 }
3766 }
3767 })
3768 .await;
3769 tracing::info!("stdin closed; shutting down");
3770}
3771
3772fn tokio_worker_threads() -> usize {
3777 std::env::var("FERROX_TOKIO_WORKERS")
3778 .ok()
3779 .and_then(|v| v.trim().parse::<usize>().ok())
3780 .filter(|n| *n > 0)
3781 .unwrap_or(2)
3782}
3783
3784fn install_ring_crypto_provider() {
3801 let _ = rustls::crypto::ring::default_provider().install_default();
3802}
3803
3804pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3810 if args.list_devices {
3811 print_available_devices();
3812 return Ok(());
3813 }
3814 apply_cli_overrides(&args)?;
3815
3816 let _instance = {
3820 use ferrox_core::instance::{register, InstancePolicy};
3821 let policy = if args.allow_multiple_instances {
3822 InstancePolicy::Multi
3823 } else {
3824 InstancePolicy::from_env_or(InstancePolicy::Single)
3825 };
3826 let model = std::env::var("FERROX_MODEL_PATH").ok();
3827 register(
3828 "server",
3829 model.as_deref(),
3830 ferrox_core::instance::current_backend(),
3831 policy,
3832 )
3833 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
3834 };
3835
3836 let journal = journal::Journal::from_env();
3837 eprintln!(
3838 "ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
3839 journal.path()
3840 );
3841 journal.append(&journal::Record::session_start(
3842 env!("CARGO_PKG_VERSION"),
3843 std::process::id(),
3844 ));
3845 journal::install_panic_hook(journal.clone());
3846
3847 let mcp_config_path = args.mcp_config.clone();
3848 let exit_on_stdin_close = args.exit_on_stdin_close
3849 || std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
3850 .map(|v| v == "1")
3851 .unwrap_or(false);
3852
3853 unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
3857 init_cpu_pool();
3858
3859 let runtime = tokio::runtime::Builder::new_multi_thread()
3860 .worker_threads(tokio_worker_threads())
3861 .enable_all()
3862 .build()?;
3863 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
3864
3865 let reason = match &result {
3866 Ok(()) => "normal".to_string(),
3867 Err(e) => e.to_string(),
3868 };
3869 journal.append(&journal::Record::session_exit(reason));
3870
3871 runtime.shutdown_background();
3876
3877 result
3878}
3879
3880async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
3881 let _ = tracing_subscriber::fmt::try_init();
3887
3888 let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
3898 let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
3899 let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
3900 .map(|v| v == "1")
3901 .unwrap_or(false);
3902 if let Err(msg) =
3903 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
3904 {
3905 anyhow::bail!(msg);
3906 }
3907
3908 let mut loaded = model::load()?;
3909 match &loaded {
3910 model::LoadedModel::Gguf(g) => tracing::info!(
3911 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
3912 g.decoder.config.name,
3913 g.is_synthetic,
3914 g.tokenizer.kind()
3915 ),
3916 model::LoadedModel::Kimi(k) => tracing::info!(
3917 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
3918 k.tokenizer.vocab_size()
3919 ),
3920 model::LoadedModel::Mla(m) => tracing::info!(
3921 "loaded MLA GGUF '{}' (tokenizer={})",
3922 m.name,
3923 m.tokenizer.kind()
3924 ),
3925 model::LoadedModel::Gemma4(m) => tracing::info!(
3926 "loaded Gemma4 GGUF '{}' (tokenizer={})",
3927 m.name,
3928 m.tokenizer.kind()
3929 ),
3930 model::LoadedModel::Glm52(g) => tracing::info!(
3931 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
3932 g.name,
3933 g.tokenizer.kind()
3934 ),
3935 }
3936 let metal_default_moe_budget = {
3944 #[cfg(feature = "metal")]
3945 {
3946 ferrox_core::metal_dense_enabled()
3947 && std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
3948 }
3949 #[cfg(not(feature = "metal"))]
3950 {
3951 false
3952 }
3953 };
3954 if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
3955 let budget: u64 = budget_str
3956 .parse()
3957 .expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
3958 match &mut loaded {
3959 model::LoadedModel::Gguf(g) => {
3960 tracing::info!(
3961 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
3962 (CUDA and/or Metal matvecs when built with the matching feature)"
3963 );
3964 g.decoder.gpu_vram_budget_bytes = Some(budget);
3965 }
3966 model::LoadedModel::Kimi(_) => {
3967 tracing::warn!(
3968 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
3969 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
3970 );
3971 }
3972 model::LoadedModel::Mla(_) => {
3973 tracing::warn!(
3974 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
3975 FFN path only today; ignoring expert VRAM budget"
3976 );
3977 }
3978 model::LoadedModel::Gemma4(_) => {
3979 tracing::warn!(
3980 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
3981 ignoring expert VRAM budget"
3982 );
3983 }
3984 model::LoadedModel::Glm52(_) => {
3985 tracing::warn!(
3986 "FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
3987 GPU expert placement not wired yet; ignoring"
3988 );
3989 }
3990 }
3991 } else if metal_default_moe_budget {
3992 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
3996 if let model::LoadedModel::Gguf(g) = &mut loaded {
3997 tracing::info!(
3998 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
3999 set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4000 );
4001 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4002 }
4003 }
4004 #[cfg(feature = "cuda")]
4005 {
4006 if ferrox_core::cuda_dense_enabled() {
4007 tracing::info!(
4008 "CUDA dense matvec enabled for WeightMatrix::apply \
4009 (FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4010 );
4011 } else {
4012 tracing::info!(
4013 "CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
4014 );
4015 }
4016 }
4017 #[cfg(feature = "metal")]
4018 {
4019 if ferrox_core::metal_dense_enabled() {
4020 tracing::info!(
4021 "Metal dense matvec enabled for WeightMatrix::apply \
4022 (FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4023 );
4024 match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
4025 Some("1") | Some("true") | Some("on") | Some("attn") => {
4026 tracing::info!(
4027 "Metal fused attention requested (FERROX_METAL_ATTN): \
4028 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4029 );
4030 }
4031 _ => {}
4032 }
4033 tracing::info!(
4034 "Metal greedy GPU argmax: temperature<=0 folds \
4035 final_norm+lm_head+argmax into the dense stack"
4036 );
4037 } else {
4038 tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
4039 }
4040 }
4041 let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
4052 let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
4053 let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
4054 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4055 panic!(
4056 "FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
4057 (set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
4058 );
4059 }
4060 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4061 (Ok(blocks), Ok(block_size), Err(_)) => {
4062 let total_blocks: usize = blocks
4063 .parse()
4064 .expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
4065 let block_size: usize = block_size
4066 .parse()
4067 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4068 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4074 .ok()
4075 .map(|v| {
4076 v.parse()
4077 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4078 })
4079 .unwrap_or(0);
4080 tracing::info!(
4081 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4082 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4083 );
4084 Some(generate::KvPoolConfig {
4085 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4086 queue_wait: Duration::from_millis(queue_wait_ms),
4087 })
4088 }
4089 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4090 let block_size: usize = block_size
4091 .parse()
4092 .expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
4093 let budget: u64 = byte_budget
4094 .parse()
4095 .expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
4096 let cfg = match &loaded {
4097 model::LoadedModel::Gguf(g) => &g.decoder.config,
4098 model::LoadedModel::Kimi(_)
4099 | model::LoadedModel::Mla(_)
4100 | model::LoadedModel::Gemma4(_)
4101 | model::LoadedModel::Glm52(_) => {
4102 panic!(
4103 "FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
4104 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4105 );
4106 }
4107 };
4108 let bytes_per_block = block_size
4109 * cfg.n_layers
4110 * cfg.n_kv_heads
4111 * cfg.head_dim
4112 * 2
4113 * std::mem::size_of::<f32>();
4114 assert!(
4115 bytes_per_block > 0,
4116 "derived KV block byte size must be positive (check model config and block size)"
4117 );
4118 let total_blocks = (budget as usize / bytes_per_block).max(1);
4119 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4120 .ok()
4121 .map(|v| {
4122 v.parse()
4123 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4124 })
4125 .unwrap_or(0);
4126 tracing::info!(
4127 "KV cache block pool enabled from byte budget: {budget} bytes / \
4128 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4129 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4130 cfg.n_layers
4131 );
4132 Some(generate::KvPoolConfig {
4133 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4134 queue_wait: Duration::from_millis(queue_wait_ms),
4135 })
4136 }
4137 (Err(_), Err(_), Err(_)) => None,
4138 (Err(_), Ok(_), Err(_)) => panic!(
4139 "FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
4140 (or unset all three to disable KV cache pooling)"
4141 ),
4142 (Ok(_), Ok(_), Ok(_)) => {
4143 unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
4144 }
4145 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4146 "FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
4147 set together (or neither, to disable KV cache pooling)"
4148 ),
4149 };
4150 let paged_kv = match (
4156 std::env::var("FERROX_PAGED_KV_BLOCKS"),
4157 std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
4158 ) {
4159 (Ok(blocks), Ok(block_size)) => {
4160 assert!(
4161 kv_pool.is_none(),
4162 "FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
4163 mutually exclusive: both bound the same KV memory, by different means. \
4164 Set one."
4165 );
4166 let blocks_per_layer: usize = blocks
4180 .parse()
4181 .expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
4182 let block_size: usize = block_size
4183 .parse()
4184 .expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
4185 let gguf = match &loaded {
4186 model::LoadedModel::Gguf(g) => g,
4187 _ => panic!(
4188 "FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
4189 (set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
4190 ),
4191 };
4192 let cfg = &gguf.decoder.config;
4193 let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
4194 .ok()
4195 .map(|v| {
4196 v.parse()
4197 .expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4198 })
4199 .unwrap_or(0);
4200 tracing::info!(
4201 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4202 layer across {} layers, shared by all concurrent requests, \
4203 {queue_wait_ms}ms admission queue wait",
4204 cfg.n_layers
4205 );
4206 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4211 block_size,
4212 ))));
4213 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4222 crate::policy::parser::ToolCallFormat::infer(
4223 &std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
4224 )
4225 .opener(),
4226 |text| {
4227 gguf.tokenizer
4228 .encode(text)
4229 .into_iter()
4230 .map(|t| t as u32)
4231 .collect()
4232 },
4233 );
4234 if let Some(id) = anchor_token {
4235 tracing::info!(
4236 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4237 window stops short of where its next turn rejoins"
4238 );
4239 }
4240 let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
4241 .ok()
4242 .map(|v| {
4243 v.parse()
4244 .expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4245 })
4246 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4247 if let Some(window) = cfg.uniform_sliding_window() {
4248 tracing::info!(
4249 "Paged KV window slide enabled: every layer slides by {window} every \
4250 {slide_interval} decode steps, so a request holds its prompt and a \
4251 window rather than its whole context"
4252 );
4253 } else if cfg.kv_block_window().is_some() {
4254 tracing::info!(
4255 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4256 and a page group holds one block in every layer"
4257 );
4258 }
4259 Some(generate::PagedKvConfig {
4260 store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
4261 cfg.n_layers,
4262 block_size,
4263 blocks_per_layer,
4264 cfg.n_kv_heads,
4265 cfg.head_dim,
4266 )),
4267 queue_wait: Duration::from_millis(queue_wait_ms),
4268 radix,
4269 anchor_token,
4270 slide_interval,
4271 })
4272 }
4273 (Err(_), Err(_)) => None,
4274 _ => panic!(
4275 "FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
4276 (or neither, to disable paged KV)"
4277 ),
4278 };
4279 let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4286 let max_entries: usize = v
4287 .parse()
4288 .expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
4289 if kv_pool.is_some() {
4290 tracing::warn!(
4291 "FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4292 caching will never be consulted while a KV pool is configured"
4293 );
4294 }
4295 assert!(
4302 paged_kv.is_none(),
4303 "FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
4304 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4305 produce, so the cache could never hit. Set one."
4306 );
4307 tracing::info!(
4308 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4309 all requests"
4310 );
4311 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4312 });
4313 if matches!(
4314 loaded,
4315 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4316 ) && (kv_pool.is_some() || prefix_cache.is_some())
4317 {
4318 tracing::warn!(
4319 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4320 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4321 ferrox_models::engine's module docs"
4322 );
4323 }
4324 let enable_cb = std::env::var("FERROX_CONTINUOUS_BATCHING")
4325 .map(|v| v == "1")
4326 .unwrap_or(false)
4327 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
4335 && matches!(loaded, model::LoadedModel::Gguf(_));
4336 if std::env::var("FERROX_CONTINUOUS_BATCHING")
4337 .map(|v| v == "1")
4338 .unwrap_or(false)
4339 && (kv_pool.is_some() || prefix_cache.is_some())
4340 {
4341 tracing::warn!(
4342 "FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4343 (those modes keep the private generate path)"
4344 );
4345 }
4346 if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
4347 if let Ok(chunk) = n.parse::<usize>() {
4348 if chunk > 0 {
4349 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4350 }
4351 }
4352 }
4353 if matches!(
4354 std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
4355 Some("1")
4356 ) {
4357 tracing::warn!(
4358 "FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4359 (minimal spill; full layer offload still planned)"
4360 );
4361 }
4362
4363 let mcp = match mcp_config_path {
4364 Some(path) => {
4365 let loaded = mcp::load_mcp_config(&path)?;
4366 tracing::info!(
4367 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4368 loaded.path,
4369 loaded.servers.len()
4370 );
4371 Some(loaded)
4372 }
4373 None => None,
4374 };
4375
4376 let detection = health::Detection::spawn();
4380
4381 let state = Arc::new(build_app_state(
4382 loaded,
4383 kv_pool,
4384 paged_kv,
4385 prefix_cache,
4386 enable_cb,
4387 mcp,
4388 detection,
4389 ));
4390
4391 use ferrox_api::routes;
4395
4396 let public = Router::new().route(routes::HEALTH, get(health));
4401
4402 let mut protected = Router::new()
4403 .route(routes::V1_MODELS, get(list_models))
4404 .route(routes::V1_RESPONSES, post(responses::responses))
4408 .route(
4409 &axum_path(routes::V1_RESPONSE),
4410 get(responses::responses_get),
4411 )
4412 .route(
4413 &axum_path(routes::V1_RESPONSE_CANCEL),
4414 post(responses::responses_cancel),
4415 )
4416 .route(routes::V1_STATS, get(serving_stats))
4417 .route(routes::V1_REQUESTS, get(recent_requests))
4418 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
4419 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
4420 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
4421 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
4422 .route(routes::V1_CANCEL, post(cancel_generation))
4426 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
4431 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
4432 .route(routes::V1_MESSAGES, post(anthropic::messages))
4433 .route(
4434 routes::V1_MESSAGES_COUNT_TOKENS,
4435 post(anthropic::count_tokens),
4436 )
4437 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
4438 .route(routes::COMPLETION, post(completion::completion))
4443 .route(routes::COMPLETIONS, post(completion::completion))
4444 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
4445 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
4446 .route(routes::TOKENIZE, post(openai_extra::tokenize))
4452 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
4453 .route(routes::V1_EMBEDDINGS, post(openai_extra::embeddings))
4454 .route(routes::CACHE_STATS, get(cache_stats))
4455 .route(routes::METRICS, get(metrics))
4456 .route(routes::ADMIN_MODELS, get(admin::models))
4461 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
4462 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
4463 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
4464 .route(routes::ADMIN_TASKS, get(admin::tasks))
4465 .route(&admin::cancel_route(), post(admin::cancel_task))
4466 .route(routes::ADMIN_STATS, get(admin::stats))
4467 .merge(conversations::router());
4471
4472 if let Ok(key) = std::env::var("FERROX_API_KEY") {
4477 tracing::info!("API key auth enabled");
4478 let auth = limits::AuthConfig {
4479 api_key: Arc::new(key),
4480 };
4481 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4482 auth,
4483 limits::require_api_key,
4484 ));
4485 }
4486 if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
4487 let rpm: u32 = rpm
4488 .parse()
4489 .expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
4490 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4491 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4492 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4493 limiter,
4494 limits::rate_limit,
4495 ));
4496 }
4497 if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
4505 let origins = security::parse_cors_origins(&spec)
4506 .unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
4507 tracing::info!(
4508 "CORS enabled: {} allow-listed origin(s) ({})",
4509 origins.len(),
4510 spec
4511 );
4512 let cors = tower_http::cors::CorsLayer::new()
4513 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4514 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4515 .allow_headers([
4516 axum::http::header::CONTENT_TYPE,
4517 axum::http::header::AUTHORIZATION,
4518 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4524 axum::http::HeaderName::from_static("last-event-id"),
4528 ]);
4529 protected = protected.route_layer(cors);
4530 }
4531
4532 let app = public
4536 .merge(protected)
4537 .layer(axum::middleware::from_fn(limits::retry_after))
4538 .with_state(state);
4539
4540 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4546 install_ring_crypto_provider();
4547 match tls_paths {
4552 Some(paths) => {
4553 let config =
4554 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4555 .await
4556 .map_err(|e| {
4557 anyhow::anyhow!(
4558 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4559 paths.cert,
4560 paths.key
4561 )
4562 })?;
4563 let socket_addr: std::net::SocketAddr = addr
4564 .parse()
4565 .map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
4566 let listener = std::net::TcpListener::bind(socket_addr)?;
4567 listener.set_nonblocking(true)?;
4575 let bound = listener.local_addr()?;
4576 tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
4577 announce_ready(bound, "https");
4578
4579 let handle = axum_server::Handle::new();
4580 let shutdown_handle = handle.clone();
4581 tokio::spawn(async move {
4582 shutdown_signal(exit_on_stdin_close).await;
4583 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4584 });
4585 axum_server::from_tcp_rustls(listener, config)?
4586 .handle(handle)
4587 .serve(app.into_make_service())
4588 .await?;
4589 }
4590 None => {
4591 let listener = tokio::net::TcpListener::bind(&addr).await?;
4592 let bound = listener.local_addr()?;
4593 tracing::info!("ferrox-server listening on {bound}");
4594 announce_ready(bound, "http");
4595 axum::serve(listener, app)
4596 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4597 .await?;
4598 }
4599 }
4600 Ok(())
4601}
4602
4603#[cfg(test)]
4604mod tests {
4605 use super::*;
4606 use ferrox_models::config::test_dense_fixture;
4607
4608 #[test]
4609 fn parses_llama_server_style_options() {
4610 let argv = [
4611 "ferrox-server",
4612 "-m",
4613 "model.gguf",
4614 "--host",
4615 "::1",
4616 "--port",
4617 "9000",
4618 "-t",
4619 "4",
4620 "-dev",
4621 "Metal",
4622 "-ngl",
4623 "all",
4624 ]
4625 .into_iter()
4626 .map(String::from)
4627 .collect();
4628 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4629
4630 assert_eq!(args.model.as_deref(), Some("model.gguf"));
4631 assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
4632 assert_eq!(args.port, Some(9000));
4633 assert_eq!(args.threads, Some(4));
4634 assert_eq!(args.device, Some(OffloadDevice::Metal));
4635 assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
4636 assert_eq!(
4637 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4638 Some("[::1]:9000")
4639 );
4640 }
4641
4642 #[test]
4643 fn port_zero_survives_argument_parsing_as_a_real_request() {
4644 let argv = ["ferrox-server", "--port", "0"]
4649 .into_iter()
4650 .map(String::from)
4651 .collect();
4652 let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
4653 assert_eq!(args.port, Some(0));
4654 assert_eq!(
4655 cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
4656 Some("127.0.0.1:0")
4657 );
4658 }
4659
4660 #[test]
4661 fn stdin_close_exit_is_opt_in() {
4662 let args =
4665 ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
4666 assert!(!args.exit_on_stdin_close);
4667 let args = ServerArgs::try_parse_from(
4668 ["ferrox-server", "--exit-on-stdin-close"]
4669 .into_iter()
4670 .map(String::from),
4671 )
4672 .unwrap();
4673 assert!(args.exit_on_stdin_close);
4674 }
4675
4676 #[test]
4677 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4678 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4679 let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4680 let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
4681 assert_eq!(parsed.port, 51999);
4682 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4683 assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
4685 }
4686
4687 fn test_model() -> Model {
4688 let cfg = test_dense_fixture();
4692 Model::Gguf(GgufModel {
4693 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4694 tokenizer: Arc::new(ServerTokenizer::Byte),
4695 stop_tokens: StopTokens::default(),
4696 bos_id: None,
4697 is_synthetic: true,
4698 chat_template: chat_template::PromptTemplate::plain(),
4699 })
4700 }
4701
4702 fn greedy_params(max_tokens: usize) -> GenerationParams {
4703 GenerationParams {
4704 max_tokens,
4705 sampling: SamplingParams::default(),
4706 seed: 1,
4707 stop: Vec::new(),
4708 stop_token_ids: Vec::new(),
4709 json_object: false,
4710 grammar: None,
4711 cancel: None,
4712 ignore_eos: false,
4713 }
4714 }
4715
4716 fn test_model_full_byte_vocab() -> Model {
4720 let mut cfg = test_dense_fixture();
4721 cfg.vocab_size = 256;
4722 Model::Gguf(GgufModel {
4723 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4724 tokenizer: Arc::new(ServerTokenizer::Byte),
4725 stop_tokens: StopTokens::default(),
4726 bos_id: None,
4727 is_synthetic: true,
4728 chat_template: chat_template::PromptTemplate::plain(),
4729 })
4730 }
4731
4732 fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4736 AppState {
4737 paged_kv: None,
4738 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4739 id: None,
4740 model: Arc::new(model),
4741 batcher: None,
4742 ceiling: None,
4743 }))),
4744 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4745 tasks: Arc::new(tasks::TaskRegistry::new()),
4746 cancels: Arc::new(cancel::CancelRegistry::new()),
4747 stats: stats::Stats::new(),
4748 streams: resume::StreamRegistry::new(),
4749 model_dir: None,
4750 response_cache: Mutex::new(response_cache),
4751 kv_pool: None,
4752 prefix_cache: None,
4753 sessions: session::SessionStore::new(),
4754 requests_total: std::sync::atomic::AtomicU64::new(0),
4755 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4756 started_at: std::time::Instant::now(),
4757 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4758 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4759 mcp: None,
4760 continuous_batching_enabled: false,
4761 loading_model: Mutex::new(None),
4762 last_load_error: Mutex::new(None),
4763 serving: Mutex::new(crate::stats::ServingStats::default()),
4764 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4765 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4766 started_unix: unix_now(),
4767 }
4768 }
4769
4770 fn test_app() -> Router {
4779 test_app_with_state(Arc::new(test_state(
4780 test_model_full_byte_vocab(),
4781 ResponseCache::new(1000, Duration::from_secs(3600)),
4782 )))
4783 }
4784
4785 fn test_app_with_state(state: Arc<AppState>) -> Router {
4788 Router::new()
4789 .route(ferrox_api::routes::HEALTH, get(health))
4790 .route(ferrox_api::routes::V1_MODELS, get(list_models))
4791 .route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
4792 .route(
4793 &axum_path(ferrox_api::routes::V1_RESPONSE),
4794 get(responses::responses_get),
4795 )
4796 .route(
4797 &axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
4798 post(responses::responses_cancel),
4799 )
4800 .route(ferrox_api::routes::V1_STATS, get(serving_stats))
4801 .route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
4802 .route(
4803 ferrox_api::routes::V1_CACHE_STATUS,
4804 get(cache_admin::cache_status),
4805 )
4806 .route(
4807 ferrox_api::routes::V1_CACHE_REBUILD,
4808 post(cache_admin::cache_rebuild),
4809 )
4810 .route(
4811 ferrox_api::routes::ADMIN_PREPARE_STOP,
4812 post(cache_admin::prepare_stop),
4813 )
4814 .route("/v1/chat/completions", post(chat_completions))
4815 .route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
4816 .route(
4817 ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
4818 post(anthropic::count_tokens),
4819 )
4820 .route("/v1/tokenize", post(openai_extra::tokenize))
4821 .route("/v1/detokenize", post(openai_extra::detokenize))
4822 .route(ferrox_api::routes::TOKENIZE, post(openai_extra::tokenize))
4826 .route(
4827 ferrox_api::routes::DETOKENIZE,
4828 post(openai_extra::detokenize),
4829 )
4830 .route("/v1/embeddings", post(openai_extra::embeddings))
4831 .route("/v1/completions", post(openai_extra::completions))
4832 .route(ferrox_api::routes::COMPLETION, post(completion::completion))
4834 .route(
4835 ferrox_api::routes::COMPLETIONS,
4836 post(completion::completion),
4837 )
4838 .route(
4839 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
4840 post(admin::unload_model),
4841 )
4842 .route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
4843 .route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
4844 .route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
4845 .route(
4846 &axum_path(ferrox_api::routes::V1_STREAM),
4847 get(resume::resume),
4848 )
4849 .route(
4850 &axum_path(ferrox_api::routes::V1_STREAM_POLL),
4851 get(resume::poll),
4852 )
4853 .with_state(state)
4854 }
4855
4856 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4857 let mut cfg = test_dense_fixture();
4858 cfg.name = name;
4859 cfg.vocab_size = vocab_size;
4860 Model::Gguf(GgufModel {
4861 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4862 tokenizer: Arc::new(ServerTokenizer::Byte),
4863 stop_tokens: StopTokens::default(),
4864 bos_id: None,
4865 is_synthetic: true,
4866 chat_template: chat_template::PromptTemplate::plain(),
4867 })
4868 }
4869
4870 fn model_with_template(name: &'static str, source: &str) -> Model {
4874 let mut cfg = test_dense_fixture();
4875 cfg.name = name;
4876 cfg.vocab_size = 256;
4877 Model::Gguf(GgufModel {
4878 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4879 tokenizer: Arc::new(ServerTokenizer::Byte),
4880 stop_tokens: StopTokens::default(),
4881 bos_id: None,
4882 is_synthetic: true,
4883 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4884 Some(source),
4885 Some("qwen3"),
4886 false,
4887 true,
4888 None,
4889 None,
4890 ),
4891 })
4892 }
4893
4894 #[tokio::test]
4903 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4904 let strict = "{% if messages | length > 1 %}\
4907 {{ raise_exception('this template takes one turn') }}\
4908 {% endif %}{{ messages[0].content }}";
4909 let state = Arc::new(test_state(
4910 model_with_template("strict", strict),
4911 ResponseCache::new(4, Duration::from_secs(60)),
4912 ));
4913 let app = test_app_with_state(state);
4914
4915 let (status, body) = post_json_uri(
4916 &app,
4917 "/v1/chat/completions",
4918 serde_json::json!({
4919 "model": "strict",
4920 "stream": true,
4921 "messages": [
4922 {"role": "user", "content": "one"},
4923 {"role": "user", "content": "two"},
4924 ],
4925 }),
4926 )
4927 .await;
4928 assert_eq!(status, StatusCode::BAD_REQUEST);
4929 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4930 assert!(
4931 body["error"]["message"]
4932 .as_str()
4933 .unwrap()
4934 .contains("one turn"),
4935 "the template's own message must reach the caller: {body}"
4936 );
4937
4938 let (status, _) = post_json_uri(
4940 &app,
4941 "/v1/chat/completions",
4942 serde_json::json!({
4943 "model": "strict",
4944 "stream": true,
4945 "max_tokens": 1,
4946 "messages": [{"role": "user", "content": "one"}],
4947 }),
4948 )
4949 .await;
4950 assert_eq!(status, StatusCode::OK);
4951 }
4952
4953 #[tokio::test]
4955 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4956 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4957 {% if reasoning_effort %}\
4958 {% if reasoning_effort not in ['low','medium','high'] %}\
4959 {{ raise_exception('bad effort') }}\
4960 {% endif %}[{{ reasoning_effort }}]\
4961 {% endif %}{{ messages[0].content }}";
4962 let state = Arc::new(test_state(
4963 model_with_template("thinker", reasoning),
4964 ResponseCache::new(4, Duration::from_secs(60)),
4965 ));
4966 let app = test_app_with_state(state);
4967 let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4968 assert_eq!(status, StatusCode::OK);
4969 let entry = &models["data"][0];
4970 assert_eq!(
4971 entry["supported_reasoning_efforts"],
4972 serde_json::json!(["off", "low", "medium", "high"])
4973 );
4974 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4975 }
4976
4977 #[tokio::test]
4982 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4983 let app = test_app();
4984 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
4985 let entry = &models["data"][0];
4986 assert!(entry.get("supported_reasoning_efforts").is_none());
4987 assert!(entry.get("default_reasoning_effort").is_none());
4988 }
4989
4990 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
4991 Arc::new(ActiveModel {
4992 id: Some(name.to_string()),
4993 model: Arc::new(named_test_model(name, 256)),
4994 batcher: None,
4995 ceiling: None,
4996 })
4997 .tap_into(state)
4998 }
4999
5000 trait TapInto {
5002 fn tap_into(self, state: &AppState) -> Self;
5003 }
5004 impl TapInto for Arc<ActiveModel> {
5005 fn tap_into(self, state: &AppState) -> Self {
5006 state.swap_active(Some(Arc::clone(&self)));
5007 self
5008 }
5009 }
5010
5011 #[test]
5017 fn an_in_flight_request_keeps_the_model_it_started_on() {
5018 let state = test_state(
5019 named_test_model("model-a", 256),
5020 ResponseCache::new(4, Duration::from_secs(60)),
5021 );
5022
5023 let in_flight = state.active().expect("a model is loaded");
5026 assert_eq!(in_flight.model.name(), "model-a");
5027
5028 active_model(&state, "model-b");
5029
5030 assert_eq!(state.active().unwrap().model.name(), "model-b");
5032 assert_eq!(in_flight.model.name(), "model-a");
5034 let (_chunks, finish, _usage) = run_generation(
5035 &in_flight.model,
5036 "hi",
5037 &greedy_params(3),
5038 None,
5039 None,
5040 None,
5041 None,
5042 None,
5043 )
5044 .expect("the old model must still decode after being swapped out");
5045 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
5046 }
5047
5048 #[test]
5053 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5054 let state = test_state(
5055 named_test_model("model-a", 256),
5056 ResponseCache::new(4, Duration::from_secs(60)),
5057 );
5058 let in_flight = state.active().expect("a model is loaded");
5059 let weights = Arc::clone(&in_flight.model);
5060 assert!(Arc::strong_count(&weights) >= 2);
5061
5062 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5063 id: Some("model-b".to_string()),
5064 model: Arc::new(named_test_model("model-b", 256)),
5065 batcher: None,
5066 ceiling: None,
5067 })));
5068 drop(previous);
5069 assert!(Arc::strong_count(&weights) >= 2);
5071 drop(in_flight);
5072 assert_eq!(Arc::strong_count(&weights), 1);
5073 }
5074
5075 #[tokio::test]
5079 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5080 let state = Arc::new(test_state(
5081 named_test_model("model-a", 256),
5082 ResponseCache::new(4, Duration::from_secs(60)),
5083 ));
5084 let app = test_app_with_state(Arc::clone(&state));
5085
5086 let (status, body) = post_json_uri(
5087 &app,
5088 ferrox_api::routes::ADMIN_MODELS_UNLOAD,
5089 serde_json::json!({}),
5090 )
5091 .await;
5092 assert_eq!(status, StatusCode::OK);
5093 assert_eq!(body["ok"], true);
5094 assert!(body["active"].is_null());
5095 assert!(state.active().is_none());
5096
5097 let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5098 assert_eq!(status, StatusCode::OK);
5099 let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
5100 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5101
5102 let (status, body) = post_json_uri(
5103 &app,
5104 "/v1/chat/completions",
5105 serde_json::json!({
5106 "model": "x",
5107 "messages": [{"role": "user", "content": "hi"}]
5108 }),
5109 )
5110 .await;
5111 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5112 assert_eq!(body["error"]["type"], "model_not_loaded");
5113 }
5114
5115 #[tokio::test]
5119 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5120 let state = Arc::new(test_state(
5121 named_test_model("model-a", 256),
5122 ResponseCache::new(4, Duration::from_secs(60)),
5123 ));
5124 let app = test_app_with_state(Arc::clone(&state));
5125 state.swap_active(None);
5126
5127 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
5128 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5131 assert_eq!(body["state"], "unavailable");
5132 assert_eq!(body["reason"], "model_not_loaded");
5133 assert!(body["model"].is_null());
5134 let real_weights = body["capabilities"]
5135 .as_array()
5136 .unwrap()
5137 .iter()
5138 .find(|c| c["id"] == "real_weights")
5139 .cloned()
5140 .expect("real_weights is always reported");
5141 assert_eq!(real_weights["available"], false);
5142 assert_eq!(real_weights["reason"], "model_not_loaded");
5143 }
5144
5145 #[tokio::test]
5149 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5150 let app = test_app();
5151
5152 let (status, completion) = post_json_uri(
5153 &app,
5154 "/v1/chat/completions",
5155 serde_json::json!({
5156 "model": "x",
5157 "messages": [{"role": "user", "content": "hi"}],
5158 "max_tokens": 4
5159 }),
5160 )
5161 .await;
5162 assert_eq!(status, StatusCode::OK);
5163 let request_id = completion["request_id"].as_str().unwrap().to_string();
5164
5165 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5166 assert_eq!(status, StatusCode::OK);
5167 let recent = stats["recent"].as_array().unwrap();
5168 assert_eq!(recent.len(), 1);
5169 let row = &recent[0];
5170 assert_eq!(row["request_id"], request_id);
5171 assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
5172 assert_eq!(row["status"], 200);
5173 assert_eq!(row["stream"], false);
5174 assert!(row["duration_ms"].is_number());
5177 assert!(row["decode_ms"].is_number());
5178 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5179 assert_eq!(
5180 stats["tokens_prompt_total"].as_u64().unwrap(),
5181 row["prompt_tokens"].as_u64().unwrap()
5182 );
5183 }
5184
5185 #[tokio::test]
5189 async fn a_rejected_request_is_recorded_too() {
5190 let state = Arc::new(test_state(
5191 named_test_model("model-a", 256),
5192 ResponseCache::new(4, Duration::from_secs(60)),
5193 ));
5194 let app = test_app_with_state(Arc::clone(&state));
5195 state.swap_active(None);
5196
5197 let (status, _) = post_json_uri(
5198 &app,
5199 "/v1/chat/completions",
5200 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5201 )
5202 .await;
5203 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5204
5205 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5206 let recent = stats["recent"].as_array().unwrap();
5207 assert_eq!(recent.len(), 1);
5208 assert_eq!(recent[0]["status"], 503);
5209 assert_eq!(recent[0]["completion_tokens"], 0);
5210 assert!(recent[0]["decode_ms"].is_null());
5211 assert_eq!(stats["errors_total"], 1);
5212 }
5213
5214 async fn post_json_with_headers(
5218 app: &Router,
5219 uri: &str,
5220 body: serde_json::Value,
5221 headers: &[(&str, &str)],
5222 ) -> (StatusCode, serde_json::Value) {
5223 use http_body_util::BodyExt;
5224 use tower::ServiceExt;
5225
5226 let mut builder = axum::http::Request::builder()
5227 .method("POST")
5228 .uri(uri)
5229 .header("content-type", "application/json");
5230 for (name, value) in headers {
5231 builder = builder.header(*name, *value);
5232 }
5233 let response = app
5234 .clone()
5235 .oneshot(
5236 builder
5237 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5238 .unwrap(),
5239 )
5240 .await
5241 .unwrap();
5242 let status = response.status();
5243 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5244 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5245 (status, json)
5246 }
5247
5248 #[tokio::test]
5252 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5253 let app = test_app();
5254
5255 let (status, _) = post_json_uri(
5256 &app,
5257 ferrox_api::routes::V1_TOKENIZE,
5258 serde_json::json!({"prompt": "hello"}),
5259 )
5260 .await;
5261 assert_eq!(status, StatusCode::OK);
5262 let (status, _) = post_json_uri(
5263 &app,
5264 ferrox_api::routes::V1_DETOKENIZE,
5265 serde_json::json!({"tokens": [104, 105]}),
5266 )
5267 .await;
5268 assert_eq!(status, StatusCode::OK);
5269 let (status, _) = post_json_uri(
5270 &app,
5271 ferrox_api::routes::V1_EMBEDDINGS,
5272 serde_json::json!({"input": "hello"}),
5273 )
5274 .await;
5275 assert_eq!(status, StatusCode::OK);
5276
5277 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5278 let routes: Vec<&str> = stats["recent"]
5279 .as_array()
5280 .unwrap()
5281 .iter()
5282 .map(|row| row["route"].as_str().unwrap())
5283 .collect();
5284 for expected in [
5285 ferrox_api::routes::V1_TOKENIZE,
5286 ferrox_api::routes::V1_DETOKENIZE,
5287 ferrox_api::routes::V1_EMBEDDINGS,
5288 ] {
5289 assert!(
5290 routes.contains(&expected),
5291 "{expected} is missing: {routes:?}"
5292 );
5293 }
5294
5295 let row = |route: &str| {
5296 stats["recent"]
5297 .as_array()
5298 .unwrap()
5299 .iter()
5300 .find(|r| r["route"] == route)
5301 .cloned()
5302 .unwrap()
5303 };
5304 let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
5308 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5309 assert!(embed["decode_ms"].is_null());
5310 assert_eq!(embed["completion_tokens"], 0);
5311 assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5315 assert_eq!(
5316 stats["tokens_prompt_total"].as_u64().unwrap(),
5317 embed["prompt_tokens"].as_u64().unwrap(),
5318 "only the forward pass counted"
5319 );
5320 }
5321
5322 fn streaming_test_app() -> Router {
5327 let mut cfg = test_dense_fixture();
5328 cfg.vocab_size = 256;
5329 let model = Model::Gguf(GgufModel {
5330 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5331 tokenizer: Arc::new(ServerTokenizer::Byte),
5332 stop_tokens: StopTokens::default(),
5333 bos_id: None,
5334 is_synthetic: false,
5335 chat_template: chat_template::PromptTemplate::plain(),
5336 });
5337 test_app_with_state(Arc::new(test_state(
5338 model,
5339 ResponseCache::new(1000, Duration::from_secs(3600)),
5340 )))
5341 }
5342
5343 #[tokio::test]
5347 async fn the_native_completion_wire_is_not_the_openai_one() {
5348 let app = test_app();
5349
5350 let (status, native) = post_json_uri(
5351 &app,
5352 ferrox_api::routes::COMPLETION,
5353 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5354 )
5355 .await;
5356 assert_eq!(status, StatusCode::OK, "{native}");
5357 assert!(native["content"].is_string(), "{native}");
5358 assert_eq!(native["stop"], true);
5359 assert_eq!(native["stop_type"], "limit");
5360 assert_eq!(native["stopping_word"], "");
5361 assert_eq!(native["truncated"], false);
5362 assert_eq!(native["id_slot"], -1);
5363 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5364 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5365 assert!(
5366 native.get("choices").is_none(),
5367 "the native shape has no `choices`: {native}"
5368 );
5369
5370 let (status, openai) = post_json_uri(
5371 &app,
5372 ferrox_api::routes::V1_COMPLETIONS,
5373 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5374 )
5375 .await;
5376 assert_eq!(status, StatusCode::OK);
5377 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5378 assert!(
5379 openai.get("content").is_none(),
5380 "the OpenAI shape has no top-level `content`: {openai}"
5381 );
5382 }
5383
5384 #[tokio::test]
5388 async fn both_native_spellings_reach_the_same_handler() {
5389 let app = test_app();
5390 for route in [
5391 ferrox_api::routes::COMPLETION,
5392 ferrox_api::routes::COMPLETIONS,
5393 ] {
5394 let (status, body) = post_json_uri(
5395 &app,
5396 route,
5397 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5398 )
5399 .await;
5400 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5401 assert_eq!(body["stop"], true, "{route}");
5402 assert!(body["content"].is_string(), "{route}");
5403 }
5404
5405 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5408 let routes: Vec<&str> = stats["recent"]
5409 .as_array()
5410 .unwrap()
5411 .iter()
5412 .map(|row| row["route"].as_str().unwrap())
5413 .collect();
5414 assert!(
5415 routes.contains(&ferrox_api::routes::COMPLETION),
5416 "{routes:?}"
5417 );
5418 assert!(
5419 routes.contains(&ferrox_api::routes::COMPLETIONS),
5420 "{routes:?}"
5421 );
5422 }
5423
5424 #[tokio::test]
5430 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5431 let app = streaming_test_app();
5432 let raw = post_sse_raw_uri(
5433 &app,
5434 ferrox_api::routes::COMPLETION,
5435 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5436 )
5437 .await;
5438
5439 assert!(
5440 !raw.contains("[DONE]"),
5441 "llama.cpp's native stream has no sentinel: {raw}"
5442 );
5443 let frames: Vec<serde_json::Value> = raw
5444 .lines()
5445 .filter_map(|line| line.strip_prefix("data: "))
5446 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5447 .collect();
5448 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5449
5450 let (last, partials) = frames.split_last().unwrap();
5451 assert_eq!(last["stop"], true, "the last frame closes the stream");
5452 assert!(last["timings"].is_object(), "{last}");
5453 assert!(last["stop_type"].is_string(), "{last}");
5454 for partial in partials {
5455 assert_eq!(partial["stop"], false, "{partial}");
5456 assert!(partial["content"].is_string(), "{partial}");
5457 assert!(partial.get("timings").is_none(), "{partial}");
5461 assert!(partial.get("generation_settings").is_none(), "{partial}");
5462 }
5463 let streamed: String = partials
5466 .iter()
5467 .filter_map(|p| p["content"].as_str())
5468 .collect();
5469 assert_eq!(last["content"].as_str().unwrap(), streamed);
5470 }
5471
5472 #[tokio::test]
5477 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5478 let app = test_app();
5479 for body in [
5480 serde_json::json!({"prompt": "hi"}),
5481 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5482 ] {
5483 let (status, refusal) =
5484 post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
5485 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5486 assert!(
5487 refusal["error"]["message"]
5488 .as_str()
5489 .unwrap()
5490 .contains("n_predict"),
5491 "{refusal}"
5492 );
5493 }
5494 let (status, _) = post_json_uri(
5497 &app,
5498 ferrox_api::routes::COMPLETION,
5499 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5500 )
5501 .await;
5502 assert_eq!(status, StatusCode::OK);
5503 }
5504
5505 #[tokio::test]
5514 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5515 let app = streaming_test_app();
5516 let ask = |stop: serde_json::Value| {
5517 let app = app.clone();
5518 async move {
5519 post_json_uri(
5520 &app,
5521 ferrox_api::routes::COMPLETION,
5522 serde_json::json!({
5523 "prompt": "hi",
5524 "n_predict": 64,
5525 "ignore_eos": true,
5526 "stop": stop,
5527 }),
5528 )
5529 .await
5530 .1
5531 }
5532 };
5533
5534 let baseline = ask(serde_json::json!([])).await;
5535 assert_eq!(baseline["stop_type"], "limit");
5536 assert_eq!(baseline["stopping_word"], "");
5537 let text = baseline["content"].as_str().unwrap().to_string();
5538 let sentinel: String = text.chars().skip(1).take(2).collect();
5544 assert_eq!(
5545 sentinel.chars().count(),
5546 2,
5547 "the fixture must produce enough output to cut: {text:?}"
5548 );
5549 let cut = text.find(&sentinel).expect("it came out of this text");
5550
5551 let stopped = ask(serde_json::json!([sentinel])).await;
5552 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5553 assert_eq!(stopped["stopping_word"], sentinel);
5554 assert_eq!(
5555 stopped["content"].as_str().unwrap(),
5556 &text[..cut],
5557 "the answer must be cut at the sentinel, not run past it"
5558 );
5559 }
5560
5561 #[tokio::test]
5567 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5568 let app = test_app();
5569
5570 let (v1_status, v1) = post_json_uri(
5571 &app,
5572 ferrox_api::routes::V1_TOKENIZE,
5573 serde_json::json!({"prompt": "hello"}),
5574 )
5575 .await;
5576 let (alias_status, alias) = post_json_uri(
5577 &app,
5578 ferrox_api::routes::TOKENIZE,
5579 serde_json::json!({"content": "hello"}),
5580 )
5581 .await;
5582 assert_eq!(v1_status, StatusCode::OK);
5583 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5584 assert_eq!(v1["tokens"], alias["tokens"]);
5585 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5586
5587 let (status, both_ways) = post_json_uri(
5590 &app,
5591 ferrox_api::routes::TOKENIZE,
5592 serde_json::json!({"prompt": "hello"}),
5593 )
5594 .await;
5595 assert_eq!(status, StatusCode::OK);
5596 assert_eq!(both_ways["tokens"], v1["tokens"]);
5597 }
5598
5599 #[tokio::test]
5604 async fn detokenize_answers_under_both_dialects_keys() {
5605 let app = test_app();
5606 for route in [
5607 ferrox_api::routes::DETOKENIZE,
5608 ferrox_api::routes::V1_DETOKENIZE,
5609 ] {
5610 let (status, body) =
5611 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5612 assert_eq!(status, StatusCode::OK, "{route}");
5613 assert_eq!(body["text"], "hi", "{route}");
5614 assert_eq!(body["content"], body["text"], "{route}");
5615 }
5616 }
5617
5618 #[tokio::test]
5622 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5623 let app = test_app();
5624 let (status, _) = post_json_uri(
5625 &app,
5626 ferrox_api::routes::TOKENIZE,
5627 serde_json::json!({"content": "hello"}),
5628 )
5629 .await;
5630 assert_eq!(status, StatusCode::OK);
5631
5632 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5633 let routes: Vec<&str> = stats["recent"]
5634 .as_array()
5635 .unwrap()
5636 .iter()
5637 .map(|row| row["route"].as_str().unwrap())
5638 .collect();
5639 assert!(
5640 routes.contains(&ferrox_api::routes::TOKENIZE),
5641 "the alias must be its own row: {routes:?}"
5642 );
5643 assert!(
5644 !routes.contains(&ferrox_api::routes::V1_TOKENIZE),
5645 "nothing called /v1/tokenize: {routes:?}"
5646 );
5647 }
5648
5649 #[tokio::test]
5654 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5655 let mut cfg = test_dense_fixture();
5656 cfg.vocab_size = 256;
5657 let model = Model::Gguf(GgufModel {
5658 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5659 tokenizer: Arc::new(ServerTokenizer::Byte),
5660 stop_tokens: StopTokens::default(),
5661 bos_id: Some(7),
5662 is_synthetic: true,
5663 chat_template: chat_template::PromptTemplate::plain(),
5664 });
5665 let app = test_app_with_state(Arc::new(test_state(
5666 model,
5667 ResponseCache::new(1000, Duration::from_secs(3600)),
5668 )));
5669
5670 let (_, plain) = post_json_uri(
5671 &app,
5672 ferrox_api::routes::TOKENIZE,
5673 serde_json::json!({"content": "hi"}),
5674 )
5675 .await;
5676 let (_, special) = post_json_uri(
5677 &app,
5678 ferrox_api::routes::TOKENIZE,
5679 serde_json::json!({"content": "hi", "add_special": true}),
5680 )
5681 .await;
5682
5683 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5684 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5685 assert_eq!(special["count"], 3);
5686 }
5687
5688 #[tokio::test]
5691 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5692 let app = test_app();
5693 let (status, _) = post_json_uri(
5694 &app,
5695 ferrox_api::routes::V1_EMBEDDINGS,
5696 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5697 )
5698 .await;
5699 assert_eq!(status, StatusCode::BAD_REQUEST);
5700
5701 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5702 let recent = stats["recent"].as_array().unwrap();
5703 assert_eq!(recent.len(), 1);
5704 assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
5705 assert_eq!(recent[0]["status"], 400);
5706 assert_eq!(
5707 recent[0]["prompt_tokens"], 0,
5708 "a rejected call embedded nothing"
5709 );
5710 }
5711
5712 #[tokio::test]
5715 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5716 let app = test_app();
5717 let key = "sk-monitor-secret";
5718 let (status, _) = post_json_with_headers(
5719 &app,
5720 "/v1/chat/completions",
5721 serde_json::json!({
5722 "model": "x",
5723 "messages": [{"role": "user", "content": "hi"}],
5724 "max_tokens": 2
5725 }),
5726 &[
5727 ("authorization", &format!("Bearer {key}")),
5728 ("x-ferrox-client", "ferrox-studio"),
5729 ],
5730 )
5731 .await;
5732 assert_eq!(status, StatusCode::OK);
5733
5734 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5735 let row = stats["recent"].as_array().unwrap()[0].clone();
5736 let fingerprint = row["via_api_key"]
5737 .as_str()
5738 .expect("the row names the key that served it")
5739 .to_string();
5740 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5741 assert!(!fingerprint.contains(key));
5742 assert!(
5743 !serde_json::to_string(&stats).unwrap().contains(key),
5744 "the stats payload must not carry the key in any form"
5745 );
5746 assert_eq!(row["client"], "ferrox-studio");
5747 }
5748
5749 #[tokio::test]
5752 async fn different_keys_are_different_callers_and_no_key_is_null() {
5753 let app = test_app();
5754 let body = serde_json::json!({
5755 "model": "x",
5756 "messages": [{"role": "user", "content": "hi"}],
5757 "max_tokens": 1
5758 });
5759 for headers in [
5760 vec![("authorization", "Bearer key-one")],
5761 vec![("authorization", "Bearer key-two")],
5762 vec![],
5763 ] {
5764 let (status, _) =
5765 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5766 assert_eq!(status, StatusCode::OK);
5767 }
5768
5769 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5770 let recent = stats["recent"].as_array().unwrap();
5771 assert_eq!(recent.len(), 3);
5772 let one = recent[0]["via_api_key"].as_str().unwrap();
5773 let two = recent[1]["via_api_key"].as_str().unwrap();
5774 assert_ne!(one, two, "two keys must not collapse into one caller");
5775 assert!(
5776 recent[2]["via_api_key"].is_null(),
5777 "an unauthenticated call is null, not a fingerprint of nothing"
5778 );
5779 assert!(recent[2]["client"].is_null());
5780 }
5781
5782 #[tokio::test]
5787 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5788 let state = Arc::new(test_state(
5789 named_test_model("really-loaded", 256),
5790 ResponseCache::new(4, Duration::from_secs(60)),
5791 ));
5792 let app = test_app_with_state(Arc::clone(&state));
5793
5794 let (status, _) = post_json_uri(
5795 &app,
5796 "/v1/chat/completions",
5797 serde_json::json!({
5798 "model": "gpt-4-turbo-that-is-not-here",
5799 "messages": [{"role": "user", "content": "hi"}],
5800 "max_tokens": 2
5801 }),
5802 )
5803 .await;
5804 assert_eq!(status, StatusCode::OK);
5805
5806 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5807 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5808
5809 state.swap_active(None);
5812 let (status, _) = post_json_uri(
5813 &app,
5814 "/v1/chat/completions",
5815 serde_json::json!({
5816 "model": "gpt-4-turbo-that-is-not-here",
5817 "messages": [{"role": "user", "content": "hi"}]
5818 }),
5819 )
5820 .await;
5821 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5822 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5823 let recent = stats["recent"].as_array().unwrap();
5824 assert!(recent[recent.len() - 1]["model"].is_null());
5825 }
5826
5827 #[tokio::test]
5831 async fn a_streamed_row_names_the_model_it_decoded_against() {
5832 let state = Arc::new(test_state(
5833 named_test_model("model-before", 256),
5834 ResponseCache::new(4, Duration::from_secs(60)),
5835 ));
5836 let app = test_app_with_state(Arc::clone(&state));
5837 let _ = post_sse_raw(&app, resumable_request()).await;
5838 active_model(&state, "model-after");
5840
5841 let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5842 assert_eq!(stats["recent"][0]["model"], "model-before");
5843 }
5844
5845 #[tokio::test]
5848 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5849 let app = test_app();
5850 let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
5851 assert_eq!(status, StatusCode::OK);
5852 assert!(
5853 stats["queue_depth"].is_null(),
5854 "without continuous batching nothing queues, so there is nothing to measure"
5855 );
5856 assert!(stats["queue_rejected_total"].is_null());
5857 assert_eq!(
5858 stats["generating_now"], 0,
5859 "work in progress is measured and really is zero here"
5860 );
5861 }
5862
5863 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5868 post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
5869 }
5870
5871 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
5875 use http_body_util::BodyExt;
5876 use tower::ServiceExt;
5877
5878 let response = app
5879 .clone()
5880 .oneshot(
5881 axum::http::Request::builder()
5882 .method("POST")
5883 .uri(uri)
5884 .header("content-type", "application/json")
5885 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5886 .unwrap(),
5887 )
5888 .await
5889 .unwrap();
5890 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5891 String::from_utf8(bytes.to_vec()).unwrap()
5892 }
5893
5894 async fn get_json_with_headers(
5895 app: &Router,
5896 uri: &str,
5897 headers: &[(&str, &str)],
5898 ) -> (StatusCode, serde_json::Value) {
5899 use http_body_util::BodyExt;
5900 use tower::ServiceExt;
5901
5902 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
5903 for (name, value) in headers {
5904 builder = builder.header(*name, *value);
5905 }
5906 let response = app
5907 .clone()
5908 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
5909 .await
5910 .unwrap();
5911 let status = response.status();
5912 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5913 (
5914 status,
5915 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
5916 )
5917 }
5918
5919 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
5920 body.lines()
5921 .filter_map(|line| line.strip_prefix(field))
5922 .map(str::trim)
5923 .collect()
5924 }
5925
5926 fn resumable_request() -> serde_json::Value {
5927 serde_json::json!({
5928 "model": "m",
5929 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5930 "max_tokens": 4,
5931 "temperature": 0,
5932 "stream": true,
5933 "stream_resumable": true,
5934 })
5935 }
5936
5937 #[tokio::test]
5942 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
5943 let app = test_app();
5944 let body = post_sse_raw(&app, resumable_request()).await;
5945
5946 let request_id = body
5947 .lines()
5948 .find_map(|l| l.strip_prefix("data: "))
5949 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
5950 .and_then(|v| v["request_id"].as_str().map(str::to_string))
5951 .expect("the first chunk names the request");
5952
5953 let ids = sse_field(&body, "id:");
5954 let datas = sse_field(&body, "data:");
5955 assert_eq!(
5956 ids.len(),
5957 datas.len(),
5958 "every event carries an id, or a reconnect cannot name where it stopped"
5959 );
5960 for (i, id) in ids.iter().enumerate() {
5961 assert_eq!(*id, format!("{request_id}:{i}"));
5962 }
5963 let retries = sse_field(&body, "retry:");
5964 assert_eq!(
5965 retries.len(),
5966 1,
5967 "the reconnect delay is stated once, not on every event"
5968 );
5969 assert_eq!(retries[0], "1500");
5970 assert!(
5971 body.contains("data: [DONE]"),
5972 "the end of stream is still stated"
5973 );
5974 }
5975
5976 #[tokio::test]
5980 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
5981 let app = test_app();
5982 let mut request = resumable_request();
5983 request["stream_resumable"] = serde_json::json!(false);
5984 let body = post_sse_raw(&app, request).await;
5985 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
5986 assert!(
5987 sse_field(&body, "id:").is_empty(),
5988 "an id promises a replay this stream cannot serve"
5989 );
5990 assert!(sse_field(&body, "retry:").is_empty());
5991 }
5992
5993 #[tokio::test]
5997 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
5998 let app = test_app();
5999 let body = post_sse_raw(&app, resumable_request()).await;
6000 let request_id = sse_field(&body, "id:")[0]
6001 .rsplit_once(':')
6002 .unwrap()
6003 .0
6004 .to_string();
6005 let streamed: Vec<String> = sse_field(&body, "data:")
6006 .iter()
6007 .map(|d| d.to_string())
6008 .collect();
6009
6010 let (status, polled) = get_json(
6011 &app,
6012 &format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
6013 )
6014 .await;
6015 assert_eq!(status, StatusCode::OK);
6016 let events: Vec<String> = polled["events"]
6017 .as_array()
6018 .unwrap()
6019 .iter()
6020 .map(|e| e["data"].as_str().unwrap().to_string())
6021 .collect();
6022 assert_eq!(
6023 events, streamed,
6024 "the fallback must deliver the same answer, not a re-run of it"
6025 );
6026 assert_eq!(polled["request_id"], request_id);
6027 assert_eq!(
6028 polled["done"], false,
6029 "events were still being handed out, so the client must ask again"
6030 );
6031
6032 let next = polled["next_index"].as_u64().unwrap();
6035 let (_, drained) = get_json(
6036 &app,
6037 &format!(
6038 "{}?from={next}",
6039 ferrox_api::routes::v1_stream_poll(&request_id)
6040 ),
6041 )
6042 .await;
6043 assert_eq!(drained["done"], true);
6044 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6045 }
6046
6047 #[tokio::test]
6051 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6052 let app = test_app();
6053 let body = post_sse_raw(&app, resumable_request()).await;
6054 let ids = sse_field(&body, "id:");
6055 let datas: Vec<String> = sse_field(&body, "data:")
6056 .iter()
6057 .map(|d| d.to_string())
6058 .collect();
6059 assert!(
6060 ids.len() >= 3,
6061 "need a few events to resume into the middle"
6062 );
6063 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6064
6065 let (status, resumed) = get_json_with_headers(
6066 &app,
6067 &format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
6068 &[],
6069 )
6070 .await;
6071 assert_eq!(status, StatusCode::OK);
6072 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6073
6074 let (_, tail) = get_json(
6076 &app,
6077 &format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
6078 )
6079 .await;
6080 let tail_events: Vec<String> = tail["events"]
6081 .as_array()
6082 .unwrap()
6083 .iter()
6084 .map(|e| e["data"].as_str().unwrap().to_string())
6085 .collect();
6086 assert_eq!(tail_events, datas[2..].to_vec());
6087 }
6088
6089 #[tokio::test]
6092 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6093 use http_body_util::BodyExt;
6094 use tower::ServiceExt;
6095
6096 let app = test_app();
6097 let body = post_sse_raw(&app, resumable_request()).await;
6098 let ids = sse_field(&body, "id:");
6099 let datas: Vec<String> = sse_field(&body, "data:")
6100 .iter()
6101 .map(|d| d.to_string())
6102 .collect();
6103 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6104
6105 let response = app
6106 .clone()
6107 .oneshot(
6108 axum::http::Request::builder()
6109 .method("GET")
6110 .uri(ferrox_api::routes::v1_stream(&request_id))
6111 .header("last-event-id", format!("{request_id}:0"))
6112 .body(axum::body::Body::empty())
6113 .unwrap(),
6114 )
6115 .await
6116 .unwrap();
6117 assert_eq!(response.status(), StatusCode::OK);
6118 assert_eq!(
6119 response
6120 .headers()
6121 .get("x-accel-buffering")
6122 .and_then(|v| v.to_str().ok()),
6123 Some("no"),
6124 "the reconnect needs the same anti-buffering header as the stream"
6125 );
6126 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6127 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6128 assert_eq!(
6129 sse_field(&resumed, "data:")
6130 .iter()
6131 .map(|d| d.to_string())
6132 .collect::<Vec<_>>(),
6133 datas[1..].to_vec()
6134 );
6135 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6136 }
6137
6138 #[tokio::test]
6142 async fn a_last_event_id_from_another_stream_is_refused() {
6143 let app = test_app();
6144 let body = post_sse_raw(&app, resumable_request()).await;
6145 let request_id = sse_field(&body, "id:")[0]
6146 .rsplit_once(':')
6147 .unwrap()
6148 .0
6149 .to_string();
6150
6151 let (status, err) = get_json_with_headers(
6152 &app,
6153 &ferrox_api::routes::v1_stream(&request_id),
6154 &[("last-event-id", "chatcmpl-someone-else:3")],
6155 )
6156 .await;
6157 assert_eq!(status, StatusCode::BAD_REQUEST);
6158 assert_eq!(err["error"]["code"], "bad_last_event_id");
6159 }
6160
6161 #[tokio::test]
6165 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6166 let app = test_app();
6167 let mut request = resumable_request();
6168 request["stream_resumable"] = serde_json::json!(false);
6169 let body = post_sse_raw(&app, request).await;
6170 let request_id = body
6171 .lines()
6172 .find_map(|l| l.strip_prefix("data: "))
6173 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6174 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6175 .unwrap();
6176
6177 let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
6178 assert_eq!(status, StatusCode::NOT_FOUND);
6179 assert_eq!(err["error"]["code"], "stream_not_found");
6180 assert!(err["error"]["message"]
6181 .as_str()
6182 .unwrap()
6183 .contains("stream_resumable"));
6184 }
6185
6186 #[test]
6190 fn the_axum_stream_patterns_match_the_published_templates() {
6191 assert_eq!(
6192 axum_path(ferrox_api::routes::V1_STREAM),
6193 "/v1/stream/:request_id"
6194 );
6195 assert_eq!(
6196 axum_path(ferrox_api::routes::V1_STREAM_POLL),
6197 "/v1/stream/:request_id/poll"
6198 );
6199 assert_eq!(
6200 ferrox_api::routes::v1_stream("abc"),
6201 axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
6202 );
6203 }
6204
6205 #[test]
6215 fn no_published_template_reaches_the_router_with_its_braces() {
6216 for template in [
6217 ferrox_api::routes::V1_STREAM,
6218 ferrox_api::routes::V1_STREAM_POLL,
6219 ferrox_api::routes::V1_RESPONSE,
6220 ferrox_api::routes::V1_RESPONSE_CANCEL,
6221 ferrox_api::routes::ADMIN_TASK_CANCEL,
6222 ] {
6223 assert!(
6224 template.contains('{'),
6225 "{template} is in the template list but has no placeholder"
6226 );
6227 let mounted = axum_path(template);
6228 assert!(
6229 !mounted.contains('{') && !mounted.contains('}'),
6230 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6231 );
6232 assert!(
6233 mounted.contains(':'),
6234 "{template} lost its placeholder entirely and would match one path only"
6235 );
6236 }
6237 }
6238
6239 #[tokio::test]
6247 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6248 let app = test_app();
6249 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6250 assert_eq!(status, StatusCode::NOT_FOUND);
6251 assert!(
6252 !body.is_null(),
6253 "empty body means axum never matched the route, so the id was read as a literal segment"
6254 );
6255 }
6256
6257 #[tokio::test]
6260 async fn the_task_list_starts_empty_rather_than_absent() {
6261 let app = test_app();
6262 let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
6263 assert_eq!(status, StatusCode::OK);
6264 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6265 }
6266
6267 async fn post_json_uri(
6268 app: &Router,
6269 uri: &str,
6270 body: serde_json::Value,
6271 ) -> (StatusCode, serde_json::Value) {
6272 use http_body_util::BodyExt;
6273 use tower::ServiceExt;
6274
6275 let response = app
6276 .clone()
6277 .oneshot(
6278 axum::http::Request::builder()
6279 .method("POST")
6280 .uri(uri)
6281 .header("content-type", "application/json")
6282 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6283 .unwrap(),
6284 )
6285 .await
6286 .unwrap();
6287 let status = response.status();
6288 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6289 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6290 (status, json)
6291 }
6292
6293 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6294 post_json_uri(app, "/v1/chat/completions", body).await.1
6295 }
6296
6297 #[tokio::test]
6306 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6307 let app = test_app();
6308 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6309 assert_eq!(status, StatusCode::OK);
6310
6311 let memory = &body["memory"];
6312 if memory.is_null() {
6313 return;
6316 }
6317 assert!(
6318 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6319 "a read that produced a zero is a broken read, not an idle \
6320 engine: {memory}"
6321 );
6322 assert!(
6323 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6324 "the quantity must travel with the number: {memory}"
6325 );
6326 }
6327
6328 #[tokio::test]
6334 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6335 let app = test_app();
6336 let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
6337 assert_eq!(status, StatusCode::OK);
6338 for pool in ["kv_pages", "window_slots", "state_slots"] {
6339 assert!(
6340 body["pools"][pool].is_null(),
6341 "{pool} must be null rather than a zero row: {}",
6342 body["pools"]
6343 );
6344 }
6345 }
6346
6347 #[tokio::test]
6355 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6356 use http_body_util::BodyExt;
6357 use tower::ServiceExt;
6358
6359 let app = test_app();
6360 let send = |body: serde_json::Value| {
6361 let app = app.clone();
6362 async move {
6363 app.oneshot(
6364 axum::http::Request::builder()
6365 .method("POST")
6366 .uri(ferrox_api::routes::V1_MESSAGES)
6367 .header("content-type", "application/json")
6368 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6369 .unwrap(),
6370 )
6371 .await
6372 .unwrap()
6373 }
6374 };
6375
6376 let ok = send(serde_json::json!({
6377 "model": "test",
6378 "max_tokens": 1,
6379 "messages": [{"role": "user", "content": "hi"}],
6380 }))
6381 .await;
6382 assert_eq!(ok.status(), StatusCode::OK);
6383 let id = ok
6384 .headers()
6385 .get("request-id")
6386 .expect("a served message names its id")
6387 .to_str()
6388 .unwrap()
6389 .to_string();
6390 assert!(!id.is_empty());
6391
6392 let bad = send(serde_json::json!({"model": "test"})).await;
6395 assert!(bad.status().is_client_error());
6396 let other = bad.headers().get("request-id").expect("errors too");
6397 assert_ne!(other.to_str().unwrap(), id);
6398 let _ = bad.into_body().collect().await.unwrap();
6399 }
6400
6401 #[tokio::test]
6409 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6410 let state = Arc::new(test_state(
6411 test_model_full_byte_vocab(),
6412 ResponseCache::new(1000, Duration::from_secs(3600)),
6413 ));
6414 let app = test_app_with_state(Arc::clone(&state));
6415 let body = serde_json::json!({
6416 "model": "test",
6417 "messages": [{"role": "user", "content": "hi"}],
6418 "max_tokens": 1,
6419 });
6420
6421 state
6422 .maintenance
6423 .lock()
6424 .unwrap()
6425 .begin_rebuild()
6426 .expect("a fresh server is serving, so the rebuild starts");
6427 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6428 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6429 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6430
6431 state.maintenance.lock().unwrap().finish_rebuild(true);
6432 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6433 assert_eq!(
6434 status,
6435 StatusCode::OK,
6436 "the gate reopens; a rebuild is not a latch"
6437 );
6438 }
6439
6440 #[tokio::test]
6445 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6446 let app = test_app();
6447 let (status, body) = post_json_uri(
6448 &app,
6449 ferrox_api::routes::V1_CANCEL,
6450 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6451 )
6452 .await;
6453 assert_eq!(status, StatusCode::NOT_FOUND);
6454 assert_eq!(body["cancelled"], serde_json::json!(false));
6455 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6456 assert!(
6457 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6458 "the verdict must carry a human reason: {body}"
6459 );
6460 }
6461
6462 #[tokio::test]
6466 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6467 let state = Arc::new(test_state(
6468 test_model_full_byte_vocab(),
6469 ResponseCache::new(1000, Duration::from_secs(3600)),
6470 ));
6471 let app = test_app_with_state(Arc::clone(&state));
6472 let (token, _guard) = state.cancels.register("chatcmpl-live");
6473
6474 let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
6475 assert_eq!(status, StatusCode::OK);
6476 assert_eq!(before["generating_now"], serde_json::json!(1));
6477
6478 let (status, body) = post_json_uri(
6479 &app,
6480 ferrox_api::routes::V1_CANCEL,
6481 serde_json::json!({ "request_id": "chatcmpl-live" }),
6482 )
6483 .await;
6484 assert_eq!(status, StatusCode::OK);
6485 assert_eq!(body["cancelled"], serde_json::json!(true));
6486 assert!(
6487 token.is_cancelled(),
6488 "the endpoint answered ok without setting the flag the decode loop reads"
6489 );
6490 }
6491
6492 #[tokio::test]
6493 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6494 let app = test_app();
6495 let (status, tok) =
6496 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6497 assert_eq!(status, StatusCode::OK);
6498 let tokens = tok["tokens"].as_array().unwrap();
6499 assert_eq!(tok["count"], tokens.len());
6500 assert!(!tokens.is_empty());
6501
6502 let (status, detok) = post_json_uri(
6503 &app,
6504 "/v1/detokenize",
6505 serde_json::json!({ "tokens": tokens }),
6506 )
6507 .await;
6508 assert_eq!(status, StatusCode::OK);
6509 assert_eq!(detok["text"], "Hi");
6510
6511 let (status, emb) = post_json_uri(
6512 &app,
6513 "/v1/embeddings",
6514 serde_json::json!({
6515 "input": "Hi",
6516 "embedding_type": "mean"
6517 }),
6518 )
6519 .await;
6520 assert_eq!(status, StatusCode::OK);
6521 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6522 assert!(!vec.is_empty());
6523 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6524 }
6525
6526 #[tokio::test]
6531 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6532 use http_body_util::BodyExt;
6533 use tower::ServiceExt;
6534
6535 let fixture = concat!(
6536 "../ferrox-models/tests/fixtures/",
6537 "ferrox_real_moe_test.gguf"
6538 );
6539 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6540 let decoder = Decoder::from_gguf_with_expert_cache(
6541 &fixture,
6542 ferrox_models::config::test_moe_fixture(),
6543 Some(1024 * 1024),
6544 )
6545 .expect("MoE fixture must load store-backed");
6546
6547 let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
6551 .layers
6552 .iter()
6553 .map(|_| {
6554 ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
6555 })
6556 .collect();
6557 decoder.forward_token(1, 0, &mut caches);
6558
6559 let model = Model::Gguf(GgufModel {
6560 decoder: Arc::new(decoder),
6561 tokenizer: Arc::new(ServerTokenizer::Byte),
6562 stop_tokens: StopTokens::default(),
6563 bos_id: None,
6564 is_synthetic: false,
6565 chat_template: chat_template::PromptTemplate::plain(),
6566 });
6567 let state = Arc::new(test_state(
6568 model,
6569 ResponseCache::new(16, Duration::from_secs(60)),
6570 ));
6571 let app = Router::new()
6572 .route("/metrics", axum::routing::get(metrics))
6573 .route("/v1/chat/completions", post(chat_completions))
6574 .with_state(state);
6575
6576 let fetch_metrics = |app: Router| async move {
6577 let resp = app
6578 .oneshot(
6579 axum::http::Request::builder()
6580 .method("GET")
6581 .uri("/metrics")
6582 .body(axum::body::Body::empty())
6583 .unwrap(),
6584 )
6585 .await
6586 .unwrap();
6587 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6588 String::from_utf8(bytes.to_vec()).unwrap()
6589 };
6590
6591 let after = fetch_metrics(app.clone()).await;
6592 assert!(
6593 after.contains("ferrox_expert_cache_misses_total"),
6594 "streaming model must expose expert-cache metrics: {after}"
6595 );
6596 let misses: u64 = after
6597 .lines()
6598 .find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
6599 .and_then(|l| l.split_whitespace().nth(1))
6600 .and_then(|v| v.parse().ok())
6601 .expect("misses metric line must parse");
6602 assert!(
6603 misses > 0,
6604 "decode must have read experts through the store: {after}"
6605 );
6606 }
6607
6608 fn weather_tool() -> serde_json::Value {
6609 serde_json::json!({
6610 "type": "function",
6611 "function": {
6612 "name": "get_weather",
6613 "description": "Get the current weather for a location.",
6614 "parameters": {
6615 "type": "object",
6616 "properties": {"location": {"type": "string"}},
6617 "required": ["location"]
6618 }
6619 }
6620 })
6621 }
6622
6623 fn weather_tool_def() -> ToolDef {
6624 ToolDef {
6625 kind: "function".to_string(),
6626 function: ToolFunctionDef {
6627 name: "get_weather".to_string(),
6628 description: Some("Get the current weather for a location.".to_string()),
6629 parameters: Some(serde_json::json!({
6630 "type": "object",
6631 "properties": {"location": {"type": "string"}},
6632 "required": ["location"]
6633 })),
6634 },
6635 }
6636 }
6637
6638 #[test]
6639 fn tool_preamble_mentions_every_tool_name_and_description() {
6640 let preamble = tool_preamble(&[weather_tool_def()]);
6641 assert!(preamble.contains("get_weather"));
6642 assert!(preamble.contains("Get the current weather for a location."));
6643 assert!(preamble.contains("<tool_call>"));
6644 assert!(preamble.contains("</tool_call>"));
6645 }
6646
6647 #[test]
6648 fn a_real_marker_becomes_a_structured_tool_call() {
6649 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6650 let (message, finish) = build_response_message(
6651 text.to_string(),
6652 &[weather_tool_def()],
6653 output::OutputPosture::for_model("test-model"),
6654 "stop",
6655 );
6656 assert_eq!(finish, "tool_calls");
6657 let calls = message.tool_calls.expect("must carry a tool call");
6658 assert_eq!(calls[0].function.name, "get_weather");
6659 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
6660 assert_eq!(parsed["location"], "Paris");
6661 }
6662
6663 #[test]
6664 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
6665 let (message, finish) = build_response_message(
6666 "just an answer".to_string(),
6667 &[weather_tool_def()],
6668 output::OutputPosture::for_model("test-model"),
6669 "stop",
6670 );
6671 assert_eq!(finish, "stop");
6672 assert!(message.tool_calls.is_none());
6673 assert_eq!(message.content.as_deref(), Some("just an answer"));
6674 }
6675
6676 #[test]
6679 fn a_malformed_payload_is_not_a_tool_call() {
6680 let (message, finish) = build_response_message(
6681 "<tool_call>not valid json at all</tool_call>".to_string(),
6682 &[weather_tool_def()],
6683 output::OutputPosture::for_model("test-model"),
6684 "stop",
6685 );
6686 assert_eq!(finish, "stop");
6687 assert!(message.tool_calls.is_none());
6688 }
6689
6690 #[test]
6693 fn a_tool_that_was_never_offered_is_not_returned() {
6694 let (message, finish) = build_response_message(
6695 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
6696 &[weather_tool_def()],
6697 output::OutputPosture::for_model("test-model"),
6698 "stop",
6699 );
6700 assert_eq!(finish, "stop");
6701 assert!(message.tool_calls.is_none());
6702 }
6703
6704 #[test]
6706 fn marker_text_with_no_tools_offered_stays_content() {
6707 let (message, finish) = build_response_message(
6708 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
6709 &[],
6710 output::OutputPosture::for_model("test-model"),
6711 "stop",
6712 );
6713 assert_eq!(finish, "stop");
6714 assert!(message.tool_calls.is_none());
6715 assert!(message.content.is_some());
6716 }
6717
6718 #[test]
6722 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
6723 let opened = std::cell::Cell::new(0usize);
6724 let mut parser = crate::policy::parser::ToolCallParser::new(
6725 crate::policy::parser::ToolCallFormat::Qwen3Coder,
6726 vec![
6727 crate::policy::parser::tool_call::ToolSchema::with_parameters(
6728 "write_file",
6729 serde_json::json!({"type": "object", "properties": {
6730 "path": {"type": "string"},
6731 "contents": {"type": "string"}
6732 }}),
6733 ),
6734 ],
6735 );
6736 let wire = "<tool_call><function=write_file>\
6737 <parameter=path>\n/tmp/x\n</parameter>\
6738 <parameter=contents>\nhello world\n</parameter>\
6739 </function></tool_call>";
6740
6741 let mut deltas = Vec::new();
6742 let mut text = String::new();
6743 for piece in wire.as_bytes().chunks(7) {
6744 let chunk = String::from_utf8_lossy(piece).into_owned();
6745 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
6746 text.push_str(&more_text);
6747 deltas.extend(more);
6748 }
6749 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
6750 text.push_str(&more_text);
6751 deltas.extend(more);
6752
6753 assert_eq!(opened.get(), 1, "one call opened");
6754 assert!(text.is_empty(), "the markers are not content: {text:?}");
6755
6756 let first = &deltas[0];
6757 assert_eq!(first.index, 0);
6758 assert_eq!(first.id.as_deref(), Some("call_0"));
6759 assert_eq!(first.kind, Some("function"));
6760 assert_eq!(first.function.name.as_deref(), Some("write_file"));
6761
6762 let joined: String = deltas
6765 .iter()
6766 .filter_map(|d| d.function.arguments.clone())
6767 .collect();
6768 let parsed: serde_json::Value =
6769 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
6770 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
6771 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
6772 assert!(
6773 deltas.len() >= 3,
6774 "the arguments arrived in pieces, not whole: {}",
6775 deltas.len()
6776 );
6777 assert!(
6778 deltas[1..].iter().all(|d| d.function.name.is_none()),
6779 "only the opening delta carries identity"
6780 );
6781 }
6782
6783 #[test]
6785 fn text_around_a_streamed_call_is_still_content() {
6786 let opened = std::cell::Cell::new(0usize);
6787 let mut parser = crate::policy::parser::ToolCallParser::new(
6788 crate::policy::parser::ToolCallFormat::Qwen25,
6789 vec![crate::policy::parser::tool_call::ToolSchema::new(
6790 "get_weather",
6791 )],
6792 );
6793 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
6794 \"arguments\": {}}</tool_call> done";
6795 let mut text = String::new();
6796 for piece in wire.as_bytes().chunks(5) {
6797 let chunk = String::from_utf8_lossy(piece).into_owned();
6798 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
6799 text.push_str(&more);
6800 }
6801 let (more, _) = tool_call_deltas(parser.finish(), &opened);
6802 text.push_str(&more);
6803
6804 assert_eq!(opened.get(), 1);
6805 assert!(text.starts_with("let me check. "), "{text:?}");
6806 assert!(text.ends_with(" done"), "{text:?}");
6807 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
6808 }
6809
6810 #[test]
6813 fn a_reasoning_block_is_split_out_of_the_answer() {
6814 let (message, finish) = build_response_message(
6815 "<think>weighing it up</think>The answer is 4.".to_string(),
6816 &[],
6817 output::OutputPosture::for_model("Qwen3-8B"),
6818 "stop",
6819 );
6820 assert_eq!(finish, "stop");
6821 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
6822 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
6823 }
6824
6825 #[test]
6828 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
6829 let (message, _) = build_response_message(
6830 "Use the <think> tag like this.".to_string(),
6831 &[],
6832 output::OutputPosture::for_model("llama-3.1-8b"),
6833 "stop",
6834 );
6835 assert_eq!(
6836 message.content.as_deref(),
6837 Some("Use the <think> tag like this.")
6838 );
6839 assert!(message.reasoning_content.is_none());
6840 }
6841
6842 #[tokio::test]
6850 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
6851 let app = test_app();
6852 let body = serde_json::json!({
6853 "model": "m",
6854 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6855 "max_tokens": 4,
6856 "temperature": 0,
6857 });
6858 let resp = post_json(&app, body).await;
6859 let message = &resp["choices"][0]["message"];
6860 assert!(message["content"].is_string());
6861 assert!(message.get("tool_calls").is_none());
6862 assert_eq!(resp["choices"][0]["finish_reason"], "length");
6863 assert_eq!(resp["usage"]["completion_tokens"], 4);
6864 assert_eq!(
6865 resp["usage"]["total_tokens"],
6866 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
6867 );
6868 }
6869
6870 async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
6871 use http_body_util::BodyExt;
6872 use tower::ServiceExt;
6873
6874 let response = app
6875 .clone()
6876 .oneshot(
6877 axum::http::Request::builder()
6878 .method("GET")
6879 .uri(uri)
6880 .body(axum::body::Body::empty())
6881 .unwrap(),
6882 )
6883 .await
6884 .unwrap();
6885 let status = response.status();
6886 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6887 (status, serde_json::from_slice(&bytes).unwrap())
6888 }
6889
6890 #[tokio::test]
6891 async fn health_answers_a_capability_handshake_not_a_boolean() {
6892 let app = test_app();
6893 let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
6894 assert_eq!(status, StatusCode::OK);
6895
6896 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
6897 assert_eq!(health.state, ferrox_api::HealthState::Ready);
6898 assert!(health.pid > 0);
6899 assert!(health.server_time_unix_ms > 0);
6900 assert_eq!(health.last_request_age_seconds, None);
6903
6904 for id in [
6907 ferrox_api::health::capability::CPU,
6908 ferrox_api::health::capability::METAL,
6909 ferrox_api::health::capability::CUDA,
6910 ferrox_api::health::capability::REAL_WEIGHTS,
6911 ferrox_api::health::capability::CONTINUOUS_BATCHING,
6912 ] {
6913 let cap = health
6914 .capability(id)
6915 .unwrap_or_else(|| panic!("{id} missing"));
6916 assert!(!cap.reason.is_empty(), "{cap:?}");
6917 assert!(!cap.detail.is_empty(), "{cap:?}");
6918 }
6919 let weights = health
6923 .capability(ferrox_api::health::capability::REAL_WEIGHTS)
6924 .unwrap();
6925 assert!(!weights.available);
6926 assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
6927 assert!(health.model.as_ref().unwrap().synthetic_weights);
6928 }
6929
6930 #[tokio::test]
6931 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
6932 let app = test_app();
6933 let _ = post_json(
6934 &app,
6935 serde_json::json!({
6936 "model": "m",
6937 "messages": [{"role": "user", "content": "\u{1}"}],
6938 "max_tokens": 1,
6939 "temperature": 0,
6940 }),
6941 )
6942 .await;
6943 let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
6944 let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
6945 let age = health
6946 .last_request_age_seconds
6947 .expect("a served request is evidence of liveness");
6948 assert!((0.0..5.0).contains(&age), "implausible age {age}");
6949 }
6950
6951 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
6953 use http_body_util::BodyExt;
6954 use tower::ServiceExt;
6955
6956 let response = app
6957 .clone()
6958 .oneshot(
6959 axum::http::Request::builder()
6960 .method("POST")
6961 .uri("/v1/chat/completions")
6962 .header("content-type", "application/json")
6963 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6964 .unwrap(),
6965 )
6966 .await
6967 .unwrap();
6968 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6969 String::from_utf8(bytes.to_vec())
6970 .unwrap()
6971 .lines()
6972 .filter_map(|line| line.strip_prefix("data: "))
6973 .filter(|payload| *payload != "[DONE]")
6974 .map(|payload| serde_json::from_str(payload).unwrap())
6975 .collect()
6976 }
6977
6978 #[tokio::test]
6979 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
6980 let app = test_app();
6981 let chunks = post_sse_chunks(
6982 &app,
6983 serde_json::json!({
6984 "model": "m",
6985 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6986 "max_tokens": 4,
6987 "temperature": 0,
6988 "stream": true,
6989 }),
6990 )
6991 .await;
6992
6993 assert!(!chunks.is_empty());
6994 let request_id = chunks[0]["request_id"]
6995 .as_str()
6996 .expect("the first chunk names the request")
6997 .to_string();
6998 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
6999 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7002 assert!(
7003 chunk.get("request_id").is_none(),
7004 "chunk {i} repeats request_id"
7005 );
7006 }
7007 for chunk in &chunks {
7010 assert_eq!(chunk["id"], serde_json::json!(request_id));
7011 }
7012
7013 let other = post_sse_chunks(
7014 &app,
7015 serde_json::json!({
7016 "model": "m",
7017 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7018 "max_tokens": 4,
7019 "temperature": 0,
7020 "stream": true,
7021 }),
7022 )
7023 .await;
7024 assert_ne!(
7025 other[0]["request_id"].as_str().unwrap(),
7026 request_id,
7027 "two concurrent chats must not share an id"
7028 );
7029 }
7030
7031 #[tokio::test]
7032 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7033 let app = test_app();
7034 let resp = post_json(
7035 &app,
7036 serde_json::json!({
7037 "model": "m",
7038 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7039 "max_tokens": 2,
7040 "temperature": 0,
7041 }),
7042 )
7043 .await;
7044 assert_eq!(resp["id"], resp["request_id"]);
7045 assert!(resp["request_id"]
7046 .as_str()
7047 .unwrap()
7048 .starts_with("chatcmpl-"));
7049 }
7050
7051 #[tokio::test]
7054 async fn usage_carries_separate_prefill_and_decode_timings() {
7055 let app = test_app();
7056 let resp = post_json(
7057 &app,
7058 serde_json::json!({
7059 "model": "m",
7060 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7061 "max_tokens": 4,
7062 "temperature": 0,
7063 }),
7064 )
7065 .await;
7066 let usage = &resp["usage"];
7067 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7068 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7069 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7070 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7071 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7073 }
7074
7075 #[tokio::test]
7083 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7084 let app = test_app();
7085 let body = serde_json::json!({
7086 "model": "m",
7087 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7088 "max_tokens": 4,
7089 "temperature": 0,
7090 "tools": [weather_tool()],
7091 });
7092 let resp = post_json(&app, body).await;
7093 let message = &resp["choices"][0]["message"];
7094 assert!(
7095 message["content"].is_string(),
7096 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7097 );
7098 assert!(message.get("tool_calls").is_none());
7099 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7103 }
7104
7105 #[tokio::test]
7109 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7110 let app = test_app();
7111 let body = serde_json::json!({
7112 "model": "m",
7113 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7114 "max_tokens": 3,
7115 "temperature": 0,
7116 });
7117 let first = post_json(&app, body.clone()).await;
7118 assert_eq!(first["ferrox_cache"], "miss");
7119 let second = post_json(&app, body).await;
7120 assert_eq!(second["ferrox_cache"], "hit");
7121 assert_eq!(
7122 first["choices"][0]["message"]["content"],
7123 second["choices"][0]["message"]["content"]
7124 );
7125 assert_eq!(
7126 first["choices"][0]["finish_reason"],
7127 second["choices"][0]["finish_reason"]
7128 );
7129 assert_eq!(first["usage"], second["usage"]);
7130 assert_eq!(second["usage"]["completion_tokens"], 3);
7131 }
7132
7133 #[tokio::test]
7139 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7140 let session_app = test_app();
7141 let manual_app = test_app();
7142
7143 let turn1 = post_json(
7145 &session_app,
7146 serde_json::json!({
7147 "model": "m",
7148 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7149 "session_id": "s1",
7150 "max_tokens": 5,
7151 "temperature": 0,
7152 }),
7153 )
7154 .await;
7155 let reply1 = turn1["choices"][0]["message"]["content"]
7156 .as_str()
7157 .unwrap()
7158 .to_string();
7159
7160 let manual_turn1 = post_json(
7165 &manual_app,
7166 serde_json::json!({
7167 "model": "m",
7168 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7169 "max_tokens": 5,
7170 "temperature": 0,
7171 }),
7172 )
7173 .await;
7174 assert_eq!(
7175 manual_turn1["choices"][0]["message"]["content"]
7176 .as_str()
7177 .unwrap(),
7178 reply1
7179 );
7180
7181 let turn2 = post_json(
7183 &session_app,
7184 serde_json::json!({
7185 "model": "m",
7186 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7187 "session_id": "s1",
7188 "max_tokens": 5,
7189 "temperature": 0,
7190 }),
7191 )
7192 .await;
7193 let reply2 = turn2["choices"][0]["message"]["content"]
7194 .as_str()
7195 .unwrap()
7196 .to_string();
7197
7198 let manual_turn2 = post_json(
7202 &manual_app,
7203 serde_json::json!({
7204 "model": "m",
7205 "messages": [
7206 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7207 {"role": "assistant", "content": reply1},
7208 {"role": "user", "content": "\u{4}\u{5}"},
7209 ],
7210 "max_tokens": 5,
7211 "temperature": 0,
7212 }),
7213 )
7214 .await;
7215 assert_eq!(
7216 manual_turn2["choices"][0]["message"]["content"]
7217 .as_str()
7218 .unwrap(),
7219 reply2,
7220 "resuming a session must produce identical output to manually resending the full history"
7221 );
7222 }
7223
7224 #[test]
7227 fn lock_cache_recovers_from_a_poisoned_mutex() {
7228 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7229
7230 let poison_cache = Arc::clone(&cache);
7231 let _ = std::thread::spawn(move || {
7232 let _guard = poison_cache.lock().unwrap();
7233 panic!("simulated panic while holding the lock");
7234 })
7235 .join();
7236
7237 let recovered = lock_cache(&cache);
7239 assert_eq!(recovered.stats().entries, 0);
7240 }
7241
7242 #[test]
7243 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7244 let mut req_body = serde_json::json!({
7245 "model": "m",
7246 "messages": [{"role": "user", "content": "hi"}],
7247 });
7248 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7249 assert!(
7250 req.is_cacheable(),
7251 "default (temperature 0) must be cacheable"
7252 );
7253
7254 req_body["temperature"] = serde_json::json!(0.8);
7255 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7256 assert!(
7257 !req.is_cacheable(),
7258 "unseeded sampling must never be cacheable"
7259 );
7260
7261 req_body["seed"] = serde_json::json!(42);
7262 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7263 assert!(
7264 req.is_cacheable(),
7265 "sampling with an explicit seed is deterministic and must be cacheable"
7266 );
7267 }
7268
7269 const GRADED: &str = "{% if reasoning_effort %}\
7273 {% if reasoning_effort not in ['low','medium','high'] %}\
7274 {{ raise_exception('unsupported effort') }}\
7275 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7276 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7277
7278 fn graded_template() -> chat_template::PromptTemplate {
7279 chat_template::PromptTemplate::from_gguf_metadata(
7280 Some(GRADED),
7281 Some("qwen3"),
7282 false,
7283 true,
7284 None,
7285 None,
7286 )
7287 }
7288
7289 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7290 serde_json::from_value(value).expect("request")
7291 }
7292
7293 #[test]
7300 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7301 let req = chat_request(serde_json::json!({
7302 "model": "m",
7303 "messages": [{"role": "user", "content": "hi"}],
7304 "grammar": "root ::= \"a\"+",
7305 }));
7306 req.validate_supported_fields()
7307 .expect("a valid grammar is a valid request");
7308 let params = req
7309 .generation_params()
7310 .expect("a valid grammar compiles at params time too");
7311 assert!(
7312 params.grammar.is_some(),
7313 "the grammar was dropped between the wire and the sampler"
7314 );
7315 assert!(
7316 params.needs_vocab_logits(),
7317 "a grammar request that may fold lm_head into a GPU argmax is \
7318 a grammar request served unconstrained"
7319 );
7320
7321 let plain = chat_request(serde_json::json!({
7322 "model": "m",
7323 "messages": [{"role": "user", "content": "hi"}],
7324 }));
7325 assert!(plain.generation_params().unwrap().grammar.is_none());
7326 }
7327
7328 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7329 chat_request(serde_json::json!({
7330 "model": "m",
7331 "messages": [{"role": "user", "content": "weather in Rome?"}],
7332 "tools": [weather_tool()],
7333 "tool_choice": tool_choice,
7334 }))
7335 }
7336
7337 #[test]
7341 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7342 for choice in [
7343 serde_json::json!("required"),
7344 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7345 ] {
7346 let req = tool_request(choice.clone());
7347 req.validate_supported_fields()
7348 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7349 let params = req
7350 .generation_params_for_template(&graded_template(), "Qwen3-8B")
7351 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7352 let grammar = params
7353 .grammar
7354 .as_ref()
7355 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7356 assert!(
7357 grammar.is_awaiting_trigger(),
7358 "the model must be free to think before it calls"
7359 );
7360 assert!(
7361 !grammar.allows_eog(),
7362 "{choice} must not be able to end the turn without a call"
7363 );
7364 assert!(
7370 params.needs_vocab_logits(),
7371 "{choice} would let a backend return a token id instead of logits"
7372 );
7373 assert!(
7374 !generate::greedy_gpu_fold_allowed(¶ms),
7375 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7376 );
7377 }
7378 }
7379
7380 #[test]
7382 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7383 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7384 let req = tool_request(choice.clone());
7385 req.validate_supported_fields().expect("still supported");
7386 let params = match req.generation_params_for_template(&graded_template(), "Qwen3-8B") {
7387 Ok(p) => p,
7388 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7389 };
7390 assert!(
7391 params.grammar.is_none(),
7392 "{choice} does not force a call and must not be constrained"
7393 );
7394 }
7395 }
7396
7397 #[test]
7400 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7401 let req = chat_request(serde_json::json!({
7403 "model": "m",
7404 "messages": [{"role": "user", "content": "hi"}],
7405 "tool_choice": "required",
7406 }));
7407 let (status, _) = req
7408 .validate_supported_fields()
7409 .expect_err("nothing to call");
7410 assert_eq!(status, StatusCode::BAD_REQUEST);
7411
7412 let req =
7414 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7415 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7416 assert_eq!(status, StatusCode::BAD_REQUEST);
7417 assert_eq!(body["error"]["param"], "tool_choice");
7418
7419 let req = tool_request(serde_json::json!({"type": "function"}));
7421 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7422 assert_eq!(status, StatusCode::BAD_REQUEST);
7423
7424 let req = chat_request(serde_json::json!({
7426 "model": "m",
7427 "messages": [{"role": "user", "content": "hi"}],
7428 "tools": [weather_tool()],
7429 "tool_choice": "required",
7430 "grammar": "root ::= \"a\"+",
7431 }));
7432 let (status, _) = req
7433 .validate_supported_fields()
7434 .expect_err("a grammar and a forced call are two constraints");
7435 assert_eq!(status, StatusCode::BAD_REQUEST);
7436
7437 let req = tool_request(serde_json::json!("required"));
7440 let (status, Json(body)) =
7441 match req.generation_params_for_template(&graded_template(), "GLM-4.7") {
7442 Err(e) => e,
7443 Ok(_) => panic!("glm calls are not JSON behind a marker"),
7444 };
7445 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7446 assert!(
7447 body["error"]["message"].as_str().unwrap().contains("glm47"),
7448 "{body}"
7449 );
7450 }
7451
7452 #[test]
7455 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
7456 let req = chat_request(serde_json::json!({
7457 "model": "m",
7458 "messages": [{"role": "user", "content": "hi"}],
7459 "grammar": "root ::= \"a",
7460 }));
7461 let (status, Json(body)) = req
7462 .validate_supported_fields()
7463 .expect_err("this does not parse");
7464 assert_eq!(status, StatusCode::BAD_REQUEST);
7465 assert_eq!(body["error"]["param"], "grammar");
7466 assert!(req.generation_params().is_err(), "and again at params time");
7467 }
7468
7469 #[test]
7474 fn response_format_json_schema_names_the_missing_converter() {
7475 let req = chat_request(serde_json::json!({
7476 "model": "m",
7477 "messages": [{"role": "user", "content": "hi"}],
7478 "response_format": {
7479 "type": "json_schema",
7480 "json_schema": {"name": "x", "schema": {"type": "object"}},
7481 },
7482 }));
7483 let (status, Json(body)) = req
7484 .validate_supported_fields()
7485 .expect_err("not implemented yet");
7486 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7487 let message = body["error"]["message"].as_str().expect("a message");
7488 assert!(
7489 message.contains("JSON schema") && message.contains("grammar"),
7490 "the refusal must say what is missing: {message}"
7491 );
7492 }
7493
7494 #[test]
7497 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
7498 let req = chat_request(serde_json::json!({
7499 "model": "m",
7500 "messages": [{"role": "user", "content": "hi"}],
7501 }));
7502 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
7503 }
7504
7505 #[test]
7510 fn min_p_reaches_the_sampler_from_the_chat_wire() {
7511 let asked = chat_request(serde_json::json!({
7512 "model": "m",
7513 "messages": [{"role": "user", "content": "hi"}],
7514 "min_p": 0.07,
7515 }));
7516 assert_eq!(asked.sampling_params().min_p, 0.07);
7517
7518 let silent = chat_request(serde_json::json!({
7519 "model": "m",
7520 "messages": [{"role": "user", "content": "hi"}],
7521 }));
7522 assert_eq!(
7523 silent.sampling_params().min_p,
7524 0.0,
7525 "an unset min_p must be off, not llama.cpp's CLI default"
7526 );
7527 }
7528
7529 #[test]
7538 fn no_sampler_knob_is_missing_from_the_cache_key() {
7539 let base = serde_json::json!({
7540 "model": "m",
7541 "messages": [{"role": "user", "content": "hi"}],
7542 "seed": 1,
7543 });
7544 let baseline = chat_request(base.clone()).cache_key("prompt");
7545 for (knob, value) in [
7546 ("temperature", serde_json::json!(0.5)),
7547 ("top_p", serde_json::json!(0.9)),
7548 ("min_p", serde_json::json!(0.05)),
7549 ("top_k", serde_json::json!(40)),
7550 ("repetition_penalty", serde_json::json!(1.1)),
7551 ("presence_penalty", serde_json::json!(0.3)),
7552 ("frequency_penalty", serde_json::json!(0.3)),
7553 ] {
7554 let mut body = base.clone();
7555 body[knob] = value;
7556 assert_ne!(
7557 chat_request(body).cache_key("prompt"),
7558 baseline,
7559 "`{knob}` is not in the cache key: two requests differing \
7560 only in it would share one cached answer"
7561 );
7562 }
7563 }
7564
7565 #[test]
7569 fn an_explicit_zero_output_budget_is_a_client_error() {
7570 let req = chat_request(serde_json::json!({
7571 "model": "m",
7572 "messages": [{"role": "user", "content": "hi"}],
7573 "max_tokens": 0,
7574 }));
7575 let (status, body) = req.validate_supported_fields().expect_err("rejected");
7576 assert_eq!(status, StatusCode::BAD_REQUEST);
7577 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
7578 }
7579
7580 #[test]
7583 fn a_request_can_turn_thinking_off() {
7584 let template = graded_template();
7585 for body in [
7586 serde_json::json!({
7587 "model": "m",
7588 "messages": [{"role": "user", "content": "hi"}],
7589 "reasoning_effort": "none",
7590 }),
7591 serde_json::json!({
7592 "model": "m",
7593 "messages": [{"role": "user", "content": "hi"}],
7594 "thinking": {"type": "disabled"},
7595 }),
7596 ] {
7597 let kwargs = chat_request(body).resolve_template_kwargs(&template);
7598 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
7599 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
7600 assert!(!kwargs.contains_key("reasoning_effort"));
7603 }
7604 }
7605
7606 #[test]
7609 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
7610 let template = graded_template();
7611 let kwargs = chat_request(serde_json::json!({
7612 "model": "m",
7613 "messages": [{"role": "user", "content": "hi"}],
7614 "reasoning_effort": "high",
7615 "thinking": {"type": "disabled"},
7616 }))
7617 .resolve_template_kwargs(&template);
7618 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
7619 assert!(!kwargs.contains_key("reasoning_effort"));
7620 }
7621
7622 #[test]
7625 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
7626 let req = chat_request(serde_json::json!({
7627 "model": "m",
7628 "messages": [{"role": "user", "content": "hi"}],
7629 "thinking": {"type": "disable"},
7630 }));
7631 let (status, _) = req.validate_supported_fields().expect_err("rejected");
7632 assert_eq!(status, StatusCode::BAD_REQUEST);
7633 }
7634
7635 #[test]
7638 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
7639 let template = graded_template();
7640 let kwargs = chat_request(serde_json::json!({
7641 "model": "m",
7642 "messages": [{"role": "user", "content": "hi"}],
7643 "reasoning_effort": "none",
7644 "chat_template_kwargs": {"enable_thinking": true},
7645 }))
7646 .resolve_template_kwargs(&template);
7647 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
7648 }
7649
7650 #[test]
7654 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
7655 let template = graded_template();
7656 let req = chat_request(serde_json::json!({
7657 "model": "m",
7658 "messages": [{"role": "user", "content": "hi"}],
7659 "reasoning_effort": "minimal",
7660 }));
7661 let kwargs = req.resolve_template_kwargs(&template);
7662 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
7663 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
7664 assert!(prompt.starts_with("E:low|"), "{prompt}");
7665 }
7666
7667 #[test]
7671 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
7672 let template = graded_template();
7673 let req = chat_request(serde_json::json!({
7674 "model": "m",
7675 "messages": [{"role": "user", "content": "hi"}],
7676 "chat_template_kwargs": {"reasoning_effort": "none"},
7677 }));
7678 let kwargs = req.resolve_template_kwargs(&template);
7679 assert!(!kwargs.contains_key("reasoning_effort"));
7680 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
7681 assert_eq!(prompt, "hi");
7682 }
7683
7684 #[test]
7687 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
7688 let template = graded_template();
7689 let req = chat_request(serde_json::json!({
7690 "model": "m",
7691 "messages": [{"role": "user", "content": "hi"}],
7692 "reasoning_effort": "low",
7693 "chat_template_kwargs": {"reasoning_effort": "high"},
7694 }));
7695 assert_eq!(
7696 req.resolve_template_kwargs(&template)["reasoning_effort"],
7697 serde_json::json!("high")
7698 );
7699 }
7700
7701 #[test]
7705 fn offering_tools_turns_thinking_on_by_itself() {
7706 let template = graded_template();
7707 let quiet = chat_request(serde_json::json!({
7708 "model": "m",
7709 "messages": [{"role": "user", "content": "hi"}],
7710 }));
7711 assert!(!quiet
7712 .resolve_template_kwargs(&template)
7713 .contains_key("enable_thinking"));
7714
7715 let with_tools = chat_request(serde_json::json!({
7716 "model": "m",
7717 "messages": [{"role": "user", "content": "hi"}],
7718 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
7719 }));
7720 let kwargs = with_tools.resolve_template_kwargs(&template);
7721 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
7722 let prompt =
7723 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
7724 assert!(prompt.starts_with("THINK|"), "{prompt}");
7725 }
7726
7727 #[test]
7732 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
7733 let opener = chat_template::PromptTemplate::from_gguf_metadata(
7734 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
7735 Some("qwen3"),
7736 false,
7737 true,
7738 None,
7739 None,
7740 );
7741 let req = chat_request(serde_json::json!({
7742 "model": "m",
7743 "messages": [{"role": "user", "content": "hi"}],
7744 "chat_template_kwargs": {"enable_thinking": true},
7745 }));
7746 let kwargs = req.resolve_template_kwargs(&opener);
7747 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
7748 assert!(prompt.ends_with("<think>"), "{prompt}");
7749
7750 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
7753 let (message, _) = build_response_message(
7754 "weighing it up</think>Paris.".to_string(),
7755 &[],
7756 posture,
7757 "stop",
7758 );
7759 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7760 assert_eq!(message.content.as_deref(), Some("Paris."));
7761
7762 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
7765 let (message, _) = build_response_message(
7766 "weighing it up</think>Paris.".to_string(),
7767 &[],
7768 closed,
7769 "stop",
7770 );
7771 assert_eq!(message.reasoning_content, None);
7772 }
7773
7774 #[test]
7775 fn stop_param_accepts_both_single_string_and_array() {
7776 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
7777 "model": "m",
7778 "messages": [{"role": "user", "content": "hi"}],
7779 "stop": "END",
7780 }))
7781 .unwrap();
7782 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
7783
7784 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
7785 "model": "m",
7786 "messages": [{"role": "user", "content": "hi"}],
7787 "stop": ["A", "B"],
7788 }))
7789 .unwrap();
7790 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
7791 }
7792
7793 #[test]
7794 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
7795 let model = test_model();
7796 let result = run_generation(
7797 &model,
7798 "hello",
7799 &greedy_params(4),
7800 None,
7801 None,
7802 None,
7803 None,
7804 None,
7805 );
7806 assert!(matches!(
7807 result,
7808 Err(generate::DecodeError::TokenOutOfVocab { .. })
7809 ));
7810 }
7811
7812 #[test]
7816 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
7817 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7819 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
7820
7821 let holder_pool = Arc::clone(&pool);
7822 let holder = std::thread::spawn(move || {
7823 let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
7824 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
7826 drop(held);
7827 });
7828 std::thread::sleep(Duration::from_millis(15));
7829
7830 let config = generate::KvPoolConfig {
7831 pool,
7832 queue_wait: Duration::ZERO,
7833 };
7834 let result = run_generation(
7835 &model,
7836 &prompt,
7837 &greedy_params(4),
7838 Some(&config),
7839 None,
7840 None,
7841 None,
7842 None,
7843 );
7844 assert!(matches!(
7845 result,
7846 Err(generate::DecodeError::KvPoolExhausted)
7847 ));
7848
7849 let (status, _body) = decode_error_response(result.unwrap_err());
7850 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7851 holder.join().unwrap();
7852 }
7853
7854 #[test]
7862 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
7863 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7865 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
7867 let config = generate::KvPoolConfig {
7868 pool,
7869 queue_wait: Duration::ZERO,
7870 };
7871
7872 let result = run_generation(
7873 &model,
7874 &prompt,
7875 &greedy_params(4),
7876 Some(&config),
7877 None,
7878 None,
7879 None,
7880 None,
7881 );
7882 let err = result.expect_err("one block cannot hold two layers' caches");
7883 assert!(
7884 matches!(
7885 &err,
7886 generate::DecodeError::KvBudgetExceeded { binding, .. }
7887 if *binding == ferrox_models::Ceiling::DeviceMemory.code()
7888 ),
7889 "expected an immovable device-memory refusal, got {err:?}"
7890 );
7891 let (status, _body) = decode_error_response(err);
7892 assert_eq!(status, StatusCode::BAD_REQUEST);
7893 }
7894
7895 #[test]
7901 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
7902 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
7903 queued: 512,
7904 cap: 512,
7905 });
7906 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7907 assert_eq!(body["error"]["retry_after_seconds"], 1);
7908 let message = body["error"]["message"].as_str().expect("message");
7909 assert!(message.contains("512"), "{message}");
7910 }
7911
7912 #[test]
7913 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
7914 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
7915 token: 99,
7916 vocab_size: 32,
7917 });
7918 assert!(
7919 body["error"]["retry_after_seconds"].is_null(),
7920 "retrying a prompt this model cannot tokenize never helps"
7921 );
7922 }
7923
7924 #[test]
7925 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
7926 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
7927 token: 99,
7928 vocab_size: 32,
7929 });
7930 assert_eq!(status, StatusCode::BAD_REQUEST);
7931 }
7932
7933 #[test]
7934 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
7935 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7937 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
7938 let config = generate::KvPoolConfig {
7939 pool: pool.clone(),
7940 queue_wait: Duration::ZERO,
7941 };
7942
7943 let (_, finish, _usage) = run_generation(
7944 &model,
7945 &prompt,
7946 &greedy_params(4),
7947 Some(&config),
7948 None,
7949 None,
7950 None,
7951 None,
7952 )
7953 .unwrap();
7954 assert_eq!(finish, FinishReason::Length);
7955 assert_eq!(
7956 pool.lock().unwrap().free_blocks(),
7957 2,
7958 "a completed request must return its blocks to the pool"
7959 );
7960 }
7961
7962 #[tokio::test]
7967 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
7968 let model = Arc::new(test_model());
7969 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
7970
7971 let mut handles = Vec::new();
7972 for _ in 0..8 {
7973 let model = Arc::clone(&model);
7974 let prompt = prompt.clone();
7975 handles.push(tokio::task::spawn_blocking(move || {
7976 run_generation(
7977 &model,
7978 &prompt,
7979 &greedy_params(6),
7980 None,
7981 None,
7982 None,
7983 None,
7984 None,
7985 )
7986 .unwrap()
7987 }));
7988 }
7989
7990 let mut results = Vec::new();
7991 for h in handles {
7992 results.push(h.await.unwrap());
7993 }
7994 for r in &results[1..] {
7999 assert_eq!(r.0, results[0].0, "decoded chunks must match");
8000 assert_eq!(r.1, results[0].1, "finish reason must match");
8001 assert_eq!(
8002 r.2.prompt_tokens, results[0].2.prompt_tokens,
8003 "prompt token count must match"
8004 );
8005 assert_eq!(
8006 r.2.completion_tokens, results[0].2.completion_tokens,
8007 "completion token count must match"
8008 );
8009 }
8010 }
8011
8012 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8019 let mut header_entries = Vec::new();
8020 let mut data = Vec::new();
8021 for (name, shape, values) in tensors {
8022 let start = data.len();
8023 for v in values {
8024 data.extend_from_slice(&v.to_le_bytes());
8025 }
8026 let end = data.len();
8027 let shape_str = shape
8028 .iter()
8029 .map(|d| d.to_string())
8030 .collect::<Vec<_>>()
8031 .join(",");
8032 header_entries.push(format!(
8033 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8034 ));
8035 }
8036 let header = format!("{{{}}}", header_entries.join(","));
8037 let header_bytes = header.as_bytes();
8038 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8039 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8040 out.extend_from_slice(header_bytes);
8041 out.extend_from_slice(&data);
8042 out
8043 }
8044
8045 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8056 use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8057 use ferrox_models::kimi_loader::KimiRealHparams;
8058 use ferrox_moe::{GatingFunction, MoeLayerConfig};
8059
8060 let hidden_dim = 8;
8061 let kda_num_heads = 2;
8062 let kda_head_dim = 3;
8063 let kda_proj = kda_num_heads * kda_head_dim;
8064 let conv_kernel = 4;
8065 let dense_intermediate = 5;
8066 let vocab_size = 256;
8070 let mla_num_heads = 1;
8071 let mla_q_lora_rank = 2;
8072 let mla_kv_lora_rank = 2;
8073 let mla_qk_nope_head_dim = 2;
8074 let mla_qk_rope_head_dim = 2;
8075 let mla_v_head_dim = 2;
8076
8077 let model_cfg = ferrox_models::ModelConfig {
8078 name: "synthetic-kimi-server-test",
8079 n_layers: 1,
8080 hidden_dim,
8081 n_heads: 1,
8082 n_kv_heads: 1,
8083 head_dim: 4,
8084 vocab_size,
8085 rope_theta: 10000.0,
8086 rms_norm_eps: 1e-5,
8087 sliding_window: None,
8088 moe: MoeLayerConfig {
8089 expert_weights_scale: 1.0,
8090 n_experts: 1,
8091 n_experts_active: 1,
8092 n_shared_experts: 0,
8093 hidden_dim,
8094 expert_ffn_dim: 4,
8095 gating: GatingFunction::Sigmoid,
8096 norm_topk_prob: true,
8097 expert_group_count: None,
8098 expert_group_used_count: None,
8099 },
8100 n_dense_leading_layers: 1,
8105 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8106 kda_layers: vec![1],
8107 full_attn_layers: vec![],
8108 mla: MlaConfig {
8109 num_heads: mla_num_heads,
8110 q_lora_rank: mla_q_lora_rank,
8111 kv_lora_rank: mla_kv_lora_rank,
8112 qk_nope_head_dim: mla_qk_nope_head_dim,
8113 qk_rope_head_dim: mla_qk_rope_head_dim,
8114 v_head_dim: mla_v_head_dim,
8115 use_output_gate: true,
8116 rope: None,
8117 },
8118 kda: KdaConfig {
8119 num_heads: kda_num_heads,
8120 head_dim: kda_head_dim,
8121 short_conv_kernel_size: conv_kernel,
8122 gate_lower_bound: -5.0,
8123 use_full_rank_gate: true,
8124 },
8125 }),
8126 rope_freqs: None,
8127 rope_attn_factor: 1.0,
8128 rope_dim: None,
8129 rope_freqs_long: None,
8130 rope_freqs_short: None,
8131 rope_orig_ctx: None,
8132 rope_layout: ferrox_models::config::RopeLayout::Neox,
8133 qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
8134 swa_pattern: None,
8135 swa_dense_first: false,
8136 attn_logit_softcap: None,
8137 final_logit_softcap: None,
8138 embedding_scale: None,
8139 attention_scale: None,
8140 rope_theta_swa: None,
8141 ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
8142 best_effort_fields: &["synthetic test config, not a real preset"],
8143 };
8144 let hp = KimiRealHparams {
8145 hidden_dim,
8146 kda_num_heads,
8147 kda_head_dim,
8148 mla_num_heads,
8149 mla_q_lora_rank,
8150 mla_kv_lora_rank,
8151 mla_qk_nope_head_dim,
8152 mla_qk_rope_head_dim,
8153 mla_v_head_dim,
8154 dense_intermediate_dim: dense_intermediate,
8155 moe_hidden_dim: hidden_dim,
8156 moe_intermediate_dim: 4,
8157 n_experts: 1,
8158 num_shared_experts: 0,
8159 };
8160
8161 let prefix = "language_model.model.layers.0";
8165 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8166 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8167 name: String,
8168 shape: Vec<usize>,
8169 n: usize| {
8170 tensors.push((name, shape, vec![0.01f32; n]));
8171 };
8172 push(
8173 &mut tensors,
8174 format!("{prefix}.input_layernorm.weight"),
8175 vec![hidden_dim],
8176 hidden_dim,
8177 );
8178 push(
8179 &mut tensors,
8180 format!("{prefix}.post_attention_layernorm.weight"),
8181 vec![hidden_dim],
8182 hidden_dim,
8183 );
8184 push(
8185 &mut tensors,
8186 format!("{prefix}.self_attention_res_norm.weight"),
8187 vec![hidden_dim],
8188 hidden_dim,
8189 );
8190 push(
8191 &mut tensors,
8192 format!("{prefix}.self_attention_res_proj.weight"),
8193 vec![1, hidden_dim],
8194 hidden_dim,
8195 );
8196 push(
8197 &mut tensors,
8198 format!("{prefix}.mlp_res_norm.weight"),
8199 vec![hidden_dim],
8200 hidden_dim,
8201 );
8202 push(
8203 &mut tensors,
8204 format!("{prefix}.mlp_res_proj.weight"),
8205 vec![1, hidden_dim],
8206 hidden_dim,
8207 );
8208 push(
8209 &mut tensors,
8210 format!("{prefix}.self_attn.q_proj.weight"),
8211 vec![kda_proj, hidden_dim],
8212 kda_proj * hidden_dim,
8213 );
8214 push(
8215 &mut tensors,
8216 format!("{prefix}.self_attn.k_proj.weight"),
8217 vec![kda_proj, hidden_dim],
8218 kda_proj * hidden_dim,
8219 );
8220 push(
8221 &mut tensors,
8222 format!("{prefix}.self_attn.v_proj.weight"),
8223 vec![kda_proj, hidden_dim],
8224 kda_proj * hidden_dim,
8225 );
8226 push(
8227 &mut tensors,
8228 format!("{prefix}.self_attn.q_conv1d.weight"),
8229 vec![kda_proj, 1, conv_kernel],
8230 kda_proj * conv_kernel,
8231 );
8232 push(
8233 &mut tensors,
8234 format!("{prefix}.self_attn.k_conv1d.weight"),
8235 vec![kda_proj, 1, conv_kernel],
8236 kda_proj * conv_kernel,
8237 );
8238 push(
8239 &mut tensors,
8240 format!("{prefix}.self_attn.v_conv1d.weight"),
8241 vec![kda_proj, 1, conv_kernel],
8242 kda_proj * conv_kernel,
8243 );
8244 push(
8245 &mut tensors,
8246 format!("{prefix}.self_attn.A_log"),
8247 vec![kda_num_heads],
8248 kda_num_heads,
8249 );
8250 push(
8251 &mut tensors,
8252 format!("{prefix}.self_attn.f_a_proj.weight"),
8253 vec![kda_head_dim, hidden_dim],
8254 kda_head_dim * hidden_dim,
8255 );
8256 push(
8257 &mut tensors,
8258 format!("{prefix}.self_attn.f_b_proj.weight"),
8259 vec![kda_proj, kda_head_dim],
8260 kda_proj * kda_head_dim,
8261 );
8262 push(
8263 &mut tensors,
8264 format!("{prefix}.self_attn.dt_bias"),
8265 vec![kda_proj],
8266 kda_proj,
8267 );
8268 push(
8269 &mut tensors,
8270 format!("{prefix}.self_attn.b_proj.weight"),
8271 vec![kda_num_heads, hidden_dim],
8272 kda_num_heads * hidden_dim,
8273 );
8274 push(
8275 &mut tensors,
8276 format!("{prefix}.self_attn.g_proj.weight"),
8277 vec![kda_proj, hidden_dim],
8278 kda_proj * hidden_dim,
8279 );
8280 push(
8281 &mut tensors,
8282 format!("{prefix}.self_attn.o_norm.weight"),
8283 vec![kda_head_dim],
8284 kda_head_dim,
8285 );
8286 push(
8287 &mut tensors,
8288 format!("{prefix}.self_attn.o_proj.weight"),
8289 vec![hidden_dim, kda_proj],
8290 hidden_dim * kda_proj,
8291 );
8292 push(
8293 &mut tensors,
8294 format!("{prefix}.mlp.gate_proj.weight"),
8295 vec![dense_intermediate, hidden_dim],
8296 dense_intermediate * hidden_dim,
8297 );
8298 push(
8299 &mut tensors,
8300 format!("{prefix}.mlp.up_proj.weight"),
8301 vec![dense_intermediate, hidden_dim],
8302 dense_intermediate * hidden_dim,
8303 );
8304 push(
8305 &mut tensors,
8306 format!("{prefix}.mlp.down_proj.weight"),
8307 vec![hidden_dim, dense_intermediate],
8308 hidden_dim * dense_intermediate,
8309 );
8310 push(
8311 &mut tensors,
8312 "language_model.model.embed_tokens.weight".to_string(),
8313 vec![vocab_size, hidden_dim],
8314 vocab_size * hidden_dim,
8315 );
8316 push(
8317 &mut tensors,
8318 "language_model.lm_head.weight".to_string(),
8319 vec![vocab_size, hidden_dim],
8320 vocab_size * hidden_dim,
8321 );
8322 push(
8323 &mut tensors,
8324 "language_model.model.norm.weight".to_string(),
8325 vec![hidden_dim],
8326 hidden_dim,
8327 );
8328 push(
8329 &mut tensors,
8330 "language_model.model.output_attn_res_norm.weight".to_string(),
8331 vec![hidden_dim],
8332 hidden_dim,
8333 );
8334 push(
8335 &mut tensors,
8336 "language_model.model.output_attn_res_proj.weight".to_string(),
8337 vec![1, hidden_dim],
8338 hidden_dim,
8339 );
8340
8341 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
8354 let dir = std::env::temp_dir().join(format!(
8355 "ferrox_server_kimi_e2e_test_{}_{}",
8356 std::process::id(),
8357 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8358 ));
8359 std::fs::create_dir_all(&dir).unwrap();
8360 let shard_bytes = write_safetensors_shard(&tensors);
8361 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
8362 let map_entries: Vec<String> = tensors
8363 .iter()
8364 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
8365 .collect();
8366 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
8367 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
8368
8369 use base64::Engine;
8373 let vocab_lines: Vec<String> = (0..vocab_size as u32)
8374 .map(|b| {
8375 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
8376 format!("{b64} {b}")
8377 })
8378 .collect();
8379 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
8380
8381 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
8382 .expect("must load the synthetic Kimi checkpoint end to end");
8383 std::fs::remove_dir_all(&dir).ok();
8384 loaded
8385 }
8386
8387 #[test]
8395 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
8396 let loaded = build_synthetic_kimi_loaded();
8397 let state = build_app_state(
8398 model::LoadedModel::Kimi(loaded),
8399 None,
8400 None,
8401 None,
8402 false,
8403 None,
8404 Arc::new(health::Detection::ready(health::probe_backends())),
8405 );
8406 let active = state.active().expect("a freshly built state has a model");
8407 assert_eq!(active.model.tokenizer_kind(), "kimi-tiktoken-bpe");
8408 assert!(!active.model.is_synthetic());
8409
8410 let (_chunks, finish, _usage) = run_generation(
8411 &active.model,
8412 "hi",
8413 &greedy_params(5),
8414 None,
8415 None,
8416 None,
8417 None,
8418 None,
8419 )
8420 .expect("a real Kimi checkpoint must generate without error");
8421 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
8422 }
8423
8424 #[test]
8437 fn a_grammar_constrains_the_engine_decode_path() {
8438 let loaded = build_synthetic_kimi_loaded();
8439 let state = build_app_state(
8440 model::LoadedModel::Kimi(loaded),
8441 None,
8442 None,
8443 None,
8444 false,
8445 None,
8446 Arc::new(health::Detection::ready(health::probe_backends())),
8447 );
8448 let active = state.active().expect("a freshly built state has a model");
8449
8450 let run = |grammar: Option<&str>| {
8451 let mut params = greedy_params(6);
8452 params.grammar = grammar.map(|src| {
8453 Arc::new(
8454 ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
8455 .expect("test grammar parses"),
8456 )
8457 });
8458 run_generation(&active.model, "hi", ¶ms, None, None, None, None, None)
8459 };
8460
8461 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
8462 let unconstrained = chunks.concat();
8463 assert!(
8464 unconstrained.chars().any(|c| c != 'a'),
8465 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
8466 constrained run below would prove nothing"
8467 );
8468
8469 let (chunks, finish, _) =
8470 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
8471 let constrained = chunks.concat();
8472 assert!(
8473 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
8474 "the engine decode path served text its grammar forbids ({constrained:?}): \
8475 the constraint was dropped between `generate_engine` and the sampler"
8476 );
8477 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
8478 }
8479
8480 #[test]
8490 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
8491 let loaded = build_synthetic_kimi_loaded();
8492 let state = build_app_state(
8493 model::LoadedModel::Kimi(loaded),
8494 None,
8495 None,
8496 None,
8497 false,
8498 None,
8499 Arc::new(health::Detection::ready(health::probe_backends())),
8500 );
8501
8502 let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
8503 let kv_pool_config = generate::KvPoolConfig {
8504 pool: pool.clone(),
8505 queue_wait: Duration::ZERO,
8506 };
8507 let pc = Mutex::new(PrefixCache::new(4));
8508
8509 run_generation(
8510 &state
8511 .active()
8512 .expect("a freshly built state has a model")
8513 .model,
8514 "hi",
8515 &greedy_params(5),
8516 Some(&kv_pool_config),
8517 None,
8518 Some(&pc),
8519 None,
8520 None,
8521 )
8522 .expect("a real Kimi checkpoint must generate without error");
8523
8524 assert_eq!(
8525 pool.lock().unwrap().free_blocks(),
8526 4,
8527 "the KV pool must be completely untouched by a Kimi request"
8528 );
8529 let stats = pc.lock().unwrap().stats();
8530 assert_eq!(
8531 stats.hits + stats.misses,
8532 0,
8533 "the prefix cache must never be consulted for a Kimi request"
8534 );
8535 }
8536}