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 request_tail;
66mod rerank;
67mod response_cache;
68pub(crate) mod responses;
69mod resume;
70mod sample_step;
71mod sampling_knobs;
72mod sampling_loop;
73mod security;
74mod serving;
75mod session;
76mod slots;
77mod sse;
78mod stats;
79mod stop;
80mod stream_events;
81mod tasks;
82mod tool_grammar;
83mod unimplemented_fields;
84mod unsupported_sampling;
85mod utf8_stream;
86
87use std::cell::RefCell;
88use std::convert::Infallible;
89use std::net::SocketAddr;
90use std::path::PathBuf;
91use std::rc::Rc;
92use std::sync::{Arc, Mutex, MutexGuard};
93use std::time::Duration;
94
95use axum::{
96 extract::State,
97 http::StatusCode,
98 response::sse::{Event, Sse},
99 response::{IntoResponse, Response},
100 routing::{get, post},
101 Json, Router,
102};
103use serde::{Deserialize, Serialize};
104
105use cli::apply_cli_overrides;
106pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
107
108use frink_core::cache::KvBlockPool;
109use frink_models::kimi_tokenizer::KimiTokenizer;
110use frink_models::sampling::SamplingParams;
111use frink_models::tokenizer::{SpecialTokens, StopTokens};
112use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
113use generate::{FinishReason, GenerationParams};
114pub(crate) use loaded::{ActiveModel, Loaded};
115use model::ServerTokenizer;
116use rerank::encoder_endpoints;
117use response_cache::ResponseCache;
118use sampling_knobs::SamplingKnobs;
119
120#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
129 Gguf(GgufModel),
130 Kimi(KimiModel),
131 Mla(MlaModel),
132 Gemma4(Gemma4Model),
133 Glm52(Glm52Model),
134}
135
136pub(crate) struct GgufModel {
137 decoder: Arc<Decoder>,
138 tokenizer: Arc<ServerTokenizer>,
139 stop_tokens: StopTokens,
140 bos_id: Option<usize>,
141 is_synthetic: bool,
142 chat_template: chat_template::PromptTemplate,
143}
144
145pub(crate) struct KimiModel {
146 engine: KimiEngine,
147 tokenizer: KimiTokenizer,
148 stop_tokens: StopTokens,
149 chat_template: chat_template::PromptTemplate,
150}
151
152pub(crate) struct MlaModel {
153 engine: MlaEngine,
154 tokenizer: ServerTokenizer,
155 stop_tokens: StopTokens,
156 bos_id: Option<usize>,
157 name: String,
158 chat_template: chat_template::PromptTemplate,
159}
160
161pub(crate) struct Gemma4Model {
162 engine: Gemma4Engine,
163 tokenizer: ServerTokenizer,
164 stop_tokens: StopTokens,
165 bos_id: Option<usize>,
166 name: String,
167 chat_template: chat_template::PromptTemplate,
168}
169
170pub(crate) struct Glm52Model {
171 engine: frink_models::Glm52Engine,
172 tokenizer: ServerTokenizer,
173 stop_tokens: StopTokens,
174 bos_id: Option<usize>,
175 name: String,
176 chat_template: chat_template::PromptTemplate,
177}
178
179impl Model {
180 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
181 match self {
182 Model::Gguf(m) => m.chat_template.clone(),
183 Model::Kimi(m) => m.chat_template.clone(),
184 Model::Mla(m) => m.chat_template.clone(),
185 Model::Gemma4(m) => m.chat_template.clone(),
186 Model::Glm52(m) => m.chat_template.clone(),
187 }
188 }
189
190 fn is_synthetic(&self) -> bool {
195 match self {
196 Model::Gguf(m) => m.is_synthetic,
197 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
198 }
199 }
200
201 fn tokenizer_kind(&self) -> &'static str {
202 match self {
203 Model::Gguf(m) => m.tokenizer.kind(),
204 Model::Kimi(_) => "kimi-tiktoken-bpe",
205 Model::Mla(m) => m.tokenizer.kind(),
206 Model::Gemma4(m) => m.tokenizer.kind(),
207 Model::Glm52(m) => m.tokenizer.kind(),
208 }
209 }
210
211 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
215 match self {
216 Model::Gguf(m) => m.decoder.expert_store_stats(),
217 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
218 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
219 }
220 }
221
222 pub(crate) fn name(&self) -> &str {
223 match self {
224 Model::Gguf(m) => m.decoder.config.name,
225 Model::Kimi(_) => "kimi-k3",
226 Model::Mla(m) => m.name.as_str(),
227 Model::Gemma4(m) => m.name.as_str(),
228 Model::Glm52(m) => m.name.as_str(),
229 }
230 }
231
232 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
259 match self {
260 Model::Gguf(m) => m.tokenizer.encode(text, specials),
261 Model::Kimi(m) => m
262 .tokenizer
263 .encode(text, specials)
264 .into_iter()
265 .map(|id| id as usize)
266 .collect(),
267 Model::Mla(m) => m.tokenizer.encode(text, specials),
268 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
269 Model::Glm52(m) => m.tokenizer.encode(text, specials),
270 }
271 }
272
273 pub(crate) fn bos_id(&self) -> Option<usize> {
282 match self {
283 Model::Gguf(m) => m.bos_id,
284 Model::Kimi(_) => None,
285 Model::Mla(m) => m.bos_id,
286 Model::Gemma4(m) => m.bos_id,
287 Model::Glm52(m) => m.bos_id,
288 }
289 }
290
291 pub(crate) fn decode(&self, ids: &[usize]) -> String {
292 match self {
293 Model::Gguf(m) => m.tokenizer.decode(ids),
294 Model::Kimi(m) => {
295 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
296 m.tokenizer.decode(&ids32)
297 }
298 Model::Mla(m) => m.tokenizer.decode(ids),
299 Model::Gemma4(m) => m.tokenizer.decode(ids),
300 Model::Glm52(m) => m.tokenizer.decode(ids),
301 }
302 }
303
304 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
307 match self {
308 Model::Gguf(m) => {
309 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
310 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
311 }
312 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
313 }
314 }
315
316 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
322 match self {
323 Model::Gguf(m) => Some(&m.decoder),
324 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
325 }
326 }
327
328 pub(crate) fn vocab_size(&self) -> Option<usize> {
329 match self {
330 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
331 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
332 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
333 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
334 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
335 }
336 }
337
338 fn has_real_vocabulary(&self) -> bool {
345 match self {
346 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
347 Model::Kimi(_) => true,
348 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
349 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
350 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
351 }
352 }
353}
354
355impl frink_models::dry::DryVocab for Model {
361 fn n_tokens(&self) -> usize {
362 self.vocab_size().unwrap_or(0)
363 }
364
365 fn detokenize(&self, token: usize) -> String {
366 self.decode(&[token])
367 }
368
369 fn tokenize(&self, text: &str) -> Vec<usize> {
370 self.encode(text, SpecialTokens::AsText)
371 }
372}
373
374pub(crate) struct AppState {
375 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
389 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
408 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
412 pub(crate) tasks: Arc<tasks::TaskRegistry>,
414 pub(crate) cancels: Arc<cancel::CancelRegistry>,
418 pub(crate) stats: stats::Stats,
421 pub(crate) streams: resume::StreamRegistry,
424 pub(crate) model_dir: Option<PathBuf>,
426 response_cache: Mutex<ResponseCache>,
430 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
442 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
455 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
465 sessions: session::SessionStore,
471 requests_total: std::sync::atomic::AtomicU64,
472 request_errors_total: std::sync::atomic::AtomicU64,
473 started_at: std::time::Instant,
474 last_request_ms: std::sync::atomic::AtomicU64,
479 detection: Arc<health::Detection>,
481 mcp: Option<mcp::LoadedMcpConfig>,
483 pub(crate) continuous_batching_enabled: bool,
487 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
491 loading_model: Mutex<Option<String>>,
495 last_load_error: Mutex<Option<(String, String)>>,
499 pub(crate) serving: Mutex<crate::stats::ServingStats>,
504 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
509 pub(crate) footprint:
515 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
516 pub(crate) started_unix: u64,
526}
527
528pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
534
535impl AppState {
536 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
540 self.active
541 .read()
542 .unwrap_or_else(|p| p.into_inner())
543 .clone()
544 }
545
546 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
552 self.active().ok_or_else(|| {
553 (
554 StatusCode::SERVICE_UNAVAILABLE,
555 Json(serde_json::json!({"error": {
556 "message": "no model is loaded; POST /admin/models/load with an id from \
557 GET /admin/models",
558 "type": "model_not_loaded"
559 }})),
560 )
561 })
562 }
563
564 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
572 Ok(Arc::clone(self.require_active()?.generative()?))
573 }
574
575 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
583 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
584 std::mem::replace(&mut *guard, next)
585 }
586
587 fn mark_request_finished(&self) {
591 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
592 self.last_request_ms
593 .store(ms, std::sync::atomic::Ordering::Relaxed);
594 }
595
596 pub(crate) fn uptime(&self) -> Duration {
597 self.started_at.elapsed()
598 }
599
600 pub(crate) fn requests_total(&self) -> u64 {
601 self.requests_total
602 .load(std::sync::atomic::Ordering::Relaxed)
603 }
604
605 pub(crate) fn errors_total(&self) -> u64 {
606 self.request_errors_total
607 .load(std::sync::atomic::Ordering::Relaxed)
608 }
609
610 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
611 lock_cache(&self.response_cache).stats()
612 }
613
614 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
617 let last = self
618 .last_request_ms
619 .load(std::sync::atomic::Ordering::Relaxed);
620 (last > 0)
621 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
622 .map(|age| age.max(0.0))
623 }
624
625 pub(crate) fn loading_model_id(&self) -> Option<String> {
626 self.loading_model
627 .lock()
628 .unwrap_or_else(|p| p.into_inner())
629 .clone()
630 }
631
632 pub(crate) fn set_loading_model(&self, id: Option<String>) {
633 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
634 }
635
636 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
637 self.last_load_error
638 .lock()
639 .unwrap_or_else(|p| p.into_inner())
640 .clone()
641 }
642
643 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
644 *self
645 .last_load_error
646 .lock()
647 .unwrap_or_else(|p| p.into_inner()) = error;
648 }
649
650 pub(crate) fn active_model_name(&self) -> Option<String> {
659 self.active().map(|a| a.name().to_string())
660 }
661
662 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
672 self.embedding
673 .clone()
674 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
675 }
676
677 pub(crate) fn embedding_model_name(&self) -> Option<String> {
681 match self.embedding_model() {
682 Some(e) => Some(e.name().to_string()),
683 None => self.active_model_name(),
684 }
685 }
686
687 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
688 self.stats.record(stats::entry(record));
689 }
690}
691
692fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
697 cache
698 .lock()
699 .unwrap_or_else(|poisoned| poisoned.into_inner())
700}
701
702#[derive(Debug, Clone, Deserialize)]
703#[serde(untagged)]
704pub(crate) enum MessageContent {
705 Text(String),
706 Parts(Vec<ContentPart>),
707}
708
709#[derive(Debug, Clone, Deserialize)]
710struct ContentPart {
711 #[serde(rename = "type")]
712 kind: String,
713 #[serde(default)]
714 text: Option<String>,
715 #[serde(default)]
716 image_url: Option<serde_json::Value>,
717}
718
719impl MessageContent {
720 fn as_text(&self) -> String {
721 match self {
722 Self::Text(s) => s.clone(),
723 Self::Parts(parts) => parts
724 .iter()
725 .filter_map(|p| p.text.as_deref())
726 .collect::<Vec<_>>()
727 .join(""),
728 }
729 }
730
731 fn has_image(&self) -> bool {
732 match self {
733 Self::Text(_) => false,
734 Self::Parts(parts) => parts
735 .iter()
736 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
737 }
738 }
739}
740
741#[derive(Debug, Clone, Deserialize)]
742pub(crate) struct ChatMessage {
743 pub(crate) role: String,
744 #[serde(default)]
748 pub(crate) content: Option<MessageContent>,
749 #[serde(default)]
753 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
754 #[serde(default)]
759 #[allow(dead_code)]
760 pub(crate) tool_call_id: Option<String>,
761 #[serde(default, alias = "reasoning")]
777 pub(crate) reasoning_content: Option<String>,
778}
779
780impl ChatMessage {
781 fn rendered_content(&self) -> String {
788 let mut out = self
789 .content
790 .as_ref()
791 .map(MessageContent::as_text)
792 .unwrap_or_default();
793 if let Some(calls) = &self.tool_calls {
794 for call in calls {
795 out.push_str(&format!(
796 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
797 call.function.name, call.function.arguments
798 ));
799 }
800 }
801 out
802 }
803}
804
805#[derive(Debug, Clone, Deserialize)]
806pub(crate) struct ToolCallIn {
807 #[serde(default)]
808 #[allow(dead_code)]
809 id: String,
810 #[serde(rename = "type", default)]
811 #[allow(dead_code)]
812 kind: String,
813 function: ToolCallFunctionIn,
814}
815
816#[derive(Debug, Clone, Deserialize)]
817struct ToolCallFunctionIn {
818 name: String,
819 arguments: String,
824}
825
826#[derive(Debug, Clone, Deserialize)]
829struct ToolDef {
830 #[serde(rename = "type", default)]
831 #[allow(dead_code)]
832 kind: String,
833 function: ToolFunctionDef,
834}
835
836#[derive(Debug, Clone, Deserialize)]
837struct ToolFunctionDef {
838 name: String,
839 #[serde(default)]
840 description: Option<String>,
841 #[serde(default)]
842 parameters: Option<serde_json::Value>,
843}
844
845#[derive(Debug, Clone, Deserialize)]
856#[serde(untagged)]
857enum ToolChoice {
858 Mode(String),
859 Specific(serde_json::Value),
860}
861
862#[derive(Deserialize)]
865#[serde(untagged)]
866enum StopParam {
867 One(String),
868 Many(Vec<String>),
869}
870
871#[derive(Deserialize)]
872struct ChatCompletionRequest {
873 model: String,
874 messages: Vec<ChatMessage>,
875 #[serde(default = "default_max_tokens")]
876 max_tokens: usize,
877 #[serde(default)]
878 temperature: Option<f32>,
879 #[serde(default)]
880 top_p: Option<f32>,
881 #[serde(default)]
886 min_p: Option<f32>,
887 #[serde(default)]
888 top_k: Option<usize>,
889 #[serde(default)]
890 repetition_penalty: Option<f32>,
891 #[serde(flatten)]
895 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
896 #[serde(flatten)]
900 unimplemented: crate::unimplemented_fields::UnimplementedFields,
901 #[serde(default)]
902 seed: Option<u64>,
903 #[serde(default)]
904 stop: Option<StopParam>,
905 #[serde(default)]
906 stream: Option<bool>,
907 #[serde(default)]
920 stream_resumable: Option<bool>,
921 #[serde(default)]
932 ignore_eos: Option<bool>,
933 #[serde(default)]
934 tools: Vec<ToolDef>,
935 #[serde(default)]
936 tool_choice: Option<ToolChoice>,
937 #[serde(default)]
945 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
946 #[serde(default)]
954 reasoning_effort: Option<String>,
955 #[serde(default)]
959 thinking: Option<ThinkingSwitch>,
960 #[serde(default)]
965 session_id: Option<String>,
966 #[serde(default, deserialize_with = "continuation::deserialize")]
974 continue_final_message: continuation::ContinueFinalMessage,
975 #[serde(default, alias = "thinking_budget_tokens")]
983 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
984 #[serde(default)]
986 logprobs: Option<bool>,
987 #[serde(default)]
988 top_logprobs: Option<u32>,
989 #[serde(default)]
990 presence_penalty: Option<f32>,
991 #[serde(default)]
992 frequency_penalty: Option<f32>,
993 #[serde(default)]
994 response_format: Option<serde_json::Value>,
995 #[serde(default)]
1002 logit_bias: Option<serde_json::Value>,
1003 #[serde(default)]
1007 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1008 #[serde(default)]
1018 samplers: Option<serde_json::Value>,
1019 #[serde(default)]
1029 grammar: Option<String>,
1030}
1031
1032const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1044
1045#[derive(Debug, Clone, Deserialize)]
1047pub(crate) struct ThinkingSwitch {
1048 #[serde(rename = "type")]
1049 pub(crate) kind: String,
1050}
1051
1052const THINKING_KWARG_KEYS: [&str; 4] = [
1056 "enable_thinking",
1057 "thinking",
1058 "thinking_mode",
1059 "reasoning_effort",
1060];
1061
1062const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1066
1067fn default_max_tokens() -> usize {
1068 DEFAULT_CHAT_MAX_TOKENS
1069}
1070
1071impl ChatCompletionRequest {
1072 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1081 let mut knobs = SamplingKnobs {
1082 temperature: self.temperature,
1083 top_p: self.top_p,
1084 min_p: self.min_p,
1085 top_k: self.top_k,
1086 repetition_penalty: self.repetition_penalty,
1087 presence_penalty: self.presence_penalty,
1088 frequency_penalty: self.frequency_penalty,
1089 penalty_last_n: None,
1093 sampler_order: unsupported_sampling::parse_sampler_order(
1094 self.samplers.as_ref(),
1095 "/v1/chat/completions",
1096 )?,
1097 ..SamplingKnobs::default()
1098 };
1099 self.extra_samplers.apply(&mut knobs);
1100 Ok(knobs)
1101 }
1102
1103 fn sampling_params(
1104 &self,
1105 model: crate::sampling_knobs::SamplerModel<'_>,
1106 ) -> Result<SamplingParams, ApiError> {
1107 self.sampling_knobs()?.resolve(model).map_err(|e| {
1108 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1109 })
1110 }
1111
1112 fn stop_sequences(&self) -> Vec<String> {
1113 self.stop
1114 .as_ref()
1115 .map(|s| match s {
1116 StopParam::One(v) => vec![v.clone()],
1117 StopParam::Many(v) => v.clone(),
1118 })
1119 .unwrap_or_default()
1120 }
1121
1122 fn tools_active(&self) -> bool {
1127 !self.tools.is_empty()
1128 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1129 }
1130
1131 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1143 match &self.tool_choice {
1144 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1145 Some(ToolChoice::Specific(value)) => {
1146 let name = value
1150 .get("function")
1151 .and_then(|f| f.get("name"))
1152 .or_else(|| value.get("name"))
1153 .and_then(|n| n.as_str());
1154 match name {
1155 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1156 None => Err(invalid_request(
1157 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1158 function.name",
1159 "tool_choice",
1160 )),
1161 }
1162 }
1163 _ => Ok(None),
1164 }
1165 }
1166
1167 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1169 self.tools
1170 .iter()
1171 .map(|t| tool_grammar::ToolSpec {
1172 name: &t.function.name,
1173 parameters: t.function.parameters.as_ref(),
1174 })
1175 .collect()
1176 }
1177
1178 fn thinking_direction(&self) -> Option<bool> {
1221 if let Some(switch) = &self.thinking {
1222 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1223 "disabled" => Some(false),
1224 "enabled" => Some(true),
1225 _ => None,
1228 };
1229 }
1230 let effort = self.reasoning_effort.as_ref()?;
1231 DISABLE_EFFORTS
1232 .contains(&effort.trim().to_ascii_lowercase().as_str())
1233 .then_some(false)
1234 }
1235
1236 fn resolve_template_kwargs(
1237 &self,
1238 template: &chat_template::PromptTemplate,
1239 ) -> serde_json::Map<String, serde_json::Value> {
1240 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1241 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1245
1246 if !caller_steered {
1247 match self.thinking_direction() {
1248 Some(false) => {
1249 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1250 kwargs.insert(k, v);
1251 }
1252 return kwargs;
1255 }
1256 Some(true) => {
1257 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1258 kwargs.insert(k, v);
1259 }
1260 }
1261 None => {}
1262 }
1263 if let Some(effort) = &self.reasoning_effort {
1264 kwargs
1265 .entry("reasoning_effort".to_string())
1266 .or_insert_with(|| serde_json::json!(effort));
1267 }
1268 }
1269
1270 let offered: Vec<serde_json::Value> = if self.tools_active() {
1271 self.tools.iter().map(chat_template::tool_json).collect()
1272 } else {
1273 Vec::new()
1274 };
1275 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1276 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1277 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1278 kwargs.entry(k).or_insert(v);
1279 }
1280 }
1281 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1282 crate::policy::effort::EffortMapping::Mapped(to) => {
1283 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1284 }
1285 crate::policy::effort::EffortMapping::Dropped => {
1286 tracing::debug!(
1287 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1288 enough, so its own default applies"
1289 );
1290 }
1291 crate::policy::effort::EffortMapping::Unchanged => {}
1292 }
1293 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1294 kwargs
1295 }
1296
1297 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1300 if self.max_tokens == 0 {
1307 return Err(invalid_request(
1308 "max_tokens must be at least 1",
1309 "max_tokens",
1310 ));
1311 }
1312 if let Some(switch) = &self.thinking {
1317 let kind = switch.kind.trim().to_ascii_lowercase();
1318 if kind != "enabled" && kind != "disabled" {
1319 return Err(invalid_request(
1320 "thinking.type must be \"enabled\" or \"disabled\"",
1321 "thinking.type",
1322 ));
1323 }
1324 }
1325 for msg in &self.messages {
1326 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1327 return Err(unsupported_feature(
1328 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1329 ));
1330 }
1331 }
1332 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1333 return Err(unsupported_feature(
1334 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1335 ));
1336 }
1337 self.unimplemented.refuse("/v1/chat/completions")?;
1341 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1342 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1346 let stated_grammar =
1361 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1362 if let Some(forced) = self.forced_tool_choice()? {
1370 if self.tools.is_empty() {
1371 return Err(invalid_request(
1372 "tool_choice forces a tool call, but no tools were offered",
1373 "tool_choice",
1374 ));
1375 }
1376 if let tool_grammar::Forced::Named(name) = forced {
1377 if !self.tools.iter().any(|t| t.function.name == name) {
1378 return Err(invalid_request(
1379 &format!(
1380 "tool_choice names {name:?}, which is not one of the tools offered"
1381 ),
1382 "tool_choice",
1383 ));
1384 }
1385 }
1386 if stated_grammar.is_some() {
1396 return Err(invalid_request(
1397 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1398 are two different constraints on the same generation; send one",
1399 "tool_choice",
1400 ));
1401 }
1402 if self.json_object_mode() {
1403 return Err(invalid_request(
1404 "a forced tool_choice cannot be combined with response_format json_object: \
1405 the tool-call markers are not JSON",
1406 "tool_choice",
1407 ));
1408 }
1409 }
1410 Ok(())
1411 }
1412
1413 fn effective_stop_sequences(&self) -> Vec<String> {
1420 let mut stop = self.stop_sequences();
1421 if self.tools_active() {
1422 stop.push("</tool_call>".to_string());
1423 }
1424 stop
1425 }
1426
1427 fn json_object_mode(&self) -> bool {
1428 self.response_format
1429 .as_ref()
1430 .and_then(|v| v.get("type"))
1431 .and_then(|v| v.as_str())
1432 == Some("json_object")
1433 }
1434}
1435
1436#[derive(Serialize)]
1437struct ChatCompletionChoice {
1438 index: usize,
1439 message: ChatCompletionResponseMessage,
1440 finish_reason: &'static str,
1441}
1442
1443#[derive(Serialize)]
1444struct ChatCompletionResponseMessage {
1445 role: &'static str,
1446 #[serde(skip_serializing_if = "Option::is_none")]
1447 content: Option<String>,
1448 #[serde(skip_serializing_if = "Option::is_none")]
1452 reasoning_content: Option<String>,
1453 #[serde(skip_serializing_if = "Option::is_none")]
1454 tool_calls: Option<Vec<ToolCallOut>>,
1455}
1456
1457#[derive(Serialize, Clone)]
1458struct ToolCallOut {
1459 id: String,
1460 #[serde(rename = "type")]
1461 kind: &'static str,
1462 function: ToolCallFunctionOut,
1463}
1464
1465#[derive(Serialize, Clone)]
1473struct ToolCallDelta {
1474 index: usize,
1475 #[serde(skip_serializing_if = "Option::is_none")]
1476 id: Option<String>,
1477 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1478 kind: Option<&'static str>,
1479 function: ToolCallFunctionDelta,
1480}
1481
1482#[derive(Serialize, Clone, Default)]
1483struct ToolCallFunctionDelta {
1484 #[serde(skip_serializing_if = "Option::is_none")]
1485 name: Option<String>,
1486 #[serde(skip_serializing_if = "Option::is_none")]
1489 arguments: Option<String>,
1490}
1491
1492impl ToolCallDelta {
1493 fn whole(index: usize, name: String, arguments: String) -> Self {
1495 ToolCallDelta {
1496 index,
1497 id: Some(format!("call_{index}")),
1498 kind: Some("function"),
1499 function: ToolCallFunctionDelta {
1500 name: Some(name),
1501 arguments: Some(arguments),
1502 },
1503 }
1504 }
1505
1506 fn opening(index: usize, name: String) -> Self {
1508 ToolCallDelta {
1509 index,
1510 id: Some(format!("call_{index}")),
1511 kind: Some("function"),
1512 function: ToolCallFunctionDelta {
1513 name: Some(name),
1514 arguments: Some(String::new()),
1515 },
1516 }
1517 }
1518
1519 fn arguments(index: usize, fragment: String) -> Self {
1521 ToolCallDelta {
1522 index,
1523 id: None,
1524 kind: None,
1525 function: ToolCallFunctionDelta {
1526 name: None,
1527 arguments: Some(fragment),
1528 },
1529 }
1530 }
1531}
1532
1533#[derive(Serialize, Clone)]
1534struct ToolCallFunctionOut {
1535 name: String,
1536 arguments: String,
1540}
1541
1542#[derive(Serialize)]
1543struct ChatCompletionResponse {
1544 id: String,
1545 request_id: String,
1551 object: &'static str,
1552 model: String,
1553 choices: Vec<ChatCompletionChoice>,
1554 usage: generate::Usage,
1559 frink_cache: &'static str,
1570}
1571
1572#[derive(Serialize)]
1573struct ChatCompletionChunkDelta {
1574 #[serde(skip_serializing_if = "Option::is_none")]
1575 role: Option<&'static str>,
1576 #[serde(skip_serializing_if = "Option::is_none")]
1577 content: Option<String>,
1578 #[serde(skip_serializing_if = "Option::is_none")]
1580 reasoning_content: Option<String>,
1581 #[serde(skip_serializing_if = "Option::is_none")]
1582 tool_calls: Option<Vec<ToolCallDelta>>,
1583}
1584
1585#[derive(Serialize)]
1586struct ChatCompletionChunkChoice {
1587 index: usize,
1588 delta: ChatCompletionChunkDelta,
1589 finish_reason: Option<&'static str>,
1590}
1591
1592#[derive(Serialize)]
1593struct ChatCompletionChunk {
1594 id: String,
1595 #[serde(skip_serializing_if = "Option::is_none")]
1602 request_id: Option<String>,
1603 object: &'static str,
1604 model: String,
1605 choices: Vec<ChatCompletionChunkChoice>,
1606 #[serde(skip_serializing_if = "Option::is_none")]
1609 usage: Option<generate::Usage>,
1610}
1611
1612async fn health(State(state): State<Arc<AppState>>) -> Response {
1618 let snapshot = state.detection.snapshot();
1619 let mut capabilities = snapshot.capabilities;
1620 let active = state.active();
1621
1622 capabilities.push(match active.as_deref() {
1625 None => frink_api::Capability::unavailable(
1630 frink_api::health::capability::REAL_WEIGHTS,
1631 frink_api::health::reason::MODEL_NOT_LOADED,
1632 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1633 ),
1634 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1635 frink_api::health::capability::REAL_WEIGHTS,
1636 frink_api::health::reason::MODEL_NOT_LOADED,
1637 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1638 checkpoint. Output from this model is noise.",
1639 ),
1640 Some(active) if active.encoder().is_some() => {
1650 let endpoints = active
1651 .encoder()
1652 .map(|e| encoder_endpoints(e))
1653 .unwrap_or_default();
1654 let served_by = match endpoints.is_empty() {
1655 true => "no endpoint in this build serves it".to_string(),
1656 false => format!("served by {}", endpoints.join(" and ")),
1657 };
1658 frink_api::Capability::available(
1659 frink_api::health::capability::REAL_WEIGHTS,
1660 format!(
1661 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1662 {served_by}; generation endpoints refuse it.",
1663 active.name(),
1664 ),
1665 )
1666 }
1667 Some(active) => frink_api::Capability::available(
1668 frink_api::health::capability::REAL_WEIGHTS,
1669 format!("Serving the real checkpoint '{}'.", active.name()),
1670 ),
1671 });
1672 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1673 frink_api::Capability::available(
1674 frink_api::health::capability::CONTINUOUS_BATCHING,
1675 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1676 "On by default on Metal. Concurrent requests share one batched decode worker."
1677 } else {
1678 "Concurrent requests share one batched decode step."
1679 },
1680 )
1681 } else if state.metal_private_decode_gate.is_some() {
1682 frink_api::Capability::unavailable(
1683 frink_api::health::capability::CONTINUOUS_BATCHING,
1684 frink_api::health::reason::DISABLED,
1685 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1686 )
1687 } else {
1688 frink_api::Capability::unavailable(
1689 frink_api::health::capability::CONTINUOUS_BATCHING,
1690 frink_api::health::reason::DISABLED,
1691 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1692 )
1693 });
1694
1695 let last_request_ms = state
1696 .last_request_ms
1697 .load(std::sync::atomic::Ordering::Relaxed);
1698 let uptime = state.started_at.elapsed();
1699 let health_state = if active.is_none() {
1707 frink_api::HealthState::Unavailable
1708 } else {
1709 snapshot.state
1710 };
1711 let body = frink_api::HealthResponse {
1712 state: health_state,
1713 reason: match health_state {
1714 frink_api::HealthState::Ready => None,
1715 frink_api::HealthState::Unavailable => {
1716 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1717 }
1718 frink_api::HealthState::Detecting => {
1719 Some(frink_api::health::reason::DETECTING.to_string())
1720 }
1721 },
1722 detail: match health_state {
1723 frink_api::HealthState::Ready => None,
1724 frink_api::HealthState::Unavailable => Some(
1725 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1726 .to_string(),
1727 ),
1728 frink_api::HealthState::Detecting => {
1729 Some("Probing available compute backends.".to_string())
1730 }
1731 },
1732 model: active
1733 .as_deref()
1734 .map(|active| frink_api::health::ModelSummary {
1735 id: active.name().to_string(),
1736 tokenizer: active.tokenizer_kind().to_string(),
1737 synthetic_weights: active.is_synthetic(),
1738 }),
1739 capabilities,
1740 version: env!("CARGO_PKG_VERSION").to_string(),
1741 pid: std::process::id(),
1742 uptime_seconds: uptime.as_secs_f64(),
1743 server_time_unix_ms: std::time::SystemTime::now()
1744 .duration_since(std::time::UNIX_EPOCH)
1745 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1746 .unwrap_or(0),
1747 last_request_age_seconds: (last_request_ms > 0)
1748 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1749 .map(|age| age.max(0.0)),
1750 };
1751
1752 let status =
1753 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1754 (status, Json(body)).into_response()
1755}
1756
1757async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1758 let Some(active) = state.active() else {
1762 return Json(serde_json::json!({ "object": "list", "data": [] }));
1763 };
1764 let mut model_entry = serde_json::json!({
1765 "id": active.name(),
1766 "object": "model",
1767 "frink_synthetic_weights": active.is_synthetic(),
1768 "frink_tokenizer": active.tokenizer_kind(),
1769 });
1770 if let Some(encoder) = active.encoder() {
1776 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1777 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1778 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1779 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1780 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1781 }
1782 if let Some(model) = active.generative_opt() {
1790 let parser_configured = active.reasoning_format().is_some();
1791 let gears = model.chat_template().think_gears(parser_configured);
1792 if !gears.is_empty() {
1793 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1794 if let Some(default) = &gears.default {
1795 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1796 }
1797 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1800 }
1801 }
1802 if let Some(mcp) = &state.mcp {
1803 model_entry["frink_mcp"] = mcp.models_metadata();
1804 }
1805 Json(serde_json::json!({
1806 "object": "list",
1807 "data": [model_entry]
1808 }))
1809}
1810
1811async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1825 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1826 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1827 let active = state.active();
1828 Json(serde_json::json!({
1829 "model": active.as_ref().map(|a| a.name()),
1830 "state": state
1831 .maintenance
1832 .lock()
1833 .unwrap_or_else(|p| p.into_inner())
1834 .state()
1835 .as_str(),
1836 "uptime_s": state.uptime().as_secs(),
1837 "throughput": {
1838 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1839 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1840 },
1841 "requests": {
1842 "active": state.cancels.live_count(),
1843 "completed": state.stats.recorded_total(),
1844 "p95_ms": state.stats.p95_duration_ms(),
1845 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1846 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1847 "completion_tokens_total": state.stats.tokens_generated_total(),
1848 },
1849 "pools": cache_admin::pool_gauges(&state),
1855 "memory": cache_admin::footprint_json(&state),
1858 }))
1859}
1860
1861#[derive(Deserialize)]
1862struct RequestsQuery {
1863 #[serde(default)]
1864 since: u64,
1865 #[serde(default = "default_requests_limit")]
1866 limit: usize,
1867}
1868
1869fn default_requests_limit() -> usize {
1870 stats::MAX_PAGE
1871}
1872
1873async fn recent_requests(
1881 State(state): State<Arc<AppState>>,
1882 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1883) -> Json<serde_json::Value> {
1884 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1885 Json(serde_json::json!({
1886 "requests": rows,
1887 "next_cursor": cursor,
1888 "missed": missed,
1889 "total": state.stats.recorded_total(),
1890 }))
1891}
1892
1893#[derive(Serialize)]
1894struct CombinedCacheStats {
1895 response_cache: response_cache::CacheStats,
1896 prefix_cache: Option<frink_models::PrefixCacheStats>,
1898}
1899
1900async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1901 Json(CombinedCacheStats {
1902 response_cache: lock_cache(&state.response_cache).stats(),
1903 prefix_cache: state
1904 .prefix_cache
1905 .as_ref()
1906 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1907 })
1908}
1909
1910async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1915 use std::sync::atomic::Ordering;
1916
1917 let cache_stats = lock_cache(&state.response_cache).stats();
1918 let active = state.active();
1919 let requests_total = state.requests_total.load(Ordering::Relaxed);
1920 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1921 let uptime = state.started_at.elapsed().as_secs_f64();
1922
1923 let body = format!(
1924 "# HELP frink_requests_total Total chat completion requests received.\n\
1925 # TYPE frink_requests_total counter\n\
1926 frink_requests_total {requests_total}\n\
1927 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1928 # TYPE frink_request_errors_total counter\n\
1929 frink_request_errors_total {errors_total}\n\
1930 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1931 # TYPE frink_cache_hits_total counter\n\
1932 frink_cache_hits_total {}\n\
1933 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1934 # TYPE frink_cache_misses_total counter\n\
1935 frink_cache_misses_total {}\n\
1936 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1937 # TYPE frink_cache_entries gauge\n\
1938 frink_cache_entries {}\n\
1939 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1940 # TYPE frink_synthetic_weights gauge\n\
1941 frink_synthetic_weights {}\n\
1942 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1943 # TYPE frink_uptime_seconds gauge\n\
1944 frink_uptime_seconds {uptime}\n",
1945 cache_stats.hits,
1946 cache_stats.misses,
1947 cache_stats.entries,
1948 active
1952 .as_ref()
1953 .map(|a| a.is_synthetic() as u8)
1954 .unwrap_or(0),
1955 );
1956
1957 let body = match active
1961 .as_ref()
1962 .and_then(|a| a.expert_store_stats())
1963 {
1964 Some(es) => format!(
1965 "{body}\
1966 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
1967 # TYPE frink_expert_cache_hits_total counter\n\
1968 frink_expert_cache_hits_total {}\n\
1969 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1970 # TYPE frink_expert_cache_misses_total counter\n\
1971 frink_expert_cache_misses_total {}\n\
1972 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
1973 # TYPE frink_expert_cache_evictions_total counter\n\
1974 frink_expert_cache_evictions_total {}\n\
1975 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1976 # TYPE frink_expert_cache_pass_throughs_total counter\n\
1977 frink_expert_cache_pass_throughs_total {}\n\
1978 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1979 # TYPE frink_expert_cache_bytes_read_total counter\n\
1980 frink_expert_cache_bytes_read_total {}\n\
1981 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1982 # TYPE frink_expert_cache_resident_bytes gauge\n\
1983 frink_expert_cache_resident_bytes {}\n",
1984 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1985 ),
1986 None => body,
1987 };
1988
1989 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
1994 Some(batcher) => {
1995 let sched = batcher.stats();
1996 format!(
1997 "{body}\
1998 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
1999 # TYPE frink_prefill_chunks_total counter\n\
2000 frink_prefill_chunks_total {}\n\
2001 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2002 # TYPE frink_prefill_tokens_total counter\n\
2003 frink_prefill_tokens_total {}\n\
2004 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2005 # TYPE frink_decode_steps_total counter\n\
2006 frink_decode_steps_total {}\n\
2007 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2008 # TYPE frink_scheduler_queue_depth gauge\n\
2009 frink_scheduler_queue_depth {}\n\
2010 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2011 # TYPE frink_scheduler_queue_rejected_total counter\n\
2012 frink_scheduler_queue_rejected_total {}\n\
2013 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2014 # TYPE frink_kv_blocks_total gauge\n\
2015 frink_kv_blocks_total {}\n\
2016 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2017 # TYPE frink_kv_blocks_free gauge\n\
2018 frink_kv_blocks_free {}\n\
2019 # HELP frink_kv_block_size Token positions per KV block.\n\
2020 # TYPE frink_kv_block_size gauge\n\
2021 frink_kv_block_size {}\n\
2022 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2023 # TYPE frink_kv_rejected_too_large_total counter\n\
2024 frink_kv_rejected_too_large_total {}\n\
2025 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2026 # TYPE frink_kv_rejected_context_length_total counter\n\
2027 frink_kv_rejected_context_length_total {}\n\
2028 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2029 # TYPE frink_scheduler_aborted_total counter\n\
2030 frink_scheduler_aborted_total {}\n\
2031 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2032 # TYPE frink_scheduler_max_seqs gauge\n\
2033 frink_scheduler_max_seqs {}\n\
2034 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2035 # TYPE frink_scheduler_prefill_chunk gauge\n\
2036 frink_scheduler_prefill_chunk {}\n",
2037 sched.prefill_chunks,
2038 sched.prefill_tokens,
2039 sched.decode_steps,
2040 sched.queue_depth,
2041 sched.queue_rejected,
2042 sched.kv_blocks_total,
2043 sched.kv_blocks_free,
2044 sched.kv_block_size,
2045 sched.kv_rejected_too_large,
2046 sched.kv_rejected_context_length,
2047 sched.aborted,
2048 sched.max_seqs,
2049 sched.prefill_chunk,
2050 )
2051 }
2052 None => body,
2053 };
2054
2055 (
2056 [(
2057 axum::http::header::CONTENT_TYPE,
2058 "text/plain; version=0.0.4",
2059 )],
2060 body,
2061 )
2062 .into_response()
2063}
2064
2065pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2066
2067pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2073 (
2074 StatusCode::BAD_REQUEST,
2075 Json(serde_json::json!({"error": {
2076 "message": message,
2077 "type": "invalid_request_error",
2078 "param": param,
2079 "code": null,
2080 }})),
2081 )
2082}
2083
2084pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2085 (
2086 StatusCode::NOT_IMPLEMENTED,
2087 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2088 )
2089}
2090
2091pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2092 let status = match e {
2093 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2094 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2098 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2103 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2109 StatusCode::SERVICE_UNAVAILABLE
2110 }
2111 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2115 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2119 };
2120 tracing::warn!("decode error: {e}");
2121 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2122 if let generate::DecodeError::KvBudgetExceeded {
2128 binding,
2129 estimated_bytes,
2130 limit_bytes,
2131 positions,
2132 positions_limit,
2133 ..
2134 } = &e
2135 {
2136 body["error"]["type"] = serde_json::json!("invalid_request_error");
2137 body["error"]["code"] = serde_json::json!(binding);
2138 body["error"]["binding"] = serde_json::json!(binding);
2139 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2140 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2141 body["error"]["positions"] = serde_json::json!(positions);
2142 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2143 }
2144 if let Some(secs) = e.retry_after_secs() {
2148 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2149 }
2150 (status, Json(body))
2151}
2152
2153pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2154 tracing::error!("generation task panicked: {e}");
2155 (
2156 StatusCode::INTERNAL_SERVER_ERROR,
2157 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2158 )
2159}
2160
2161#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2171 model: &Model,
2172 prompt: &str,
2173 params: &GenerationParams,
2174 kv_pool: Option<&generate::KvPoolConfig>,
2175 paged_kv: Option<&generate::PagedKvConfig>,
2176 prefix_cache: Option<&Mutex<PrefixCache>>,
2177 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2178 ceiling: Option<&budget::ContextCeiling>,
2179 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2180 mut emit: impl FnMut(&str),
2181 ) -> Result<(Vec<(FinishReason, String)>, generate::Usage), generate::DecodeError> {
2187 let synthetic = model.is_synthetic();
2188 let _lora_lease = lora::lease(model, params.lora.as_deref());
2193 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2194 let params = &{
2200 let mut resolved = params.clone();
2201 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2202 model.encode(text, SpecialTokens::Parse)
2203 });
2204 resolved.reasoning_budget = resolved
2210 .reasoning_budget
2211 .armed(resolved.reasoning, prompt, |text| {
2212 model.encode(text, SpecialTokens::Parse)
2213 })
2214 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2215 resolved
2216 };
2217 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2218 let _metal_private_guard =
2219 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2220 let (finishes, usage) = match model {
2221 Model::Gguf(m) => {
2222 if let Some(batcher) = continuous_batcher {
2223 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2224 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2225 let (finish, _generated_ids, text, usage) = if synthetic {
2226 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2227 } else {
2228 batcher.generate_streaming(
2229 tokens,
2230 params.clone(),
2231 m.stop_tokens.clone(),
2232 Some(|chunk: &str| {
2233 if !chunk.is_empty() {
2234 chunks[0].push(chunk.to_string());
2235 emit(chunk);
2236 }
2237 }),
2238 )?
2239 };
2240 if !text.is_empty() && chunks[0].is_empty() {
2241 chunks[0].push(text);
2242 }
2243 (vec![finish], usage)
2247 } else {
2248 generate::generate(
2249 &m.decoder,
2250 m.tokenizer.as_ref(),
2251 &m.stop_tokens,
2252 m.bos_id,
2253 prompt,
2254 params,
2255 kv_pool,
2256 paged_kv,
2257 prefix_cache,
2258 ceiling,
2259 |choice, chunk| {
2260 chunks[choice].push(chunk.to_string());
2261 if !synthetic && choice == 0 {
2265 emit(chunk);
2266 }
2267 },
2268 )?
2269 }
2270 }
2271 Model::Kimi(m) => generate::generate_engine(
2272 &m.engine,
2273 &m.tokenizer,
2274 &m.stop_tokens,
2275 None,
2276 prompt,
2277 params,
2278 |chunk| {
2279 chunks[0].push(chunk.to_string());
2280 if !synthetic {
2281 emit(chunk);
2282 }
2283 },
2284 )?,
2285 Model::Mla(m) => generate::generate_engine(
2286 &m.engine,
2287 &m.tokenizer,
2288 &m.stop_tokens,
2289 m.bos_id,
2290 prompt,
2291 params,
2292 |chunk| {
2293 chunks[0].push(chunk.to_string());
2294 if !synthetic {
2295 emit(chunk);
2296 }
2297 },
2298 )?,
2299 Model::Gemma4(m) => generate::generate_engine(
2300 &m.engine,
2301 &m.tokenizer,
2302 &m.stop_tokens,
2303 m.bos_id,
2304 prompt,
2305 params,
2306 |chunk| {
2307 chunks[0].push(chunk.to_string());
2308 if !synthetic {
2309 emit(chunk);
2310 }
2311 },
2312 )?,
2313 Model::Glm52(m) => generate::generate_engine(
2314 &m.engine,
2315 &m.tokenizer,
2316 &m.stop_tokens,
2317 m.bos_id,
2318 prompt,
2319 params,
2320 |chunk| {
2321 chunks[0].push(chunk.to_string());
2322 if !synthetic {
2323 emit(chunk);
2324 }
2325 },
2326 )?,
2327 };
2328
2329 let mut full = chunks[0].concat();
2330 if synthetic {
2331 full = format!(
2332 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2333 to serve a real model. Decoded ids -> {full:?}]"
2334 );
2335 emit(&full);
2336 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2337 emit(&full);
2338 }
2339
2340 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2346 let mut out: Vec<(FinishReason, String)> = finishes
2347 .into_iter()
2348 .zip(chunks.into_iter().map(|c| c.concat()))
2349 .collect();
2350 if let Some(first) = out.first_mut() {
2351 first.1 = full;
2352 }
2353 Ok((out, usage))
2354}
2355
2356#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2361 model: &Model,
2362 prompt: &str,
2363 params: &GenerationParams,
2364 kv_pool: Option<&generate::KvPoolConfig>,
2365 paged_kv: Option<&generate::PagedKvConfig>,
2366 prefix_cache: Option<&Mutex<PrefixCache>>,
2367 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2368 ceiling: Option<&budget::ContextCeiling>,
2369 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2370 ) -> Result<(Vec<(FinishReason, String)>, generate::Usage), generate::DecodeError> {
2373 run_generation_emit(
2374 model,
2375 prompt,
2376 params,
2377 kv_pool,
2378 paged_kv,
2379 prefix_cache,
2380 continuous_batcher,
2381 ceiling,
2382 metal_private_decode_gate,
2383 |_| {},
2384 )
2385}
2386
2387pub(crate) fn prompt_from_messages(
2398 messages: &[ChatMessage],
2399 template: &chat_template::PromptTemplate,
2400 tools: &[ToolDef],
2401 extra: serde_json::Map<String, serde_json::Value>,
2402) -> Result<String, ApiError> {
2403 let rendered = if tools.is_empty() || template.handles_tools() {
2404 template.render(messages, tools, extra)
2405 } else {
2406 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2407 with_preamble.push(ChatMessage {
2408 role: "system".to_string(),
2409 content: Some(MessageContent::Text(tool_preamble(tools))),
2410 tool_calls: None,
2411 tool_call_id: None,
2412 reasoning_content: None,
2413 });
2414 with_preamble.extend_from_slice(messages);
2415 template.render(&with_preamble, &[], extra)
2416 };
2417 rendered.map_err(template_error_response)
2418}
2419
2420fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2425 (
2426 StatusCode::BAD_REQUEST,
2427 Json(serde_json::json!({
2428 "error": {
2429 "message": format!("chat template failed to render: {err}"),
2430 "type": "invalid_request_error",
2431 "param": "messages",
2432 "code": null,
2433 }
2434 })),
2435 )
2436}
2437
2438fn tool_preamble(tools: &[ToolDef]) -> String {
2451 let mut out = String::from(
2452 "You can call tools to help answer the user. To call a tool, respond with \
2453 EXACTLY one line in this format and nothing else:\n\
2454 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2455 object matching that tool's parameters>}}</tool_call>\n\n\
2456 Available tools:\n",
2457 );
2458 for t in tools {
2459 out.push_str(&format!(
2460 "- {}: {}\n parameters (JSON schema): {}\n",
2461 t.function.name,
2462 t.function.description.as_deref().unwrap_or(""),
2463 t.function
2464 .parameters
2465 .as_ref()
2466 .map(|v| v.to_string())
2467 .unwrap_or_else(|| "{}".to_string()),
2468 ));
2469 }
2470 out
2471}
2472
2473fn tool_call_deltas(
2483 events: Vec<crate::policy::parser::ToolCallEvent>,
2484 opened: &std::cell::Cell<usize>,
2485) -> (String, Vec<ToolCallDelta>) {
2486 let mut text = String::new();
2487 let mut deltas = Vec::new();
2488 for event in events {
2489 match event {
2490 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2491 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2492 opened.set(opened.get().max(index + 1));
2493 deltas.push(ToolCallDelta::opening(index, name));
2494 }
2495 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2496 if !fragment.is_empty() {
2497 deltas.push(ToolCallDelta::arguments(index, fragment));
2498 }
2499 }
2500 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2501 }
2502 }
2503 (text, deltas)
2504}
2505
2506fn build_response_message(
2517 text: String,
2518 tools: &[ToolDef],
2519 posture: output::OutputPosture,
2520 base_finish: &'static str,
2521) -> (ChatCompletionResponseMessage, &'static str) {
2522 let parsed = output::parse_output(&text, tools, posture);
2523 let calls: Vec<ToolCallOut> = parsed
2524 .calls
2525 .into_iter()
2526 .enumerate()
2527 .map(|(index, call)| ToolCallOut {
2528 id: format!("call_{index}"),
2529 kind: "function",
2530 function: ToolCallFunctionOut {
2531 name: call.name,
2532 arguments: call.arguments,
2533 },
2534 })
2535 .collect();
2536 if !calls.is_empty() {
2537 return (
2538 ChatCompletionResponseMessage {
2539 role: "assistant",
2540 content: None,
2541 reasoning_content: parsed.reasoning,
2542 tool_calls: Some(calls),
2543 },
2544 "tool_calls",
2545 );
2546 }
2547 (
2548 ChatCompletionResponseMessage {
2549 role: "assistant",
2550 content: Some(parsed.content),
2551 reasoning_content: parsed.reasoning,
2552 tool_calls: None,
2553 },
2554 base_finish,
2555 )
2556}
2557
2558fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2563 let mut history = match &req.session_id {
2564 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2565 None => req.messages.clone(),
2566 };
2567 if req.json_object_mode() {
2568 inject_json_object_system_hint(&mut history);
2569 }
2570 history
2571}
2572
2573fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2574 const HINT: &str =
2575 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2576 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2577 match &mut sys.content {
2578 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2579 s.push_str("\n\n");
2580 s.push_str(HINT);
2581 }
2582 None => {
2583 sys.content = Some(MessageContent::Text(HINT.to_string()));
2584 }
2585 _ => {}
2586 }
2587 } else {
2588 messages.insert(
2589 0,
2590 ChatMessage {
2591 role: "system".to_string(),
2592 content: Some(MessageContent::Text(HINT.to_string())),
2593 tool_calls: None,
2594 tool_call_id: None,
2595 reasoning_content: None,
2596 },
2597 );
2598 }
2599}
2600
2601async fn chat_completions(
2602 State(state): State<Arc<AppState>>,
2603 headers: axum::http::HeaderMap,
2604 Json(req): Json<ChatCompletionRequest>,
2605) -> Response {
2606 let attribution = attribution::Attribution::from_headers(&headers);
2607 state
2608 .requests_total
2609 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2610 let started = std::time::Instant::now();
2611
2612 let request_id = frink_api::next_request_id();
2616 let stream = req.stream.unwrap_or(false);
2617
2618 let refusal = cache_admin::check_admission(&state)
2624 .err()
2625 .or_else(|| req.validate_supported_fields().err());
2626 if let Some(err) = refusal {
2627 state
2628 .request_errors_total
2629 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2630 let response = err.into_response();
2631 state.record_request(stats::Record {
2632 request_id: &request_id,
2633 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2634 model: state.active_model_name(),
2635 status: response.status().as_u16(),
2636 stream,
2637 duration_ms: started.elapsed().as_millis() as u64,
2638 usage: None,
2639 attribution: &attribution,
2640 });
2641 return response;
2642 }
2643
2644 let response = if stream {
2645 chat_completions_stream(
2646 Arc::clone(&state),
2647 req,
2648 request_id.clone(),
2649 started,
2650 attribution.clone(),
2651 )
2652 .await
2653 .into_response()
2654 } else {
2655 chat_completions_full(
2656 Arc::clone(&state),
2657 req,
2658 request_id.clone(),
2659 started,
2660 attribution.clone(),
2661 )
2662 .await
2663 .into_response()
2664 };
2665
2666 if response.status().is_client_error() || response.status().is_server_error() {
2667 state
2668 .request_errors_total
2669 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2670 state.record_request(stats::Record {
2674 request_id: &request_id,
2675 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2676 model: state.active_model_name(),
2679 status: response.status().as_u16(),
2680 stream,
2681 duration_ms: started.elapsed().as_millis() as u64,
2682 usage: None,
2683 attribution: &attribution,
2684 });
2685 }
2686 state.mark_request_finished();
2687
2688 response
2689}
2690
2691async fn chat_completions_full(
2692 state: Arc<AppState>,
2693 req: ChatCompletionRequest,
2694 request_id: String,
2695 started: std::time::Instant,
2696 attribution: attribution::Attribution,
2697) -> Result<Json<ChatCompletionResponse>, ApiError> {
2698 let tools_active = req.tools_active();
2699 let active = state.require_active()?;
2703 let history = resolve_history(&state, &req);
2704 let template = active.generative()?.chat_template();
2705 let kwargs = req.resolve_template_kwargs(&template);
2706 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2707 let mut params =
2715 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2716 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2717 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2718
2719 let (completion, cache_status) = if let Some(cached) = key
2720 .as_ref()
2721 .and_then(|key| lock_cache(&state.response_cache).get(key))
2722 {
2723 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2724 (cached, "hit")
2725 } else {
2726 let (choices, usage) = decode_task::buffered(
2727 decode_task::DecodeHandles::take(&state, &active)?,
2728 prompt.clone(),
2729 params,
2730 )
2731 .await?;
2732
2733 let (finish, content) = choices
2736 .into_iter()
2737 .next()
2738 .expect("a generation produces at least one choice");
2739 let completion = response_cache::CachedCompletion {
2740 content,
2741 finish,
2742 usage,
2743 };
2744 let cache_status = match key {
2750 Some(key) => match completion.clone().cacheable() {
2754 Some(cacheable) => {
2755 tracing::debug!("cache miss for key {}", key.digest());
2756 lock_cache(&state.response_cache).put(key, cacheable);
2757 "miss"
2758 }
2759 None => "skip",
2760 },
2761 None => "skip",
2762 };
2763 (completion, cache_status)
2764 };
2765 let content = completion.content;
2766
2767 if req.json_object_mode() {
2768 json_mode::validate_json_object_output(&content)?;
2769 }
2770
2771 if let Some(id) = &req.session_id {
2775 state.sessions.store_reply(
2776 id,
2777 ChatMessage {
2778 role: "assistant".to_string(),
2779 content: Some(MessageContent::Text(content.clone())),
2780 tool_calls: None,
2781 tool_call_id: None,
2782 reasoning_content: None,
2783 },
2784 );
2785 }
2786
2787 let (message, finish_reason) = build_response_message(
2788 content,
2789 if tools_active { &req.tools } else { &[] },
2790 output::OutputPosture::resolve_full(
2791 active.reasoning_format(),
2792 active.tool_call_format(),
2793 &prompt,
2794 ),
2795 completion.finish.as_str(),
2796 );
2797
2798 state.record_request(stats::Record {
2799 request_id: &request_id,
2800 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2801 model: Some(active.name().to_string()),
2804 status: 200,
2805 stream: false,
2806 duration_ms: started.elapsed().as_millis() as u64,
2807 usage: Some(&completion.usage),
2808 attribution: &attribution,
2809 });
2810
2811 Ok(Json(ChatCompletionResponse {
2812 id: request_id.clone(),
2813 request_id,
2814 object: "chat.completion",
2815 model: req.model,
2816 choices: vec![ChatCompletionChoice {
2817 index: 0,
2818 message,
2819 finish_reason,
2820 }],
2821 usage: completion.usage,
2822 frink_cache: cache_status,
2823 }))
2824}
2825
2826async fn chat_completions_stream(
2827 state: Arc<AppState>,
2828 req: ChatCompletionRequest,
2829 request_id: String,
2830 started: std::time::Instant,
2831 attribution: attribution::Attribution,
2832) -> Result<Response, ApiError> {
2833 let tools_active = req.tools_active();
2835 let active = state.require_active()?;
2839 let history = resolve_history(&state, &req);
2840 let template = active.generative()?.chat_template();
2841 let kwargs = req.resolve_template_kwargs(&template);
2842 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2843 let model_name = req.model.clone();
2844 let session_id = req.session_id.clone();
2845 let sessions = state.sessions.clone();
2846
2847 let model = Arc::clone(active.generative()?);
2848 let kv_pool = state.kv_pool.clone();
2849 let paged_kv = state.paged_kv.clone();
2850 let prefix_cache = state.prefix_cache.clone();
2851 let batcher = active.batcher.clone();
2852 let ceiling = active.ceiling.clone();
2853 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2854 let mut params =
2855 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2856 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2857 let stats_state = Arc::clone(&state);
2858 let served_model = active.name().to_string();
2861 let posture = output::OutputPosture::resolve_full(
2865 active.reasoning_format(),
2866 active.tool_call_format(),
2867 &prompt,
2868 );
2869 let offered_tools: Vec<ToolDef> = if tools_active {
2872 req.tools.clone()
2873 } else {
2874 Vec::new()
2875 };
2876
2877 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2882 params.cancel = Some(cancel_token.clone());
2883
2884 let overlap = true;
2893
2894 let slot = req
2898 .stream_resumable
2899 .unwrap_or(false)
2900 .then(|| state.streams.register(&request_id));
2901 let emitter = resume::Emitter::new(slot);
2902
2903 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2904 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2908 id: request_id.clone(),
2909 request_id: None,
2910 object: "chat.completion.chunk",
2911 model: model_name.clone(),
2912 choices: vec![ChatCompletionChunkChoice {
2913 index: 0,
2914 delta: ChatCompletionChunkDelta {
2915 role: None,
2916 content: None,
2917 reasoning_content: None,
2918 tool_calls: None,
2919 },
2920 finish_reason: None,
2921 }],
2922 usage: None,
2923 });
2924
2925 tokio::task::spawn_blocking(move || {
2926 let _cancel_guard = cancel_guard;
2929 let tx_chunks = tx.clone();
2930 let orphan_timeout = sse::orphan_timeout_from_env();
2935 let mut first = true;
2936 let head_request_id = request_id.clone();
2937 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
2947 Rc::new(RefCell::new(posture.reasoning_parser()));
2948 let emit_reasoning = Rc::clone(&stream_reasoning);
2949 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
2953 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
2954 );
2955 let emit_tools = Rc::clone(&stream_tools);
2956 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
2960 let emit_streamed_calls = Rc::clone(&streamed_calls);
2961 let result = run_generation_emit(
2962 &model,
2963 &prompt,
2964 ¶ms,
2965 kv_pool.as_ref(),
2966 paged_kv.as_ref(),
2967 prefix_cache.as_deref(),
2968 batcher.as_ref(),
2969 ceiling.as_deref(),
2970 metal_private_decode_gate.as_deref(),
2971 |chunk| {
2972 if !overlap || chunk.is_empty() {
2973 return;
2974 }
2975 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
2976 Some(parser) => {
2977 let delta = parser.push(chunk);
2978 (delta.reasoning, delta.content)
2979 }
2980 None => (String::new(), chunk.to_string()),
2981 };
2982 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
2986 Some(parser) => {
2987 let (text, calls) =
2988 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
2989 (text, calls)
2990 }
2991 None => (content, Vec::new()),
2992 };
2993 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
2996 return;
2997 }
2998 let role = if first { Some("assistant") } else { None };
2999 let request_id = first.then(|| head_request_id.clone());
3000 first = false;
3001 let payload = ChatCompletionChunk {
3002 id: head_request_id.clone(),
3003 request_id,
3004 object: "chat.completion.chunk",
3005 model: model_name.clone(),
3006 choices: vec![ChatCompletionChunkChoice {
3007 index: 0,
3008 delta: ChatCompletionChunkDelta {
3009 role,
3010 content: (!content.is_empty()).then_some(content),
3011 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3012 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3013 },
3014 finish_reason: None,
3015 }],
3016 usage: None,
3017 };
3018 if let Err(why) =
3027 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3028 {
3029 if why == sse::SendFailure::Orphaned {
3030 tracing::warn!(
3031 "SSE stream {head_request_id} accepted nothing for the orphan \
3032 deadline; treating it as abandoned"
3033 );
3034 }
3035 if !emitter.is_resumable() {
3047 cancel_token.cancel();
3048 }
3049 }
3050 },
3051 );
3052
3053 let mut pending_request_id = first.then(|| request_id.clone());
3059
3060 match result {
3061 Ok((choices, usage)) => {
3064 let (finish, full_text) = choices
3065 .into_iter()
3066 .next()
3067 .expect("a generation produces at least one choice");
3068 if let Some(id) = &session_id {
3069 sessions.store_reply(
3070 id,
3071 ChatMessage {
3072 role: "assistant".to_string(),
3073 content: Some(MessageContent::Text(full_text.clone())),
3074 tool_calls: None,
3075 tool_call_id: None,
3076 reasoning_content: None,
3077 },
3078 );
3079 }
3080 let mut streamed_finish: Option<&'static str> = None;
3086 if overlap {
3087 let tail = stream_reasoning
3088 .borrow_mut()
3089 .as_mut()
3090 .map(|parser| parser.flush())
3091 .unwrap_or_default();
3092 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3093 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3094 let mut events = parser.push(&content);
3095 events.extend(parser.finish());
3096 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3097 content = text;
3098 tool_calls = calls;
3099 }
3100 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3101 let payload = ChatCompletionChunk {
3102 id: request_id.clone(),
3103 request_id: pending_request_id.take(),
3104 object: "chat.completion.chunk",
3105 model: model_name.clone(),
3106 choices: vec![ChatCompletionChunkChoice {
3107 index: 0,
3108 delta: ChatCompletionChunkDelta {
3109 role: None,
3110 content: (!content.is_empty()).then_some(content),
3111 reasoning_content: (!tail.reasoning.is_empty())
3112 .then_some(tail.reasoning),
3113 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3114 },
3115 finish_reason: None,
3116 }],
3117 usage: None,
3118 };
3119 let _ =
3120 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3121 }
3122 if streamed_calls.get() > 0 {
3123 streamed_finish = Some("tool_calls");
3124 }
3125 } else {
3126 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3129 let tool_calls: Vec<ToolCallDelta> = parsed
3130 .calls
3131 .iter()
3132 .enumerate()
3133 .map(|(index, call)| {
3134 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3135 })
3136 .collect();
3137 if !tool_calls.is_empty() {
3138 streamed_finish = Some("tool_calls");
3139 }
3140 if !tool_calls.is_empty()
3141 || !parsed.content.is_empty()
3142 || parsed.reasoning.is_some()
3143 {
3144 let payload = ChatCompletionChunk {
3145 id: request_id.clone(),
3146 request_id: pending_request_id.take(),
3147 object: "chat.completion.chunk",
3148 model: model_name.clone(),
3149 choices: vec![ChatCompletionChunkChoice {
3150 index: 0,
3151 delta: ChatCompletionChunkDelta {
3152 role: Some("assistant"),
3153 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3154 .then(|| parsed.content.clone()),
3155 reasoning_content: parsed.reasoning.clone(),
3156 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3157 },
3158 finish_reason: None,
3159 }],
3160 usage: None,
3161 };
3162 let _ =
3163 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3164 }
3165 }
3166 let final_finish_reason = match streamed_finish {
3170 Some(reason) if finish.as_str() != "length" => reason,
3171 _ => finish.as_str(),
3172 };
3173 let final_payload = ChatCompletionChunk {
3174 id: request_id.clone(),
3175 request_id: pending_request_id.take(),
3176 object: "chat.completion.chunk",
3177 model: model_name,
3178 choices: vec![ChatCompletionChunkChoice {
3179 index: 0,
3180 delta: ChatCompletionChunkDelta {
3181 role: None,
3182 content: None,
3183 reasoning_content: None,
3184 tool_calls: None,
3185 },
3186 finish_reason: Some(final_finish_reason),
3187 }],
3188 usage: Some(usage.clone()),
3189 };
3190 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3191 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3192 stats_state.record_request(stats::Record {
3197 request_id: &request_id,
3198 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3199 model: Some(served_model.clone()),
3200 status: 200,
3201 stream: true,
3202 duration_ms: started.elapsed().as_millis() as u64,
3203 usage: Some(&usage),
3204 attribution: &attribution,
3205 });
3206 }
3207 Err(e) => {
3208 tracing::warn!("decode error on streamed request {request_id}: {e}");
3209 stats_state.record_request(stats::Record {
3215 request_id: &request_id,
3216 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3217 model: Some(served_model.clone()),
3218 status: 500,
3219 stream: true,
3220 duration_ms: started.elapsed().as_millis() as u64,
3221 usage: None,
3222 attribution: &attribution,
3223 });
3224 let payload = ChatCompletionChunk {
3225 id: request_id.clone(),
3226 request_id: pending_request_id.take(),
3227 object: "chat.completion.chunk",
3228 model: model_name,
3229 choices: vec![ChatCompletionChunkChoice {
3230 index: 0,
3231 delta: ChatCompletionChunkDelta {
3232 role: Some("assistant"),
3233 content: Some(format!("[error: {e}]")),
3234 reasoning_content: None,
3235 tool_calls: None,
3236 },
3237 finish_reason: Some("stop"),
3238 }],
3239 usage: None,
3240 };
3241 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3242 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3243 }
3244 }
3245 drop(emitter);
3249 });
3250
3251 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3252 Ok((
3275 [(
3276 axum::http::HeaderName::from_static("x-accel-buffering"),
3277 axum::http::HeaderValue::from_static("no"),
3278 )],
3279 Sse::new(stream),
3280 )
3281 .into_response())
3282}
3283
3284fn protected_routes() -> Router<Arc<AppState>> {
3315 use frink_api::routes;
3316
3317 Router::new()
3318 .route(routes::V1_MODELS, get(list_models))
3319 .route(routes::V1_RESPONSES, post(responses::responses))
3323 .route(
3324 &axum_path(routes::V1_RESPONSE),
3325 get(responses::responses_get),
3326 )
3327 .route(
3328 &axum_path(routes::V1_RESPONSE_CANCEL),
3329 post(responses::responses_cancel),
3330 )
3331 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3332 .route(routes::V1_STATS, get(serving_stats))
3333 .route(routes::V1_REQUESTS, get(recent_requests))
3334 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3335 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3336 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3337 .route(
3338 routes::LORA_ADAPTERS,
3339 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3340 )
3341 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3342 .route(routes::V1_CANCEL, post(cancel_generation))
3346 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3351 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3352 .route(routes::V1_MESSAGES, post(anthropic::messages))
3353 .route(
3354 routes::V1_MESSAGES_COUNT_TOKENS,
3355 post(anthropic::count_tokens),
3356 )
3357 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3358 .route(routes::COMPLETION, post(completion::completion))
3363 .route(routes::COMPLETIONS, post(completion::completion))
3364 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3365 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3366 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3372 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3373 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3374 .route(routes::V1_RERANK, post(rerank::rerank))
3378 .route(routes::RERANK, post(rerank::rerank))
3379 .route(routes::CACHE_STATS, get(cache_stats))
3380 .route(routes::METRICS, get(metrics))
3381 .route(routes::ADMIN_MODELS, get(admin::models))
3386 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3387 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3388 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3389 .route(routes::ADMIN_TASKS, get(admin::tasks))
3390 .route(&admin::cancel_route(), post(admin::cancel_task))
3391 .route(routes::ADMIN_STATS, get(admin::stats))
3392 .merge(conversations::router())
3396}
3397
3398fn axum_path(template: &str) -> String {
3399 let mut out = String::with_capacity(template.len());
3400 let mut rest = template;
3401 while let Some(open) = rest.find('{') {
3402 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3403 break;
3404 };
3405 out.push_str(&rest[..open]);
3406 out.push(':');
3407 out.push_str(&rest[open + 1..close]);
3408 rest = &rest[close + 1..];
3409 }
3410 out.push_str(rest);
3411 out
3412}
3413
3414async fn cancel_generation(
3423 State(state): State<Arc<AppState>>,
3424 Json(req): Json<frink_api::CancelGenerationRequest>,
3425) -> Response {
3426 let cancelled = state.cancels.cancel(&req.request_id);
3427 let status = if cancelled {
3428 StatusCode::OK
3429 } else {
3430 StatusCode::NOT_FOUND
3431 };
3432 let detail = if cancelled {
3433 "the generation was asked to stop; it ends at its next token".to_string()
3434 } else {
3435 "no generation with that request_id is running -- it has already \
3436 finished, was never issued, or was served by a path that does \
3437 not register for cancellation"
3438 .to_string()
3439 };
3440 (
3441 status,
3442 Json(frink_api::CancelGenerationResponse {
3443 request_id: req.request_id,
3444 cancelled,
3445 detail,
3446 }),
3447 )
3448 .into_response()
3449}
3450
3451type Activated = (
3455 Loaded,
3456 Option<serving::batch::ContinuousBatcher>,
3457 Option<Arc<budget::ContextCeiling>>,
3458);
3459
3460fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3478 let mut batcher = serving::batch::BatcherConfig::from_env();
3479 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3480 return batcher;
3483 }
3484 let Some(path) = path else {
3485 return batcher;
3486 };
3487 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3493 let Some((priced, gguf_ctx, source)) = priced else {
3494 return batcher;
3495 };
3496 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3497 tracing::warn!(
3501 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3502 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3503 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3504 admit on a number you choose.",
3505 priced.weights_bytes,
3506 priced.device_budget_bytes,
3507 );
3508 return batcher;
3509 };
3510 tracing::info!("{source}");
3511 tracing::info!("{}", derived.fit);
3512 let adopted = budget::apply_derived(&mut batcher, &derived);
3513 if adopted.max_context {
3514 tracing::info!(
3515 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3516 override with FRINK_CB_MAX_CONTEXT",
3517 derived.max_context
3518 );
3519 }
3520 if adopted.kv_blocks {
3521 tracing::info!(
3522 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3523 derived.kv_blocks,
3524 batcher.kv_block_size
3525 );
3526 }
3527 if let Some(narrowed) = adopted.max_context_narrowed {
3528 tracing::info!(
3529 "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",
3530 batcher.kv_blocks.unwrap_or_default(),
3531 batcher.kv_block_size
3532 );
3533 }
3534 batcher
3535}
3536
3537pub(crate) fn activate_loaded_model(
3546 loaded: model::LoadedModel,
3547 enable_continuous_batching: bool,
3548 path: Option<&str>,
3549 paged_kv: Option<&generate::PagedKvConfig>,
3550) -> Activated {
3551 match loaded {
3552 model::LoadedModel::Gguf(g) => {
3553 let decoder = Arc::new(g.decoder);
3554 let tokenizer = Arc::new(g.tokenizer);
3555 let config = price_batcher_config(path);
3556 let shape =
3562 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3563 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3564 let batcher = if enable_continuous_batching {
3565 tracing::info!(
3566 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3567 (stop sequences use the same pending-buffer trim as the private generate loop)"
3568 );
3569 let tok = Arc::clone(&tokenizer);
3570 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3571 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3572 Arc::clone(&decoder),
3573 decode,
3574 config,
3575 Arc::clone(&ceiling),
3576 paged_kv.cloned(),
3577 ))
3578 } else {
3579 None
3580 };
3581 (
3582 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3583 decoder,
3584 tokenizer,
3585 stop_tokens: g.stop_tokens,
3586 bos_id: g.bos_id,
3587 is_synthetic: g.is_synthetic,
3588 chat_template: g.chat_template,
3589 }))),
3590 batcher,
3591 Some(ceiling),
3592 )
3593 }
3594 model::LoadedModel::Kimi(k) => (
3595 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3596 engine: k.engine,
3597 tokenizer: k.tokenizer,
3598 stop_tokens: k.stop_tokens,
3599 chat_template: k.chat_template,
3600 }))),
3601 None,
3602 None,
3603 ),
3604 model::LoadedModel::Mla(m) => (
3605 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3606 engine: m.engine,
3607 tokenizer: m.tokenizer,
3608 stop_tokens: m.stop_tokens,
3609 bos_id: m.bos_id,
3610 name: m.name,
3611 chat_template: m.chat_template,
3612 }))),
3613 None,
3614 None,
3615 ),
3616 model::LoadedModel::Gemma4(m) => (
3617 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3618 engine: m.engine,
3619 tokenizer: m.tokenizer,
3620 stop_tokens: m.stop_tokens,
3621 bos_id: m.bos_id,
3622 name: m.name,
3623 chat_template: m.chat_template,
3624 }))),
3625 None,
3626 None,
3627 ),
3628 model::LoadedModel::Glm52(g) => (
3629 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3630 engine: g.engine,
3631 tokenizer: g.tokenizer,
3632 stop_tokens: g.stop_tokens,
3633 bos_id: g.bos_id,
3634 name: g.name,
3635 chat_template: g.chat_template,
3636 }))),
3637 None,
3638 None,
3639 ),
3640 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3645 }
3646}
3647
3648struct StartupModels {
3655 loaded: model::LoadedModel,
3656 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3657}
3658
3659fn continuous_batching_env() -> Option<bool> {
3660 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3661 .ok()
3662 .map(|v| v.trim().to_ascii_lowercase())
3663 .as_deref()
3664 {
3665 None => None,
3666 Some("1" | "true" | "yes" | "on") => Some(true),
3667 Some("0" | "false" | "no" | "off") => Some(false),
3668 _ => None,
3669 }
3670}
3671
3672fn metal_private_decode_active() -> bool {
3673 #[cfg(feature = "metal")]
3674 {
3675 BUILT_WITH_METAL
3676 && frink_metal::attn::metal_attn_enabled()
3677 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3678 }
3679 #[cfg(not(feature = "metal"))]
3680 {
3681 false
3682 }
3683}
3684
3685fn continuous_batching_compatible(
3686 loaded: &model::LoadedModel,
3687 kv_pool: &Option<generate::KvPoolConfig>,
3688 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3689 paged_kv: &Option<generate::PagedKvConfig>,
3690) -> bool {
3691 matches!(loaded, model::LoadedModel::Gguf(_))
3692 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3693}
3694
3695fn resolve_continuous_batching_enabled(
3696 loaded: &model::LoadedModel,
3697 kv_pool: &Option<generate::KvPoolConfig>,
3698 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3699 paged_kv: &Option<generate::PagedKvConfig>,
3700) -> bool {
3701 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3702 return false;
3703 }
3704 match continuous_batching_env() {
3705 Some(true) => true,
3706 Some(false) => false,
3707 None => metal_private_decode_active(),
3708 }
3709}
3710
3711fn acquire_metal_private_decode_gate(
3712 gate: Option<&std::sync::Mutex<()>>,
3713 used_batcher: bool,
3714) -> Option<std::sync::MutexGuard<'_, ()>> {
3715 if used_batcher {
3716 None
3717 } else {
3718 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3719 }
3720}
3721
3722fn build_app_state(
3723 models: StartupModels,
3724 kv_pool: Option<generate::KvPoolConfig>,
3725 paged_kv: Option<generate::PagedKvConfig>,
3726 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3727 enable_continuous_batching: bool,
3728 mcp: Option<mcp::LoadedMcpConfig>,
3729 detection: Arc<health::Detection>,
3730) -> AppState {
3731 let StartupModels { loaded, embedding } = models;
3732 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3733 let (loaded, batcher, ceiling) = activate_loaded_model(
3734 loaded,
3735 enable_continuous_batching,
3736 configured_path.as_deref(),
3737 paged_kv.as_ref(),
3738 );
3739 let id = startup_model_id();
3745 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3746 {
3747 None
3748 } else {
3749 tracing::info!(
3750 "Metal private-loop decode will serialize concurrent requests until \
3751 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3752 );
3753 Some(Arc::new(std::sync::Mutex::new(())))
3754 };
3755 AppState {
3756 embedding,
3757 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3758 id,
3759 loaded,
3760 batcher,
3761 ceiling,
3762 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3763 }))),
3764 paged_kv,
3765 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3766 tasks: Arc::new(tasks::TaskRegistry::new()),
3767 cancels: Arc::new(cancel::CancelRegistry::new()),
3768 stats: stats::Stats::new(),
3769 streams: resume::StreamRegistry::new(),
3770 model_dir: admin::model_dirs().into_iter().next(),
3771 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3772 kv_pool,
3773 prefix_cache,
3774 sessions: session::SessionStore::new(),
3775 requests_total: std::sync::atomic::AtomicU64::new(0),
3776 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3777 started_at: std::time::Instant::now(),
3778 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3779 detection,
3780 mcp,
3781 continuous_batching_enabled: enable_continuous_batching,
3782 metal_private_decode_gate,
3783 loading_model: Mutex::new(None),
3784 last_load_error: Mutex::new(None),
3785 serving: Mutex::new(crate::stats::ServingStats::default()),
3786 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3787 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3788 started_unix: unix_now(),
3789 }
3790}
3791
3792fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3800 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3801 return Ok(None);
3802 };
3803 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3804 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3805 tracing::info!(
3806 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3807 model.name(),
3808 model.architecture(),
3809 model.n_embd(),
3810 model.pooling_type().name(),
3811 model.n_ctx_train(),
3812 );
3813 Ok(Some(Arc::new(model)))
3814}
3815
3816fn unix_now() -> u64 {
3821 std::time::SystemTime::now()
3822 .duration_since(std::time::UNIX_EPOCH)
3823 .map(|d| d.as_secs())
3824 .unwrap_or(0)
3825}
3826
3827fn startup_model_id() -> Option<String> {
3832 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3833 let configured = std::fs::canonicalize(&configured).ok()?;
3834 admin::discover(&admin::model_dirs())
3835 .into_iter()
3836 .find(|d| {
3837 std::fs::canonicalize(&d.path)
3838 .map(|p| p == configured)
3839 .unwrap_or(false)
3840 })
3841 .map(|d| d.id)
3842}
3843
3844fn init_cpu_pool() {
3853 match frink_core::threads::init_cpu_pool() {
3854 Some(n) => eprintln!(
3855 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
3856 frink_core::threads::perf_core_count()
3857 ),
3858 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
3859 }
3860}
3861
3862fn announce_ready(addr: SocketAddr, scheme: &str) {
3875 use std::io::Write;
3876 let ready =
3877 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3878 let mut stdout = std::io::stdout().lock();
3879 let _ = writeln!(stdout, "{}", ready.to_line());
3880 let _ = stdout.flush();
3881}
3882
3883async fn shutdown_signal(exit_on_stdin_close: bool) {
3895 if !exit_on_stdin_close {
3896 std::future::pending::<()>().await;
3897 return;
3898 }
3899 let _ = tokio::task::spawn_blocking(|| {
3900 use std::io::Read;
3901 let mut sink = [0u8; 256];
3902 let mut stdin = std::io::stdin().lock();
3903 loop {
3904 match stdin.read(&mut sink) {
3905 Ok(0) => break,
3907 Ok(_) => continue,
3909 Err(e) => {
3910 tracing::warn!("stdin read failed ({e}); treating it as closed");
3911 break;
3912 }
3913 }
3914 }
3915 })
3916 .await;
3917 tracing::info!("stdin closed; shutting down");
3918}
3919
3920fn tokio_worker_threads() -> usize {
3925 std::env::var("FRINK_TOKIO_WORKERS")
3926 .ok()
3927 .and_then(|v| v.trim().parse::<usize>().ok())
3928 .filter(|n| *n > 0)
3929 .unwrap_or(2)
3930}
3931
3932fn install_ring_crypto_provider() {
3949 let _ = rustls::crypto::ring::default_provider().install_default();
3950}
3951
3952pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3958 if args.list_devices {
3959 frink_models::devices::print_available_devices();
3960 return Ok(());
3961 }
3962 apply_cli_overrides(&args)?;
3963
3964 let _instance = {
3968 use frink_core::instance::{register, InstancePolicy};
3969 let policy = if args.allow_multiple_instances {
3970 InstancePolicy::Multi
3971 } else {
3972 InstancePolicy::from_env_or(InstancePolicy::Single)
3973 };
3974 let model = std::env::var("FRINK_MODEL_PATH").ok();
3975 register(
3976 "server",
3977 model.as_deref(),
3978 frink_core::instance::current_backend(),
3979 policy,
3980 )
3981 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
3982 };
3983
3984 let journal = journal::Journal::from_env();
3985 eprintln!(
3986 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
3987 journal.path()
3988 );
3989 journal.append(&journal::Record::session_start(
3990 env!("CARGO_PKG_VERSION"),
3991 std::process::id(),
3992 ));
3993 journal::install_panic_hook(journal.clone());
3994
3995 let mcp_config_path = args.mcp_config.clone();
3996 let exit_on_stdin_close = args.exit_on_stdin_close
3997 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
3998 .map(|v| v == "1")
3999 .unwrap_or(false);
4000
4001 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4005 init_cpu_pool();
4006
4007 let runtime = tokio::runtime::Builder::new_multi_thread()
4008 .worker_threads(tokio_worker_threads())
4009 .enable_all()
4010 .build()?;
4011 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4012
4013 let reason = match &result {
4014 Ok(()) => "normal".to_string(),
4015 Err(e) => e.to_string(),
4016 };
4017 journal.append(&journal::Record::session_exit(reason));
4018
4019 runtime.shutdown_background();
4024
4025 result
4026}
4027
4028async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4029 let _ = tracing_subscriber::fmt::try_init();
4035
4036 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4046 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4047 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4048 .map(|v| v == "1")
4049 .unwrap_or(false);
4050 if let Err(msg) =
4051 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4052 {
4053 anyhow::bail!(msg);
4054 }
4055
4056 let embedding_model = load_embedding_model()?;
4063
4064 let mut loaded = model::load()?;
4065 match &loaded {
4066 model::LoadedModel::Gguf(g) => tracing::info!(
4067 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4068 g.decoder.config.name,
4069 g.is_synthetic,
4070 g.tokenizer.kind()
4071 ),
4072 model::LoadedModel::Kimi(k) => tracing::info!(
4073 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4074 k.tokenizer.vocab_size()
4075 ),
4076 model::LoadedModel::Mla(m) => tracing::info!(
4077 "loaded MLA GGUF '{}' (tokenizer={})",
4078 m.name,
4079 m.tokenizer.kind()
4080 ),
4081 model::LoadedModel::Gemma4(m) => tracing::info!(
4082 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4083 m.name,
4084 m.tokenizer.kind()
4085 ),
4086 model::LoadedModel::Glm52(g) => tracing::info!(
4087 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4088 g.name,
4089 g.tokenizer.kind()
4090 ),
4091 model::LoadedModel::Encoder(_) => {}
4094 }
4095 let metal_default_moe_budget = {
4103 #[cfg(feature = "metal")]
4104 {
4105 frink_core::metal_dense_enabled()
4106 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4107 }
4108 #[cfg(not(feature = "metal"))]
4109 {
4110 false
4111 }
4112 };
4113 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4114 let budget: u64 = budget_str
4115 .parse()
4116 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4117 match &mut loaded {
4118 model::LoadedModel::Gguf(g) => {
4119 tracing::info!(
4120 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4121 (CUDA and/or Metal matvecs when built with the matching feature)"
4122 );
4123 g.decoder.gpu_vram_budget_bytes = Some(budget);
4124 }
4125 model::LoadedModel::Kimi(_) => {
4126 tracing::warn!(
4127 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4128 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4129 );
4130 }
4131 model::LoadedModel::Mla(_) => {
4132 tracing::warn!(
4133 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4134 FFN path only today; ignoring expert VRAM budget"
4135 );
4136 }
4137 model::LoadedModel::Gemma4(_) => {
4138 tracing::warn!(
4139 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4140 ignoring expert VRAM budget"
4141 );
4142 }
4143 model::LoadedModel::Glm52(_) => {
4144 tracing::warn!(
4145 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4146 GPU expert placement not wired yet; ignoring"
4147 );
4148 }
4149 model::LoadedModel::Encoder(_) => {
4150 tracing::warn!(
4151 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4152 it has no routed experts to place; ignoring"
4153 );
4154 }
4155 }
4156 } else if metal_default_moe_budget {
4157 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4161 if let model::LoadedModel::Gguf(g) = &mut loaded {
4162 tracing::info!(
4163 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4164 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4165 );
4166 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4167 }
4168 }
4169 #[cfg(feature = "cuda")]
4170 {
4171 if frink_core::cuda_dense_enabled() {
4172 tracing::info!(
4173 "CUDA dense matvec enabled for WeightMatrix::apply \
4174 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4175 );
4176 } else {
4177 tracing::info!(
4178 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4179 );
4180 }
4181 }
4182 #[cfg(feature = "metal")]
4183 {
4184 if frink_core::metal_dense_enabled() {
4185 tracing::info!(
4186 "Metal dense matvec enabled for WeightMatrix::apply \
4187 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4188 );
4189 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4190 Some("1") | Some("true") | Some("on") | Some("attn") => {
4191 tracing::info!(
4192 "Metal fused attention requested (FRINK_METAL_ATTN): \
4193 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4194 );
4195 }
4196 _ => {}
4197 }
4198 tracing::info!(
4199 "Metal greedy GPU argmax: temperature<=0 folds \
4200 final_norm+lm_head+argmax into the dense stack"
4201 );
4202 } else {
4203 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4204 }
4205 }
4206 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4217 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4218 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4219 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4220 panic!(
4221 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4222 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4223 );
4224 }
4225 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4226 (Ok(blocks), Ok(block_size), Err(_)) => {
4227 let total_blocks: usize = blocks
4228 .parse()
4229 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4230 let block_size: usize = block_size
4231 .parse()
4232 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4233 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4239 .ok()
4240 .map(|v| {
4241 v.parse()
4242 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4243 })
4244 .unwrap_or(0);
4245 tracing::info!(
4246 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4247 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4248 );
4249 Some(generate::KvPoolConfig {
4250 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4251 queue_wait: Duration::from_millis(queue_wait_ms),
4252 })
4253 }
4254 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4255 let block_size: usize = block_size
4256 .parse()
4257 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4258 let budget: u64 = byte_budget
4259 .parse()
4260 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4261 let cfg = match &loaded {
4262 model::LoadedModel::Gguf(g) => &g.decoder.config,
4263 model::LoadedModel::Kimi(_)
4264 | model::LoadedModel::Mla(_)
4265 | model::LoadedModel::Gemma4(_)
4266 | model::LoadedModel::Glm52(_)
4267 | model::LoadedModel::Encoder(_) => {
4268 panic!(
4269 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4270 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4271 );
4272 }
4273 };
4274 let bytes_per_block = block_size
4275 * cfg.kv_heads_all_layers()
4276 * (cfg.head_dim + cfg.v_head_dim())
4277 * std::mem::size_of::<f32>();
4278 assert!(
4279 bytes_per_block > 0,
4280 "derived KV block byte size must be positive (check model config and block size)"
4281 );
4282 let total_blocks = (budget as usize / bytes_per_block).max(1);
4283 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4284 .ok()
4285 .map(|v| {
4286 v.parse()
4287 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4288 })
4289 .unwrap_or(0);
4290 tracing::info!(
4291 "KV cache block pool enabled from byte budget: {budget} bytes / \
4292 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4293 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4294 cfg.n_layers
4295 );
4296 Some(generate::KvPoolConfig {
4297 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4298 queue_wait: Duration::from_millis(queue_wait_ms),
4299 })
4300 }
4301 (Err(_), Err(_), Err(_)) => None,
4302 (Err(_), Ok(_), Err(_)) => panic!(
4303 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4304 (or unset all three to disable KV cache pooling)"
4305 ),
4306 (Ok(_), Ok(_), Ok(_)) => {
4307 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4308 }
4309 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4310 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4311 set together (or neither, to disable KV cache pooling)"
4312 ),
4313 };
4314 let paged_kv = match (
4320 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4321 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4322 ) {
4323 (Ok(blocks), Ok(block_size)) => {
4324 assert!(
4325 kv_pool.is_none(),
4326 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4327 mutually exclusive: both bound the same KV memory, by different means. \
4328 Set one."
4329 );
4330 let blocks_per_layer: usize = blocks
4344 .parse()
4345 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4346 let block_size: usize = block_size
4347 .parse()
4348 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4349 let gguf = match &loaded {
4350 model::LoadedModel::Gguf(g) => g,
4351 _ => panic!(
4352 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4353 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4354 ),
4355 };
4356 let cfg = &gguf.decoder.config;
4357 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4358 .ok()
4359 .map(|v| {
4360 v.parse()
4361 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4362 })
4363 .unwrap_or(0);
4364 tracing::info!(
4365 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4366 layer across {} layers, shared by all concurrent requests, \
4367 {queue_wait_ms}ms admission queue wait",
4368 cfg.n_layers
4369 );
4370 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4375 block_size,
4376 ))));
4377 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4386 crate::policy::parser::ToolCallFormat::infer(
4387 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4388 )
4389 .opener(),
4390 |text| {
4391 gguf.tokenizer
4392 .encode(text, SpecialTokens::Parse)
4393 .into_iter()
4394 .map(|t| t as u32)
4395 .collect()
4396 },
4397 );
4398 if let Some(id) = anchor_token {
4399 tracing::info!(
4400 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4401 window stops short of where its next turn rejoins"
4402 );
4403 }
4404 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4405 .ok()
4406 .map(|v| {
4407 v.parse()
4408 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4409 })
4410 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4411 if let Some(window) = cfg.uniform_sliding_window() {
4412 tracing::info!(
4413 "Paged KV window slide enabled: every layer slides by {window} every \
4414 {slide_interval} decode steps, so a request holds its prompt and a \
4415 window rather than its whole context"
4416 );
4417 } else if cfg.kv_block_window().is_some() {
4418 tracing::info!(
4419 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4420 and a page group holds one block in every layer"
4421 );
4422 }
4423 Some(generate::PagedKvConfig {
4424 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4427 queue_wait: Duration::from_millis(queue_wait_ms),
4428 radix,
4429 anchor_token,
4430 slide_interval,
4431 })
4432 }
4433 (Err(_), Err(_)) => None,
4434 _ => panic!(
4435 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4436 (or neither, to disable paged KV)"
4437 ),
4438 };
4439 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4446 let max_entries: usize = v
4447 .parse()
4448 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4449 if kv_pool.is_some() {
4450 tracing::warn!(
4451 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4452 caching will never be consulted while a KV pool is configured"
4453 );
4454 }
4455 assert!(
4462 paged_kv.is_none(),
4463 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4464 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4465 produce, so the cache could never hit. Set one."
4466 );
4467 tracing::info!(
4468 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4469 all requests"
4470 );
4471 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4472 });
4473 if matches!(
4474 loaded,
4475 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4476 ) && (kv_pool.is_some() || prefix_cache.is_some())
4477 {
4478 tracing::warn!(
4479 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4480 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4481 frink_models::engine's module docs"
4482 );
4483 }
4484 let enable_cb =
4485 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4486 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4487 tracing::info!(
4488 "continuous batching enabled by default on Metal for safe parallel serving \
4489 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4490 );
4491 }
4492 if continuous_batching_env() == Some(true)
4493 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4494 && (kv_pool.is_some() || prefix_cache.is_some())
4495 {
4496 tracing::warn!(
4497 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4498 (those modes keep the private generate path)"
4499 );
4500 }
4501 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4502 if let Ok(chunk) = n.parse::<usize>() {
4503 if chunk > 0 {
4504 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4505 }
4506 }
4507 }
4508 if matches!(
4509 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4510 Some("1")
4511 ) {
4512 tracing::warn!(
4513 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4514 (minimal spill; full layer offload still planned)"
4515 );
4516 }
4517
4518 let mcp = match mcp_config_path {
4519 Some(path) => {
4520 let loaded = mcp::load_mcp_config(&path)?;
4521 tracing::info!(
4522 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4523 loaded.path,
4524 loaded.servers.len()
4525 );
4526 Some(loaded)
4527 }
4528 None => None,
4529 };
4530
4531 let detection = health::Detection::spawn();
4535
4536 let state = Arc::new(build_app_state(
4537 StartupModels {
4538 loaded,
4539 embedding: embedding_model,
4540 },
4541 kv_pool,
4542 paged_kv,
4543 prefix_cache,
4544 enable_cb,
4545 mcp,
4546 detection,
4547 ));
4548
4549 use frink_api::routes;
4553
4554 let public = Router::new().route(routes::HEALTH, get(health));
4559
4560 let mut protected = protected_routes();
4561
4562 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4567 tracing::info!("API key auth enabled");
4568 let auth = limits::AuthConfig {
4569 api_key: Arc::new(key),
4570 };
4571 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4572 auth,
4573 limits::require_api_key,
4574 ));
4575 }
4576 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4577 let rpm: u32 = rpm
4578 .parse()
4579 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4580 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4581 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4582 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4583 limiter,
4584 limits::rate_limit,
4585 ));
4586 }
4587 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4595 let origins = security::parse_cors_origins(&spec)
4596 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4597 tracing::info!(
4598 "CORS enabled: {} allow-listed origin(s) ({})",
4599 origins.len(),
4600 spec
4601 );
4602 let cors = tower_http::cors::CorsLayer::new()
4603 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4604 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4605 .allow_headers([
4606 axum::http::header::CONTENT_TYPE,
4607 axum::http::header::AUTHORIZATION,
4608 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4614 axum::http::HeaderName::from_static("last-event-id"),
4618 ]);
4619 protected = protected.route_layer(cors);
4620 }
4621
4622 let app = public
4626 .merge(protected)
4627 .layer(axum::middleware::from_fn(limits::retry_after))
4628 .with_state(state);
4629
4630 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4636 install_ring_crypto_provider();
4637 match tls_paths {
4642 Some(paths) => {
4643 let config =
4644 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4645 .await
4646 .map_err(|e| {
4647 anyhow::anyhow!(
4648 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4649 paths.cert,
4650 paths.key
4651 )
4652 })?;
4653 let socket_addr: std::net::SocketAddr = addr
4654 .parse()
4655 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4656 let listener = std::net::TcpListener::bind(socket_addr)?;
4657 listener.set_nonblocking(true)?;
4665 let bound = listener.local_addr()?;
4666 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4667 announce_ready(bound, "https");
4668
4669 let handle = axum_server::Handle::new();
4670 let shutdown_handle = handle.clone();
4671 tokio::spawn(async move {
4672 shutdown_signal(exit_on_stdin_close).await;
4673 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4674 });
4675 axum_server::from_tcp_rustls(listener, config)?
4676 .handle(handle)
4677 .serve(app.into_make_service())
4678 .await?;
4679 }
4680 None => {
4681 let listener = tokio::net::TcpListener::bind(&addr).await?;
4682 let bound = listener.local_addr()?;
4683 tracing::info!("frink-server listening on {bound}");
4684 announce_ready(bound, "http");
4685 axum::serve(listener, app)
4686 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4687 .await?;
4688 }
4689 }
4690 Ok(())
4691}
4692
4693#[cfg(test)]
4694pub(crate) mod tests {
4695 use super::*;
4696 use frink_models::config::test_dense_fixture;
4697
4698 #[test]
4699 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4700 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4701 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4702 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4703 assert_eq!(parsed.port, 51999);
4704 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4705 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4707 }
4708
4709 fn test_model() -> Model {
4710 let cfg = test_dense_fixture();
4714 Model::Gguf(GgufModel {
4715 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4716 tokenizer: Arc::new(ServerTokenizer::Byte),
4717 stop_tokens: StopTokens::default(),
4718 bos_id: None,
4719 is_synthetic: true,
4720 chat_template: chat_template::PromptTemplate::plain(),
4721 })
4722 }
4723
4724 fn greedy_params(max_tokens: usize) -> GenerationParams {
4725 GenerationParams {
4726 n: 1,
4727 reasoning: None,
4728 max_tokens,
4729 sampling: SamplingParams::default(),
4730 seed: 1,
4731 stop: Vec::new(),
4732 stop_token_ids: Vec::new(),
4733 json_object: false,
4734 grammar: None,
4735 cancel: None,
4736 ignore_eos: false,
4737 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4738 lora: None,
4739 }
4740 }
4741
4742 fn test_model_full_byte_vocab() -> Model {
4746 test_model_full_byte_vocab_with_eos(None)
4747 }
4748
4749 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4756 let mut cfg = test_dense_fixture();
4757 cfg.vocab_size = 256;
4758 Model::Gguf(GgufModel {
4759 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4760 tokenizer: Arc::new(ServerTokenizer::Byte),
4761 stop_tokens: StopTokens::from_eos(eos),
4762 bos_id: None,
4763 is_synthetic: true,
4764 chat_template: chat_template::PromptTemplate::plain(),
4765 })
4766 }
4767
4768 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4772 AppState {
4773 embedding: None,
4774 paged_kv: None,
4775 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4776 id: None,
4777 loaded: Loaded::Generative(Arc::new(model)),
4778 batcher: None,
4779 ceiling: None,
4780 checkpoint_path: None,
4781 }))),
4782 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4783 tasks: Arc::new(tasks::TaskRegistry::new()),
4784 cancels: Arc::new(cancel::CancelRegistry::new()),
4785 stats: stats::Stats::new(),
4786 streams: resume::StreamRegistry::new(),
4787 model_dir: None,
4788 response_cache: Mutex::new(response_cache),
4789 kv_pool: None,
4790 prefix_cache: None,
4791 sessions: session::SessionStore::new(),
4792 requests_total: std::sync::atomic::AtomicU64::new(0),
4793 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4794 started_at: std::time::Instant::now(),
4795 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4796 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4797 mcp: None,
4798 continuous_batching_enabled: false,
4799 metal_private_decode_gate: None,
4800 loading_model: Mutex::new(None),
4801 last_load_error: Mutex::new(None),
4802 serving: Mutex::new(crate::stats::ServingStats::default()),
4803 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4804 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4805 started_unix: unix_now(),
4806 }
4807 }
4808
4809 pub(crate) fn test_app() -> Router {
4818 test_app_with_state(Arc::new(test_state(
4819 test_model_full_byte_vocab(),
4820 ResponseCache::new(1000, Duration::from_secs(3600)),
4821 )))
4822 }
4823
4824 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4827 Router::new()
4837 .route(frink_api::routes::HEALTH, get(health))
4838 .merge(protected_routes())
4839 .with_state(state)
4840 }
4841
4842 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4843 let mut cfg = test_dense_fixture();
4844 cfg.name = name;
4845 cfg.vocab_size = vocab_size;
4846 Model::Gguf(GgufModel {
4847 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4848 tokenizer: Arc::new(ServerTokenizer::Byte),
4849 stop_tokens: StopTokens::default(),
4850 bos_id: None,
4851 is_synthetic: true,
4852 chat_template: chat_template::PromptTemplate::plain(),
4853 })
4854 }
4855
4856 fn model_with_template(name: &'static str, source: &str) -> Model {
4860 let mut cfg = test_dense_fixture();
4861 cfg.name = name;
4862 cfg.vocab_size = 256;
4863 Model::Gguf(GgufModel {
4864 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4865 tokenizer: Arc::new(ServerTokenizer::Byte),
4866 stop_tokens: StopTokens::default(),
4867 bos_id: None,
4868 is_synthetic: true,
4869 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4870 Some(source),
4871 Some("qwen3"),
4872 false,
4873 true,
4874 None,
4875 None,
4876 ),
4877 })
4878 }
4879
4880 #[tokio::test]
4889 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4890 let strict = "{% if messages | length > 1 %}\
4893 {{ raise_exception('this template takes one turn') }}\
4894 {% endif %}{{ messages[0].content }}";
4895 let state = Arc::new(test_state(
4896 model_with_template("strict", strict),
4897 ResponseCache::new(4, Duration::from_secs(60)),
4898 ));
4899 let app = test_app_with_state(state);
4900
4901 let (status, body) = post_json_uri(
4902 &app,
4903 "/v1/chat/completions",
4904 serde_json::json!({
4905 "model": "strict",
4906 "stream": true,
4907 "messages": [
4908 {"role": "user", "content": "one"},
4909 {"role": "user", "content": "two"},
4910 ],
4911 }),
4912 )
4913 .await;
4914 assert_eq!(status, StatusCode::BAD_REQUEST);
4915 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4916 assert!(
4917 body["error"]["message"]
4918 .as_str()
4919 .unwrap()
4920 .contains("one turn"),
4921 "the template's own message must reach the caller: {body}"
4922 );
4923
4924 let (status, _) = post_json_uri(
4926 &app,
4927 "/v1/chat/completions",
4928 serde_json::json!({
4929 "model": "strict",
4930 "stream": true,
4931 "max_tokens": 1,
4932 "messages": [{"role": "user", "content": "one"}],
4933 }),
4934 )
4935 .await;
4936 assert_eq!(status, StatusCode::OK);
4937 }
4938
4939 #[tokio::test]
4941 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4942 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4943 {% if reasoning_effort %}\
4944 {% if reasoning_effort not in ['low','medium','high'] %}\
4945 {{ raise_exception('bad effort') }}\
4946 {% endif %}[{{ reasoning_effort }}]\
4947 {% endif %}{{ messages[0].content }}";
4948 let state = Arc::new(test_state(
4949 model_with_template("thinker", reasoning),
4950 ResponseCache::new(4, Duration::from_secs(60)),
4951 ));
4952 let app = test_app_with_state(state);
4953 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4954 assert_eq!(status, StatusCode::OK);
4955 let entry = &models["data"][0];
4956 assert_eq!(
4957 entry["supported_reasoning_efforts"],
4958 serde_json::json!(["off", "low", "medium", "high"])
4959 );
4960 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4961 }
4962
4963 #[tokio::test]
4968 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4969 let app = test_app();
4970 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4971 let entry = &models["data"][0];
4972 assert!(entry.get("supported_reasoning_efforts").is_none());
4973 assert!(entry.get("default_reasoning_effort").is_none());
4974 }
4975
4976 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
4977 Arc::new(ActiveModel {
4978 id: Some(name.to_string()),
4979 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
4980 batcher: None,
4981 ceiling: None,
4982 checkpoint_path: None,
4983 })
4984 .tap_into(state)
4985 }
4986
4987 trait TapInto {
4989 fn tap_into(self, state: &AppState) -> Self;
4990 }
4991 impl TapInto for Arc<ActiveModel> {
4992 fn tap_into(self, state: &AppState) -> Self {
4993 state.swap_active(Some(Arc::clone(&self)));
4994 self
4995 }
4996 }
4997
4998 #[test]
5004 fn an_in_flight_request_keeps_the_model_it_started_on() {
5005 let state = test_state(
5006 named_test_model("model-a", 256),
5007 ResponseCache::new(4, Duration::from_secs(60)),
5008 );
5009
5010 let in_flight = state.active().expect("a model is loaded");
5013 assert_eq!(in_flight.name(), "model-a");
5014
5015 active_model(&state, "model-b");
5016
5017 assert_eq!(state.active().unwrap().name(), "model-b");
5019 assert_eq!(in_flight.name(), "model-a");
5021 let (choices, _usage) = run_generation(
5022 in_flight.generative().unwrap(),
5023 "hi",
5024 &greedy_params(3),
5025 None,
5026 None,
5027 None,
5028 None,
5029 None,
5030 None,
5031 )
5032 .expect("the old model must still decode after being swapped out");
5033 assert!(matches!(
5034 choices[0].0,
5035 FinishReason::Length | FinishReason::Stop
5036 ));
5037 }
5038
5039 #[test]
5044 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5045 let state = test_state(
5046 named_test_model("model-a", 256),
5047 ResponseCache::new(4, Duration::from_secs(60)),
5048 );
5049 let in_flight = state.active().expect("a model is loaded");
5050 let weights = Arc::clone(in_flight.generative().unwrap());
5051 assert!(Arc::strong_count(&weights) >= 2);
5052
5053 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5054 id: Some("model-b".to_string()),
5055 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5056 batcher: None,
5057 ceiling: None,
5058 checkpoint_path: None,
5059 })));
5060 drop(previous);
5061 assert!(Arc::strong_count(&weights) >= 2);
5063 drop(in_flight);
5064 assert_eq!(Arc::strong_count(&weights), 1);
5065 }
5066
5067 #[tokio::test]
5071 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5072 let state = Arc::new(test_state(
5073 named_test_model("model-a", 256),
5074 ResponseCache::new(4, Duration::from_secs(60)),
5075 ));
5076 let app = test_app_with_state(Arc::clone(&state));
5077
5078 let (status, body) = post_json_uri(
5079 &app,
5080 frink_api::routes::ADMIN_MODELS_UNLOAD,
5081 serde_json::json!({}),
5082 )
5083 .await;
5084 assert_eq!(status, StatusCode::OK);
5085 assert_eq!(body["ok"], true);
5086 assert!(body["active"].is_null());
5087 assert!(state.active().is_none());
5088
5089 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5090 assert_eq!(status, StatusCode::OK);
5091 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5092 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5093
5094 let (status, body) = post_json_uri(
5095 &app,
5096 "/v1/chat/completions",
5097 serde_json::json!({
5098 "model": "x",
5099 "messages": [{"role": "user", "content": "hi"}]
5100 }),
5101 )
5102 .await;
5103 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5104 assert_eq!(body["error"]["type"], "model_not_loaded");
5105 }
5106
5107 #[tokio::test]
5111 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5112 let state = Arc::new(test_state(
5113 named_test_model("model-a", 256),
5114 ResponseCache::new(4, Duration::from_secs(60)),
5115 ));
5116 let app = test_app_with_state(Arc::clone(&state));
5117 state.swap_active(None);
5118
5119 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5120 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5123 assert_eq!(body["state"], "unavailable");
5124 assert_eq!(body["reason"], "model_not_loaded");
5125 assert!(body["model"].is_null());
5126 let real_weights = body["capabilities"]
5127 .as_array()
5128 .unwrap()
5129 .iter()
5130 .find(|c| c["id"] == "real_weights")
5131 .cloned()
5132 .expect("real_weights is always reported");
5133 assert_eq!(real_weights["available"], false);
5134 assert_eq!(real_weights["reason"], "model_not_loaded");
5135 }
5136
5137 #[tokio::test]
5141 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5142 let app = test_app();
5143
5144 let (status, completion) = post_json_uri(
5145 &app,
5146 "/v1/chat/completions",
5147 serde_json::json!({
5148 "model": "x",
5149 "messages": [{"role": "user", "content": "hi"}],
5150 "max_tokens": 4
5151 }),
5152 )
5153 .await;
5154 assert_eq!(status, StatusCode::OK);
5155 let request_id = completion["request_id"].as_str().unwrap().to_string();
5156
5157 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5158 assert_eq!(status, StatusCode::OK);
5159 let recent = stats["recent"].as_array().unwrap();
5160 assert_eq!(recent.len(), 1);
5161 let row = &recent[0];
5162 assert_eq!(row["request_id"], request_id);
5163 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5164 assert_eq!(row["status"], 200);
5165 assert_eq!(row["stream"], false);
5166 assert!(row["duration_ms"].is_number());
5169 assert!(row["decode_ms"].is_number());
5170 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5171 assert_eq!(
5172 stats["tokens_prompt_total"].as_u64().unwrap(),
5173 row["prompt_tokens"].as_u64().unwrap()
5174 );
5175 }
5176
5177 #[tokio::test]
5181 async fn a_rejected_request_is_recorded_too() {
5182 let state = Arc::new(test_state(
5183 named_test_model("model-a", 256),
5184 ResponseCache::new(4, Duration::from_secs(60)),
5185 ));
5186 let app = test_app_with_state(Arc::clone(&state));
5187 state.swap_active(None);
5188
5189 let (status, _) = post_json_uri(
5190 &app,
5191 "/v1/chat/completions",
5192 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5193 )
5194 .await;
5195 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5196
5197 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5198 let recent = stats["recent"].as_array().unwrap();
5199 assert_eq!(recent.len(), 1);
5200 assert_eq!(recent[0]["status"], 503);
5201 assert_eq!(recent[0]["completion_tokens"], 0);
5202 assert!(recent[0]["decode_ms"].is_null());
5203 assert_eq!(stats["errors_total"], 1);
5204 }
5205
5206 async fn post_json_with_headers(
5210 app: &Router,
5211 uri: &str,
5212 body: serde_json::Value,
5213 headers: &[(&str, &str)],
5214 ) -> (StatusCode, serde_json::Value) {
5215 use http_body_util::BodyExt;
5216 use tower::ServiceExt;
5217
5218 let mut builder = axum::http::Request::builder()
5219 .method("POST")
5220 .uri(uri)
5221 .header("content-type", "application/json");
5222 for (name, value) in headers {
5223 builder = builder.header(*name, *value);
5224 }
5225 let response = app
5226 .clone()
5227 .oneshot(
5228 builder
5229 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5230 .unwrap(),
5231 )
5232 .await
5233 .unwrap();
5234 let status = response.status();
5235 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5236 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5237 (status, json)
5238 }
5239
5240 #[tokio::test]
5244 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5245 let app = test_app();
5246
5247 let (status, _) = post_json_uri(
5248 &app,
5249 frink_api::routes::V1_TOKENIZE,
5250 serde_json::json!({"prompt": "hello"}),
5251 )
5252 .await;
5253 assert_eq!(status, StatusCode::OK);
5254 let (status, _) = post_json_uri(
5255 &app,
5256 frink_api::routes::V1_DETOKENIZE,
5257 serde_json::json!({"tokens": [104, 105]}),
5258 )
5259 .await;
5260 assert_eq!(status, StatusCode::OK);
5261 let (status, _) = post_json_uri(
5262 &app,
5263 frink_api::routes::V1_EMBEDDINGS,
5264 serde_json::json!({"input": "hello"}),
5265 )
5266 .await;
5267 assert_eq!(status, StatusCode::OK);
5268
5269 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5270 let routes: Vec<&str> = stats["recent"]
5271 .as_array()
5272 .unwrap()
5273 .iter()
5274 .map(|row| row["route"].as_str().unwrap())
5275 .collect();
5276 for expected in [
5277 frink_api::routes::V1_TOKENIZE,
5278 frink_api::routes::V1_DETOKENIZE,
5279 frink_api::routes::V1_EMBEDDINGS,
5280 ] {
5281 assert!(
5282 routes.contains(&expected),
5283 "{expected} is missing: {routes:?}"
5284 );
5285 }
5286
5287 let row = |route: &str| {
5288 stats["recent"]
5289 .as_array()
5290 .unwrap()
5291 .iter()
5292 .find(|r| r["route"] == route)
5293 .cloned()
5294 .unwrap()
5295 };
5296 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5300 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5301 assert!(embed["decode_ms"].is_null());
5302 assert_eq!(embed["completion_tokens"], 0);
5303 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5307 assert_eq!(
5308 stats["tokens_prompt_total"].as_u64().unwrap(),
5309 embed["prompt_tokens"].as_u64().unwrap(),
5310 "only the forward pass counted"
5311 );
5312 }
5313
5314 fn streaming_test_app() -> Router {
5319 let mut cfg = test_dense_fixture();
5320 cfg.vocab_size = 256;
5321 let model = Model::Gguf(GgufModel {
5322 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5323 tokenizer: Arc::new(ServerTokenizer::Byte),
5324 stop_tokens: StopTokens::default(),
5325 bos_id: None,
5326 is_synthetic: false,
5327 chat_template: chat_template::PromptTemplate::plain(),
5328 });
5329 test_app_with_state(Arc::new(test_state(
5330 model,
5331 ResponseCache::new(1000, Duration::from_secs(3600)),
5332 )))
5333 }
5334
5335 #[tokio::test]
5348 async fn every_route_refuses_the_same_unimplemented_fields() {
5349 let app = test_app();
5350 let fields = [
5351 ("n", serde_json::json!(3)),
5352 ("best_of", serde_json::json!(2)),
5353 ("prompt_logprobs", serde_json::json!(1)),
5354 ("echo", serde_json::json!(true)),
5355 ("use_beam_search", serde_json::json!(true)),
5356 ("truncate_prompt_tokens", serde_json::json!(8)),
5357 ("prompt_embeds", serde_json::json!("AA==")),
5358 ("allowed_token_ids", serde_json::json!([1, 2])),
5359 ("bad_words", serde_json::json!(["x"])),
5360 ("skip_special_tokens", serde_json::json!(false)),
5361 ("return_tokens_as_token_ids", serde_json::json!(true)),
5362 ];
5363 for (field, value) in fields {
5364 for (uri, base) in [
5365 (
5366 frink_api::routes::V1_CHAT_COMPLETIONS,
5367 serde_json::json!({
5368 "model": "x",
5369 "messages": [{"role": "user", "content": "hi"}],
5370 "max_tokens": 2
5371 }),
5372 ),
5373 (
5374 frink_api::routes::V1_COMPLETIONS,
5375 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
5376 ),
5377 (
5378 frink_api::routes::COMPLETION,
5379 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5380 ),
5381 ] {
5382 let mut body = base;
5383 body[field] = value.clone();
5384 if field == "n" && uri == frink_api::routes::V1_COMPLETIONS {
5389 let (status, answer) = post_json_uri(&app, uri, body).await;
5390 assert_eq!(
5391 status,
5392 StatusCode::OK,
5393 "{uri} refused a served `n`: {answer}"
5394 );
5395 assert_eq!(
5396 answer["choices"].as_array().map(Vec::len),
5397 Some(3),
5398 "{answer}"
5399 );
5400 continue;
5401 }
5402 let (status, answer) = post_json_uri(&app, uri, body).await;
5403 assert_eq!(
5404 status,
5405 StatusCode::NOT_IMPLEMENTED,
5406 "{uri} served `{field}` instead of refusing it: {answer}"
5407 );
5408 assert!(
5409 answer["error"]["message"]
5410 .as_str()
5411 .is_some_and(|m| m.contains(field)),
5412 "{uri} refused `{field}` without naming it: {answer}"
5413 );
5414 }
5415 }
5416 }
5417
5418 #[tokio::test]
5419 async fn the_native_completion_wire_is_not_the_openai_one() {
5420 let app = test_app();
5421
5422 let (status, native) = post_json_uri(
5423 &app,
5424 frink_api::routes::COMPLETION,
5425 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5426 )
5427 .await;
5428 assert_eq!(status, StatusCode::OK, "{native}");
5429 assert!(native["content"].is_string(), "{native}");
5430 assert_eq!(native["stop"], true);
5431 assert_eq!(native["stop_type"], "limit");
5432 assert_eq!(native["stopping_word"], "");
5433 assert_eq!(native["truncated"], false);
5434 assert_eq!(native["id_slot"], -1);
5435 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5436 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5437 assert!(
5438 native.get("choices").is_none(),
5439 "the native shape has no `choices`: {native}"
5440 );
5441
5442 let (status, openai) = post_json_uri(
5443 &app,
5444 frink_api::routes::V1_COMPLETIONS,
5445 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5446 )
5447 .await;
5448 assert_eq!(status, StatusCode::OK);
5449 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5450 assert!(
5451 openai.get("content").is_none(),
5452 "the OpenAI shape has no top-level `content`: {openai}"
5453 );
5454 }
5455
5456 #[tokio::test]
5460 async fn both_native_spellings_reach_the_same_handler() {
5461 let app = test_app();
5462 for route in [
5463 frink_api::routes::COMPLETION,
5464 frink_api::routes::COMPLETIONS,
5465 ] {
5466 let (status, body) = post_json_uri(
5467 &app,
5468 route,
5469 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5470 )
5471 .await;
5472 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5473 assert_eq!(body["stop"], true, "{route}");
5474 assert!(body["content"].is_string(), "{route}");
5475 }
5476
5477 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5480 let routes: Vec<&str> = stats["recent"]
5481 .as_array()
5482 .unwrap()
5483 .iter()
5484 .map(|row| row["route"].as_str().unwrap())
5485 .collect();
5486 assert!(
5487 routes.contains(&frink_api::routes::COMPLETION),
5488 "{routes:?}"
5489 );
5490 assert!(
5491 routes.contains(&frink_api::routes::COMPLETIONS),
5492 "{routes:?}"
5493 );
5494 }
5495
5496 #[tokio::test]
5502 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5503 let app = streaming_test_app();
5504 let raw = post_sse_raw_uri(
5505 &app,
5506 frink_api::routes::COMPLETION,
5507 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5508 )
5509 .await;
5510
5511 assert!(
5512 !raw.contains("[DONE]"),
5513 "llama.cpp's native stream has no sentinel: {raw}"
5514 );
5515 let frames: Vec<serde_json::Value> = raw
5516 .lines()
5517 .filter_map(|line| line.strip_prefix("data: "))
5518 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5519 .collect();
5520 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5521
5522 let (last, partials) = frames.split_last().unwrap();
5523 assert_eq!(last["stop"], true, "the last frame closes the stream");
5524 assert!(last["timings"].is_object(), "{last}");
5525 assert!(last["stop_type"].is_string(), "{last}");
5526 for partial in partials {
5527 assert_eq!(partial["stop"], false, "{partial}");
5528 assert!(partial["content"].is_string(), "{partial}");
5529 assert!(partial.get("timings").is_none(), "{partial}");
5533 assert!(partial.get("generation_settings").is_none(), "{partial}");
5534 }
5535 let streamed: String = partials
5538 .iter()
5539 .filter_map(|p| p["content"].as_str())
5540 .collect();
5541 assert_eq!(last["content"].as_str().unwrap(), streamed);
5542 }
5543
5544 #[tokio::test]
5549 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5550 let app = test_app();
5551 for body in [
5552 serde_json::json!({"prompt": "hi"}),
5553 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5554 ] {
5555 let (status, refusal) =
5556 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5557 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5558 assert!(
5559 refusal["error"]["message"]
5560 .as_str()
5561 .unwrap()
5562 .contains("n_predict"),
5563 "{refusal}"
5564 );
5565 }
5566 let (status, _) = post_json_uri(
5569 &app,
5570 frink_api::routes::COMPLETION,
5571 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5572 )
5573 .await;
5574 assert_eq!(status, StatusCode::OK);
5575 }
5576
5577 #[tokio::test]
5586 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5587 let app = streaming_test_app();
5588 let ask = |stop: serde_json::Value| {
5589 let app = app.clone();
5590 async move {
5591 post_json_uri(
5592 &app,
5593 frink_api::routes::COMPLETION,
5594 serde_json::json!({
5595 "prompt": "hi",
5596 "n_predict": 64,
5597 "ignore_eos": true,
5598 "stop": stop,
5599 }),
5600 )
5601 .await
5602 .1
5603 }
5604 };
5605
5606 let baseline = ask(serde_json::json!([])).await;
5607 assert_eq!(baseline["stop_type"], "limit");
5608 assert_eq!(baseline["stopping_word"], "");
5609 let text = baseline["content"].as_str().unwrap().to_string();
5610 let sentinel: String = text.chars().skip(1).take(2).collect();
5616 assert_eq!(
5617 sentinel.chars().count(),
5618 2,
5619 "the fixture must produce enough output to cut: {text:?}"
5620 );
5621 let cut = text.find(&sentinel).expect("it came out of this text");
5622
5623 let stopped = ask(serde_json::json!([sentinel])).await;
5624 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5625 assert_eq!(stopped["stopping_word"], sentinel);
5626 assert_eq!(
5627 stopped["content"].as_str().unwrap(),
5628 &text[..cut],
5629 "the answer must be cut at the sentinel, not run past it"
5630 );
5631 }
5632
5633 #[tokio::test]
5639 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5640 let app = test_app();
5641
5642 let (v1_status, v1) = post_json_uri(
5643 &app,
5644 frink_api::routes::V1_TOKENIZE,
5645 serde_json::json!({"prompt": "hello"}),
5646 )
5647 .await;
5648 let (alias_status, alias) = post_json_uri(
5649 &app,
5650 frink_api::routes::TOKENIZE,
5651 serde_json::json!({"content": "hello"}),
5652 )
5653 .await;
5654 assert_eq!(v1_status, StatusCode::OK);
5655 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5656 assert_eq!(v1["tokens"], alias["tokens"]);
5657 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5658
5659 let (status, both_ways) = post_json_uri(
5662 &app,
5663 frink_api::routes::TOKENIZE,
5664 serde_json::json!({"prompt": "hello"}),
5665 )
5666 .await;
5667 assert_eq!(status, StatusCode::OK);
5668 assert_eq!(both_ways["tokens"], v1["tokens"]);
5669 }
5670
5671 #[tokio::test]
5676 async fn detokenize_answers_under_both_dialects_keys() {
5677 let app = test_app();
5678 for route in [
5679 frink_api::routes::DETOKENIZE,
5680 frink_api::routes::V1_DETOKENIZE,
5681 ] {
5682 let (status, body) =
5683 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5684 assert_eq!(status, StatusCode::OK, "{route}");
5685 assert_eq!(body["text"], "hi", "{route}");
5686 assert_eq!(body["content"], body["text"], "{route}");
5687 }
5688 }
5689
5690 #[tokio::test]
5694 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5695 let app = test_app();
5696 let (status, _) = post_json_uri(
5697 &app,
5698 frink_api::routes::TOKENIZE,
5699 serde_json::json!({"content": "hello"}),
5700 )
5701 .await;
5702 assert_eq!(status, StatusCode::OK);
5703
5704 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5705 let routes: Vec<&str> = stats["recent"]
5706 .as_array()
5707 .unwrap()
5708 .iter()
5709 .map(|row| row["route"].as_str().unwrap())
5710 .collect();
5711 assert!(
5712 routes.contains(&frink_api::routes::TOKENIZE),
5713 "the alias must be its own row: {routes:?}"
5714 );
5715 assert!(
5716 !routes.contains(&frink_api::routes::V1_TOKENIZE),
5717 "nothing called /v1/tokenize: {routes:?}"
5718 );
5719 }
5720
5721 #[tokio::test]
5726 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5727 let mut cfg = test_dense_fixture();
5728 cfg.vocab_size = 256;
5729 let model = Model::Gguf(GgufModel {
5730 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5731 tokenizer: Arc::new(ServerTokenizer::Byte),
5732 stop_tokens: StopTokens::default(),
5733 bos_id: Some(7),
5734 is_synthetic: true,
5735 chat_template: chat_template::PromptTemplate::plain(),
5736 });
5737 let app = test_app_with_state(Arc::new(test_state(
5738 model,
5739 ResponseCache::new(1000, Duration::from_secs(3600)),
5740 )));
5741
5742 let (_, plain) = post_json_uri(
5743 &app,
5744 frink_api::routes::TOKENIZE,
5745 serde_json::json!({"content": "hi"}),
5746 )
5747 .await;
5748 let (_, special) = post_json_uri(
5749 &app,
5750 frink_api::routes::TOKENIZE,
5751 serde_json::json!({"content": "hi", "add_special": true}),
5752 )
5753 .await;
5754
5755 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5756 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5757 assert_eq!(special["count"], 3);
5758 }
5759
5760 #[tokio::test]
5763 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5764 let app = test_app();
5765 let (status, _) = post_json_uri(
5766 &app,
5767 frink_api::routes::V1_EMBEDDINGS,
5768 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5769 )
5770 .await;
5771 assert_eq!(status, StatusCode::BAD_REQUEST);
5772
5773 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5774 let recent = stats["recent"].as_array().unwrap();
5775 assert_eq!(recent.len(), 1);
5776 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
5777 assert_eq!(recent[0]["status"], 400);
5778 assert_eq!(
5779 recent[0]["prompt_tokens"], 0,
5780 "a rejected call embedded nothing"
5781 );
5782 }
5783
5784 #[tokio::test]
5787 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5788 let app = test_app();
5789 let key = "sk-monitor-secret";
5790 let (status, _) = post_json_with_headers(
5791 &app,
5792 "/v1/chat/completions",
5793 serde_json::json!({
5794 "model": "x",
5795 "messages": [{"role": "user", "content": "hi"}],
5796 "max_tokens": 2
5797 }),
5798 &[
5799 ("authorization", &format!("Bearer {key}")),
5800 ("x-frink-client", "frink-studio"),
5801 ],
5802 )
5803 .await;
5804 assert_eq!(status, StatusCode::OK);
5805
5806 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5807 let row = stats["recent"].as_array().unwrap()[0].clone();
5808 let fingerprint = row["via_api_key"]
5809 .as_str()
5810 .expect("the row names the key that served it")
5811 .to_string();
5812 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5813 assert!(!fingerprint.contains(key));
5814 assert!(
5815 !serde_json::to_string(&stats).unwrap().contains(key),
5816 "the stats payload must not carry the key in any form"
5817 );
5818 assert_eq!(row["client"], "frink-studio");
5819 }
5820
5821 #[tokio::test]
5824 async fn different_keys_are_different_callers_and_no_key_is_null() {
5825 let app = test_app();
5826 let body = serde_json::json!({
5827 "model": "x",
5828 "messages": [{"role": "user", "content": "hi"}],
5829 "max_tokens": 1
5830 });
5831 for headers in [
5832 vec![("authorization", "Bearer key-one")],
5833 vec![("authorization", "Bearer key-two")],
5834 vec![],
5835 ] {
5836 let (status, _) =
5837 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5838 assert_eq!(status, StatusCode::OK);
5839 }
5840
5841 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5842 let recent = stats["recent"].as_array().unwrap();
5843 assert_eq!(recent.len(), 3);
5844 let one = recent[0]["via_api_key"].as_str().unwrap();
5845 let two = recent[1]["via_api_key"].as_str().unwrap();
5846 assert_ne!(one, two, "two keys must not collapse into one caller");
5847 assert!(
5848 recent[2]["via_api_key"].is_null(),
5849 "an unauthenticated call is null, not a fingerprint of nothing"
5850 );
5851 assert!(recent[2]["client"].is_null());
5852 }
5853
5854 #[tokio::test]
5859 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5860 let state = Arc::new(test_state(
5861 named_test_model("really-loaded", 256),
5862 ResponseCache::new(4, Duration::from_secs(60)),
5863 ));
5864 let app = test_app_with_state(Arc::clone(&state));
5865
5866 let (status, _) = post_json_uri(
5867 &app,
5868 "/v1/chat/completions",
5869 serde_json::json!({
5870 "model": "gpt-4-turbo-that-is-not-here",
5871 "messages": [{"role": "user", "content": "hi"}],
5872 "max_tokens": 2
5873 }),
5874 )
5875 .await;
5876 assert_eq!(status, StatusCode::OK);
5877
5878 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5879 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5880
5881 state.swap_active(None);
5884 let (status, _) = post_json_uri(
5885 &app,
5886 "/v1/chat/completions",
5887 serde_json::json!({
5888 "model": "gpt-4-turbo-that-is-not-here",
5889 "messages": [{"role": "user", "content": "hi"}]
5890 }),
5891 )
5892 .await;
5893 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5894 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5895 let recent = stats["recent"].as_array().unwrap();
5896 assert!(recent[recent.len() - 1]["model"].is_null());
5897 }
5898
5899 #[tokio::test]
5903 async fn a_streamed_row_names_the_model_it_decoded_against() {
5904 let state = Arc::new(test_state(
5905 named_test_model("model-before", 256),
5906 ResponseCache::new(4, Duration::from_secs(60)),
5907 ));
5908 let app = test_app_with_state(Arc::clone(&state));
5909 let _ = post_sse_raw(&app, resumable_request()).await;
5910 active_model(&state, "model-after");
5912
5913 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5914 assert_eq!(stats["recent"][0]["model"], "model-before");
5915 }
5916
5917 #[tokio::test]
5920 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5921 let app = test_app();
5922 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5923 assert_eq!(status, StatusCode::OK);
5924 assert!(
5925 stats["queue_depth"].is_null(),
5926 "without continuous batching nothing queues, so there is nothing to measure"
5927 );
5928 assert!(stats["queue_rejected_total"].is_null());
5929 assert_eq!(
5930 stats["generating_now"], 0,
5931 "work in progress is measured and really is zero here"
5932 );
5933 }
5934
5935 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
5940 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
5941 }
5942
5943 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
5947 use http_body_util::BodyExt;
5948 use tower::ServiceExt;
5949
5950 let response = app
5951 .clone()
5952 .oneshot(
5953 axum::http::Request::builder()
5954 .method("POST")
5955 .uri(uri)
5956 .header("content-type", "application/json")
5957 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5958 .unwrap(),
5959 )
5960 .await
5961 .unwrap();
5962 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5963 String::from_utf8(bytes.to_vec()).unwrap()
5964 }
5965
5966 async fn get_json_with_headers(
5967 app: &Router,
5968 uri: &str,
5969 headers: &[(&str, &str)],
5970 ) -> (StatusCode, serde_json::Value) {
5971 use http_body_util::BodyExt;
5972 use tower::ServiceExt;
5973
5974 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
5975 for (name, value) in headers {
5976 builder = builder.header(*name, *value);
5977 }
5978 let response = app
5979 .clone()
5980 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
5981 .await
5982 .unwrap();
5983 let status = response.status();
5984 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5985 (
5986 status,
5987 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
5988 )
5989 }
5990
5991 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
5992 body.lines()
5993 .filter_map(|line| line.strip_prefix(field))
5994 .map(str::trim)
5995 .collect()
5996 }
5997
5998 fn resumable_request() -> serde_json::Value {
5999 serde_json::json!({
6000 "model": "m",
6001 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6002 "max_tokens": 4,
6003 "temperature": 0,
6004 "stream": true,
6005 "stream_resumable": true,
6006 })
6007 }
6008
6009 #[tokio::test]
6014 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6015 let app = test_app();
6016 let body = post_sse_raw(&app, resumable_request()).await;
6017
6018 let request_id = body
6019 .lines()
6020 .find_map(|l| l.strip_prefix("data: "))
6021 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6022 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6023 .expect("the first chunk names the request");
6024
6025 let ids = sse_field(&body, "id:");
6026 let datas = sse_field(&body, "data:");
6027 assert_eq!(
6028 ids.len(),
6029 datas.len(),
6030 "every event carries an id, or a reconnect cannot name where it stopped"
6031 );
6032 for (i, id) in ids.iter().enumerate() {
6033 assert_eq!(*id, format!("{request_id}:{i}"));
6034 }
6035 let retries = sse_field(&body, "retry:");
6036 assert_eq!(
6037 retries.len(),
6038 1,
6039 "the reconnect delay is stated once, not on every event"
6040 );
6041 assert_eq!(retries[0], "1500");
6042 assert!(
6043 body.contains("data: [DONE]"),
6044 "the end of stream is still stated"
6045 );
6046 }
6047
6048 #[tokio::test]
6052 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6053 let app = test_app();
6054 let mut request = resumable_request();
6055 request["stream_resumable"] = serde_json::json!(false);
6056 let body = post_sse_raw(&app, request).await;
6057 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6058 assert!(
6059 sse_field(&body, "id:").is_empty(),
6060 "an id promises a replay this stream cannot serve"
6061 );
6062 assert!(sse_field(&body, "retry:").is_empty());
6063 }
6064
6065 #[tokio::test]
6069 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6070 let app = test_app();
6071 let body = post_sse_raw(&app, resumable_request()).await;
6072 let request_id = sse_field(&body, "id:")[0]
6073 .rsplit_once(':')
6074 .unwrap()
6075 .0
6076 .to_string();
6077 let streamed: Vec<String> = sse_field(&body, "data:")
6078 .iter()
6079 .map(|d| d.to_string())
6080 .collect();
6081
6082 let (status, polled) = get_json(
6083 &app,
6084 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6085 )
6086 .await;
6087 assert_eq!(status, StatusCode::OK);
6088 let events: Vec<String> = polled["events"]
6089 .as_array()
6090 .unwrap()
6091 .iter()
6092 .map(|e| e["data"].as_str().unwrap().to_string())
6093 .collect();
6094 assert_eq!(
6095 events, streamed,
6096 "the fallback must deliver the same answer, not a re-run of it"
6097 );
6098 assert_eq!(polled["request_id"], request_id);
6099 assert_eq!(
6100 polled["done"], false,
6101 "events were still being handed out, so the client must ask again"
6102 );
6103
6104 let next = polled["next_index"].as_u64().unwrap();
6107 let (_, drained) = get_json(
6108 &app,
6109 &format!(
6110 "{}?from={next}",
6111 frink_api::routes::v1_stream_poll(&request_id)
6112 ),
6113 )
6114 .await;
6115 assert_eq!(drained["done"], true);
6116 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6117 }
6118
6119 #[tokio::test]
6123 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6124 let app = test_app();
6125 let body = post_sse_raw(&app, resumable_request()).await;
6126 let ids = sse_field(&body, "id:");
6127 let datas: Vec<String> = sse_field(&body, "data:")
6128 .iter()
6129 .map(|d| d.to_string())
6130 .collect();
6131 assert!(
6132 ids.len() >= 3,
6133 "need a few events to resume into the middle"
6134 );
6135 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6136
6137 let (status, resumed) = get_json_with_headers(
6138 &app,
6139 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6140 &[],
6141 )
6142 .await;
6143 assert_eq!(status, StatusCode::OK);
6144 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6145
6146 let (_, tail) = get_json(
6148 &app,
6149 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6150 )
6151 .await;
6152 let tail_events: Vec<String> = tail["events"]
6153 .as_array()
6154 .unwrap()
6155 .iter()
6156 .map(|e| e["data"].as_str().unwrap().to_string())
6157 .collect();
6158 assert_eq!(tail_events, datas[2..].to_vec());
6159 }
6160
6161 #[tokio::test]
6164 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6165 use http_body_util::BodyExt;
6166 use tower::ServiceExt;
6167
6168 let app = test_app();
6169 let body = post_sse_raw(&app, resumable_request()).await;
6170 let ids = sse_field(&body, "id:");
6171 let datas: Vec<String> = sse_field(&body, "data:")
6172 .iter()
6173 .map(|d| d.to_string())
6174 .collect();
6175 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6176
6177 let response = app
6178 .clone()
6179 .oneshot(
6180 axum::http::Request::builder()
6181 .method("GET")
6182 .uri(frink_api::routes::v1_stream(&request_id))
6183 .header("last-event-id", format!("{request_id}:0"))
6184 .body(axum::body::Body::empty())
6185 .unwrap(),
6186 )
6187 .await
6188 .unwrap();
6189 assert_eq!(response.status(), StatusCode::OK);
6190 assert_eq!(
6191 response
6192 .headers()
6193 .get("x-accel-buffering")
6194 .and_then(|v| v.to_str().ok()),
6195 Some("no"),
6196 "the reconnect needs the same anti-buffering header as the stream"
6197 );
6198 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6199 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6200 assert_eq!(
6201 sse_field(&resumed, "data:")
6202 .iter()
6203 .map(|d| d.to_string())
6204 .collect::<Vec<_>>(),
6205 datas[1..].to_vec()
6206 );
6207 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6208 }
6209
6210 #[tokio::test]
6214 async fn a_last_event_id_from_another_stream_is_refused() {
6215 let app = test_app();
6216 let body = post_sse_raw(&app, resumable_request()).await;
6217 let request_id = sse_field(&body, "id:")[0]
6218 .rsplit_once(':')
6219 .unwrap()
6220 .0
6221 .to_string();
6222
6223 let (status, err) = get_json_with_headers(
6224 &app,
6225 &frink_api::routes::v1_stream(&request_id),
6226 &[("last-event-id", "chatcmpl-someone-else:3")],
6227 )
6228 .await;
6229 assert_eq!(status, StatusCode::BAD_REQUEST);
6230 assert_eq!(err["error"]["code"], "bad_last_event_id");
6231 }
6232
6233 #[tokio::test]
6237 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6238 let app = test_app();
6239 let mut request = resumable_request();
6240 request["stream_resumable"] = serde_json::json!(false);
6241 let body = post_sse_raw(&app, request).await;
6242 let request_id = body
6243 .lines()
6244 .find_map(|l| l.strip_prefix("data: "))
6245 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6246 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6247 .unwrap();
6248
6249 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6250 assert_eq!(status, StatusCode::NOT_FOUND);
6251 assert_eq!(err["error"]["code"], "stream_not_found");
6252 assert!(err["error"]["message"]
6253 .as_str()
6254 .unwrap()
6255 .contains("stream_resumable"));
6256 }
6257
6258 #[test]
6262 fn the_axum_stream_patterns_match_the_published_templates() {
6263 assert_eq!(
6264 axum_path(frink_api::routes::V1_STREAM),
6265 "/v1/stream/:request_id"
6266 );
6267 assert_eq!(
6268 axum_path(frink_api::routes::V1_STREAM_POLL),
6269 "/v1/stream/:request_id/poll"
6270 );
6271 assert_eq!(
6272 frink_api::routes::v1_stream("abc"),
6273 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6274 );
6275 }
6276
6277 #[test]
6287 fn no_published_template_reaches_the_router_with_its_braces() {
6288 for template in [
6289 frink_api::routes::V1_STREAM,
6290 frink_api::routes::V1_STREAM_POLL,
6291 frink_api::routes::V1_RESPONSE,
6292 frink_api::routes::V1_RESPONSE_CANCEL,
6293 frink_api::routes::ADMIN_TASK_CANCEL,
6294 ] {
6295 assert!(
6296 template.contains('{'),
6297 "{template} is in the template list but has no placeholder"
6298 );
6299 let mounted = axum_path(template);
6300 assert!(
6301 !mounted.contains('{') && !mounted.contains('}'),
6302 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6303 );
6304 assert!(
6305 mounted.contains(':'),
6306 "{template} lost its placeholder entirely and would match one path only"
6307 );
6308 }
6309 }
6310
6311 #[tokio::test]
6319 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6320 let app = test_app();
6321 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6322 assert_eq!(status, StatusCode::NOT_FOUND);
6323 assert!(
6324 !body.is_null(),
6325 "empty body means axum never matched the route, so the id was read as a literal segment"
6326 );
6327 }
6328
6329 #[tokio::test]
6332 async fn the_task_list_starts_empty_rather_than_absent() {
6333 let app = test_app();
6334 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6335 assert_eq!(status, StatusCode::OK);
6336 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6337 }
6338
6339 #[tokio::test]
6348 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6349 assert!(
6350 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6351 "this test asserts the unconfigured behaviour"
6352 );
6353 let app = test_app();
6354 let (status, body) = post_json_uri(
6355 &app,
6356 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6357 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6358 )
6359 .await;
6360 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6361 assert!(
6362 body["error"]["message"]
6363 .as_str()
6364 .unwrap()
6365 .contains("--slot-save-path"),
6366 "{body}"
6367 );
6368 }
6369
6370 pub(crate) async fn post_json_uri(
6371 app: &Router,
6372 uri: &str,
6373 body: serde_json::Value,
6374 ) -> (StatusCode, serde_json::Value) {
6375 use http_body_util::BodyExt;
6376 use tower::ServiceExt;
6377
6378 let response = app
6379 .clone()
6380 .oneshot(
6381 axum::http::Request::builder()
6382 .method("POST")
6383 .uri(uri)
6384 .header("content-type", "application/json")
6385 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6386 .unwrap(),
6387 )
6388 .await
6389 .unwrap();
6390 let status = response.status();
6391 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6392 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6393 (status, json)
6394 }
6395
6396 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6397 post_json_uri(app, "/v1/chat/completions", body).await.1
6398 }
6399
6400 #[tokio::test]
6409 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6410 let app = test_app();
6411 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6412 assert_eq!(status, StatusCode::OK);
6413
6414 let memory = &body["memory"];
6415 if memory.is_null() {
6416 return;
6419 }
6420 assert!(
6421 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6422 "a read that produced a zero is a broken read, not an idle \
6423 engine: {memory}"
6424 );
6425 assert!(
6426 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6427 "the quantity must travel with the number: {memory}"
6428 );
6429 }
6430
6431 #[tokio::test]
6437 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6438 let app = test_app();
6439 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6440 assert_eq!(status, StatusCode::OK);
6441 for pool in ["kv_pages", "window_slots", "state_slots"] {
6442 assert!(
6443 body["pools"][pool].is_null(),
6444 "{pool} must be null rather than a zero row: {}",
6445 body["pools"]
6446 );
6447 }
6448 }
6449
6450 #[tokio::test]
6458 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6459 use http_body_util::BodyExt;
6460 use tower::ServiceExt;
6461
6462 let app = test_app();
6463 let send = |body: serde_json::Value| {
6464 let app = app.clone();
6465 async move {
6466 app.oneshot(
6467 axum::http::Request::builder()
6468 .method("POST")
6469 .uri(frink_api::routes::V1_MESSAGES)
6470 .header("content-type", "application/json")
6471 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6472 .unwrap(),
6473 )
6474 .await
6475 .unwrap()
6476 }
6477 };
6478
6479 let ok = send(serde_json::json!({
6480 "model": "test",
6481 "max_tokens": 1,
6482 "messages": [{"role": "user", "content": "hi"}],
6483 }))
6484 .await;
6485 assert_eq!(ok.status(), StatusCode::OK);
6486 let id = ok
6487 .headers()
6488 .get("request-id")
6489 .expect("a served message names its id")
6490 .to_str()
6491 .unwrap()
6492 .to_string();
6493 assert!(!id.is_empty());
6494
6495 let bad = send(serde_json::json!({"model": "test"})).await;
6498 assert!(bad.status().is_client_error());
6499 let other = bad.headers().get("request-id").expect("errors too");
6500 assert_ne!(other.to_str().unwrap(), id);
6501 let _ = bad.into_body().collect().await.unwrap();
6502 }
6503
6504 #[tokio::test]
6512 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6513 let state = Arc::new(test_state(
6514 test_model_full_byte_vocab(),
6515 ResponseCache::new(1000, Duration::from_secs(3600)),
6516 ));
6517 let app = test_app_with_state(Arc::clone(&state));
6518 let body = serde_json::json!({
6519 "model": "test",
6520 "messages": [{"role": "user", "content": "hi"}],
6521 "max_tokens": 1,
6522 });
6523
6524 state
6525 .maintenance
6526 .lock()
6527 .unwrap()
6528 .begin_rebuild()
6529 .expect("a fresh server is serving, so the rebuild starts");
6530 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6531 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6532 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6533
6534 state.maintenance.lock().unwrap().finish_rebuild(true);
6535 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6536 assert_eq!(
6537 status,
6538 StatusCode::OK,
6539 "the gate reopens; a rebuild is not a latch"
6540 );
6541 }
6542
6543 #[tokio::test]
6548 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6549 let app = test_app();
6550 let (status, body) = post_json_uri(
6551 &app,
6552 frink_api::routes::V1_CANCEL,
6553 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6554 )
6555 .await;
6556 assert_eq!(status, StatusCode::NOT_FOUND);
6557 assert_eq!(body["cancelled"], serde_json::json!(false));
6558 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6559 assert!(
6560 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6561 "the verdict must carry a human reason: {body}"
6562 );
6563 }
6564
6565 #[tokio::test]
6569 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6570 let state = Arc::new(test_state(
6571 test_model_full_byte_vocab(),
6572 ResponseCache::new(1000, Duration::from_secs(3600)),
6573 ));
6574 let app = test_app_with_state(Arc::clone(&state));
6575 let (token, _guard) = state.cancels.register("chatcmpl-live");
6576
6577 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6578 assert_eq!(status, StatusCode::OK);
6579 assert_eq!(before["generating_now"], serde_json::json!(1));
6580
6581 let (status, body) = post_json_uri(
6582 &app,
6583 frink_api::routes::V1_CANCEL,
6584 serde_json::json!({ "request_id": "chatcmpl-live" }),
6585 )
6586 .await;
6587 assert_eq!(status, StatusCode::OK);
6588 assert_eq!(body["cancelled"], serde_json::json!(true));
6589 assert!(
6590 token.is_cancelled(),
6591 "the endpoint answered ok without setting the flag the decode loop reads"
6592 );
6593 }
6594
6595 #[tokio::test]
6596 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6597 let app = test_app();
6598 let (status, tok) =
6599 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6600 assert_eq!(status, StatusCode::OK);
6601 let tokens = tok["tokens"].as_array().unwrap();
6602 assert_eq!(tok["count"], tokens.len());
6603 assert!(!tokens.is_empty());
6604
6605 let (status, detok) = post_json_uri(
6606 &app,
6607 "/v1/detokenize",
6608 serde_json::json!({ "tokens": tokens }),
6609 )
6610 .await;
6611 assert_eq!(status, StatusCode::OK);
6612 assert_eq!(detok["text"], "Hi");
6613
6614 let (status, emb) = post_json_uri(
6615 &app,
6616 "/v1/embeddings",
6617 serde_json::json!({
6618 "input": "Hi",
6619 "embedding_type": "mean"
6620 }),
6621 )
6622 .await;
6623 assert_eq!(status, StatusCode::OK);
6624 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6625 assert!(!vec.is_empty());
6626 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6627 }
6628
6629 #[tokio::test]
6635 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6636 let app = test_app();
6637 let (status, body) = post_json_uri(
6638 &app,
6639 "/v1/embeddings",
6640 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6641 )
6642 .await;
6643 assert_eq!(status, StatusCode::BAD_REQUEST);
6644 let msg = body["error"]["message"].as_str().unwrap();
6645 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6646 }
6647
6648 #[tokio::test]
6653 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6654 async fn a_real_embedding_model_serves_v1_embeddings() {
6655 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6656 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6657 if !path.exists() {
6658 eprintln!("SKIP: {} not present", path.display());
6659 return;
6660 }
6661 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6662 let mut state = test_state(
6663 test_model_full_byte_vocab(),
6664 ResponseCache::new(1000, Duration::from_secs(3600)),
6665 );
6666 state.embedding = Some(Arc::new(encoder));
6667 let app = test_app_with_state(Arc::new(state));
6668
6669 let (status, body) = post_json_uri(
6670 &app,
6671 "/v1/embeddings",
6672 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6673 )
6674 .await;
6675 assert_eq!(status, StatusCode::OK, "{body}");
6676 assert_eq!(body["model"], "bge-small-en-v1.5");
6677 let data = body["data"].as_array().unwrap();
6678 assert_eq!(data.len(), 2);
6679 for (i, row) in data.iter().enumerate() {
6680 assert_eq!(row["index"], i);
6681 let v: Vec<f64> = row["embedding"]
6682 .as_array()
6683 .unwrap()
6684 .iter()
6685 .map(|x| x.as_f64().unwrap())
6686 .collect();
6687 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6688 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6689 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6690 }
6691 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6694
6695 let (status, mean) = post_json_uri(
6699 &app,
6700 "/v1/embeddings",
6701 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6702 )
6703 .await;
6704 assert_eq!(status, StatusCode::OK);
6705 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6706 }
6707
6708 #[tokio::test]
6720 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6721 async fn an_encoder_can_be_the_loaded_model() {
6722 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6723 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6724 if !path.exists() {
6725 eprintln!("SKIP: {} not present", path.display());
6726 return;
6727 }
6728
6729 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6733 assert!(
6734 matches!(loaded, model::LoadedModel::Encoder(_)),
6735 "an encoder-only GGUF reached a decoder loader"
6736 );
6737 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6738 assert!(
6739 matches!(loaded, Loaded::Encoder(_)),
6740 "the encoder did not stay an encoder through activation"
6741 );
6742 assert!(
6743 batcher.is_none() && ceiling.is_none(),
6744 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6745 );
6746
6747 let state = test_state(
6748 test_model_full_byte_vocab(),
6749 ResponseCache::new(1000, Duration::from_secs(3600)),
6750 );
6751 state.swap_active(Some(Arc::new(ActiveModel {
6752 id: None,
6753 loaded,
6754 batcher,
6755 ceiling,
6756 checkpoint_path: None,
6757 })));
6758 let app = test_app_with_state(Arc::new(state));
6759
6760 let (status, body) = post_json_uri(
6762 &app,
6763 "/v1/embeddings",
6764 serde_json::json!({ "input": "Hello world" }),
6765 )
6766 .await;
6767 assert_eq!(status, StatusCode::OK, "{body}");
6768 assert_eq!(body["model"], "bge-small-en-v1.5");
6769 let v = body["data"][0]["embedding"].as_array().unwrap();
6770 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6771
6772 let (status, body) = post_json_uri(
6774 &app,
6775 "/v1/chat/completions",
6776 serde_json::json!({
6777 "model": "bge-small-en-v1.5",
6778 "messages": [{"role": "user", "content": "hi"}],
6779 }),
6780 )
6781 .await;
6782 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6783 let msg = body["error"]["message"].as_str().unwrap();
6784 for fact in [
6785 "bge-small-en-v1.5",
6786 "bert",
6787 "embedding model",
6788 "/v1/embeddings",
6789 ] {
6790 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6791 }
6792
6793 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
6795 assert_eq!(status, StatusCode::OK);
6796 let entry = &models["data"][0];
6797 assert_eq!(entry["id"], "bge-small-en-v1.5");
6798 assert_eq!(entry["frink_model_kind"], "embedding");
6799 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
6800 assert_eq!(entry["frink_n_embd"], 384);
6801 assert_eq!(entry["frink_pooling"], "CLS");
6802 assert_eq!(
6803 entry["frink_endpoints"],
6804 serde_json::json!(["/v1/embeddings"])
6805 );
6806 assert!(entry.get("supported_reasoning_efforts").is_none());
6809
6810 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
6812 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6813 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6814 assert_eq!(health["model"]["synthetic_weights"], false);
6815 let weights = health["capabilities"]
6816 .as_array()
6817 .unwrap()
6818 .iter()
6819 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
6820 .expect("a real-weights capability row");
6821 let detail = weights["detail"].as_str().unwrap_or_default();
6822 assert!(detail.contains("ENCODER"), "{detail}");
6823 let (status, body) = post_json_uri(
6831 &app,
6832 frink_api::routes::V1_TOKENIZE,
6833 serde_json::json!({ "content": "hello world" }),
6834 )
6835 .await;
6836 assert_eq!(
6837 status,
6838 StatusCode::OK,
6839 "an encoder has a real tokenizer: {body}"
6840 );
6841 let tokens = body["tokens"].as_array().expect("tokens array").clone();
6842 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
6843
6844 let (status, body) = post_json_uri(
6845 &app,
6846 frink_api::routes::V1_DETOKENIZE,
6847 serde_json::json!({ "tokens": tokens }),
6848 )
6849 .await;
6850 assert_eq!(status, StatusCode::OK, "{body}");
6851 let round_tripped = body["content"].as_str().expect("content").to_string();
6852 assert!(
6853 round_tripped.contains("hello") && round_tripped.contains("world"),
6854 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
6855 );
6856
6857 let (status, _) = post_json_uri(
6860 &app,
6861 "/v1/completions",
6862 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
6863 )
6864 .await;
6865 assert_eq!(
6866 status,
6867 StatusCode::NOT_IMPLEMENTED,
6868 "tokenizing an encoder must not have opened a path to generating with one"
6869 );
6870 }
6871
6872 #[tokio::test]
6877 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6878 use http_body_util::BodyExt;
6879 use tower::ServiceExt;
6880
6881 let fixture = concat!(
6882 "../frink-models/tests/fixtures/",
6883 "frink_real_moe_test.gguf"
6884 );
6885 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6886 let decoder = Decoder::from_gguf_with_expert_cache(
6887 &fixture,
6888 frink_models::config::test_moe_fixture(),
6889 Some(1024 * 1024),
6890 )
6891 .expect("MoE fixture must load store-backed");
6892
6893 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
6897 decoder.forward_token(1, 0, &mut caches);
6898
6899 let model = Model::Gguf(GgufModel {
6900 decoder: Arc::new(decoder),
6901 tokenizer: Arc::new(ServerTokenizer::Byte),
6902 stop_tokens: StopTokens::default(),
6903 bos_id: None,
6904 is_synthetic: false,
6905 chat_template: chat_template::PromptTemplate::plain(),
6906 });
6907 let state = Arc::new(test_state(
6908 model,
6909 ResponseCache::new(16, Duration::from_secs(60)),
6910 ));
6911 let app = Router::new()
6912 .route("/metrics", axum::routing::get(metrics))
6913 .route("/v1/chat/completions", post(chat_completions))
6914 .with_state(state);
6915
6916 let fetch_metrics = |app: Router| async move {
6917 let resp = app
6918 .oneshot(
6919 axum::http::Request::builder()
6920 .method("GET")
6921 .uri("/metrics")
6922 .body(axum::body::Body::empty())
6923 .unwrap(),
6924 )
6925 .await
6926 .unwrap();
6927 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
6928 String::from_utf8(bytes.to_vec()).unwrap()
6929 };
6930
6931 let after = fetch_metrics(app.clone()).await;
6932 assert!(
6933 after.contains("frink_expert_cache_misses_total"),
6934 "streaming model must expose expert-cache metrics: {after}"
6935 );
6936 let misses: u64 = after
6937 .lines()
6938 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
6939 .and_then(|l| l.split_whitespace().nth(1))
6940 .and_then(|v| v.parse().ok())
6941 .expect("misses metric line must parse");
6942 assert!(
6943 misses > 0,
6944 "decode must have read experts through the store: {after}"
6945 );
6946 }
6947
6948 fn weather_tool() -> serde_json::Value {
6949 serde_json::json!({
6950 "type": "function",
6951 "function": {
6952 "name": "get_weather",
6953 "description": "Get the current weather for a location.",
6954 "parameters": {
6955 "type": "object",
6956 "properties": {"location": {"type": "string"}},
6957 "required": ["location"]
6958 }
6959 }
6960 })
6961 }
6962
6963 fn weather_tool_def() -> ToolDef {
6964 ToolDef {
6965 kind: "function".to_string(),
6966 function: ToolFunctionDef {
6967 name: "get_weather".to_string(),
6968 description: Some("Get the current weather for a location.".to_string()),
6969 parameters: Some(serde_json::json!({
6970 "type": "object",
6971 "properties": {"location": {"type": "string"}},
6972 "required": ["location"]
6973 })),
6974 },
6975 }
6976 }
6977
6978 #[test]
6979 fn tool_preamble_mentions_every_tool_name_and_description() {
6980 let preamble = tool_preamble(&[weather_tool_def()]);
6981 assert!(preamble.contains("get_weather"));
6982 assert!(preamble.contains("Get the current weather for a location."));
6983 assert!(preamble.contains("<tool_call>"));
6984 assert!(preamble.contains("</tool_call>"));
6985 }
6986
6987 #[test]
6988 fn a_real_marker_becomes_a_structured_tool_call() {
6989 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
6990 let (message, finish) = build_response_message(
6991 text.to_string(),
6992 &[weather_tool_def()],
6993 output::OutputPosture::for_model("test-model"),
6994 "stop",
6995 );
6996 assert_eq!(finish, "tool_calls");
6997 let calls = message.tool_calls.expect("must carry a tool call");
6998 assert_eq!(calls[0].function.name, "get_weather");
6999 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7000 assert_eq!(parsed["location"], "Paris");
7001 }
7002
7003 #[test]
7004 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7005 let (message, finish) = build_response_message(
7006 "just an answer".to_string(),
7007 &[weather_tool_def()],
7008 output::OutputPosture::for_model("test-model"),
7009 "stop",
7010 );
7011 assert_eq!(finish, "stop");
7012 assert!(message.tool_calls.is_none());
7013 assert_eq!(message.content.as_deref(), Some("just an answer"));
7014 }
7015
7016 #[test]
7019 fn a_malformed_payload_is_not_a_tool_call() {
7020 let (message, finish) = build_response_message(
7021 "<tool_call>not valid json at all</tool_call>".to_string(),
7022 &[weather_tool_def()],
7023 output::OutputPosture::for_model("test-model"),
7024 "stop",
7025 );
7026 assert_eq!(finish, "stop");
7027 assert!(message.tool_calls.is_none());
7028 }
7029
7030 #[test]
7033 fn a_tool_that_was_never_offered_is_not_returned() {
7034 let (message, finish) = build_response_message(
7035 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7036 &[weather_tool_def()],
7037 output::OutputPosture::for_model("test-model"),
7038 "stop",
7039 );
7040 assert_eq!(finish, "stop");
7041 assert!(message.tool_calls.is_none());
7042 }
7043
7044 #[test]
7046 fn marker_text_with_no_tools_offered_stays_content() {
7047 let (message, finish) = build_response_message(
7048 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7049 &[],
7050 output::OutputPosture::for_model("test-model"),
7051 "stop",
7052 );
7053 assert_eq!(finish, "stop");
7054 assert!(message.tool_calls.is_none());
7055 assert!(message.content.is_some());
7056 }
7057
7058 #[test]
7062 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7063 let opened = std::cell::Cell::new(0usize);
7064 let mut parser = crate::policy::parser::ToolCallParser::new(
7065 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7066 vec![
7067 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7068 "write_file",
7069 serde_json::json!({"type": "object", "properties": {
7070 "path": {"type": "string"},
7071 "contents": {"type": "string"}
7072 }}),
7073 ),
7074 ],
7075 );
7076 let wire = "<tool_call><function=write_file>\
7077 <parameter=path>\n/tmp/x\n</parameter>\
7078 <parameter=contents>\nhello world\n</parameter>\
7079 </function></tool_call>";
7080
7081 let mut deltas = Vec::new();
7082 let mut text = String::new();
7083 for piece in wire.as_bytes().chunks(7) {
7084 let chunk = String::from_utf8_lossy(piece).into_owned();
7085 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7086 text.push_str(&more_text);
7087 deltas.extend(more);
7088 }
7089 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7090 text.push_str(&more_text);
7091 deltas.extend(more);
7092
7093 assert_eq!(opened.get(), 1, "one call opened");
7094 assert!(text.is_empty(), "the markers are not content: {text:?}");
7095
7096 let first = &deltas[0];
7097 assert_eq!(first.index, 0);
7098 assert_eq!(first.id.as_deref(), Some("call_0"));
7099 assert_eq!(first.kind, Some("function"));
7100 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7101
7102 let joined: String = deltas
7105 .iter()
7106 .filter_map(|d| d.function.arguments.clone())
7107 .collect();
7108 let parsed: serde_json::Value =
7109 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7110 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7111 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7112 assert!(
7113 deltas.len() >= 3,
7114 "the arguments arrived in pieces, not whole: {}",
7115 deltas.len()
7116 );
7117 assert!(
7118 deltas[1..].iter().all(|d| d.function.name.is_none()),
7119 "only the opening delta carries identity"
7120 );
7121 }
7122
7123 #[test]
7125 fn text_around_a_streamed_call_is_still_content() {
7126 let opened = std::cell::Cell::new(0usize);
7127 let mut parser = crate::policy::parser::ToolCallParser::new(
7128 crate::policy::parser::ToolCallFormat::Qwen25,
7129 vec![crate::policy::parser::tool_call::ToolSchema::new(
7130 "get_weather",
7131 )],
7132 );
7133 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7134 \"arguments\": {}}</tool_call> done";
7135 let mut text = String::new();
7136 for piece in wire.as_bytes().chunks(5) {
7137 let chunk = String::from_utf8_lossy(piece).into_owned();
7138 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7139 text.push_str(&more);
7140 }
7141 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7142 text.push_str(&more);
7143
7144 assert_eq!(opened.get(), 1);
7145 assert!(text.starts_with("let me check. "), "{text:?}");
7146 assert!(text.ends_with(" done"), "{text:?}");
7147 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7148 }
7149
7150 #[test]
7153 fn a_reasoning_block_is_split_out_of_the_answer() {
7154 let (message, finish) = build_response_message(
7155 "<think>weighing it up</think>The answer is 4.".to_string(),
7156 &[],
7157 output::OutputPosture::for_model("Qwen3-8B"),
7158 "stop",
7159 );
7160 assert_eq!(finish, "stop");
7161 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7162 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7163 }
7164
7165 #[test]
7168 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7169 let (message, _) = build_response_message(
7170 "Use the <think> tag like this.".to_string(),
7171 &[],
7172 output::OutputPosture::for_model("llama-3.1-8b"),
7173 "stop",
7174 );
7175 assert_eq!(
7176 message.content.as_deref(),
7177 Some("Use the <think> tag like this.")
7178 );
7179 assert!(message.reasoning_content.is_none());
7180 }
7181
7182 #[tokio::test]
7190 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7191 let app = test_app();
7192 let body = serde_json::json!({
7193 "model": "m",
7194 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7195 "max_tokens": 4,
7196 "temperature": 0,
7197 });
7198 let resp = post_json(&app, body).await;
7199 let message = &resp["choices"][0]["message"];
7200 assert!(message["content"].is_string());
7201 assert!(message.get("tool_calls").is_none());
7202 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7203 assert_eq!(resp["usage"]["completion_tokens"], 4);
7204 assert_eq!(
7205 resp["usage"]["total_tokens"],
7206 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7207 );
7208 }
7209
7210 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7211 use http_body_util::BodyExt;
7212 use tower::ServiceExt;
7213
7214 let response = app
7215 .clone()
7216 .oneshot(
7217 axum::http::Request::builder()
7218 .method("GET")
7219 .uri(uri)
7220 .body(axum::body::Body::empty())
7221 .unwrap(),
7222 )
7223 .await
7224 .unwrap();
7225 let status = response.status();
7226 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7227 (status, serde_json::from_slice(&bytes).unwrap())
7228 }
7229
7230 #[tokio::test]
7231 async fn health_answers_a_capability_handshake_not_a_boolean() {
7232 let app = test_app();
7233 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7234 assert_eq!(status, StatusCode::OK);
7235
7236 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7237 assert_eq!(health.state, frink_api::HealthState::Ready);
7238 assert!(health.pid > 0);
7239 assert!(health.server_time_unix_ms > 0);
7240 assert_eq!(health.last_request_age_seconds, None);
7243
7244 for id in [
7247 frink_api::health::capability::CPU,
7248 frink_api::health::capability::METAL,
7249 frink_api::health::capability::CUDA,
7250 frink_api::health::capability::REAL_WEIGHTS,
7251 frink_api::health::capability::CONTINUOUS_BATCHING,
7252 ] {
7253 let cap = health
7254 .capability(id)
7255 .unwrap_or_else(|| panic!("{id} missing"));
7256 assert!(!cap.reason.is_empty(), "{cap:?}");
7257 assert!(!cap.detail.is_empty(), "{cap:?}");
7258 }
7259 let weights = health
7263 .capability(frink_api::health::capability::REAL_WEIGHTS)
7264 .unwrap();
7265 assert!(!weights.available);
7266 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7267 assert!(health.model.as_ref().unwrap().synthetic_weights);
7268 }
7269
7270 #[tokio::test]
7271 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7272 let app = test_app();
7273 let _ = post_json(
7274 &app,
7275 serde_json::json!({
7276 "model": "m",
7277 "messages": [{"role": "user", "content": "\u{1}"}],
7278 "max_tokens": 1,
7279 "temperature": 0,
7280 }),
7281 )
7282 .await;
7283 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7284 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7285 let age = health
7286 .last_request_age_seconds
7287 .expect("a served request is evidence of liveness");
7288 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7289 }
7290
7291 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7293 use http_body_util::BodyExt;
7294 use tower::ServiceExt;
7295
7296 let response = app
7297 .clone()
7298 .oneshot(
7299 axum::http::Request::builder()
7300 .method("POST")
7301 .uri("/v1/chat/completions")
7302 .header("content-type", "application/json")
7303 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7304 .unwrap(),
7305 )
7306 .await
7307 .unwrap();
7308 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7309 String::from_utf8(bytes.to_vec())
7310 .unwrap()
7311 .lines()
7312 .filter_map(|line| line.strip_prefix("data: "))
7313 .filter(|payload| *payload != "[DONE]")
7314 .map(|payload| serde_json::from_str(payload).unwrap())
7315 .collect()
7316 }
7317
7318 #[tokio::test]
7319 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7320 let app = test_app();
7321 let chunks = post_sse_chunks(
7322 &app,
7323 serde_json::json!({
7324 "model": "m",
7325 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7326 "max_tokens": 4,
7327 "temperature": 0,
7328 "stream": true,
7329 }),
7330 )
7331 .await;
7332
7333 assert!(!chunks.is_empty());
7334 let request_id = chunks[0]["request_id"]
7335 .as_str()
7336 .expect("the first chunk names the request")
7337 .to_string();
7338 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7339 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7342 assert!(
7343 chunk.get("request_id").is_none(),
7344 "chunk {i} repeats request_id"
7345 );
7346 }
7347 for chunk in &chunks {
7350 assert_eq!(chunk["id"], serde_json::json!(request_id));
7351 }
7352
7353 let other = post_sse_chunks(
7354 &app,
7355 serde_json::json!({
7356 "model": "m",
7357 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7358 "max_tokens": 4,
7359 "temperature": 0,
7360 "stream": true,
7361 }),
7362 )
7363 .await;
7364 assert_ne!(
7365 other[0]["request_id"].as_str().unwrap(),
7366 request_id,
7367 "two concurrent chats must not share an id"
7368 );
7369 }
7370
7371 #[tokio::test]
7372 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7373 let app = test_app();
7374 let resp = post_json(
7375 &app,
7376 serde_json::json!({
7377 "model": "m",
7378 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7379 "max_tokens": 2,
7380 "temperature": 0,
7381 }),
7382 )
7383 .await;
7384 assert_eq!(resp["id"], resp["request_id"]);
7385 assert!(resp["request_id"]
7386 .as_str()
7387 .unwrap()
7388 .starts_with("chatcmpl-"));
7389 }
7390
7391 #[tokio::test]
7394 async fn usage_carries_separate_prefill_and_decode_timings() {
7395 let app = test_app();
7396 let resp = post_json(
7397 &app,
7398 serde_json::json!({
7399 "model": "m",
7400 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7401 "max_tokens": 4,
7402 "temperature": 0,
7403 }),
7404 )
7405 .await;
7406 let usage = &resp["usage"];
7407 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7408 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7409 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7410 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7411 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7413 }
7414
7415 #[tokio::test]
7423 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7424 let app = test_app();
7425 let body = serde_json::json!({
7426 "model": "m",
7427 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7428 "max_tokens": 4,
7429 "temperature": 0,
7430 "tools": [weather_tool()],
7431 });
7432 let resp = post_json(&app, body).await;
7433 let message = &resp["choices"][0]["message"];
7434 assert!(
7435 message["content"].is_string(),
7436 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7437 );
7438 assert!(message.get("tool_calls").is_none());
7439 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7443 }
7444
7445 #[tokio::test]
7449 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7450 let app = test_app();
7451 let body = serde_json::json!({
7452 "model": "m",
7453 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7454 "max_tokens": 3,
7455 "temperature": 0,
7456 });
7457 let first = post_json(&app, body.clone()).await;
7458 assert_eq!(first["frink_cache"], "miss");
7459 let second = post_json(&app, body).await;
7460 assert_eq!(second["frink_cache"], "hit");
7461 assert_eq!(
7462 first["choices"][0]["message"]["content"],
7463 second["choices"][0]["message"]["content"]
7464 );
7465 assert_eq!(
7466 first["choices"][0]["finish_reason"],
7467 second["choices"][0]["finish_reason"]
7468 );
7469 assert_eq!(first["usage"], second["usage"]);
7470 assert_eq!(second["usage"]["completion_tokens"], 3);
7471 }
7472
7473 #[tokio::test]
7484 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7485 let app = test_app();
7486 let plain = serde_json::json!({
7487 "model": "m",
7488 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7489 "max_tokens": 3,
7490 "temperature": 0,
7491 });
7492
7493 let first = post_json(&app, plain.clone()).await;
7494 assert_eq!(first["frink_cache"], "miss");
7495 let unconstrained = first["choices"][0]["message"]["content"]
7496 .as_str()
7497 .expect("content")
7498 .to_string();
7499
7500 let mut constrained = plain.clone();
7501 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7502 let second = post_json(&app, constrained).await;
7503 assert_eq!(
7504 second["frink_cache"], "miss",
7505 "a grammar is part of the key, so this body has never been answered"
7506 );
7507 let constrained_answer = second["choices"][0]["message"]["content"]
7511 .as_str()
7512 .expect("content")
7513 .to_string();
7514 assert!(
7515 constrained_answer.contains("-> \"yes\"]"),
7516 "the grammar must have been compiled AND applied, not skipped \
7517 by a cache hit: {constrained_answer}"
7518 );
7519 assert_ne!(
7520 constrained_answer, unconstrained,
7521 "the constrained request was served the unconstrained answer"
7522 );
7523
7524 let third = post_json(&app, plain).await;
7528 assert_eq!(third["frink_cache"], "hit");
7529 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7530 }
7531
7532 #[tokio::test]
7554 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7555 let state = Arc::new(test_state(
7556 test_model_full_byte_vocab(),
7557 ResponseCache::new(1000, Duration::from_secs(3600)),
7558 ));
7559 let app = test_app_with_state(state.clone());
7560 let plain = serde_json::json!({
7561 "model": "m",
7562 "messages": [
7563 {"role": "system", "content": "Answer in JSON when it helps."},
7564 {"role": "user", "content": "\u{1}\u{2}"},
7565 ],
7566 "max_tokens": 3,
7567 "temperature": 0,
7568 });
7569
7570 let first = post_json(&app, plain.clone()).await;
7571 assert_eq!(first["frink_cache"], "miss");
7572 assert_eq!(state.cache_stats().entries, 1);
7573
7574 let mut as_json = plain.clone();
7575 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7576 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7577 assert_eq!(
7578 status,
7579 StatusCode::BAD_REQUEST,
7580 "the demo banner is not a JSON object, whoever generated it"
7581 );
7582 assert_eq!(
7583 state.cache_stats().hits,
7584 0,
7585 "a json_object request must not be answered from an entry the \
7586 JSON mask never produced"
7587 );
7588 assert_eq!(
7589 state.cache_stats().entries,
7590 2,
7591 "json_object must key its own entry, not reuse the unconstrained \
7592 one it happens to render the same prompt as"
7593 );
7594 }
7595
7596 #[tokio::test]
7608 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7609 let app = test_app_with_state(Arc::new(test_state(
7610 test_model_full_byte_vocab_with_eos(Some(0x77)),
7611 ResponseCache::new(1000, Duration::from_secs(3600)),
7612 )));
7613 let body = serde_json::json!({
7614 "model": "m",
7615 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7616 "max_tokens": 6,
7617 "temperature": 0,
7618 });
7619
7620 let stopped = post_json(&app, body.clone()).await;
7621 assert_eq!(stopped["frink_cache"], "miss");
7622 assert_eq!(
7623 stopped["choices"][0]["finish_reason"], "stop",
7624 "the fixture is only meaningful if the model's EOS really fires here"
7625 );
7626 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7627
7628 let mut ignoring = body.clone();
7629 ignoring["ignore_eos"] = serde_json::json!(true);
7630 let ran_on = post_json(&app, ignoring).await;
7631 assert_eq!(
7632 ran_on["frink_cache"], "miss",
7633 "ignore_eos is part of the key, so this body has never been answered"
7634 );
7635 assert_eq!(
7636 ran_on["usage"]["completion_tokens"], 6,
7637 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7638 );
7639 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7640 assert_ne!(
7641 ran_on["choices"][0]["message"]["content"],
7642 stopped["choices"][0]["message"]["content"]
7643 );
7644 }
7645
7646 #[tokio::test]
7652 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7653 let session_app = test_app();
7654 let manual_app = test_app();
7655
7656 let turn1 = post_json(
7658 &session_app,
7659 serde_json::json!({
7660 "model": "m",
7661 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7662 "session_id": "s1",
7663 "max_tokens": 5,
7664 "temperature": 0,
7665 }),
7666 )
7667 .await;
7668 let reply1 = turn1["choices"][0]["message"]["content"]
7669 .as_str()
7670 .unwrap()
7671 .to_string();
7672
7673 let manual_turn1 = post_json(
7678 &manual_app,
7679 serde_json::json!({
7680 "model": "m",
7681 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7682 "max_tokens": 5,
7683 "temperature": 0,
7684 }),
7685 )
7686 .await;
7687 assert_eq!(
7688 manual_turn1["choices"][0]["message"]["content"]
7689 .as_str()
7690 .unwrap(),
7691 reply1
7692 );
7693
7694 let turn2 = post_json(
7696 &session_app,
7697 serde_json::json!({
7698 "model": "m",
7699 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7700 "session_id": "s1",
7701 "max_tokens": 5,
7702 "temperature": 0,
7703 }),
7704 )
7705 .await;
7706 let reply2 = turn2["choices"][0]["message"]["content"]
7707 .as_str()
7708 .unwrap()
7709 .to_string();
7710
7711 let manual_turn2 = post_json(
7715 &manual_app,
7716 serde_json::json!({
7717 "model": "m",
7718 "messages": [
7719 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7720 {"role": "assistant", "content": reply1},
7721 {"role": "user", "content": "\u{4}\u{5}"},
7722 ],
7723 "max_tokens": 5,
7724 "temperature": 0,
7725 }),
7726 )
7727 .await;
7728 assert_eq!(
7729 manual_turn2["choices"][0]["message"]["content"]
7730 .as_str()
7731 .unwrap(),
7732 reply2,
7733 "resuming a session must produce identical output to manually resending the full history"
7734 );
7735 }
7736
7737 #[test]
7740 fn lock_cache_recovers_from_a_poisoned_mutex() {
7741 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7742
7743 let poison_cache = Arc::clone(&cache);
7744 let _ = std::thread::spawn(move || {
7745 let _guard = poison_cache.lock().unwrap();
7746 panic!("simulated panic while holding the lock");
7747 })
7748 .join();
7749
7750 let recovered = lock_cache(&cache);
7752 assert_eq!(recovered.stats().entries, 0);
7753 }
7754
7755 #[test]
7756 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7757 let mut req_body = serde_json::json!({
7758 "model": "m",
7759 "messages": [{"role": "user", "content": "hi"}],
7760 });
7761 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7762 assert!(
7763 req.is_cacheable(),
7764 "default (temperature 0) must be cacheable"
7765 );
7766
7767 req_body["temperature"] = serde_json::json!(0.8);
7768 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7769 assert!(
7770 !req.is_cacheable(),
7771 "unseeded sampling must never be cacheable"
7772 );
7773
7774 req_body["seed"] = serde_json::json!(42);
7775 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7776 assert!(
7777 req.is_cacheable(),
7778 "sampling with an explicit seed is deterministic and must be cacheable"
7779 );
7780 }
7781
7782 const GRADED: &str = "{% if reasoning_effort %}\
7786 {% if reasoning_effort not in ['low','medium','high'] %}\
7787 {{ raise_exception('unsupported effort') }}\
7788 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7789 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7790
7791 fn graded_template() -> chat_template::PromptTemplate {
7792 chat_template::PromptTemplate::from_gguf_metadata(
7793 Some(GRADED),
7794 Some("qwen3"),
7795 false,
7796 true,
7797 None,
7798 None,
7799 )
7800 }
7801
7802 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7803 serde_json::from_value(value).expect("request")
7804 }
7805
7806 #[test]
7813 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7814 let req = chat_request(serde_json::json!({
7815 "model": "m",
7816 "messages": [{"role": "user", "content": "hi"}],
7817 "grammar": "root ::= \"a\"+",
7818 }));
7819 req.validate_supported_fields()
7820 .expect("a valid grammar is a valid request");
7821 let params = req
7822 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7823 .expect("a valid grammar compiles at params time too");
7824 assert!(
7825 params.grammar.is_some(),
7826 "the grammar was dropped between the wire and the sampler"
7827 );
7828 assert!(
7829 params.needs_vocab_logits(),
7830 "a grammar request that may fold lm_head into a GPU argmax is \
7831 a grammar request served unconstrained"
7832 );
7833
7834 let plain = chat_request(serde_json::json!({
7835 "model": "m",
7836 "messages": [{"role": "user", "content": "hi"}],
7837 }));
7838 assert!(plain
7839 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7840 .unwrap()
7841 .grammar
7842 .is_none());
7843 }
7844
7845 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7846 chat_request(serde_json::json!({
7847 "model": "m",
7848 "messages": [{"role": "user", "content": "weather in Rome?"}],
7849 "tools": [weather_tool()],
7850 "tool_choice": tool_choice,
7851 }))
7852 }
7853
7854 #[test]
7858 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7859 for choice in [
7860 serde_json::json!("required"),
7861 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7862 ] {
7863 let req = tool_request(choice.clone());
7864 req.validate_supported_fields()
7865 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7866 let params = req
7867 .generation_params_for_template(
7868 &graded_template(),
7869 "Qwen3-8B",
7870 crate::sampling_knobs::SamplerModel::absent(),
7871 )
7872 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7873 let grammar = params
7874 .grammar
7875 .as_ref()
7876 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7877 assert!(
7878 grammar.is_awaiting_trigger(),
7879 "the model must be free to think before it calls"
7880 );
7881 assert!(
7882 !grammar.allows_eog(),
7883 "{choice} must not be able to end the turn without a call"
7884 );
7885 assert!(
7891 params.needs_vocab_logits(),
7892 "{choice} would let a backend return a token id instead of logits"
7893 );
7894 assert!(
7895 !generate::greedy_gpu_fold_allowed(¶ms),
7896 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7897 );
7898 }
7899 }
7900
7901 #[test]
7903 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7904 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7905 let req = tool_request(choice.clone());
7906 req.validate_supported_fields().expect("still supported");
7907 let params = match req.generation_params_for_template(
7908 &graded_template(),
7909 "Qwen3-8B",
7910 crate::sampling_knobs::SamplerModel::absent(),
7911 ) {
7912 Ok(p) => p,
7913 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7914 };
7915 assert!(
7916 params.grammar.is_none(),
7917 "{choice} does not force a call and must not be constrained"
7918 );
7919 }
7920 }
7921
7922 #[test]
7925 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
7926 let req = chat_request(serde_json::json!({
7928 "model": "m",
7929 "messages": [{"role": "user", "content": "hi"}],
7930 "tool_choice": "required",
7931 }));
7932 let (status, _) = req
7933 .validate_supported_fields()
7934 .expect_err("nothing to call");
7935 assert_eq!(status, StatusCode::BAD_REQUEST);
7936
7937 let req =
7939 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
7940 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
7941 assert_eq!(status, StatusCode::BAD_REQUEST);
7942 assert_eq!(body["error"]["param"], "tool_choice");
7943
7944 let req = tool_request(serde_json::json!({"type": "function"}));
7946 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
7947 assert_eq!(status, StatusCode::BAD_REQUEST);
7948
7949 let req = chat_request(serde_json::json!({
7951 "model": "m",
7952 "messages": [{"role": "user", "content": "hi"}],
7953 "tools": [weather_tool()],
7954 "tool_choice": "required",
7955 "grammar": "root ::= \"a\"+",
7956 }));
7957 let (status, _) = req
7958 .validate_supported_fields()
7959 .expect_err("a grammar and a forced call are two constraints");
7960 assert_eq!(status, StatusCode::BAD_REQUEST);
7961
7962 let req = tool_request(serde_json::json!("required"));
7968 let (status, Json(body)) = match req.generation_params_for_template(
7969 &graded_template(),
7970 "muse-glimmer-8b",
7971 crate::sampling_knobs::SamplerModel::absent(),
7972 ) {
7973 Err(e) => e,
7974 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
7975 };
7976 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7977 assert!(
7978 body["error"]["message"]
7979 .as_str()
7980 .unwrap()
7981 .contains("muse_glimmer"),
7982 "{body}"
7983 );
7984
7985 let req = tool_request(serde_json::json!("required"));
7991 let params = req
7992 .generation_params_for_template(
7993 &graded_template(),
7994 "gemma-4-E2B-it",
7995 crate::sampling_knobs::SamplerModel::absent(),
7996 )
7997 .expect("a gemma4 forced tool_choice is served");
7998 assert!(
7999 params.grammar.is_some(),
8000 "a forced tool_choice must arrive as the generation's grammar"
8001 );
8002 }
8003
8004 #[test]
8007 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8008 let req = chat_request(serde_json::json!({
8009 "model": "m",
8010 "messages": [{"role": "user", "content": "hi"}],
8011 "grammar": "root ::= \"a",
8012 }));
8013 let (status, Json(body)) = req
8014 .validate_supported_fields()
8015 .expect_err("this does not parse");
8016 assert_eq!(status, StatusCode::BAD_REQUEST);
8017 assert_eq!(body["error"]["param"], "grammar");
8018 assert!(
8019 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8020 .is_err(),
8021 "and again at params time"
8022 );
8023 }
8024
8025 #[test]
8032 fn response_format_json_schema_becomes_the_requests_grammar() {
8033 let req = chat_request(serde_json::json!({
8034 "model": "m",
8035 "messages": [{"role": "user", "content": "hi"}],
8036 "response_format": {
8037 "type": "json_schema",
8038 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8039 },
8040 }));
8041 req.validate_supported_fields()
8042 .expect("a boolean schema converts");
8043 let params = req
8044 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8045 .expect("and compiles");
8046 let grammar = params.grammar.expect("the schema is the grammar");
8047 let mut g = (*grammar).clone();
8048 g.accept_token(0, b"true").expect("a boolean is accepted");
8049 assert!(g.allows_eog(), "and completes the parse");
8050 assert!(
8051 !params.json_object,
8052 "a schema is not the json_object character-class mask"
8053 );
8054 }
8055
8056 #[test]
8060 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8061 let req = chat_request(serde_json::json!({
8062 "model": "m",
8063 "messages": [{"role": "user", "content": "hi"}],
8064 "response_format": {
8065 "type": "json_schema",
8066 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8067 },
8068 }));
8069 let (status, Json(body)) = req
8070 .validate_supported_fields()
8071 .expect_err("minimum has no grammar in this port");
8072 assert_eq!(status, StatusCode::BAD_REQUEST);
8073 assert!(
8074 body["error"]["message"]
8075 .as_str()
8076 .expect("a message")
8077 .contains("minimum"),
8078 "the refusal must name the keyword: {body}"
8079 );
8080 assert!(
8081 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8082 .is_err(),
8083 "and again at params time"
8084 );
8085 }
8086
8087 #[test]
8093 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8094 let req = chat_request(serde_json::json!({
8095 "model": "m",
8096 "messages": [{"role": "user", "content": "hi"}],
8097 "tool_choice": "required",
8098 "tools": [{
8099 "type": "function",
8100 "function": {"name": "f", "parameters": {"type": "object"}},
8101 }],
8102 "response_format": {
8103 "type": "json_schema",
8104 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8105 },
8106 }));
8107 let (status, Json(body)) = req
8108 .validate_supported_fields()
8109 .expect_err("two constraints, one generation");
8110 assert_eq!(status, StatusCode::BAD_REQUEST);
8111 assert_eq!(body["error"]["param"], "tool_choice");
8112 }
8113
8114 #[test]
8117 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8118 let req = chat_request(serde_json::json!({
8119 "model": "m",
8120 "messages": [{"role": "user", "content": "hi"}],
8121 }));
8122 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8123 }
8124
8125 #[test]
8130 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8131 let asked = chat_request(serde_json::json!({
8132 "model": "m",
8133 "messages": [{"role": "user", "content": "hi"}],
8134 "min_p": 0.07,
8135 }));
8136 assert_eq!(
8137 asked
8138 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8139 .expect("knobs")
8140 .min_p,
8141 0.07
8142 );
8143
8144 let silent = chat_request(serde_json::json!({
8145 "model": "m",
8146 "messages": [{"role": "user", "content": "hi"}],
8147 }));
8148 assert_eq!(
8149 silent
8150 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8151 .expect("knobs")
8152 .min_p,
8153 0.0,
8154 "an unset min_p must be off, not llama.cpp's CLI default"
8155 );
8156 }
8157
8158 #[test]
8167 fn no_sampler_knob_is_missing_from_the_cache_key() {
8168 let base = serde_json::json!({
8169 "model": "m",
8170 "messages": [{"role": "user", "content": "hi"}],
8171 "seed": 1,
8172 });
8173 let key_for = |body: serde_json::Value| {
8174 let req = chat_request(body);
8175 let params = req
8176 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8177 .expect("params");
8178 req.cache_key("prompt", ¶ms)
8179 };
8180 let baseline = key_for(base.clone());
8181 for (knob, value) in [
8182 ("temperature", serde_json::json!(0.5)),
8183 ("top_p", serde_json::json!(0.9)),
8184 ("min_p", serde_json::json!(0.05)),
8185 ("top_k", serde_json::json!(40)),
8186 ("repetition_penalty", serde_json::json!(1.1)),
8187 ("presence_penalty", serde_json::json!(0.3)),
8188 ("frequency_penalty", serde_json::json!(0.3)),
8189 (
8190 "samplers",
8191 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8192 ),
8193 ] {
8194 let mut body = base.clone();
8195 body[knob] = value;
8196 assert_ne!(
8197 key_for(body),
8198 baseline,
8199 "`{knob}` is not in the cache key: two requests differing \
8200 only in it would share one cached answer"
8201 );
8202 }
8203 }
8204
8205 #[test]
8215 fn no_constraint_is_missing_from_the_cache_key() {
8216 let base = serde_json::json!({
8217 "model": "m",
8218 "messages": [{"role": "user", "content": "pick one"}],
8219 });
8220 let key_for = |body: serde_json::Value| {
8221 let req = chat_request(body);
8222 let params = req
8223 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8224 .expect("params");
8225 req.cache_key("prompt", ¶ms)
8226 };
8227 let baseline = key_for(base.clone());
8228 for (field, value) in [
8229 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8230 (
8231 "response_format",
8232 serde_json::json!({"type": "json_object"}),
8233 ),
8234 (
8235 "response_format",
8236 serde_json::json!({"type": "json_schema", "json_schema": {
8237 "name": "answer",
8238 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8239 }}),
8240 ),
8241 ("ignore_eos", serde_json::json!(true)),
8242 ("stop", serde_json::json!(["\n"])),
8243 ("max_tokens", serde_json::json!(7)),
8244 ] {
8245 let mut body = base.clone();
8246 body[field] = value.clone();
8247 assert_ne!(
8248 key_for(body),
8249 baseline,
8250 "`{field}: {value}` is not in the cache key: two requests \
8251 differing only in it would share one cached answer"
8252 );
8253 }
8254 }
8255
8256 #[test]
8260 fn an_explicit_zero_output_budget_is_a_client_error() {
8261 let req = chat_request(serde_json::json!({
8262 "model": "m",
8263 "messages": [{"role": "user", "content": "hi"}],
8264 "max_tokens": 0,
8265 }));
8266 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8267 assert_eq!(status, StatusCode::BAD_REQUEST);
8268 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8269 }
8270
8271 #[test]
8274 fn a_request_can_turn_thinking_off() {
8275 let template = graded_template();
8276 for body in [
8277 serde_json::json!({
8278 "model": "m",
8279 "messages": [{"role": "user", "content": "hi"}],
8280 "reasoning_effort": "none",
8281 }),
8282 serde_json::json!({
8283 "model": "m",
8284 "messages": [{"role": "user", "content": "hi"}],
8285 "thinking": {"type": "disabled"},
8286 }),
8287 ] {
8288 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8289 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8290 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8291 assert!(!kwargs.contains_key("reasoning_effort"));
8294 }
8295 }
8296
8297 #[test]
8300 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8301 let template = graded_template();
8302 let kwargs = chat_request(serde_json::json!({
8303 "model": "m",
8304 "messages": [{"role": "user", "content": "hi"}],
8305 "reasoning_effort": "high",
8306 "thinking": {"type": "disabled"},
8307 }))
8308 .resolve_template_kwargs(&template);
8309 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8310 assert!(!kwargs.contains_key("reasoning_effort"));
8311 }
8312
8313 #[test]
8316 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8317 let req = chat_request(serde_json::json!({
8318 "model": "m",
8319 "messages": [{"role": "user", "content": "hi"}],
8320 "thinking": {"type": "disable"},
8321 }));
8322 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8323 assert_eq!(status, StatusCode::BAD_REQUEST);
8324 }
8325
8326 #[test]
8329 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8330 let template = graded_template();
8331 let kwargs = chat_request(serde_json::json!({
8332 "model": "m",
8333 "messages": [{"role": "user", "content": "hi"}],
8334 "reasoning_effort": "none",
8335 "chat_template_kwargs": {"enable_thinking": true},
8336 }))
8337 .resolve_template_kwargs(&template);
8338 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8339 }
8340
8341 #[test]
8345 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8346 let template = graded_template();
8347 let req = chat_request(serde_json::json!({
8348 "model": "m",
8349 "messages": [{"role": "user", "content": "hi"}],
8350 "reasoning_effort": "minimal",
8351 }));
8352 let kwargs = req.resolve_template_kwargs(&template);
8353 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8354 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8355 assert!(prompt.starts_with("E:low|"), "{prompt}");
8356 }
8357
8358 #[test]
8362 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8363 let template = graded_template();
8364 let req = chat_request(serde_json::json!({
8365 "model": "m",
8366 "messages": [{"role": "user", "content": "hi"}],
8367 "chat_template_kwargs": {"reasoning_effort": "none"},
8368 }));
8369 let kwargs = req.resolve_template_kwargs(&template);
8370 assert!(!kwargs.contains_key("reasoning_effort"));
8371 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8372 assert_eq!(prompt, "hi");
8373 }
8374
8375 #[test]
8378 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8379 let template = graded_template();
8380 let req = chat_request(serde_json::json!({
8381 "model": "m",
8382 "messages": [{"role": "user", "content": "hi"}],
8383 "reasoning_effort": "low",
8384 "chat_template_kwargs": {"reasoning_effort": "high"},
8385 }));
8386 assert_eq!(
8387 req.resolve_template_kwargs(&template)["reasoning_effort"],
8388 serde_json::json!("high")
8389 );
8390 }
8391
8392 #[test]
8396 fn offering_tools_turns_thinking_on_by_itself() {
8397 let template = graded_template();
8398 let quiet = chat_request(serde_json::json!({
8399 "model": "m",
8400 "messages": [{"role": "user", "content": "hi"}],
8401 }));
8402 assert!(!quiet
8403 .resolve_template_kwargs(&template)
8404 .contains_key("enable_thinking"));
8405
8406 let with_tools = chat_request(serde_json::json!({
8407 "model": "m",
8408 "messages": [{"role": "user", "content": "hi"}],
8409 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8410 }));
8411 let kwargs = with_tools.resolve_template_kwargs(&template);
8412 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8413 let prompt =
8414 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8415 assert!(prompt.starts_with("THINK|"), "{prompt}");
8416 }
8417
8418 #[test]
8423 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8424 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8425 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8426 Some("qwen3"),
8427 false,
8428 true,
8429 None,
8430 None,
8431 );
8432 let req = chat_request(serde_json::json!({
8433 "model": "m",
8434 "messages": [{"role": "user", "content": "hi"}],
8435 "chat_template_kwargs": {"enable_thinking": true},
8436 }));
8437 let kwargs = req.resolve_template_kwargs(&opener);
8438 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8439 assert!(prompt.ends_with("<think>"), "{prompt}");
8440
8441 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8444 let (message, _) = build_response_message(
8445 "weighing it up</think>Paris.".to_string(),
8446 &[],
8447 posture,
8448 "stop",
8449 );
8450 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8451 assert_eq!(message.content.as_deref(), Some("Paris."));
8452
8453 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8456 let (message, _) = build_response_message(
8457 "weighing it up</think>Paris.".to_string(),
8458 &[],
8459 closed,
8460 "stop",
8461 );
8462 assert_eq!(message.reasoning_content, None);
8463 }
8464
8465 #[test]
8466 fn stop_param_accepts_both_single_string_and_array() {
8467 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8468 "model": "m",
8469 "messages": [{"role": "user", "content": "hi"}],
8470 "stop": "END",
8471 }))
8472 .unwrap();
8473 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8474
8475 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8476 "model": "m",
8477 "messages": [{"role": "user", "content": "hi"}],
8478 "stop": ["A", "B"],
8479 }))
8480 .unwrap();
8481 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8482 }
8483
8484 #[test]
8485 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8486 let model = test_model();
8487 let result = run_generation(
8488 &model,
8489 "hello",
8490 &greedy_params(4),
8491 None,
8492 None,
8493 None,
8494 None,
8495 None,
8496 None,
8497 );
8498 assert!(matches!(
8499 result,
8500 Err(generate::DecodeError::TokenOutOfVocab { .. })
8501 ));
8502 }
8503
8504 #[test]
8508 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8509 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8511 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8512
8513 let holder_pool = Arc::clone(&pool);
8514 let holder = std::thread::spawn(move || {
8515 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8516 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8518 drop(held);
8519 });
8520 std::thread::sleep(Duration::from_millis(15));
8521
8522 let config = generate::KvPoolConfig {
8523 pool,
8524 queue_wait: Duration::ZERO,
8525 };
8526 let result = run_generation(
8527 &model,
8528 &prompt,
8529 &greedy_params(4),
8530 Some(&config),
8531 None,
8532 None,
8533 None,
8534 None,
8535 None,
8536 );
8537 assert!(matches!(
8538 result,
8539 Err(generate::DecodeError::KvPoolExhausted)
8540 ));
8541
8542 let (status, _body) = decode_error_response(result.unwrap_err());
8543 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8544 holder.join().unwrap();
8545 }
8546
8547 #[test]
8555 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8556 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8558 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8560 let config = generate::KvPoolConfig {
8561 pool,
8562 queue_wait: Duration::ZERO,
8563 };
8564
8565 let result = run_generation(
8566 &model,
8567 &prompt,
8568 &greedy_params(4),
8569 Some(&config),
8570 None,
8571 None,
8572 None,
8573 None,
8574 None,
8575 );
8576 let err = result.expect_err("one block cannot hold two layers' caches");
8577 assert!(
8578 matches!(
8579 &err,
8580 generate::DecodeError::KvBudgetExceeded { binding, .. }
8581 if *binding == frink_models::Ceiling::DeviceMemory.code()
8582 ),
8583 "expected an immovable device-memory refusal, got {err:?}"
8584 );
8585 let (status, _body) = decode_error_response(err);
8586 assert_eq!(status, StatusCode::BAD_REQUEST);
8587 }
8588
8589 #[test]
8595 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8596 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8597 queued: 512,
8598 cap: 512,
8599 });
8600 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8601 assert_eq!(body["error"]["retry_after_seconds"], 1);
8602 let message = body["error"]["message"].as_str().expect("message");
8603 assert!(message.contains("512"), "{message}");
8604 }
8605
8606 #[test]
8607 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8608 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8609 token: 99,
8610 vocab_size: 32,
8611 });
8612 assert!(
8613 body["error"]["retry_after_seconds"].is_null(),
8614 "retrying a prompt this model cannot tokenize never helps"
8615 );
8616 }
8617
8618 #[test]
8619 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8620 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8621 token: 99,
8622 vocab_size: 32,
8623 });
8624 assert_eq!(status, StatusCode::BAD_REQUEST);
8625 }
8626
8627 #[test]
8628 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8629 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8631 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8632 let config = generate::KvPoolConfig {
8633 pool: pool.clone(),
8634 queue_wait: Duration::ZERO,
8635 };
8636
8637 let (choices, _usage) = run_generation(
8638 &model,
8639 &prompt,
8640 &greedy_params(4),
8641 Some(&config),
8642 None,
8643 None,
8644 None,
8645 None,
8646 None,
8647 )
8648 .unwrap();
8649 assert_eq!(choices[0].0, FinishReason::Length);
8650 assert_eq!(
8651 pool.lock().unwrap().free_blocks(),
8652 2,
8653 "a completed request must return its blocks to the pool"
8654 );
8655 }
8656
8657 #[tokio::test]
8662 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8663 let model = Arc::new(test_model());
8664 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8665
8666 let mut handles = Vec::new();
8667 for _ in 0..8 {
8668 let model = Arc::clone(&model);
8669 let prompt = prompt.clone();
8670 handles.push(tokio::task::spawn_blocking(move || {
8671 run_generation(
8672 &model,
8673 &prompt,
8674 &greedy_params(6),
8675 None,
8676 None,
8677 None,
8678 None,
8679 None,
8680 None,
8681 )
8682 .unwrap()
8683 }));
8684 }
8685
8686 let mut results = Vec::new();
8687 for h in handles {
8688 results.push(h.await.unwrap());
8689 }
8690 for r in &results[1..] {
8695 assert_eq!(r.0, results[0].0, "choices must match");
8699 assert_eq!(
8700 r.1.prompt_tokens, results[0].1.prompt_tokens,
8701 "prompt token count must match"
8702 );
8703 assert_eq!(
8704 r.1.completion_tokens, results[0].1.completion_tokens,
8705 "completion token count must match"
8706 );
8707 }
8708 }
8709
8710 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8717 let mut header_entries = Vec::new();
8718 let mut data = Vec::new();
8719 for (name, shape, values) in tensors {
8720 let start = data.len();
8721 for v in values {
8722 data.extend_from_slice(&v.to_le_bytes());
8723 }
8724 let end = data.len();
8725 let shape_str = shape
8726 .iter()
8727 .map(|d| d.to_string())
8728 .collect::<Vec<_>>()
8729 .join(",");
8730 header_entries.push(format!(
8731 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8732 ));
8733 }
8734 let header = format!("{{{}}}", header_entries.join(","));
8735 let header_bytes = header.as_bytes();
8736 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8737 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8738 out.extend_from_slice(header_bytes);
8739 out.extend_from_slice(&data);
8740 out
8741 }
8742
8743 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8754 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8755 use frink_models::kimi_loader::KimiRealHparams;
8756 use frink_moe::{GatingFunction, MoeLayerConfig};
8757
8758 let hidden_dim = 8;
8759 let kda_num_heads = 2;
8760 let kda_head_dim = 3;
8761 let kda_proj = kda_num_heads * kda_head_dim;
8762 let conv_kernel = 4;
8763 let dense_intermediate = 5;
8764 let vocab_size = 256;
8768 let mla_num_heads = 1;
8769 let mla_q_lora_rank = 2;
8770 let mla_kv_lora_rank = 2;
8771 let mla_qk_nope_head_dim = 2;
8772 let mla_qk_rope_head_dim = 2;
8773 let mla_v_head_dim = 2;
8774
8775 let model_cfg = frink_models::ModelConfig {
8776 rope_layers: frink_models::rope_layers::RopeLayers::All,
8777 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
8778 name: "synthetic-kimi-server-test",
8779 n_layers: 1,
8780 n_mtp_blocks: 0,
8781 hidden_dim,
8782 n_heads: 1,
8783 n_kv_heads: 1,
8784 head_dim: 4,
8785 v_head_dim: None,
8786 vocab_size,
8787 rope_theta: 10000.0,
8788 rms_norm_eps: 1e-5,
8789 post_norm_eps: 1e-5,
8790 sliding_window: None,
8791 moe: MoeLayerConfig {
8792 expert_weights_scale: 1.0,
8793 routed_weight_before_ffn: false,
8794 n_experts: 1,
8795 n_experts_active: 1,
8796 n_shared_experts: 0,
8797 hidden_dim,
8798 expert_ffn_dim: 4,
8799 gating: GatingFunction::Sigmoid,
8800 norm_topk_prob: true,
8801 expert_group_count: None,
8802 expert_group_used_count: None,
8803 },
8804 n_dense_leading_layers: 1,
8809 moe_interleave_step: None,
8810 norm_function: frink_models::norm::NormFunction::Rms,
8811 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8812 kda_layers: vec![1],
8813 full_attn_layers: vec![],
8814 mla: MlaConfig {
8815 num_heads: mla_num_heads,
8816 q_lora_rank: mla_q_lora_rank,
8817 kv_lora_rank: mla_kv_lora_rank,
8818 qk_nope_head_dim: mla_qk_nope_head_dim,
8819 qk_rope_head_dim: mla_qk_rope_head_dim,
8820 v_head_dim: mla_v_head_dim,
8821 use_output_gate: true,
8822 rope: None,
8823 },
8824 kda: KdaConfig {
8825 num_heads: kda_num_heads,
8826 head_dim: kda_head_dim,
8827 short_conv_kernel_size: conv_kernel,
8828 gate_lower_bound: -5.0,
8829 use_full_rank_gate: true,
8830 },
8831 }),
8832 rope_freqs: None,
8833 rope_attn_factor: 1.0,
8834 rope_dim: None,
8835 rope_dim_swa: None,
8836 rope_freqs_long: None,
8837 rope_freqs_short: None,
8838 rope_orig_ctx: None,
8839 rope_layout: frink_models::config::RopeLayout::Neox,
8840 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
8841 swa_layers: frink_models::swa_layers::SwaLayers::All,
8842 attn_logit_softcap: None,
8843 final_logit_softcap: None,
8844 embedding_scale: None,
8845 residual_scale: None,
8846 normed_residual_scale: None,
8847 clamp_kqv: None,
8848 attn_temperature: None,
8849 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
8850 block_sub_norms: false,
8851 parallel_residual: false,
8852 learned_positions: false,
8853 attn_value_scale: None,
8854 alibi_max_bias: None,
8855 layer_loops: None,
8856 skip_stream: false,
8857 parallel_ssm: false,
8858 swa_chunked: false,
8859 weightless_qk_norm: false,
8860 logit_multiplier: None,
8861 attention_scale: None,
8862 rope_theta_swa: None,
8863 ffn_activation: frink_models::config::FfnActivation::Swiglu,
8864 best_effort_fields: &["synthetic test config, not a real preset"],
8865 };
8866 let hp = KimiRealHparams {
8867 hidden_dim,
8868 kda_num_heads,
8869 kda_head_dim,
8870 mla_num_heads,
8871 mla_q_lora_rank,
8872 mla_kv_lora_rank,
8873 mla_qk_nope_head_dim,
8874 mla_qk_rope_head_dim,
8875 mla_v_head_dim,
8876 dense_intermediate_dim: dense_intermediate,
8877 moe_hidden_dim: hidden_dim,
8878 moe_intermediate_dim: 4,
8879 n_experts: 1,
8880 num_shared_experts: 0,
8881 };
8882
8883 let prefix = "language_model.model.layers.0";
8887 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8888 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8889 name: String,
8890 shape: Vec<usize>,
8891 n: usize| {
8892 tensors.push((name, shape, vec![0.01f32; n]));
8893 };
8894 push(
8895 &mut tensors,
8896 format!("{prefix}.input_layernorm.weight"),
8897 vec![hidden_dim],
8898 hidden_dim,
8899 );
8900 push(
8901 &mut tensors,
8902 format!("{prefix}.post_attention_layernorm.weight"),
8903 vec![hidden_dim],
8904 hidden_dim,
8905 );
8906 push(
8907 &mut tensors,
8908 format!("{prefix}.self_attention_res_norm.weight"),
8909 vec![hidden_dim],
8910 hidden_dim,
8911 );
8912 push(
8913 &mut tensors,
8914 format!("{prefix}.self_attention_res_proj.weight"),
8915 vec![1, hidden_dim],
8916 hidden_dim,
8917 );
8918 push(
8919 &mut tensors,
8920 format!("{prefix}.mlp_res_norm.weight"),
8921 vec![hidden_dim],
8922 hidden_dim,
8923 );
8924 push(
8925 &mut tensors,
8926 format!("{prefix}.mlp_res_proj.weight"),
8927 vec![1, hidden_dim],
8928 hidden_dim,
8929 );
8930 push(
8931 &mut tensors,
8932 format!("{prefix}.self_attn.q_proj.weight"),
8933 vec![kda_proj, hidden_dim],
8934 kda_proj * hidden_dim,
8935 );
8936 push(
8937 &mut tensors,
8938 format!("{prefix}.self_attn.k_proj.weight"),
8939 vec![kda_proj, hidden_dim],
8940 kda_proj * hidden_dim,
8941 );
8942 push(
8943 &mut tensors,
8944 format!("{prefix}.self_attn.v_proj.weight"),
8945 vec![kda_proj, hidden_dim],
8946 kda_proj * hidden_dim,
8947 );
8948 push(
8949 &mut tensors,
8950 format!("{prefix}.self_attn.q_conv1d.weight"),
8951 vec![kda_proj, 1, conv_kernel],
8952 kda_proj * conv_kernel,
8953 );
8954 push(
8955 &mut tensors,
8956 format!("{prefix}.self_attn.k_conv1d.weight"),
8957 vec![kda_proj, 1, conv_kernel],
8958 kda_proj * conv_kernel,
8959 );
8960 push(
8961 &mut tensors,
8962 format!("{prefix}.self_attn.v_conv1d.weight"),
8963 vec![kda_proj, 1, conv_kernel],
8964 kda_proj * conv_kernel,
8965 );
8966 push(
8967 &mut tensors,
8968 format!("{prefix}.self_attn.A_log"),
8969 vec![kda_num_heads],
8970 kda_num_heads,
8971 );
8972 push(
8973 &mut tensors,
8974 format!("{prefix}.self_attn.f_a_proj.weight"),
8975 vec![kda_head_dim, hidden_dim],
8976 kda_head_dim * hidden_dim,
8977 );
8978 push(
8979 &mut tensors,
8980 format!("{prefix}.self_attn.f_b_proj.weight"),
8981 vec![kda_proj, kda_head_dim],
8982 kda_proj * kda_head_dim,
8983 );
8984 push(
8985 &mut tensors,
8986 format!("{prefix}.self_attn.dt_bias"),
8987 vec![kda_proj],
8988 kda_proj,
8989 );
8990 push(
8991 &mut tensors,
8992 format!("{prefix}.self_attn.b_proj.weight"),
8993 vec![kda_num_heads, hidden_dim],
8994 kda_num_heads * hidden_dim,
8995 );
8996 push(
8997 &mut tensors,
8998 format!("{prefix}.self_attn.g_proj.weight"),
8999 vec![kda_proj, hidden_dim],
9000 kda_proj * hidden_dim,
9001 );
9002 push(
9003 &mut tensors,
9004 format!("{prefix}.self_attn.o_norm.weight"),
9005 vec![kda_head_dim],
9006 kda_head_dim,
9007 );
9008 push(
9009 &mut tensors,
9010 format!("{prefix}.self_attn.o_proj.weight"),
9011 vec![hidden_dim, kda_proj],
9012 hidden_dim * kda_proj,
9013 );
9014 push(
9015 &mut tensors,
9016 format!("{prefix}.mlp.gate_proj.weight"),
9017 vec![dense_intermediate, hidden_dim],
9018 dense_intermediate * hidden_dim,
9019 );
9020 push(
9021 &mut tensors,
9022 format!("{prefix}.mlp.up_proj.weight"),
9023 vec![dense_intermediate, hidden_dim],
9024 dense_intermediate * hidden_dim,
9025 );
9026 push(
9027 &mut tensors,
9028 format!("{prefix}.mlp.down_proj.weight"),
9029 vec![hidden_dim, dense_intermediate],
9030 hidden_dim * dense_intermediate,
9031 );
9032 push(
9033 &mut tensors,
9034 "language_model.model.embed_tokens.weight".to_string(),
9035 vec![vocab_size, hidden_dim],
9036 vocab_size * hidden_dim,
9037 );
9038 push(
9039 &mut tensors,
9040 "language_model.lm_head.weight".to_string(),
9041 vec![vocab_size, hidden_dim],
9042 vocab_size * hidden_dim,
9043 );
9044 push(
9045 &mut tensors,
9046 "language_model.model.norm.weight".to_string(),
9047 vec![hidden_dim],
9048 hidden_dim,
9049 );
9050 push(
9051 &mut tensors,
9052 "language_model.model.output_attn_res_norm.weight".to_string(),
9053 vec![hidden_dim],
9054 hidden_dim,
9055 );
9056 push(
9057 &mut tensors,
9058 "language_model.model.output_attn_res_proj.weight".to_string(),
9059 vec![1, hidden_dim],
9060 hidden_dim,
9061 );
9062
9063 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9076 let dir = std::env::temp_dir().join(format!(
9077 "frink_server_kimi_e2e_test_{}_{}",
9078 std::process::id(),
9079 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9080 ));
9081 std::fs::create_dir_all(&dir).unwrap();
9082 let shard_bytes = write_safetensors_shard(&tensors);
9083 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9084 let map_entries: Vec<String> = tensors
9085 .iter()
9086 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9087 .collect();
9088 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9089 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9090
9091 use base64::Engine;
9095 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9096 .map(|b| {
9097 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9098 format!("{b64} {b}")
9099 })
9100 .collect();
9101 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9102
9103 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9104 .expect("must load the synthetic Kimi checkpoint end to end");
9105 std::fs::remove_dir_all(&dir).ok();
9106 loaded
9107 }
9108
9109 #[test]
9117 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9118 let loaded = build_synthetic_kimi_loaded();
9119 let state = build_app_state(
9120 StartupModels {
9121 loaded: model::LoadedModel::Kimi(loaded),
9122 embedding: None,
9123 },
9124 None,
9125 None,
9126 None,
9127 false,
9128 None,
9129 Arc::new(health::Detection::ready(health::probe_backends())),
9130 );
9131 let active = state.active().expect("a freshly built state has a model");
9132 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9133 assert!(!active.is_synthetic());
9134
9135 let (choices, _usage) = run_generation(
9136 active.generative().unwrap(),
9137 "hi",
9138 &greedy_params(5),
9139 None,
9140 None,
9141 None,
9142 None,
9143 None,
9144 None,
9145 )
9146 .expect("a real Kimi checkpoint must generate without error");
9147 assert!(matches!(
9148 choices[0].0,
9149 FinishReason::Length | FinishReason::Stop
9150 ));
9151 }
9152
9153 #[test]
9166 fn a_grammar_constrains_the_engine_decode_path() {
9167 let loaded = build_synthetic_kimi_loaded();
9168 let state = build_app_state(
9169 StartupModels {
9170 loaded: model::LoadedModel::Kimi(loaded),
9171 embedding: None,
9172 },
9173 None,
9174 None,
9175 None,
9176 false,
9177 None,
9178 Arc::new(health::Detection::ready(health::probe_backends())),
9179 );
9180 let active = state.active().expect("a freshly built state has a model");
9181
9182 let run = |grammar: Option<&str>| {
9183 let mut params = greedy_params(6);
9184 params.grammar = grammar.map(|src| {
9185 Arc::new(
9186 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9187 .expect("test grammar parses"),
9188 )
9189 });
9190 run_generation(
9191 active.generative().unwrap(),
9192 "hi",
9193 ¶ms,
9194 None,
9195 None,
9196 None,
9197 None,
9198 None,
9199 None,
9200 )
9201 };
9202
9203 let (choices, _) = run(None).expect("the unconstrained run must serve");
9204 let unconstrained = choices[0].1.clone();
9205 assert!(
9206 unconstrained.chars().any(|c| c != 'a'),
9207 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9208 constrained run below would prove nothing"
9209 );
9210
9211 let (choices, _) =
9212 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9213 let (finish, constrained) = choices.into_iter().next().unwrap();
9214 assert!(
9215 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9216 "the engine decode path served text its grammar forbids ({constrained:?}): \
9217 the constraint was dropped between `generate_engine` and the sampler"
9218 );
9219 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9220 }
9221
9222 #[test]
9232 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9233 let loaded = build_synthetic_kimi_loaded();
9234 let state = build_app_state(
9235 StartupModels {
9236 loaded: model::LoadedModel::Kimi(loaded),
9237 embedding: None,
9238 },
9239 None,
9240 None,
9241 None,
9242 false,
9243 None,
9244 Arc::new(health::Detection::ready(health::probe_backends())),
9245 );
9246
9247 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9248 let kv_pool_config = generate::KvPoolConfig {
9249 pool: pool.clone(),
9250 queue_wait: Duration::ZERO,
9251 };
9252 let pc = Mutex::new(PrefixCache::new(4));
9253
9254 run_generation(
9255 state
9256 .active()
9257 .expect("a freshly built state has a model")
9258 .generative()
9259 .unwrap(),
9260 "hi",
9261 &greedy_params(5),
9262 Some(&kv_pool_config),
9263 None,
9264 Some(&pc),
9265 None,
9266 None,
9267 None,
9268 )
9269 .expect("a real Kimi checkpoint must generate without error");
9270
9271 assert_eq!(
9272 pool.lock().unwrap().free_blocks(),
9273 4,
9274 "the KV pool must be completely untouched by a Kimi request"
9275 );
9276 let stats = pc.lock().unwrap().stats();
9277 assert_eq!(
9278 stats.hits + stats.misses,
9279 0,
9280 "the prefix cache must never be consulted for a Kimi request"
9281 );
9282 }
9283}