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 sampling_loop;
72mod security;
73mod serving;
74mod session;
75mod slots;
76mod sse;
77mod stats;
78mod stop;
79mod stream_events;
80mod tasks;
81mod tool_grammar;
82mod unimplemented_fields;
83mod unsupported_sampling;
84mod utf8_stream;
85
86use std::cell::RefCell;
87use std::convert::Infallible;
88use std::net::SocketAddr;
89use std::path::PathBuf;
90use std::rc::Rc;
91use std::sync::{Arc, Mutex, MutexGuard};
92use std::time::Duration;
93
94use axum::{
95 extract::State,
96 http::StatusCode,
97 response::sse::{Event, Sse},
98 response::{IntoResponse, Response},
99 routing::{get, post},
100 Json, Router,
101};
102use serde::{Deserialize, Serialize};
103
104use cli::apply_cli_overrides;
105pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
106
107use frink_core::cache::KvBlockPool;
108use frink_models::kimi_tokenizer::KimiTokenizer;
109use frink_models::sampling::SamplingParams;
110use frink_models::tokenizer::{SpecialTokens, StopTokens};
111use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
112use generate::{FinishReason, GenerationParams};
113pub(crate) use loaded::{ActiveModel, Loaded};
114use model::ServerTokenizer;
115use rerank::encoder_endpoints;
116use response_cache::ResponseCache;
117use sampling_knobs::SamplingKnobs;
118
119#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
128 Gguf(GgufModel),
129 Kimi(KimiModel),
130 Mla(MlaModel),
131 Gemma4(Gemma4Model),
132 Glm52(Glm52Model),
133}
134
135pub(crate) struct GgufModel {
136 decoder: Arc<Decoder>,
137 tokenizer: Arc<ServerTokenizer>,
138 stop_tokens: StopTokens,
139 bos_id: Option<usize>,
140 is_synthetic: bool,
141 chat_template: chat_template::PromptTemplate,
142}
143
144pub(crate) struct KimiModel {
145 engine: KimiEngine,
146 tokenizer: KimiTokenizer,
147 stop_tokens: StopTokens,
148 chat_template: chat_template::PromptTemplate,
149}
150
151pub(crate) struct MlaModel {
152 engine: MlaEngine,
153 tokenizer: ServerTokenizer,
154 stop_tokens: StopTokens,
155 bos_id: Option<usize>,
156 name: String,
157 chat_template: chat_template::PromptTemplate,
158}
159
160pub(crate) struct Gemma4Model {
161 engine: Gemma4Engine,
162 tokenizer: ServerTokenizer,
163 stop_tokens: StopTokens,
164 bos_id: Option<usize>,
165 name: String,
166 chat_template: chat_template::PromptTemplate,
167}
168
169pub(crate) struct Glm52Model {
170 engine: frink_models::Glm52Engine,
171 tokenizer: ServerTokenizer,
172 stop_tokens: StopTokens,
173 bos_id: Option<usize>,
174 name: String,
175 chat_template: chat_template::PromptTemplate,
176}
177
178impl Model {
179 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
180 match self {
181 Model::Gguf(m) => m.chat_template.clone(),
182 Model::Kimi(m) => m.chat_template.clone(),
183 Model::Mla(m) => m.chat_template.clone(),
184 Model::Gemma4(m) => m.chat_template.clone(),
185 Model::Glm52(m) => m.chat_template.clone(),
186 }
187 }
188
189 fn is_synthetic(&self) -> bool {
194 match self {
195 Model::Gguf(m) => m.is_synthetic,
196 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
197 }
198 }
199
200 fn tokenizer_kind(&self) -> &'static str {
201 match self {
202 Model::Gguf(m) => m.tokenizer.kind(),
203 Model::Kimi(_) => "kimi-tiktoken-bpe",
204 Model::Mla(m) => m.tokenizer.kind(),
205 Model::Gemma4(m) => m.tokenizer.kind(),
206 Model::Glm52(m) => m.tokenizer.kind(),
207 }
208 }
209
210 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
214 match self {
215 Model::Gguf(m) => m.decoder.expert_store_stats(),
216 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
217 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
218 }
219 }
220
221 pub(crate) fn name(&self) -> &str {
222 match self {
223 Model::Gguf(m) => m.decoder.config.name,
224 Model::Kimi(_) => "kimi-k3",
225 Model::Mla(m) => m.name.as_str(),
226 Model::Gemma4(m) => m.name.as_str(),
227 Model::Glm52(m) => m.name.as_str(),
228 }
229 }
230
231 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
258 match self {
259 Model::Gguf(m) => m.tokenizer.encode(text, specials),
260 Model::Kimi(m) => m
261 .tokenizer
262 .encode(text, specials)
263 .into_iter()
264 .map(|id| id as usize)
265 .collect(),
266 Model::Mla(m) => m.tokenizer.encode(text, specials),
267 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
268 Model::Glm52(m) => m.tokenizer.encode(text, specials),
269 }
270 }
271
272 pub(crate) fn bos_id(&self) -> Option<usize> {
281 match self {
282 Model::Gguf(m) => m.bos_id,
283 Model::Kimi(_) => None,
284 Model::Mla(m) => m.bos_id,
285 Model::Gemma4(m) => m.bos_id,
286 Model::Glm52(m) => m.bos_id,
287 }
288 }
289
290 pub(crate) fn decode(&self, ids: &[usize]) -> String {
291 match self {
292 Model::Gguf(m) => m.tokenizer.decode(ids),
293 Model::Kimi(m) => {
294 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
295 m.tokenizer.decode(&ids32)
296 }
297 Model::Mla(m) => m.tokenizer.decode(ids),
298 Model::Gemma4(m) => m.tokenizer.decode(ids),
299 Model::Glm52(m) => m.tokenizer.decode(ids),
300 }
301 }
302
303 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
306 match self {
307 Model::Gguf(m) => {
308 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
309 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
310 }
311 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
312 }
313 }
314
315 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
321 match self {
322 Model::Gguf(m) => Some(&m.decoder),
323 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
324 }
325 }
326
327 pub(crate) fn vocab_size(&self) -> Option<usize> {
328 match self {
329 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
330 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
331 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
332 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
333 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
334 }
335 }
336
337 fn has_real_vocabulary(&self) -> bool {
344 match self {
345 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
346 Model::Kimi(_) => true,
347 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
348 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
349 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
350 }
351 }
352}
353
354impl frink_models::dry::DryVocab for Model {
360 fn n_tokens(&self) -> usize {
361 self.vocab_size().unwrap_or(0)
362 }
363
364 fn detokenize(&self, token: usize) -> String {
365 self.decode(&[token])
366 }
367
368 fn tokenize(&self, text: &str) -> Vec<usize> {
369 self.encode(text, SpecialTokens::AsText)
370 }
371}
372
373pub(crate) struct AppState {
374 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
388 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
407 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
411 pub(crate) tasks: Arc<tasks::TaskRegistry>,
413 pub(crate) cancels: Arc<cancel::CancelRegistry>,
417 pub(crate) stats: stats::Stats,
420 pub(crate) streams: resume::StreamRegistry,
423 pub(crate) model_dir: Option<PathBuf>,
425 response_cache: Mutex<ResponseCache>,
429 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
441 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
454 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
464 sessions: session::SessionStore,
470 requests_total: std::sync::atomic::AtomicU64,
471 request_errors_total: std::sync::atomic::AtomicU64,
472 started_at: std::time::Instant,
473 last_request_ms: std::sync::atomic::AtomicU64,
478 detection: Arc<health::Detection>,
480 mcp: Option<mcp::LoadedMcpConfig>,
482 pub(crate) continuous_batching_enabled: bool,
486 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
490 loading_model: Mutex<Option<String>>,
494 last_load_error: Mutex<Option<(String, String)>>,
498 pub(crate) serving: Mutex<crate::stats::ServingStats>,
503 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
508 pub(crate) footprint:
514 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
515 pub(crate) started_unix: u64,
525}
526
527pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
533
534impl AppState {
535 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
539 self.active
540 .read()
541 .unwrap_or_else(|p| p.into_inner())
542 .clone()
543 }
544
545 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
551 self.active().ok_or_else(|| {
552 (
553 StatusCode::SERVICE_UNAVAILABLE,
554 Json(serde_json::json!({"error": {
555 "message": "no model is loaded; POST /admin/models/load with an id from \
556 GET /admin/models",
557 "type": "model_not_loaded"
558 }})),
559 )
560 })
561 }
562
563 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
571 Ok(Arc::clone(self.require_active()?.generative()?))
572 }
573
574 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
582 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
583 std::mem::replace(&mut *guard, next)
584 }
585
586 fn mark_request_finished(&self) {
590 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
591 self.last_request_ms
592 .store(ms, std::sync::atomic::Ordering::Relaxed);
593 }
594
595 pub(crate) fn uptime(&self) -> Duration {
596 self.started_at.elapsed()
597 }
598
599 pub(crate) fn requests_total(&self) -> u64 {
600 self.requests_total
601 .load(std::sync::atomic::Ordering::Relaxed)
602 }
603
604 pub(crate) fn errors_total(&self) -> u64 {
605 self.request_errors_total
606 .load(std::sync::atomic::Ordering::Relaxed)
607 }
608
609 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
610 lock_cache(&self.response_cache).stats()
611 }
612
613 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
616 let last = self
617 .last_request_ms
618 .load(std::sync::atomic::Ordering::Relaxed);
619 (last > 0)
620 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
621 .map(|age| age.max(0.0))
622 }
623
624 pub(crate) fn loading_model_id(&self) -> Option<String> {
625 self.loading_model
626 .lock()
627 .unwrap_or_else(|p| p.into_inner())
628 .clone()
629 }
630
631 pub(crate) fn set_loading_model(&self, id: Option<String>) {
632 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
633 }
634
635 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
636 self.last_load_error
637 .lock()
638 .unwrap_or_else(|p| p.into_inner())
639 .clone()
640 }
641
642 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
643 *self
644 .last_load_error
645 .lock()
646 .unwrap_or_else(|p| p.into_inner()) = error;
647 }
648
649 pub(crate) fn active_model_name(&self) -> Option<String> {
658 self.active().map(|a| a.name().to_string())
659 }
660
661 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
671 self.embedding
672 .clone()
673 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
674 }
675
676 pub(crate) fn embedding_model_name(&self) -> Option<String> {
680 match self.embedding_model() {
681 Some(e) => Some(e.name().to_string()),
682 None => self.active_model_name(),
683 }
684 }
685
686 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
687 self.stats.record(stats::entry(record));
688 }
689}
690
691fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
696 cache
697 .lock()
698 .unwrap_or_else(|poisoned| poisoned.into_inner())
699}
700
701#[derive(Debug, Clone, Deserialize)]
702#[serde(untagged)]
703pub(crate) enum MessageContent {
704 Text(String),
705 Parts(Vec<ContentPart>),
706}
707
708#[derive(Debug, Clone, Deserialize)]
709struct ContentPart {
710 #[serde(rename = "type")]
711 kind: String,
712 #[serde(default)]
713 text: Option<String>,
714 #[serde(default)]
715 image_url: Option<serde_json::Value>,
716}
717
718impl MessageContent {
719 fn as_text(&self) -> String {
720 match self {
721 Self::Text(s) => s.clone(),
722 Self::Parts(parts) => parts
723 .iter()
724 .filter_map(|p| p.text.as_deref())
725 .collect::<Vec<_>>()
726 .join(""),
727 }
728 }
729
730 fn has_image(&self) -> bool {
731 match self {
732 Self::Text(_) => false,
733 Self::Parts(parts) => parts
734 .iter()
735 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
736 }
737 }
738}
739
740#[derive(Debug, Clone, Deserialize)]
741pub(crate) struct ChatMessage {
742 pub(crate) role: String,
743 #[serde(default)]
747 pub(crate) content: Option<MessageContent>,
748 #[serde(default)]
752 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
753 #[serde(default)]
758 #[allow(dead_code)]
759 pub(crate) tool_call_id: Option<String>,
760 #[serde(default, alias = "reasoning")]
776 pub(crate) reasoning_content: Option<String>,
777}
778
779impl ChatMessage {
780 fn rendered_content(&self) -> String {
787 let mut out = self
788 .content
789 .as_ref()
790 .map(MessageContent::as_text)
791 .unwrap_or_default();
792 if let Some(calls) = &self.tool_calls {
793 for call in calls {
794 out.push_str(&format!(
795 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
796 call.function.name, call.function.arguments
797 ));
798 }
799 }
800 out
801 }
802}
803
804#[derive(Debug, Clone, Deserialize)]
805pub(crate) struct ToolCallIn {
806 #[serde(default)]
807 #[allow(dead_code)]
808 id: String,
809 #[serde(rename = "type", default)]
810 #[allow(dead_code)]
811 kind: String,
812 function: ToolCallFunctionIn,
813}
814
815#[derive(Debug, Clone, Deserialize)]
816struct ToolCallFunctionIn {
817 name: String,
818 arguments: String,
823}
824
825#[derive(Debug, Clone, Deserialize)]
828struct ToolDef {
829 #[serde(rename = "type", default)]
830 #[allow(dead_code)]
831 kind: String,
832 function: ToolFunctionDef,
833}
834
835#[derive(Debug, Clone, Deserialize)]
836struct ToolFunctionDef {
837 name: String,
838 #[serde(default)]
839 description: Option<String>,
840 #[serde(default)]
841 parameters: Option<serde_json::Value>,
842}
843
844#[derive(Debug, Clone, Deserialize)]
855#[serde(untagged)]
856enum ToolChoice {
857 Mode(String),
858 Specific(serde_json::Value),
859}
860
861#[derive(Deserialize)]
864#[serde(untagged)]
865enum StopParam {
866 One(String),
867 Many(Vec<String>),
868}
869
870#[derive(Deserialize)]
871struct ChatCompletionRequest {
872 model: String,
873 messages: Vec<ChatMessage>,
874 #[serde(default = "default_max_tokens")]
875 max_tokens: usize,
876 #[serde(default)]
877 temperature: Option<f32>,
878 #[serde(default)]
879 top_p: Option<f32>,
880 #[serde(default)]
885 min_p: Option<f32>,
886 #[serde(default)]
887 top_k: Option<usize>,
888 #[serde(default)]
889 repetition_penalty: Option<f32>,
890 #[serde(flatten)]
894 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
895 #[serde(flatten)]
899 unimplemented: crate::unimplemented_fields::UnimplementedFields,
900 #[serde(default)]
901 seed: Option<u64>,
902 #[serde(default)]
903 stop: Option<StopParam>,
904 #[serde(default)]
905 stream: Option<bool>,
906 #[serde(default)]
919 stream_resumable: Option<bool>,
920 #[serde(default)]
931 ignore_eos: Option<bool>,
932 #[serde(default)]
933 tools: Vec<ToolDef>,
934 #[serde(default)]
935 tool_choice: Option<ToolChoice>,
936 #[serde(default)]
944 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
945 #[serde(default)]
953 reasoning_effort: Option<String>,
954 #[serde(default)]
958 thinking: Option<ThinkingSwitch>,
959 #[serde(default)]
964 session_id: Option<String>,
965 #[serde(default, deserialize_with = "continuation::deserialize")]
973 continue_final_message: continuation::ContinueFinalMessage,
974 #[serde(default, alias = "thinking_budget_tokens")]
982 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
983 #[serde(default)]
985 logprobs: Option<bool>,
986 #[serde(default)]
987 top_logprobs: Option<u32>,
988 #[serde(default)]
989 presence_penalty: Option<f32>,
990 #[serde(default)]
991 frequency_penalty: Option<f32>,
992 #[serde(default)]
993 response_format: Option<serde_json::Value>,
994 #[serde(default)]
1001 logit_bias: Option<serde_json::Value>,
1002 #[serde(default)]
1006 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1007 #[serde(default)]
1017 samplers: Option<serde_json::Value>,
1018 #[serde(default)]
1028 grammar: Option<String>,
1029}
1030
1031const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1043
1044#[derive(Debug, Clone, Deserialize)]
1046pub(crate) struct ThinkingSwitch {
1047 #[serde(rename = "type")]
1048 pub(crate) kind: String,
1049}
1050
1051const THINKING_KWARG_KEYS: [&str; 4] = [
1055 "enable_thinking",
1056 "thinking",
1057 "thinking_mode",
1058 "reasoning_effort",
1059];
1060
1061const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1065
1066fn default_max_tokens() -> usize {
1067 DEFAULT_CHAT_MAX_TOKENS
1068}
1069
1070impl ChatCompletionRequest {
1071 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1080 let mut knobs = SamplingKnobs {
1081 temperature: self.temperature,
1082 top_p: self.top_p,
1083 min_p: self.min_p,
1084 top_k: self.top_k,
1085 repetition_penalty: self.repetition_penalty,
1086 presence_penalty: self.presence_penalty,
1087 frequency_penalty: self.frequency_penalty,
1088 penalty_last_n: None,
1092 sampler_order: unsupported_sampling::parse_sampler_order(
1093 self.samplers.as_ref(),
1094 "/v1/chat/completions",
1095 )?,
1096 ..SamplingKnobs::default()
1097 };
1098 self.extra_samplers.apply(&mut knobs);
1099 Ok(knobs)
1100 }
1101
1102 fn sampling_params(
1103 &self,
1104 model: crate::sampling_knobs::SamplerModel<'_>,
1105 ) -> Result<SamplingParams, ApiError> {
1106 self.sampling_knobs()?.resolve(model).map_err(|e| {
1107 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1108 })
1109 }
1110
1111 fn stop_sequences(&self) -> Vec<String> {
1112 self.stop
1113 .as_ref()
1114 .map(|s| match s {
1115 StopParam::One(v) => vec![v.clone()],
1116 StopParam::Many(v) => v.clone(),
1117 })
1118 .unwrap_or_default()
1119 }
1120
1121 fn tools_active(&self) -> bool {
1126 !self.tools.is_empty()
1127 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1128 }
1129
1130 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1142 match &self.tool_choice {
1143 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1144 Some(ToolChoice::Specific(value)) => {
1145 let name = value
1149 .get("function")
1150 .and_then(|f| f.get("name"))
1151 .or_else(|| value.get("name"))
1152 .and_then(|n| n.as_str());
1153 match name {
1154 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1155 None => Err(invalid_request(
1156 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1157 function.name",
1158 "tool_choice",
1159 )),
1160 }
1161 }
1162 _ => Ok(None),
1163 }
1164 }
1165
1166 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1168 self.tools
1169 .iter()
1170 .map(|t| tool_grammar::ToolSpec {
1171 name: &t.function.name,
1172 parameters: t.function.parameters.as_ref(),
1173 })
1174 .collect()
1175 }
1176
1177 fn thinking_direction(&self) -> Option<bool> {
1220 if let Some(switch) = &self.thinking {
1221 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1222 "disabled" => Some(false),
1223 "enabled" => Some(true),
1224 _ => None,
1227 };
1228 }
1229 let effort = self.reasoning_effort.as_ref()?;
1230 DISABLE_EFFORTS
1231 .contains(&effort.trim().to_ascii_lowercase().as_str())
1232 .then_some(false)
1233 }
1234
1235 fn resolve_template_kwargs(
1236 &self,
1237 template: &chat_template::PromptTemplate,
1238 ) -> serde_json::Map<String, serde_json::Value> {
1239 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1240 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1244
1245 if !caller_steered {
1246 match self.thinking_direction() {
1247 Some(false) => {
1248 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1249 kwargs.insert(k, v);
1250 }
1251 return kwargs;
1254 }
1255 Some(true) => {
1256 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1257 kwargs.insert(k, v);
1258 }
1259 }
1260 None => {}
1261 }
1262 if let Some(effort) = &self.reasoning_effort {
1263 kwargs
1264 .entry("reasoning_effort".to_string())
1265 .or_insert_with(|| serde_json::json!(effort));
1266 }
1267 }
1268
1269 let offered: Vec<serde_json::Value> = if self.tools_active() {
1270 self.tools.iter().map(chat_template::tool_json).collect()
1271 } else {
1272 Vec::new()
1273 };
1274 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1275 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1276 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1277 kwargs.entry(k).or_insert(v);
1278 }
1279 }
1280 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1281 crate::policy::effort::EffortMapping::Mapped(to) => {
1282 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1283 }
1284 crate::policy::effort::EffortMapping::Dropped => {
1285 tracing::debug!(
1286 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1287 enough, so its own default applies"
1288 );
1289 }
1290 crate::policy::effort::EffortMapping::Unchanged => {}
1291 }
1292 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1293 kwargs
1294 }
1295
1296 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1299 if self.max_tokens == 0 {
1306 return Err(invalid_request(
1307 "max_tokens must be at least 1",
1308 "max_tokens",
1309 ));
1310 }
1311 if let Some(switch) = &self.thinking {
1316 let kind = switch.kind.trim().to_ascii_lowercase();
1317 if kind != "enabled" && kind != "disabled" {
1318 return Err(invalid_request(
1319 "thinking.type must be \"enabled\" or \"disabled\"",
1320 "thinking.type",
1321 ));
1322 }
1323 }
1324 for msg in &self.messages {
1325 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1326 return Err(unsupported_feature(
1327 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1328 ));
1329 }
1330 }
1331 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1332 return Err(unsupported_feature(
1333 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1334 ));
1335 }
1336 self.unimplemented.refuse("/v1/chat/completions")?;
1340 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1341 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1345 let stated_grammar =
1360 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1361 if let Some(forced) = self.forced_tool_choice()? {
1369 if self.tools.is_empty() {
1370 return Err(invalid_request(
1371 "tool_choice forces a tool call, but no tools were offered",
1372 "tool_choice",
1373 ));
1374 }
1375 if let tool_grammar::Forced::Named(name) = forced {
1376 if !self.tools.iter().any(|t| t.function.name == name) {
1377 return Err(invalid_request(
1378 &format!(
1379 "tool_choice names {name:?}, which is not one of the tools offered"
1380 ),
1381 "tool_choice",
1382 ));
1383 }
1384 }
1385 if stated_grammar.is_some() {
1395 return Err(invalid_request(
1396 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1397 are two different constraints on the same generation; send one",
1398 "tool_choice",
1399 ));
1400 }
1401 if self.json_object_mode() {
1402 return Err(invalid_request(
1403 "a forced tool_choice cannot be combined with response_format json_object: \
1404 the tool-call markers are not JSON",
1405 "tool_choice",
1406 ));
1407 }
1408 }
1409 Ok(())
1410 }
1411
1412 fn effective_stop_sequences(&self) -> Vec<String> {
1419 let mut stop = self.stop_sequences();
1420 if self.tools_active() {
1421 stop.push("</tool_call>".to_string());
1422 }
1423 stop
1424 }
1425
1426 fn json_object_mode(&self) -> bool {
1427 self.response_format
1428 .as_ref()
1429 .and_then(|v| v.get("type"))
1430 .and_then(|v| v.as_str())
1431 == Some("json_object")
1432 }
1433}
1434
1435#[derive(Serialize)]
1436struct ChatCompletionChoice {
1437 index: usize,
1438 message: ChatCompletionResponseMessage,
1439 finish_reason: &'static str,
1440}
1441
1442#[derive(Serialize)]
1443struct ChatCompletionResponseMessage {
1444 role: &'static str,
1445 #[serde(skip_serializing_if = "Option::is_none")]
1446 content: Option<String>,
1447 #[serde(skip_serializing_if = "Option::is_none")]
1451 reasoning_content: Option<String>,
1452 #[serde(skip_serializing_if = "Option::is_none")]
1453 tool_calls: Option<Vec<ToolCallOut>>,
1454}
1455
1456#[derive(Serialize, Clone)]
1457struct ToolCallOut {
1458 id: String,
1459 #[serde(rename = "type")]
1460 kind: &'static str,
1461 function: ToolCallFunctionOut,
1462}
1463
1464#[derive(Serialize, Clone)]
1472struct ToolCallDelta {
1473 index: usize,
1474 #[serde(skip_serializing_if = "Option::is_none")]
1475 id: Option<String>,
1476 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1477 kind: Option<&'static str>,
1478 function: ToolCallFunctionDelta,
1479}
1480
1481#[derive(Serialize, Clone, Default)]
1482struct ToolCallFunctionDelta {
1483 #[serde(skip_serializing_if = "Option::is_none")]
1484 name: Option<String>,
1485 #[serde(skip_serializing_if = "Option::is_none")]
1488 arguments: Option<String>,
1489}
1490
1491impl ToolCallDelta {
1492 fn whole(index: usize, name: String, arguments: String) -> Self {
1494 ToolCallDelta {
1495 index,
1496 id: Some(format!("call_{index}")),
1497 kind: Some("function"),
1498 function: ToolCallFunctionDelta {
1499 name: Some(name),
1500 arguments: Some(arguments),
1501 },
1502 }
1503 }
1504
1505 fn opening(index: usize, name: String) -> Self {
1507 ToolCallDelta {
1508 index,
1509 id: Some(format!("call_{index}")),
1510 kind: Some("function"),
1511 function: ToolCallFunctionDelta {
1512 name: Some(name),
1513 arguments: Some(String::new()),
1514 },
1515 }
1516 }
1517
1518 fn arguments(index: usize, fragment: String) -> Self {
1520 ToolCallDelta {
1521 index,
1522 id: None,
1523 kind: None,
1524 function: ToolCallFunctionDelta {
1525 name: None,
1526 arguments: Some(fragment),
1527 },
1528 }
1529 }
1530}
1531
1532#[derive(Serialize, Clone)]
1533struct ToolCallFunctionOut {
1534 name: String,
1535 arguments: String,
1539}
1540
1541#[derive(Serialize)]
1542struct ChatCompletionResponse {
1543 id: String,
1544 request_id: String,
1550 object: &'static str,
1551 model: String,
1552 choices: Vec<ChatCompletionChoice>,
1553 usage: generate::Usage,
1558 frink_cache: &'static str,
1569}
1570
1571#[derive(Serialize)]
1572struct ChatCompletionChunkDelta {
1573 #[serde(skip_serializing_if = "Option::is_none")]
1574 role: Option<&'static str>,
1575 #[serde(skip_serializing_if = "Option::is_none")]
1576 content: Option<String>,
1577 #[serde(skip_serializing_if = "Option::is_none")]
1579 reasoning_content: Option<String>,
1580 #[serde(skip_serializing_if = "Option::is_none")]
1581 tool_calls: Option<Vec<ToolCallDelta>>,
1582}
1583
1584#[derive(Serialize)]
1585struct ChatCompletionChunkChoice {
1586 index: usize,
1587 delta: ChatCompletionChunkDelta,
1588 finish_reason: Option<&'static str>,
1589}
1590
1591#[derive(Serialize)]
1592struct ChatCompletionChunk {
1593 id: String,
1594 #[serde(skip_serializing_if = "Option::is_none")]
1601 request_id: Option<String>,
1602 object: &'static str,
1603 model: String,
1604 choices: Vec<ChatCompletionChunkChoice>,
1605 #[serde(skip_serializing_if = "Option::is_none")]
1608 usage: Option<generate::Usage>,
1609}
1610
1611async fn health(State(state): State<Arc<AppState>>) -> Response {
1617 let snapshot = state.detection.snapshot();
1618 let mut capabilities = snapshot.capabilities;
1619 let active = state.active();
1620
1621 capabilities.push(match active.as_deref() {
1624 None => frink_api::Capability::unavailable(
1629 frink_api::health::capability::REAL_WEIGHTS,
1630 frink_api::health::reason::MODEL_NOT_LOADED,
1631 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1632 ),
1633 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1634 frink_api::health::capability::REAL_WEIGHTS,
1635 frink_api::health::reason::MODEL_NOT_LOADED,
1636 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1637 checkpoint. Output from this model is noise.",
1638 ),
1639 Some(active) if active.encoder().is_some() => {
1649 let endpoints = active
1650 .encoder()
1651 .map(|e| encoder_endpoints(e))
1652 .unwrap_or_default();
1653 let served_by = match endpoints.is_empty() {
1654 true => "no endpoint in this build serves it".to_string(),
1655 false => format!("served by {}", endpoints.join(" and ")),
1656 };
1657 frink_api::Capability::available(
1658 frink_api::health::capability::REAL_WEIGHTS,
1659 format!(
1660 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1661 {served_by}; generation endpoints refuse it.",
1662 active.name(),
1663 ),
1664 )
1665 }
1666 Some(active) => frink_api::Capability::available(
1667 frink_api::health::capability::REAL_WEIGHTS,
1668 format!("Serving the real checkpoint '{}'.", active.name()),
1669 ),
1670 });
1671 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1672 frink_api::Capability::available(
1673 frink_api::health::capability::CONTINUOUS_BATCHING,
1674 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1675 "On by default on Metal. Concurrent requests share one batched decode worker."
1676 } else {
1677 "Concurrent requests share one batched decode step."
1678 },
1679 )
1680 } else if state.metal_private_decode_gate.is_some() {
1681 frink_api::Capability::unavailable(
1682 frink_api::health::capability::CONTINUOUS_BATCHING,
1683 frink_api::health::reason::DISABLED,
1684 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1685 )
1686 } else {
1687 frink_api::Capability::unavailable(
1688 frink_api::health::capability::CONTINUOUS_BATCHING,
1689 frink_api::health::reason::DISABLED,
1690 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1691 )
1692 });
1693
1694 let last_request_ms = state
1695 .last_request_ms
1696 .load(std::sync::atomic::Ordering::Relaxed);
1697 let uptime = state.started_at.elapsed();
1698 let health_state = if active.is_none() {
1706 frink_api::HealthState::Unavailable
1707 } else {
1708 snapshot.state
1709 };
1710 let body = frink_api::HealthResponse {
1711 state: health_state,
1712 reason: match health_state {
1713 frink_api::HealthState::Ready => None,
1714 frink_api::HealthState::Unavailable => {
1715 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1716 }
1717 frink_api::HealthState::Detecting => {
1718 Some(frink_api::health::reason::DETECTING.to_string())
1719 }
1720 },
1721 detail: match health_state {
1722 frink_api::HealthState::Ready => None,
1723 frink_api::HealthState::Unavailable => Some(
1724 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1725 .to_string(),
1726 ),
1727 frink_api::HealthState::Detecting => {
1728 Some("Probing available compute backends.".to_string())
1729 }
1730 },
1731 model: active
1732 .as_deref()
1733 .map(|active| frink_api::health::ModelSummary {
1734 id: active.name().to_string(),
1735 tokenizer: active.tokenizer_kind().to_string(),
1736 synthetic_weights: active.is_synthetic(),
1737 }),
1738 capabilities,
1739 version: env!("CARGO_PKG_VERSION").to_string(),
1740 pid: std::process::id(),
1741 uptime_seconds: uptime.as_secs_f64(),
1742 server_time_unix_ms: std::time::SystemTime::now()
1743 .duration_since(std::time::UNIX_EPOCH)
1744 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1745 .unwrap_or(0),
1746 last_request_age_seconds: (last_request_ms > 0)
1747 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1748 .map(|age| age.max(0.0)),
1749 };
1750
1751 let status =
1752 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1753 (status, Json(body)).into_response()
1754}
1755
1756async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1757 let Some(active) = state.active() else {
1761 return Json(serde_json::json!({ "object": "list", "data": [] }));
1762 };
1763 let mut model_entry = serde_json::json!({
1764 "id": active.name(),
1765 "object": "model",
1766 "frink_synthetic_weights": active.is_synthetic(),
1767 "frink_tokenizer": active.tokenizer_kind(),
1768 });
1769 if let Some(encoder) = active.encoder() {
1775 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1776 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1777 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1778 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1779 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1780 }
1781 if let Some(model) = active.generative_opt() {
1789 let parser_configured = active.reasoning_format().is_some();
1790 let gears = model.chat_template().think_gears(parser_configured);
1791 if !gears.is_empty() {
1792 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1793 if let Some(default) = &gears.default {
1794 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1795 }
1796 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1799 }
1800 }
1801 if let Some(mcp) = &state.mcp {
1802 model_entry["frink_mcp"] = mcp.models_metadata();
1803 }
1804 Json(serde_json::json!({
1805 "object": "list",
1806 "data": [model_entry]
1807 }))
1808}
1809
1810async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1824 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1825 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1826 let active = state.active();
1827 Json(serde_json::json!({
1828 "model": active.as_ref().map(|a| a.name()),
1829 "state": state
1830 .maintenance
1831 .lock()
1832 .unwrap_or_else(|p| p.into_inner())
1833 .state()
1834 .as_str(),
1835 "uptime_s": state.uptime().as_secs(),
1836 "throughput": {
1837 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1838 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1839 },
1840 "requests": {
1841 "active": state.cancels.live_count(),
1842 "completed": state.stats.recorded_total(),
1843 "p95_ms": state.stats.p95_duration_ms(),
1844 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1845 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1846 "completion_tokens_total": state.stats.tokens_generated_total(),
1847 },
1848 "pools": cache_admin::pool_gauges(&state),
1854 "memory": cache_admin::footprint_json(&state),
1857 }))
1858}
1859
1860#[derive(Deserialize)]
1861struct RequestsQuery {
1862 #[serde(default)]
1863 since: u64,
1864 #[serde(default = "default_requests_limit")]
1865 limit: usize,
1866}
1867
1868fn default_requests_limit() -> usize {
1869 stats::MAX_PAGE
1870}
1871
1872async fn recent_requests(
1880 State(state): State<Arc<AppState>>,
1881 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1882) -> Json<serde_json::Value> {
1883 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1884 Json(serde_json::json!({
1885 "requests": rows,
1886 "next_cursor": cursor,
1887 "missed": missed,
1888 "total": state.stats.recorded_total(),
1889 }))
1890}
1891
1892#[derive(Serialize)]
1893struct CombinedCacheStats {
1894 response_cache: response_cache::CacheStats,
1895 prefix_cache: Option<frink_models::PrefixCacheStats>,
1897}
1898
1899async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1900 Json(CombinedCacheStats {
1901 response_cache: lock_cache(&state.response_cache).stats(),
1902 prefix_cache: state
1903 .prefix_cache
1904 .as_ref()
1905 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1906 })
1907}
1908
1909async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1914 use std::sync::atomic::Ordering;
1915
1916 let cache_stats = lock_cache(&state.response_cache).stats();
1917 let active = state.active();
1918 let requests_total = state.requests_total.load(Ordering::Relaxed);
1919 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1920 let uptime = state.started_at.elapsed().as_secs_f64();
1921
1922 let body = format!(
1923 "# HELP frink_requests_total Total chat completion requests received.\n\
1924 # TYPE frink_requests_total counter\n\
1925 frink_requests_total {requests_total}\n\
1926 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1927 # TYPE frink_request_errors_total counter\n\
1928 frink_request_errors_total {errors_total}\n\
1929 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1930 # TYPE frink_cache_hits_total counter\n\
1931 frink_cache_hits_total {}\n\
1932 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1933 # TYPE frink_cache_misses_total counter\n\
1934 frink_cache_misses_total {}\n\
1935 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1936 # TYPE frink_cache_entries gauge\n\
1937 frink_cache_entries {}\n\
1938 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1939 # TYPE frink_synthetic_weights gauge\n\
1940 frink_synthetic_weights {}\n\
1941 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1942 # TYPE frink_uptime_seconds gauge\n\
1943 frink_uptime_seconds {uptime}\n",
1944 cache_stats.hits,
1945 cache_stats.misses,
1946 cache_stats.entries,
1947 active
1951 .as_ref()
1952 .map(|a| a.is_synthetic() as u8)
1953 .unwrap_or(0),
1954 );
1955
1956 let body = match active
1960 .as_ref()
1961 .and_then(|a| a.expert_store_stats())
1962 {
1963 Some(es) => format!(
1964 "{body}\
1965 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
1966 # TYPE frink_expert_cache_hits_total counter\n\
1967 frink_expert_cache_hits_total {}\n\
1968 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1969 # TYPE frink_expert_cache_misses_total counter\n\
1970 frink_expert_cache_misses_total {}\n\
1971 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
1972 # TYPE frink_expert_cache_evictions_total counter\n\
1973 frink_expert_cache_evictions_total {}\n\
1974 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1975 # TYPE frink_expert_cache_pass_throughs_total counter\n\
1976 frink_expert_cache_pass_throughs_total {}\n\
1977 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1978 # TYPE frink_expert_cache_bytes_read_total counter\n\
1979 frink_expert_cache_bytes_read_total {}\n\
1980 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1981 # TYPE frink_expert_cache_resident_bytes gauge\n\
1982 frink_expert_cache_resident_bytes {}\n",
1983 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1984 ),
1985 None => body,
1986 };
1987
1988 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
1993 Some(batcher) => {
1994 let sched = batcher.stats();
1995 format!(
1996 "{body}\
1997 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
1998 # TYPE frink_prefill_chunks_total counter\n\
1999 frink_prefill_chunks_total {}\n\
2000 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2001 # TYPE frink_prefill_tokens_total counter\n\
2002 frink_prefill_tokens_total {}\n\
2003 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2004 # TYPE frink_decode_steps_total counter\n\
2005 frink_decode_steps_total {}\n\
2006 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2007 # TYPE frink_scheduler_queue_depth gauge\n\
2008 frink_scheduler_queue_depth {}\n\
2009 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2010 # TYPE frink_scheduler_queue_rejected_total counter\n\
2011 frink_scheduler_queue_rejected_total {}\n\
2012 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2013 # TYPE frink_kv_blocks_total gauge\n\
2014 frink_kv_blocks_total {}\n\
2015 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2016 # TYPE frink_kv_blocks_free gauge\n\
2017 frink_kv_blocks_free {}\n\
2018 # HELP frink_kv_block_size Token positions per KV block.\n\
2019 # TYPE frink_kv_block_size gauge\n\
2020 frink_kv_block_size {}\n\
2021 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2022 # TYPE frink_kv_rejected_too_large_total counter\n\
2023 frink_kv_rejected_too_large_total {}\n\
2024 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2025 # TYPE frink_kv_rejected_context_length_total counter\n\
2026 frink_kv_rejected_context_length_total {}\n\
2027 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2028 # TYPE frink_scheduler_aborted_total counter\n\
2029 frink_scheduler_aborted_total {}\n\
2030 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2031 # TYPE frink_scheduler_max_seqs gauge\n\
2032 frink_scheduler_max_seqs {}\n\
2033 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2034 # TYPE frink_scheduler_prefill_chunk gauge\n\
2035 frink_scheduler_prefill_chunk {}\n",
2036 sched.prefill_chunks,
2037 sched.prefill_tokens,
2038 sched.decode_steps,
2039 sched.queue_depth,
2040 sched.queue_rejected,
2041 sched.kv_blocks_total,
2042 sched.kv_blocks_free,
2043 sched.kv_block_size,
2044 sched.kv_rejected_too_large,
2045 sched.kv_rejected_context_length,
2046 sched.aborted,
2047 sched.max_seqs,
2048 sched.prefill_chunk,
2049 )
2050 }
2051 None => body,
2052 };
2053
2054 (
2055 [(
2056 axum::http::header::CONTENT_TYPE,
2057 "text/plain; version=0.0.4",
2058 )],
2059 body,
2060 )
2061 .into_response()
2062}
2063
2064pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2065
2066pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2072 (
2073 StatusCode::BAD_REQUEST,
2074 Json(serde_json::json!({"error": {
2075 "message": message,
2076 "type": "invalid_request_error",
2077 "param": param,
2078 "code": null,
2079 }})),
2080 )
2081}
2082
2083pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2084 (
2085 StatusCode::NOT_IMPLEMENTED,
2086 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2087 )
2088}
2089
2090pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2091 let status = match e {
2092 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2093 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2098 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2104 StatusCode::SERVICE_UNAVAILABLE
2105 }
2106 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2110 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2114 };
2115 tracing::warn!("decode error: {e}");
2116 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2117 if let generate::DecodeError::KvBudgetExceeded {
2123 binding,
2124 estimated_bytes,
2125 limit_bytes,
2126 positions,
2127 positions_limit,
2128 ..
2129 } = &e
2130 {
2131 body["error"]["type"] = serde_json::json!("invalid_request_error");
2132 body["error"]["code"] = serde_json::json!(binding);
2133 body["error"]["binding"] = serde_json::json!(binding);
2134 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2135 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2136 body["error"]["positions"] = serde_json::json!(positions);
2137 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2138 }
2139 if let Some(secs) = e.retry_after_secs() {
2143 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2144 }
2145 (status, Json(body))
2146}
2147
2148pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2149 tracing::error!("generation task panicked: {e}");
2150 (
2151 StatusCode::INTERNAL_SERVER_ERROR,
2152 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2153 )
2154}
2155
2156#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2166 model: &Model,
2167 prompt: &str,
2168 params: &GenerationParams,
2169 kv_pool: Option<&generate::KvPoolConfig>,
2170 paged_kv: Option<&generate::PagedKvConfig>,
2171 prefix_cache: Option<&Mutex<PrefixCache>>,
2172 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2173 ceiling: Option<&budget::ContextCeiling>,
2174 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2175 mut emit: impl FnMut(&str),
2176) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
2177 let synthetic = model.is_synthetic();
2178 let _lora_lease = lora::lease(model, params.lora.as_deref());
2183 let mut chunks = Vec::new();
2184 let params = &{
2190 let mut resolved = params.clone();
2191 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2192 model.encode(text, SpecialTokens::Parse)
2193 });
2194 resolved.reasoning_budget = resolved
2200 .reasoning_budget
2201 .armed(resolved.reasoning, prompt, |text| {
2202 model.encode(text, SpecialTokens::Parse)
2203 })
2204 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2205 resolved
2206 };
2207 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2208 let _metal_private_guard =
2209 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2210 let (finish, usage) = match model {
2211 Model::Gguf(m) => {
2212 if let Some(batcher) = continuous_batcher {
2213 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2214 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2215 let (finish, _generated_ids, text, usage) = if synthetic {
2216 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2217 } else {
2218 batcher.generate_streaming(
2219 tokens,
2220 params.clone(),
2221 m.stop_tokens.clone(),
2222 Some(|chunk: &str| {
2223 if !chunk.is_empty() {
2224 chunks.push(chunk.to_string());
2225 emit(chunk);
2226 }
2227 }),
2228 )?
2229 };
2230 if !text.is_empty() && chunks.is_empty() {
2231 chunks.push(text);
2232 }
2233 (finish, usage)
2234 } else {
2235 generate::generate(
2236 &m.decoder,
2237 m.tokenizer.as_ref(),
2238 &m.stop_tokens,
2239 m.bos_id,
2240 prompt,
2241 params,
2242 kv_pool,
2243 paged_kv,
2244 prefix_cache,
2245 ceiling,
2246 |chunk| {
2247 chunks.push(chunk.to_string());
2248 if !synthetic {
2249 emit(chunk);
2250 }
2251 },
2252 )?
2253 }
2254 }
2255 Model::Kimi(m) => generate::generate_engine(
2256 &m.engine,
2257 &m.tokenizer,
2258 &m.stop_tokens,
2259 None,
2260 prompt,
2261 params,
2262 |chunk| {
2263 chunks.push(chunk.to_string());
2264 if !synthetic {
2265 emit(chunk);
2266 }
2267 },
2268 )?,
2269 Model::Mla(m) => generate::generate_engine(
2270 &m.engine,
2271 &m.tokenizer,
2272 &m.stop_tokens,
2273 m.bos_id,
2274 prompt,
2275 params,
2276 |chunk| {
2277 chunks.push(chunk.to_string());
2278 if !synthetic {
2279 emit(chunk);
2280 }
2281 },
2282 )?,
2283 Model::Gemma4(m) => generate::generate_engine(
2284 &m.engine,
2285 &m.tokenizer,
2286 &m.stop_tokens,
2287 m.bos_id,
2288 prompt,
2289 params,
2290 |chunk| {
2291 chunks.push(chunk.to_string());
2292 if !synthetic {
2293 emit(chunk);
2294 }
2295 },
2296 )?,
2297 Model::Glm52(m) => generate::generate_engine(
2298 &m.engine,
2299 &m.tokenizer,
2300 &m.stop_tokens,
2301 m.bos_id,
2302 prompt,
2303 params,
2304 |chunk| {
2305 chunks.push(chunk.to_string());
2306 if !synthetic {
2307 emit(chunk);
2308 }
2309 },
2310 )?,
2311 };
2312
2313 let mut full = chunks.concat();
2314 if synthetic {
2315 full = format!(
2316 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2317 to serve a real model. Decoded ids -> {full:?}]"
2318 );
2319 emit(&full);
2320 } else if used_batcher && !full.is_empty() && chunks.is_empty() {
2321 emit(&full);
2322 }
2323
2324 Ok((finish, usage, full))
2325}
2326
2327#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2332 model: &Model,
2333 prompt: &str,
2334 params: &GenerationParams,
2335 kv_pool: Option<&generate::KvPoolConfig>,
2336 paged_kv: Option<&generate::PagedKvConfig>,
2337 prefix_cache: Option<&Mutex<PrefixCache>>,
2338 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2339 ceiling: Option<&budget::ContextCeiling>,
2340 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2341) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
2342 let (finish, usage, full) = run_generation_emit(
2343 model,
2344 prompt,
2345 params,
2346 kv_pool,
2347 paged_kv,
2348 prefix_cache,
2349 continuous_batcher,
2350 ceiling,
2351 metal_private_decode_gate,
2352 |_| {},
2353 )?;
2354 Ok((
2355 if full.is_empty() {
2356 Vec::new()
2357 } else {
2358 vec![full]
2359 },
2360 finish,
2361 usage,
2362 ))
2363}
2364
2365pub(crate) fn prompt_from_messages(
2376 messages: &[ChatMessage],
2377 template: &chat_template::PromptTemplate,
2378 tools: &[ToolDef],
2379 extra: serde_json::Map<String, serde_json::Value>,
2380) -> Result<String, ApiError> {
2381 let rendered = if tools.is_empty() || template.handles_tools() {
2382 template.render(messages, tools, extra)
2383 } else {
2384 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2385 with_preamble.push(ChatMessage {
2386 role: "system".to_string(),
2387 content: Some(MessageContent::Text(tool_preamble(tools))),
2388 tool_calls: None,
2389 tool_call_id: None,
2390 reasoning_content: None,
2391 });
2392 with_preamble.extend_from_slice(messages);
2393 template.render(&with_preamble, &[], extra)
2394 };
2395 rendered.map_err(template_error_response)
2396}
2397
2398fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2403 (
2404 StatusCode::BAD_REQUEST,
2405 Json(serde_json::json!({
2406 "error": {
2407 "message": format!("chat template failed to render: {err}"),
2408 "type": "invalid_request_error",
2409 "param": "messages",
2410 "code": null,
2411 }
2412 })),
2413 )
2414}
2415
2416fn tool_preamble(tools: &[ToolDef]) -> String {
2429 let mut out = String::from(
2430 "You can call tools to help answer the user. To call a tool, respond with \
2431 EXACTLY one line in this format and nothing else:\n\
2432 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2433 object matching that tool's parameters>}}</tool_call>\n\n\
2434 Available tools:\n",
2435 );
2436 for t in tools {
2437 out.push_str(&format!(
2438 "- {}: {}\n parameters (JSON schema): {}\n",
2439 t.function.name,
2440 t.function.description.as_deref().unwrap_or(""),
2441 t.function
2442 .parameters
2443 .as_ref()
2444 .map(|v| v.to_string())
2445 .unwrap_or_else(|| "{}".to_string()),
2446 ));
2447 }
2448 out
2449}
2450
2451fn tool_call_deltas(
2461 events: Vec<crate::policy::parser::ToolCallEvent>,
2462 opened: &std::cell::Cell<usize>,
2463) -> (String, Vec<ToolCallDelta>) {
2464 let mut text = String::new();
2465 let mut deltas = Vec::new();
2466 for event in events {
2467 match event {
2468 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2469 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2470 opened.set(opened.get().max(index + 1));
2471 deltas.push(ToolCallDelta::opening(index, name));
2472 }
2473 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2474 if !fragment.is_empty() {
2475 deltas.push(ToolCallDelta::arguments(index, fragment));
2476 }
2477 }
2478 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2479 }
2480 }
2481 (text, deltas)
2482}
2483
2484fn build_response_message(
2495 text: String,
2496 tools: &[ToolDef],
2497 posture: output::OutputPosture,
2498 base_finish: &'static str,
2499) -> (ChatCompletionResponseMessage, &'static str) {
2500 let parsed = output::parse_output(&text, tools, posture);
2501 let calls: Vec<ToolCallOut> = parsed
2502 .calls
2503 .into_iter()
2504 .enumerate()
2505 .map(|(index, call)| ToolCallOut {
2506 id: format!("call_{index}"),
2507 kind: "function",
2508 function: ToolCallFunctionOut {
2509 name: call.name,
2510 arguments: call.arguments,
2511 },
2512 })
2513 .collect();
2514 if !calls.is_empty() {
2515 return (
2516 ChatCompletionResponseMessage {
2517 role: "assistant",
2518 content: None,
2519 reasoning_content: parsed.reasoning,
2520 tool_calls: Some(calls),
2521 },
2522 "tool_calls",
2523 );
2524 }
2525 (
2526 ChatCompletionResponseMessage {
2527 role: "assistant",
2528 content: Some(parsed.content),
2529 reasoning_content: parsed.reasoning,
2530 tool_calls: None,
2531 },
2532 base_finish,
2533 )
2534}
2535
2536fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2541 let mut history = match &req.session_id {
2542 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2543 None => req.messages.clone(),
2544 };
2545 if req.json_object_mode() {
2546 inject_json_object_system_hint(&mut history);
2547 }
2548 history
2549}
2550
2551fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2552 const HINT: &str =
2553 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2554 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2555 match &mut sys.content {
2556 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2557 s.push_str("\n\n");
2558 s.push_str(HINT);
2559 }
2560 None => {
2561 sys.content = Some(MessageContent::Text(HINT.to_string()));
2562 }
2563 _ => {}
2564 }
2565 } else {
2566 messages.insert(
2567 0,
2568 ChatMessage {
2569 role: "system".to_string(),
2570 content: Some(MessageContent::Text(HINT.to_string())),
2571 tool_calls: None,
2572 tool_call_id: None,
2573 reasoning_content: None,
2574 },
2575 );
2576 }
2577}
2578
2579async fn chat_completions(
2580 State(state): State<Arc<AppState>>,
2581 headers: axum::http::HeaderMap,
2582 Json(req): Json<ChatCompletionRequest>,
2583) -> Response {
2584 let attribution = attribution::Attribution::from_headers(&headers);
2585 state
2586 .requests_total
2587 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2588 let started = std::time::Instant::now();
2589
2590 let request_id = frink_api::next_request_id();
2594 let stream = req.stream.unwrap_or(false);
2595
2596 let refusal = cache_admin::check_admission(&state)
2602 .err()
2603 .or_else(|| req.validate_supported_fields().err());
2604 if let Some(err) = refusal {
2605 state
2606 .request_errors_total
2607 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2608 let response = err.into_response();
2609 state.record_request(stats::Record {
2610 request_id: &request_id,
2611 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2612 model: state.active_model_name(),
2613 status: response.status().as_u16(),
2614 stream,
2615 duration_ms: started.elapsed().as_millis() as u64,
2616 usage: None,
2617 attribution: &attribution,
2618 });
2619 return response;
2620 }
2621
2622 let response = if stream {
2623 chat_completions_stream(
2624 Arc::clone(&state),
2625 req,
2626 request_id.clone(),
2627 started,
2628 attribution.clone(),
2629 )
2630 .await
2631 .into_response()
2632 } else {
2633 chat_completions_full(
2634 Arc::clone(&state),
2635 req,
2636 request_id.clone(),
2637 started,
2638 attribution.clone(),
2639 )
2640 .await
2641 .into_response()
2642 };
2643
2644 if response.status().is_client_error() || response.status().is_server_error() {
2645 state
2646 .request_errors_total
2647 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2648 state.record_request(stats::Record {
2652 request_id: &request_id,
2653 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2654 model: state.active_model_name(),
2657 status: response.status().as_u16(),
2658 stream,
2659 duration_ms: started.elapsed().as_millis() as u64,
2660 usage: None,
2661 attribution: &attribution,
2662 });
2663 }
2664 state.mark_request_finished();
2665
2666 response
2667}
2668
2669async fn chat_completions_full(
2670 state: Arc<AppState>,
2671 req: ChatCompletionRequest,
2672 request_id: String,
2673 started: std::time::Instant,
2674 attribution: attribution::Attribution,
2675) -> Result<Json<ChatCompletionResponse>, ApiError> {
2676 let tools_active = req.tools_active();
2677 let active = state.require_active()?;
2681 let history = resolve_history(&state, &req);
2682 let template = active.generative()?.chat_template();
2683 let kwargs = req.resolve_template_kwargs(&template);
2684 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2685 let mut params =
2693 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2694 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2695 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2696
2697 let (completion, cache_status) = if let Some(cached) = key
2698 .as_ref()
2699 .and_then(|key| lock_cache(&state.response_cache).get(key))
2700 {
2701 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2702 (cached, "hit")
2703 } else {
2704 let (chunks, finish, usage) = decode_task::buffered(
2705 decode_task::DecodeHandles::take(&state, &active)?,
2706 prompt.clone(),
2707 params,
2708 )
2709 .await?;
2710
2711 let completion = response_cache::CachedCompletion {
2712 content: chunks.concat(),
2713 finish,
2714 usage,
2715 };
2716 let cache_status = match key {
2722 Some(key) => match completion.clone().cacheable() {
2726 Some(cacheable) => {
2727 tracing::debug!("cache miss for key {}", key.digest());
2728 lock_cache(&state.response_cache).put(key, cacheable);
2729 "miss"
2730 }
2731 None => "skip",
2732 },
2733 None => "skip",
2734 };
2735 (completion, cache_status)
2736 };
2737 let content = completion.content;
2738
2739 if req.json_object_mode() {
2740 json_mode::validate_json_object_output(&content)?;
2741 }
2742
2743 if let Some(id) = &req.session_id {
2747 state.sessions.store_reply(
2748 id,
2749 ChatMessage {
2750 role: "assistant".to_string(),
2751 content: Some(MessageContent::Text(content.clone())),
2752 tool_calls: None,
2753 tool_call_id: None,
2754 reasoning_content: None,
2755 },
2756 );
2757 }
2758
2759 let (message, finish_reason) = build_response_message(
2760 content,
2761 if tools_active { &req.tools } else { &[] },
2762 output::OutputPosture::resolve_full(
2763 active.reasoning_format(),
2764 active.tool_call_format(),
2765 &prompt,
2766 ),
2767 completion.finish.as_str(),
2768 );
2769
2770 state.record_request(stats::Record {
2771 request_id: &request_id,
2772 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2773 model: Some(active.name().to_string()),
2776 status: 200,
2777 stream: false,
2778 duration_ms: started.elapsed().as_millis() as u64,
2779 usage: Some(&completion.usage),
2780 attribution: &attribution,
2781 });
2782
2783 Ok(Json(ChatCompletionResponse {
2784 id: request_id.clone(),
2785 request_id,
2786 object: "chat.completion",
2787 model: req.model,
2788 choices: vec![ChatCompletionChoice {
2789 index: 0,
2790 message,
2791 finish_reason,
2792 }],
2793 usage: completion.usage,
2794 frink_cache: cache_status,
2795 }))
2796}
2797
2798async fn chat_completions_stream(
2799 state: Arc<AppState>,
2800 req: ChatCompletionRequest,
2801 request_id: String,
2802 started: std::time::Instant,
2803 attribution: attribution::Attribution,
2804) -> Result<Response, ApiError> {
2805 let tools_active = req.tools_active();
2807 let active = state.require_active()?;
2811 let history = resolve_history(&state, &req);
2812 let template = active.generative()?.chat_template();
2813 let kwargs = req.resolve_template_kwargs(&template);
2814 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2815 let model_name = req.model.clone();
2816 let session_id = req.session_id.clone();
2817 let sessions = state.sessions.clone();
2818
2819 let model = Arc::clone(active.generative()?);
2820 let kv_pool = state.kv_pool.clone();
2821 let paged_kv = state.paged_kv.clone();
2822 let prefix_cache = state.prefix_cache.clone();
2823 let batcher = active.batcher.clone();
2824 let ceiling = active.ceiling.clone();
2825 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2826 let mut params =
2827 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2828 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2829 let stats_state = Arc::clone(&state);
2830 let served_model = active.name().to_string();
2833 let posture = output::OutputPosture::resolve_full(
2837 active.reasoning_format(),
2838 active.tool_call_format(),
2839 &prompt,
2840 );
2841 let offered_tools: Vec<ToolDef> = if tools_active {
2844 req.tools.clone()
2845 } else {
2846 Vec::new()
2847 };
2848
2849 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2854 params.cancel = Some(cancel_token.clone());
2855
2856 let overlap = true;
2865
2866 let slot = req
2870 .stream_resumable
2871 .unwrap_or(false)
2872 .then(|| state.streams.register(&request_id));
2873 let emitter = resume::Emitter::new(slot);
2874
2875 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2876 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2880 id: request_id.clone(),
2881 request_id: None,
2882 object: "chat.completion.chunk",
2883 model: model_name.clone(),
2884 choices: vec![ChatCompletionChunkChoice {
2885 index: 0,
2886 delta: ChatCompletionChunkDelta {
2887 role: None,
2888 content: None,
2889 reasoning_content: None,
2890 tool_calls: None,
2891 },
2892 finish_reason: None,
2893 }],
2894 usage: None,
2895 });
2896
2897 tokio::task::spawn_blocking(move || {
2898 let _cancel_guard = cancel_guard;
2901 let tx_chunks = tx.clone();
2902 let orphan_timeout = sse::orphan_timeout_from_env();
2907 let mut first = true;
2908 let head_request_id = request_id.clone();
2909 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
2919 Rc::new(RefCell::new(posture.reasoning_parser()));
2920 let emit_reasoning = Rc::clone(&stream_reasoning);
2921 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
2925 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
2926 );
2927 let emit_tools = Rc::clone(&stream_tools);
2928 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
2932 let emit_streamed_calls = Rc::clone(&streamed_calls);
2933 let result = run_generation_emit(
2934 &model,
2935 &prompt,
2936 ¶ms,
2937 kv_pool.as_ref(),
2938 paged_kv.as_ref(),
2939 prefix_cache.as_deref(),
2940 batcher.as_ref(),
2941 ceiling.as_deref(),
2942 metal_private_decode_gate.as_deref(),
2943 |chunk| {
2944 if !overlap || chunk.is_empty() {
2945 return;
2946 }
2947 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
2948 Some(parser) => {
2949 let delta = parser.push(chunk);
2950 (delta.reasoning, delta.content)
2951 }
2952 None => (String::new(), chunk.to_string()),
2953 };
2954 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
2958 Some(parser) => {
2959 let (text, calls) =
2960 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
2961 (text, calls)
2962 }
2963 None => (content, Vec::new()),
2964 };
2965 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
2968 return;
2969 }
2970 let role = if first { Some("assistant") } else { None };
2971 let request_id = first.then(|| head_request_id.clone());
2972 first = false;
2973 let payload = ChatCompletionChunk {
2974 id: head_request_id.clone(),
2975 request_id,
2976 object: "chat.completion.chunk",
2977 model: model_name.clone(),
2978 choices: vec![ChatCompletionChunkChoice {
2979 index: 0,
2980 delta: ChatCompletionChunkDelta {
2981 role,
2982 content: (!content.is_empty()).then_some(content),
2983 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
2984 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
2985 },
2986 finish_reason: None,
2987 }],
2988 usage: None,
2989 };
2990 if let Err(why) =
2999 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3000 {
3001 if why == sse::SendFailure::Orphaned {
3002 tracing::warn!(
3003 "SSE stream {head_request_id} accepted nothing for the orphan \
3004 deadline; treating it as abandoned"
3005 );
3006 }
3007 if !emitter.is_resumable() {
3019 cancel_token.cancel();
3020 }
3021 }
3022 },
3023 );
3024
3025 let mut pending_request_id = first.then(|| request_id.clone());
3031
3032 match result {
3033 Ok((finish, usage, full_text)) => {
3034 if let Some(id) = &session_id {
3035 sessions.store_reply(
3036 id,
3037 ChatMessage {
3038 role: "assistant".to_string(),
3039 content: Some(MessageContent::Text(full_text.clone())),
3040 tool_calls: None,
3041 tool_call_id: None,
3042 reasoning_content: None,
3043 },
3044 );
3045 }
3046 let mut streamed_finish: Option<&'static str> = None;
3052 if overlap {
3053 let tail = stream_reasoning
3054 .borrow_mut()
3055 .as_mut()
3056 .map(|parser| parser.flush())
3057 .unwrap_or_default();
3058 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3059 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3060 let mut events = parser.push(&content);
3061 events.extend(parser.finish());
3062 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3063 content = text;
3064 tool_calls = calls;
3065 }
3066 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3067 let payload = ChatCompletionChunk {
3068 id: request_id.clone(),
3069 request_id: pending_request_id.take(),
3070 object: "chat.completion.chunk",
3071 model: model_name.clone(),
3072 choices: vec![ChatCompletionChunkChoice {
3073 index: 0,
3074 delta: ChatCompletionChunkDelta {
3075 role: None,
3076 content: (!content.is_empty()).then_some(content),
3077 reasoning_content: (!tail.reasoning.is_empty())
3078 .then_some(tail.reasoning),
3079 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3080 },
3081 finish_reason: None,
3082 }],
3083 usage: None,
3084 };
3085 let _ =
3086 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3087 }
3088 if streamed_calls.get() > 0 {
3089 streamed_finish = Some("tool_calls");
3090 }
3091 } else {
3092 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3095 let tool_calls: Vec<ToolCallDelta> = parsed
3096 .calls
3097 .iter()
3098 .enumerate()
3099 .map(|(index, call)| {
3100 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3101 })
3102 .collect();
3103 if !tool_calls.is_empty() {
3104 streamed_finish = Some("tool_calls");
3105 }
3106 if !tool_calls.is_empty()
3107 || !parsed.content.is_empty()
3108 || parsed.reasoning.is_some()
3109 {
3110 let payload = ChatCompletionChunk {
3111 id: request_id.clone(),
3112 request_id: pending_request_id.take(),
3113 object: "chat.completion.chunk",
3114 model: model_name.clone(),
3115 choices: vec![ChatCompletionChunkChoice {
3116 index: 0,
3117 delta: ChatCompletionChunkDelta {
3118 role: Some("assistant"),
3119 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3120 .then(|| parsed.content.clone()),
3121 reasoning_content: parsed.reasoning.clone(),
3122 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3123 },
3124 finish_reason: None,
3125 }],
3126 usage: None,
3127 };
3128 let _ =
3129 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3130 }
3131 }
3132 let final_finish_reason = match streamed_finish {
3136 Some(reason) if finish.as_str() != "length" => reason,
3137 _ => finish.as_str(),
3138 };
3139 let final_payload = ChatCompletionChunk {
3140 id: request_id.clone(),
3141 request_id: pending_request_id.take(),
3142 object: "chat.completion.chunk",
3143 model: model_name,
3144 choices: vec![ChatCompletionChunkChoice {
3145 index: 0,
3146 delta: ChatCompletionChunkDelta {
3147 role: None,
3148 content: None,
3149 reasoning_content: None,
3150 tool_calls: None,
3151 },
3152 finish_reason: Some(final_finish_reason),
3153 }],
3154 usage: Some(usage.clone()),
3155 };
3156 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3157 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3158 stats_state.record_request(stats::Record {
3163 request_id: &request_id,
3164 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3165 model: Some(served_model.clone()),
3166 status: 200,
3167 stream: true,
3168 duration_ms: started.elapsed().as_millis() as u64,
3169 usage: Some(&usage),
3170 attribution: &attribution,
3171 });
3172 }
3173 Err(e) => {
3174 tracing::warn!("decode error on streamed request {request_id}: {e}");
3175 stats_state.record_request(stats::Record {
3181 request_id: &request_id,
3182 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3183 model: Some(served_model.clone()),
3184 status: 500,
3185 stream: true,
3186 duration_ms: started.elapsed().as_millis() as u64,
3187 usage: None,
3188 attribution: &attribution,
3189 });
3190 let payload = ChatCompletionChunk {
3191 id: request_id.clone(),
3192 request_id: pending_request_id.take(),
3193 object: "chat.completion.chunk",
3194 model: model_name,
3195 choices: vec![ChatCompletionChunkChoice {
3196 index: 0,
3197 delta: ChatCompletionChunkDelta {
3198 role: Some("assistant"),
3199 content: Some(format!("[error: {e}]")),
3200 reasoning_content: None,
3201 tool_calls: None,
3202 },
3203 finish_reason: Some("stop"),
3204 }],
3205 usage: None,
3206 };
3207 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3208 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3209 }
3210 }
3211 drop(emitter);
3215 });
3216
3217 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3218 Ok((
3241 [(
3242 axum::http::HeaderName::from_static("x-accel-buffering"),
3243 axum::http::HeaderValue::from_static("no"),
3244 )],
3245 Sse::new(stream),
3246 )
3247 .into_response())
3248}
3249
3250fn protected_routes() -> Router<Arc<AppState>> {
3281 use frink_api::routes;
3282
3283 Router::new()
3284 .route(routes::V1_MODELS, get(list_models))
3285 .route(routes::V1_RESPONSES, post(responses::responses))
3289 .route(
3290 &axum_path(routes::V1_RESPONSE),
3291 get(responses::responses_get),
3292 )
3293 .route(
3294 &axum_path(routes::V1_RESPONSE_CANCEL),
3295 post(responses::responses_cancel),
3296 )
3297 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3298 .route(routes::V1_STATS, get(serving_stats))
3299 .route(routes::V1_REQUESTS, get(recent_requests))
3300 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3301 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3302 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3303 .route(
3304 routes::LORA_ADAPTERS,
3305 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3306 )
3307 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3308 .route(routes::V1_CANCEL, post(cancel_generation))
3312 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3317 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3318 .route(routes::V1_MESSAGES, post(anthropic::messages))
3319 .route(
3320 routes::V1_MESSAGES_COUNT_TOKENS,
3321 post(anthropic::count_tokens),
3322 )
3323 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3324 .route(routes::COMPLETION, post(completion::completion))
3329 .route(routes::COMPLETIONS, post(completion::completion))
3330 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3331 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3332 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3338 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3339 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3340 .route(routes::V1_RERANK, post(rerank::rerank))
3344 .route(routes::RERANK, post(rerank::rerank))
3345 .route(routes::CACHE_STATS, get(cache_stats))
3346 .route(routes::METRICS, get(metrics))
3347 .route(routes::ADMIN_MODELS, get(admin::models))
3352 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3353 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3354 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3355 .route(routes::ADMIN_TASKS, get(admin::tasks))
3356 .route(&admin::cancel_route(), post(admin::cancel_task))
3357 .route(routes::ADMIN_STATS, get(admin::stats))
3358 .merge(conversations::router())
3362}
3363
3364fn axum_path(template: &str) -> String {
3365 let mut out = String::with_capacity(template.len());
3366 let mut rest = template;
3367 while let Some(open) = rest.find('{') {
3368 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3369 break;
3370 };
3371 out.push_str(&rest[..open]);
3372 out.push(':');
3373 out.push_str(&rest[open + 1..close]);
3374 rest = &rest[close + 1..];
3375 }
3376 out.push_str(rest);
3377 out
3378}
3379
3380async fn cancel_generation(
3389 State(state): State<Arc<AppState>>,
3390 Json(req): Json<frink_api::CancelGenerationRequest>,
3391) -> Response {
3392 let cancelled = state.cancels.cancel(&req.request_id);
3393 let status = if cancelled {
3394 StatusCode::OK
3395 } else {
3396 StatusCode::NOT_FOUND
3397 };
3398 let detail = if cancelled {
3399 "the generation was asked to stop; it ends at its next token".to_string()
3400 } else {
3401 "no generation with that request_id is running -- it has already \
3402 finished, was never issued, or was served by a path that does \
3403 not register for cancellation"
3404 .to_string()
3405 };
3406 (
3407 status,
3408 Json(frink_api::CancelGenerationResponse {
3409 request_id: req.request_id,
3410 cancelled,
3411 detail,
3412 }),
3413 )
3414 .into_response()
3415}
3416
3417type Activated = (
3421 Loaded,
3422 Option<serving::batch::ContinuousBatcher>,
3423 Option<Arc<budget::ContextCeiling>>,
3424);
3425
3426fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3444 let mut batcher = serving::batch::BatcherConfig::from_env();
3445 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3446 return batcher;
3449 }
3450 let Some(path) = path else {
3451 return batcher;
3452 };
3453 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3459 let Some((priced, gguf_ctx, source)) = priced else {
3460 return batcher;
3461 };
3462 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3463 tracing::warn!(
3467 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3468 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3469 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3470 admit on a number you choose.",
3471 priced.weights_bytes,
3472 priced.device_budget_bytes,
3473 );
3474 return batcher;
3475 };
3476 tracing::info!("{source}");
3477 tracing::info!("{}", derived.fit);
3478 let adopted = budget::apply_derived(&mut batcher, &derived);
3479 if adopted.max_context {
3480 tracing::info!(
3481 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3482 override with FRINK_CB_MAX_CONTEXT",
3483 derived.max_context
3484 );
3485 }
3486 if adopted.kv_blocks {
3487 tracing::info!(
3488 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3489 derived.kv_blocks,
3490 batcher.kv_block_size
3491 );
3492 }
3493 if let Some(narrowed) = adopted.max_context_narrowed {
3494 tracing::info!(
3495 "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",
3496 batcher.kv_blocks.unwrap_or_default(),
3497 batcher.kv_block_size
3498 );
3499 }
3500 batcher
3501}
3502
3503pub(crate) fn activate_loaded_model(
3512 loaded: model::LoadedModel,
3513 enable_continuous_batching: bool,
3514 path: Option<&str>,
3515 paged_kv: Option<&generate::PagedKvConfig>,
3516) -> Activated {
3517 match loaded {
3518 model::LoadedModel::Gguf(g) => {
3519 let decoder = Arc::new(g.decoder);
3520 let tokenizer = Arc::new(g.tokenizer);
3521 let config = price_batcher_config(path);
3522 let shape =
3528 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3529 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3530 let batcher = if enable_continuous_batching {
3531 tracing::info!(
3532 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3533 (stop sequences use the same pending-buffer trim as the private generate loop)"
3534 );
3535 let tok = Arc::clone(&tokenizer);
3536 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3537 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3538 Arc::clone(&decoder),
3539 decode,
3540 config,
3541 Arc::clone(&ceiling),
3542 paged_kv.cloned(),
3543 ))
3544 } else {
3545 None
3546 };
3547 (
3548 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3549 decoder,
3550 tokenizer,
3551 stop_tokens: g.stop_tokens,
3552 bos_id: g.bos_id,
3553 is_synthetic: g.is_synthetic,
3554 chat_template: g.chat_template,
3555 }))),
3556 batcher,
3557 Some(ceiling),
3558 )
3559 }
3560 model::LoadedModel::Kimi(k) => (
3561 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3562 engine: k.engine,
3563 tokenizer: k.tokenizer,
3564 stop_tokens: k.stop_tokens,
3565 chat_template: k.chat_template,
3566 }))),
3567 None,
3568 None,
3569 ),
3570 model::LoadedModel::Mla(m) => (
3571 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3572 engine: m.engine,
3573 tokenizer: m.tokenizer,
3574 stop_tokens: m.stop_tokens,
3575 bos_id: m.bos_id,
3576 name: m.name,
3577 chat_template: m.chat_template,
3578 }))),
3579 None,
3580 None,
3581 ),
3582 model::LoadedModel::Gemma4(m) => (
3583 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3584 engine: m.engine,
3585 tokenizer: m.tokenizer,
3586 stop_tokens: m.stop_tokens,
3587 bos_id: m.bos_id,
3588 name: m.name,
3589 chat_template: m.chat_template,
3590 }))),
3591 None,
3592 None,
3593 ),
3594 model::LoadedModel::Glm52(g) => (
3595 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3596 engine: g.engine,
3597 tokenizer: g.tokenizer,
3598 stop_tokens: g.stop_tokens,
3599 bos_id: g.bos_id,
3600 name: g.name,
3601 chat_template: g.chat_template,
3602 }))),
3603 None,
3604 None,
3605 ),
3606 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3611 }
3612}
3613
3614struct StartupModels {
3621 loaded: model::LoadedModel,
3622 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3623}
3624
3625fn continuous_batching_env() -> Option<bool> {
3626 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3627 .ok()
3628 .map(|v| v.trim().to_ascii_lowercase())
3629 .as_deref()
3630 {
3631 None => None,
3632 Some("1" | "true" | "yes" | "on") => Some(true),
3633 Some("0" | "false" | "no" | "off") => Some(false),
3634 _ => None,
3635 }
3636}
3637
3638fn metal_private_decode_active() -> bool {
3639 #[cfg(feature = "metal")]
3640 {
3641 BUILT_WITH_METAL
3642 && frink_metal::attn::metal_attn_enabled()
3643 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3644 }
3645 #[cfg(not(feature = "metal"))]
3646 {
3647 false
3648 }
3649}
3650
3651fn continuous_batching_compatible(
3652 loaded: &model::LoadedModel,
3653 kv_pool: &Option<generate::KvPoolConfig>,
3654 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3655 paged_kv: &Option<generate::PagedKvConfig>,
3656) -> bool {
3657 matches!(loaded, model::LoadedModel::Gguf(_))
3658 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3659}
3660
3661fn resolve_continuous_batching_enabled(
3662 loaded: &model::LoadedModel,
3663 kv_pool: &Option<generate::KvPoolConfig>,
3664 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3665 paged_kv: &Option<generate::PagedKvConfig>,
3666) -> bool {
3667 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3668 return false;
3669 }
3670 match continuous_batching_env() {
3671 Some(true) => true,
3672 Some(false) => false,
3673 None => metal_private_decode_active(),
3674 }
3675}
3676
3677fn acquire_metal_private_decode_gate(
3678 gate: Option<&std::sync::Mutex<()>>,
3679 used_batcher: bool,
3680) -> Option<std::sync::MutexGuard<'_, ()>> {
3681 if used_batcher {
3682 None
3683 } else {
3684 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3685 }
3686}
3687
3688fn build_app_state(
3689 models: StartupModels,
3690 kv_pool: Option<generate::KvPoolConfig>,
3691 paged_kv: Option<generate::PagedKvConfig>,
3692 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3693 enable_continuous_batching: bool,
3694 mcp: Option<mcp::LoadedMcpConfig>,
3695 detection: Arc<health::Detection>,
3696) -> AppState {
3697 let StartupModels { loaded, embedding } = models;
3698 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3699 let (loaded, batcher, ceiling) = activate_loaded_model(
3700 loaded,
3701 enable_continuous_batching,
3702 configured_path.as_deref(),
3703 paged_kv.as_ref(),
3704 );
3705 let id = startup_model_id();
3711 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3712 {
3713 None
3714 } else {
3715 tracing::info!(
3716 "Metal private-loop decode will serialize concurrent requests until \
3717 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3718 );
3719 Some(Arc::new(std::sync::Mutex::new(())))
3720 };
3721 AppState {
3722 embedding,
3723 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3724 id,
3725 loaded,
3726 batcher,
3727 ceiling,
3728 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3729 }))),
3730 paged_kv,
3731 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3732 tasks: Arc::new(tasks::TaskRegistry::new()),
3733 cancels: Arc::new(cancel::CancelRegistry::new()),
3734 stats: stats::Stats::new(),
3735 streams: resume::StreamRegistry::new(),
3736 model_dir: admin::model_dirs().into_iter().next(),
3737 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3738 kv_pool,
3739 prefix_cache,
3740 sessions: session::SessionStore::new(),
3741 requests_total: std::sync::atomic::AtomicU64::new(0),
3742 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3743 started_at: std::time::Instant::now(),
3744 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3745 detection,
3746 mcp,
3747 continuous_batching_enabled: enable_continuous_batching,
3748 metal_private_decode_gate,
3749 loading_model: Mutex::new(None),
3750 last_load_error: Mutex::new(None),
3751 serving: Mutex::new(crate::stats::ServingStats::default()),
3752 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3753 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3754 started_unix: unix_now(),
3755 }
3756}
3757
3758fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3766 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3767 return Ok(None);
3768 };
3769 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3770 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3771 tracing::info!(
3772 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3773 model.name(),
3774 model.architecture(),
3775 model.n_embd(),
3776 model.pooling_type().name(),
3777 model.n_ctx_train(),
3778 );
3779 Ok(Some(Arc::new(model)))
3780}
3781
3782fn unix_now() -> u64 {
3787 std::time::SystemTime::now()
3788 .duration_since(std::time::UNIX_EPOCH)
3789 .map(|d| d.as_secs())
3790 .unwrap_or(0)
3791}
3792
3793fn startup_model_id() -> Option<String> {
3798 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3799 let configured = std::fs::canonicalize(&configured).ok()?;
3800 admin::discover(&admin::model_dirs())
3801 .into_iter()
3802 .find(|d| {
3803 std::fs::canonicalize(&d.path)
3804 .map(|p| p == configured)
3805 .unwrap_or(false)
3806 })
3807 .map(|d| d.id)
3808}
3809
3810fn init_cpu_pool() {
3819 match frink_core::threads::init_cpu_pool() {
3820 Some(n) => eprintln!(
3821 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
3822 frink_core::threads::perf_core_count()
3823 ),
3824 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
3825 }
3826}
3827
3828fn announce_ready(addr: SocketAddr, scheme: &str) {
3841 use std::io::Write;
3842 let ready =
3843 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3844 let mut stdout = std::io::stdout().lock();
3845 let _ = writeln!(stdout, "{}", ready.to_line());
3846 let _ = stdout.flush();
3847}
3848
3849async fn shutdown_signal(exit_on_stdin_close: bool) {
3861 if !exit_on_stdin_close {
3862 std::future::pending::<()>().await;
3863 return;
3864 }
3865 let _ = tokio::task::spawn_blocking(|| {
3866 use std::io::Read;
3867 let mut sink = [0u8; 256];
3868 let mut stdin = std::io::stdin().lock();
3869 loop {
3870 match stdin.read(&mut sink) {
3871 Ok(0) => break,
3873 Ok(_) => continue,
3875 Err(e) => {
3876 tracing::warn!("stdin read failed ({e}); treating it as closed");
3877 break;
3878 }
3879 }
3880 }
3881 })
3882 .await;
3883 tracing::info!("stdin closed; shutting down");
3884}
3885
3886fn tokio_worker_threads() -> usize {
3891 std::env::var("FRINK_TOKIO_WORKERS")
3892 .ok()
3893 .and_then(|v| v.trim().parse::<usize>().ok())
3894 .filter(|n| *n > 0)
3895 .unwrap_or(2)
3896}
3897
3898fn install_ring_crypto_provider() {
3915 let _ = rustls::crypto::ring::default_provider().install_default();
3916}
3917
3918pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3924 if args.list_devices {
3925 frink_models::devices::print_available_devices();
3926 return Ok(());
3927 }
3928 apply_cli_overrides(&args)?;
3929
3930 let _instance = {
3934 use frink_core::instance::{register, InstancePolicy};
3935 let policy = if args.allow_multiple_instances {
3936 InstancePolicy::Multi
3937 } else {
3938 InstancePolicy::from_env_or(InstancePolicy::Single)
3939 };
3940 let model = std::env::var("FRINK_MODEL_PATH").ok();
3941 register(
3942 "server",
3943 model.as_deref(),
3944 frink_core::instance::current_backend(),
3945 policy,
3946 )
3947 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
3948 };
3949
3950 let journal = journal::Journal::from_env();
3951 eprintln!(
3952 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
3953 journal.path()
3954 );
3955 journal.append(&journal::Record::session_start(
3956 env!("CARGO_PKG_VERSION"),
3957 std::process::id(),
3958 ));
3959 journal::install_panic_hook(journal.clone());
3960
3961 let mcp_config_path = args.mcp_config.clone();
3962 let exit_on_stdin_close = args.exit_on_stdin_close
3963 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
3964 .map(|v| v == "1")
3965 .unwrap_or(false);
3966
3967 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
3971 init_cpu_pool();
3972
3973 let runtime = tokio::runtime::Builder::new_multi_thread()
3974 .worker_threads(tokio_worker_threads())
3975 .enable_all()
3976 .build()?;
3977 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
3978
3979 let reason = match &result {
3980 Ok(()) => "normal".to_string(),
3981 Err(e) => e.to_string(),
3982 };
3983 journal.append(&journal::Record::session_exit(reason));
3984
3985 runtime.shutdown_background();
3990
3991 result
3992}
3993
3994async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
3995 let _ = tracing_subscriber::fmt::try_init();
4001
4002 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4012 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4013 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4014 .map(|v| v == "1")
4015 .unwrap_or(false);
4016 if let Err(msg) =
4017 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4018 {
4019 anyhow::bail!(msg);
4020 }
4021
4022 let embedding_model = load_embedding_model()?;
4029
4030 let mut loaded = model::load()?;
4031 match &loaded {
4032 model::LoadedModel::Gguf(g) => tracing::info!(
4033 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4034 g.decoder.config.name,
4035 g.is_synthetic,
4036 g.tokenizer.kind()
4037 ),
4038 model::LoadedModel::Kimi(k) => tracing::info!(
4039 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4040 k.tokenizer.vocab_size()
4041 ),
4042 model::LoadedModel::Mla(m) => tracing::info!(
4043 "loaded MLA GGUF '{}' (tokenizer={})",
4044 m.name,
4045 m.tokenizer.kind()
4046 ),
4047 model::LoadedModel::Gemma4(m) => tracing::info!(
4048 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4049 m.name,
4050 m.tokenizer.kind()
4051 ),
4052 model::LoadedModel::Glm52(g) => tracing::info!(
4053 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4054 g.name,
4055 g.tokenizer.kind()
4056 ),
4057 model::LoadedModel::Encoder(_) => {}
4060 }
4061 let metal_default_moe_budget = {
4069 #[cfg(feature = "metal")]
4070 {
4071 frink_core::metal_dense_enabled()
4072 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4073 }
4074 #[cfg(not(feature = "metal"))]
4075 {
4076 false
4077 }
4078 };
4079 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4080 let budget: u64 = budget_str
4081 .parse()
4082 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4083 match &mut loaded {
4084 model::LoadedModel::Gguf(g) => {
4085 tracing::info!(
4086 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4087 (CUDA and/or Metal matvecs when built with the matching feature)"
4088 );
4089 g.decoder.gpu_vram_budget_bytes = Some(budget);
4090 }
4091 model::LoadedModel::Kimi(_) => {
4092 tracing::warn!(
4093 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4094 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4095 );
4096 }
4097 model::LoadedModel::Mla(_) => {
4098 tracing::warn!(
4099 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4100 FFN path only today; ignoring expert VRAM budget"
4101 );
4102 }
4103 model::LoadedModel::Gemma4(_) => {
4104 tracing::warn!(
4105 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4106 ignoring expert VRAM budget"
4107 );
4108 }
4109 model::LoadedModel::Glm52(_) => {
4110 tracing::warn!(
4111 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4112 GPU expert placement not wired yet; ignoring"
4113 );
4114 }
4115 model::LoadedModel::Encoder(_) => {
4116 tracing::warn!(
4117 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4118 it has no routed experts to place; ignoring"
4119 );
4120 }
4121 }
4122 } else if metal_default_moe_budget {
4123 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4127 if let model::LoadedModel::Gguf(g) = &mut loaded {
4128 tracing::info!(
4129 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4130 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4131 );
4132 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4133 }
4134 }
4135 #[cfg(feature = "cuda")]
4136 {
4137 if frink_core::cuda_dense_enabled() {
4138 tracing::info!(
4139 "CUDA dense matvec enabled for WeightMatrix::apply \
4140 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4141 );
4142 } else {
4143 tracing::info!(
4144 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4145 );
4146 }
4147 }
4148 #[cfg(feature = "metal")]
4149 {
4150 if frink_core::metal_dense_enabled() {
4151 tracing::info!(
4152 "Metal dense matvec enabled for WeightMatrix::apply \
4153 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4154 );
4155 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4156 Some("1") | Some("true") | Some("on") | Some("attn") => {
4157 tracing::info!(
4158 "Metal fused attention requested (FRINK_METAL_ATTN): \
4159 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4160 );
4161 }
4162 _ => {}
4163 }
4164 tracing::info!(
4165 "Metal greedy GPU argmax: temperature<=0 folds \
4166 final_norm+lm_head+argmax into the dense stack"
4167 );
4168 } else {
4169 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4170 }
4171 }
4172 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4183 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4184 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4185 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4186 panic!(
4187 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4188 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4189 );
4190 }
4191 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4192 (Ok(blocks), Ok(block_size), Err(_)) => {
4193 let total_blocks: usize = blocks
4194 .parse()
4195 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4196 let block_size: usize = block_size
4197 .parse()
4198 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4199 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4205 .ok()
4206 .map(|v| {
4207 v.parse()
4208 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4209 })
4210 .unwrap_or(0);
4211 tracing::info!(
4212 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4213 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4214 );
4215 Some(generate::KvPoolConfig {
4216 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4217 queue_wait: Duration::from_millis(queue_wait_ms),
4218 })
4219 }
4220 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4221 let block_size: usize = block_size
4222 .parse()
4223 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4224 let budget: u64 = byte_budget
4225 .parse()
4226 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4227 let cfg = match &loaded {
4228 model::LoadedModel::Gguf(g) => &g.decoder.config,
4229 model::LoadedModel::Kimi(_)
4230 | model::LoadedModel::Mla(_)
4231 | model::LoadedModel::Gemma4(_)
4232 | model::LoadedModel::Glm52(_)
4233 | model::LoadedModel::Encoder(_) => {
4234 panic!(
4235 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4236 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4237 );
4238 }
4239 };
4240 let bytes_per_block = block_size
4241 * cfg.kv_heads_all_layers()
4242 * (cfg.head_dim + cfg.v_head_dim())
4243 * std::mem::size_of::<f32>();
4244 assert!(
4245 bytes_per_block > 0,
4246 "derived KV block byte size must be positive (check model config and block size)"
4247 );
4248 let total_blocks = (budget as usize / bytes_per_block).max(1);
4249 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4250 .ok()
4251 .map(|v| {
4252 v.parse()
4253 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4254 })
4255 .unwrap_or(0);
4256 tracing::info!(
4257 "KV cache block pool enabled from byte budget: {budget} bytes / \
4258 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4259 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4260 cfg.n_layers
4261 );
4262 Some(generate::KvPoolConfig {
4263 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4264 queue_wait: Duration::from_millis(queue_wait_ms),
4265 })
4266 }
4267 (Err(_), Err(_), Err(_)) => None,
4268 (Err(_), Ok(_), Err(_)) => panic!(
4269 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4270 (or unset all three to disable KV cache pooling)"
4271 ),
4272 (Ok(_), Ok(_), Ok(_)) => {
4273 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4274 }
4275 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4276 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4277 set together (or neither, to disable KV cache pooling)"
4278 ),
4279 };
4280 let paged_kv = match (
4286 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4287 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4288 ) {
4289 (Ok(blocks), Ok(block_size)) => {
4290 assert!(
4291 kv_pool.is_none(),
4292 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4293 mutually exclusive: both bound the same KV memory, by different means. \
4294 Set one."
4295 );
4296 let blocks_per_layer: usize = blocks
4310 .parse()
4311 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4312 let block_size: usize = block_size
4313 .parse()
4314 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4315 let gguf = match &loaded {
4316 model::LoadedModel::Gguf(g) => g,
4317 _ => panic!(
4318 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4319 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4320 ),
4321 };
4322 let cfg = &gguf.decoder.config;
4323 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4324 .ok()
4325 .map(|v| {
4326 v.parse()
4327 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4328 })
4329 .unwrap_or(0);
4330 tracing::info!(
4331 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4332 layer across {} layers, shared by all concurrent requests, \
4333 {queue_wait_ms}ms admission queue wait",
4334 cfg.n_layers
4335 );
4336 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4341 block_size,
4342 ))));
4343 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4352 crate::policy::parser::ToolCallFormat::infer(
4353 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4354 )
4355 .opener(),
4356 |text| {
4357 gguf.tokenizer
4358 .encode(text, SpecialTokens::Parse)
4359 .into_iter()
4360 .map(|t| t as u32)
4361 .collect()
4362 },
4363 );
4364 if let Some(id) = anchor_token {
4365 tracing::info!(
4366 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4367 window stops short of where its next turn rejoins"
4368 );
4369 }
4370 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4371 .ok()
4372 .map(|v| {
4373 v.parse()
4374 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4375 })
4376 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4377 if let Some(window) = cfg.uniform_sliding_window() {
4378 tracing::info!(
4379 "Paged KV window slide enabled: every layer slides by {window} every \
4380 {slide_interval} decode steps, so a request holds its prompt and a \
4381 window rather than its whole context"
4382 );
4383 } else if cfg.kv_block_window().is_some() {
4384 tracing::info!(
4385 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4386 and a page group holds one block in every layer"
4387 );
4388 }
4389 Some(generate::PagedKvConfig {
4390 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4393 queue_wait: Duration::from_millis(queue_wait_ms),
4394 radix,
4395 anchor_token,
4396 slide_interval,
4397 })
4398 }
4399 (Err(_), Err(_)) => None,
4400 _ => panic!(
4401 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4402 (or neither, to disable paged KV)"
4403 ),
4404 };
4405 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4412 let max_entries: usize = v
4413 .parse()
4414 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4415 if kv_pool.is_some() {
4416 tracing::warn!(
4417 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4418 caching will never be consulted while a KV pool is configured"
4419 );
4420 }
4421 assert!(
4428 paged_kv.is_none(),
4429 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4430 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4431 produce, so the cache could never hit. Set one."
4432 );
4433 tracing::info!(
4434 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4435 all requests"
4436 );
4437 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4438 });
4439 if matches!(
4440 loaded,
4441 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4442 ) && (kv_pool.is_some() || prefix_cache.is_some())
4443 {
4444 tracing::warn!(
4445 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4446 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4447 frink_models::engine's module docs"
4448 );
4449 }
4450 let enable_cb =
4451 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4452 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4453 tracing::info!(
4454 "continuous batching enabled by default on Metal for safe parallel serving \
4455 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4456 );
4457 }
4458 if continuous_batching_env() == Some(true)
4459 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4460 && (kv_pool.is_some() || prefix_cache.is_some())
4461 {
4462 tracing::warn!(
4463 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4464 (those modes keep the private generate path)"
4465 );
4466 }
4467 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4468 if let Ok(chunk) = n.parse::<usize>() {
4469 if chunk > 0 {
4470 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4471 }
4472 }
4473 }
4474 if matches!(
4475 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4476 Some("1")
4477 ) {
4478 tracing::warn!(
4479 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4480 (minimal spill; full layer offload still planned)"
4481 );
4482 }
4483
4484 let mcp = match mcp_config_path {
4485 Some(path) => {
4486 let loaded = mcp::load_mcp_config(&path)?;
4487 tracing::info!(
4488 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4489 loaded.path,
4490 loaded.servers.len()
4491 );
4492 Some(loaded)
4493 }
4494 None => None,
4495 };
4496
4497 let detection = health::Detection::spawn();
4501
4502 let state = Arc::new(build_app_state(
4503 StartupModels {
4504 loaded,
4505 embedding: embedding_model,
4506 },
4507 kv_pool,
4508 paged_kv,
4509 prefix_cache,
4510 enable_cb,
4511 mcp,
4512 detection,
4513 ));
4514
4515 use frink_api::routes;
4519
4520 let public = Router::new().route(routes::HEALTH, get(health));
4525
4526 let mut protected = protected_routes();
4527
4528 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4533 tracing::info!("API key auth enabled");
4534 let auth = limits::AuthConfig {
4535 api_key: Arc::new(key),
4536 };
4537 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4538 auth,
4539 limits::require_api_key,
4540 ));
4541 }
4542 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4543 let rpm: u32 = rpm
4544 .parse()
4545 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4546 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4547 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4548 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4549 limiter,
4550 limits::rate_limit,
4551 ));
4552 }
4553 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4561 let origins = security::parse_cors_origins(&spec)
4562 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4563 tracing::info!(
4564 "CORS enabled: {} allow-listed origin(s) ({})",
4565 origins.len(),
4566 spec
4567 );
4568 let cors = tower_http::cors::CorsLayer::new()
4569 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4570 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4571 .allow_headers([
4572 axum::http::header::CONTENT_TYPE,
4573 axum::http::header::AUTHORIZATION,
4574 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4580 axum::http::HeaderName::from_static("last-event-id"),
4584 ]);
4585 protected = protected.route_layer(cors);
4586 }
4587
4588 let app = public
4592 .merge(protected)
4593 .layer(axum::middleware::from_fn(limits::retry_after))
4594 .with_state(state);
4595
4596 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4602 install_ring_crypto_provider();
4603 match tls_paths {
4608 Some(paths) => {
4609 let config =
4610 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4611 .await
4612 .map_err(|e| {
4613 anyhow::anyhow!(
4614 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4615 paths.cert,
4616 paths.key
4617 )
4618 })?;
4619 let socket_addr: std::net::SocketAddr = addr
4620 .parse()
4621 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4622 let listener = std::net::TcpListener::bind(socket_addr)?;
4623 listener.set_nonblocking(true)?;
4631 let bound = listener.local_addr()?;
4632 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4633 announce_ready(bound, "https");
4634
4635 let handle = axum_server::Handle::new();
4636 let shutdown_handle = handle.clone();
4637 tokio::spawn(async move {
4638 shutdown_signal(exit_on_stdin_close).await;
4639 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4640 });
4641 axum_server::from_tcp_rustls(listener, config)?
4642 .handle(handle)
4643 .serve(app.into_make_service())
4644 .await?;
4645 }
4646 None => {
4647 let listener = tokio::net::TcpListener::bind(&addr).await?;
4648 let bound = listener.local_addr()?;
4649 tracing::info!("frink-server listening on {bound}");
4650 announce_ready(bound, "http");
4651 axum::serve(listener, app)
4652 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4653 .await?;
4654 }
4655 }
4656 Ok(())
4657}
4658
4659#[cfg(test)]
4660pub(crate) mod tests {
4661 use super::*;
4662 use frink_models::config::test_dense_fixture;
4663
4664 #[test]
4665 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4666 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4667 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4668 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4669 assert_eq!(parsed.port, 51999);
4670 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4671 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4673 }
4674
4675 fn test_model() -> Model {
4676 let cfg = test_dense_fixture();
4680 Model::Gguf(GgufModel {
4681 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4682 tokenizer: Arc::new(ServerTokenizer::Byte),
4683 stop_tokens: StopTokens::default(),
4684 bos_id: None,
4685 is_synthetic: true,
4686 chat_template: chat_template::PromptTemplate::plain(),
4687 })
4688 }
4689
4690 fn greedy_params(max_tokens: usize) -> GenerationParams {
4691 GenerationParams {
4692 reasoning: None,
4693 max_tokens,
4694 sampling: SamplingParams::default(),
4695 seed: 1,
4696 stop: Vec::new(),
4697 stop_token_ids: Vec::new(),
4698 json_object: false,
4699 grammar: None,
4700 cancel: None,
4701 ignore_eos: false,
4702 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4703 lora: None,
4704 }
4705 }
4706
4707 fn test_model_full_byte_vocab() -> Model {
4711 test_model_full_byte_vocab_with_eos(None)
4712 }
4713
4714 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4721 let mut cfg = test_dense_fixture();
4722 cfg.vocab_size = 256;
4723 Model::Gguf(GgufModel {
4724 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4725 tokenizer: Arc::new(ServerTokenizer::Byte),
4726 stop_tokens: StopTokens::from_eos(eos),
4727 bos_id: None,
4728 is_synthetic: true,
4729 chat_template: chat_template::PromptTemplate::plain(),
4730 })
4731 }
4732
4733 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4737 AppState {
4738 embedding: None,
4739 paged_kv: None,
4740 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4741 id: None,
4742 loaded: Loaded::Generative(Arc::new(model)),
4743 batcher: None,
4744 ceiling: None,
4745 checkpoint_path: None,
4746 }))),
4747 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4748 tasks: Arc::new(tasks::TaskRegistry::new()),
4749 cancels: Arc::new(cancel::CancelRegistry::new()),
4750 stats: stats::Stats::new(),
4751 streams: resume::StreamRegistry::new(),
4752 model_dir: None,
4753 response_cache: Mutex::new(response_cache),
4754 kv_pool: None,
4755 prefix_cache: None,
4756 sessions: session::SessionStore::new(),
4757 requests_total: std::sync::atomic::AtomicU64::new(0),
4758 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4759 started_at: std::time::Instant::now(),
4760 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4761 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4762 mcp: None,
4763 continuous_batching_enabled: false,
4764 metal_private_decode_gate: None,
4765 loading_model: Mutex::new(None),
4766 last_load_error: Mutex::new(None),
4767 serving: Mutex::new(crate::stats::ServingStats::default()),
4768 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4769 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4770 started_unix: unix_now(),
4771 }
4772 }
4773
4774 fn test_app() -> Router {
4783 test_app_with_state(Arc::new(test_state(
4784 test_model_full_byte_vocab(),
4785 ResponseCache::new(1000, Duration::from_secs(3600)),
4786 )))
4787 }
4788
4789 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4792 Router::new()
4802 .route(frink_api::routes::HEALTH, get(health))
4803 .merge(protected_routes())
4804 .with_state(state)
4805 }
4806
4807 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4808 let mut cfg = test_dense_fixture();
4809 cfg.name = name;
4810 cfg.vocab_size = vocab_size;
4811 Model::Gguf(GgufModel {
4812 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4813 tokenizer: Arc::new(ServerTokenizer::Byte),
4814 stop_tokens: StopTokens::default(),
4815 bos_id: None,
4816 is_synthetic: true,
4817 chat_template: chat_template::PromptTemplate::plain(),
4818 })
4819 }
4820
4821 fn model_with_template(name: &'static str, source: &str) -> Model {
4825 let mut cfg = test_dense_fixture();
4826 cfg.name = name;
4827 cfg.vocab_size = 256;
4828 Model::Gguf(GgufModel {
4829 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4830 tokenizer: Arc::new(ServerTokenizer::Byte),
4831 stop_tokens: StopTokens::default(),
4832 bos_id: None,
4833 is_synthetic: true,
4834 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4835 Some(source),
4836 Some("qwen3"),
4837 false,
4838 true,
4839 None,
4840 None,
4841 ),
4842 })
4843 }
4844
4845 #[tokio::test]
4854 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4855 let strict = "{% if messages | length > 1 %}\
4858 {{ raise_exception('this template takes one turn') }}\
4859 {% endif %}{{ messages[0].content }}";
4860 let state = Arc::new(test_state(
4861 model_with_template("strict", strict),
4862 ResponseCache::new(4, Duration::from_secs(60)),
4863 ));
4864 let app = test_app_with_state(state);
4865
4866 let (status, body) = post_json_uri(
4867 &app,
4868 "/v1/chat/completions",
4869 serde_json::json!({
4870 "model": "strict",
4871 "stream": true,
4872 "messages": [
4873 {"role": "user", "content": "one"},
4874 {"role": "user", "content": "two"},
4875 ],
4876 }),
4877 )
4878 .await;
4879 assert_eq!(status, StatusCode::BAD_REQUEST);
4880 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4881 assert!(
4882 body["error"]["message"]
4883 .as_str()
4884 .unwrap()
4885 .contains("one turn"),
4886 "the template's own message must reach the caller: {body}"
4887 );
4888
4889 let (status, _) = post_json_uri(
4891 &app,
4892 "/v1/chat/completions",
4893 serde_json::json!({
4894 "model": "strict",
4895 "stream": true,
4896 "max_tokens": 1,
4897 "messages": [{"role": "user", "content": "one"}],
4898 }),
4899 )
4900 .await;
4901 assert_eq!(status, StatusCode::OK);
4902 }
4903
4904 #[tokio::test]
4906 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4907 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4908 {% if reasoning_effort %}\
4909 {% if reasoning_effort not in ['low','medium','high'] %}\
4910 {{ raise_exception('bad effort') }}\
4911 {% endif %}[{{ reasoning_effort }}]\
4912 {% endif %}{{ messages[0].content }}";
4913 let state = Arc::new(test_state(
4914 model_with_template("thinker", reasoning),
4915 ResponseCache::new(4, Duration::from_secs(60)),
4916 ));
4917 let app = test_app_with_state(state);
4918 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4919 assert_eq!(status, StatusCode::OK);
4920 let entry = &models["data"][0];
4921 assert_eq!(
4922 entry["supported_reasoning_efforts"],
4923 serde_json::json!(["off", "low", "medium", "high"])
4924 );
4925 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4926 }
4927
4928 #[tokio::test]
4933 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4934 let app = test_app();
4935 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4936 let entry = &models["data"][0];
4937 assert!(entry.get("supported_reasoning_efforts").is_none());
4938 assert!(entry.get("default_reasoning_effort").is_none());
4939 }
4940
4941 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
4942 Arc::new(ActiveModel {
4943 id: Some(name.to_string()),
4944 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
4945 batcher: None,
4946 ceiling: None,
4947 checkpoint_path: None,
4948 })
4949 .tap_into(state)
4950 }
4951
4952 trait TapInto {
4954 fn tap_into(self, state: &AppState) -> Self;
4955 }
4956 impl TapInto for Arc<ActiveModel> {
4957 fn tap_into(self, state: &AppState) -> Self {
4958 state.swap_active(Some(Arc::clone(&self)));
4959 self
4960 }
4961 }
4962
4963 #[test]
4969 fn an_in_flight_request_keeps_the_model_it_started_on() {
4970 let state = test_state(
4971 named_test_model("model-a", 256),
4972 ResponseCache::new(4, Duration::from_secs(60)),
4973 );
4974
4975 let in_flight = state.active().expect("a model is loaded");
4978 assert_eq!(in_flight.name(), "model-a");
4979
4980 active_model(&state, "model-b");
4981
4982 assert_eq!(state.active().unwrap().name(), "model-b");
4984 assert_eq!(in_flight.name(), "model-a");
4986 let (_chunks, finish, _usage) = run_generation(
4987 in_flight.generative().unwrap(),
4988 "hi",
4989 &greedy_params(3),
4990 None,
4991 None,
4992 None,
4993 None,
4994 None,
4995 None,
4996 )
4997 .expect("the old model must still decode after being swapped out");
4998 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
4999 }
5000
5001 #[test]
5006 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5007 let state = test_state(
5008 named_test_model("model-a", 256),
5009 ResponseCache::new(4, Duration::from_secs(60)),
5010 );
5011 let in_flight = state.active().expect("a model is loaded");
5012 let weights = Arc::clone(in_flight.generative().unwrap());
5013 assert!(Arc::strong_count(&weights) >= 2);
5014
5015 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5016 id: Some("model-b".to_string()),
5017 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5018 batcher: None,
5019 ceiling: None,
5020 checkpoint_path: None,
5021 })));
5022 drop(previous);
5023 assert!(Arc::strong_count(&weights) >= 2);
5025 drop(in_flight);
5026 assert_eq!(Arc::strong_count(&weights), 1);
5027 }
5028
5029 #[tokio::test]
5033 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5034 let state = Arc::new(test_state(
5035 named_test_model("model-a", 256),
5036 ResponseCache::new(4, Duration::from_secs(60)),
5037 ));
5038 let app = test_app_with_state(Arc::clone(&state));
5039
5040 let (status, body) = post_json_uri(
5041 &app,
5042 frink_api::routes::ADMIN_MODELS_UNLOAD,
5043 serde_json::json!({}),
5044 )
5045 .await;
5046 assert_eq!(status, StatusCode::OK);
5047 assert_eq!(body["ok"], true);
5048 assert!(body["active"].is_null());
5049 assert!(state.active().is_none());
5050
5051 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5052 assert_eq!(status, StatusCode::OK);
5053 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5054 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5055
5056 let (status, body) = post_json_uri(
5057 &app,
5058 "/v1/chat/completions",
5059 serde_json::json!({
5060 "model": "x",
5061 "messages": [{"role": "user", "content": "hi"}]
5062 }),
5063 )
5064 .await;
5065 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5066 assert_eq!(body["error"]["type"], "model_not_loaded");
5067 }
5068
5069 #[tokio::test]
5073 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5074 let state = Arc::new(test_state(
5075 named_test_model("model-a", 256),
5076 ResponseCache::new(4, Duration::from_secs(60)),
5077 ));
5078 let app = test_app_with_state(Arc::clone(&state));
5079 state.swap_active(None);
5080
5081 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5082 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5085 assert_eq!(body["state"], "unavailable");
5086 assert_eq!(body["reason"], "model_not_loaded");
5087 assert!(body["model"].is_null());
5088 let real_weights = body["capabilities"]
5089 .as_array()
5090 .unwrap()
5091 .iter()
5092 .find(|c| c["id"] == "real_weights")
5093 .cloned()
5094 .expect("real_weights is always reported");
5095 assert_eq!(real_weights["available"], false);
5096 assert_eq!(real_weights["reason"], "model_not_loaded");
5097 }
5098
5099 #[tokio::test]
5103 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5104 let app = test_app();
5105
5106 let (status, completion) = post_json_uri(
5107 &app,
5108 "/v1/chat/completions",
5109 serde_json::json!({
5110 "model": "x",
5111 "messages": [{"role": "user", "content": "hi"}],
5112 "max_tokens": 4
5113 }),
5114 )
5115 .await;
5116 assert_eq!(status, StatusCode::OK);
5117 let request_id = completion["request_id"].as_str().unwrap().to_string();
5118
5119 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5120 assert_eq!(status, StatusCode::OK);
5121 let recent = stats["recent"].as_array().unwrap();
5122 assert_eq!(recent.len(), 1);
5123 let row = &recent[0];
5124 assert_eq!(row["request_id"], request_id);
5125 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5126 assert_eq!(row["status"], 200);
5127 assert_eq!(row["stream"], false);
5128 assert!(row["duration_ms"].is_number());
5131 assert!(row["decode_ms"].is_number());
5132 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5133 assert_eq!(
5134 stats["tokens_prompt_total"].as_u64().unwrap(),
5135 row["prompt_tokens"].as_u64().unwrap()
5136 );
5137 }
5138
5139 #[tokio::test]
5143 async fn a_rejected_request_is_recorded_too() {
5144 let state = Arc::new(test_state(
5145 named_test_model("model-a", 256),
5146 ResponseCache::new(4, Duration::from_secs(60)),
5147 ));
5148 let app = test_app_with_state(Arc::clone(&state));
5149 state.swap_active(None);
5150
5151 let (status, _) = post_json_uri(
5152 &app,
5153 "/v1/chat/completions",
5154 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5155 )
5156 .await;
5157 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5158
5159 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5160 let recent = stats["recent"].as_array().unwrap();
5161 assert_eq!(recent.len(), 1);
5162 assert_eq!(recent[0]["status"], 503);
5163 assert_eq!(recent[0]["completion_tokens"], 0);
5164 assert!(recent[0]["decode_ms"].is_null());
5165 assert_eq!(stats["errors_total"], 1);
5166 }
5167
5168 async fn post_json_with_headers(
5172 app: &Router,
5173 uri: &str,
5174 body: serde_json::Value,
5175 headers: &[(&str, &str)],
5176 ) -> (StatusCode, serde_json::Value) {
5177 use http_body_util::BodyExt;
5178 use tower::ServiceExt;
5179
5180 let mut builder = axum::http::Request::builder()
5181 .method("POST")
5182 .uri(uri)
5183 .header("content-type", "application/json");
5184 for (name, value) in headers {
5185 builder = builder.header(*name, *value);
5186 }
5187 let response = app
5188 .clone()
5189 .oneshot(
5190 builder
5191 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5192 .unwrap(),
5193 )
5194 .await
5195 .unwrap();
5196 let status = response.status();
5197 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5198 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5199 (status, json)
5200 }
5201
5202 #[tokio::test]
5206 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5207 let app = test_app();
5208
5209 let (status, _) = post_json_uri(
5210 &app,
5211 frink_api::routes::V1_TOKENIZE,
5212 serde_json::json!({"prompt": "hello"}),
5213 )
5214 .await;
5215 assert_eq!(status, StatusCode::OK);
5216 let (status, _) = post_json_uri(
5217 &app,
5218 frink_api::routes::V1_DETOKENIZE,
5219 serde_json::json!({"tokens": [104, 105]}),
5220 )
5221 .await;
5222 assert_eq!(status, StatusCode::OK);
5223 let (status, _) = post_json_uri(
5224 &app,
5225 frink_api::routes::V1_EMBEDDINGS,
5226 serde_json::json!({"input": "hello"}),
5227 )
5228 .await;
5229 assert_eq!(status, StatusCode::OK);
5230
5231 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5232 let routes: Vec<&str> = stats["recent"]
5233 .as_array()
5234 .unwrap()
5235 .iter()
5236 .map(|row| row["route"].as_str().unwrap())
5237 .collect();
5238 for expected in [
5239 frink_api::routes::V1_TOKENIZE,
5240 frink_api::routes::V1_DETOKENIZE,
5241 frink_api::routes::V1_EMBEDDINGS,
5242 ] {
5243 assert!(
5244 routes.contains(&expected),
5245 "{expected} is missing: {routes:?}"
5246 );
5247 }
5248
5249 let row = |route: &str| {
5250 stats["recent"]
5251 .as_array()
5252 .unwrap()
5253 .iter()
5254 .find(|r| r["route"] == route)
5255 .cloned()
5256 .unwrap()
5257 };
5258 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5262 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5263 assert!(embed["decode_ms"].is_null());
5264 assert_eq!(embed["completion_tokens"], 0);
5265 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5269 assert_eq!(
5270 stats["tokens_prompt_total"].as_u64().unwrap(),
5271 embed["prompt_tokens"].as_u64().unwrap(),
5272 "only the forward pass counted"
5273 );
5274 }
5275
5276 fn streaming_test_app() -> Router {
5281 let mut cfg = test_dense_fixture();
5282 cfg.vocab_size = 256;
5283 let model = Model::Gguf(GgufModel {
5284 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5285 tokenizer: Arc::new(ServerTokenizer::Byte),
5286 stop_tokens: StopTokens::default(),
5287 bos_id: None,
5288 is_synthetic: false,
5289 chat_template: chat_template::PromptTemplate::plain(),
5290 });
5291 test_app_with_state(Arc::new(test_state(
5292 model,
5293 ResponseCache::new(1000, Duration::from_secs(3600)),
5294 )))
5295 }
5296
5297 #[tokio::test]
5310 async fn every_route_refuses_the_same_unimplemented_fields() {
5311 let app = test_app();
5312 let fields = [
5313 ("n", serde_json::json!(3)),
5314 ("best_of", serde_json::json!(2)),
5315 ("prompt_logprobs", serde_json::json!(1)),
5316 ("echo", serde_json::json!(true)),
5317 ("use_beam_search", serde_json::json!(true)),
5318 ("truncate_prompt_tokens", serde_json::json!(8)),
5319 ("prompt_embeds", serde_json::json!("AA==")),
5320 ("allowed_token_ids", serde_json::json!([1, 2])),
5321 ("bad_words", serde_json::json!(["x"])),
5322 ("skip_special_tokens", serde_json::json!(false)),
5323 ("return_tokens_as_token_ids", serde_json::json!(true)),
5324 ];
5325 for (field, value) in fields {
5326 for (uri, base) in [
5327 (
5328 frink_api::routes::V1_CHAT_COMPLETIONS,
5329 serde_json::json!({
5330 "model": "x",
5331 "messages": [{"role": "user", "content": "hi"}],
5332 "max_tokens": 2
5333 }),
5334 ),
5335 (
5336 frink_api::routes::V1_COMPLETIONS,
5337 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
5338 ),
5339 (
5340 frink_api::routes::COMPLETION,
5341 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5342 ),
5343 ] {
5344 let mut body = base;
5345 body[field] = value.clone();
5346 let (status, answer) = post_json_uri(&app, uri, body).await;
5347 assert_eq!(
5348 status,
5349 StatusCode::NOT_IMPLEMENTED,
5350 "{uri} served `{field}` instead of refusing it: {answer}"
5351 );
5352 assert!(
5353 answer["error"]["message"]
5354 .as_str()
5355 .is_some_and(|m| m.contains(field)),
5356 "{uri} refused `{field}` without naming it: {answer}"
5357 );
5358 }
5359 }
5360 }
5361
5362 #[tokio::test]
5363 async fn the_native_completion_wire_is_not_the_openai_one() {
5364 let app = test_app();
5365
5366 let (status, native) = post_json_uri(
5367 &app,
5368 frink_api::routes::COMPLETION,
5369 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5370 )
5371 .await;
5372 assert_eq!(status, StatusCode::OK, "{native}");
5373 assert!(native["content"].is_string(), "{native}");
5374 assert_eq!(native["stop"], true);
5375 assert_eq!(native["stop_type"], "limit");
5376 assert_eq!(native["stopping_word"], "");
5377 assert_eq!(native["truncated"], false);
5378 assert_eq!(native["id_slot"], -1);
5379 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5380 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5381 assert!(
5382 native.get("choices").is_none(),
5383 "the native shape has no `choices`: {native}"
5384 );
5385
5386 let (status, openai) = post_json_uri(
5387 &app,
5388 frink_api::routes::V1_COMPLETIONS,
5389 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5390 )
5391 .await;
5392 assert_eq!(status, StatusCode::OK);
5393 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5394 assert!(
5395 openai.get("content").is_none(),
5396 "the OpenAI shape has no top-level `content`: {openai}"
5397 );
5398 }
5399
5400 #[tokio::test]
5404 async fn both_native_spellings_reach_the_same_handler() {
5405 let app = test_app();
5406 for route in [
5407 frink_api::routes::COMPLETION,
5408 frink_api::routes::COMPLETIONS,
5409 ] {
5410 let (status, body) = post_json_uri(
5411 &app,
5412 route,
5413 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5414 )
5415 .await;
5416 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5417 assert_eq!(body["stop"], true, "{route}");
5418 assert!(body["content"].is_string(), "{route}");
5419 }
5420
5421 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5424 let routes: Vec<&str> = stats["recent"]
5425 .as_array()
5426 .unwrap()
5427 .iter()
5428 .map(|row| row["route"].as_str().unwrap())
5429 .collect();
5430 assert!(
5431 routes.contains(&frink_api::routes::COMPLETION),
5432 "{routes:?}"
5433 );
5434 assert!(
5435 routes.contains(&frink_api::routes::COMPLETIONS),
5436 "{routes:?}"
5437 );
5438 }
5439
5440 #[tokio::test]
5446 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5447 let app = streaming_test_app();
5448 let raw = post_sse_raw_uri(
5449 &app,
5450 frink_api::routes::COMPLETION,
5451 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5452 )
5453 .await;
5454
5455 assert!(
5456 !raw.contains("[DONE]"),
5457 "llama.cpp's native stream has no sentinel: {raw}"
5458 );
5459 let frames: Vec<serde_json::Value> = raw
5460 .lines()
5461 .filter_map(|line| line.strip_prefix("data: "))
5462 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5463 .collect();
5464 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5465
5466 let (last, partials) = frames.split_last().unwrap();
5467 assert_eq!(last["stop"], true, "the last frame closes the stream");
5468 assert!(last["timings"].is_object(), "{last}");
5469 assert!(last["stop_type"].is_string(), "{last}");
5470 for partial in partials {
5471 assert_eq!(partial["stop"], false, "{partial}");
5472 assert!(partial["content"].is_string(), "{partial}");
5473 assert!(partial.get("timings").is_none(), "{partial}");
5477 assert!(partial.get("generation_settings").is_none(), "{partial}");
5478 }
5479 let streamed: String = partials
5482 .iter()
5483 .filter_map(|p| p["content"].as_str())
5484 .collect();
5485 assert_eq!(last["content"].as_str().unwrap(), streamed);
5486 }
5487
5488 #[tokio::test]
5493 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5494 let app = test_app();
5495 for body in [
5496 serde_json::json!({"prompt": "hi"}),
5497 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5498 ] {
5499 let (status, refusal) =
5500 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5501 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5502 assert!(
5503 refusal["error"]["message"]
5504 .as_str()
5505 .unwrap()
5506 .contains("n_predict"),
5507 "{refusal}"
5508 );
5509 }
5510 let (status, _) = post_json_uri(
5513 &app,
5514 frink_api::routes::COMPLETION,
5515 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5516 )
5517 .await;
5518 assert_eq!(status, StatusCode::OK);
5519 }
5520
5521 #[tokio::test]
5530 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5531 let app = streaming_test_app();
5532 let ask = |stop: serde_json::Value| {
5533 let app = app.clone();
5534 async move {
5535 post_json_uri(
5536 &app,
5537 frink_api::routes::COMPLETION,
5538 serde_json::json!({
5539 "prompt": "hi",
5540 "n_predict": 64,
5541 "ignore_eos": true,
5542 "stop": stop,
5543 }),
5544 )
5545 .await
5546 .1
5547 }
5548 };
5549
5550 let baseline = ask(serde_json::json!([])).await;
5551 assert_eq!(baseline["stop_type"], "limit");
5552 assert_eq!(baseline["stopping_word"], "");
5553 let text = baseline["content"].as_str().unwrap().to_string();
5554 let sentinel: String = text.chars().skip(1).take(2).collect();
5560 assert_eq!(
5561 sentinel.chars().count(),
5562 2,
5563 "the fixture must produce enough output to cut: {text:?}"
5564 );
5565 let cut = text.find(&sentinel).expect("it came out of this text");
5566
5567 let stopped = ask(serde_json::json!([sentinel])).await;
5568 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5569 assert_eq!(stopped["stopping_word"], sentinel);
5570 assert_eq!(
5571 stopped["content"].as_str().unwrap(),
5572 &text[..cut],
5573 "the answer must be cut at the sentinel, not run past it"
5574 );
5575 }
5576
5577 #[tokio::test]
5583 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5584 let app = test_app();
5585
5586 let (v1_status, v1) = post_json_uri(
5587 &app,
5588 frink_api::routes::V1_TOKENIZE,
5589 serde_json::json!({"prompt": "hello"}),
5590 )
5591 .await;
5592 let (alias_status, alias) = post_json_uri(
5593 &app,
5594 frink_api::routes::TOKENIZE,
5595 serde_json::json!({"content": "hello"}),
5596 )
5597 .await;
5598 assert_eq!(v1_status, StatusCode::OK);
5599 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5600 assert_eq!(v1["tokens"], alias["tokens"]);
5601 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5602
5603 let (status, both_ways) = post_json_uri(
5606 &app,
5607 frink_api::routes::TOKENIZE,
5608 serde_json::json!({"prompt": "hello"}),
5609 )
5610 .await;
5611 assert_eq!(status, StatusCode::OK);
5612 assert_eq!(both_ways["tokens"], v1["tokens"]);
5613 }
5614
5615 #[tokio::test]
5620 async fn detokenize_answers_under_both_dialects_keys() {
5621 let app = test_app();
5622 for route in [
5623 frink_api::routes::DETOKENIZE,
5624 frink_api::routes::V1_DETOKENIZE,
5625 ] {
5626 let (status, body) =
5627 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5628 assert_eq!(status, StatusCode::OK, "{route}");
5629 assert_eq!(body["text"], "hi", "{route}");
5630 assert_eq!(body["content"], body["text"], "{route}");
5631 }
5632 }
5633
5634 #[tokio::test]
5638 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5639 let app = test_app();
5640 let (status, _) = post_json_uri(
5641 &app,
5642 frink_api::routes::TOKENIZE,
5643 serde_json::json!({"content": "hello"}),
5644 )
5645 .await;
5646 assert_eq!(status, StatusCode::OK);
5647
5648 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5649 let routes: Vec<&str> = stats["recent"]
5650 .as_array()
5651 .unwrap()
5652 .iter()
5653 .map(|row| row["route"].as_str().unwrap())
5654 .collect();
5655 assert!(
5656 routes.contains(&frink_api::routes::TOKENIZE),
5657 "the alias must be its own row: {routes:?}"
5658 );
5659 assert!(
5660 !routes.contains(&frink_api::routes::V1_TOKENIZE),
5661 "nothing called /v1/tokenize: {routes:?}"
5662 );
5663 }
5664
5665 #[tokio::test]
5670 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5671 let mut cfg = test_dense_fixture();
5672 cfg.vocab_size = 256;
5673 let model = Model::Gguf(GgufModel {
5674 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5675 tokenizer: Arc::new(ServerTokenizer::Byte),
5676 stop_tokens: StopTokens::default(),
5677 bos_id: Some(7),
5678 is_synthetic: true,
5679 chat_template: chat_template::PromptTemplate::plain(),
5680 });
5681 let app = test_app_with_state(Arc::new(test_state(
5682 model,
5683 ResponseCache::new(1000, Duration::from_secs(3600)),
5684 )));
5685
5686 let (_, plain) = post_json_uri(
5687 &app,
5688 frink_api::routes::TOKENIZE,
5689 serde_json::json!({"content": "hi"}),
5690 )
5691 .await;
5692 let (_, special) = post_json_uri(
5693 &app,
5694 frink_api::routes::TOKENIZE,
5695 serde_json::json!({"content": "hi", "add_special": true}),
5696 )
5697 .await;
5698
5699 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5700 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5701 assert_eq!(special["count"], 3);
5702 }
5703
5704 #[tokio::test]
5707 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5708 let app = test_app();
5709 let (status, _) = post_json_uri(
5710 &app,
5711 frink_api::routes::V1_EMBEDDINGS,
5712 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5713 )
5714 .await;
5715 assert_eq!(status, StatusCode::BAD_REQUEST);
5716
5717 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5718 let recent = stats["recent"].as_array().unwrap();
5719 assert_eq!(recent.len(), 1);
5720 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
5721 assert_eq!(recent[0]["status"], 400);
5722 assert_eq!(
5723 recent[0]["prompt_tokens"], 0,
5724 "a rejected call embedded nothing"
5725 );
5726 }
5727
5728 #[tokio::test]
5731 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5732 let app = test_app();
5733 let key = "sk-monitor-secret";
5734 let (status, _) = post_json_with_headers(
5735 &app,
5736 "/v1/chat/completions",
5737 serde_json::json!({
5738 "model": "x",
5739 "messages": [{"role": "user", "content": "hi"}],
5740 "max_tokens": 2
5741 }),
5742 &[
5743 ("authorization", &format!("Bearer {key}")),
5744 ("x-frink-client", "frink-studio"),
5745 ],
5746 )
5747 .await;
5748 assert_eq!(status, StatusCode::OK);
5749
5750 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5751 let row = stats["recent"].as_array().unwrap()[0].clone();
5752 let fingerprint = row["via_api_key"]
5753 .as_str()
5754 .expect("the row names the key that served it")
5755 .to_string();
5756 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5757 assert!(!fingerprint.contains(key));
5758 assert!(
5759 !serde_json::to_string(&stats).unwrap().contains(key),
5760 "the stats payload must not carry the key in any form"
5761 );
5762 assert_eq!(row["client"], "frink-studio");
5763 }
5764
5765 #[tokio::test]
5768 async fn different_keys_are_different_callers_and_no_key_is_null() {
5769 let app = test_app();
5770 let body = serde_json::json!({
5771 "model": "x",
5772 "messages": [{"role": "user", "content": "hi"}],
5773 "max_tokens": 1
5774 });
5775 for headers in [
5776 vec![("authorization", "Bearer key-one")],
5777 vec![("authorization", "Bearer key-two")],
5778 vec![],
5779 ] {
5780 let (status, _) =
5781 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5782 assert_eq!(status, StatusCode::OK);
5783 }
5784
5785 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5786 let recent = stats["recent"].as_array().unwrap();
5787 assert_eq!(recent.len(), 3);
5788 let one = recent[0]["via_api_key"].as_str().unwrap();
5789 let two = recent[1]["via_api_key"].as_str().unwrap();
5790 assert_ne!(one, two, "two keys must not collapse into one caller");
5791 assert!(
5792 recent[2]["via_api_key"].is_null(),
5793 "an unauthenticated call is null, not a fingerprint of nothing"
5794 );
5795 assert!(recent[2]["client"].is_null());
5796 }
5797
5798 #[tokio::test]
5803 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5804 let state = Arc::new(test_state(
5805 named_test_model("really-loaded", 256),
5806 ResponseCache::new(4, Duration::from_secs(60)),
5807 ));
5808 let app = test_app_with_state(Arc::clone(&state));
5809
5810 let (status, _) = post_json_uri(
5811 &app,
5812 "/v1/chat/completions",
5813 serde_json::json!({
5814 "model": "gpt-4-turbo-that-is-not-here",
5815 "messages": [{"role": "user", "content": "hi"}],
5816 "max_tokens": 2
5817 }),
5818 )
5819 .await;
5820 assert_eq!(status, StatusCode::OK);
5821
5822 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5823 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5824
5825 state.swap_active(None);
5828 let (status, _) = post_json_uri(
5829 &app,
5830 "/v1/chat/completions",
5831 serde_json::json!({
5832 "model": "gpt-4-turbo-that-is-not-here",
5833 "messages": [{"role": "user", "content": "hi"}]
5834 }),
5835 )
5836 .await;
5837 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5838 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5839 let recent = stats["recent"].as_array().unwrap();
5840 assert!(recent[recent.len() - 1]["model"].is_null());
5841 }
5842
5843 #[tokio::test]
5847 async fn a_streamed_row_names_the_model_it_decoded_against() {
5848 let state = Arc::new(test_state(
5849 named_test_model("model-before", 256),
5850 ResponseCache::new(4, Duration::from_secs(60)),
5851 ));
5852 let app = test_app_with_state(Arc::clone(&state));
5853 let _ = post_sse_raw(&app, resumable_request()).await;
5854 active_model(&state, "model-after");
5856
5857 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5858 assert_eq!(stats["recent"][0]["model"], "model-before");
5859 }
5860
5861 #[tokio::test]
5864 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5865 let app = test_app();
5866 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5867 assert_eq!(status, StatusCode::OK);
5868 assert!(
5869 stats["queue_depth"].is_null(),
5870 "without continuous batching nothing queues, so there is nothing to measure"
5871 );
5872 assert!(stats["queue_rejected_total"].is_null());
5873 assert_eq!(
5874 stats["generating_now"], 0,
5875 "work in progress is measured and really is zero here"
5876 );
5877 }
5878
5879 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5884 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
5885 }
5886
5887 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
5891 use http_body_util::BodyExt;
5892 use tower::ServiceExt;
5893
5894 let response = app
5895 .clone()
5896 .oneshot(
5897 axum::http::Request::builder()
5898 .method("POST")
5899 .uri(uri)
5900 .header("content-type", "application/json")
5901 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5902 .unwrap(),
5903 )
5904 .await
5905 .unwrap();
5906 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5907 String::from_utf8(bytes.to_vec()).unwrap()
5908 }
5909
5910 async fn get_json_with_headers(
5911 app: &Router,
5912 uri: &str,
5913 headers: &[(&str, &str)],
5914 ) -> (StatusCode, serde_json::Value) {
5915 use http_body_util::BodyExt;
5916 use tower::ServiceExt;
5917
5918 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
5919 for (name, value) in headers {
5920 builder = builder.header(*name, *value);
5921 }
5922 let response = app
5923 .clone()
5924 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
5925 .await
5926 .unwrap();
5927 let status = response.status();
5928 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5929 (
5930 status,
5931 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
5932 )
5933 }
5934
5935 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
5936 body.lines()
5937 .filter_map(|line| line.strip_prefix(field))
5938 .map(str::trim)
5939 .collect()
5940 }
5941
5942 fn resumable_request() -> serde_json::Value {
5943 serde_json::json!({
5944 "model": "m",
5945 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
5946 "max_tokens": 4,
5947 "temperature": 0,
5948 "stream": true,
5949 "stream_resumable": true,
5950 })
5951 }
5952
5953 #[tokio::test]
5958 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
5959 let app = test_app();
5960 let body = post_sse_raw(&app, resumable_request()).await;
5961
5962 let request_id = body
5963 .lines()
5964 .find_map(|l| l.strip_prefix("data: "))
5965 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
5966 .and_then(|v| v["request_id"].as_str().map(str::to_string))
5967 .expect("the first chunk names the request");
5968
5969 let ids = sse_field(&body, "id:");
5970 let datas = sse_field(&body, "data:");
5971 assert_eq!(
5972 ids.len(),
5973 datas.len(),
5974 "every event carries an id, or a reconnect cannot name where it stopped"
5975 );
5976 for (i, id) in ids.iter().enumerate() {
5977 assert_eq!(*id, format!("{request_id}:{i}"));
5978 }
5979 let retries = sse_field(&body, "retry:");
5980 assert_eq!(
5981 retries.len(),
5982 1,
5983 "the reconnect delay is stated once, not on every event"
5984 );
5985 assert_eq!(retries[0], "1500");
5986 assert!(
5987 body.contains("data: [DONE]"),
5988 "the end of stream is still stated"
5989 );
5990 }
5991
5992 #[tokio::test]
5996 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
5997 let app = test_app();
5998 let mut request = resumable_request();
5999 request["stream_resumable"] = serde_json::json!(false);
6000 let body = post_sse_raw(&app, request).await;
6001 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6002 assert!(
6003 sse_field(&body, "id:").is_empty(),
6004 "an id promises a replay this stream cannot serve"
6005 );
6006 assert!(sse_field(&body, "retry:").is_empty());
6007 }
6008
6009 #[tokio::test]
6013 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6014 let app = test_app();
6015 let body = post_sse_raw(&app, resumable_request()).await;
6016 let request_id = sse_field(&body, "id:")[0]
6017 .rsplit_once(':')
6018 .unwrap()
6019 .0
6020 .to_string();
6021 let streamed: Vec<String> = sse_field(&body, "data:")
6022 .iter()
6023 .map(|d| d.to_string())
6024 .collect();
6025
6026 let (status, polled) = get_json(
6027 &app,
6028 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6029 )
6030 .await;
6031 assert_eq!(status, StatusCode::OK);
6032 let events: Vec<String> = polled["events"]
6033 .as_array()
6034 .unwrap()
6035 .iter()
6036 .map(|e| e["data"].as_str().unwrap().to_string())
6037 .collect();
6038 assert_eq!(
6039 events, streamed,
6040 "the fallback must deliver the same answer, not a re-run of it"
6041 );
6042 assert_eq!(polled["request_id"], request_id);
6043 assert_eq!(
6044 polled["done"], false,
6045 "events were still being handed out, so the client must ask again"
6046 );
6047
6048 let next = polled["next_index"].as_u64().unwrap();
6051 let (_, drained) = get_json(
6052 &app,
6053 &format!(
6054 "{}?from={next}",
6055 frink_api::routes::v1_stream_poll(&request_id)
6056 ),
6057 )
6058 .await;
6059 assert_eq!(drained["done"], true);
6060 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6061 }
6062
6063 #[tokio::test]
6067 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6068 let app = test_app();
6069 let body = post_sse_raw(&app, resumable_request()).await;
6070 let ids = sse_field(&body, "id:");
6071 let datas: Vec<String> = sse_field(&body, "data:")
6072 .iter()
6073 .map(|d| d.to_string())
6074 .collect();
6075 assert!(
6076 ids.len() >= 3,
6077 "need a few events to resume into the middle"
6078 );
6079 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6080
6081 let (status, resumed) = get_json_with_headers(
6082 &app,
6083 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6084 &[],
6085 )
6086 .await;
6087 assert_eq!(status, StatusCode::OK);
6088 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6089
6090 let (_, tail) = get_json(
6092 &app,
6093 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6094 )
6095 .await;
6096 let tail_events: Vec<String> = tail["events"]
6097 .as_array()
6098 .unwrap()
6099 .iter()
6100 .map(|e| e["data"].as_str().unwrap().to_string())
6101 .collect();
6102 assert_eq!(tail_events, datas[2..].to_vec());
6103 }
6104
6105 #[tokio::test]
6108 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6109 use http_body_util::BodyExt;
6110 use tower::ServiceExt;
6111
6112 let app = test_app();
6113 let body = post_sse_raw(&app, resumable_request()).await;
6114 let ids = sse_field(&body, "id:");
6115 let datas: Vec<String> = sse_field(&body, "data:")
6116 .iter()
6117 .map(|d| d.to_string())
6118 .collect();
6119 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6120
6121 let response = app
6122 .clone()
6123 .oneshot(
6124 axum::http::Request::builder()
6125 .method("GET")
6126 .uri(frink_api::routes::v1_stream(&request_id))
6127 .header("last-event-id", format!("{request_id}:0"))
6128 .body(axum::body::Body::empty())
6129 .unwrap(),
6130 )
6131 .await
6132 .unwrap();
6133 assert_eq!(response.status(), StatusCode::OK);
6134 assert_eq!(
6135 response
6136 .headers()
6137 .get("x-accel-buffering")
6138 .and_then(|v| v.to_str().ok()),
6139 Some("no"),
6140 "the reconnect needs the same anti-buffering header as the stream"
6141 );
6142 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6143 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6144 assert_eq!(
6145 sse_field(&resumed, "data:")
6146 .iter()
6147 .map(|d| d.to_string())
6148 .collect::<Vec<_>>(),
6149 datas[1..].to_vec()
6150 );
6151 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6152 }
6153
6154 #[tokio::test]
6158 async fn a_last_event_id_from_another_stream_is_refused() {
6159 let app = test_app();
6160 let body = post_sse_raw(&app, resumable_request()).await;
6161 let request_id = sse_field(&body, "id:")[0]
6162 .rsplit_once(':')
6163 .unwrap()
6164 .0
6165 .to_string();
6166
6167 let (status, err) = get_json_with_headers(
6168 &app,
6169 &frink_api::routes::v1_stream(&request_id),
6170 &[("last-event-id", "chatcmpl-someone-else:3")],
6171 )
6172 .await;
6173 assert_eq!(status, StatusCode::BAD_REQUEST);
6174 assert_eq!(err["error"]["code"], "bad_last_event_id");
6175 }
6176
6177 #[tokio::test]
6181 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6182 let app = test_app();
6183 let mut request = resumable_request();
6184 request["stream_resumable"] = serde_json::json!(false);
6185 let body = post_sse_raw(&app, request).await;
6186 let request_id = body
6187 .lines()
6188 .find_map(|l| l.strip_prefix("data: "))
6189 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6190 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6191 .unwrap();
6192
6193 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6194 assert_eq!(status, StatusCode::NOT_FOUND);
6195 assert_eq!(err["error"]["code"], "stream_not_found");
6196 assert!(err["error"]["message"]
6197 .as_str()
6198 .unwrap()
6199 .contains("stream_resumable"));
6200 }
6201
6202 #[test]
6206 fn the_axum_stream_patterns_match_the_published_templates() {
6207 assert_eq!(
6208 axum_path(frink_api::routes::V1_STREAM),
6209 "/v1/stream/:request_id"
6210 );
6211 assert_eq!(
6212 axum_path(frink_api::routes::V1_STREAM_POLL),
6213 "/v1/stream/:request_id/poll"
6214 );
6215 assert_eq!(
6216 frink_api::routes::v1_stream("abc"),
6217 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6218 );
6219 }
6220
6221 #[test]
6231 fn no_published_template_reaches_the_router_with_its_braces() {
6232 for template in [
6233 frink_api::routes::V1_STREAM,
6234 frink_api::routes::V1_STREAM_POLL,
6235 frink_api::routes::V1_RESPONSE,
6236 frink_api::routes::V1_RESPONSE_CANCEL,
6237 frink_api::routes::ADMIN_TASK_CANCEL,
6238 ] {
6239 assert!(
6240 template.contains('{'),
6241 "{template} is in the template list but has no placeholder"
6242 );
6243 let mounted = axum_path(template);
6244 assert!(
6245 !mounted.contains('{') && !mounted.contains('}'),
6246 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6247 );
6248 assert!(
6249 mounted.contains(':'),
6250 "{template} lost its placeholder entirely and would match one path only"
6251 );
6252 }
6253 }
6254
6255 #[tokio::test]
6263 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6264 let app = test_app();
6265 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6266 assert_eq!(status, StatusCode::NOT_FOUND);
6267 assert!(
6268 !body.is_null(),
6269 "empty body means axum never matched the route, so the id was read as a literal segment"
6270 );
6271 }
6272
6273 #[tokio::test]
6276 async fn the_task_list_starts_empty_rather_than_absent() {
6277 let app = test_app();
6278 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6279 assert_eq!(status, StatusCode::OK);
6280 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6281 }
6282
6283 #[tokio::test]
6292 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6293 assert!(
6294 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6295 "this test asserts the unconfigured behaviour"
6296 );
6297 let app = test_app();
6298 let (status, body) = post_json_uri(
6299 &app,
6300 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6301 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6302 )
6303 .await;
6304 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6305 assert!(
6306 body["error"]["message"]
6307 .as_str()
6308 .unwrap()
6309 .contains("--slot-save-path"),
6310 "{body}"
6311 );
6312 }
6313
6314 pub(crate) async fn post_json_uri(
6315 app: &Router,
6316 uri: &str,
6317 body: serde_json::Value,
6318 ) -> (StatusCode, serde_json::Value) {
6319 use http_body_util::BodyExt;
6320 use tower::ServiceExt;
6321
6322 let response = app
6323 .clone()
6324 .oneshot(
6325 axum::http::Request::builder()
6326 .method("POST")
6327 .uri(uri)
6328 .header("content-type", "application/json")
6329 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6330 .unwrap(),
6331 )
6332 .await
6333 .unwrap();
6334 let status = response.status();
6335 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6336 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6337 (status, json)
6338 }
6339
6340 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6341 post_json_uri(app, "/v1/chat/completions", body).await.1
6342 }
6343
6344 #[tokio::test]
6353 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6354 let app = test_app();
6355 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6356 assert_eq!(status, StatusCode::OK);
6357
6358 let memory = &body["memory"];
6359 if memory.is_null() {
6360 return;
6363 }
6364 assert!(
6365 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6366 "a read that produced a zero is a broken read, not an idle \
6367 engine: {memory}"
6368 );
6369 assert!(
6370 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6371 "the quantity must travel with the number: {memory}"
6372 );
6373 }
6374
6375 #[tokio::test]
6381 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6382 let app = test_app();
6383 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6384 assert_eq!(status, StatusCode::OK);
6385 for pool in ["kv_pages", "window_slots", "state_slots"] {
6386 assert!(
6387 body["pools"][pool].is_null(),
6388 "{pool} must be null rather than a zero row: {}",
6389 body["pools"]
6390 );
6391 }
6392 }
6393
6394 #[tokio::test]
6402 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6403 use http_body_util::BodyExt;
6404 use tower::ServiceExt;
6405
6406 let app = test_app();
6407 let send = |body: serde_json::Value| {
6408 let app = app.clone();
6409 async move {
6410 app.oneshot(
6411 axum::http::Request::builder()
6412 .method("POST")
6413 .uri(frink_api::routes::V1_MESSAGES)
6414 .header("content-type", "application/json")
6415 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6416 .unwrap(),
6417 )
6418 .await
6419 .unwrap()
6420 }
6421 };
6422
6423 let ok = send(serde_json::json!({
6424 "model": "test",
6425 "max_tokens": 1,
6426 "messages": [{"role": "user", "content": "hi"}],
6427 }))
6428 .await;
6429 assert_eq!(ok.status(), StatusCode::OK);
6430 let id = ok
6431 .headers()
6432 .get("request-id")
6433 .expect("a served message names its id")
6434 .to_str()
6435 .unwrap()
6436 .to_string();
6437 assert!(!id.is_empty());
6438
6439 let bad = send(serde_json::json!({"model": "test"})).await;
6442 assert!(bad.status().is_client_error());
6443 let other = bad.headers().get("request-id").expect("errors too");
6444 assert_ne!(other.to_str().unwrap(), id);
6445 let _ = bad.into_body().collect().await.unwrap();
6446 }
6447
6448 #[tokio::test]
6456 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6457 let state = Arc::new(test_state(
6458 test_model_full_byte_vocab(),
6459 ResponseCache::new(1000, Duration::from_secs(3600)),
6460 ));
6461 let app = test_app_with_state(Arc::clone(&state));
6462 let body = serde_json::json!({
6463 "model": "test",
6464 "messages": [{"role": "user", "content": "hi"}],
6465 "max_tokens": 1,
6466 });
6467
6468 state
6469 .maintenance
6470 .lock()
6471 .unwrap()
6472 .begin_rebuild()
6473 .expect("a fresh server is serving, so the rebuild starts");
6474 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6475 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6476 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6477
6478 state.maintenance.lock().unwrap().finish_rebuild(true);
6479 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6480 assert_eq!(
6481 status,
6482 StatusCode::OK,
6483 "the gate reopens; a rebuild is not a latch"
6484 );
6485 }
6486
6487 #[tokio::test]
6492 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6493 let app = test_app();
6494 let (status, body) = post_json_uri(
6495 &app,
6496 frink_api::routes::V1_CANCEL,
6497 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6498 )
6499 .await;
6500 assert_eq!(status, StatusCode::NOT_FOUND);
6501 assert_eq!(body["cancelled"], serde_json::json!(false));
6502 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6503 assert!(
6504 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6505 "the verdict must carry a human reason: {body}"
6506 );
6507 }
6508
6509 #[tokio::test]
6513 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6514 let state = Arc::new(test_state(
6515 test_model_full_byte_vocab(),
6516 ResponseCache::new(1000, Duration::from_secs(3600)),
6517 ));
6518 let app = test_app_with_state(Arc::clone(&state));
6519 let (token, _guard) = state.cancels.register("chatcmpl-live");
6520
6521 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6522 assert_eq!(status, StatusCode::OK);
6523 assert_eq!(before["generating_now"], serde_json::json!(1));
6524
6525 let (status, body) = post_json_uri(
6526 &app,
6527 frink_api::routes::V1_CANCEL,
6528 serde_json::json!({ "request_id": "chatcmpl-live" }),
6529 )
6530 .await;
6531 assert_eq!(status, StatusCode::OK);
6532 assert_eq!(body["cancelled"], serde_json::json!(true));
6533 assert!(
6534 token.is_cancelled(),
6535 "the endpoint answered ok without setting the flag the decode loop reads"
6536 );
6537 }
6538
6539 #[tokio::test]
6540 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6541 let app = test_app();
6542 let (status, tok) =
6543 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6544 assert_eq!(status, StatusCode::OK);
6545 let tokens = tok["tokens"].as_array().unwrap();
6546 assert_eq!(tok["count"], tokens.len());
6547 assert!(!tokens.is_empty());
6548
6549 let (status, detok) = post_json_uri(
6550 &app,
6551 "/v1/detokenize",
6552 serde_json::json!({ "tokens": tokens }),
6553 )
6554 .await;
6555 assert_eq!(status, StatusCode::OK);
6556 assert_eq!(detok["text"], "Hi");
6557
6558 let (status, emb) = post_json_uri(
6559 &app,
6560 "/v1/embeddings",
6561 serde_json::json!({
6562 "input": "Hi",
6563 "embedding_type": "mean"
6564 }),
6565 )
6566 .await;
6567 assert_eq!(status, StatusCode::OK);
6568 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6569 assert!(!vec.is_empty());
6570 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6571 }
6572
6573 #[tokio::test]
6579 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6580 let app = test_app();
6581 let (status, body) = post_json_uri(
6582 &app,
6583 "/v1/embeddings",
6584 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6585 )
6586 .await;
6587 assert_eq!(status, StatusCode::BAD_REQUEST);
6588 let msg = body["error"]["message"].as_str().unwrap();
6589 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6590 }
6591
6592 #[tokio::test]
6597 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6598 async fn a_real_embedding_model_serves_v1_embeddings() {
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 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6606 let mut state = test_state(
6607 test_model_full_byte_vocab(),
6608 ResponseCache::new(1000, Duration::from_secs(3600)),
6609 );
6610 state.embedding = Some(Arc::new(encoder));
6611 let app = test_app_with_state(Arc::new(state));
6612
6613 let (status, body) = post_json_uri(
6614 &app,
6615 "/v1/embeddings",
6616 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6617 )
6618 .await;
6619 assert_eq!(status, StatusCode::OK, "{body}");
6620 assert_eq!(body["model"], "bge-small-en-v1.5");
6621 let data = body["data"].as_array().unwrap();
6622 assert_eq!(data.len(), 2);
6623 for (i, row) in data.iter().enumerate() {
6624 assert_eq!(row["index"], i);
6625 let v: Vec<f64> = row["embedding"]
6626 .as_array()
6627 .unwrap()
6628 .iter()
6629 .map(|x| x.as_f64().unwrap())
6630 .collect();
6631 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6632 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6633 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6634 }
6635 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6638
6639 let (status, mean) = post_json_uri(
6643 &app,
6644 "/v1/embeddings",
6645 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6646 )
6647 .await;
6648 assert_eq!(status, StatusCode::OK);
6649 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6650 }
6651
6652 #[tokio::test]
6664 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6665 async fn an_encoder_can_be_the_loaded_model() {
6666 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6667 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6668 if !path.exists() {
6669 eprintln!("SKIP: {} not present", path.display());
6670 return;
6671 }
6672
6673 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6677 assert!(
6678 matches!(loaded, model::LoadedModel::Encoder(_)),
6679 "an encoder-only GGUF reached a decoder loader"
6680 );
6681 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6682 assert!(
6683 matches!(loaded, Loaded::Encoder(_)),
6684 "the encoder did not stay an encoder through activation"
6685 );
6686 assert!(
6687 batcher.is_none() && ceiling.is_none(),
6688 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6689 );
6690
6691 let state = test_state(
6692 test_model_full_byte_vocab(),
6693 ResponseCache::new(1000, Duration::from_secs(3600)),
6694 );
6695 state.swap_active(Some(Arc::new(ActiveModel {
6696 id: None,
6697 loaded,
6698 batcher,
6699 ceiling,
6700 checkpoint_path: None,
6701 })));
6702 let app = test_app_with_state(Arc::new(state));
6703
6704 let (status, body) = post_json_uri(
6706 &app,
6707 "/v1/embeddings",
6708 serde_json::json!({ "input": "Hello world" }),
6709 )
6710 .await;
6711 assert_eq!(status, StatusCode::OK, "{body}");
6712 assert_eq!(body["model"], "bge-small-en-v1.5");
6713 let v = body["data"][0]["embedding"].as_array().unwrap();
6714 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6715
6716 let (status, body) = post_json_uri(
6718 &app,
6719 "/v1/chat/completions",
6720 serde_json::json!({
6721 "model": "bge-small-en-v1.5",
6722 "messages": [{"role": "user", "content": "hi"}],
6723 }),
6724 )
6725 .await;
6726 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6727 let msg = body["error"]["message"].as_str().unwrap();
6728 for fact in [
6729 "bge-small-en-v1.5",
6730 "bert",
6731 "embedding model",
6732 "/v1/embeddings",
6733 ] {
6734 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6735 }
6736
6737 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
6739 assert_eq!(status, StatusCode::OK);
6740 let entry = &models["data"][0];
6741 assert_eq!(entry["id"], "bge-small-en-v1.5");
6742 assert_eq!(entry["frink_model_kind"], "embedding");
6743 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
6744 assert_eq!(entry["frink_n_embd"], 384);
6745 assert_eq!(entry["frink_pooling"], "CLS");
6746 assert_eq!(
6747 entry["frink_endpoints"],
6748 serde_json::json!(["/v1/embeddings"])
6749 );
6750 assert!(entry.get("supported_reasoning_efforts").is_none());
6753
6754 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
6756 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6757 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6758 assert_eq!(health["model"]["synthetic_weights"], false);
6759 let weights = health["capabilities"]
6760 .as_array()
6761 .unwrap()
6762 .iter()
6763 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
6764 .expect("a real-weights capability row");
6765 let detail = weights["detail"].as_str().unwrap_or_default();
6766 assert!(detail.contains("ENCODER"), "{detail}");
6767 let (status, body) = post_json_uri(
6775 &app,
6776 frink_api::routes::V1_TOKENIZE,
6777 serde_json::json!({ "content": "hello world" }),
6778 )
6779 .await;
6780 assert_eq!(
6781 status,
6782 StatusCode::OK,
6783 "an encoder has a real tokenizer: {body}"
6784 );
6785 let tokens = body["tokens"].as_array().expect("tokens array").clone();
6786 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
6787
6788 let (status, body) = post_json_uri(
6789 &app,
6790 frink_api::routes::V1_DETOKENIZE,
6791 serde_json::json!({ "tokens": tokens }),
6792 )
6793 .await;
6794 assert_eq!(status, StatusCode::OK, "{body}");
6795 let round_tripped = body["content"].as_str().expect("content").to_string();
6796 assert!(
6797 round_tripped.contains("hello") && round_tripped.contains("world"),
6798 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
6799 );
6800
6801 let (status, _) = post_json_uri(
6804 &app,
6805 "/v1/completions",
6806 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
6807 )
6808 .await;
6809 assert_eq!(
6810 status,
6811 StatusCode::NOT_IMPLEMENTED,
6812 "tokenizing an encoder must not have opened a path to generating with one"
6813 );
6814 }
6815
6816 #[tokio::test]
6821 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6822 use http_body_util::BodyExt;
6823 use tower::ServiceExt;
6824
6825 let fixture = concat!(
6826 "../frink-models/tests/fixtures/",
6827 "frink_real_moe_test.gguf"
6828 );
6829 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6830 let decoder = Decoder::from_gguf_with_expert_cache(
6831 &fixture,
6832 frink_models::config::test_moe_fixture(),
6833 Some(1024 * 1024),
6834 )
6835 .expect("MoE fixture must load store-backed");
6836
6837 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
6841 decoder.forward_token(1, 0, &mut caches);
6842
6843 let model = Model::Gguf(GgufModel {
6844 decoder: Arc::new(decoder),
6845 tokenizer: Arc::new(ServerTokenizer::Byte),
6846 stop_tokens: StopTokens::default(),
6847 bos_id: None,
6848 is_synthetic: false,
6849 chat_template: chat_template::PromptTemplate::plain(),
6850 });
6851 let state = Arc::new(test_state(
6852 model,
6853 ResponseCache::new(16, Duration::from_secs(60)),
6854 ));
6855 let app = Router::new()
6856 .route("/metrics", axum::routing::get(metrics))
6857 .route("/v1/chat/completions", post(chat_completions))
6858 .with_state(state);
6859
6860 let fetch_metrics = |app: Router| async move {
6861 let resp = app
6862 .oneshot(
6863 axum::http::Request::builder()
6864 .method("GET")
6865 .uri("/metrics")
6866 .body(axum::body::Body::empty())
6867 .unwrap(),
6868 )
6869 .await
6870 .unwrap();
6871 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6872 String::from_utf8(bytes.to_vec()).unwrap()
6873 };
6874
6875 let after = fetch_metrics(app.clone()).await;
6876 assert!(
6877 after.contains("frink_expert_cache_misses_total"),
6878 "streaming model must expose expert-cache metrics: {after}"
6879 );
6880 let misses: u64 = after
6881 .lines()
6882 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
6883 .and_then(|l| l.split_whitespace().nth(1))
6884 .and_then(|v| v.parse().ok())
6885 .expect("misses metric line must parse");
6886 assert!(
6887 misses > 0,
6888 "decode must have read experts through the store: {after}"
6889 );
6890 }
6891
6892 fn weather_tool() -> serde_json::Value {
6893 serde_json::json!({
6894 "type": "function",
6895 "function": {
6896 "name": "get_weather",
6897 "description": "Get the current weather for a location.",
6898 "parameters": {
6899 "type": "object",
6900 "properties": {"location": {"type": "string"}},
6901 "required": ["location"]
6902 }
6903 }
6904 })
6905 }
6906
6907 fn weather_tool_def() -> ToolDef {
6908 ToolDef {
6909 kind: "function".to_string(),
6910 function: ToolFunctionDef {
6911 name: "get_weather".to_string(),
6912 description: Some("Get the current weather for a location.".to_string()),
6913 parameters: Some(serde_json::json!({
6914 "type": "object",
6915 "properties": {"location": {"type": "string"}},
6916 "required": ["location"]
6917 })),
6918 },
6919 }
6920 }
6921
6922 #[test]
6923 fn tool_preamble_mentions_every_tool_name_and_description() {
6924 let preamble = tool_preamble(&[weather_tool_def()]);
6925 assert!(preamble.contains("get_weather"));
6926 assert!(preamble.contains("Get the current weather for a location."));
6927 assert!(preamble.contains("<tool_call>"));
6928 assert!(preamble.contains("</tool_call>"));
6929 }
6930
6931 #[test]
6932 fn a_real_marker_becomes_a_structured_tool_call() {
6933 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6934 let (message, finish) = build_response_message(
6935 text.to_string(),
6936 &[weather_tool_def()],
6937 output::OutputPosture::for_model("test-model"),
6938 "stop",
6939 );
6940 assert_eq!(finish, "tool_calls");
6941 let calls = message.tool_calls.expect("must carry a tool call");
6942 assert_eq!(calls[0].function.name, "get_weather");
6943 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
6944 assert_eq!(parsed["location"], "Paris");
6945 }
6946
6947 #[test]
6948 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
6949 let (message, finish) = build_response_message(
6950 "just an answer".to_string(),
6951 &[weather_tool_def()],
6952 output::OutputPosture::for_model("test-model"),
6953 "stop",
6954 );
6955 assert_eq!(finish, "stop");
6956 assert!(message.tool_calls.is_none());
6957 assert_eq!(message.content.as_deref(), Some("just an answer"));
6958 }
6959
6960 #[test]
6963 fn a_malformed_payload_is_not_a_tool_call() {
6964 let (message, finish) = build_response_message(
6965 "<tool_call>not valid json at all</tool_call>".to_string(),
6966 &[weather_tool_def()],
6967 output::OutputPosture::for_model("test-model"),
6968 "stop",
6969 );
6970 assert_eq!(finish, "stop");
6971 assert!(message.tool_calls.is_none());
6972 }
6973
6974 #[test]
6977 fn a_tool_that_was_never_offered_is_not_returned() {
6978 let (message, finish) = build_response_message(
6979 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
6980 &[weather_tool_def()],
6981 output::OutputPosture::for_model("test-model"),
6982 "stop",
6983 );
6984 assert_eq!(finish, "stop");
6985 assert!(message.tool_calls.is_none());
6986 }
6987
6988 #[test]
6990 fn marker_text_with_no_tools_offered_stays_content() {
6991 let (message, finish) = build_response_message(
6992 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
6993 &[],
6994 output::OutputPosture::for_model("test-model"),
6995 "stop",
6996 );
6997 assert_eq!(finish, "stop");
6998 assert!(message.tool_calls.is_none());
6999 assert!(message.content.is_some());
7000 }
7001
7002 #[test]
7006 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7007 let opened = std::cell::Cell::new(0usize);
7008 let mut parser = crate::policy::parser::ToolCallParser::new(
7009 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7010 vec![
7011 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7012 "write_file",
7013 serde_json::json!({"type": "object", "properties": {
7014 "path": {"type": "string"},
7015 "contents": {"type": "string"}
7016 }}),
7017 ),
7018 ],
7019 );
7020 let wire = "<tool_call><function=write_file>\
7021 <parameter=path>\n/tmp/x\n</parameter>\
7022 <parameter=contents>\nhello world\n</parameter>\
7023 </function></tool_call>";
7024
7025 let mut deltas = Vec::new();
7026 let mut text = String::new();
7027 for piece in wire.as_bytes().chunks(7) {
7028 let chunk = String::from_utf8_lossy(piece).into_owned();
7029 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7030 text.push_str(&more_text);
7031 deltas.extend(more);
7032 }
7033 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7034 text.push_str(&more_text);
7035 deltas.extend(more);
7036
7037 assert_eq!(opened.get(), 1, "one call opened");
7038 assert!(text.is_empty(), "the markers are not content: {text:?}");
7039
7040 let first = &deltas[0];
7041 assert_eq!(first.index, 0);
7042 assert_eq!(first.id.as_deref(), Some("call_0"));
7043 assert_eq!(first.kind, Some("function"));
7044 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7045
7046 let joined: String = deltas
7049 .iter()
7050 .filter_map(|d| d.function.arguments.clone())
7051 .collect();
7052 let parsed: serde_json::Value =
7053 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7054 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7055 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7056 assert!(
7057 deltas.len() >= 3,
7058 "the arguments arrived in pieces, not whole: {}",
7059 deltas.len()
7060 );
7061 assert!(
7062 deltas[1..].iter().all(|d| d.function.name.is_none()),
7063 "only the opening delta carries identity"
7064 );
7065 }
7066
7067 #[test]
7069 fn text_around_a_streamed_call_is_still_content() {
7070 let opened = std::cell::Cell::new(0usize);
7071 let mut parser = crate::policy::parser::ToolCallParser::new(
7072 crate::policy::parser::ToolCallFormat::Qwen25,
7073 vec![crate::policy::parser::tool_call::ToolSchema::new(
7074 "get_weather",
7075 )],
7076 );
7077 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7078 \"arguments\": {}}</tool_call> done";
7079 let mut text = String::new();
7080 for piece in wire.as_bytes().chunks(5) {
7081 let chunk = String::from_utf8_lossy(piece).into_owned();
7082 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7083 text.push_str(&more);
7084 }
7085 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7086 text.push_str(&more);
7087
7088 assert_eq!(opened.get(), 1);
7089 assert!(text.starts_with("let me check. "), "{text:?}");
7090 assert!(text.ends_with(" done"), "{text:?}");
7091 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7092 }
7093
7094 #[test]
7097 fn a_reasoning_block_is_split_out_of_the_answer() {
7098 let (message, finish) = build_response_message(
7099 "<think>weighing it up</think>The answer is 4.".to_string(),
7100 &[],
7101 output::OutputPosture::for_model("Qwen3-8B"),
7102 "stop",
7103 );
7104 assert_eq!(finish, "stop");
7105 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7106 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7107 }
7108
7109 #[test]
7112 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7113 let (message, _) = build_response_message(
7114 "Use the <think> tag like this.".to_string(),
7115 &[],
7116 output::OutputPosture::for_model("llama-3.1-8b"),
7117 "stop",
7118 );
7119 assert_eq!(
7120 message.content.as_deref(),
7121 Some("Use the <think> tag like this.")
7122 );
7123 assert!(message.reasoning_content.is_none());
7124 }
7125
7126 #[tokio::test]
7134 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7135 let app = test_app();
7136 let body = serde_json::json!({
7137 "model": "m",
7138 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7139 "max_tokens": 4,
7140 "temperature": 0,
7141 });
7142 let resp = post_json(&app, body).await;
7143 let message = &resp["choices"][0]["message"];
7144 assert!(message["content"].is_string());
7145 assert!(message.get("tool_calls").is_none());
7146 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7147 assert_eq!(resp["usage"]["completion_tokens"], 4);
7148 assert_eq!(
7149 resp["usage"]["total_tokens"],
7150 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7151 );
7152 }
7153
7154 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7155 use http_body_util::BodyExt;
7156 use tower::ServiceExt;
7157
7158 let response = app
7159 .clone()
7160 .oneshot(
7161 axum::http::Request::builder()
7162 .method("GET")
7163 .uri(uri)
7164 .body(axum::body::Body::empty())
7165 .unwrap(),
7166 )
7167 .await
7168 .unwrap();
7169 let status = response.status();
7170 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7171 (status, serde_json::from_slice(&bytes).unwrap())
7172 }
7173
7174 #[tokio::test]
7175 async fn health_answers_a_capability_handshake_not_a_boolean() {
7176 let app = test_app();
7177 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7178 assert_eq!(status, StatusCode::OK);
7179
7180 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7181 assert_eq!(health.state, frink_api::HealthState::Ready);
7182 assert!(health.pid > 0);
7183 assert!(health.server_time_unix_ms > 0);
7184 assert_eq!(health.last_request_age_seconds, None);
7187
7188 for id in [
7191 frink_api::health::capability::CPU,
7192 frink_api::health::capability::METAL,
7193 frink_api::health::capability::CUDA,
7194 frink_api::health::capability::REAL_WEIGHTS,
7195 frink_api::health::capability::CONTINUOUS_BATCHING,
7196 ] {
7197 let cap = health
7198 .capability(id)
7199 .unwrap_or_else(|| panic!("{id} missing"));
7200 assert!(!cap.reason.is_empty(), "{cap:?}");
7201 assert!(!cap.detail.is_empty(), "{cap:?}");
7202 }
7203 let weights = health
7207 .capability(frink_api::health::capability::REAL_WEIGHTS)
7208 .unwrap();
7209 assert!(!weights.available);
7210 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7211 assert!(health.model.as_ref().unwrap().synthetic_weights);
7212 }
7213
7214 #[tokio::test]
7215 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7216 let app = test_app();
7217 let _ = post_json(
7218 &app,
7219 serde_json::json!({
7220 "model": "m",
7221 "messages": [{"role": "user", "content": "\u{1}"}],
7222 "max_tokens": 1,
7223 "temperature": 0,
7224 }),
7225 )
7226 .await;
7227 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7228 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7229 let age = health
7230 .last_request_age_seconds
7231 .expect("a served request is evidence of liveness");
7232 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7233 }
7234
7235 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7237 use http_body_util::BodyExt;
7238 use tower::ServiceExt;
7239
7240 let response = app
7241 .clone()
7242 .oneshot(
7243 axum::http::Request::builder()
7244 .method("POST")
7245 .uri("/v1/chat/completions")
7246 .header("content-type", "application/json")
7247 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7248 .unwrap(),
7249 )
7250 .await
7251 .unwrap();
7252 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7253 String::from_utf8(bytes.to_vec())
7254 .unwrap()
7255 .lines()
7256 .filter_map(|line| line.strip_prefix("data: "))
7257 .filter(|payload| *payload != "[DONE]")
7258 .map(|payload| serde_json::from_str(payload).unwrap())
7259 .collect()
7260 }
7261
7262 #[tokio::test]
7263 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7264 let app = test_app();
7265 let chunks = post_sse_chunks(
7266 &app,
7267 serde_json::json!({
7268 "model": "m",
7269 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7270 "max_tokens": 4,
7271 "temperature": 0,
7272 "stream": true,
7273 }),
7274 )
7275 .await;
7276
7277 assert!(!chunks.is_empty());
7278 let request_id = chunks[0]["request_id"]
7279 .as_str()
7280 .expect("the first chunk names the request")
7281 .to_string();
7282 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7283 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7286 assert!(
7287 chunk.get("request_id").is_none(),
7288 "chunk {i} repeats request_id"
7289 );
7290 }
7291 for chunk in &chunks {
7294 assert_eq!(chunk["id"], serde_json::json!(request_id));
7295 }
7296
7297 let other = post_sse_chunks(
7298 &app,
7299 serde_json::json!({
7300 "model": "m",
7301 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7302 "max_tokens": 4,
7303 "temperature": 0,
7304 "stream": true,
7305 }),
7306 )
7307 .await;
7308 assert_ne!(
7309 other[0]["request_id"].as_str().unwrap(),
7310 request_id,
7311 "two concurrent chats must not share an id"
7312 );
7313 }
7314
7315 #[tokio::test]
7316 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7317 let app = test_app();
7318 let resp = post_json(
7319 &app,
7320 serde_json::json!({
7321 "model": "m",
7322 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7323 "max_tokens": 2,
7324 "temperature": 0,
7325 }),
7326 )
7327 .await;
7328 assert_eq!(resp["id"], resp["request_id"]);
7329 assert!(resp["request_id"]
7330 .as_str()
7331 .unwrap()
7332 .starts_with("chatcmpl-"));
7333 }
7334
7335 #[tokio::test]
7338 async fn usage_carries_separate_prefill_and_decode_timings() {
7339 let app = test_app();
7340 let resp = post_json(
7341 &app,
7342 serde_json::json!({
7343 "model": "m",
7344 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7345 "max_tokens": 4,
7346 "temperature": 0,
7347 }),
7348 )
7349 .await;
7350 let usage = &resp["usage"];
7351 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7352 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7353 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7354 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7355 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7357 }
7358
7359 #[tokio::test]
7367 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7368 let app = test_app();
7369 let body = serde_json::json!({
7370 "model": "m",
7371 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7372 "max_tokens": 4,
7373 "temperature": 0,
7374 "tools": [weather_tool()],
7375 });
7376 let resp = post_json(&app, body).await;
7377 let message = &resp["choices"][0]["message"];
7378 assert!(
7379 message["content"].is_string(),
7380 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7381 );
7382 assert!(message.get("tool_calls").is_none());
7383 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7387 }
7388
7389 #[tokio::test]
7393 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7394 let app = test_app();
7395 let body = serde_json::json!({
7396 "model": "m",
7397 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7398 "max_tokens": 3,
7399 "temperature": 0,
7400 });
7401 let first = post_json(&app, body.clone()).await;
7402 assert_eq!(first["frink_cache"], "miss");
7403 let second = post_json(&app, body).await;
7404 assert_eq!(second["frink_cache"], "hit");
7405 assert_eq!(
7406 first["choices"][0]["message"]["content"],
7407 second["choices"][0]["message"]["content"]
7408 );
7409 assert_eq!(
7410 first["choices"][0]["finish_reason"],
7411 second["choices"][0]["finish_reason"]
7412 );
7413 assert_eq!(first["usage"], second["usage"]);
7414 assert_eq!(second["usage"]["completion_tokens"], 3);
7415 }
7416
7417 #[tokio::test]
7428 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7429 let app = test_app();
7430 let plain = serde_json::json!({
7431 "model": "m",
7432 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7433 "max_tokens": 3,
7434 "temperature": 0,
7435 });
7436
7437 let first = post_json(&app, plain.clone()).await;
7438 assert_eq!(first["frink_cache"], "miss");
7439 let unconstrained = first["choices"][0]["message"]["content"]
7440 .as_str()
7441 .expect("content")
7442 .to_string();
7443
7444 let mut constrained = plain.clone();
7445 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7446 let second = post_json(&app, constrained).await;
7447 assert_eq!(
7448 second["frink_cache"], "miss",
7449 "a grammar is part of the key, so this body has never been answered"
7450 );
7451 let constrained_answer = second["choices"][0]["message"]["content"]
7455 .as_str()
7456 .expect("content")
7457 .to_string();
7458 assert!(
7459 constrained_answer.contains("-> \"yes\"]"),
7460 "the grammar must have been compiled AND applied, not skipped \
7461 by a cache hit: {constrained_answer}"
7462 );
7463 assert_ne!(
7464 constrained_answer, unconstrained,
7465 "the constrained request was served the unconstrained answer"
7466 );
7467
7468 let third = post_json(&app, plain).await;
7472 assert_eq!(third["frink_cache"], "hit");
7473 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7474 }
7475
7476 #[tokio::test]
7498 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7499 let state = Arc::new(test_state(
7500 test_model_full_byte_vocab(),
7501 ResponseCache::new(1000, Duration::from_secs(3600)),
7502 ));
7503 let app = test_app_with_state(state.clone());
7504 let plain = serde_json::json!({
7505 "model": "m",
7506 "messages": [
7507 {"role": "system", "content": "Answer in JSON when it helps."},
7508 {"role": "user", "content": "\u{1}\u{2}"},
7509 ],
7510 "max_tokens": 3,
7511 "temperature": 0,
7512 });
7513
7514 let first = post_json(&app, plain.clone()).await;
7515 assert_eq!(first["frink_cache"], "miss");
7516 assert_eq!(state.cache_stats().entries, 1);
7517
7518 let mut as_json = plain.clone();
7519 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7520 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7521 assert_eq!(
7522 status,
7523 StatusCode::BAD_REQUEST,
7524 "the demo banner is not a JSON object, whoever generated it"
7525 );
7526 assert_eq!(
7527 state.cache_stats().hits,
7528 0,
7529 "a json_object request must not be answered from an entry the \
7530 JSON mask never produced"
7531 );
7532 assert_eq!(
7533 state.cache_stats().entries,
7534 2,
7535 "json_object must key its own entry, not reuse the unconstrained \
7536 one it happens to render the same prompt as"
7537 );
7538 }
7539
7540 #[tokio::test]
7552 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7553 let app = test_app_with_state(Arc::new(test_state(
7554 test_model_full_byte_vocab_with_eos(Some(0x77)),
7555 ResponseCache::new(1000, Duration::from_secs(3600)),
7556 )));
7557 let body = serde_json::json!({
7558 "model": "m",
7559 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7560 "max_tokens": 6,
7561 "temperature": 0,
7562 });
7563
7564 let stopped = post_json(&app, body.clone()).await;
7565 assert_eq!(stopped["frink_cache"], "miss");
7566 assert_eq!(
7567 stopped["choices"][0]["finish_reason"], "stop",
7568 "the fixture is only meaningful if the model's EOS really fires here"
7569 );
7570 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7571
7572 let mut ignoring = body.clone();
7573 ignoring["ignore_eos"] = serde_json::json!(true);
7574 let ran_on = post_json(&app, ignoring).await;
7575 assert_eq!(
7576 ran_on["frink_cache"], "miss",
7577 "ignore_eos is part of the key, so this body has never been answered"
7578 );
7579 assert_eq!(
7580 ran_on["usage"]["completion_tokens"], 6,
7581 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7582 );
7583 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7584 assert_ne!(
7585 ran_on["choices"][0]["message"]["content"],
7586 stopped["choices"][0]["message"]["content"]
7587 );
7588 }
7589
7590 #[tokio::test]
7596 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7597 let session_app = test_app();
7598 let manual_app = test_app();
7599
7600 let turn1 = post_json(
7602 &session_app,
7603 serde_json::json!({
7604 "model": "m",
7605 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7606 "session_id": "s1",
7607 "max_tokens": 5,
7608 "temperature": 0,
7609 }),
7610 )
7611 .await;
7612 let reply1 = turn1["choices"][0]["message"]["content"]
7613 .as_str()
7614 .unwrap()
7615 .to_string();
7616
7617 let manual_turn1 = post_json(
7622 &manual_app,
7623 serde_json::json!({
7624 "model": "m",
7625 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7626 "max_tokens": 5,
7627 "temperature": 0,
7628 }),
7629 )
7630 .await;
7631 assert_eq!(
7632 manual_turn1["choices"][0]["message"]["content"]
7633 .as_str()
7634 .unwrap(),
7635 reply1
7636 );
7637
7638 let turn2 = post_json(
7640 &session_app,
7641 serde_json::json!({
7642 "model": "m",
7643 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7644 "session_id": "s1",
7645 "max_tokens": 5,
7646 "temperature": 0,
7647 }),
7648 )
7649 .await;
7650 let reply2 = turn2["choices"][0]["message"]["content"]
7651 .as_str()
7652 .unwrap()
7653 .to_string();
7654
7655 let manual_turn2 = post_json(
7659 &manual_app,
7660 serde_json::json!({
7661 "model": "m",
7662 "messages": [
7663 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7664 {"role": "assistant", "content": reply1},
7665 {"role": "user", "content": "\u{4}\u{5}"},
7666 ],
7667 "max_tokens": 5,
7668 "temperature": 0,
7669 }),
7670 )
7671 .await;
7672 assert_eq!(
7673 manual_turn2["choices"][0]["message"]["content"]
7674 .as_str()
7675 .unwrap(),
7676 reply2,
7677 "resuming a session must produce identical output to manually resending the full history"
7678 );
7679 }
7680
7681 #[test]
7684 fn lock_cache_recovers_from_a_poisoned_mutex() {
7685 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7686
7687 let poison_cache = Arc::clone(&cache);
7688 let _ = std::thread::spawn(move || {
7689 let _guard = poison_cache.lock().unwrap();
7690 panic!("simulated panic while holding the lock");
7691 })
7692 .join();
7693
7694 let recovered = lock_cache(&cache);
7696 assert_eq!(recovered.stats().entries, 0);
7697 }
7698
7699 #[test]
7700 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7701 let mut req_body = serde_json::json!({
7702 "model": "m",
7703 "messages": [{"role": "user", "content": "hi"}],
7704 });
7705 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7706 assert!(
7707 req.is_cacheable(),
7708 "default (temperature 0) must be cacheable"
7709 );
7710
7711 req_body["temperature"] = serde_json::json!(0.8);
7712 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7713 assert!(
7714 !req.is_cacheable(),
7715 "unseeded sampling must never be cacheable"
7716 );
7717
7718 req_body["seed"] = serde_json::json!(42);
7719 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7720 assert!(
7721 req.is_cacheable(),
7722 "sampling with an explicit seed is deterministic and must be cacheable"
7723 );
7724 }
7725
7726 const GRADED: &str = "{% if reasoning_effort %}\
7730 {% if reasoning_effort not in ['low','medium','high'] %}\
7731 {{ raise_exception('unsupported effort') }}\
7732 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7733 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7734
7735 fn graded_template() -> chat_template::PromptTemplate {
7736 chat_template::PromptTemplate::from_gguf_metadata(
7737 Some(GRADED),
7738 Some("qwen3"),
7739 false,
7740 true,
7741 None,
7742 None,
7743 )
7744 }
7745
7746 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7747 serde_json::from_value(value).expect("request")
7748 }
7749
7750 #[test]
7757 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7758 let req = chat_request(serde_json::json!({
7759 "model": "m",
7760 "messages": [{"role": "user", "content": "hi"}],
7761 "grammar": "root ::= \"a\"+",
7762 }));
7763 req.validate_supported_fields()
7764 .expect("a valid grammar is a valid request");
7765 let params = req
7766 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7767 .expect("a valid grammar compiles at params time too");
7768 assert!(
7769 params.grammar.is_some(),
7770 "the grammar was dropped between the wire and the sampler"
7771 );
7772 assert!(
7773 params.needs_vocab_logits(),
7774 "a grammar request that may fold lm_head into a GPU argmax is \
7775 a grammar request served unconstrained"
7776 );
7777
7778 let plain = chat_request(serde_json::json!({
7779 "model": "m",
7780 "messages": [{"role": "user", "content": "hi"}],
7781 }));
7782 assert!(plain
7783 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7784 .unwrap()
7785 .grammar
7786 .is_none());
7787 }
7788
7789 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7790 chat_request(serde_json::json!({
7791 "model": "m",
7792 "messages": [{"role": "user", "content": "weather in Rome?"}],
7793 "tools": [weather_tool()],
7794 "tool_choice": tool_choice,
7795 }))
7796 }
7797
7798 #[test]
7802 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7803 for choice in [
7804 serde_json::json!("required"),
7805 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7806 ] {
7807 let req = tool_request(choice.clone());
7808 req.validate_supported_fields()
7809 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7810 let params = req
7811 .generation_params_for_template(
7812 &graded_template(),
7813 "Qwen3-8B",
7814 crate::sampling_knobs::SamplerModel::absent(),
7815 )
7816 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7817 let grammar = params
7818 .grammar
7819 .as_ref()
7820 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7821 assert!(
7822 grammar.is_awaiting_trigger(),
7823 "the model must be free to think before it calls"
7824 );
7825 assert!(
7826 !grammar.allows_eog(),
7827 "{choice} must not be able to end the turn without a call"
7828 );
7829 assert!(
7835 params.needs_vocab_logits(),
7836 "{choice} would let a backend return a token id instead of logits"
7837 );
7838 assert!(
7839 !generate::greedy_gpu_fold_allowed(¶ms),
7840 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7841 );
7842 }
7843 }
7844
7845 #[test]
7847 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7848 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7849 let req = tool_request(choice.clone());
7850 req.validate_supported_fields().expect("still supported");
7851 let params = match req.generation_params_for_template(
7852 &graded_template(),
7853 "Qwen3-8B",
7854 crate::sampling_knobs::SamplerModel::absent(),
7855 ) {
7856 Ok(p) => p,
7857 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7858 };
7859 assert!(
7860 params.grammar.is_none(),
7861 "{choice} does not force a call and must not be constrained"
7862 );
7863 }
7864 }
7865
7866 #[test]
7869 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7870 let req = chat_request(serde_json::json!({
7872 "model": "m",
7873 "messages": [{"role": "user", "content": "hi"}],
7874 "tool_choice": "required",
7875 }));
7876 let (status, _) = req
7877 .validate_supported_fields()
7878 .expect_err("nothing to call");
7879 assert_eq!(status, StatusCode::BAD_REQUEST);
7880
7881 let req =
7883 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7884 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7885 assert_eq!(status, StatusCode::BAD_REQUEST);
7886 assert_eq!(body["error"]["param"], "tool_choice");
7887
7888 let req = tool_request(serde_json::json!({"type": "function"}));
7890 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7891 assert_eq!(status, StatusCode::BAD_REQUEST);
7892
7893 let req = chat_request(serde_json::json!({
7895 "model": "m",
7896 "messages": [{"role": "user", "content": "hi"}],
7897 "tools": [weather_tool()],
7898 "tool_choice": "required",
7899 "grammar": "root ::= \"a\"+",
7900 }));
7901 let (status, _) = req
7902 .validate_supported_fields()
7903 .expect_err("a grammar and a forced call are two constraints");
7904 assert_eq!(status, StatusCode::BAD_REQUEST);
7905
7906 let req = tool_request(serde_json::json!("required"));
7912 let (status, Json(body)) = match req.generation_params_for_template(
7913 &graded_template(),
7914 "muse-glimmer-8b",
7915 crate::sampling_knobs::SamplerModel::absent(),
7916 ) {
7917 Err(e) => e,
7918 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
7919 };
7920 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7921 assert!(
7922 body["error"]["message"]
7923 .as_str()
7924 .unwrap()
7925 .contains("muse_glimmer"),
7926 "{body}"
7927 );
7928
7929 let req = tool_request(serde_json::json!("required"));
7935 let params = req
7936 .generation_params_for_template(
7937 &graded_template(),
7938 "gemma-4-E2B-it",
7939 crate::sampling_knobs::SamplerModel::absent(),
7940 )
7941 .expect("a gemma4 forced tool_choice is served");
7942 assert!(
7943 params.grammar.is_some(),
7944 "a forced tool_choice must arrive as the generation's grammar"
7945 );
7946 }
7947
7948 #[test]
7951 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
7952 let req = chat_request(serde_json::json!({
7953 "model": "m",
7954 "messages": [{"role": "user", "content": "hi"}],
7955 "grammar": "root ::= \"a",
7956 }));
7957 let (status, Json(body)) = req
7958 .validate_supported_fields()
7959 .expect_err("this does not parse");
7960 assert_eq!(status, StatusCode::BAD_REQUEST);
7961 assert_eq!(body["error"]["param"], "grammar");
7962 assert!(
7963 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
7964 .is_err(),
7965 "and again at params time"
7966 );
7967 }
7968
7969 #[test]
7976 fn response_format_json_schema_becomes_the_requests_grammar() {
7977 let req = chat_request(serde_json::json!({
7978 "model": "m",
7979 "messages": [{"role": "user", "content": "hi"}],
7980 "response_format": {
7981 "type": "json_schema",
7982 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
7983 },
7984 }));
7985 req.validate_supported_fields()
7986 .expect("a boolean schema converts");
7987 let params = req
7988 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7989 .expect("and compiles");
7990 let grammar = params.grammar.expect("the schema is the grammar");
7991 let mut g = (*grammar).clone();
7992 g.accept_token(0, b"true").expect("a boolean is accepted");
7993 assert!(g.allows_eog(), "and completes the parse");
7994 assert!(
7995 !params.json_object,
7996 "a schema is not the json_object character-class mask"
7997 );
7998 }
7999
8000 #[test]
8004 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8005 let req = chat_request(serde_json::json!({
8006 "model": "m",
8007 "messages": [{"role": "user", "content": "hi"}],
8008 "response_format": {
8009 "type": "json_schema",
8010 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8011 },
8012 }));
8013 let (status, Json(body)) = req
8014 .validate_supported_fields()
8015 .expect_err("minimum has no grammar in this port");
8016 assert_eq!(status, StatusCode::BAD_REQUEST);
8017 assert!(
8018 body["error"]["message"]
8019 .as_str()
8020 .expect("a message")
8021 .contains("minimum"),
8022 "the refusal must name the keyword: {body}"
8023 );
8024 assert!(
8025 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8026 .is_err(),
8027 "and again at params time"
8028 );
8029 }
8030
8031 #[test]
8037 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8038 let req = chat_request(serde_json::json!({
8039 "model": "m",
8040 "messages": [{"role": "user", "content": "hi"}],
8041 "tool_choice": "required",
8042 "tools": [{
8043 "type": "function",
8044 "function": {"name": "f", "parameters": {"type": "object"}},
8045 }],
8046 "response_format": {
8047 "type": "json_schema",
8048 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8049 },
8050 }));
8051 let (status, Json(body)) = req
8052 .validate_supported_fields()
8053 .expect_err("two constraints, one generation");
8054 assert_eq!(status, StatusCode::BAD_REQUEST);
8055 assert_eq!(body["error"]["param"], "tool_choice");
8056 }
8057
8058 #[test]
8061 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8062 let req = chat_request(serde_json::json!({
8063 "model": "m",
8064 "messages": [{"role": "user", "content": "hi"}],
8065 }));
8066 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8067 }
8068
8069 #[test]
8074 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8075 let asked = chat_request(serde_json::json!({
8076 "model": "m",
8077 "messages": [{"role": "user", "content": "hi"}],
8078 "min_p": 0.07,
8079 }));
8080 assert_eq!(
8081 asked
8082 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8083 .expect("knobs")
8084 .min_p,
8085 0.07
8086 );
8087
8088 let silent = chat_request(serde_json::json!({
8089 "model": "m",
8090 "messages": [{"role": "user", "content": "hi"}],
8091 }));
8092 assert_eq!(
8093 silent
8094 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8095 .expect("knobs")
8096 .min_p,
8097 0.0,
8098 "an unset min_p must be off, not llama.cpp's CLI default"
8099 );
8100 }
8101
8102 #[test]
8111 fn no_sampler_knob_is_missing_from_the_cache_key() {
8112 let base = serde_json::json!({
8113 "model": "m",
8114 "messages": [{"role": "user", "content": "hi"}],
8115 "seed": 1,
8116 });
8117 let key_for = |body: serde_json::Value| {
8118 let req = chat_request(body);
8119 let params = req
8120 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8121 .expect("params");
8122 req.cache_key("prompt", ¶ms)
8123 };
8124 let baseline = key_for(base.clone());
8125 for (knob, value) in [
8126 ("temperature", serde_json::json!(0.5)),
8127 ("top_p", serde_json::json!(0.9)),
8128 ("min_p", serde_json::json!(0.05)),
8129 ("top_k", serde_json::json!(40)),
8130 ("repetition_penalty", serde_json::json!(1.1)),
8131 ("presence_penalty", serde_json::json!(0.3)),
8132 ("frequency_penalty", serde_json::json!(0.3)),
8133 (
8134 "samplers",
8135 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8136 ),
8137 ] {
8138 let mut body = base.clone();
8139 body[knob] = value;
8140 assert_ne!(
8141 key_for(body),
8142 baseline,
8143 "`{knob}` is not in the cache key: two requests differing \
8144 only in it would share one cached answer"
8145 );
8146 }
8147 }
8148
8149 #[test]
8159 fn no_constraint_is_missing_from_the_cache_key() {
8160 let base = serde_json::json!({
8161 "model": "m",
8162 "messages": [{"role": "user", "content": "pick one"}],
8163 });
8164 let key_for = |body: serde_json::Value| {
8165 let req = chat_request(body);
8166 let params = req
8167 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8168 .expect("params");
8169 req.cache_key("prompt", ¶ms)
8170 };
8171 let baseline = key_for(base.clone());
8172 for (field, value) in [
8173 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8174 (
8175 "response_format",
8176 serde_json::json!({"type": "json_object"}),
8177 ),
8178 (
8179 "response_format",
8180 serde_json::json!({"type": "json_schema", "json_schema": {
8181 "name": "answer",
8182 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8183 }}),
8184 ),
8185 ("ignore_eos", serde_json::json!(true)),
8186 ("stop", serde_json::json!(["\n"])),
8187 ("max_tokens", serde_json::json!(7)),
8188 ] {
8189 let mut body = base.clone();
8190 body[field] = value.clone();
8191 assert_ne!(
8192 key_for(body),
8193 baseline,
8194 "`{field}: {value}` is not in the cache key: two requests \
8195 differing only in it would share one cached answer"
8196 );
8197 }
8198 }
8199
8200 #[test]
8204 fn an_explicit_zero_output_budget_is_a_client_error() {
8205 let req = chat_request(serde_json::json!({
8206 "model": "m",
8207 "messages": [{"role": "user", "content": "hi"}],
8208 "max_tokens": 0,
8209 }));
8210 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8211 assert_eq!(status, StatusCode::BAD_REQUEST);
8212 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8213 }
8214
8215 #[test]
8218 fn a_request_can_turn_thinking_off() {
8219 let template = graded_template();
8220 for body in [
8221 serde_json::json!({
8222 "model": "m",
8223 "messages": [{"role": "user", "content": "hi"}],
8224 "reasoning_effort": "none",
8225 }),
8226 serde_json::json!({
8227 "model": "m",
8228 "messages": [{"role": "user", "content": "hi"}],
8229 "thinking": {"type": "disabled"},
8230 }),
8231 ] {
8232 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8233 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8234 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8235 assert!(!kwargs.contains_key("reasoning_effort"));
8238 }
8239 }
8240
8241 #[test]
8244 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8245 let template = graded_template();
8246 let kwargs = chat_request(serde_json::json!({
8247 "model": "m",
8248 "messages": [{"role": "user", "content": "hi"}],
8249 "reasoning_effort": "high",
8250 "thinking": {"type": "disabled"},
8251 }))
8252 .resolve_template_kwargs(&template);
8253 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8254 assert!(!kwargs.contains_key("reasoning_effort"));
8255 }
8256
8257 #[test]
8260 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8261 let req = chat_request(serde_json::json!({
8262 "model": "m",
8263 "messages": [{"role": "user", "content": "hi"}],
8264 "thinking": {"type": "disable"},
8265 }));
8266 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8267 assert_eq!(status, StatusCode::BAD_REQUEST);
8268 }
8269
8270 #[test]
8273 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8274 let template = graded_template();
8275 let kwargs = chat_request(serde_json::json!({
8276 "model": "m",
8277 "messages": [{"role": "user", "content": "hi"}],
8278 "reasoning_effort": "none",
8279 "chat_template_kwargs": {"enable_thinking": true},
8280 }))
8281 .resolve_template_kwargs(&template);
8282 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8283 }
8284
8285 #[test]
8289 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8290 let template = graded_template();
8291 let req = chat_request(serde_json::json!({
8292 "model": "m",
8293 "messages": [{"role": "user", "content": "hi"}],
8294 "reasoning_effort": "minimal",
8295 }));
8296 let kwargs = req.resolve_template_kwargs(&template);
8297 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8298 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8299 assert!(prompt.starts_with("E:low|"), "{prompt}");
8300 }
8301
8302 #[test]
8306 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8307 let template = graded_template();
8308 let req = chat_request(serde_json::json!({
8309 "model": "m",
8310 "messages": [{"role": "user", "content": "hi"}],
8311 "chat_template_kwargs": {"reasoning_effort": "none"},
8312 }));
8313 let kwargs = req.resolve_template_kwargs(&template);
8314 assert!(!kwargs.contains_key("reasoning_effort"));
8315 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8316 assert_eq!(prompt, "hi");
8317 }
8318
8319 #[test]
8322 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8323 let template = graded_template();
8324 let req = chat_request(serde_json::json!({
8325 "model": "m",
8326 "messages": [{"role": "user", "content": "hi"}],
8327 "reasoning_effort": "low",
8328 "chat_template_kwargs": {"reasoning_effort": "high"},
8329 }));
8330 assert_eq!(
8331 req.resolve_template_kwargs(&template)["reasoning_effort"],
8332 serde_json::json!("high")
8333 );
8334 }
8335
8336 #[test]
8340 fn offering_tools_turns_thinking_on_by_itself() {
8341 let template = graded_template();
8342 let quiet = chat_request(serde_json::json!({
8343 "model": "m",
8344 "messages": [{"role": "user", "content": "hi"}],
8345 }));
8346 assert!(!quiet
8347 .resolve_template_kwargs(&template)
8348 .contains_key("enable_thinking"));
8349
8350 let with_tools = chat_request(serde_json::json!({
8351 "model": "m",
8352 "messages": [{"role": "user", "content": "hi"}],
8353 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8354 }));
8355 let kwargs = with_tools.resolve_template_kwargs(&template);
8356 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8357 let prompt =
8358 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8359 assert!(prompt.starts_with("THINK|"), "{prompt}");
8360 }
8361
8362 #[test]
8367 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8368 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8369 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8370 Some("qwen3"),
8371 false,
8372 true,
8373 None,
8374 None,
8375 );
8376 let req = chat_request(serde_json::json!({
8377 "model": "m",
8378 "messages": [{"role": "user", "content": "hi"}],
8379 "chat_template_kwargs": {"enable_thinking": true},
8380 }));
8381 let kwargs = req.resolve_template_kwargs(&opener);
8382 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8383 assert!(prompt.ends_with("<think>"), "{prompt}");
8384
8385 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8388 let (message, _) = build_response_message(
8389 "weighing it up</think>Paris.".to_string(),
8390 &[],
8391 posture,
8392 "stop",
8393 );
8394 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8395 assert_eq!(message.content.as_deref(), Some("Paris."));
8396
8397 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8400 let (message, _) = build_response_message(
8401 "weighing it up</think>Paris.".to_string(),
8402 &[],
8403 closed,
8404 "stop",
8405 );
8406 assert_eq!(message.reasoning_content, None);
8407 }
8408
8409 #[test]
8410 fn stop_param_accepts_both_single_string_and_array() {
8411 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8412 "model": "m",
8413 "messages": [{"role": "user", "content": "hi"}],
8414 "stop": "END",
8415 }))
8416 .unwrap();
8417 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8418
8419 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8420 "model": "m",
8421 "messages": [{"role": "user", "content": "hi"}],
8422 "stop": ["A", "B"],
8423 }))
8424 .unwrap();
8425 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8426 }
8427
8428 #[test]
8429 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8430 let model = test_model();
8431 let result = run_generation(
8432 &model,
8433 "hello",
8434 &greedy_params(4),
8435 None,
8436 None,
8437 None,
8438 None,
8439 None,
8440 None,
8441 );
8442 assert!(matches!(
8443 result,
8444 Err(generate::DecodeError::TokenOutOfVocab { .. })
8445 ));
8446 }
8447
8448 #[test]
8452 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8453 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8455 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8456
8457 let holder_pool = Arc::clone(&pool);
8458 let holder = std::thread::spawn(move || {
8459 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8460 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8462 drop(held);
8463 });
8464 std::thread::sleep(Duration::from_millis(15));
8465
8466 let config = generate::KvPoolConfig {
8467 pool,
8468 queue_wait: Duration::ZERO,
8469 };
8470 let result = run_generation(
8471 &model,
8472 &prompt,
8473 &greedy_params(4),
8474 Some(&config),
8475 None,
8476 None,
8477 None,
8478 None,
8479 None,
8480 );
8481 assert!(matches!(
8482 result,
8483 Err(generate::DecodeError::KvPoolExhausted)
8484 ));
8485
8486 let (status, _body) = decode_error_response(result.unwrap_err());
8487 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8488 holder.join().unwrap();
8489 }
8490
8491 #[test]
8499 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8500 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8502 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8504 let config = generate::KvPoolConfig {
8505 pool,
8506 queue_wait: Duration::ZERO,
8507 };
8508
8509 let result = run_generation(
8510 &model,
8511 &prompt,
8512 &greedy_params(4),
8513 Some(&config),
8514 None,
8515 None,
8516 None,
8517 None,
8518 None,
8519 );
8520 let err = result.expect_err("one block cannot hold two layers' caches");
8521 assert!(
8522 matches!(
8523 &err,
8524 generate::DecodeError::KvBudgetExceeded { binding, .. }
8525 if *binding == frink_models::Ceiling::DeviceMemory.code()
8526 ),
8527 "expected an immovable device-memory refusal, got {err:?}"
8528 );
8529 let (status, _body) = decode_error_response(err);
8530 assert_eq!(status, StatusCode::BAD_REQUEST);
8531 }
8532
8533 #[test]
8539 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8540 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8541 queued: 512,
8542 cap: 512,
8543 });
8544 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8545 assert_eq!(body["error"]["retry_after_seconds"], 1);
8546 let message = body["error"]["message"].as_str().expect("message");
8547 assert!(message.contains("512"), "{message}");
8548 }
8549
8550 #[test]
8551 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8552 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8553 token: 99,
8554 vocab_size: 32,
8555 });
8556 assert!(
8557 body["error"]["retry_after_seconds"].is_null(),
8558 "retrying a prompt this model cannot tokenize never helps"
8559 );
8560 }
8561
8562 #[test]
8563 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8564 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8565 token: 99,
8566 vocab_size: 32,
8567 });
8568 assert_eq!(status, StatusCode::BAD_REQUEST);
8569 }
8570
8571 #[test]
8572 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8573 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8575 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8576 let config = generate::KvPoolConfig {
8577 pool: pool.clone(),
8578 queue_wait: Duration::ZERO,
8579 };
8580
8581 let (_, finish, _usage) = run_generation(
8582 &model,
8583 &prompt,
8584 &greedy_params(4),
8585 Some(&config),
8586 None,
8587 None,
8588 None,
8589 None,
8590 None,
8591 )
8592 .unwrap();
8593 assert_eq!(finish, FinishReason::Length);
8594 assert_eq!(
8595 pool.lock().unwrap().free_blocks(),
8596 2,
8597 "a completed request must return its blocks to the pool"
8598 );
8599 }
8600
8601 #[tokio::test]
8606 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8607 let model = Arc::new(test_model());
8608 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8609
8610 let mut handles = Vec::new();
8611 for _ in 0..8 {
8612 let model = Arc::clone(&model);
8613 let prompt = prompt.clone();
8614 handles.push(tokio::task::spawn_blocking(move || {
8615 run_generation(
8616 &model,
8617 &prompt,
8618 &greedy_params(6),
8619 None,
8620 None,
8621 None,
8622 None,
8623 None,
8624 None,
8625 )
8626 .unwrap()
8627 }));
8628 }
8629
8630 let mut results = Vec::new();
8631 for h in handles {
8632 results.push(h.await.unwrap());
8633 }
8634 for r in &results[1..] {
8639 assert_eq!(r.0, results[0].0, "decoded chunks must match");
8640 assert_eq!(r.1, results[0].1, "finish reason must match");
8641 assert_eq!(
8642 r.2.prompt_tokens, results[0].2.prompt_tokens,
8643 "prompt token count must match"
8644 );
8645 assert_eq!(
8646 r.2.completion_tokens, results[0].2.completion_tokens,
8647 "completion token count must match"
8648 );
8649 }
8650 }
8651
8652 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8659 let mut header_entries = Vec::new();
8660 let mut data = Vec::new();
8661 for (name, shape, values) in tensors {
8662 let start = data.len();
8663 for v in values {
8664 data.extend_from_slice(&v.to_le_bytes());
8665 }
8666 let end = data.len();
8667 let shape_str = shape
8668 .iter()
8669 .map(|d| d.to_string())
8670 .collect::<Vec<_>>()
8671 .join(",");
8672 header_entries.push(format!(
8673 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8674 ));
8675 }
8676 let header = format!("{{{}}}", header_entries.join(","));
8677 let header_bytes = header.as_bytes();
8678 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8679 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8680 out.extend_from_slice(header_bytes);
8681 out.extend_from_slice(&data);
8682 out
8683 }
8684
8685 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8696 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8697 use frink_models::kimi_loader::KimiRealHparams;
8698 use frink_moe::{GatingFunction, MoeLayerConfig};
8699
8700 let hidden_dim = 8;
8701 let kda_num_heads = 2;
8702 let kda_head_dim = 3;
8703 let kda_proj = kda_num_heads * kda_head_dim;
8704 let conv_kernel = 4;
8705 let dense_intermediate = 5;
8706 let vocab_size = 256;
8710 let mla_num_heads = 1;
8711 let mla_q_lora_rank = 2;
8712 let mla_kv_lora_rank = 2;
8713 let mla_qk_nope_head_dim = 2;
8714 let mla_qk_rope_head_dim = 2;
8715 let mla_v_head_dim = 2;
8716
8717 let model_cfg = frink_models::ModelConfig {
8718 rope_layers: frink_models::rope_layers::RopeLayers::All,
8719 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
8720 name: "synthetic-kimi-server-test",
8721 n_layers: 1,
8722 n_mtp_blocks: 0,
8723 hidden_dim,
8724 n_heads: 1,
8725 n_kv_heads: 1,
8726 head_dim: 4,
8727 v_head_dim: None,
8728 vocab_size,
8729 rope_theta: 10000.0,
8730 rms_norm_eps: 1e-5,
8731 post_norm_eps: 1e-5,
8732 sliding_window: None,
8733 moe: MoeLayerConfig {
8734 expert_weights_scale: 1.0,
8735 routed_weight_before_ffn: false,
8736 n_experts: 1,
8737 n_experts_active: 1,
8738 n_shared_experts: 0,
8739 hidden_dim,
8740 expert_ffn_dim: 4,
8741 gating: GatingFunction::Sigmoid,
8742 norm_topk_prob: true,
8743 expert_group_count: None,
8744 expert_group_used_count: None,
8745 },
8746 n_dense_leading_layers: 1,
8751 moe_interleave_step: None,
8752 norm_function: frink_models::norm::NormFunction::Rms,
8753 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8754 kda_layers: vec![1],
8755 full_attn_layers: vec![],
8756 mla: MlaConfig {
8757 num_heads: mla_num_heads,
8758 q_lora_rank: mla_q_lora_rank,
8759 kv_lora_rank: mla_kv_lora_rank,
8760 qk_nope_head_dim: mla_qk_nope_head_dim,
8761 qk_rope_head_dim: mla_qk_rope_head_dim,
8762 v_head_dim: mla_v_head_dim,
8763 use_output_gate: true,
8764 rope: None,
8765 },
8766 kda: KdaConfig {
8767 num_heads: kda_num_heads,
8768 head_dim: kda_head_dim,
8769 short_conv_kernel_size: conv_kernel,
8770 gate_lower_bound: -5.0,
8771 use_full_rank_gate: true,
8772 },
8773 }),
8774 rope_freqs: None,
8775 rope_attn_factor: 1.0,
8776 rope_dim: None,
8777 rope_dim_swa: None,
8778 rope_freqs_long: None,
8779 rope_freqs_short: None,
8780 rope_orig_ctx: None,
8781 rope_layout: frink_models::config::RopeLayout::Neox,
8782 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
8783 swa_layers: frink_models::swa_layers::SwaLayers::All,
8784 attn_logit_softcap: None,
8785 final_logit_softcap: None,
8786 embedding_scale: None,
8787 residual_scale: None,
8788 normed_residual_scale: None,
8789 clamp_kqv: None,
8790 attn_temperature: None,
8791 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
8792 block_sub_norms: false,
8793 parallel_residual: false,
8794 learned_positions: false,
8795 attn_value_scale: None,
8796 alibi_max_bias: None,
8797 layer_loops: None,
8798 skip_stream: false,
8799 parallel_ssm: false,
8800 swa_chunked: false,
8801 weightless_qk_norm: false,
8802 logit_multiplier: None,
8803 attention_scale: None,
8804 rope_theta_swa: None,
8805 ffn_activation: frink_models::config::FfnActivation::Swiglu,
8806 best_effort_fields: &["synthetic test config, not a real preset"],
8807 };
8808 let hp = KimiRealHparams {
8809 hidden_dim,
8810 kda_num_heads,
8811 kda_head_dim,
8812 mla_num_heads,
8813 mla_q_lora_rank,
8814 mla_kv_lora_rank,
8815 mla_qk_nope_head_dim,
8816 mla_qk_rope_head_dim,
8817 mla_v_head_dim,
8818 dense_intermediate_dim: dense_intermediate,
8819 moe_hidden_dim: hidden_dim,
8820 moe_intermediate_dim: 4,
8821 n_experts: 1,
8822 num_shared_experts: 0,
8823 };
8824
8825 let prefix = "language_model.model.layers.0";
8829 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8830 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8831 name: String,
8832 shape: Vec<usize>,
8833 n: usize| {
8834 tensors.push((name, shape, vec![0.01f32; n]));
8835 };
8836 push(
8837 &mut tensors,
8838 format!("{prefix}.input_layernorm.weight"),
8839 vec![hidden_dim],
8840 hidden_dim,
8841 );
8842 push(
8843 &mut tensors,
8844 format!("{prefix}.post_attention_layernorm.weight"),
8845 vec![hidden_dim],
8846 hidden_dim,
8847 );
8848 push(
8849 &mut tensors,
8850 format!("{prefix}.self_attention_res_norm.weight"),
8851 vec![hidden_dim],
8852 hidden_dim,
8853 );
8854 push(
8855 &mut tensors,
8856 format!("{prefix}.self_attention_res_proj.weight"),
8857 vec![1, hidden_dim],
8858 hidden_dim,
8859 );
8860 push(
8861 &mut tensors,
8862 format!("{prefix}.mlp_res_norm.weight"),
8863 vec![hidden_dim],
8864 hidden_dim,
8865 );
8866 push(
8867 &mut tensors,
8868 format!("{prefix}.mlp_res_proj.weight"),
8869 vec![1, hidden_dim],
8870 hidden_dim,
8871 );
8872 push(
8873 &mut tensors,
8874 format!("{prefix}.self_attn.q_proj.weight"),
8875 vec![kda_proj, hidden_dim],
8876 kda_proj * hidden_dim,
8877 );
8878 push(
8879 &mut tensors,
8880 format!("{prefix}.self_attn.k_proj.weight"),
8881 vec![kda_proj, hidden_dim],
8882 kda_proj * hidden_dim,
8883 );
8884 push(
8885 &mut tensors,
8886 format!("{prefix}.self_attn.v_proj.weight"),
8887 vec![kda_proj, hidden_dim],
8888 kda_proj * hidden_dim,
8889 );
8890 push(
8891 &mut tensors,
8892 format!("{prefix}.self_attn.q_conv1d.weight"),
8893 vec![kda_proj, 1, conv_kernel],
8894 kda_proj * conv_kernel,
8895 );
8896 push(
8897 &mut tensors,
8898 format!("{prefix}.self_attn.k_conv1d.weight"),
8899 vec![kda_proj, 1, conv_kernel],
8900 kda_proj * conv_kernel,
8901 );
8902 push(
8903 &mut tensors,
8904 format!("{prefix}.self_attn.v_conv1d.weight"),
8905 vec![kda_proj, 1, conv_kernel],
8906 kda_proj * conv_kernel,
8907 );
8908 push(
8909 &mut tensors,
8910 format!("{prefix}.self_attn.A_log"),
8911 vec![kda_num_heads],
8912 kda_num_heads,
8913 );
8914 push(
8915 &mut tensors,
8916 format!("{prefix}.self_attn.f_a_proj.weight"),
8917 vec![kda_head_dim, hidden_dim],
8918 kda_head_dim * hidden_dim,
8919 );
8920 push(
8921 &mut tensors,
8922 format!("{prefix}.self_attn.f_b_proj.weight"),
8923 vec![kda_proj, kda_head_dim],
8924 kda_proj * kda_head_dim,
8925 );
8926 push(
8927 &mut tensors,
8928 format!("{prefix}.self_attn.dt_bias"),
8929 vec![kda_proj],
8930 kda_proj,
8931 );
8932 push(
8933 &mut tensors,
8934 format!("{prefix}.self_attn.b_proj.weight"),
8935 vec![kda_num_heads, hidden_dim],
8936 kda_num_heads * hidden_dim,
8937 );
8938 push(
8939 &mut tensors,
8940 format!("{prefix}.self_attn.g_proj.weight"),
8941 vec![kda_proj, hidden_dim],
8942 kda_proj * hidden_dim,
8943 );
8944 push(
8945 &mut tensors,
8946 format!("{prefix}.self_attn.o_norm.weight"),
8947 vec![kda_head_dim],
8948 kda_head_dim,
8949 );
8950 push(
8951 &mut tensors,
8952 format!("{prefix}.self_attn.o_proj.weight"),
8953 vec![hidden_dim, kda_proj],
8954 hidden_dim * kda_proj,
8955 );
8956 push(
8957 &mut tensors,
8958 format!("{prefix}.mlp.gate_proj.weight"),
8959 vec![dense_intermediate, hidden_dim],
8960 dense_intermediate * hidden_dim,
8961 );
8962 push(
8963 &mut tensors,
8964 format!("{prefix}.mlp.up_proj.weight"),
8965 vec![dense_intermediate, hidden_dim],
8966 dense_intermediate * hidden_dim,
8967 );
8968 push(
8969 &mut tensors,
8970 format!("{prefix}.mlp.down_proj.weight"),
8971 vec![hidden_dim, dense_intermediate],
8972 hidden_dim * dense_intermediate,
8973 );
8974 push(
8975 &mut tensors,
8976 "language_model.model.embed_tokens.weight".to_string(),
8977 vec![vocab_size, hidden_dim],
8978 vocab_size * hidden_dim,
8979 );
8980 push(
8981 &mut tensors,
8982 "language_model.lm_head.weight".to_string(),
8983 vec![vocab_size, hidden_dim],
8984 vocab_size * hidden_dim,
8985 );
8986 push(
8987 &mut tensors,
8988 "language_model.model.norm.weight".to_string(),
8989 vec![hidden_dim],
8990 hidden_dim,
8991 );
8992 push(
8993 &mut tensors,
8994 "language_model.model.output_attn_res_norm.weight".to_string(),
8995 vec![hidden_dim],
8996 hidden_dim,
8997 );
8998 push(
8999 &mut tensors,
9000 "language_model.model.output_attn_res_proj.weight".to_string(),
9001 vec![1, hidden_dim],
9002 hidden_dim,
9003 );
9004
9005 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9018 let dir = std::env::temp_dir().join(format!(
9019 "frink_server_kimi_e2e_test_{}_{}",
9020 std::process::id(),
9021 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9022 ));
9023 std::fs::create_dir_all(&dir).unwrap();
9024 let shard_bytes = write_safetensors_shard(&tensors);
9025 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9026 let map_entries: Vec<String> = tensors
9027 .iter()
9028 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9029 .collect();
9030 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9031 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9032
9033 use base64::Engine;
9037 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9038 .map(|b| {
9039 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9040 format!("{b64} {b}")
9041 })
9042 .collect();
9043 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9044
9045 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9046 .expect("must load the synthetic Kimi checkpoint end to end");
9047 std::fs::remove_dir_all(&dir).ok();
9048 loaded
9049 }
9050
9051 #[test]
9059 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9060 let loaded = build_synthetic_kimi_loaded();
9061 let state = build_app_state(
9062 StartupModels {
9063 loaded: model::LoadedModel::Kimi(loaded),
9064 embedding: None,
9065 },
9066 None,
9067 None,
9068 None,
9069 false,
9070 None,
9071 Arc::new(health::Detection::ready(health::probe_backends())),
9072 );
9073 let active = state.active().expect("a freshly built state has a model");
9074 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9075 assert!(!active.is_synthetic());
9076
9077 let (_chunks, finish, _usage) = run_generation(
9078 active.generative().unwrap(),
9079 "hi",
9080 &greedy_params(5),
9081 None,
9082 None,
9083 None,
9084 None,
9085 None,
9086 None,
9087 )
9088 .expect("a real Kimi checkpoint must generate without error");
9089 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9090 }
9091
9092 #[test]
9105 fn a_grammar_constrains_the_engine_decode_path() {
9106 let loaded = build_synthetic_kimi_loaded();
9107 let state = build_app_state(
9108 StartupModels {
9109 loaded: model::LoadedModel::Kimi(loaded),
9110 embedding: None,
9111 },
9112 None,
9113 None,
9114 None,
9115 false,
9116 None,
9117 Arc::new(health::Detection::ready(health::probe_backends())),
9118 );
9119 let active = state.active().expect("a freshly built state has a model");
9120
9121 let run = |grammar: Option<&str>| {
9122 let mut params = greedy_params(6);
9123 params.grammar = grammar.map(|src| {
9124 Arc::new(
9125 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9126 .expect("test grammar parses"),
9127 )
9128 });
9129 run_generation(
9130 active.generative().unwrap(),
9131 "hi",
9132 ¶ms,
9133 None,
9134 None,
9135 None,
9136 None,
9137 None,
9138 None,
9139 )
9140 };
9141
9142 let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
9143 let unconstrained = chunks.concat();
9144 assert!(
9145 unconstrained.chars().any(|c| c != 'a'),
9146 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9147 constrained run below would prove nothing"
9148 );
9149
9150 let (chunks, finish, _) =
9151 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9152 let constrained = chunks.concat();
9153 assert!(
9154 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9155 "the engine decode path served text its grammar forbids ({constrained:?}): \
9156 the constraint was dropped between `generate_engine` and the sampler"
9157 );
9158 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9159 }
9160
9161 #[test]
9171 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9172 let loaded = build_synthetic_kimi_loaded();
9173 let state = build_app_state(
9174 StartupModels {
9175 loaded: model::LoadedModel::Kimi(loaded),
9176 embedding: None,
9177 },
9178 None,
9179 None,
9180 None,
9181 false,
9182 None,
9183 Arc::new(health::Detection::ready(health::probe_backends())),
9184 );
9185
9186 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9187 let kv_pool_config = generate::KvPoolConfig {
9188 pool: pool.clone(),
9189 queue_wait: Duration::ZERO,
9190 };
9191 let pc = Mutex::new(PrefixCache::new(4));
9192
9193 run_generation(
9194 state
9195 .active()
9196 .expect("a freshly built state has a model")
9197 .generative()
9198 .unwrap(),
9199 "hi",
9200 &greedy_params(5),
9201 Some(&kv_pool_config),
9202 None,
9203 Some(&pc),
9204 None,
9205 None,
9206 None,
9207 )
9208 .expect("a real Kimi checkpoint must generate without error");
9209
9210 assert_eq!(
9211 pool.lock().unwrap().free_blocks(),
9212 4,
9213 "the KV pool must be completely untouched by a Kimi request"
9214 );
9215 let stats = pc.lock().unwrap().stats();
9216 assert_eq!(
9217 stats.hits + stats.misses,
9218 0,
9219 "the prefix cache must never be consulted for a Kimi request"
9220 );
9221 }
9222}