1mod admin;
36mod anthropic;
37mod attribution;
38mod budget;
39mod cache_admin;
40mod cancel;
41mod chat_params;
42mod chat_template;
43mod cli;
44mod completion;
45mod continuation;
46mod conversations;
47mod decode_task;
48mod embeddings;
49mod generate;
50mod grammar_request;
51mod health;
52mod journal;
53mod json_mode;
54mod limits;
55mod loaded;
56mod lora;
57mod mcp;
58mod model;
59mod openai_extra;
60mod output;
61mod policy;
62mod prefill_batch;
63mod reasoning_budget;
64mod reasoning_tokens;
65mod rerank;
66mod response_cache;
67pub(crate) mod responses;
68mod resume;
69mod sample_step;
70mod sampling_knobs;
71mod security;
72mod serving;
73mod session;
74mod slots;
75mod sse;
76mod stats;
77mod stop;
78mod stream_events;
79mod tasks;
80mod tool_grammar;
81mod unsupported_sampling;
82mod utf8_stream;
83
84use std::cell::RefCell;
85use std::convert::Infallible;
86use std::net::SocketAddr;
87use std::path::PathBuf;
88use std::rc::Rc;
89use std::sync::{Arc, Mutex, MutexGuard};
90use std::time::Duration;
91
92use axum::{
93 extract::State,
94 http::StatusCode,
95 response::sse::{Event, Sse},
96 response::{IntoResponse, Response},
97 routing::{get, post},
98 Json, Router,
99};
100use serde::{Deserialize, Serialize};
101
102use cli::{apply_cli_overrides, print_available_devices};
103pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
104
105use frink_core::cache::KvBlockPool;
106use frink_models::kimi_tokenizer::KimiTokenizer;
107use frink_models::sampling::SamplingParams;
108use frink_models::tokenizer::{SpecialTokens, StopTokens};
109use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
110use generate::{FinishReason, GenerationParams};
111pub(crate) use loaded::{ActiveModel, Loaded};
112use model::ServerTokenizer;
113use rerank::encoder_endpoints;
114use response_cache::ResponseCache;
115use sampling_knobs::SamplingKnobs;
116
117#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
126 Gguf(GgufModel),
127 Kimi(KimiModel),
128 Mla(MlaModel),
129 Gemma4(Gemma4Model),
130 Glm52(Glm52Model),
131}
132
133pub(crate) struct GgufModel {
134 decoder: Arc<Decoder>,
135 tokenizer: Arc<ServerTokenizer>,
136 stop_tokens: StopTokens,
137 bos_id: Option<usize>,
138 is_synthetic: bool,
139 chat_template: chat_template::PromptTemplate,
140}
141
142pub(crate) struct KimiModel {
143 engine: KimiEngine,
144 tokenizer: KimiTokenizer,
145 stop_tokens: StopTokens,
146 chat_template: chat_template::PromptTemplate,
147}
148
149pub(crate) struct MlaModel {
150 engine: MlaEngine,
151 tokenizer: ServerTokenizer,
152 stop_tokens: StopTokens,
153 bos_id: Option<usize>,
154 name: String,
155 chat_template: chat_template::PromptTemplate,
156}
157
158pub(crate) struct Gemma4Model {
159 engine: Gemma4Engine,
160 tokenizer: ServerTokenizer,
161 stop_tokens: StopTokens,
162 bos_id: Option<usize>,
163 name: String,
164 chat_template: chat_template::PromptTemplate,
165}
166
167pub(crate) struct Glm52Model {
168 engine: frink_models::Glm52Engine,
169 tokenizer: ServerTokenizer,
170 stop_tokens: StopTokens,
171 bos_id: Option<usize>,
172 name: String,
173 chat_template: chat_template::PromptTemplate,
174}
175
176impl Model {
177 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
178 match self {
179 Model::Gguf(m) => m.chat_template.clone(),
180 Model::Kimi(m) => m.chat_template.clone(),
181 Model::Mla(m) => m.chat_template.clone(),
182 Model::Gemma4(m) => m.chat_template.clone(),
183 Model::Glm52(m) => m.chat_template.clone(),
184 }
185 }
186
187 fn is_synthetic(&self) -> bool {
192 match self {
193 Model::Gguf(m) => m.is_synthetic,
194 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
195 }
196 }
197
198 fn tokenizer_kind(&self) -> &'static str {
199 match self {
200 Model::Gguf(m) => m.tokenizer.kind(),
201 Model::Kimi(_) => "kimi-tiktoken-bpe",
202 Model::Mla(m) => m.tokenizer.kind(),
203 Model::Gemma4(m) => m.tokenizer.kind(),
204 Model::Glm52(m) => m.tokenizer.kind(),
205 }
206 }
207
208 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
212 match self {
213 Model::Gguf(m) => m.decoder.expert_store_stats(),
214 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
215 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
216 }
217 }
218
219 pub(crate) fn name(&self) -> &str {
220 match self {
221 Model::Gguf(m) => m.decoder.config.name,
222 Model::Kimi(_) => "kimi-k3",
223 Model::Mla(m) => m.name.as_str(),
224 Model::Gemma4(m) => m.name.as_str(),
225 Model::Glm52(m) => m.name.as_str(),
226 }
227 }
228
229 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
256 match self {
257 Model::Gguf(m) => m.tokenizer.encode(text, specials),
258 Model::Kimi(m) => m
259 .tokenizer
260 .encode(text, specials)
261 .into_iter()
262 .map(|id| id as usize)
263 .collect(),
264 Model::Mla(m) => m.tokenizer.encode(text, specials),
265 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
266 Model::Glm52(m) => m.tokenizer.encode(text, specials),
267 }
268 }
269
270 pub(crate) fn bos_id(&self) -> Option<usize> {
279 match self {
280 Model::Gguf(m) => m.bos_id,
281 Model::Kimi(_) => None,
282 Model::Mla(m) => m.bos_id,
283 Model::Gemma4(m) => m.bos_id,
284 Model::Glm52(m) => m.bos_id,
285 }
286 }
287
288 pub(crate) fn decode(&self, ids: &[usize]) -> String {
289 match self {
290 Model::Gguf(m) => m.tokenizer.decode(ids),
291 Model::Kimi(m) => {
292 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
293 m.tokenizer.decode(&ids32)
294 }
295 Model::Mla(m) => m.tokenizer.decode(ids),
296 Model::Gemma4(m) => m.tokenizer.decode(ids),
297 Model::Glm52(m) => m.tokenizer.decode(ids),
298 }
299 }
300
301 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
304 match self {
305 Model::Gguf(m) => {
306 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
307 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
308 }
309 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
310 }
311 }
312
313 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
319 match self {
320 Model::Gguf(m) => Some(&m.decoder),
321 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
322 }
323 }
324
325 pub(crate) fn vocab_size(&self) -> Option<usize> {
326 match self {
327 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
328 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
329 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
330 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
331 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
332 }
333 }
334
335 fn has_real_vocabulary(&self) -> bool {
342 match self {
343 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
344 Model::Kimi(_) => true,
345 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
346 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
347 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
348 }
349 }
350}
351
352impl frink_models::dry::DryVocab for Model {
358 fn n_tokens(&self) -> usize {
359 self.vocab_size().unwrap_or(0)
360 }
361
362 fn detokenize(&self, token: usize) -> String {
363 self.decode(&[token])
364 }
365
366 fn tokenize(&self, text: &str) -> Vec<usize> {
367 self.encode(text, SpecialTokens::AsText)
368 }
369}
370
371pub(crate) struct AppState {
372 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
386 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
405 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
409 pub(crate) tasks: Arc<tasks::TaskRegistry>,
411 pub(crate) cancels: Arc<cancel::CancelRegistry>,
415 pub(crate) stats: stats::Stats,
418 pub(crate) streams: resume::StreamRegistry,
421 pub(crate) model_dir: Option<PathBuf>,
423 response_cache: Mutex<ResponseCache>,
427 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
439 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
452 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
462 sessions: session::SessionStore,
468 requests_total: std::sync::atomic::AtomicU64,
469 request_errors_total: std::sync::atomic::AtomicU64,
470 started_at: std::time::Instant,
471 last_request_ms: std::sync::atomic::AtomicU64,
476 detection: Arc<health::Detection>,
478 mcp: Option<mcp::LoadedMcpConfig>,
480 pub(crate) continuous_batching_enabled: bool,
484 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
488 loading_model: Mutex<Option<String>>,
492 last_load_error: Mutex<Option<(String, String)>>,
496 pub(crate) serving: Mutex<crate::stats::ServingStats>,
501 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
506 pub(crate) footprint:
512 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
513 pub(crate) started_unix: u64,
523}
524
525pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
531
532impl AppState {
533 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
537 self.active
538 .read()
539 .unwrap_or_else(|p| p.into_inner())
540 .clone()
541 }
542
543 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
549 self.active().ok_or_else(|| {
550 (
551 StatusCode::SERVICE_UNAVAILABLE,
552 Json(serde_json::json!({"error": {
553 "message": "no model is loaded; POST /admin/models/load with an id from \
554 GET /admin/models",
555 "type": "model_not_loaded"
556 }})),
557 )
558 })
559 }
560
561 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
569 Ok(Arc::clone(self.require_active()?.generative()?))
570 }
571
572 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
580 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
581 std::mem::replace(&mut *guard, next)
582 }
583
584 fn mark_request_finished(&self) {
588 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
589 self.last_request_ms
590 .store(ms, std::sync::atomic::Ordering::Relaxed);
591 }
592
593 pub(crate) fn uptime(&self) -> Duration {
594 self.started_at.elapsed()
595 }
596
597 pub(crate) fn requests_total(&self) -> u64 {
598 self.requests_total
599 .load(std::sync::atomic::Ordering::Relaxed)
600 }
601
602 pub(crate) fn errors_total(&self) -> u64 {
603 self.request_errors_total
604 .load(std::sync::atomic::Ordering::Relaxed)
605 }
606
607 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
608 lock_cache(&self.response_cache).stats()
609 }
610
611 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
614 let last = self
615 .last_request_ms
616 .load(std::sync::atomic::Ordering::Relaxed);
617 (last > 0)
618 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
619 .map(|age| age.max(0.0))
620 }
621
622 pub(crate) fn loading_model_id(&self) -> Option<String> {
623 self.loading_model
624 .lock()
625 .unwrap_or_else(|p| p.into_inner())
626 .clone()
627 }
628
629 pub(crate) fn set_loading_model(&self, id: Option<String>) {
630 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
631 }
632
633 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
634 self.last_load_error
635 .lock()
636 .unwrap_or_else(|p| p.into_inner())
637 .clone()
638 }
639
640 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
641 *self
642 .last_load_error
643 .lock()
644 .unwrap_or_else(|p| p.into_inner()) = error;
645 }
646
647 pub(crate) fn active_model_name(&self) -> Option<String> {
656 self.active().map(|a| a.name().to_string())
657 }
658
659 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
669 self.embedding
670 .clone()
671 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
672 }
673
674 pub(crate) fn embedding_model_name(&self) -> Option<String> {
678 match self.embedding_model() {
679 Some(e) => Some(e.name().to_string()),
680 None => self.active_model_name(),
681 }
682 }
683
684 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
685 self.stats.record(stats::entry(record));
686 }
687}
688
689fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
694 cache
695 .lock()
696 .unwrap_or_else(|poisoned| poisoned.into_inner())
697}
698
699#[derive(Debug, Clone, Deserialize)]
700#[serde(untagged)]
701pub(crate) enum MessageContent {
702 Text(String),
703 Parts(Vec<ContentPart>),
704}
705
706#[derive(Debug, Clone, Deserialize)]
707struct ContentPart {
708 #[serde(rename = "type")]
709 kind: String,
710 #[serde(default)]
711 text: Option<String>,
712 #[serde(default)]
713 image_url: Option<serde_json::Value>,
714}
715
716impl MessageContent {
717 fn as_text(&self) -> String {
718 match self {
719 Self::Text(s) => s.clone(),
720 Self::Parts(parts) => parts
721 .iter()
722 .filter_map(|p| p.text.as_deref())
723 .collect::<Vec<_>>()
724 .join(""),
725 }
726 }
727
728 fn has_image(&self) -> bool {
729 match self {
730 Self::Text(_) => false,
731 Self::Parts(parts) => parts
732 .iter()
733 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
734 }
735 }
736}
737
738#[derive(Debug, Clone, Deserialize)]
739pub(crate) struct ChatMessage {
740 pub(crate) role: String,
741 #[serde(default)]
745 pub(crate) content: Option<MessageContent>,
746 #[serde(default)]
750 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
751 #[serde(default)]
756 #[allow(dead_code)]
757 pub(crate) tool_call_id: Option<String>,
758 #[serde(default, alias = "reasoning")]
774 pub(crate) reasoning_content: Option<String>,
775}
776
777impl ChatMessage {
778 fn rendered_content(&self) -> String {
785 let mut out = self
786 .content
787 .as_ref()
788 .map(MessageContent::as_text)
789 .unwrap_or_default();
790 if let Some(calls) = &self.tool_calls {
791 for call in calls {
792 out.push_str(&format!(
793 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
794 call.function.name, call.function.arguments
795 ));
796 }
797 }
798 out
799 }
800}
801
802#[derive(Debug, Clone, Deserialize)]
803pub(crate) struct ToolCallIn {
804 #[serde(default)]
805 #[allow(dead_code)]
806 id: String,
807 #[serde(rename = "type", default)]
808 #[allow(dead_code)]
809 kind: String,
810 function: ToolCallFunctionIn,
811}
812
813#[derive(Debug, Clone, Deserialize)]
814struct ToolCallFunctionIn {
815 name: String,
816 arguments: String,
821}
822
823#[derive(Debug, Clone, Deserialize)]
826struct ToolDef {
827 #[serde(rename = "type", default)]
828 #[allow(dead_code)]
829 kind: String,
830 function: ToolFunctionDef,
831}
832
833#[derive(Debug, Clone, Deserialize)]
834struct ToolFunctionDef {
835 name: String,
836 #[serde(default)]
837 description: Option<String>,
838 #[serde(default)]
839 parameters: Option<serde_json::Value>,
840}
841
842#[derive(Debug, Clone, Deserialize)]
853#[serde(untagged)]
854enum ToolChoice {
855 Mode(String),
856 Specific(serde_json::Value),
857}
858
859#[derive(Deserialize)]
862#[serde(untagged)]
863enum StopParam {
864 One(String),
865 Many(Vec<String>),
866}
867
868#[derive(Deserialize)]
869struct ChatCompletionRequest {
870 model: String,
871 messages: Vec<ChatMessage>,
872 #[serde(default = "default_max_tokens")]
873 max_tokens: usize,
874 #[serde(default)]
875 temperature: Option<f32>,
876 #[serde(default)]
877 top_p: Option<f32>,
878 #[serde(default)]
883 min_p: Option<f32>,
884 #[serde(default)]
885 top_k: Option<usize>,
886 #[serde(default)]
887 repetition_penalty: Option<f32>,
888 #[serde(flatten)]
892 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
893 #[serde(default)]
894 seed: Option<u64>,
895 #[serde(default)]
896 stop: Option<StopParam>,
897 #[serde(default)]
898 stream: Option<bool>,
899 #[serde(default)]
912 stream_resumable: Option<bool>,
913 #[serde(default)]
923 ignore_eos: Option<bool>,
924 #[serde(default)]
925 tools: Vec<ToolDef>,
926 #[serde(default)]
927 tool_choice: Option<ToolChoice>,
928 #[serde(default)]
936 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
937 #[serde(default)]
945 reasoning_effort: Option<String>,
946 #[serde(default)]
950 thinking: Option<ThinkingSwitch>,
951 #[serde(default)]
956 session_id: Option<String>,
957 #[serde(default, deserialize_with = "continuation::deserialize")]
965 continue_final_message: continuation::ContinueFinalMessage,
966 #[serde(default, alias = "thinking_budget_tokens")]
974 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
975 #[serde(default)]
977 logprobs: Option<bool>,
978 #[serde(default)]
979 top_logprobs: Option<u32>,
980 #[serde(default)]
981 n: Option<u32>,
982 #[serde(default)]
983 presence_penalty: Option<f32>,
984 #[serde(default)]
985 frequency_penalty: Option<f32>,
986 #[serde(default)]
987 response_format: Option<serde_json::Value>,
988 #[serde(default)]
995 logit_bias: Option<serde_json::Value>,
996 #[serde(default)]
1000 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1001 #[serde(default)]
1011 samplers: Option<serde_json::Value>,
1012 #[serde(default)]
1022 grammar: Option<String>,
1023}
1024
1025const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1037
1038#[derive(Debug, Clone, Deserialize)]
1040pub(crate) struct ThinkingSwitch {
1041 #[serde(rename = "type")]
1042 pub(crate) kind: String,
1043}
1044
1045const THINKING_KWARG_KEYS: [&str; 4] = [
1049 "enable_thinking",
1050 "thinking",
1051 "thinking_mode",
1052 "reasoning_effort",
1053];
1054
1055const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1059
1060fn default_max_tokens() -> usize {
1061 DEFAULT_CHAT_MAX_TOKENS
1062}
1063
1064impl ChatCompletionRequest {
1065 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1074 let mut knobs = SamplingKnobs {
1075 temperature: self.temperature,
1076 top_p: self.top_p,
1077 min_p: self.min_p,
1078 top_k: self.top_k,
1079 repetition_penalty: self.repetition_penalty,
1080 presence_penalty: self.presence_penalty,
1081 frequency_penalty: self.frequency_penalty,
1082 penalty_last_n: None,
1086 sampler_order: unsupported_sampling::parse_sampler_order(
1087 self.samplers.as_ref(),
1088 "/v1/chat/completions",
1089 )?,
1090 ..SamplingKnobs::default()
1091 };
1092 self.extra_samplers.apply(&mut knobs);
1093 Ok(knobs)
1094 }
1095
1096 fn sampling_params(
1097 &self,
1098 model: crate::sampling_knobs::SamplerModel<'_>,
1099 ) -> Result<SamplingParams, ApiError> {
1100 self.sampling_knobs()?.resolve(model).map_err(|e| {
1101 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1102 })
1103 }
1104
1105 fn stop_sequences(&self) -> Vec<String> {
1106 self.stop
1107 .as_ref()
1108 .map(|s| match s {
1109 StopParam::One(v) => vec![v.clone()],
1110 StopParam::Many(v) => v.clone(),
1111 })
1112 .unwrap_or_default()
1113 }
1114
1115 fn tools_active(&self) -> bool {
1120 !self.tools.is_empty()
1121 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1122 }
1123
1124 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1136 match &self.tool_choice {
1137 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1138 Some(ToolChoice::Specific(value)) => {
1139 let name = value
1143 .get("function")
1144 .and_then(|f| f.get("name"))
1145 .or_else(|| value.get("name"))
1146 .and_then(|n| n.as_str());
1147 match name {
1148 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1149 None => Err(invalid_request(
1150 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1151 function.name",
1152 "tool_choice",
1153 )),
1154 }
1155 }
1156 _ => Ok(None),
1157 }
1158 }
1159
1160 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1162 self.tools
1163 .iter()
1164 .map(|t| tool_grammar::ToolSpec {
1165 name: &t.function.name,
1166 parameters: t.function.parameters.as_ref(),
1167 })
1168 .collect()
1169 }
1170
1171 fn thinking_direction(&self) -> Option<bool> {
1214 if let Some(switch) = &self.thinking {
1215 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1216 "disabled" => Some(false),
1217 "enabled" => Some(true),
1218 _ => None,
1221 };
1222 }
1223 let effort = self.reasoning_effort.as_ref()?;
1224 DISABLE_EFFORTS
1225 .contains(&effort.trim().to_ascii_lowercase().as_str())
1226 .then_some(false)
1227 }
1228
1229 fn resolve_template_kwargs(
1230 &self,
1231 template: &chat_template::PromptTemplate,
1232 ) -> serde_json::Map<String, serde_json::Value> {
1233 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1234 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1238
1239 if !caller_steered {
1240 match self.thinking_direction() {
1241 Some(false) => {
1242 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1243 kwargs.insert(k, v);
1244 }
1245 return kwargs;
1248 }
1249 Some(true) => {
1250 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1251 kwargs.insert(k, v);
1252 }
1253 }
1254 None => {}
1255 }
1256 if let Some(effort) = &self.reasoning_effort {
1257 kwargs
1258 .entry("reasoning_effort".to_string())
1259 .or_insert_with(|| serde_json::json!(effort));
1260 }
1261 }
1262
1263 let offered: Vec<serde_json::Value> = if self.tools_active() {
1264 self.tools.iter().map(chat_template::tool_json).collect()
1265 } else {
1266 Vec::new()
1267 };
1268 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1269 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1270 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1271 kwargs.entry(k).or_insert(v);
1272 }
1273 }
1274 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1275 crate::policy::effort::EffortMapping::Mapped(to) => {
1276 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1277 }
1278 crate::policy::effort::EffortMapping::Dropped => {
1279 tracing::debug!(
1280 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1281 enough, so its own default applies"
1282 );
1283 }
1284 crate::policy::effort::EffortMapping::Unchanged => {}
1285 }
1286 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1287 kwargs
1288 }
1289
1290 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1293 if self.max_tokens == 0 {
1300 return Err(invalid_request(
1301 "max_tokens must be at least 1",
1302 "max_tokens",
1303 ));
1304 }
1305 if let Some(switch) = &self.thinking {
1310 let kind = switch.kind.trim().to_ascii_lowercase();
1311 if kind != "enabled" && kind != "disabled" {
1312 return Err(invalid_request(
1313 "thinking.type must be \"enabled\" or \"disabled\"",
1314 "thinking.type",
1315 ));
1316 }
1317 }
1318 for msg in &self.messages {
1319 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1320 return Err(unsupported_feature(
1321 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1322 ));
1323 }
1324 }
1325 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1326 return Err(unsupported_feature(
1327 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1328 ));
1329 }
1330 if self.n.is_some_and(|n| n > 1) {
1331 return Err(unsupported_feature(
1332 "n > 1 is not implemented (single completion only)",
1333 ));
1334 }
1335 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1336 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1340 let stated_grammar =
1355 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1356 if let Some(forced) = self.forced_tool_choice()? {
1364 if self.tools.is_empty() {
1365 return Err(invalid_request(
1366 "tool_choice forces a tool call, but no tools were offered",
1367 "tool_choice",
1368 ));
1369 }
1370 if let tool_grammar::Forced::Named(name) = forced {
1371 if !self.tools.iter().any(|t| t.function.name == name) {
1372 return Err(invalid_request(
1373 &format!(
1374 "tool_choice names {name:?}, which is not one of the tools offered"
1375 ),
1376 "tool_choice",
1377 ));
1378 }
1379 }
1380 if stated_grammar.is_some() {
1390 return Err(invalid_request(
1391 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1392 are two different constraints on the same generation; send one",
1393 "tool_choice",
1394 ));
1395 }
1396 if self.json_object_mode() {
1397 return Err(invalid_request(
1398 "a forced tool_choice cannot be combined with response_format json_object: \
1399 the tool-call markers are not JSON",
1400 "tool_choice",
1401 ));
1402 }
1403 }
1404 Ok(())
1405 }
1406
1407 fn effective_stop_sequences(&self) -> Vec<String> {
1414 let mut stop = self.stop_sequences();
1415 if self.tools_active() {
1416 stop.push("</tool_call>".to_string());
1417 }
1418 stop
1419 }
1420
1421 fn json_object_mode(&self) -> bool {
1422 self.response_format
1423 .as_ref()
1424 .and_then(|v| v.get("type"))
1425 .and_then(|v| v.as_str())
1426 == Some("json_object")
1427 }
1428}
1429
1430#[derive(Serialize)]
1431struct ChatCompletionChoice {
1432 index: usize,
1433 message: ChatCompletionResponseMessage,
1434 finish_reason: &'static str,
1435}
1436
1437#[derive(Serialize)]
1438struct ChatCompletionResponseMessage {
1439 role: &'static str,
1440 #[serde(skip_serializing_if = "Option::is_none")]
1441 content: Option<String>,
1442 #[serde(skip_serializing_if = "Option::is_none")]
1446 reasoning_content: Option<String>,
1447 #[serde(skip_serializing_if = "Option::is_none")]
1448 tool_calls: Option<Vec<ToolCallOut>>,
1449}
1450
1451#[derive(Serialize, Clone)]
1452struct ToolCallOut {
1453 id: String,
1454 #[serde(rename = "type")]
1455 kind: &'static str,
1456 function: ToolCallFunctionOut,
1457}
1458
1459#[derive(Serialize, Clone)]
1467struct ToolCallDelta {
1468 index: usize,
1469 #[serde(skip_serializing_if = "Option::is_none")]
1470 id: Option<String>,
1471 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1472 kind: Option<&'static str>,
1473 function: ToolCallFunctionDelta,
1474}
1475
1476#[derive(Serialize, Clone, Default)]
1477struct ToolCallFunctionDelta {
1478 #[serde(skip_serializing_if = "Option::is_none")]
1479 name: Option<String>,
1480 #[serde(skip_serializing_if = "Option::is_none")]
1483 arguments: Option<String>,
1484}
1485
1486impl ToolCallDelta {
1487 fn whole(index: usize, name: String, arguments: String) -> Self {
1489 ToolCallDelta {
1490 index,
1491 id: Some(format!("call_{index}")),
1492 kind: Some("function"),
1493 function: ToolCallFunctionDelta {
1494 name: Some(name),
1495 arguments: Some(arguments),
1496 },
1497 }
1498 }
1499
1500 fn opening(index: usize, name: String) -> Self {
1502 ToolCallDelta {
1503 index,
1504 id: Some(format!("call_{index}")),
1505 kind: Some("function"),
1506 function: ToolCallFunctionDelta {
1507 name: Some(name),
1508 arguments: Some(String::new()),
1509 },
1510 }
1511 }
1512
1513 fn arguments(index: usize, fragment: String) -> Self {
1515 ToolCallDelta {
1516 index,
1517 id: None,
1518 kind: None,
1519 function: ToolCallFunctionDelta {
1520 name: None,
1521 arguments: Some(fragment),
1522 },
1523 }
1524 }
1525}
1526
1527#[derive(Serialize, Clone)]
1528struct ToolCallFunctionOut {
1529 name: String,
1530 arguments: String,
1534}
1535
1536#[derive(Serialize)]
1537struct ChatCompletionResponse {
1538 id: String,
1539 request_id: String,
1545 object: &'static str,
1546 model: String,
1547 choices: Vec<ChatCompletionChoice>,
1548 usage: generate::Usage,
1553 frink_cache: &'static str,
1564}
1565
1566#[derive(Serialize)]
1567struct ChatCompletionChunkDelta {
1568 #[serde(skip_serializing_if = "Option::is_none")]
1569 role: Option<&'static str>,
1570 #[serde(skip_serializing_if = "Option::is_none")]
1571 content: Option<String>,
1572 #[serde(skip_serializing_if = "Option::is_none")]
1574 reasoning_content: Option<String>,
1575 #[serde(skip_serializing_if = "Option::is_none")]
1576 tool_calls: Option<Vec<ToolCallDelta>>,
1577}
1578
1579#[derive(Serialize)]
1580struct ChatCompletionChunkChoice {
1581 index: usize,
1582 delta: ChatCompletionChunkDelta,
1583 finish_reason: Option<&'static str>,
1584}
1585
1586#[derive(Serialize)]
1587struct ChatCompletionChunk {
1588 id: String,
1589 #[serde(skip_serializing_if = "Option::is_none")]
1596 request_id: Option<String>,
1597 object: &'static str,
1598 model: String,
1599 choices: Vec<ChatCompletionChunkChoice>,
1600 #[serde(skip_serializing_if = "Option::is_none")]
1603 usage: Option<generate::Usage>,
1604}
1605
1606async fn health(State(state): State<Arc<AppState>>) -> Response {
1612 let snapshot = state.detection.snapshot();
1613 let mut capabilities = snapshot.capabilities;
1614 let active = state.active();
1615
1616 capabilities.push(match active.as_deref() {
1619 None => frink_api::Capability::unavailable(
1624 frink_api::health::capability::REAL_WEIGHTS,
1625 frink_api::health::reason::MODEL_NOT_LOADED,
1626 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1627 ),
1628 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1629 frink_api::health::capability::REAL_WEIGHTS,
1630 frink_api::health::reason::MODEL_NOT_LOADED,
1631 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1632 checkpoint. Output from this model is noise.",
1633 ),
1634 Some(active) if active.encoder().is_some() => {
1644 let endpoints = active
1645 .encoder()
1646 .map(|e| encoder_endpoints(e))
1647 .unwrap_or_default();
1648 let served_by = match endpoints.is_empty() {
1649 true => "no endpoint in this build serves it".to_string(),
1650 false => format!("served by {}", endpoints.join(" and ")),
1651 };
1652 frink_api::Capability::available(
1653 frink_api::health::capability::REAL_WEIGHTS,
1654 format!(
1655 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1656 {served_by}; generation endpoints refuse it.",
1657 active.name(),
1658 ),
1659 )
1660 }
1661 Some(active) => frink_api::Capability::available(
1662 frink_api::health::capability::REAL_WEIGHTS,
1663 format!("Serving the real checkpoint '{}'.", active.name()),
1664 ),
1665 });
1666 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1667 frink_api::Capability::available(
1668 frink_api::health::capability::CONTINUOUS_BATCHING,
1669 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1670 "On by default on Metal. Concurrent requests share one batched decode worker."
1671 } else {
1672 "Concurrent requests share one batched decode step."
1673 },
1674 )
1675 } else if state.metal_private_decode_gate.is_some() {
1676 frink_api::Capability::unavailable(
1677 frink_api::health::capability::CONTINUOUS_BATCHING,
1678 frink_api::health::reason::DISABLED,
1679 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1680 )
1681 } else {
1682 frink_api::Capability::unavailable(
1683 frink_api::health::capability::CONTINUOUS_BATCHING,
1684 frink_api::health::reason::DISABLED,
1685 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1686 )
1687 });
1688
1689 let last_request_ms = state
1690 .last_request_ms
1691 .load(std::sync::atomic::Ordering::Relaxed);
1692 let uptime = state.started_at.elapsed();
1693 let health_state = if active.is_none() {
1701 frink_api::HealthState::Unavailable
1702 } else {
1703 snapshot.state
1704 };
1705 let body = frink_api::HealthResponse {
1706 state: health_state,
1707 reason: match health_state {
1708 frink_api::HealthState::Ready => None,
1709 frink_api::HealthState::Unavailable => {
1710 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1711 }
1712 frink_api::HealthState::Detecting => {
1713 Some(frink_api::health::reason::DETECTING.to_string())
1714 }
1715 },
1716 detail: match health_state {
1717 frink_api::HealthState::Ready => None,
1718 frink_api::HealthState::Unavailable => Some(
1719 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1720 .to_string(),
1721 ),
1722 frink_api::HealthState::Detecting => {
1723 Some("Probing available compute backends.".to_string())
1724 }
1725 },
1726 model: active
1727 .as_deref()
1728 .map(|active| frink_api::health::ModelSummary {
1729 id: active.name().to_string(),
1730 tokenizer: active.tokenizer_kind().to_string(),
1731 synthetic_weights: active.is_synthetic(),
1732 }),
1733 capabilities,
1734 version: env!("CARGO_PKG_VERSION").to_string(),
1735 pid: std::process::id(),
1736 uptime_seconds: uptime.as_secs_f64(),
1737 server_time_unix_ms: std::time::SystemTime::now()
1738 .duration_since(std::time::UNIX_EPOCH)
1739 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1740 .unwrap_or(0),
1741 last_request_age_seconds: (last_request_ms > 0)
1742 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1743 .map(|age| age.max(0.0)),
1744 };
1745
1746 let status =
1747 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1748 (status, Json(body)).into_response()
1749}
1750
1751async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1752 let Some(active) = state.active() else {
1756 return Json(serde_json::json!({ "object": "list", "data": [] }));
1757 };
1758 let mut model_entry = serde_json::json!({
1759 "id": active.name(),
1760 "object": "model",
1761 "frink_synthetic_weights": active.is_synthetic(),
1762 "frink_tokenizer": active.tokenizer_kind(),
1763 });
1764 if let Some(encoder) = active.encoder() {
1770 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1771 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1772 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1773 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1774 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1775 }
1776 if let Some(model) = active.generative_opt() {
1784 let parser_configured = active.reasoning_format().is_some();
1785 let gears = model.chat_template().think_gears(parser_configured);
1786 if !gears.is_empty() {
1787 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1788 if let Some(default) = &gears.default {
1789 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1790 }
1791 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1794 }
1795 }
1796 if let Some(mcp) = &state.mcp {
1797 model_entry["frink_mcp"] = mcp.models_metadata();
1798 }
1799 Json(serde_json::json!({
1800 "object": "list",
1801 "data": [model_entry]
1802 }))
1803}
1804
1805async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1819 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1820 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1821 let active = state.active();
1822 Json(serde_json::json!({
1823 "model": active.as_ref().map(|a| a.name()),
1824 "state": state
1825 .maintenance
1826 .lock()
1827 .unwrap_or_else(|p| p.into_inner())
1828 .state()
1829 .as_str(),
1830 "uptime_s": state.uptime().as_secs(),
1831 "throughput": {
1832 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1833 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1834 },
1835 "requests": {
1836 "active": state.cancels.live_count(),
1837 "completed": state.stats.recorded_total(),
1838 "p95_ms": state.stats.p95_duration_ms(),
1839 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1840 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1841 "completion_tokens_total": state.stats.tokens_generated_total(),
1842 },
1843 "pools": cache_admin::pool_gauges(&state),
1849 "memory": cache_admin::footprint_json(&state),
1852 }))
1853}
1854
1855#[derive(Deserialize)]
1856struct RequestsQuery {
1857 #[serde(default)]
1858 since: u64,
1859 #[serde(default = "default_requests_limit")]
1860 limit: usize,
1861}
1862
1863fn default_requests_limit() -> usize {
1864 stats::MAX_PAGE
1865}
1866
1867async fn recent_requests(
1875 State(state): State<Arc<AppState>>,
1876 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1877) -> Json<serde_json::Value> {
1878 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1879 Json(serde_json::json!({
1880 "requests": rows,
1881 "next_cursor": cursor,
1882 "missed": missed,
1883 "total": state.stats.recorded_total(),
1884 }))
1885}
1886
1887#[derive(Serialize)]
1888struct CombinedCacheStats {
1889 response_cache: response_cache::CacheStats,
1890 prefix_cache: Option<frink_models::PrefixCacheStats>,
1892}
1893
1894async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1895 Json(CombinedCacheStats {
1896 response_cache: lock_cache(&state.response_cache).stats(),
1897 prefix_cache: state
1898 .prefix_cache
1899 .as_ref()
1900 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1901 })
1902}
1903
1904async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1909 use std::sync::atomic::Ordering;
1910
1911 let cache_stats = lock_cache(&state.response_cache).stats();
1912 let active = state.active();
1913 let requests_total = state.requests_total.load(Ordering::Relaxed);
1914 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1915 let uptime = state.started_at.elapsed().as_secs_f64();
1916
1917 let body = format!(
1918 "# HELP frink_requests_total Total chat completion requests received.\n\
1919 # TYPE frink_requests_total counter\n\
1920 frink_requests_total {requests_total}\n\
1921 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1922 # TYPE frink_request_errors_total counter\n\
1923 frink_request_errors_total {errors_total}\n\
1924 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1925 # TYPE frink_cache_hits_total counter\n\
1926 frink_cache_hits_total {}\n\
1927 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1928 # TYPE frink_cache_misses_total counter\n\
1929 frink_cache_misses_total {}\n\
1930 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1931 # TYPE frink_cache_entries gauge\n\
1932 frink_cache_entries {}\n\
1933 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1934 # TYPE frink_synthetic_weights gauge\n\
1935 frink_synthetic_weights {}\n\
1936 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1937 # TYPE frink_uptime_seconds gauge\n\
1938 frink_uptime_seconds {uptime}\n",
1939 cache_stats.hits,
1940 cache_stats.misses,
1941 cache_stats.entries,
1942 active
1946 .as_ref()
1947 .map(|a| a.is_synthetic() as u8)
1948 .unwrap_or(0),
1949 );
1950
1951 let body = match active
1955 .as_ref()
1956 .and_then(|a| a.expert_store_stats())
1957 {
1958 Some(es) => format!(
1959 "{body}\
1960 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
1961 # TYPE frink_expert_cache_hits_total counter\n\
1962 frink_expert_cache_hits_total {}\n\
1963 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1964 # TYPE frink_expert_cache_misses_total counter\n\
1965 frink_expert_cache_misses_total {}\n\
1966 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
1967 # TYPE frink_expert_cache_evictions_total counter\n\
1968 frink_expert_cache_evictions_total {}\n\
1969 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1970 # TYPE frink_expert_cache_pass_throughs_total counter\n\
1971 frink_expert_cache_pass_throughs_total {}\n\
1972 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1973 # TYPE frink_expert_cache_bytes_read_total counter\n\
1974 frink_expert_cache_bytes_read_total {}\n\
1975 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1976 # TYPE frink_expert_cache_resident_bytes gauge\n\
1977 frink_expert_cache_resident_bytes {}\n",
1978 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1979 ),
1980 None => body,
1981 };
1982
1983 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
1988 Some(batcher) => {
1989 let sched = batcher.stats();
1990 format!(
1991 "{body}\
1992 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
1993 # TYPE frink_prefill_chunks_total counter\n\
1994 frink_prefill_chunks_total {}\n\
1995 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
1996 # TYPE frink_prefill_tokens_total counter\n\
1997 frink_prefill_tokens_total {}\n\
1998 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
1999 # TYPE frink_decode_steps_total counter\n\
2000 frink_decode_steps_total {}\n\
2001 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2002 # TYPE frink_scheduler_queue_depth gauge\n\
2003 frink_scheduler_queue_depth {}\n\
2004 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2005 # TYPE frink_scheduler_queue_rejected_total counter\n\
2006 frink_scheduler_queue_rejected_total {}\n\
2007 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2008 # TYPE frink_kv_blocks_total gauge\n\
2009 frink_kv_blocks_total {}\n\
2010 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2011 # TYPE frink_kv_blocks_free gauge\n\
2012 frink_kv_blocks_free {}\n\
2013 # HELP frink_kv_block_size Token positions per KV block.\n\
2014 # TYPE frink_kv_block_size gauge\n\
2015 frink_kv_block_size {}\n\
2016 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2017 # TYPE frink_kv_rejected_too_large_total counter\n\
2018 frink_kv_rejected_too_large_total {}\n\
2019 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2020 # TYPE frink_kv_rejected_context_length_total counter\n\
2021 frink_kv_rejected_context_length_total {}\n\
2022 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2023 # TYPE frink_scheduler_aborted_total counter\n\
2024 frink_scheduler_aborted_total {}\n\
2025 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2026 # TYPE frink_scheduler_max_seqs gauge\n\
2027 frink_scheduler_max_seqs {}\n\
2028 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2029 # TYPE frink_scheduler_prefill_chunk gauge\n\
2030 frink_scheduler_prefill_chunk {}\n",
2031 sched.prefill_chunks,
2032 sched.prefill_tokens,
2033 sched.decode_steps,
2034 sched.queue_depth,
2035 sched.queue_rejected,
2036 sched.kv_blocks_total,
2037 sched.kv_blocks_free,
2038 sched.kv_block_size,
2039 sched.kv_rejected_too_large,
2040 sched.kv_rejected_context_length,
2041 sched.aborted,
2042 sched.max_seqs,
2043 sched.prefill_chunk,
2044 )
2045 }
2046 None => body,
2047 };
2048
2049 (
2050 [(
2051 axum::http::header::CONTENT_TYPE,
2052 "text/plain; version=0.0.4",
2053 )],
2054 body,
2055 )
2056 .into_response()
2057}
2058
2059pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2060
2061pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2067 (
2068 StatusCode::BAD_REQUEST,
2069 Json(serde_json::json!({"error": {
2070 "message": message,
2071 "type": "invalid_request_error",
2072 "param": param,
2073 "code": null,
2074 }})),
2075 )
2076}
2077
2078pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2079 (
2080 StatusCode::NOT_IMPLEMENTED,
2081 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2082 )
2083}
2084
2085pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2086 let status = match e {
2087 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2088 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2093 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2099 StatusCode::SERVICE_UNAVAILABLE
2100 }
2101 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2105 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2109 };
2110 tracing::warn!("decode error: {e}");
2111 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2112 if let generate::DecodeError::KvBudgetExceeded {
2118 binding,
2119 estimated_bytes,
2120 limit_bytes,
2121 positions,
2122 positions_limit,
2123 ..
2124 } = &e
2125 {
2126 body["error"]["type"] = serde_json::json!("invalid_request_error");
2127 body["error"]["code"] = serde_json::json!(binding);
2128 body["error"]["binding"] = serde_json::json!(binding);
2129 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2130 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2131 body["error"]["positions"] = serde_json::json!(positions);
2132 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2133 }
2134 if let Some(secs) = e.retry_after_secs() {
2138 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2139 }
2140 (status, Json(body))
2141}
2142
2143pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2144 tracing::error!("generation task panicked: {e}");
2145 (
2146 StatusCode::INTERNAL_SERVER_ERROR,
2147 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2148 )
2149}
2150
2151#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2161 model: &Model,
2162 prompt: &str,
2163 params: &GenerationParams,
2164 kv_pool: Option<&generate::KvPoolConfig>,
2165 paged_kv: Option<&generate::PagedKvConfig>,
2166 prefix_cache: Option<&Mutex<PrefixCache>>,
2167 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2168 ceiling: Option<&budget::ContextCeiling>,
2169 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2170 mut emit: impl FnMut(&str),
2171) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2172 let synthetic = model.is_synthetic();
2173 let _lora_lease = lora::lease(model, params.lora.as_deref());
2178 let mut chunks = Vec::new();
2179 let params = &{
2185 let mut resolved = params.clone();
2186 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2187 model.encode(text, SpecialTokens::Parse)
2188 });
2189 resolved.reasoning_budget = resolved
2195 .reasoning_budget
2196 .armed(resolved.reasoning, prompt, |text| {
2197 model.encode(text, SpecialTokens::Parse)
2198 })
2199 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2200 resolved
2201 };
2202 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2203 let _metal_private_guard =
2204 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2205 let (finish, usage) = match model {
2206 Model::Gguf(m) => {
2207 if let Some(batcher) = continuous_batcher {
2208 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2209 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2210 let (finish, _generated_ids, text, usage) = if synthetic {
2211 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2212 } else {
2213 batcher.generate_streaming(
2214 tokens,
2215 params.clone(),
2216 m.stop_tokens.clone(),
2217 Some(|chunk: &str| {
2218 if !chunk.is_empty() {
2219 chunks.push(chunk.to_string());
2220 emit(chunk);
2221 }
2222 }),
2223 )?
2224 };
2225 if !text.is_empty() && chunks.is_empty() {
2226 chunks.push(text);
2227 }
2228 (finish, usage)
2229 } else {
2230 generate::generate(
2231 &m.decoder,
2232 m.tokenizer.as_ref(),
2233 &m.stop_tokens,
2234 m.bos_id,
2235 prompt,
2236 params,
2237 kv_pool,
2238 paged_kv,
2239 prefix_cache,
2240 ceiling,
2241 |chunk| {
2242 chunks.push(chunk.to_string());
2243 if !synthetic {
2244 emit(chunk);
2245 }
2246 },
2247 )?
2248 }
2249 }
2250 Model::Kimi(m) => generate::generate_engine(
2251 &m.engine,
2252 &m.tokenizer,
2253 &m.stop_tokens,
2254 None,
2255 prompt,
2256 params,
2257 |chunk| {
2258 chunks.push(chunk.to_string());
2259 if !synthetic {
2260 emit(chunk);
2261 }
2262 },
2263 )?,
2264 Model::Mla(m) => generate::generate_engine(
2265 &m.engine,
2266 &m.tokenizer,
2267 &m.stop_tokens,
2268 m.bos_id,
2269 prompt,
2270 params,
2271 |chunk| {
2272 chunks.push(chunk.to_string());
2273 if !synthetic {
2274 emit(chunk);
2275 }
2276 },
2277 )?,
2278 Model::Gemma4(m) => generate::generate_engine(
2279 &m.engine,
2280 &m.tokenizer,
2281 &m.stop_tokens,
2282 m.bos_id,
2283 prompt,
2284 params,
2285 |chunk| {
2286 chunks.push(chunk.to_string());
2287 if !synthetic {
2288 emit(chunk);
2289 }
2290 },
2291 )?,
2292 Model::Glm52(m) => generate::generate_engine(
2293 &m.engine,
2294 &m.tokenizer,
2295 &m.stop_tokens,
2296 m.bos_id,
2297 prompt,
2298 params,
2299 |chunk| {
2300 chunks.push(chunk.to_string());
2301 if !synthetic {
2302 emit(chunk);
2303 }
2304 },
2305 )?,
2306 };
2307
2308 let mut full = chunks.concat();
2309 if synthetic {
2310 full = format!(
2311 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2312 to serve a real model. Decoded ids -> {full:?}]"
2313 );
2314 emit(&full);
2315 } else if used_batcher && !full.is_empty() && chunks.is_empty() {
2316 emit(&full);
2317 }
2318
2319 Ok((finish, usage, full))
2320}
2321
2322#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2327 model: &Model,
2328 prompt: &str,
2329 params: &GenerationParams,
2330 kv_pool: Option<&generate::KvPoolConfig>,
2331 paged_kv: Option<&generate::PagedKvConfig>,
2332 prefix_cache: Option<&Mutex<PrefixCache>>,
2333 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2334 ceiling: Option<&budget::ContextCeiling>,
2335 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2336) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2337 let (finish, usage, full) = run_generation_emit(
2338 model,
2339 prompt,
2340 params,
2341 kv_pool,
2342 paged_kv,
2343 prefix_cache,
2344 continuous_batcher,
2345 ceiling,
2346 metal_private_decode_gate,
2347 |_| {},
2348 )?;
2349 Ok((
2350 if full.is_empty() {
2351 Vec::new()
2352 } else {
2353 vec![full]
2354 },
2355 finish,
2356 usage,
2357 ))
2358}
2359
2360pub(crate) fn prompt_from_messages(
2371 messages: &[ChatMessage],
2372 template: &chat_template::PromptTemplate,
2373 tools: &[ToolDef],
2374 extra: serde_json::Map<String, serde_json::Value>,
2375) -> Result<String, ApiError> {
2376 let rendered = if tools.is_empty() || template.handles_tools() {
2377 template.render(messages, tools, extra)
2378 } else {
2379 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2380 with_preamble.push(ChatMessage {
2381 role: "system".to_string(),
2382 content: Some(MessageContent::Text(tool_preamble(tools))),
2383 tool_calls: None,
2384 tool_call_id: None,
2385 reasoning_content: None,
2386 });
2387 with_preamble.extend_from_slice(messages);
2388 template.render(&with_preamble, &[], extra)
2389 };
2390 rendered.map_err(template_error_response)
2391}
2392
2393fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2398 (
2399 StatusCode::BAD_REQUEST,
2400 Json(serde_json::json!({
2401 "error": {
2402 "message": format!("chat template failed to render: {err}"),
2403 "type": "invalid_request_error",
2404 "param": "messages",
2405 "code": null,
2406 }
2407 })),
2408 )
2409}
2410
2411fn tool_preamble(tools: &[ToolDef]) -> String {
2424 let mut out = String::from(
2425 "You can call tools to help answer the user. To call a tool, respond with \
2426 EXACTLY one line in this format and nothing else:\n\
2427 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2428 object matching that tool's parameters>}}</tool_call>\n\n\
2429 Available tools:\n",
2430 );
2431 for t in tools {
2432 out.push_str(&format!(
2433 "- {}: {}\n parameters (JSON schema): {}\n",
2434 t.function.name,
2435 t.function.description.as_deref().unwrap_or(""),
2436 t.function
2437 .parameters
2438 .as_ref()
2439 .map(|v| v.to_string())
2440 .unwrap_or_else(|| "{}".to_string()),
2441 ));
2442 }
2443 out
2444}
2445
2446fn tool_call_deltas(
2456 events: Vec<crate::policy::parser::ToolCallEvent>,
2457 opened: &std::cell::Cell<usize>,
2458) -> (String, Vec<ToolCallDelta>) {
2459 let mut text = String::new();
2460 let mut deltas = Vec::new();
2461 for event in events {
2462 match event {
2463 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2464 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2465 opened.set(opened.get().max(index + 1));
2466 deltas.push(ToolCallDelta::opening(index, name));
2467 }
2468 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2469 if !fragment.is_empty() {
2470 deltas.push(ToolCallDelta::arguments(index, fragment));
2471 }
2472 }
2473 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2474 }
2475 }
2476 (text, deltas)
2477}
2478
2479fn build_response_message(
2490 text: String,
2491 tools: &[ToolDef],
2492 posture: output::OutputPosture,
2493 base_finish: &'static str,
2494) -> (ChatCompletionResponseMessage, &'static str) {
2495 let parsed = output::parse_output(&text, tools, posture);
2496 let calls: Vec<ToolCallOut> = parsed
2497 .calls
2498 .into_iter()
2499 .enumerate()
2500 .map(|(index, call)| ToolCallOut {
2501 id: format!("call_{index}"),
2502 kind: "function",
2503 function: ToolCallFunctionOut {
2504 name: call.name,
2505 arguments: call.arguments,
2506 },
2507 })
2508 .collect();
2509 if !calls.is_empty() {
2510 return (
2511 ChatCompletionResponseMessage {
2512 role: "assistant",
2513 content: None,
2514 reasoning_content: parsed.reasoning,
2515 tool_calls: Some(calls),
2516 },
2517 "tool_calls",
2518 );
2519 }
2520 (
2521 ChatCompletionResponseMessage {
2522 role: "assistant",
2523 content: Some(parsed.content),
2524 reasoning_content: parsed.reasoning,
2525 tool_calls: None,
2526 },
2527 base_finish,
2528 )
2529}
2530
2531fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2536 let mut history = match &req.session_id {
2537 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2538 None => req.messages.clone(),
2539 };
2540 if req.json_object_mode() {
2541 inject_json_object_system_hint(&mut history);
2542 }
2543 history
2544}
2545
2546fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2547 const HINT: &str =
2548 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2549 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2550 match &mut sys.content {
2551 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2552 s.push_str("\n\n");
2553 s.push_str(HINT);
2554 }
2555 None => {
2556 sys.content = Some(MessageContent::Text(HINT.to_string()));
2557 }
2558 _ => {}
2559 }
2560 } else {
2561 messages.insert(
2562 0,
2563 ChatMessage {
2564 role: "system".to_string(),
2565 content: Some(MessageContent::Text(HINT.to_string())),
2566 tool_calls: None,
2567 tool_call_id: None,
2568 reasoning_content: None,
2569 },
2570 );
2571 }
2572}
2573
2574async fn chat_completions(
2575 State(state): State<Arc<AppState>>,
2576 headers: axum::http::HeaderMap,
2577 Json(req): Json<ChatCompletionRequest>,
2578) -> Response {
2579 let attribution = attribution::Attribution::from_headers(&headers);
2580 state
2581 .requests_total
2582 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2583 let started = std::time::Instant::now();
2584
2585 let request_id = frink_api::next_request_id();
2589 let stream = req.stream.unwrap_or(false);
2590
2591 let refusal = cache_admin::check_admission(&state)
2597 .err()
2598 .or_else(|| req.validate_supported_fields().err());
2599 if let Some(err) = refusal {
2600 state
2601 .request_errors_total
2602 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2603 let response = err.into_response();
2604 state.record_request(stats::Record {
2605 request_id: &request_id,
2606 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2607 model: state.active_model_name(),
2608 status: response.status().as_u16(),
2609 stream,
2610 duration_ms: started.elapsed().as_millis() as u64,
2611 usage: None,
2612 attribution: &attribution,
2613 });
2614 return response;
2615 }
2616
2617 let response = if stream {
2618 chat_completions_stream(
2619 Arc::clone(&state),
2620 req,
2621 request_id.clone(),
2622 started,
2623 attribution.clone(),
2624 )
2625 .await
2626 .into_response()
2627 } else {
2628 chat_completions_full(
2629 Arc::clone(&state),
2630 req,
2631 request_id.clone(),
2632 started,
2633 attribution.clone(),
2634 )
2635 .await
2636 .into_response()
2637 };
2638
2639 if response.status().is_client_error() || response.status().is_server_error() {
2640 state
2641 .request_errors_total
2642 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2643 state.record_request(stats::Record {
2647 request_id: &request_id,
2648 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2649 model: state.active_model_name(),
2652 status: response.status().as_u16(),
2653 stream,
2654 duration_ms: started.elapsed().as_millis() as u64,
2655 usage: None,
2656 attribution: &attribution,
2657 });
2658 }
2659 state.mark_request_finished();
2660
2661 response
2662}
2663
2664async fn chat_completions_full(
2665 state: Arc<AppState>,
2666 req: ChatCompletionRequest,
2667 request_id: String,
2668 started: std::time::Instant,
2669 attribution: attribution::Attribution,
2670) -> Result<Json<ChatCompletionResponse>, ApiError> {
2671 let tools_active = req.tools_active();
2672 let active = state.require_active()?;
2676 let history = resolve_history(&state, &req);
2677 let template = active.generative()?.chat_template();
2678 let kwargs = req.resolve_template_kwargs(&template);
2679 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2680 let mut params =
2688 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2689 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2690 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2691
2692 let (completion, cache_status) = if let Some(cached) = key
2693 .as_ref()
2694 .and_then(|key| lock_cache(&state.response_cache).get(key))
2695 {
2696 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2697 (cached, "hit")
2698 } else {
2699 let (chunks, finish, usage) = decode_task::buffered(
2700 decode_task::DecodeHandles::take(&state, &active)?,
2701 prompt.clone(),
2702 params,
2703 )
2704 .await?;
2705
2706 let completion = response_cache::CachedCompletion {
2707 content: chunks.concat(),
2708 finish,
2709 usage,
2710 };
2711 let cache_status = match key {
2717 Some(key) => match completion.clone().cacheable() {
2721 Some(cacheable) => {
2722 tracing::debug!("cache miss for key {}", key.digest());
2723 lock_cache(&state.response_cache).put(key, cacheable);
2724 "miss"
2725 }
2726 None => "skip",
2727 },
2728 None => "skip",
2729 };
2730 (completion, cache_status)
2731 };
2732 let content = completion.content;
2733
2734 if req.json_object_mode() {
2735 json_mode::validate_json_object_output(&content)?;
2736 }
2737
2738 if let Some(id) = &req.session_id {
2742 state.sessions.store_reply(
2743 id,
2744 ChatMessage {
2745 role: "assistant".to_string(),
2746 content: Some(MessageContent::Text(content.clone())),
2747 tool_calls: None,
2748 tool_call_id: None,
2749 reasoning_content: None,
2750 },
2751 );
2752 }
2753
2754 let (message, finish_reason) = build_response_message(
2755 content,
2756 if tools_active { &req.tools } else { &[] },
2757 output::OutputPosture::resolve_full(
2758 active.reasoning_format(),
2759 active.tool_call_format(),
2760 &prompt,
2761 ),
2762 completion.finish.as_str(),
2763 );
2764
2765 state.record_request(stats::Record {
2766 request_id: &request_id,
2767 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2768 model: Some(active.name().to_string()),
2771 status: 200,
2772 stream: false,
2773 duration_ms: started.elapsed().as_millis() as u64,
2774 usage: Some(&completion.usage),
2775 attribution: &attribution,
2776 });
2777
2778 Ok(Json(ChatCompletionResponse {
2779 id: request_id.clone(),
2780 request_id,
2781 object: "chat.completion",
2782 model: req.model,
2783 choices: vec![ChatCompletionChoice {
2784 index: 0,
2785 message,
2786 finish_reason,
2787 }],
2788 usage: completion.usage,
2789 frink_cache: cache_status,
2790 }))
2791}
2792
2793async fn chat_completions_stream(
2794 state: Arc<AppState>,
2795 req: ChatCompletionRequest,
2796 request_id: String,
2797 started: std::time::Instant,
2798 attribution: attribution::Attribution,
2799) -> Result<Response, ApiError> {
2800 let tools_active = req.tools_active();
2802 let active = state.require_active()?;
2806 let history = resolve_history(&state, &req);
2807 let template = active.generative()?.chat_template();
2808 let kwargs = req.resolve_template_kwargs(&template);
2809 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2810 let model_name = req.model.clone();
2811 let session_id = req.session_id.clone();
2812 let sessions = state.sessions.clone();
2813
2814 let model = Arc::clone(active.generative()?);
2815 let kv_pool = state.kv_pool.clone();
2816 let paged_kv = state.paged_kv.clone();
2817 let prefix_cache = state.prefix_cache.clone();
2818 let batcher = active.batcher.clone();
2819 let ceiling = active.ceiling.clone();
2820 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2821 let mut params =
2822 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2823 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2824 let stats_state = Arc::clone(&state);
2825 let served_model = active.name().to_string();
2828 let posture = output::OutputPosture::resolve_full(
2832 active.reasoning_format(),
2833 active.tool_call_format(),
2834 &prompt,
2835 );
2836 let offered_tools: Vec<ToolDef> = if tools_active {
2839 req.tools.clone()
2840 } else {
2841 Vec::new()
2842 };
2843
2844 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2849 params.cancel = Some(cancel_token.clone());
2850
2851 let overlap = true;
2860
2861 let slot = req
2865 .stream_resumable
2866 .unwrap_or(false)
2867 .then(|| state.streams.register(&request_id));
2868 let emitter = resume::Emitter::new(slot);
2869
2870 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2871 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2875 id: request_id.clone(),
2876 request_id: None,
2877 object: "chat.completion.chunk",
2878 model: model_name.clone(),
2879 choices: vec![ChatCompletionChunkChoice {
2880 index: 0,
2881 delta: ChatCompletionChunkDelta {
2882 role: None,
2883 content: None,
2884 reasoning_content: None,
2885 tool_calls: None,
2886 },
2887 finish_reason: None,
2888 }],
2889 usage: None,
2890 });
2891
2892 tokio::task::spawn_blocking(move || {
2893 let _cancel_guard = cancel_guard;
2896 let tx_chunks = tx.clone();
2897 let orphan_timeout = sse::orphan_timeout_from_env();
2902 let mut first = true;
2903 let head_request_id = request_id.clone();
2904 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
2914 Rc::new(RefCell::new(posture.reasoning_parser()));
2915 let emit_reasoning = Rc::clone(&stream_reasoning);
2916 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
2920 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
2921 );
2922 let emit_tools = Rc::clone(&stream_tools);
2923 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
2927 let emit_streamed_calls = Rc::clone(&streamed_calls);
2928 let result = run_generation_emit(
2929 &model,
2930 &prompt,
2931 ¶ms,
2932 kv_pool.as_ref(),
2933 paged_kv.as_ref(),
2934 prefix_cache.as_deref(),
2935 batcher.as_ref(),
2936 ceiling.as_deref(),
2937 metal_private_decode_gate.as_deref(),
2938 |chunk| {
2939 if !overlap || chunk.is_empty() {
2940 return;
2941 }
2942 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
2943 Some(parser) => {
2944 let delta = parser.push(chunk);
2945 (delta.reasoning, delta.content)
2946 }
2947 None => (String::new(), chunk.to_string()),
2948 };
2949 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
2953 Some(parser) => {
2954 let (text, calls) =
2955 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
2956 (text, calls)
2957 }
2958 None => (content, Vec::new()),
2959 };
2960 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
2963 return;
2964 }
2965 let role = if first { Some("assistant") } else { None };
2966 let request_id = first.then(|| head_request_id.clone());
2967 first = false;
2968 let payload = ChatCompletionChunk {
2969 id: head_request_id.clone(),
2970 request_id,
2971 object: "chat.completion.chunk",
2972 model: model_name.clone(),
2973 choices: vec![ChatCompletionChunkChoice {
2974 index: 0,
2975 delta: ChatCompletionChunkDelta {
2976 role,
2977 content: (!content.is_empty()).then_some(content),
2978 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
2979 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
2980 },
2981 finish_reason: None,
2982 }],
2983 usage: None,
2984 };
2985 if let Err(why) =
2994 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
2995 {
2996 if why == sse::SendFailure::Orphaned {
2997 tracing::warn!(
2998 "SSE stream {head_request_id} accepted nothing for the orphan \
2999 deadline; treating it as abandoned"
3000 );
3001 }
3002 if !emitter.is_resumable() {
3014 cancel_token.cancel();
3015 }
3016 }
3017 },
3018 );
3019
3020 let mut pending_request_id = first.then(|| request_id.clone());
3026
3027 match result {
3028 Ok((finish, usage, full_text)) => {
3029 if let Some(id) = &session_id {
3030 sessions.store_reply(
3031 id,
3032 ChatMessage {
3033 role: "assistant".to_string(),
3034 content: Some(MessageContent::Text(full_text.clone())),
3035 tool_calls: None,
3036 tool_call_id: None,
3037 reasoning_content: None,
3038 },
3039 );
3040 }
3041 let mut streamed_finish: Option<&'static str> = None;
3047 if overlap {
3048 let tail = stream_reasoning
3049 .borrow_mut()
3050 .as_mut()
3051 .map(|parser| parser.flush())
3052 .unwrap_or_default();
3053 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3054 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3055 let mut events = parser.push(&content);
3056 events.extend(parser.finish());
3057 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3058 content = text;
3059 tool_calls = calls;
3060 }
3061 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3062 let payload = ChatCompletionChunk {
3063 id: request_id.clone(),
3064 request_id: pending_request_id.take(),
3065 object: "chat.completion.chunk",
3066 model: model_name.clone(),
3067 choices: vec![ChatCompletionChunkChoice {
3068 index: 0,
3069 delta: ChatCompletionChunkDelta {
3070 role: None,
3071 content: (!content.is_empty()).then_some(content),
3072 reasoning_content: (!tail.reasoning.is_empty())
3073 .then_some(tail.reasoning),
3074 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3075 },
3076 finish_reason: None,
3077 }],
3078 usage: None,
3079 };
3080 let _ =
3081 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3082 }
3083 if streamed_calls.get() > 0 {
3084 streamed_finish = Some("tool_calls");
3085 }
3086 } else {
3087 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3090 let tool_calls: Vec<ToolCallDelta> = parsed
3091 .calls
3092 .iter()
3093 .enumerate()
3094 .map(|(index, call)| {
3095 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3096 })
3097 .collect();
3098 if !tool_calls.is_empty() {
3099 streamed_finish = Some("tool_calls");
3100 }
3101 if !tool_calls.is_empty()
3102 || !parsed.content.is_empty()
3103 || parsed.reasoning.is_some()
3104 {
3105 let payload = ChatCompletionChunk {
3106 id: request_id.clone(),
3107 request_id: pending_request_id.take(),
3108 object: "chat.completion.chunk",
3109 model: model_name.clone(),
3110 choices: vec![ChatCompletionChunkChoice {
3111 index: 0,
3112 delta: ChatCompletionChunkDelta {
3113 role: Some("assistant"),
3114 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3115 .then(|| parsed.content.clone()),
3116 reasoning_content: parsed.reasoning.clone(),
3117 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3118 },
3119 finish_reason: None,
3120 }],
3121 usage: None,
3122 };
3123 let _ =
3124 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3125 }
3126 }
3127 let final_finish_reason = match streamed_finish {
3131 Some(reason) if finish.as_str() != "length" => reason,
3132 _ => finish.as_str(),
3133 };
3134 let final_payload = ChatCompletionChunk {
3135 id: request_id.clone(),
3136 request_id: pending_request_id.take(),
3137 object: "chat.completion.chunk",
3138 model: model_name,
3139 choices: vec![ChatCompletionChunkChoice {
3140 index: 0,
3141 delta: ChatCompletionChunkDelta {
3142 role: None,
3143 content: None,
3144 reasoning_content: None,
3145 tool_calls: None,
3146 },
3147 finish_reason: Some(final_finish_reason),
3148 }],
3149 usage: Some(usage.clone()),
3150 };
3151 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3152 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3153 stats_state.record_request(stats::Record {
3158 request_id: &request_id,
3159 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3160 model: Some(served_model.clone()),
3161 status: 200,
3162 stream: true,
3163 duration_ms: started.elapsed().as_millis() as u64,
3164 usage: Some(&usage),
3165 attribution: &attribution,
3166 });
3167 }
3168 Err(e) => {
3169 tracing::warn!("decode error on streamed request {request_id}: {e}");
3170 stats_state.record_request(stats::Record {
3176 request_id: &request_id,
3177 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3178 model: Some(served_model.clone()),
3179 status: 500,
3180 stream: true,
3181 duration_ms: started.elapsed().as_millis() as u64,
3182 usage: None,
3183 attribution: &attribution,
3184 });
3185 let payload = ChatCompletionChunk {
3186 id: request_id.clone(),
3187 request_id: pending_request_id.take(),
3188 object: "chat.completion.chunk",
3189 model: model_name,
3190 choices: vec![ChatCompletionChunkChoice {
3191 index: 0,
3192 delta: ChatCompletionChunkDelta {
3193 role: Some("assistant"),
3194 content: Some(format!("[error: {e}]")),
3195 reasoning_content: None,
3196 tool_calls: None,
3197 },
3198 finish_reason: Some("stop"),
3199 }],
3200 usage: None,
3201 };
3202 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3203 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3204 }
3205 }
3206 drop(emitter);
3210 });
3211
3212 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3213 Ok((
3236 [(
3237 axum::http::HeaderName::from_static("x-accel-buffering"),
3238 axum::http::HeaderValue::from_static("no"),
3239 )],
3240 Sse::new(stream),
3241 )
3242 .into_response())
3243}
3244
3245fn protected_routes() -> Router<Arc<AppState>> {
3276 use frink_api::routes;
3277
3278 Router::new()
3279 .route(routes::V1_MODELS, get(list_models))
3280 .route(routes::V1_RESPONSES, post(responses::responses))
3284 .route(
3285 &axum_path(routes::V1_RESPONSE),
3286 get(responses::responses_get),
3287 )
3288 .route(
3289 &axum_path(routes::V1_RESPONSE_CANCEL),
3290 post(responses::responses_cancel),
3291 )
3292 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3293 .route(routes::V1_STATS, get(serving_stats))
3294 .route(routes::V1_REQUESTS, get(recent_requests))
3295 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3296 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3297 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3298 .route(
3299 routes::LORA_ADAPTERS,
3300 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3301 )
3302 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3303 .route(routes::V1_CANCEL, post(cancel_generation))
3307 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3312 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3313 .route(routes::V1_MESSAGES, post(anthropic::messages))
3314 .route(
3315 routes::V1_MESSAGES_COUNT_TOKENS,
3316 post(anthropic::count_tokens),
3317 )
3318 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3319 .route(routes::COMPLETION, post(completion::completion))
3324 .route(routes::COMPLETIONS, post(completion::completion))
3325 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3326 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3327 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3333 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3334 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3335 .route(routes::V1_RERANK, post(rerank::rerank))
3339 .route(routes::RERANK, post(rerank::rerank))
3340 .route(routes::CACHE_STATS, get(cache_stats))
3341 .route(routes::METRICS, get(metrics))
3342 .route(routes::ADMIN_MODELS, get(admin::models))
3347 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3348 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3349 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3350 .route(routes::ADMIN_TASKS, get(admin::tasks))
3351 .route(&admin::cancel_route(), post(admin::cancel_task))
3352 .route(routes::ADMIN_STATS, get(admin::stats))
3353 .merge(conversations::router())
3357}
3358
3359fn axum_path(template: &str) -> String {
3360 let mut out = String::with_capacity(template.len());
3361 let mut rest = template;
3362 while let Some(open) = rest.find('{') {
3363 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3364 break;
3365 };
3366 out.push_str(&rest[..open]);
3367 out.push(':');
3368 out.push_str(&rest[open + 1..close]);
3369 rest = &rest[close + 1..];
3370 }
3371 out.push_str(rest);
3372 out
3373}
3374
3375async fn cancel_generation(
3384 State(state): State<Arc<AppState>>,
3385 Json(req): Json<frink_api::CancelGenerationRequest>,
3386) -> Response {
3387 let cancelled = state.cancels.cancel(&req.request_id);
3388 let status = if cancelled {
3389 StatusCode::OK
3390 } else {
3391 StatusCode::NOT_FOUND
3392 };
3393 let detail = if cancelled {
3394 "the generation was asked to stop; it ends at its next token".to_string()
3395 } else {
3396 "no generation with that request_id is running -- it has already \
3397 finished, was never issued, or was served by a path that does \
3398 not register for cancellation"
3399 .to_string()
3400 };
3401 (
3402 status,
3403 Json(frink_api::CancelGenerationResponse {
3404 request_id: req.request_id,
3405 cancelled,
3406 detail,
3407 }),
3408 )
3409 .into_response()
3410}
3411
3412type Activated = (
3416 Loaded,
3417 Option<serving::batch::ContinuousBatcher>,
3418 Option<Arc<budget::ContextCeiling>>,
3419);
3420
3421fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3439 let mut batcher = serving::batch::BatcherConfig::from_env();
3440 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3441 return batcher;
3444 }
3445 let Some(path) = path else {
3446 return batcher;
3447 };
3448 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3454 let Some((priced, gguf_ctx, source)) = priced else {
3455 return batcher;
3456 };
3457 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3458 tracing::warn!(
3462 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3463 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3464 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3465 admit on a number you choose.",
3466 priced.weights_bytes,
3467 priced.device_budget_bytes,
3468 );
3469 return batcher;
3470 };
3471 tracing::info!("{source}");
3472 tracing::info!("{}", derived.fit);
3473 let adopted = budget::apply_derived(&mut batcher, &derived);
3474 if adopted.max_context {
3475 tracing::info!(
3476 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3477 override with FRINK_CB_MAX_CONTEXT",
3478 derived.max_context
3479 );
3480 }
3481 if adopted.kv_blocks {
3482 tracing::info!(
3483 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3484 derived.kv_blocks,
3485 batcher.kv_block_size
3486 );
3487 }
3488 if let Some(narrowed) = adopted.max_context_narrowed {
3489 tracing::info!(
3490 "per-request context ceiling narrowed to {narrowed} token positions: the whole KV ledger is {} blocks x {} positions, so a longer request could never be admitted",
3491 batcher.kv_blocks.unwrap_or_default(),
3492 batcher.kv_block_size
3493 );
3494 }
3495 batcher
3496}
3497
3498pub(crate) fn activate_loaded_model(
3507 loaded: model::LoadedModel,
3508 enable_continuous_batching: bool,
3509 path: Option<&str>,
3510 paged_kv: Option<&generate::PagedKvConfig>,
3511) -> Activated {
3512 match loaded {
3513 model::LoadedModel::Gguf(g) => {
3514 let decoder = Arc::new(g.decoder);
3515 let tokenizer = Arc::new(g.tokenizer);
3516 let config = price_batcher_config(path);
3517 let shape =
3523 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3524 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3525 let batcher = if enable_continuous_batching {
3526 tracing::info!(
3527 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3528 (stop sequences use the same pending-buffer trim as the private generate loop)"
3529 );
3530 let tok = Arc::clone(&tokenizer);
3531 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3532 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3533 Arc::clone(&decoder),
3534 decode,
3535 config,
3536 Arc::clone(&ceiling),
3537 paged_kv.cloned(),
3538 ))
3539 } else {
3540 None
3541 };
3542 (
3543 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3544 decoder,
3545 tokenizer,
3546 stop_tokens: g.stop_tokens,
3547 bos_id: g.bos_id,
3548 is_synthetic: g.is_synthetic,
3549 chat_template: g.chat_template,
3550 }))),
3551 batcher,
3552 Some(ceiling),
3553 )
3554 }
3555 model::LoadedModel::Kimi(k) => (
3556 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3557 engine: k.engine,
3558 tokenizer: k.tokenizer,
3559 stop_tokens: k.stop_tokens,
3560 chat_template: k.chat_template,
3561 }))),
3562 None,
3563 None,
3564 ),
3565 model::LoadedModel::Mla(m) => (
3566 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3567 engine: m.engine,
3568 tokenizer: m.tokenizer,
3569 stop_tokens: m.stop_tokens,
3570 bos_id: m.bos_id,
3571 name: m.name,
3572 chat_template: m.chat_template,
3573 }))),
3574 None,
3575 None,
3576 ),
3577 model::LoadedModel::Gemma4(m) => (
3578 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3579 engine: m.engine,
3580 tokenizer: m.tokenizer,
3581 stop_tokens: m.stop_tokens,
3582 bos_id: m.bos_id,
3583 name: m.name,
3584 chat_template: m.chat_template,
3585 }))),
3586 None,
3587 None,
3588 ),
3589 model::LoadedModel::Glm52(g) => (
3590 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3591 engine: g.engine,
3592 tokenizer: g.tokenizer,
3593 stop_tokens: g.stop_tokens,
3594 bos_id: g.bos_id,
3595 name: g.name,
3596 chat_template: g.chat_template,
3597 }))),
3598 None,
3599 None,
3600 ),
3601 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3606 }
3607}
3608
3609struct StartupModels {
3616 loaded: model::LoadedModel,
3617 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3618}
3619
3620fn continuous_batching_env() -> Option<bool> {
3621 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3622 .ok()
3623 .map(|v| v.trim().to_ascii_lowercase())
3624 .as_deref()
3625 {
3626 None => None,
3627 Some("1" | "true" | "yes" | "on") => Some(true),
3628 Some("0" | "false" | "no" | "off") => Some(false),
3629 _ => None,
3630 }
3631}
3632
3633fn metal_private_decode_active() -> bool {
3634 #[cfg(feature = "metal")]
3635 {
3636 BUILT_WITH_METAL
3637 && frink_metal::attn::metal_attn_enabled()
3638 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3639 }
3640 #[cfg(not(feature = "metal"))]
3641 {
3642 false
3643 }
3644}
3645
3646fn continuous_batching_compatible(
3647 loaded: &model::LoadedModel,
3648 kv_pool: &Option<generate::KvPoolConfig>,
3649 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3650 paged_kv: &Option<generate::PagedKvConfig>,
3651) -> bool {
3652 matches!(loaded, model::LoadedModel::Gguf(_))
3653 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3654}
3655
3656fn resolve_continuous_batching_enabled(
3657 loaded: &model::LoadedModel,
3658 kv_pool: &Option<generate::KvPoolConfig>,
3659 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3660 paged_kv: &Option<generate::PagedKvConfig>,
3661) -> bool {
3662 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3663 return false;
3664 }
3665 match continuous_batching_env() {
3666 Some(true) => true,
3667 Some(false) => false,
3668 None => metal_private_decode_active(),
3669 }
3670}
3671
3672fn acquire_metal_private_decode_gate(
3673 gate: Option<&std::sync::Mutex<()>>,
3674 used_batcher: bool,
3675) -> Option<std::sync::MutexGuard<'_, ()>> {
3676 if used_batcher {
3677 None
3678 } else {
3679 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3680 }
3681}
3682
3683fn build_app_state(
3684 models: StartupModels,
3685 kv_pool: Option<generate::KvPoolConfig>,
3686 paged_kv: Option<generate::PagedKvConfig>,
3687 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3688 enable_continuous_batching: bool,
3689 mcp: Option<mcp::LoadedMcpConfig>,
3690 detection: Arc<health::Detection>,
3691) -> AppState {
3692 let StartupModels { loaded, embedding } = models;
3693 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3694 let (loaded, batcher, ceiling) = activate_loaded_model(
3695 loaded,
3696 enable_continuous_batching,
3697 configured_path.as_deref(),
3698 paged_kv.as_ref(),
3699 );
3700 let id = startup_model_id();
3706 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3707 {
3708 None
3709 } else {
3710 tracing::info!(
3711 "Metal private-loop decode will serialize concurrent requests until \
3712 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3713 );
3714 Some(Arc::new(std::sync::Mutex::new(())))
3715 };
3716 AppState {
3717 embedding,
3718 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3719 id,
3720 loaded,
3721 batcher,
3722 ceiling,
3723 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3724 }))),
3725 paged_kv,
3726 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3727 tasks: Arc::new(tasks::TaskRegistry::new()),
3728 cancels: Arc::new(cancel::CancelRegistry::new()),
3729 stats: stats::Stats::new(),
3730 streams: resume::StreamRegistry::new(),
3731 model_dir: admin::model_dirs().into_iter().next(),
3732 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3733 kv_pool,
3734 prefix_cache,
3735 sessions: session::SessionStore::new(),
3736 requests_total: std::sync::atomic::AtomicU64::new(0),
3737 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3738 started_at: std::time::Instant::now(),
3739 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3740 detection,
3741 mcp,
3742 continuous_batching_enabled: enable_continuous_batching,
3743 metal_private_decode_gate,
3744 loading_model: Mutex::new(None),
3745 last_load_error: Mutex::new(None),
3746 serving: Mutex::new(crate::stats::ServingStats::default()),
3747 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3748 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3749 started_unix: unix_now(),
3750 }
3751}
3752
3753fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3761 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3762 return Ok(None);
3763 };
3764 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3765 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3766 tracing::info!(
3767 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3768 model.name(),
3769 model.architecture(),
3770 model.n_embd(),
3771 model.pooling_type().name(),
3772 model.n_ctx_train(),
3773 );
3774 Ok(Some(Arc::new(model)))
3775}
3776
3777fn unix_now() -> u64 {
3782 std::time::SystemTime::now()
3783 .duration_since(std::time::UNIX_EPOCH)
3784 .map(|d| d.as_secs())
3785 .unwrap_or(0)
3786}
3787
3788fn startup_model_id() -> Option<String> {
3793 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3794 let configured = std::fs::canonicalize(&configured).ok()?;
3795 admin::discover(&admin::model_dirs())
3796 .into_iter()
3797 .find(|d| {
3798 std::fs::canonicalize(&d.path)
3799 .map(|p| p == configured)
3800 .unwrap_or(false)
3801 })
3802 .map(|d| d.id)
3803}
3804
3805fn init_cpu_pool() {
3814 match frink_core::threads::init_cpu_pool() {
3815 Some(n) => eprintln!(
3816 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
3817 frink_core::threads::perf_core_count()
3818 ),
3819 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
3820 }
3821}
3822
3823fn announce_ready(addr: SocketAddr, scheme: &str) {
3836 use std::io::Write;
3837 let ready =
3838 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3839 let mut stdout = std::io::stdout().lock();
3840 let _ = writeln!(stdout, "{}", ready.to_line());
3841 let _ = stdout.flush();
3842}
3843
3844async fn shutdown_signal(exit_on_stdin_close: bool) {
3856 if !exit_on_stdin_close {
3857 std::future::pending::<()>().await;
3858 return;
3859 }
3860 let _ = tokio::task::spawn_blocking(|| {
3861 use std::io::Read;
3862 let mut sink = [0u8; 256];
3863 let mut stdin = std::io::stdin().lock();
3864 loop {
3865 match stdin.read(&mut sink) {
3866 Ok(0) => break,
3868 Ok(_) => continue,
3870 Err(e) => {
3871 tracing::warn!("stdin read failed ({e}); treating it as closed");
3872 break;
3873 }
3874 }
3875 }
3876 })
3877 .await;
3878 tracing::info!("stdin closed; shutting down");
3879}
3880
3881fn tokio_worker_threads() -> usize {
3886 std::env::var("FRINK_TOKIO_WORKERS")
3887 .ok()
3888 .and_then(|v| v.trim().parse::<usize>().ok())
3889 .filter(|n| *n > 0)
3890 .unwrap_or(2)
3891}
3892
3893fn install_ring_crypto_provider() {
3910 let _ = rustls::crypto::ring::default_provider().install_default();
3911}
3912
3913pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3919 if args.list_devices {
3920 print_available_devices();
3921 return Ok(());
3922 }
3923 apply_cli_overrides(&args)?;
3924
3925 let _instance = {
3929 use frink_core::instance::{register, InstancePolicy};
3930 let policy = if args.allow_multiple_instances {
3931 InstancePolicy::Multi
3932 } else {
3933 InstancePolicy::from_env_or(InstancePolicy::Single)
3934 };
3935 let model = std::env::var("FRINK_MODEL_PATH").ok();
3936 register(
3937 "server",
3938 model.as_deref(),
3939 frink_core::instance::current_backend(),
3940 policy,
3941 )
3942 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
3943 };
3944
3945 let journal = journal::Journal::from_env();
3946 eprintln!(
3947 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
3948 journal.path()
3949 );
3950 journal.append(&journal::Record::session_start(
3951 env!("CARGO_PKG_VERSION"),
3952 std::process::id(),
3953 ));
3954 journal::install_panic_hook(journal.clone());
3955
3956 let mcp_config_path = args.mcp_config.clone();
3957 let exit_on_stdin_close = args.exit_on_stdin_close
3958 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
3959 .map(|v| v == "1")
3960 .unwrap_or(false);
3961
3962 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
3966 init_cpu_pool();
3967
3968 let runtime = tokio::runtime::Builder::new_multi_thread()
3969 .worker_threads(tokio_worker_threads())
3970 .enable_all()
3971 .build()?;
3972 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
3973
3974 let reason = match &result {
3975 Ok(()) => "normal".to_string(),
3976 Err(e) => e.to_string(),
3977 };
3978 journal.append(&journal::Record::session_exit(reason));
3979
3980 runtime.shutdown_background();
3985
3986 result
3987}
3988
3989async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
3990 let _ = tracing_subscriber::fmt::try_init();
3996
3997 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4007 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4008 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4009 .map(|v| v == "1")
4010 .unwrap_or(false);
4011 if let Err(msg) =
4012 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4013 {
4014 anyhow::bail!(msg);
4015 }
4016
4017 let embedding_model = load_embedding_model()?;
4024
4025 let mut loaded = model::load()?;
4026 match &loaded {
4027 model::LoadedModel::Gguf(g) => tracing::info!(
4028 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4029 g.decoder.config.name,
4030 g.is_synthetic,
4031 g.tokenizer.kind()
4032 ),
4033 model::LoadedModel::Kimi(k) => tracing::info!(
4034 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4035 k.tokenizer.vocab_size()
4036 ),
4037 model::LoadedModel::Mla(m) => tracing::info!(
4038 "loaded MLA GGUF '{}' (tokenizer={})",
4039 m.name,
4040 m.tokenizer.kind()
4041 ),
4042 model::LoadedModel::Gemma4(m) => tracing::info!(
4043 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4044 m.name,
4045 m.tokenizer.kind()
4046 ),
4047 model::LoadedModel::Glm52(g) => tracing::info!(
4048 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4049 g.name,
4050 g.tokenizer.kind()
4051 ),
4052 model::LoadedModel::Encoder(_) => {}
4055 }
4056 let metal_default_moe_budget = {
4064 #[cfg(feature = "metal")]
4065 {
4066 frink_core::metal_dense_enabled()
4067 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4068 }
4069 #[cfg(not(feature = "metal"))]
4070 {
4071 false
4072 }
4073 };
4074 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4075 let budget: u64 = budget_str
4076 .parse()
4077 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4078 match &mut loaded {
4079 model::LoadedModel::Gguf(g) => {
4080 tracing::info!(
4081 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4082 (CUDA and/or Metal matvecs when built with the matching feature)"
4083 );
4084 g.decoder.gpu_vram_budget_bytes = Some(budget);
4085 }
4086 model::LoadedModel::Kimi(_) => {
4087 tracing::warn!(
4088 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4089 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4090 );
4091 }
4092 model::LoadedModel::Mla(_) => {
4093 tracing::warn!(
4094 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4095 FFN path only today; ignoring expert VRAM budget"
4096 );
4097 }
4098 model::LoadedModel::Gemma4(_) => {
4099 tracing::warn!(
4100 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4101 ignoring expert VRAM budget"
4102 );
4103 }
4104 model::LoadedModel::Glm52(_) => {
4105 tracing::warn!(
4106 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4107 GPU expert placement not wired yet; ignoring"
4108 );
4109 }
4110 model::LoadedModel::Encoder(_) => {
4111 tracing::warn!(
4112 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4113 it has no routed experts to place; ignoring"
4114 );
4115 }
4116 }
4117 } else if metal_default_moe_budget {
4118 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4122 if let model::LoadedModel::Gguf(g) = &mut loaded {
4123 tracing::info!(
4124 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4125 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4126 );
4127 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4128 }
4129 }
4130 #[cfg(feature = "cuda")]
4131 {
4132 if frink_core::cuda_dense_enabled() {
4133 tracing::info!(
4134 "CUDA dense matvec enabled for WeightMatrix::apply \
4135 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4136 );
4137 } else {
4138 tracing::info!(
4139 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4140 );
4141 }
4142 }
4143 #[cfg(feature = "metal")]
4144 {
4145 if frink_core::metal_dense_enabled() {
4146 tracing::info!(
4147 "Metal dense matvec enabled for WeightMatrix::apply \
4148 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4149 );
4150 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4151 Some("1") | Some("true") | Some("on") | Some("attn") => {
4152 tracing::info!(
4153 "Metal fused attention requested (FRINK_METAL_ATTN): \
4154 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4155 );
4156 }
4157 _ => {}
4158 }
4159 tracing::info!(
4160 "Metal greedy GPU argmax: temperature<=0 folds \
4161 final_norm+lm_head+argmax into the dense stack"
4162 );
4163 } else {
4164 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4165 }
4166 }
4167 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4178 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4179 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4180 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4181 panic!(
4182 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4183 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4184 );
4185 }
4186 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4187 (Ok(blocks), Ok(block_size), Err(_)) => {
4188 let total_blocks: usize = blocks
4189 .parse()
4190 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4191 let block_size: usize = block_size
4192 .parse()
4193 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4194 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4200 .ok()
4201 .map(|v| {
4202 v.parse()
4203 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4204 })
4205 .unwrap_or(0);
4206 tracing::info!(
4207 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4208 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4209 );
4210 Some(generate::KvPoolConfig {
4211 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4212 queue_wait: Duration::from_millis(queue_wait_ms),
4213 })
4214 }
4215 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4216 let block_size: usize = block_size
4217 .parse()
4218 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4219 let budget: u64 = byte_budget
4220 .parse()
4221 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4222 let cfg = match &loaded {
4223 model::LoadedModel::Gguf(g) => &g.decoder.config,
4224 model::LoadedModel::Kimi(_)
4225 | model::LoadedModel::Mla(_)
4226 | model::LoadedModel::Gemma4(_)
4227 | model::LoadedModel::Glm52(_)
4228 | model::LoadedModel::Encoder(_) => {
4229 panic!(
4230 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4231 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4232 );
4233 }
4234 };
4235 let bytes_per_block = block_size
4236 * cfg.kv_heads_all_layers()
4237 * (cfg.head_dim + cfg.v_head_dim())
4238 * std::mem::size_of::<f32>();
4239 assert!(
4240 bytes_per_block > 0,
4241 "derived KV block byte size must be positive (check model config and block size)"
4242 );
4243 let total_blocks = (budget as usize / bytes_per_block).max(1);
4244 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4245 .ok()
4246 .map(|v| {
4247 v.parse()
4248 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4249 })
4250 .unwrap_or(0);
4251 tracing::info!(
4252 "KV cache block pool enabled from byte budget: {budget} bytes / \
4253 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4254 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4255 cfg.n_layers
4256 );
4257 Some(generate::KvPoolConfig {
4258 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4259 queue_wait: Duration::from_millis(queue_wait_ms),
4260 })
4261 }
4262 (Err(_), Err(_), Err(_)) => None,
4263 (Err(_), Ok(_), Err(_)) => panic!(
4264 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4265 (or unset all three to disable KV cache pooling)"
4266 ),
4267 (Ok(_), Ok(_), Ok(_)) => {
4268 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4269 }
4270 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4271 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4272 set together (or neither, to disable KV cache pooling)"
4273 ),
4274 };
4275 let paged_kv = match (
4281 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4282 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4283 ) {
4284 (Ok(blocks), Ok(block_size)) => {
4285 assert!(
4286 kv_pool.is_none(),
4287 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4288 mutually exclusive: both bound the same KV memory, by different means. \
4289 Set one."
4290 );
4291 let blocks_per_layer: usize = blocks
4305 .parse()
4306 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4307 let block_size: usize = block_size
4308 .parse()
4309 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4310 let gguf = match &loaded {
4311 model::LoadedModel::Gguf(g) => g,
4312 _ => panic!(
4313 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4314 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4315 ),
4316 };
4317 let cfg = &gguf.decoder.config;
4318 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4319 .ok()
4320 .map(|v| {
4321 v.parse()
4322 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4323 })
4324 .unwrap_or(0);
4325 tracing::info!(
4326 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4327 layer across {} layers, shared by all concurrent requests, \
4328 {queue_wait_ms}ms admission queue wait",
4329 cfg.n_layers
4330 );
4331 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4336 block_size,
4337 ))));
4338 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4347 crate::policy::parser::ToolCallFormat::infer(
4348 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4349 )
4350 .opener(),
4351 |text| {
4352 gguf.tokenizer
4353 .encode(text, SpecialTokens::Parse)
4354 .into_iter()
4355 .map(|t| t as u32)
4356 .collect()
4357 },
4358 );
4359 if let Some(id) = anchor_token {
4360 tracing::info!(
4361 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4362 window stops short of where its next turn rejoins"
4363 );
4364 }
4365 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4366 .ok()
4367 .map(|v| {
4368 v.parse()
4369 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4370 })
4371 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4372 if let Some(window) = cfg.uniform_sliding_window() {
4373 tracing::info!(
4374 "Paged KV window slide enabled: every layer slides by {window} every \
4375 {slide_interval} decode steps, so a request holds its prompt and a \
4376 window rather than its whole context"
4377 );
4378 } else if cfg.kv_block_window().is_some() {
4379 tracing::info!(
4380 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4381 and a page group holds one block in every layer"
4382 );
4383 }
4384 Some(generate::PagedKvConfig {
4385 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4388 queue_wait: Duration::from_millis(queue_wait_ms),
4389 radix,
4390 anchor_token,
4391 slide_interval,
4392 })
4393 }
4394 (Err(_), Err(_)) => None,
4395 _ => panic!(
4396 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4397 (or neither, to disable paged KV)"
4398 ),
4399 };
4400 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4407 let max_entries: usize = v
4408 .parse()
4409 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4410 if kv_pool.is_some() {
4411 tracing::warn!(
4412 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4413 caching will never be consulted while a KV pool is configured"
4414 );
4415 }
4416 assert!(
4423 paged_kv.is_none(),
4424 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4425 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4426 produce, so the cache could never hit. Set one."
4427 );
4428 tracing::info!(
4429 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4430 all requests"
4431 );
4432 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4433 });
4434 if matches!(
4435 loaded,
4436 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4437 ) && (kv_pool.is_some() || prefix_cache.is_some())
4438 {
4439 tracing::warn!(
4440 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4441 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4442 frink_models::engine's module docs"
4443 );
4444 }
4445 let enable_cb =
4446 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4447 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4448 tracing::info!(
4449 "continuous batching enabled by default on Metal for safe parallel serving \
4450 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4451 );
4452 }
4453 if continuous_batching_env() == Some(true)
4454 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4455 && (kv_pool.is_some() || prefix_cache.is_some())
4456 {
4457 tracing::warn!(
4458 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4459 (those modes keep the private generate path)"
4460 );
4461 }
4462 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4463 if let Ok(chunk) = n.parse::<usize>() {
4464 if chunk > 0 {
4465 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4466 }
4467 }
4468 }
4469 if matches!(
4470 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4471 Some("1")
4472 ) {
4473 tracing::warn!(
4474 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4475 (minimal spill; full layer offload still planned)"
4476 );
4477 }
4478
4479 let mcp = match mcp_config_path {
4480 Some(path) => {
4481 let loaded = mcp::load_mcp_config(&path)?;
4482 tracing::info!(
4483 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4484 loaded.path,
4485 loaded.servers.len()
4486 );
4487 Some(loaded)
4488 }
4489 None => None,
4490 };
4491
4492 let detection = health::Detection::spawn();
4496
4497 let state = Arc::new(build_app_state(
4498 StartupModels {
4499 loaded,
4500 embedding: embedding_model,
4501 },
4502 kv_pool,
4503 paged_kv,
4504 prefix_cache,
4505 enable_cb,
4506 mcp,
4507 detection,
4508 ));
4509
4510 use frink_api::routes;
4514
4515 let public = Router::new().route(routes::HEALTH, get(health));
4520
4521 let mut protected = protected_routes();
4522
4523 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4528 tracing::info!("API key auth enabled");
4529 let auth = limits::AuthConfig {
4530 api_key: Arc::new(key),
4531 };
4532 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4533 auth,
4534 limits::require_api_key,
4535 ));
4536 }
4537 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4538 let rpm: u32 = rpm
4539 .parse()
4540 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4541 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4542 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4543 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4544 limiter,
4545 limits::rate_limit,
4546 ));
4547 }
4548 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4556 let origins = security::parse_cors_origins(&spec)
4557 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4558 tracing::info!(
4559 "CORS enabled: {} allow-listed origin(s) ({})",
4560 origins.len(),
4561 spec
4562 );
4563 let cors = tower_http::cors::CorsLayer::new()
4564 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4565 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4566 .allow_headers([
4567 axum::http::header::CONTENT_TYPE,
4568 axum::http::header::AUTHORIZATION,
4569 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4575 axum::http::HeaderName::from_static("last-event-id"),
4579 ]);
4580 protected = protected.route_layer(cors);
4581 }
4582
4583 let app = public
4587 .merge(protected)
4588 .layer(axum::middleware::from_fn(limits::retry_after))
4589 .with_state(state);
4590
4591 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4597 install_ring_crypto_provider();
4598 match tls_paths {
4603 Some(paths) => {
4604 let config =
4605 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4606 .await
4607 .map_err(|e| {
4608 anyhow::anyhow!(
4609 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4610 paths.cert,
4611 paths.key
4612 )
4613 })?;
4614 let socket_addr: std::net::SocketAddr = addr
4615 .parse()
4616 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4617 let listener = std::net::TcpListener::bind(socket_addr)?;
4618 listener.set_nonblocking(true)?;
4626 let bound = listener.local_addr()?;
4627 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4628 announce_ready(bound, "https");
4629
4630 let handle = axum_server::Handle::new();
4631 let shutdown_handle = handle.clone();
4632 tokio::spawn(async move {
4633 shutdown_signal(exit_on_stdin_close).await;
4634 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4635 });
4636 axum_server::from_tcp_rustls(listener, config)?
4637 .handle(handle)
4638 .serve(app.into_make_service())
4639 .await?;
4640 }
4641 None => {
4642 let listener = tokio::net::TcpListener::bind(&addr).await?;
4643 let bound = listener.local_addr()?;
4644 tracing::info!("frink-server listening on {bound}");
4645 announce_ready(bound, "http");
4646 axum::serve(listener, app)
4647 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4648 .await?;
4649 }
4650 }
4651 Ok(())
4652}
4653
4654#[cfg(test)]
4655pub(crate) mod tests {
4656 use super::*;
4657 use frink_models::config::test_dense_fixture;
4658
4659 #[test]
4660 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4661 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4662 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4663 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4664 assert_eq!(parsed.port, 51999);
4665 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4666 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4668 }
4669
4670 fn test_model() -> Model {
4671 let cfg = test_dense_fixture();
4675 Model::Gguf(GgufModel {
4676 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4677 tokenizer: Arc::new(ServerTokenizer::Byte),
4678 stop_tokens: StopTokens::default(),
4679 bos_id: None,
4680 is_synthetic: true,
4681 chat_template: chat_template::PromptTemplate::plain(),
4682 })
4683 }
4684
4685 fn greedy_params(max_tokens: usize) -> GenerationParams {
4686 GenerationParams {
4687 reasoning: None,
4688 max_tokens,
4689 sampling: SamplingParams::default(),
4690 seed: 1,
4691 stop: Vec::new(),
4692 stop_token_ids: Vec::new(),
4693 json_object: false,
4694 grammar: None,
4695 cancel: None,
4696 ignore_eos: false,
4697 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4698 lora: None,
4699 }
4700 }
4701
4702 fn test_model_full_byte_vocab() -> Model {
4706 test_model_full_byte_vocab_with_eos(None)
4707 }
4708
4709 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4716 let mut cfg = test_dense_fixture();
4717 cfg.vocab_size = 256;
4718 Model::Gguf(GgufModel {
4719 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4720 tokenizer: Arc::new(ServerTokenizer::Byte),
4721 stop_tokens: StopTokens::from_eos(eos),
4722 bos_id: None,
4723 is_synthetic: true,
4724 chat_template: chat_template::PromptTemplate::plain(),
4725 })
4726 }
4727
4728 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4732 AppState {
4733 embedding: None,
4734 paged_kv: None,
4735 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4736 id: None,
4737 loaded: Loaded::Generative(Arc::new(model)),
4738 batcher: None,
4739 ceiling: None,
4740 checkpoint_path: None,
4741 }))),
4742 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4743 tasks: Arc::new(tasks::TaskRegistry::new()),
4744 cancels: Arc::new(cancel::CancelRegistry::new()),
4745 stats: stats::Stats::new(),
4746 streams: resume::StreamRegistry::new(),
4747 model_dir: None,
4748 response_cache: Mutex::new(response_cache),
4749 kv_pool: None,
4750 prefix_cache: None,
4751 sessions: session::SessionStore::new(),
4752 requests_total: std::sync::atomic::AtomicU64::new(0),
4753 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4754 started_at: std::time::Instant::now(),
4755 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4756 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4757 mcp: None,
4758 continuous_batching_enabled: false,
4759 metal_private_decode_gate: None,
4760 loading_model: Mutex::new(None),
4761 last_load_error: Mutex::new(None),
4762 serving: Mutex::new(crate::stats::ServingStats::default()),
4763 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4764 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4765 started_unix: unix_now(),
4766 }
4767 }
4768
4769 fn test_app() -> Router {
4778 test_app_with_state(Arc::new(test_state(
4779 test_model_full_byte_vocab(),
4780 ResponseCache::new(1000, Duration::from_secs(3600)),
4781 )))
4782 }
4783
4784 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4787 Router::new()
4797 .route(frink_api::routes::HEALTH, get(health))
4798 .merge(protected_routes())
4799 .with_state(state)
4800 }
4801
4802 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4803 let mut cfg = test_dense_fixture();
4804 cfg.name = name;
4805 cfg.vocab_size = vocab_size;
4806 Model::Gguf(GgufModel {
4807 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4808 tokenizer: Arc::new(ServerTokenizer::Byte),
4809 stop_tokens: StopTokens::default(),
4810 bos_id: None,
4811 is_synthetic: true,
4812 chat_template: chat_template::PromptTemplate::plain(),
4813 })
4814 }
4815
4816 fn model_with_template(name: &'static str, source: &str) -> Model {
4820 let mut cfg = test_dense_fixture();
4821 cfg.name = name;
4822 cfg.vocab_size = 256;
4823 Model::Gguf(GgufModel {
4824 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4825 tokenizer: Arc::new(ServerTokenizer::Byte),
4826 stop_tokens: StopTokens::default(),
4827 bos_id: None,
4828 is_synthetic: true,
4829 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4830 Some(source),
4831 Some("qwen3"),
4832 false,
4833 true,
4834 None,
4835 None,
4836 ),
4837 })
4838 }
4839
4840 #[tokio::test]
4849 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4850 let strict = "{% if messages | length > 1 %}\
4853 {{ raise_exception('this template takes one turn') }}\
4854 {% endif %}{{ messages[0].content }}";
4855 let state = Arc::new(test_state(
4856 model_with_template("strict", strict),
4857 ResponseCache::new(4, Duration::from_secs(60)),
4858 ));
4859 let app = test_app_with_state(state);
4860
4861 let (status, body) = post_json_uri(
4862 &app,
4863 "/v1/chat/completions",
4864 serde_json::json!({
4865 "model": "strict",
4866 "stream": true,
4867 "messages": [
4868 {"role": "user", "content": "one"},
4869 {"role": "user", "content": "two"},
4870 ],
4871 }),
4872 )
4873 .await;
4874 assert_eq!(status, StatusCode::BAD_REQUEST);
4875 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4876 assert!(
4877 body["error"]["message"]
4878 .as_str()
4879 .unwrap()
4880 .contains("one turn"),
4881 "the template's own message must reach the caller: {body}"
4882 );
4883
4884 let (status, _) = post_json_uri(
4886 &app,
4887 "/v1/chat/completions",
4888 serde_json::json!({
4889 "model": "strict",
4890 "stream": true,
4891 "max_tokens": 1,
4892 "messages": [{"role": "user", "content": "one"}],
4893 }),
4894 )
4895 .await;
4896 assert_eq!(status, StatusCode::OK);
4897 }
4898
4899 #[tokio::test]
4901 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4902 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4903 {% if reasoning_effort %}\
4904 {% if reasoning_effort not in ['low','medium','high'] %}\
4905 {{ raise_exception('bad effort') }}\
4906 {% endif %}[{{ reasoning_effort }}]\
4907 {% endif %}{{ messages[0].content }}";
4908 let state = Arc::new(test_state(
4909 model_with_template("thinker", reasoning),
4910 ResponseCache::new(4, Duration::from_secs(60)),
4911 ));
4912 let app = test_app_with_state(state);
4913 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4914 assert_eq!(status, StatusCode::OK);
4915 let entry = &models["data"][0];
4916 assert_eq!(
4917 entry["supported_reasoning_efforts"],
4918 serde_json::json!(["off", "low", "medium", "high"])
4919 );
4920 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4921 }
4922
4923 #[tokio::test]
4928 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4929 let app = test_app();
4930 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4931 let entry = &models["data"][0];
4932 assert!(entry.get("supported_reasoning_efforts").is_none());
4933 assert!(entry.get("default_reasoning_effort").is_none());
4934 }
4935
4936 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
4937 Arc::new(ActiveModel {
4938 id: Some(name.to_string()),
4939 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
4940 batcher: None,
4941 ceiling: None,
4942 checkpoint_path: None,
4943 })
4944 .tap_into(state)
4945 }
4946
4947 trait TapInto {
4949 fn tap_into(self, state: &AppState) -> Self;
4950 }
4951 impl TapInto for Arc<ActiveModel> {
4952 fn tap_into(self, state: &AppState) -> Self {
4953 state.swap_active(Some(Arc::clone(&self)));
4954 self
4955 }
4956 }
4957
4958 #[test]
4964 fn an_in_flight_request_keeps_the_model_it_started_on() {
4965 let state = test_state(
4966 named_test_model("model-a", 256),
4967 ResponseCache::new(4, Duration::from_secs(60)),
4968 );
4969
4970 let in_flight = state.active().expect("a model is loaded");
4973 assert_eq!(in_flight.name(), "model-a");
4974
4975 active_model(&state, "model-b");
4976
4977 assert_eq!(state.active().unwrap().name(), "model-b");
4979 assert_eq!(in_flight.name(), "model-a");
4981 let (_chunks, finish, _usage) = run_generation(
4982 in_flight.generative().unwrap(),
4983 "hi",
4984 &greedy_params(3),
4985 None,
4986 None,
4987 None,
4988 None,
4989 None,
4990 None,
4991 )
4992 .expect("the old model must still decode after being swapped out");
4993 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
4994 }
4995
4996 #[test]
5001 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5002 let state = test_state(
5003 named_test_model("model-a", 256),
5004 ResponseCache::new(4, Duration::from_secs(60)),
5005 );
5006 let in_flight = state.active().expect("a model is loaded");
5007 let weights = Arc::clone(in_flight.generative().unwrap());
5008 assert!(Arc::strong_count(&weights) >= 2);
5009
5010 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5011 id: Some("model-b".to_string()),
5012 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5013 batcher: None,
5014 ceiling: None,
5015 checkpoint_path: None,
5016 })));
5017 drop(previous);
5018 assert!(Arc::strong_count(&weights) >= 2);
5020 drop(in_flight);
5021 assert_eq!(Arc::strong_count(&weights), 1);
5022 }
5023
5024 #[tokio::test]
5028 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5029 let state = Arc::new(test_state(
5030 named_test_model("model-a", 256),
5031 ResponseCache::new(4, Duration::from_secs(60)),
5032 ));
5033 let app = test_app_with_state(Arc::clone(&state));
5034
5035 let (status, body) = post_json_uri(
5036 &app,
5037 frink_api::routes::ADMIN_MODELS_UNLOAD,
5038 serde_json::json!({}),
5039 )
5040 .await;
5041 assert_eq!(status, StatusCode::OK);
5042 assert_eq!(body["ok"], true);
5043 assert!(body["active"].is_null());
5044 assert!(state.active().is_none());
5045
5046 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5047 assert_eq!(status, StatusCode::OK);
5048 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5049 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5050
5051 let (status, body) = post_json_uri(
5052 &app,
5053 "/v1/chat/completions",
5054 serde_json::json!({
5055 "model": "x",
5056 "messages": [{"role": "user", "content": "hi"}]
5057 }),
5058 )
5059 .await;
5060 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5061 assert_eq!(body["error"]["type"], "model_not_loaded");
5062 }
5063
5064 #[tokio::test]
5068 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5069 let state = Arc::new(test_state(
5070 named_test_model("model-a", 256),
5071 ResponseCache::new(4, Duration::from_secs(60)),
5072 ));
5073 let app = test_app_with_state(Arc::clone(&state));
5074 state.swap_active(None);
5075
5076 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5077 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5080 assert_eq!(body["state"], "unavailable");
5081 assert_eq!(body["reason"], "model_not_loaded");
5082 assert!(body["model"].is_null());
5083 let real_weights = body["capabilities"]
5084 .as_array()
5085 .unwrap()
5086 .iter()
5087 .find(|c| c["id"] == "real_weights")
5088 .cloned()
5089 .expect("real_weights is always reported");
5090 assert_eq!(real_weights["available"], false);
5091 assert_eq!(real_weights["reason"], "model_not_loaded");
5092 }
5093
5094 #[tokio::test]
5098 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5099 let app = test_app();
5100
5101 let (status, completion) = post_json_uri(
5102 &app,
5103 "/v1/chat/completions",
5104 serde_json::json!({
5105 "model": "x",
5106 "messages": [{"role": "user", "content": "hi"}],
5107 "max_tokens": 4
5108 }),
5109 )
5110 .await;
5111 assert_eq!(status, StatusCode::OK);
5112 let request_id = completion["request_id"].as_str().unwrap().to_string();
5113
5114 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5115 assert_eq!(status, StatusCode::OK);
5116 let recent = stats["recent"].as_array().unwrap();
5117 assert_eq!(recent.len(), 1);
5118 let row = &recent[0];
5119 assert_eq!(row["request_id"], request_id);
5120 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5121 assert_eq!(row["status"], 200);
5122 assert_eq!(row["stream"], false);
5123 assert!(row["duration_ms"].is_number());
5126 assert!(row["decode_ms"].is_number());
5127 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5128 assert_eq!(
5129 stats["tokens_prompt_total"].as_u64().unwrap(),
5130 row["prompt_tokens"].as_u64().unwrap()
5131 );
5132 }
5133
5134 #[tokio::test]
5138 async fn a_rejected_request_is_recorded_too() {
5139 let state = Arc::new(test_state(
5140 named_test_model("model-a", 256),
5141 ResponseCache::new(4, Duration::from_secs(60)),
5142 ));
5143 let app = test_app_with_state(Arc::clone(&state));
5144 state.swap_active(None);
5145
5146 let (status, _) = post_json_uri(
5147 &app,
5148 "/v1/chat/completions",
5149 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5150 )
5151 .await;
5152 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5153
5154 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5155 let recent = stats["recent"].as_array().unwrap();
5156 assert_eq!(recent.len(), 1);
5157 assert_eq!(recent[0]["status"], 503);
5158 assert_eq!(recent[0]["completion_tokens"], 0);
5159 assert!(recent[0]["decode_ms"].is_null());
5160 assert_eq!(stats["errors_total"], 1);
5161 }
5162
5163 async fn post_json_with_headers(
5167 app: &Router,
5168 uri: &str,
5169 body: serde_json::Value,
5170 headers: &[(&str, &str)],
5171 ) -> (StatusCode, serde_json::Value) {
5172 use http_body_util::BodyExt;
5173 use tower::ServiceExt;
5174
5175 let mut builder = axum::http::Request::builder()
5176 .method("POST")
5177 .uri(uri)
5178 .header("content-type", "application/json");
5179 for (name, value) in headers {
5180 builder = builder.header(*name, *value);
5181 }
5182 let response = app
5183 .clone()
5184 .oneshot(
5185 builder
5186 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5187 .unwrap(),
5188 )
5189 .await
5190 .unwrap();
5191 let status = response.status();
5192 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5193 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5194 (status, json)
5195 }
5196
5197 #[tokio::test]
5201 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5202 let app = test_app();
5203
5204 let (status, _) = post_json_uri(
5205 &app,
5206 frink_api::routes::V1_TOKENIZE,
5207 serde_json::json!({"prompt": "hello"}),
5208 )
5209 .await;
5210 assert_eq!(status, StatusCode::OK);
5211 let (status, _) = post_json_uri(
5212 &app,
5213 frink_api::routes::V1_DETOKENIZE,
5214 serde_json::json!({"tokens": [104, 105]}),
5215 )
5216 .await;
5217 assert_eq!(status, StatusCode::OK);
5218 let (status, _) = post_json_uri(
5219 &app,
5220 frink_api::routes::V1_EMBEDDINGS,
5221 serde_json::json!({"input": "hello"}),
5222 )
5223 .await;
5224 assert_eq!(status, StatusCode::OK);
5225
5226 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5227 let routes: Vec<&str> = stats["recent"]
5228 .as_array()
5229 .unwrap()
5230 .iter()
5231 .map(|row| row["route"].as_str().unwrap())
5232 .collect();
5233 for expected in [
5234 frink_api::routes::V1_TOKENIZE,
5235 frink_api::routes::V1_DETOKENIZE,
5236 frink_api::routes::V1_EMBEDDINGS,
5237 ] {
5238 assert!(
5239 routes.contains(&expected),
5240 "{expected} is missing: {routes:?}"
5241 );
5242 }
5243
5244 let row = |route: &str| {
5245 stats["recent"]
5246 .as_array()
5247 .unwrap()
5248 .iter()
5249 .find(|r| r["route"] == route)
5250 .cloned()
5251 .unwrap()
5252 };
5253 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5257 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5258 assert!(embed["decode_ms"].is_null());
5259 assert_eq!(embed["completion_tokens"], 0);
5260 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5264 assert_eq!(
5265 stats["tokens_prompt_total"].as_u64().unwrap(),
5266 embed["prompt_tokens"].as_u64().unwrap(),
5267 "only the forward pass counted"
5268 );
5269 }
5270
5271 fn streaming_test_app() -> Router {
5276 let mut cfg = test_dense_fixture();
5277 cfg.vocab_size = 256;
5278 let model = Model::Gguf(GgufModel {
5279 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5280 tokenizer: Arc::new(ServerTokenizer::Byte),
5281 stop_tokens: StopTokens::default(),
5282 bos_id: None,
5283 is_synthetic: false,
5284 chat_template: chat_template::PromptTemplate::plain(),
5285 });
5286 test_app_with_state(Arc::new(test_state(
5287 model,
5288 ResponseCache::new(1000, Duration::from_secs(3600)),
5289 )))
5290 }
5291
5292 #[tokio::test]
5296 async fn the_native_completion_wire_is_not_the_openai_one() {
5297 let app = test_app();
5298
5299 let (status, native) = post_json_uri(
5300 &app,
5301 frink_api::routes::COMPLETION,
5302 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5303 )
5304 .await;
5305 assert_eq!(status, StatusCode::OK, "{native}");
5306 assert!(native["content"].is_string(), "{native}");
5307 assert_eq!(native["stop"], true);
5308 assert_eq!(native["stop_type"], "limit");
5309 assert_eq!(native["stopping_word"], "");
5310 assert_eq!(native["truncated"], false);
5311 assert_eq!(native["id_slot"], -1);
5312 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5313 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5314 assert!(
5315 native.get("choices").is_none(),
5316 "the native shape has no `choices`: {native}"
5317 );
5318
5319 let (status, openai) = post_json_uri(
5320 &app,
5321 frink_api::routes::V1_COMPLETIONS,
5322 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5323 )
5324 .await;
5325 assert_eq!(status, StatusCode::OK);
5326 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5327 assert!(
5328 openai.get("content").is_none(),
5329 "the OpenAI shape has no top-level `content`: {openai}"
5330 );
5331 }
5332
5333 #[tokio::test]
5337 async fn both_native_spellings_reach_the_same_handler() {
5338 let app = test_app();
5339 for route in [
5340 frink_api::routes::COMPLETION,
5341 frink_api::routes::COMPLETIONS,
5342 ] {
5343 let (status, body) = post_json_uri(
5344 &app,
5345 route,
5346 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5347 )
5348 .await;
5349 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5350 assert_eq!(body["stop"], true, "{route}");
5351 assert!(body["content"].is_string(), "{route}");
5352 }
5353
5354 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5357 let routes: Vec<&str> = stats["recent"]
5358 .as_array()
5359 .unwrap()
5360 .iter()
5361 .map(|row| row["route"].as_str().unwrap())
5362 .collect();
5363 assert!(
5364 routes.contains(&frink_api::routes::COMPLETION),
5365 "{routes:?}"
5366 );
5367 assert!(
5368 routes.contains(&frink_api::routes::COMPLETIONS),
5369 "{routes:?}"
5370 );
5371 }
5372
5373 #[tokio::test]
5379 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5380 let app = streaming_test_app();
5381 let raw = post_sse_raw_uri(
5382 &app,
5383 frink_api::routes::COMPLETION,
5384 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5385 )
5386 .await;
5387
5388 assert!(
5389 !raw.contains("[DONE]"),
5390 "llama.cpp's native stream has no sentinel: {raw}"
5391 );
5392 let frames: Vec<serde_json::Value> = raw
5393 .lines()
5394 .filter_map(|line| line.strip_prefix("data: "))
5395 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5396 .collect();
5397 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5398
5399 let (last, partials) = frames.split_last().unwrap();
5400 assert_eq!(last["stop"], true, "the last frame closes the stream");
5401 assert!(last["timings"].is_object(), "{last}");
5402 assert!(last["stop_type"].is_string(), "{last}");
5403 for partial in partials {
5404 assert_eq!(partial["stop"], false, "{partial}");
5405 assert!(partial["content"].is_string(), "{partial}");
5406 assert!(partial.get("timings").is_none(), "{partial}");
5410 assert!(partial.get("generation_settings").is_none(), "{partial}");
5411 }
5412 let streamed: String = partials
5415 .iter()
5416 .filter_map(|p| p["content"].as_str())
5417 .collect();
5418 assert_eq!(last["content"].as_str().unwrap(), streamed);
5419 }
5420
5421 #[tokio::test]
5426 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5427 let app = test_app();
5428 for body in [
5429 serde_json::json!({"prompt": "hi"}),
5430 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5431 ] {
5432 let (status, refusal) =
5433 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5434 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5435 assert!(
5436 refusal["error"]["message"]
5437 .as_str()
5438 .unwrap()
5439 .contains("n_predict"),
5440 "{refusal}"
5441 );
5442 }
5443 let (status, _) = post_json_uri(
5446 &app,
5447 frink_api::routes::COMPLETION,
5448 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5449 )
5450 .await;
5451 assert_eq!(status, StatusCode::OK);
5452 }
5453
5454 #[tokio::test]
5463 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5464 let app = streaming_test_app();
5465 let ask = |stop: serde_json::Value| {
5466 let app = app.clone();
5467 async move {
5468 post_json_uri(
5469 &app,
5470 frink_api::routes::COMPLETION,
5471 serde_json::json!({
5472 "prompt": "hi",
5473 "n_predict": 64,
5474 "ignore_eos": true,
5475 "stop": stop,
5476 }),
5477 )
5478 .await
5479 .1
5480 }
5481 };
5482
5483 let baseline = ask(serde_json::json!([])).await;
5484 assert_eq!(baseline["stop_type"], "limit");
5485 assert_eq!(baseline["stopping_word"], "");
5486 let text = baseline["content"].as_str().unwrap().to_string();
5487 let sentinel: String = text.chars().skip(1).take(2).collect();
5493 assert_eq!(
5494 sentinel.chars().count(),
5495 2,
5496 "the fixture must produce enough output to cut: {text:?}"
5497 );
5498 let cut = text.find(&sentinel).expect("it came out of this text");
5499
5500 let stopped = ask(serde_json::json!([sentinel])).await;
5501 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5502 assert_eq!(stopped["stopping_word"], sentinel);
5503 assert_eq!(
5504 stopped["content"].as_str().unwrap(),
5505 &text[..cut],
5506 "the answer must be cut at the sentinel, not run past it"
5507 );
5508 }
5509
5510 #[tokio::test]
5516 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5517 let app = test_app();
5518
5519 let (v1_status, v1) = post_json_uri(
5520 &app,
5521 frink_api::routes::V1_TOKENIZE,
5522 serde_json::json!({"prompt": "hello"}),
5523 )
5524 .await;
5525 let (alias_status, alias) = post_json_uri(
5526 &app,
5527 frink_api::routes::TOKENIZE,
5528 serde_json::json!({"content": "hello"}),
5529 )
5530 .await;
5531 assert_eq!(v1_status, StatusCode::OK);
5532 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5533 assert_eq!(v1["tokens"], alias["tokens"]);
5534 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5535
5536 let (status, both_ways) = post_json_uri(
5539 &app,
5540 frink_api::routes::TOKENIZE,
5541 serde_json::json!({"prompt": "hello"}),
5542 )
5543 .await;
5544 assert_eq!(status, StatusCode::OK);
5545 assert_eq!(both_ways["tokens"], v1["tokens"]);
5546 }
5547
5548 #[tokio::test]
5553 async fn detokenize_answers_under_both_dialects_keys() {
5554 let app = test_app();
5555 for route in [
5556 frink_api::routes::DETOKENIZE,
5557 frink_api::routes::V1_DETOKENIZE,
5558 ] {
5559 let (status, body) =
5560 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5561 assert_eq!(status, StatusCode::OK, "{route}");
5562 assert_eq!(body["text"], "hi", "{route}");
5563 assert_eq!(body["content"], body["text"], "{route}");
5564 }
5565 }
5566
5567 #[tokio::test]
5571 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5572 let app = test_app();
5573 let (status, _) = post_json_uri(
5574 &app,
5575 frink_api::routes::TOKENIZE,
5576 serde_json::json!({"content": "hello"}),
5577 )
5578 .await;
5579 assert_eq!(status, StatusCode::OK);
5580
5581 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5582 let routes: Vec<&str> = stats["recent"]
5583 .as_array()
5584 .unwrap()
5585 .iter()
5586 .map(|row| row["route"].as_str().unwrap())
5587 .collect();
5588 assert!(
5589 routes.contains(&frink_api::routes::TOKENIZE),
5590 "the alias must be its own row: {routes:?}"
5591 );
5592 assert!(
5593 !routes.contains(&frink_api::routes::V1_TOKENIZE),
5594 "nothing called /v1/tokenize: {routes:?}"
5595 );
5596 }
5597
5598 #[tokio::test]
5603 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5604 let mut cfg = test_dense_fixture();
5605 cfg.vocab_size = 256;
5606 let model = Model::Gguf(GgufModel {
5607 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5608 tokenizer: Arc::new(ServerTokenizer::Byte),
5609 stop_tokens: StopTokens::default(),
5610 bos_id: Some(7),
5611 is_synthetic: true,
5612 chat_template: chat_template::PromptTemplate::plain(),
5613 });
5614 let app = test_app_with_state(Arc::new(test_state(
5615 model,
5616 ResponseCache::new(1000, Duration::from_secs(3600)),
5617 )));
5618
5619 let (_, plain) = post_json_uri(
5620 &app,
5621 frink_api::routes::TOKENIZE,
5622 serde_json::json!({"content": "hi"}),
5623 )
5624 .await;
5625 let (_, special) = post_json_uri(
5626 &app,
5627 frink_api::routes::TOKENIZE,
5628 serde_json::json!({"content": "hi", "add_special": true}),
5629 )
5630 .await;
5631
5632 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5633 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5634 assert_eq!(special["count"], 3);
5635 }
5636
5637 #[tokio::test]
5640 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5641 let app = test_app();
5642 let (status, _) = post_json_uri(
5643 &app,
5644 frink_api::routes::V1_EMBEDDINGS,
5645 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5646 )
5647 .await;
5648 assert_eq!(status, StatusCode::BAD_REQUEST);
5649
5650 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5651 let recent = stats["recent"].as_array().unwrap();
5652 assert_eq!(recent.len(), 1);
5653 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
5654 assert_eq!(recent[0]["status"], 400);
5655 assert_eq!(
5656 recent[0]["prompt_tokens"], 0,
5657 "a rejected call embedded nothing"
5658 );
5659 }
5660
5661 #[tokio::test]
5664 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5665 let app = test_app();
5666 let key = "sk-monitor-secret";
5667 let (status, _) = post_json_with_headers(
5668 &app,
5669 "/v1/chat/completions",
5670 serde_json::json!({
5671 "model": "x",
5672 "messages": [{"role": "user", "content": "hi"}],
5673 "max_tokens": 2
5674 }),
5675 &[
5676 ("authorization", &format!("Bearer {key}")),
5677 ("x-frink-client", "frink-studio"),
5678 ],
5679 )
5680 .await;
5681 assert_eq!(status, StatusCode::OK);
5682
5683 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5684 let row = stats["recent"].as_array().unwrap()[0].clone();
5685 let fingerprint = row["via_api_key"]
5686 .as_str()
5687 .expect("the row names the key that served it")
5688 .to_string();
5689 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5690 assert!(!fingerprint.contains(key));
5691 assert!(
5692 !serde_json::to_string(&stats).unwrap().contains(key),
5693 "the stats payload must not carry the key in any form"
5694 );
5695 assert_eq!(row["client"], "frink-studio");
5696 }
5697
5698 #[tokio::test]
5701 async fn different_keys_are_different_callers_and_no_key_is_null() {
5702 let app = test_app();
5703 let body = serde_json::json!({
5704 "model": "x",
5705 "messages": [{"role": "user", "content": "hi"}],
5706 "max_tokens": 1
5707 });
5708 for headers in [
5709 vec![("authorization", "Bearer key-one")],
5710 vec![("authorization", "Bearer key-two")],
5711 vec![],
5712 ] {
5713 let (status, _) =
5714 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5715 assert_eq!(status, StatusCode::OK);
5716 }
5717
5718 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5719 let recent = stats["recent"].as_array().unwrap();
5720 assert_eq!(recent.len(), 3);
5721 let one = recent[0]["via_api_key"].as_str().unwrap();
5722 let two = recent[1]["via_api_key"].as_str().unwrap();
5723 assert_ne!(one, two, "two keys must not collapse into one caller");
5724 assert!(
5725 recent[2]["via_api_key"].is_null(),
5726 "an unauthenticated call is null, not a fingerprint of nothing"
5727 );
5728 assert!(recent[2]["client"].is_null());
5729 }
5730
5731 #[tokio::test]
5736 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5737 let state = Arc::new(test_state(
5738 named_test_model("really-loaded", 256),
5739 ResponseCache::new(4, Duration::from_secs(60)),
5740 ));
5741 let app = test_app_with_state(Arc::clone(&state));
5742
5743 let (status, _) = post_json_uri(
5744 &app,
5745 "/v1/chat/completions",
5746 serde_json::json!({
5747 "model": "gpt-4-turbo-that-is-not-here",
5748 "messages": [{"role": "user", "content": "hi"}],
5749 "max_tokens": 2
5750 }),
5751 )
5752 .await;
5753 assert_eq!(status, StatusCode::OK);
5754
5755 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5756 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5757
5758 state.swap_active(None);
5761 let (status, _) = post_json_uri(
5762 &app,
5763 "/v1/chat/completions",
5764 serde_json::json!({
5765 "model": "gpt-4-turbo-that-is-not-here",
5766 "messages": [{"role": "user", "content": "hi"}]
5767 }),
5768 )
5769 .await;
5770 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5771 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5772 let recent = stats["recent"].as_array().unwrap();
5773 assert!(recent[recent.len() - 1]["model"].is_null());
5774 }
5775
5776 #[tokio::test]
5780 async fn a_streamed_row_names_the_model_it_decoded_against() {
5781 let state = Arc::new(test_state(
5782 named_test_model("model-before", 256),
5783 ResponseCache::new(4, Duration::from_secs(60)),
5784 ));
5785 let app = test_app_with_state(Arc::clone(&state));
5786 let _ = post_sse_raw(&app, resumable_request()).await;
5787 active_model(&state, "model-after");
5789
5790 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5791 assert_eq!(stats["recent"][0]["model"], "model-before");
5792 }
5793
5794 #[tokio::test]
5797 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5798 let app = test_app();
5799 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5800 assert_eq!(status, StatusCode::OK);
5801 assert!(
5802 stats["queue_depth"].is_null(),
5803 "without continuous batching nothing queues, so there is nothing to measure"
5804 );
5805 assert!(stats["queue_rejected_total"].is_null());
5806 assert_eq!(
5807 stats["generating_now"], 0,
5808 "work in progress is measured and really is zero here"
5809 );
5810 }
5811
5812 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5817 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
5818 }
5819
5820 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
5824 use http_body_util::BodyExt;
5825 use tower::ServiceExt;
5826
5827 let response = app
5828 .clone()
5829 .oneshot(
5830 axum::http::Request::builder()
5831 .method("POST")
5832 .uri(uri)
5833 .header("content-type", "application/json")
5834 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5835 .unwrap(),
5836 )
5837 .await
5838 .unwrap();
5839 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5840 String::from_utf8(bytes.to_vec()).unwrap()
5841 }
5842
5843 async fn get_json_with_headers(
5844 app: &Router,
5845 uri: &str,
5846 headers: &[(&str, &str)],
5847 ) -> (StatusCode, serde_json::Value) {
5848 use http_body_util::BodyExt;
5849 use tower::ServiceExt;
5850
5851 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
5852 for (name, value) in headers {
5853 builder = builder.header(*name, *value);
5854 }
5855 let response = app
5856 .clone()
5857 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
5858 .await
5859 .unwrap();
5860 let status = response.status();
5861 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5862 (
5863 status,
5864 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
5865 )
5866 }
5867
5868 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
5869 body.lines()
5870 .filter_map(|line| line.strip_prefix(field))
5871 .map(str::trim)
5872 .collect()
5873 }
5874
5875 fn resumable_request() -> serde_json::Value {
5876 serde_json::json!({
5877 "model": "m",
5878 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5879 "max_tokens": 4,
5880 "temperature": 0,
5881 "stream": true,
5882 "stream_resumable": true,
5883 })
5884 }
5885
5886 #[tokio::test]
5891 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
5892 let app = test_app();
5893 let body = post_sse_raw(&app, resumable_request()).await;
5894
5895 let request_id = body
5896 .lines()
5897 .find_map(|l| l.strip_prefix("data: "))
5898 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
5899 .and_then(|v| v["request_id"].as_str().map(str::to_string))
5900 .expect("the first chunk names the request");
5901
5902 let ids = sse_field(&body, "id:");
5903 let datas = sse_field(&body, "data:");
5904 assert_eq!(
5905 ids.len(),
5906 datas.len(),
5907 "every event carries an id, or a reconnect cannot name where it stopped"
5908 );
5909 for (i, id) in ids.iter().enumerate() {
5910 assert_eq!(*id, format!("{request_id}:{i}"));
5911 }
5912 let retries = sse_field(&body, "retry:");
5913 assert_eq!(
5914 retries.len(),
5915 1,
5916 "the reconnect delay is stated once, not on every event"
5917 );
5918 assert_eq!(retries[0], "1500");
5919 assert!(
5920 body.contains("data: [DONE]"),
5921 "the end of stream is still stated"
5922 );
5923 }
5924
5925 #[tokio::test]
5929 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
5930 let app = test_app();
5931 let mut request = resumable_request();
5932 request["stream_resumable"] = serde_json::json!(false);
5933 let body = post_sse_raw(&app, request).await;
5934 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
5935 assert!(
5936 sse_field(&body, "id:").is_empty(),
5937 "an id promises a replay this stream cannot serve"
5938 );
5939 assert!(sse_field(&body, "retry:").is_empty());
5940 }
5941
5942 #[tokio::test]
5946 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
5947 let app = test_app();
5948 let body = post_sse_raw(&app, resumable_request()).await;
5949 let request_id = sse_field(&body, "id:")[0]
5950 .rsplit_once(':')
5951 .unwrap()
5952 .0
5953 .to_string();
5954 let streamed: Vec<String> = sse_field(&body, "data:")
5955 .iter()
5956 .map(|d| d.to_string())
5957 .collect();
5958
5959 let (status, polled) = get_json(
5960 &app,
5961 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
5962 )
5963 .await;
5964 assert_eq!(status, StatusCode::OK);
5965 let events: Vec<String> = polled["events"]
5966 .as_array()
5967 .unwrap()
5968 .iter()
5969 .map(|e| e["data"].as_str().unwrap().to_string())
5970 .collect();
5971 assert_eq!(
5972 events, streamed,
5973 "the fallback must deliver the same answer, not a re-run of it"
5974 );
5975 assert_eq!(polled["request_id"], request_id);
5976 assert_eq!(
5977 polled["done"], false,
5978 "events were still being handed out, so the client must ask again"
5979 );
5980
5981 let next = polled["next_index"].as_u64().unwrap();
5984 let (_, drained) = get_json(
5985 &app,
5986 &format!(
5987 "{}?from={next}",
5988 frink_api::routes::v1_stream_poll(&request_id)
5989 ),
5990 )
5991 .await;
5992 assert_eq!(drained["done"], true);
5993 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
5994 }
5995
5996 #[tokio::test]
6000 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6001 let app = test_app();
6002 let body = post_sse_raw(&app, resumable_request()).await;
6003 let ids = sse_field(&body, "id:");
6004 let datas: Vec<String> = sse_field(&body, "data:")
6005 .iter()
6006 .map(|d| d.to_string())
6007 .collect();
6008 assert!(
6009 ids.len() >= 3,
6010 "need a few events to resume into the middle"
6011 );
6012 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6013
6014 let (status, resumed) = get_json_with_headers(
6015 &app,
6016 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6017 &[],
6018 )
6019 .await;
6020 assert_eq!(status, StatusCode::OK);
6021 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6022
6023 let (_, tail) = get_json(
6025 &app,
6026 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6027 )
6028 .await;
6029 let tail_events: Vec<String> = tail["events"]
6030 .as_array()
6031 .unwrap()
6032 .iter()
6033 .map(|e| e["data"].as_str().unwrap().to_string())
6034 .collect();
6035 assert_eq!(tail_events, datas[2..].to_vec());
6036 }
6037
6038 #[tokio::test]
6041 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6042 use http_body_util::BodyExt;
6043 use tower::ServiceExt;
6044
6045 let app = test_app();
6046 let body = post_sse_raw(&app, resumable_request()).await;
6047 let ids = sse_field(&body, "id:");
6048 let datas: Vec<String> = sse_field(&body, "data:")
6049 .iter()
6050 .map(|d| d.to_string())
6051 .collect();
6052 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6053
6054 let response = app
6055 .clone()
6056 .oneshot(
6057 axum::http::Request::builder()
6058 .method("GET")
6059 .uri(frink_api::routes::v1_stream(&request_id))
6060 .header("last-event-id", format!("{request_id}:0"))
6061 .body(axum::body::Body::empty())
6062 .unwrap(),
6063 )
6064 .await
6065 .unwrap();
6066 assert_eq!(response.status(), StatusCode::OK);
6067 assert_eq!(
6068 response
6069 .headers()
6070 .get("x-accel-buffering")
6071 .and_then(|v| v.to_str().ok()),
6072 Some("no"),
6073 "the reconnect needs the same anti-buffering header as the stream"
6074 );
6075 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6076 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6077 assert_eq!(
6078 sse_field(&resumed, "data:")
6079 .iter()
6080 .map(|d| d.to_string())
6081 .collect::<Vec<_>>(),
6082 datas[1..].to_vec()
6083 );
6084 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6085 }
6086
6087 #[tokio::test]
6091 async fn a_last_event_id_from_another_stream_is_refused() {
6092 let app = test_app();
6093 let body = post_sse_raw(&app, resumable_request()).await;
6094 let request_id = sse_field(&body, "id:")[0]
6095 .rsplit_once(':')
6096 .unwrap()
6097 .0
6098 .to_string();
6099
6100 let (status, err) = get_json_with_headers(
6101 &app,
6102 &frink_api::routes::v1_stream(&request_id),
6103 &[("last-event-id", "chatcmpl-someone-else:3")],
6104 )
6105 .await;
6106 assert_eq!(status, StatusCode::BAD_REQUEST);
6107 assert_eq!(err["error"]["code"], "bad_last_event_id");
6108 }
6109
6110 #[tokio::test]
6114 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6115 let app = test_app();
6116 let mut request = resumable_request();
6117 request["stream_resumable"] = serde_json::json!(false);
6118 let body = post_sse_raw(&app, request).await;
6119 let request_id = body
6120 .lines()
6121 .find_map(|l| l.strip_prefix("data: "))
6122 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6123 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6124 .unwrap();
6125
6126 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6127 assert_eq!(status, StatusCode::NOT_FOUND);
6128 assert_eq!(err["error"]["code"], "stream_not_found");
6129 assert!(err["error"]["message"]
6130 .as_str()
6131 .unwrap()
6132 .contains("stream_resumable"));
6133 }
6134
6135 #[test]
6139 fn the_axum_stream_patterns_match_the_published_templates() {
6140 assert_eq!(
6141 axum_path(frink_api::routes::V1_STREAM),
6142 "/v1/stream/:request_id"
6143 );
6144 assert_eq!(
6145 axum_path(frink_api::routes::V1_STREAM_POLL),
6146 "/v1/stream/:request_id/poll"
6147 );
6148 assert_eq!(
6149 frink_api::routes::v1_stream("abc"),
6150 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6151 );
6152 }
6153
6154 #[test]
6164 fn no_published_template_reaches_the_router_with_its_braces() {
6165 for template in [
6166 frink_api::routes::V1_STREAM,
6167 frink_api::routes::V1_STREAM_POLL,
6168 frink_api::routes::V1_RESPONSE,
6169 frink_api::routes::V1_RESPONSE_CANCEL,
6170 frink_api::routes::ADMIN_TASK_CANCEL,
6171 ] {
6172 assert!(
6173 template.contains('{'),
6174 "{template} is in the template list but has no placeholder"
6175 );
6176 let mounted = axum_path(template);
6177 assert!(
6178 !mounted.contains('{') && !mounted.contains('}'),
6179 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6180 );
6181 assert!(
6182 mounted.contains(':'),
6183 "{template} lost its placeholder entirely and would match one path only"
6184 );
6185 }
6186 }
6187
6188 #[tokio::test]
6196 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6197 let app = test_app();
6198 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6199 assert_eq!(status, StatusCode::NOT_FOUND);
6200 assert!(
6201 !body.is_null(),
6202 "empty body means axum never matched the route, so the id was read as a literal segment"
6203 );
6204 }
6205
6206 #[tokio::test]
6209 async fn the_task_list_starts_empty_rather_than_absent() {
6210 let app = test_app();
6211 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6212 assert_eq!(status, StatusCode::OK);
6213 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6214 }
6215
6216 #[tokio::test]
6225 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6226 assert!(
6227 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6228 "this test asserts the unconfigured behaviour"
6229 );
6230 let app = test_app();
6231 let (status, body) = post_json_uri(
6232 &app,
6233 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6234 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6235 )
6236 .await;
6237 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6238 assert!(
6239 body["error"]["message"]
6240 .as_str()
6241 .unwrap()
6242 .contains("--slot-save-path"),
6243 "{body}"
6244 );
6245 }
6246
6247 pub(crate) async fn post_json_uri(
6248 app: &Router,
6249 uri: &str,
6250 body: serde_json::Value,
6251 ) -> (StatusCode, serde_json::Value) {
6252 use http_body_util::BodyExt;
6253 use tower::ServiceExt;
6254
6255 let response = app
6256 .clone()
6257 .oneshot(
6258 axum::http::Request::builder()
6259 .method("POST")
6260 .uri(uri)
6261 .header("content-type", "application/json")
6262 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6263 .unwrap(),
6264 )
6265 .await
6266 .unwrap();
6267 let status = response.status();
6268 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6269 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6270 (status, json)
6271 }
6272
6273 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6274 post_json_uri(app, "/v1/chat/completions", body).await.1
6275 }
6276
6277 #[tokio::test]
6286 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6287 let app = test_app();
6288 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6289 assert_eq!(status, StatusCode::OK);
6290
6291 let memory = &body["memory"];
6292 if memory.is_null() {
6293 return;
6296 }
6297 assert!(
6298 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6299 "a read that produced a zero is a broken read, not an idle \
6300 engine: {memory}"
6301 );
6302 assert!(
6303 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6304 "the quantity must travel with the number: {memory}"
6305 );
6306 }
6307
6308 #[tokio::test]
6314 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6315 let app = test_app();
6316 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6317 assert_eq!(status, StatusCode::OK);
6318 for pool in ["kv_pages", "window_slots", "state_slots"] {
6319 assert!(
6320 body["pools"][pool].is_null(),
6321 "{pool} must be null rather than a zero row: {}",
6322 body["pools"]
6323 );
6324 }
6325 }
6326
6327 #[tokio::test]
6335 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6336 use http_body_util::BodyExt;
6337 use tower::ServiceExt;
6338
6339 let app = test_app();
6340 let send = |body: serde_json::Value| {
6341 let app = app.clone();
6342 async move {
6343 app.oneshot(
6344 axum::http::Request::builder()
6345 .method("POST")
6346 .uri(frink_api::routes::V1_MESSAGES)
6347 .header("content-type", "application/json")
6348 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6349 .unwrap(),
6350 )
6351 .await
6352 .unwrap()
6353 }
6354 };
6355
6356 let ok = send(serde_json::json!({
6357 "model": "test",
6358 "max_tokens": 1,
6359 "messages": [{"role": "user", "content": "hi"}],
6360 }))
6361 .await;
6362 assert_eq!(ok.status(), StatusCode::OK);
6363 let id = ok
6364 .headers()
6365 .get("request-id")
6366 .expect("a served message names its id")
6367 .to_str()
6368 .unwrap()
6369 .to_string();
6370 assert!(!id.is_empty());
6371
6372 let bad = send(serde_json::json!({"model": "test"})).await;
6375 assert!(bad.status().is_client_error());
6376 let other = bad.headers().get("request-id").expect("errors too");
6377 assert_ne!(other.to_str().unwrap(), id);
6378 let _ = bad.into_body().collect().await.unwrap();
6379 }
6380
6381 #[tokio::test]
6389 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6390 let state = Arc::new(test_state(
6391 test_model_full_byte_vocab(),
6392 ResponseCache::new(1000, Duration::from_secs(3600)),
6393 ));
6394 let app = test_app_with_state(Arc::clone(&state));
6395 let body = serde_json::json!({
6396 "model": "test",
6397 "messages": [{"role": "user", "content": "hi"}],
6398 "max_tokens": 1,
6399 });
6400
6401 state
6402 .maintenance
6403 .lock()
6404 .unwrap()
6405 .begin_rebuild()
6406 .expect("a fresh server is serving, so the rebuild starts");
6407 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6408 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6409 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6410
6411 state.maintenance.lock().unwrap().finish_rebuild(true);
6412 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6413 assert_eq!(
6414 status,
6415 StatusCode::OK,
6416 "the gate reopens; a rebuild is not a latch"
6417 );
6418 }
6419
6420 #[tokio::test]
6425 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6426 let app = test_app();
6427 let (status, body) = post_json_uri(
6428 &app,
6429 frink_api::routes::V1_CANCEL,
6430 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6431 )
6432 .await;
6433 assert_eq!(status, StatusCode::NOT_FOUND);
6434 assert_eq!(body["cancelled"], serde_json::json!(false));
6435 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6436 assert!(
6437 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6438 "the verdict must carry a human reason: {body}"
6439 );
6440 }
6441
6442 #[tokio::test]
6446 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6447 let state = Arc::new(test_state(
6448 test_model_full_byte_vocab(),
6449 ResponseCache::new(1000, Duration::from_secs(3600)),
6450 ));
6451 let app = test_app_with_state(Arc::clone(&state));
6452 let (token, _guard) = state.cancels.register("chatcmpl-live");
6453
6454 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6455 assert_eq!(status, StatusCode::OK);
6456 assert_eq!(before["generating_now"], serde_json::json!(1));
6457
6458 let (status, body) = post_json_uri(
6459 &app,
6460 frink_api::routes::V1_CANCEL,
6461 serde_json::json!({ "request_id": "chatcmpl-live" }),
6462 )
6463 .await;
6464 assert_eq!(status, StatusCode::OK);
6465 assert_eq!(body["cancelled"], serde_json::json!(true));
6466 assert!(
6467 token.is_cancelled(),
6468 "the endpoint answered ok without setting the flag the decode loop reads"
6469 );
6470 }
6471
6472 #[tokio::test]
6473 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6474 let app = test_app();
6475 let (status, tok) =
6476 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6477 assert_eq!(status, StatusCode::OK);
6478 let tokens = tok["tokens"].as_array().unwrap();
6479 assert_eq!(tok["count"], tokens.len());
6480 assert!(!tokens.is_empty());
6481
6482 let (status, detok) = post_json_uri(
6483 &app,
6484 "/v1/detokenize",
6485 serde_json::json!({ "tokens": tokens }),
6486 )
6487 .await;
6488 assert_eq!(status, StatusCode::OK);
6489 assert_eq!(detok["text"], "Hi");
6490
6491 let (status, emb) = post_json_uri(
6492 &app,
6493 "/v1/embeddings",
6494 serde_json::json!({
6495 "input": "Hi",
6496 "embedding_type": "mean"
6497 }),
6498 )
6499 .await;
6500 assert_eq!(status, StatusCode::OK);
6501 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6502 assert!(!vec.is_empty());
6503 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6504 }
6505
6506 #[tokio::test]
6512 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6513 let app = test_app();
6514 let (status, body) = post_json_uri(
6515 &app,
6516 "/v1/embeddings",
6517 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6518 )
6519 .await;
6520 assert_eq!(status, StatusCode::BAD_REQUEST);
6521 let msg = body["error"]["message"].as_str().unwrap();
6522 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6523 }
6524
6525 #[tokio::test]
6530 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6531 async fn a_real_embedding_model_serves_v1_embeddings() {
6532 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6533 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6534 if !path.exists() {
6535 eprintln!("SKIP: {} not present", path.display());
6536 return;
6537 }
6538 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6539 let mut state = test_state(
6540 test_model_full_byte_vocab(),
6541 ResponseCache::new(1000, Duration::from_secs(3600)),
6542 );
6543 state.embedding = Some(Arc::new(encoder));
6544 let app = test_app_with_state(Arc::new(state));
6545
6546 let (status, body) = post_json_uri(
6547 &app,
6548 "/v1/embeddings",
6549 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6550 )
6551 .await;
6552 assert_eq!(status, StatusCode::OK, "{body}");
6553 assert_eq!(body["model"], "bge-small-en-v1.5");
6554 let data = body["data"].as_array().unwrap();
6555 assert_eq!(data.len(), 2);
6556 for (i, row) in data.iter().enumerate() {
6557 assert_eq!(row["index"], i);
6558 let v: Vec<f64> = row["embedding"]
6559 .as_array()
6560 .unwrap()
6561 .iter()
6562 .map(|x| x.as_f64().unwrap())
6563 .collect();
6564 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6565 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6566 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6567 }
6568 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6571
6572 let (status, mean) = post_json_uri(
6576 &app,
6577 "/v1/embeddings",
6578 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6579 )
6580 .await;
6581 assert_eq!(status, StatusCode::OK);
6582 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6583 }
6584
6585 #[tokio::test]
6597 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6598 async fn an_encoder_can_be_the_loaded_model() {
6599 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6600 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6601 if !path.exists() {
6602 eprintln!("SKIP: {} not present", path.display());
6603 return;
6604 }
6605
6606 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6610 assert!(
6611 matches!(loaded, model::LoadedModel::Encoder(_)),
6612 "an encoder-only GGUF reached a decoder loader"
6613 );
6614 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6615 assert!(
6616 matches!(loaded, Loaded::Encoder(_)),
6617 "the encoder did not stay an encoder through activation"
6618 );
6619 assert!(
6620 batcher.is_none() && ceiling.is_none(),
6621 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6622 );
6623
6624 let state = test_state(
6625 test_model_full_byte_vocab(),
6626 ResponseCache::new(1000, Duration::from_secs(3600)),
6627 );
6628 state.swap_active(Some(Arc::new(ActiveModel {
6629 id: None,
6630 loaded,
6631 batcher,
6632 ceiling,
6633 checkpoint_path: None,
6634 })));
6635 let app = test_app_with_state(Arc::new(state));
6636
6637 let (status, body) = post_json_uri(
6639 &app,
6640 "/v1/embeddings",
6641 serde_json::json!({ "input": "Hello world" }),
6642 )
6643 .await;
6644 assert_eq!(status, StatusCode::OK, "{body}");
6645 assert_eq!(body["model"], "bge-small-en-v1.5");
6646 let v = body["data"][0]["embedding"].as_array().unwrap();
6647 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6648
6649 let (status, body) = post_json_uri(
6651 &app,
6652 "/v1/chat/completions",
6653 serde_json::json!({
6654 "model": "bge-small-en-v1.5",
6655 "messages": [{"role": "user", "content": "hi"}],
6656 }),
6657 )
6658 .await;
6659 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6660 let msg = body["error"]["message"].as_str().unwrap();
6661 for fact in [
6662 "bge-small-en-v1.5",
6663 "bert",
6664 "embedding model",
6665 "/v1/embeddings",
6666 ] {
6667 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6668 }
6669
6670 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
6672 assert_eq!(status, StatusCode::OK);
6673 let entry = &models["data"][0];
6674 assert_eq!(entry["id"], "bge-small-en-v1.5");
6675 assert_eq!(entry["frink_model_kind"], "embedding");
6676 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
6677 assert_eq!(entry["frink_n_embd"], 384);
6678 assert_eq!(entry["frink_pooling"], "CLS");
6679 assert_eq!(
6680 entry["frink_endpoints"],
6681 serde_json::json!(["/v1/embeddings"])
6682 );
6683 assert!(entry.get("supported_reasoning_efforts").is_none());
6686
6687 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
6689 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6690 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6691 assert_eq!(health["model"]["synthetic_weights"], false);
6692 let weights = health["capabilities"]
6693 .as_array()
6694 .unwrap()
6695 .iter()
6696 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
6697 .expect("a real-weights capability row");
6698 let detail = weights["detail"].as_str().unwrap_or_default();
6699 assert!(detail.contains("ENCODER"), "{detail}");
6700 let (status, body) = post_json_uri(
6708 &app,
6709 frink_api::routes::V1_TOKENIZE,
6710 serde_json::json!({ "content": "hello world" }),
6711 )
6712 .await;
6713 assert_eq!(
6714 status,
6715 StatusCode::OK,
6716 "an encoder has a real tokenizer: {body}"
6717 );
6718 let tokens = body["tokens"].as_array().expect("tokens array").clone();
6719 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
6720
6721 let (status, body) = post_json_uri(
6722 &app,
6723 frink_api::routes::V1_DETOKENIZE,
6724 serde_json::json!({ "tokens": tokens }),
6725 )
6726 .await;
6727 assert_eq!(status, StatusCode::OK, "{body}");
6728 let round_tripped = body["content"].as_str().expect("content").to_string();
6729 assert!(
6730 round_tripped.contains("hello") && round_tripped.contains("world"),
6731 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
6732 );
6733
6734 let (status, _) = post_json_uri(
6737 &app,
6738 "/v1/completions",
6739 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
6740 )
6741 .await;
6742 assert_eq!(
6743 status,
6744 StatusCode::NOT_IMPLEMENTED,
6745 "tokenizing an encoder must not have opened a path to generating with one"
6746 );
6747 }
6748
6749 #[tokio::test]
6754 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6755 use http_body_util::BodyExt;
6756 use tower::ServiceExt;
6757
6758 let fixture = concat!(
6759 "../frink-models/tests/fixtures/",
6760 "frink_real_moe_test.gguf"
6761 );
6762 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6763 let decoder = Decoder::from_gguf_with_expert_cache(
6764 &fixture,
6765 frink_models::config::test_moe_fixture(),
6766 Some(1024 * 1024),
6767 )
6768 .expect("MoE fixture must load store-backed");
6769
6770 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
6774 decoder.forward_token(1, 0, &mut caches);
6775
6776 let model = Model::Gguf(GgufModel {
6777 decoder: Arc::new(decoder),
6778 tokenizer: Arc::new(ServerTokenizer::Byte),
6779 stop_tokens: StopTokens::default(),
6780 bos_id: None,
6781 is_synthetic: false,
6782 chat_template: chat_template::PromptTemplate::plain(),
6783 });
6784 let state = Arc::new(test_state(
6785 model,
6786 ResponseCache::new(16, Duration::from_secs(60)),
6787 ));
6788 let app = Router::new()
6789 .route("/metrics", axum::routing::get(metrics))
6790 .route("/v1/chat/completions", post(chat_completions))
6791 .with_state(state);
6792
6793 let fetch_metrics = |app: Router| async move {
6794 let resp = app
6795 .oneshot(
6796 axum::http::Request::builder()
6797 .method("GET")
6798 .uri("/metrics")
6799 .body(axum::body::Body::empty())
6800 .unwrap(),
6801 )
6802 .await
6803 .unwrap();
6804 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6805 String::from_utf8(bytes.to_vec()).unwrap()
6806 };
6807
6808 let after = fetch_metrics(app.clone()).await;
6809 assert!(
6810 after.contains("frink_expert_cache_misses_total"),
6811 "streaming model must expose expert-cache metrics: {after}"
6812 );
6813 let misses: u64 = after
6814 .lines()
6815 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
6816 .and_then(|l| l.split_whitespace().nth(1))
6817 .and_then(|v| v.parse().ok())
6818 .expect("misses metric line must parse");
6819 assert!(
6820 misses > 0,
6821 "decode must have read experts through the store: {after}"
6822 );
6823 }
6824
6825 fn weather_tool() -> serde_json::Value {
6826 serde_json::json!({
6827 "type": "function",
6828 "function": {
6829 "name": "get_weather",
6830 "description": "Get the current weather for a location.",
6831 "parameters": {
6832 "type": "object",
6833 "properties": {"location": {"type": "string"}},
6834 "required": ["location"]
6835 }
6836 }
6837 })
6838 }
6839
6840 fn weather_tool_def() -> ToolDef {
6841 ToolDef {
6842 kind: "function".to_string(),
6843 function: ToolFunctionDef {
6844 name: "get_weather".to_string(),
6845 description: Some("Get the current weather for a location.".to_string()),
6846 parameters: Some(serde_json::json!({
6847 "type": "object",
6848 "properties": {"location": {"type": "string"}},
6849 "required": ["location"]
6850 })),
6851 },
6852 }
6853 }
6854
6855 #[test]
6856 fn tool_preamble_mentions_every_tool_name_and_description() {
6857 let preamble = tool_preamble(&[weather_tool_def()]);
6858 assert!(preamble.contains("get_weather"));
6859 assert!(preamble.contains("Get the current weather for a location."));
6860 assert!(preamble.contains("<tool_call>"));
6861 assert!(preamble.contains("</tool_call>"));
6862 }
6863
6864 #[test]
6865 fn a_real_marker_becomes_a_structured_tool_call() {
6866 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6867 let (message, finish) = build_response_message(
6868 text.to_string(),
6869 &[weather_tool_def()],
6870 output::OutputPosture::for_model("test-model"),
6871 "stop",
6872 );
6873 assert_eq!(finish, "tool_calls");
6874 let calls = message.tool_calls.expect("must carry a tool call");
6875 assert_eq!(calls[0].function.name, "get_weather");
6876 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
6877 assert_eq!(parsed["location"], "Paris");
6878 }
6879
6880 #[test]
6881 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
6882 let (message, finish) = build_response_message(
6883 "just an answer".to_string(),
6884 &[weather_tool_def()],
6885 output::OutputPosture::for_model("test-model"),
6886 "stop",
6887 );
6888 assert_eq!(finish, "stop");
6889 assert!(message.tool_calls.is_none());
6890 assert_eq!(message.content.as_deref(), Some("just an answer"));
6891 }
6892
6893 #[test]
6896 fn a_malformed_payload_is_not_a_tool_call() {
6897 let (message, finish) = build_response_message(
6898 "<tool_call>not valid json at all</tool_call>".to_string(),
6899 &[weather_tool_def()],
6900 output::OutputPosture::for_model("test-model"),
6901 "stop",
6902 );
6903 assert_eq!(finish, "stop");
6904 assert!(message.tool_calls.is_none());
6905 }
6906
6907 #[test]
6910 fn a_tool_that_was_never_offered_is_not_returned() {
6911 let (message, finish) = build_response_message(
6912 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
6913 &[weather_tool_def()],
6914 output::OutputPosture::for_model("test-model"),
6915 "stop",
6916 );
6917 assert_eq!(finish, "stop");
6918 assert!(message.tool_calls.is_none());
6919 }
6920
6921 #[test]
6923 fn marker_text_with_no_tools_offered_stays_content() {
6924 let (message, finish) = build_response_message(
6925 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
6926 &[],
6927 output::OutputPosture::for_model("test-model"),
6928 "stop",
6929 );
6930 assert_eq!(finish, "stop");
6931 assert!(message.tool_calls.is_none());
6932 assert!(message.content.is_some());
6933 }
6934
6935 #[test]
6939 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
6940 let opened = std::cell::Cell::new(0usize);
6941 let mut parser = crate::policy::parser::ToolCallParser::new(
6942 crate::policy::parser::ToolCallFormat::Qwen3Coder,
6943 vec![
6944 crate::policy::parser::tool_call::ToolSchema::with_parameters(
6945 "write_file",
6946 serde_json::json!({"type": "object", "properties": {
6947 "path": {"type": "string"},
6948 "contents": {"type": "string"}
6949 }}),
6950 ),
6951 ],
6952 );
6953 let wire = "<tool_call><function=write_file>\
6954 <parameter=path>\n/tmp/x\n</parameter>\
6955 <parameter=contents>\nhello world\n</parameter>\
6956 </function></tool_call>";
6957
6958 let mut deltas = Vec::new();
6959 let mut text = String::new();
6960 for piece in wire.as_bytes().chunks(7) {
6961 let chunk = String::from_utf8_lossy(piece).into_owned();
6962 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
6963 text.push_str(&more_text);
6964 deltas.extend(more);
6965 }
6966 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
6967 text.push_str(&more_text);
6968 deltas.extend(more);
6969
6970 assert_eq!(opened.get(), 1, "one call opened");
6971 assert!(text.is_empty(), "the markers are not content: {text:?}");
6972
6973 let first = &deltas[0];
6974 assert_eq!(first.index, 0);
6975 assert_eq!(first.id.as_deref(), Some("call_0"));
6976 assert_eq!(first.kind, Some("function"));
6977 assert_eq!(first.function.name.as_deref(), Some("write_file"));
6978
6979 let joined: String = deltas
6982 .iter()
6983 .filter_map(|d| d.function.arguments.clone())
6984 .collect();
6985 let parsed: serde_json::Value =
6986 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
6987 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
6988 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
6989 assert!(
6990 deltas.len() >= 3,
6991 "the arguments arrived in pieces, not whole: {}",
6992 deltas.len()
6993 );
6994 assert!(
6995 deltas[1..].iter().all(|d| d.function.name.is_none()),
6996 "only the opening delta carries identity"
6997 );
6998 }
6999
7000 #[test]
7002 fn text_around_a_streamed_call_is_still_content() {
7003 let opened = std::cell::Cell::new(0usize);
7004 let mut parser = crate::policy::parser::ToolCallParser::new(
7005 crate::policy::parser::ToolCallFormat::Qwen25,
7006 vec![crate::policy::parser::tool_call::ToolSchema::new(
7007 "get_weather",
7008 )],
7009 );
7010 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7011 \"arguments\": {}}</tool_call> done";
7012 let mut text = String::new();
7013 for piece in wire.as_bytes().chunks(5) {
7014 let chunk = String::from_utf8_lossy(piece).into_owned();
7015 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7016 text.push_str(&more);
7017 }
7018 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7019 text.push_str(&more);
7020
7021 assert_eq!(opened.get(), 1);
7022 assert!(text.starts_with("let me check. "), "{text:?}");
7023 assert!(text.ends_with(" done"), "{text:?}");
7024 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7025 }
7026
7027 #[test]
7030 fn a_reasoning_block_is_split_out_of_the_answer() {
7031 let (message, finish) = build_response_message(
7032 "<think>weighing it up</think>The answer is 4.".to_string(),
7033 &[],
7034 output::OutputPosture::for_model("Qwen3-8B"),
7035 "stop",
7036 );
7037 assert_eq!(finish, "stop");
7038 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7039 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7040 }
7041
7042 #[test]
7045 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7046 let (message, _) = build_response_message(
7047 "Use the <think> tag like this.".to_string(),
7048 &[],
7049 output::OutputPosture::for_model("llama-3.1-8b"),
7050 "stop",
7051 );
7052 assert_eq!(
7053 message.content.as_deref(),
7054 Some("Use the <think> tag like this.")
7055 );
7056 assert!(message.reasoning_content.is_none());
7057 }
7058
7059 #[tokio::test]
7067 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7068 let app = test_app();
7069 let body = serde_json::json!({
7070 "model": "m",
7071 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7072 "max_tokens": 4,
7073 "temperature": 0,
7074 });
7075 let resp = post_json(&app, body).await;
7076 let message = &resp["choices"][0]["message"];
7077 assert!(message["content"].is_string());
7078 assert!(message.get("tool_calls").is_none());
7079 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7080 assert_eq!(resp["usage"]["completion_tokens"], 4);
7081 assert_eq!(
7082 resp["usage"]["total_tokens"],
7083 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7084 );
7085 }
7086
7087 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7088 use http_body_util::BodyExt;
7089 use tower::ServiceExt;
7090
7091 let response = app
7092 .clone()
7093 .oneshot(
7094 axum::http::Request::builder()
7095 .method("GET")
7096 .uri(uri)
7097 .body(axum::body::Body::empty())
7098 .unwrap(),
7099 )
7100 .await
7101 .unwrap();
7102 let status = response.status();
7103 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7104 (status, serde_json::from_slice(&bytes).unwrap())
7105 }
7106
7107 #[tokio::test]
7108 async fn health_answers_a_capability_handshake_not_a_boolean() {
7109 let app = test_app();
7110 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7111 assert_eq!(status, StatusCode::OK);
7112
7113 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7114 assert_eq!(health.state, frink_api::HealthState::Ready);
7115 assert!(health.pid > 0);
7116 assert!(health.server_time_unix_ms > 0);
7117 assert_eq!(health.last_request_age_seconds, None);
7120
7121 for id in [
7124 frink_api::health::capability::CPU,
7125 frink_api::health::capability::METAL,
7126 frink_api::health::capability::CUDA,
7127 frink_api::health::capability::REAL_WEIGHTS,
7128 frink_api::health::capability::CONTINUOUS_BATCHING,
7129 ] {
7130 let cap = health
7131 .capability(id)
7132 .unwrap_or_else(|| panic!("{id} missing"));
7133 assert!(!cap.reason.is_empty(), "{cap:?}");
7134 assert!(!cap.detail.is_empty(), "{cap:?}");
7135 }
7136 let weights = health
7140 .capability(frink_api::health::capability::REAL_WEIGHTS)
7141 .unwrap();
7142 assert!(!weights.available);
7143 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7144 assert!(health.model.as_ref().unwrap().synthetic_weights);
7145 }
7146
7147 #[tokio::test]
7148 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7149 let app = test_app();
7150 let _ = post_json(
7151 &app,
7152 serde_json::json!({
7153 "model": "m",
7154 "messages": [{"role": "user", "content": "\u{1}"}],
7155 "max_tokens": 1,
7156 "temperature": 0,
7157 }),
7158 )
7159 .await;
7160 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7161 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7162 let age = health
7163 .last_request_age_seconds
7164 .expect("a served request is evidence of liveness");
7165 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7166 }
7167
7168 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7170 use http_body_util::BodyExt;
7171 use tower::ServiceExt;
7172
7173 let response = app
7174 .clone()
7175 .oneshot(
7176 axum::http::Request::builder()
7177 .method("POST")
7178 .uri("/v1/chat/completions")
7179 .header("content-type", "application/json")
7180 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7181 .unwrap(),
7182 )
7183 .await
7184 .unwrap();
7185 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7186 String::from_utf8(bytes.to_vec())
7187 .unwrap()
7188 .lines()
7189 .filter_map(|line| line.strip_prefix("data: "))
7190 .filter(|payload| *payload != "[DONE]")
7191 .map(|payload| serde_json::from_str(payload).unwrap())
7192 .collect()
7193 }
7194
7195 #[tokio::test]
7196 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7197 let app = test_app();
7198 let chunks = post_sse_chunks(
7199 &app,
7200 serde_json::json!({
7201 "model": "m",
7202 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7203 "max_tokens": 4,
7204 "temperature": 0,
7205 "stream": true,
7206 }),
7207 )
7208 .await;
7209
7210 assert!(!chunks.is_empty());
7211 let request_id = chunks[0]["request_id"]
7212 .as_str()
7213 .expect("the first chunk names the request")
7214 .to_string();
7215 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7216 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7219 assert!(
7220 chunk.get("request_id").is_none(),
7221 "chunk {i} repeats request_id"
7222 );
7223 }
7224 for chunk in &chunks {
7227 assert_eq!(chunk["id"], serde_json::json!(request_id));
7228 }
7229
7230 let other = post_sse_chunks(
7231 &app,
7232 serde_json::json!({
7233 "model": "m",
7234 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7235 "max_tokens": 4,
7236 "temperature": 0,
7237 "stream": true,
7238 }),
7239 )
7240 .await;
7241 assert_ne!(
7242 other[0]["request_id"].as_str().unwrap(),
7243 request_id,
7244 "two concurrent chats must not share an id"
7245 );
7246 }
7247
7248 #[tokio::test]
7249 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7250 let app = test_app();
7251 let resp = post_json(
7252 &app,
7253 serde_json::json!({
7254 "model": "m",
7255 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7256 "max_tokens": 2,
7257 "temperature": 0,
7258 }),
7259 )
7260 .await;
7261 assert_eq!(resp["id"], resp["request_id"]);
7262 assert!(resp["request_id"]
7263 .as_str()
7264 .unwrap()
7265 .starts_with("chatcmpl-"));
7266 }
7267
7268 #[tokio::test]
7271 async fn usage_carries_separate_prefill_and_decode_timings() {
7272 let app = test_app();
7273 let resp = post_json(
7274 &app,
7275 serde_json::json!({
7276 "model": "m",
7277 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7278 "max_tokens": 4,
7279 "temperature": 0,
7280 }),
7281 )
7282 .await;
7283 let usage = &resp["usage"];
7284 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7285 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7286 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7287 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7288 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7290 }
7291
7292 #[tokio::test]
7300 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7301 let app = test_app();
7302 let body = serde_json::json!({
7303 "model": "m",
7304 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7305 "max_tokens": 4,
7306 "temperature": 0,
7307 "tools": [weather_tool()],
7308 });
7309 let resp = post_json(&app, body).await;
7310 let message = &resp["choices"][0]["message"];
7311 assert!(
7312 message["content"].is_string(),
7313 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7314 );
7315 assert!(message.get("tool_calls").is_none());
7316 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7320 }
7321
7322 #[tokio::test]
7326 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7327 let app = test_app();
7328 let body = serde_json::json!({
7329 "model": "m",
7330 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7331 "max_tokens": 3,
7332 "temperature": 0,
7333 });
7334 let first = post_json(&app, body.clone()).await;
7335 assert_eq!(first["frink_cache"], "miss");
7336 let second = post_json(&app, body).await;
7337 assert_eq!(second["frink_cache"], "hit");
7338 assert_eq!(
7339 first["choices"][0]["message"]["content"],
7340 second["choices"][0]["message"]["content"]
7341 );
7342 assert_eq!(
7343 first["choices"][0]["finish_reason"],
7344 second["choices"][0]["finish_reason"]
7345 );
7346 assert_eq!(first["usage"], second["usage"]);
7347 assert_eq!(second["usage"]["completion_tokens"], 3);
7348 }
7349
7350 #[tokio::test]
7361 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7362 let app = test_app();
7363 let plain = serde_json::json!({
7364 "model": "m",
7365 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7366 "max_tokens": 3,
7367 "temperature": 0,
7368 });
7369
7370 let first = post_json(&app, plain.clone()).await;
7371 assert_eq!(first["frink_cache"], "miss");
7372 let unconstrained = first["choices"][0]["message"]["content"]
7373 .as_str()
7374 .expect("content")
7375 .to_string();
7376
7377 let mut constrained = plain.clone();
7378 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7379 let second = post_json(&app, constrained).await;
7380 assert_eq!(
7381 second["frink_cache"], "miss",
7382 "a grammar is part of the key, so this body has never been answered"
7383 );
7384 let constrained_answer = second["choices"][0]["message"]["content"]
7388 .as_str()
7389 .expect("content")
7390 .to_string();
7391 assert!(
7392 constrained_answer.contains("-> \"yes\"]"),
7393 "the grammar must have been compiled AND applied, not skipped \
7394 by a cache hit: {constrained_answer}"
7395 );
7396 assert_ne!(
7397 constrained_answer, unconstrained,
7398 "the constrained request was served the unconstrained answer"
7399 );
7400
7401 let third = post_json(&app, plain).await;
7405 assert_eq!(third["frink_cache"], "hit");
7406 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7407 }
7408
7409 #[tokio::test]
7431 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7432 let state = Arc::new(test_state(
7433 test_model_full_byte_vocab(),
7434 ResponseCache::new(1000, Duration::from_secs(3600)),
7435 ));
7436 let app = test_app_with_state(state.clone());
7437 let plain = serde_json::json!({
7438 "model": "m",
7439 "messages": [
7440 {"role": "system", "content": "Answer in JSON when it helps."},
7441 {"role": "user", "content": "\u{1}\u{2}"},
7442 ],
7443 "max_tokens": 3,
7444 "temperature": 0,
7445 });
7446
7447 let first = post_json(&app, plain.clone()).await;
7448 assert_eq!(first["frink_cache"], "miss");
7449 assert_eq!(state.cache_stats().entries, 1);
7450
7451 let mut as_json = plain.clone();
7452 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7453 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7454 assert_eq!(
7455 status,
7456 StatusCode::BAD_REQUEST,
7457 "the demo banner is not a JSON object, whoever generated it"
7458 );
7459 assert_eq!(
7460 state.cache_stats().hits,
7461 0,
7462 "a json_object request must not be answered from an entry the \
7463 JSON mask never produced"
7464 );
7465 assert_eq!(
7466 state.cache_stats().entries,
7467 2,
7468 "json_object must key its own entry, not reuse the unconstrained \
7469 one it happens to render the same prompt as"
7470 );
7471 }
7472
7473 #[tokio::test]
7485 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7486 let app = test_app_with_state(Arc::new(test_state(
7487 test_model_full_byte_vocab_with_eos(Some(0x77)),
7488 ResponseCache::new(1000, Duration::from_secs(3600)),
7489 )));
7490 let body = serde_json::json!({
7491 "model": "m",
7492 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7493 "max_tokens": 6,
7494 "temperature": 0,
7495 });
7496
7497 let stopped = post_json(&app, body.clone()).await;
7498 assert_eq!(stopped["frink_cache"], "miss");
7499 assert_eq!(
7500 stopped["choices"][0]["finish_reason"], "stop",
7501 "the fixture is only meaningful if the model's EOS really fires here"
7502 );
7503 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7504
7505 let mut ignoring = body.clone();
7506 ignoring["ignore_eos"] = serde_json::json!(true);
7507 let ran_on = post_json(&app, ignoring).await;
7508 assert_eq!(
7509 ran_on["frink_cache"], "miss",
7510 "ignore_eos is part of the key, so this body has never been answered"
7511 );
7512 assert_eq!(
7513 ran_on["usage"]["completion_tokens"], 6,
7514 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7515 );
7516 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7517 assert_ne!(
7518 ran_on["choices"][0]["message"]["content"],
7519 stopped["choices"][0]["message"]["content"]
7520 );
7521 }
7522
7523 #[tokio::test]
7529 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7530 let session_app = test_app();
7531 let manual_app = test_app();
7532
7533 let turn1 = post_json(
7535 &session_app,
7536 serde_json::json!({
7537 "model": "m",
7538 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7539 "session_id": "s1",
7540 "max_tokens": 5,
7541 "temperature": 0,
7542 }),
7543 )
7544 .await;
7545 let reply1 = turn1["choices"][0]["message"]["content"]
7546 .as_str()
7547 .unwrap()
7548 .to_string();
7549
7550 let manual_turn1 = post_json(
7555 &manual_app,
7556 serde_json::json!({
7557 "model": "m",
7558 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7559 "max_tokens": 5,
7560 "temperature": 0,
7561 }),
7562 )
7563 .await;
7564 assert_eq!(
7565 manual_turn1["choices"][0]["message"]["content"]
7566 .as_str()
7567 .unwrap(),
7568 reply1
7569 );
7570
7571 let turn2 = post_json(
7573 &session_app,
7574 serde_json::json!({
7575 "model": "m",
7576 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7577 "session_id": "s1",
7578 "max_tokens": 5,
7579 "temperature": 0,
7580 }),
7581 )
7582 .await;
7583 let reply2 = turn2["choices"][0]["message"]["content"]
7584 .as_str()
7585 .unwrap()
7586 .to_string();
7587
7588 let manual_turn2 = post_json(
7592 &manual_app,
7593 serde_json::json!({
7594 "model": "m",
7595 "messages": [
7596 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7597 {"role": "assistant", "content": reply1},
7598 {"role": "user", "content": "\u{4}\u{5}"},
7599 ],
7600 "max_tokens": 5,
7601 "temperature": 0,
7602 }),
7603 )
7604 .await;
7605 assert_eq!(
7606 manual_turn2["choices"][0]["message"]["content"]
7607 .as_str()
7608 .unwrap(),
7609 reply2,
7610 "resuming a session must produce identical output to manually resending the full history"
7611 );
7612 }
7613
7614 #[test]
7617 fn lock_cache_recovers_from_a_poisoned_mutex() {
7618 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7619
7620 let poison_cache = Arc::clone(&cache);
7621 let _ = std::thread::spawn(move || {
7622 let _guard = poison_cache.lock().unwrap();
7623 panic!("simulated panic while holding the lock");
7624 })
7625 .join();
7626
7627 let recovered = lock_cache(&cache);
7629 assert_eq!(recovered.stats().entries, 0);
7630 }
7631
7632 #[test]
7633 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7634 let mut req_body = serde_json::json!({
7635 "model": "m",
7636 "messages": [{"role": "user", "content": "hi"}],
7637 });
7638 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7639 assert!(
7640 req.is_cacheable(),
7641 "default (temperature 0) must be cacheable"
7642 );
7643
7644 req_body["temperature"] = serde_json::json!(0.8);
7645 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7646 assert!(
7647 !req.is_cacheable(),
7648 "unseeded sampling must never be cacheable"
7649 );
7650
7651 req_body["seed"] = serde_json::json!(42);
7652 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7653 assert!(
7654 req.is_cacheable(),
7655 "sampling with an explicit seed is deterministic and must be cacheable"
7656 );
7657 }
7658
7659 const GRADED: &str = "{% if reasoning_effort %}\
7663 {% if reasoning_effort not in ['low','medium','high'] %}\
7664 {{ raise_exception('unsupported effort') }}\
7665 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7666 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7667
7668 fn graded_template() -> chat_template::PromptTemplate {
7669 chat_template::PromptTemplate::from_gguf_metadata(
7670 Some(GRADED),
7671 Some("qwen3"),
7672 false,
7673 true,
7674 None,
7675 None,
7676 )
7677 }
7678
7679 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7680 serde_json::from_value(value).expect("request")
7681 }
7682
7683 #[test]
7690 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7691 let req = chat_request(serde_json::json!({
7692 "model": "m",
7693 "messages": [{"role": "user", "content": "hi"}],
7694 "grammar": "root ::= \"a\"+",
7695 }));
7696 req.validate_supported_fields()
7697 .expect("a valid grammar is a valid request");
7698 let params = req
7699 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7700 .expect("a valid grammar compiles at params time too");
7701 assert!(
7702 params.grammar.is_some(),
7703 "the grammar was dropped between the wire and the sampler"
7704 );
7705 assert!(
7706 params.needs_vocab_logits(),
7707 "a grammar request that may fold lm_head into a GPU argmax is \
7708 a grammar request served unconstrained"
7709 );
7710
7711 let plain = chat_request(serde_json::json!({
7712 "model": "m",
7713 "messages": [{"role": "user", "content": "hi"}],
7714 }));
7715 assert!(plain
7716 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7717 .unwrap()
7718 .grammar
7719 .is_none());
7720 }
7721
7722 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7723 chat_request(serde_json::json!({
7724 "model": "m",
7725 "messages": [{"role": "user", "content": "weather in Rome?"}],
7726 "tools": [weather_tool()],
7727 "tool_choice": tool_choice,
7728 }))
7729 }
7730
7731 #[test]
7735 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7736 for choice in [
7737 serde_json::json!("required"),
7738 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7739 ] {
7740 let req = tool_request(choice.clone());
7741 req.validate_supported_fields()
7742 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7743 let params = req
7744 .generation_params_for_template(
7745 &graded_template(),
7746 "Qwen3-8B",
7747 crate::sampling_knobs::SamplerModel::absent(),
7748 )
7749 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7750 let grammar = params
7751 .grammar
7752 .as_ref()
7753 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7754 assert!(
7755 grammar.is_awaiting_trigger(),
7756 "the model must be free to think before it calls"
7757 );
7758 assert!(
7759 !grammar.allows_eog(),
7760 "{choice} must not be able to end the turn without a call"
7761 );
7762 assert!(
7768 params.needs_vocab_logits(),
7769 "{choice} would let a backend return a token id instead of logits"
7770 );
7771 assert!(
7772 !generate::greedy_gpu_fold_allowed(¶ms),
7773 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7774 );
7775 }
7776 }
7777
7778 #[test]
7780 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7781 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7782 let req = tool_request(choice.clone());
7783 req.validate_supported_fields().expect("still supported");
7784 let params = match req.generation_params_for_template(
7785 &graded_template(),
7786 "Qwen3-8B",
7787 crate::sampling_knobs::SamplerModel::absent(),
7788 ) {
7789 Ok(p) => p,
7790 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7791 };
7792 assert!(
7793 params.grammar.is_none(),
7794 "{choice} does not force a call and must not be constrained"
7795 );
7796 }
7797 }
7798
7799 #[test]
7802 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7803 let req = chat_request(serde_json::json!({
7805 "model": "m",
7806 "messages": [{"role": "user", "content": "hi"}],
7807 "tool_choice": "required",
7808 }));
7809 let (status, _) = req
7810 .validate_supported_fields()
7811 .expect_err("nothing to call");
7812 assert_eq!(status, StatusCode::BAD_REQUEST);
7813
7814 let req =
7816 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7817 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7818 assert_eq!(status, StatusCode::BAD_REQUEST);
7819 assert_eq!(body["error"]["param"], "tool_choice");
7820
7821 let req = tool_request(serde_json::json!({"type": "function"}));
7823 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7824 assert_eq!(status, StatusCode::BAD_REQUEST);
7825
7826 let req = chat_request(serde_json::json!({
7828 "model": "m",
7829 "messages": [{"role": "user", "content": "hi"}],
7830 "tools": [weather_tool()],
7831 "tool_choice": "required",
7832 "grammar": "root ::= \"a\"+",
7833 }));
7834 let (status, _) = req
7835 .validate_supported_fields()
7836 .expect_err("a grammar and a forced call are two constraints");
7837 assert_eq!(status, StatusCode::BAD_REQUEST);
7838
7839 let req = tool_request(serde_json::json!("required"));
7845 let (status, Json(body)) = match req.generation_params_for_template(
7846 &graded_template(),
7847 "muse-glimmer-8b",
7848 crate::sampling_knobs::SamplerModel::absent(),
7849 ) {
7850 Err(e) => e,
7851 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
7852 };
7853 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7854 assert!(
7855 body["error"]["message"]
7856 .as_str()
7857 .unwrap()
7858 .contains("muse_glimmer"),
7859 "{body}"
7860 );
7861
7862 let req = tool_request(serde_json::json!("required"));
7868 let params = req
7869 .generation_params_for_template(
7870 &graded_template(),
7871 "gemma-4-E2B-it",
7872 crate::sampling_knobs::SamplerModel::absent(),
7873 )
7874 .expect("a gemma4 forced tool_choice is served");
7875 assert!(
7876 params.grammar.is_some(),
7877 "a forced tool_choice must arrive as the generation's grammar"
7878 );
7879 }
7880
7881 #[test]
7884 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
7885 let req = chat_request(serde_json::json!({
7886 "model": "m",
7887 "messages": [{"role": "user", "content": "hi"}],
7888 "grammar": "root ::= \"a",
7889 }));
7890 let (status, Json(body)) = req
7891 .validate_supported_fields()
7892 .expect_err("this does not parse");
7893 assert_eq!(status, StatusCode::BAD_REQUEST);
7894 assert_eq!(body["error"]["param"], "grammar");
7895 assert!(
7896 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
7897 .is_err(),
7898 "and again at params time"
7899 );
7900 }
7901
7902 #[test]
7909 fn response_format_json_schema_becomes_the_requests_grammar() {
7910 let req = chat_request(serde_json::json!({
7911 "model": "m",
7912 "messages": [{"role": "user", "content": "hi"}],
7913 "response_format": {
7914 "type": "json_schema",
7915 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
7916 },
7917 }));
7918 req.validate_supported_fields()
7919 .expect("a boolean schema converts");
7920 let params = req
7921 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7922 .expect("and compiles");
7923 let grammar = params.grammar.expect("the schema is the grammar");
7924 let mut g = (*grammar).clone();
7925 g.accept_token(0, b"true").expect("a boolean is accepted");
7926 assert!(g.allows_eog(), "and completes the parse");
7927 assert!(
7928 !params.json_object,
7929 "a schema is not the json_object character-class mask"
7930 );
7931 }
7932
7933 #[test]
7937 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
7938 let req = chat_request(serde_json::json!({
7939 "model": "m",
7940 "messages": [{"role": "user", "content": "hi"}],
7941 "response_format": {
7942 "type": "json_schema",
7943 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
7944 },
7945 }));
7946 let (status, Json(body)) = req
7947 .validate_supported_fields()
7948 .expect_err("minimum has no grammar in this port");
7949 assert_eq!(status, StatusCode::BAD_REQUEST);
7950 assert!(
7951 body["error"]["message"]
7952 .as_str()
7953 .expect("a message")
7954 .contains("minimum"),
7955 "the refusal must name the keyword: {body}"
7956 );
7957 assert!(
7958 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
7959 .is_err(),
7960 "and again at params time"
7961 );
7962 }
7963
7964 #[test]
7970 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
7971 let req = chat_request(serde_json::json!({
7972 "model": "m",
7973 "messages": [{"role": "user", "content": "hi"}],
7974 "tool_choice": "required",
7975 "tools": [{
7976 "type": "function",
7977 "function": {"name": "f", "parameters": {"type": "object"}},
7978 }],
7979 "response_format": {
7980 "type": "json_schema",
7981 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
7982 },
7983 }));
7984 let (status, Json(body)) = req
7985 .validate_supported_fields()
7986 .expect_err("two constraints, one generation");
7987 assert_eq!(status, StatusCode::BAD_REQUEST);
7988 assert_eq!(body["error"]["param"], "tool_choice");
7989 }
7990
7991 #[test]
7994 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
7995 let req = chat_request(serde_json::json!({
7996 "model": "m",
7997 "messages": [{"role": "user", "content": "hi"}],
7998 }));
7999 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8000 }
8001
8002 #[test]
8007 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8008 let asked = chat_request(serde_json::json!({
8009 "model": "m",
8010 "messages": [{"role": "user", "content": "hi"}],
8011 "min_p": 0.07,
8012 }));
8013 assert_eq!(
8014 asked
8015 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8016 .expect("knobs")
8017 .min_p,
8018 0.07
8019 );
8020
8021 let silent = chat_request(serde_json::json!({
8022 "model": "m",
8023 "messages": [{"role": "user", "content": "hi"}],
8024 }));
8025 assert_eq!(
8026 silent
8027 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8028 .expect("knobs")
8029 .min_p,
8030 0.0,
8031 "an unset min_p must be off, not llama.cpp's CLI default"
8032 );
8033 }
8034
8035 #[test]
8044 fn no_sampler_knob_is_missing_from_the_cache_key() {
8045 let base = serde_json::json!({
8046 "model": "m",
8047 "messages": [{"role": "user", "content": "hi"}],
8048 "seed": 1,
8049 });
8050 let key_for = |body: serde_json::Value| {
8051 let req = chat_request(body);
8052 let params = req
8053 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8054 .expect("params");
8055 req.cache_key("prompt", ¶ms)
8056 };
8057 let baseline = key_for(base.clone());
8058 for (knob, value) in [
8059 ("temperature", serde_json::json!(0.5)),
8060 ("top_p", serde_json::json!(0.9)),
8061 ("min_p", serde_json::json!(0.05)),
8062 ("top_k", serde_json::json!(40)),
8063 ("repetition_penalty", serde_json::json!(1.1)),
8064 ("presence_penalty", serde_json::json!(0.3)),
8065 ("frequency_penalty", serde_json::json!(0.3)),
8066 (
8067 "samplers",
8068 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8069 ),
8070 ] {
8071 let mut body = base.clone();
8072 body[knob] = value;
8073 assert_ne!(
8074 key_for(body),
8075 baseline,
8076 "`{knob}` is not in the cache key: two requests differing \
8077 only in it would share one cached answer"
8078 );
8079 }
8080 }
8081
8082 #[test]
8092 fn no_constraint_is_missing_from_the_cache_key() {
8093 let base = serde_json::json!({
8094 "model": "m",
8095 "messages": [{"role": "user", "content": "pick one"}],
8096 });
8097 let key_for = |body: serde_json::Value| {
8098 let req = chat_request(body);
8099 let params = req
8100 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8101 .expect("params");
8102 req.cache_key("prompt", ¶ms)
8103 };
8104 let baseline = key_for(base.clone());
8105 for (field, value) in [
8106 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8107 (
8108 "response_format",
8109 serde_json::json!({"type": "json_object"}),
8110 ),
8111 (
8112 "response_format",
8113 serde_json::json!({"type": "json_schema", "json_schema": {
8114 "name": "answer",
8115 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8116 }}),
8117 ),
8118 ("ignore_eos", serde_json::json!(true)),
8119 ("stop", serde_json::json!(["\n"])),
8120 ("max_tokens", serde_json::json!(7)),
8121 ] {
8122 let mut body = base.clone();
8123 body[field] = value.clone();
8124 assert_ne!(
8125 key_for(body),
8126 baseline,
8127 "`{field}: {value}` is not in the cache key: two requests \
8128 differing only in it would share one cached answer"
8129 );
8130 }
8131 }
8132
8133 #[test]
8137 fn an_explicit_zero_output_budget_is_a_client_error() {
8138 let req = chat_request(serde_json::json!({
8139 "model": "m",
8140 "messages": [{"role": "user", "content": "hi"}],
8141 "max_tokens": 0,
8142 }));
8143 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8144 assert_eq!(status, StatusCode::BAD_REQUEST);
8145 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8146 }
8147
8148 #[test]
8151 fn a_request_can_turn_thinking_off() {
8152 let template = graded_template();
8153 for body in [
8154 serde_json::json!({
8155 "model": "m",
8156 "messages": [{"role": "user", "content": "hi"}],
8157 "reasoning_effort": "none",
8158 }),
8159 serde_json::json!({
8160 "model": "m",
8161 "messages": [{"role": "user", "content": "hi"}],
8162 "thinking": {"type": "disabled"},
8163 }),
8164 ] {
8165 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8166 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8167 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8168 assert!(!kwargs.contains_key("reasoning_effort"));
8171 }
8172 }
8173
8174 #[test]
8177 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8178 let template = graded_template();
8179 let kwargs = chat_request(serde_json::json!({
8180 "model": "m",
8181 "messages": [{"role": "user", "content": "hi"}],
8182 "reasoning_effort": "high",
8183 "thinking": {"type": "disabled"},
8184 }))
8185 .resolve_template_kwargs(&template);
8186 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8187 assert!(!kwargs.contains_key("reasoning_effort"));
8188 }
8189
8190 #[test]
8193 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8194 let req = chat_request(serde_json::json!({
8195 "model": "m",
8196 "messages": [{"role": "user", "content": "hi"}],
8197 "thinking": {"type": "disable"},
8198 }));
8199 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8200 assert_eq!(status, StatusCode::BAD_REQUEST);
8201 }
8202
8203 #[test]
8206 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8207 let template = graded_template();
8208 let kwargs = chat_request(serde_json::json!({
8209 "model": "m",
8210 "messages": [{"role": "user", "content": "hi"}],
8211 "reasoning_effort": "none",
8212 "chat_template_kwargs": {"enable_thinking": true},
8213 }))
8214 .resolve_template_kwargs(&template);
8215 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8216 }
8217
8218 #[test]
8222 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8223 let template = graded_template();
8224 let req = chat_request(serde_json::json!({
8225 "model": "m",
8226 "messages": [{"role": "user", "content": "hi"}],
8227 "reasoning_effort": "minimal",
8228 }));
8229 let kwargs = req.resolve_template_kwargs(&template);
8230 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8231 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8232 assert!(prompt.starts_with("E:low|"), "{prompt}");
8233 }
8234
8235 #[test]
8239 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8240 let template = graded_template();
8241 let req = chat_request(serde_json::json!({
8242 "model": "m",
8243 "messages": [{"role": "user", "content": "hi"}],
8244 "chat_template_kwargs": {"reasoning_effort": "none"},
8245 }));
8246 let kwargs = req.resolve_template_kwargs(&template);
8247 assert!(!kwargs.contains_key("reasoning_effort"));
8248 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8249 assert_eq!(prompt, "hi");
8250 }
8251
8252 #[test]
8255 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8256 let template = graded_template();
8257 let req = chat_request(serde_json::json!({
8258 "model": "m",
8259 "messages": [{"role": "user", "content": "hi"}],
8260 "reasoning_effort": "low",
8261 "chat_template_kwargs": {"reasoning_effort": "high"},
8262 }));
8263 assert_eq!(
8264 req.resolve_template_kwargs(&template)["reasoning_effort"],
8265 serde_json::json!("high")
8266 );
8267 }
8268
8269 #[test]
8273 fn offering_tools_turns_thinking_on_by_itself() {
8274 let template = graded_template();
8275 let quiet = chat_request(serde_json::json!({
8276 "model": "m",
8277 "messages": [{"role": "user", "content": "hi"}],
8278 }));
8279 assert!(!quiet
8280 .resolve_template_kwargs(&template)
8281 .contains_key("enable_thinking"));
8282
8283 let with_tools = chat_request(serde_json::json!({
8284 "model": "m",
8285 "messages": [{"role": "user", "content": "hi"}],
8286 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8287 }));
8288 let kwargs = with_tools.resolve_template_kwargs(&template);
8289 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8290 let prompt =
8291 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8292 assert!(prompt.starts_with("THINK|"), "{prompt}");
8293 }
8294
8295 #[test]
8300 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8301 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8302 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8303 Some("qwen3"),
8304 false,
8305 true,
8306 None,
8307 None,
8308 );
8309 let req = chat_request(serde_json::json!({
8310 "model": "m",
8311 "messages": [{"role": "user", "content": "hi"}],
8312 "chat_template_kwargs": {"enable_thinking": true},
8313 }));
8314 let kwargs = req.resolve_template_kwargs(&opener);
8315 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8316 assert!(prompt.ends_with("<think>"), "{prompt}");
8317
8318 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8321 let (message, _) = build_response_message(
8322 "weighing it up</think>Paris.".to_string(),
8323 &[],
8324 posture,
8325 "stop",
8326 );
8327 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8328 assert_eq!(message.content.as_deref(), Some("Paris."));
8329
8330 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8333 let (message, _) = build_response_message(
8334 "weighing it up</think>Paris.".to_string(),
8335 &[],
8336 closed,
8337 "stop",
8338 );
8339 assert_eq!(message.reasoning_content, None);
8340 }
8341
8342 #[test]
8343 fn stop_param_accepts_both_single_string_and_array() {
8344 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8345 "model": "m",
8346 "messages": [{"role": "user", "content": "hi"}],
8347 "stop": "END",
8348 }))
8349 .unwrap();
8350 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8351
8352 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8353 "model": "m",
8354 "messages": [{"role": "user", "content": "hi"}],
8355 "stop": ["A", "B"],
8356 }))
8357 .unwrap();
8358 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8359 }
8360
8361 #[test]
8362 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8363 let model = test_model();
8364 let result = run_generation(
8365 &model,
8366 "hello",
8367 &greedy_params(4),
8368 None,
8369 None,
8370 None,
8371 None,
8372 None,
8373 None,
8374 );
8375 assert!(matches!(
8376 result,
8377 Err(generate::DecodeError::TokenOutOfVocab { .. })
8378 ));
8379 }
8380
8381 #[test]
8385 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8386 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8388 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8389
8390 let holder_pool = Arc::clone(&pool);
8391 let holder = std::thread::spawn(move || {
8392 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8393 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8395 drop(held);
8396 });
8397 std::thread::sleep(Duration::from_millis(15));
8398
8399 let config = generate::KvPoolConfig {
8400 pool,
8401 queue_wait: Duration::ZERO,
8402 };
8403 let result = run_generation(
8404 &model,
8405 &prompt,
8406 &greedy_params(4),
8407 Some(&config),
8408 None,
8409 None,
8410 None,
8411 None,
8412 None,
8413 );
8414 assert!(matches!(
8415 result,
8416 Err(generate::DecodeError::KvPoolExhausted)
8417 ));
8418
8419 let (status, _body) = decode_error_response(result.unwrap_err());
8420 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8421 holder.join().unwrap();
8422 }
8423
8424 #[test]
8432 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8433 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8435 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8437 let config = generate::KvPoolConfig {
8438 pool,
8439 queue_wait: Duration::ZERO,
8440 };
8441
8442 let result = run_generation(
8443 &model,
8444 &prompt,
8445 &greedy_params(4),
8446 Some(&config),
8447 None,
8448 None,
8449 None,
8450 None,
8451 None,
8452 );
8453 let err = result.expect_err("one block cannot hold two layers' caches");
8454 assert!(
8455 matches!(
8456 &err,
8457 generate::DecodeError::KvBudgetExceeded { binding, .. }
8458 if *binding == frink_models::Ceiling::DeviceMemory.code()
8459 ),
8460 "expected an immovable device-memory refusal, got {err:?}"
8461 );
8462 let (status, _body) = decode_error_response(err);
8463 assert_eq!(status, StatusCode::BAD_REQUEST);
8464 }
8465
8466 #[test]
8472 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8473 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8474 queued: 512,
8475 cap: 512,
8476 });
8477 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8478 assert_eq!(body["error"]["retry_after_seconds"], 1);
8479 let message = body["error"]["message"].as_str().expect("message");
8480 assert!(message.contains("512"), "{message}");
8481 }
8482
8483 #[test]
8484 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8485 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8486 token: 99,
8487 vocab_size: 32,
8488 });
8489 assert!(
8490 body["error"]["retry_after_seconds"].is_null(),
8491 "retrying a prompt this model cannot tokenize never helps"
8492 );
8493 }
8494
8495 #[test]
8496 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8497 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8498 token: 99,
8499 vocab_size: 32,
8500 });
8501 assert_eq!(status, StatusCode::BAD_REQUEST);
8502 }
8503
8504 #[test]
8505 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8506 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8508 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8509 let config = generate::KvPoolConfig {
8510 pool: pool.clone(),
8511 queue_wait: Duration::ZERO,
8512 };
8513
8514 let (_, finish, _usage) = run_generation(
8515 &model,
8516 &prompt,
8517 &greedy_params(4),
8518 Some(&config),
8519 None,
8520 None,
8521 None,
8522 None,
8523 None,
8524 )
8525 .unwrap();
8526 assert_eq!(finish, FinishReason::Length);
8527 assert_eq!(
8528 pool.lock().unwrap().free_blocks(),
8529 2,
8530 "a completed request must return its blocks to the pool"
8531 );
8532 }
8533
8534 #[tokio::test]
8539 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8540 let model = Arc::new(test_model());
8541 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8542
8543 let mut handles = Vec::new();
8544 for _ in 0..8 {
8545 let model = Arc::clone(&model);
8546 let prompt = prompt.clone();
8547 handles.push(tokio::task::spawn_blocking(move || {
8548 run_generation(
8549 &model,
8550 &prompt,
8551 &greedy_params(6),
8552 None,
8553 None,
8554 None,
8555 None,
8556 None,
8557 None,
8558 )
8559 .unwrap()
8560 }));
8561 }
8562
8563 let mut results = Vec::new();
8564 for h in handles {
8565 results.push(h.await.unwrap());
8566 }
8567 for r in &results[1..] {
8572 assert_eq!(r.0, results[0].0, "decoded chunks must match");
8573 assert_eq!(r.1, results[0].1, "finish reason must match");
8574 assert_eq!(
8575 r.2.prompt_tokens, results[0].2.prompt_tokens,
8576 "prompt token count must match"
8577 );
8578 assert_eq!(
8579 r.2.completion_tokens, results[0].2.completion_tokens,
8580 "completion token count must match"
8581 );
8582 }
8583 }
8584
8585 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8592 let mut header_entries = Vec::new();
8593 let mut data = Vec::new();
8594 for (name, shape, values) in tensors {
8595 let start = data.len();
8596 for v in values {
8597 data.extend_from_slice(&v.to_le_bytes());
8598 }
8599 let end = data.len();
8600 let shape_str = shape
8601 .iter()
8602 .map(|d| d.to_string())
8603 .collect::<Vec<_>>()
8604 .join(",");
8605 header_entries.push(format!(
8606 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8607 ));
8608 }
8609 let header = format!("{{{}}}", header_entries.join(","));
8610 let header_bytes = header.as_bytes();
8611 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8612 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8613 out.extend_from_slice(header_bytes);
8614 out.extend_from_slice(&data);
8615 out
8616 }
8617
8618 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8629 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8630 use frink_models::kimi_loader::KimiRealHparams;
8631 use frink_moe::{GatingFunction, MoeLayerConfig};
8632
8633 let hidden_dim = 8;
8634 let kda_num_heads = 2;
8635 let kda_head_dim = 3;
8636 let kda_proj = kda_num_heads * kda_head_dim;
8637 let conv_kernel = 4;
8638 let dense_intermediate = 5;
8639 let vocab_size = 256;
8643 let mla_num_heads = 1;
8644 let mla_q_lora_rank = 2;
8645 let mla_kv_lora_rank = 2;
8646 let mla_qk_nope_head_dim = 2;
8647 let mla_qk_rope_head_dim = 2;
8648 let mla_v_head_dim = 2;
8649
8650 let model_cfg = frink_models::ModelConfig {
8651 rope_layers: frink_models::rope_layers::RopeLayers::All,
8652 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
8653 name: "synthetic-kimi-server-test",
8654 n_layers: 1,
8655 n_mtp_blocks: 0,
8656 hidden_dim,
8657 n_heads: 1,
8658 n_kv_heads: 1,
8659 head_dim: 4,
8660 v_head_dim: None,
8661 vocab_size,
8662 rope_theta: 10000.0,
8663 rms_norm_eps: 1e-5,
8664 post_norm_eps: 1e-5,
8665 sliding_window: None,
8666 moe: MoeLayerConfig {
8667 expert_weights_scale: 1.0,
8668 routed_weight_before_ffn: false,
8669 n_experts: 1,
8670 n_experts_active: 1,
8671 n_shared_experts: 0,
8672 hidden_dim,
8673 expert_ffn_dim: 4,
8674 gating: GatingFunction::Sigmoid,
8675 norm_topk_prob: true,
8676 expert_group_count: None,
8677 expert_group_used_count: None,
8678 },
8679 n_dense_leading_layers: 1,
8684 moe_interleave_step: None,
8685 norm_function: frink_models::norm::NormFunction::Rms,
8686 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8687 kda_layers: vec![1],
8688 full_attn_layers: vec![],
8689 mla: MlaConfig {
8690 num_heads: mla_num_heads,
8691 q_lora_rank: mla_q_lora_rank,
8692 kv_lora_rank: mla_kv_lora_rank,
8693 qk_nope_head_dim: mla_qk_nope_head_dim,
8694 qk_rope_head_dim: mla_qk_rope_head_dim,
8695 v_head_dim: mla_v_head_dim,
8696 use_output_gate: true,
8697 rope: None,
8698 },
8699 kda: KdaConfig {
8700 num_heads: kda_num_heads,
8701 head_dim: kda_head_dim,
8702 short_conv_kernel_size: conv_kernel,
8703 gate_lower_bound: -5.0,
8704 use_full_rank_gate: true,
8705 },
8706 }),
8707 rope_freqs: None,
8708 rope_attn_factor: 1.0,
8709 rope_dim: None,
8710 rope_dim_swa: None,
8711 rope_freqs_long: None,
8712 rope_freqs_short: None,
8713 rope_orig_ctx: None,
8714 rope_layout: frink_models::config::RopeLayout::Neox,
8715 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
8716 swa_layers: frink_models::swa_layers::SwaLayers::All,
8717 attn_logit_softcap: None,
8718 final_logit_softcap: None,
8719 embedding_scale: None,
8720 residual_scale: None,
8721 clamp_kqv: None,
8722 attn_temperature: None,
8723 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
8724 block_sub_norms: false,
8725 parallel_residual: false,
8726 learned_positions: false,
8727 attn_value_scale: None,
8728 alibi_max_bias: None,
8729 layer_loops: None,
8730 skip_stream: false,
8731 parallel_ssm: false,
8732 swa_chunked: false,
8733 weightless_qk_norm: false,
8734 logit_multiplier: None,
8735 attention_scale: None,
8736 rope_theta_swa: None,
8737 ffn_activation: frink_models::config::FfnActivation::Swiglu,
8738 best_effort_fields: &["synthetic test config, not a real preset"],
8739 };
8740 let hp = KimiRealHparams {
8741 hidden_dim,
8742 kda_num_heads,
8743 kda_head_dim,
8744 mla_num_heads,
8745 mla_q_lora_rank,
8746 mla_kv_lora_rank,
8747 mla_qk_nope_head_dim,
8748 mla_qk_rope_head_dim,
8749 mla_v_head_dim,
8750 dense_intermediate_dim: dense_intermediate,
8751 moe_hidden_dim: hidden_dim,
8752 moe_intermediate_dim: 4,
8753 n_experts: 1,
8754 num_shared_experts: 0,
8755 };
8756
8757 let prefix = "language_model.model.layers.0";
8761 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8762 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8763 name: String,
8764 shape: Vec<usize>,
8765 n: usize| {
8766 tensors.push((name, shape, vec![0.01f32; n]));
8767 };
8768 push(
8769 &mut tensors,
8770 format!("{prefix}.input_layernorm.weight"),
8771 vec![hidden_dim],
8772 hidden_dim,
8773 );
8774 push(
8775 &mut tensors,
8776 format!("{prefix}.post_attention_layernorm.weight"),
8777 vec![hidden_dim],
8778 hidden_dim,
8779 );
8780 push(
8781 &mut tensors,
8782 format!("{prefix}.self_attention_res_norm.weight"),
8783 vec![hidden_dim],
8784 hidden_dim,
8785 );
8786 push(
8787 &mut tensors,
8788 format!("{prefix}.self_attention_res_proj.weight"),
8789 vec![1, hidden_dim],
8790 hidden_dim,
8791 );
8792 push(
8793 &mut tensors,
8794 format!("{prefix}.mlp_res_norm.weight"),
8795 vec![hidden_dim],
8796 hidden_dim,
8797 );
8798 push(
8799 &mut tensors,
8800 format!("{prefix}.mlp_res_proj.weight"),
8801 vec![1, hidden_dim],
8802 hidden_dim,
8803 );
8804 push(
8805 &mut tensors,
8806 format!("{prefix}.self_attn.q_proj.weight"),
8807 vec![kda_proj, hidden_dim],
8808 kda_proj * hidden_dim,
8809 );
8810 push(
8811 &mut tensors,
8812 format!("{prefix}.self_attn.k_proj.weight"),
8813 vec![kda_proj, hidden_dim],
8814 kda_proj * hidden_dim,
8815 );
8816 push(
8817 &mut tensors,
8818 format!("{prefix}.self_attn.v_proj.weight"),
8819 vec![kda_proj, hidden_dim],
8820 kda_proj * hidden_dim,
8821 );
8822 push(
8823 &mut tensors,
8824 format!("{prefix}.self_attn.q_conv1d.weight"),
8825 vec![kda_proj, 1, conv_kernel],
8826 kda_proj * conv_kernel,
8827 );
8828 push(
8829 &mut tensors,
8830 format!("{prefix}.self_attn.k_conv1d.weight"),
8831 vec![kda_proj, 1, conv_kernel],
8832 kda_proj * conv_kernel,
8833 );
8834 push(
8835 &mut tensors,
8836 format!("{prefix}.self_attn.v_conv1d.weight"),
8837 vec![kda_proj, 1, conv_kernel],
8838 kda_proj * conv_kernel,
8839 );
8840 push(
8841 &mut tensors,
8842 format!("{prefix}.self_attn.A_log"),
8843 vec![kda_num_heads],
8844 kda_num_heads,
8845 );
8846 push(
8847 &mut tensors,
8848 format!("{prefix}.self_attn.f_a_proj.weight"),
8849 vec![kda_head_dim, hidden_dim],
8850 kda_head_dim * hidden_dim,
8851 );
8852 push(
8853 &mut tensors,
8854 format!("{prefix}.self_attn.f_b_proj.weight"),
8855 vec![kda_proj, kda_head_dim],
8856 kda_proj * kda_head_dim,
8857 );
8858 push(
8859 &mut tensors,
8860 format!("{prefix}.self_attn.dt_bias"),
8861 vec![kda_proj],
8862 kda_proj,
8863 );
8864 push(
8865 &mut tensors,
8866 format!("{prefix}.self_attn.b_proj.weight"),
8867 vec![kda_num_heads, hidden_dim],
8868 kda_num_heads * hidden_dim,
8869 );
8870 push(
8871 &mut tensors,
8872 format!("{prefix}.self_attn.g_proj.weight"),
8873 vec![kda_proj, hidden_dim],
8874 kda_proj * hidden_dim,
8875 );
8876 push(
8877 &mut tensors,
8878 format!("{prefix}.self_attn.o_norm.weight"),
8879 vec![kda_head_dim],
8880 kda_head_dim,
8881 );
8882 push(
8883 &mut tensors,
8884 format!("{prefix}.self_attn.o_proj.weight"),
8885 vec![hidden_dim, kda_proj],
8886 hidden_dim * kda_proj,
8887 );
8888 push(
8889 &mut tensors,
8890 format!("{prefix}.mlp.gate_proj.weight"),
8891 vec![dense_intermediate, hidden_dim],
8892 dense_intermediate * hidden_dim,
8893 );
8894 push(
8895 &mut tensors,
8896 format!("{prefix}.mlp.up_proj.weight"),
8897 vec![dense_intermediate, hidden_dim],
8898 dense_intermediate * hidden_dim,
8899 );
8900 push(
8901 &mut tensors,
8902 format!("{prefix}.mlp.down_proj.weight"),
8903 vec![hidden_dim, dense_intermediate],
8904 hidden_dim * dense_intermediate,
8905 );
8906 push(
8907 &mut tensors,
8908 "language_model.model.embed_tokens.weight".to_string(),
8909 vec![vocab_size, hidden_dim],
8910 vocab_size * hidden_dim,
8911 );
8912 push(
8913 &mut tensors,
8914 "language_model.lm_head.weight".to_string(),
8915 vec![vocab_size, hidden_dim],
8916 vocab_size * hidden_dim,
8917 );
8918 push(
8919 &mut tensors,
8920 "language_model.model.norm.weight".to_string(),
8921 vec![hidden_dim],
8922 hidden_dim,
8923 );
8924 push(
8925 &mut tensors,
8926 "language_model.model.output_attn_res_norm.weight".to_string(),
8927 vec![hidden_dim],
8928 hidden_dim,
8929 );
8930 push(
8931 &mut tensors,
8932 "language_model.model.output_attn_res_proj.weight".to_string(),
8933 vec![1, hidden_dim],
8934 hidden_dim,
8935 );
8936
8937 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
8950 let dir = std::env::temp_dir().join(format!(
8951 "frink_server_kimi_e2e_test_{}_{}",
8952 std::process::id(),
8953 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
8954 ));
8955 std::fs::create_dir_all(&dir).unwrap();
8956 let shard_bytes = write_safetensors_shard(&tensors);
8957 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
8958 let map_entries: Vec<String> = tensors
8959 .iter()
8960 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
8961 .collect();
8962 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
8963 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
8964
8965 use base64::Engine;
8969 let vocab_lines: Vec<String> = (0..vocab_size as u32)
8970 .map(|b| {
8971 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
8972 format!("{b64} {b}")
8973 })
8974 .collect();
8975 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
8976
8977 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
8978 .expect("must load the synthetic Kimi checkpoint end to end");
8979 std::fs::remove_dir_all(&dir).ok();
8980 loaded
8981 }
8982
8983 #[test]
8991 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
8992 let loaded = build_synthetic_kimi_loaded();
8993 let state = build_app_state(
8994 StartupModels {
8995 loaded: model::LoadedModel::Kimi(loaded),
8996 embedding: None,
8997 },
8998 None,
8999 None,
9000 None,
9001 false,
9002 None,
9003 Arc::new(health::Detection::ready(health::probe_backends())),
9004 );
9005 let active = state.active().expect("a freshly built state has a model");
9006 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9007 assert!(!active.is_synthetic());
9008
9009 let (_chunks, finish, _usage) = run_generation(
9010 active.generative().unwrap(),
9011 "hi",
9012 &greedy_params(5),
9013 None,
9014 None,
9015 None,
9016 None,
9017 None,
9018 None,
9019 )
9020 .expect("a real Kimi checkpoint must generate without error");
9021 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9022 }
9023
9024 #[test]
9037 fn a_grammar_constrains_the_engine_decode_path() {
9038 let loaded = build_synthetic_kimi_loaded();
9039 let state = build_app_state(
9040 StartupModels {
9041 loaded: model::LoadedModel::Kimi(loaded),
9042 embedding: None,
9043 },
9044 None,
9045 None,
9046 None,
9047 false,
9048 None,
9049 Arc::new(health::Detection::ready(health::probe_backends())),
9050 );
9051 let active = state.active().expect("a freshly built state has a model");
9052
9053 let run = |grammar: Option<&str>| {
9054 let mut params = greedy_params(6);
9055 params.grammar = grammar.map(|src| {
9056 Arc::new(
9057 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9058 .expect("test grammar parses"),
9059 )
9060 });
9061 run_generation(
9062 active.generative().unwrap(),
9063 "hi",
9064 ¶ms,
9065 None,
9066 None,
9067 None,
9068 None,
9069 None,
9070 None,
9071 )
9072 };
9073
9074 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
9075 let unconstrained = chunks.concat();
9076 assert!(
9077 unconstrained.chars().any(|c| c != 'a'),
9078 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9079 constrained run below would prove nothing"
9080 );
9081
9082 let (chunks, finish, _) =
9083 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9084 let constrained = chunks.concat();
9085 assert!(
9086 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9087 "the engine decode path served text its grammar forbids ({constrained:?}): \
9088 the constraint was dropped between `generate_engine` and the sampler"
9089 );
9090 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9091 }
9092
9093 #[test]
9103 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9104 let loaded = build_synthetic_kimi_loaded();
9105 let state = build_app_state(
9106 StartupModels {
9107 loaded: model::LoadedModel::Kimi(loaded),
9108 embedding: None,
9109 },
9110 None,
9111 None,
9112 None,
9113 false,
9114 None,
9115 Arc::new(health::Detection::ready(health::probe_backends())),
9116 );
9117
9118 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9119 let kv_pool_config = generate::KvPoolConfig {
9120 pool: pool.clone(),
9121 queue_wait: Duration::ZERO,
9122 };
9123 let pc = Mutex::new(PrefixCache::new(4));
9124
9125 run_generation(
9126 state
9127 .active()
9128 .expect("a freshly built state has a model")
9129 .generative()
9130 .unwrap(),
9131 "hi",
9132 &greedy_params(5),
9133 Some(&kv_pool_config),
9134 None,
9135 Some(&pc),
9136 None,
9137 None,
9138 None,
9139 )
9140 .expect("a real Kimi checkpoint must generate without error");
9141
9142 assert_eq!(
9143 pool.lock().unwrap().free_blocks(),
9144 4,
9145 "the KV pool must be completely untouched by a Kimi request"
9146 );
9147 let stats = pc.lock().unwrap().stats();
9148 assert_eq!(
9149 stats.hits + stats.misses,
9150 0,
9151 "the prefix cache must never be consulted for a Kimi request"
9152 );
9153 }
9154}