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 logprobs;
57mod lora;
58mod mcp;
59mod model;
60mod openai_extra;
61mod output;
62mod policy;
63mod prefill_batch;
64mod reasoning_budget;
65mod reasoning_tokens;
66mod request_tail;
67mod rerank;
68mod response_cache;
69pub(crate) mod responses;
70mod resume;
71mod sample_step;
72mod sampling_knobs;
73mod sampling_loop;
74mod security;
75mod serving;
76mod session;
77mod slots;
78mod sse;
79mod stats;
80mod stop;
81mod stream_events;
82mod tasks;
83mod tool_grammar;
84mod unimplemented_fields;
85mod unsupported_sampling;
86mod utf8_stream;
87
88use std::cell::RefCell;
89use std::convert::Infallible;
90use std::net::SocketAddr;
91use std::path::PathBuf;
92use std::rc::Rc;
93use std::sync::{Arc, Mutex, MutexGuard};
94use std::time::Duration;
95
96use axum::{
97 extract::State,
98 http::StatusCode,
99 response::sse::{Event, Sse},
100 response::{IntoResponse, Response},
101 routing::{get, post},
102 Json, Router,
103};
104use serde::{Deserialize, Serialize};
105
106use cli::apply_cli_overrides;
107pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
108
109use frink_core::cache::KvBlockPool;
110use frink_models::kimi_tokenizer::KimiTokenizer;
111use frink_models::sampling::SamplingParams;
112use frink_models::tokenizer::{SpecialTokens, StopTokens};
113use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
114#[cfg(test)]
115use generate::FinishReason;
116use generate::GenerationParams;
117pub(crate) use loaded::{ActiveModel, Loaded};
118use model::ServerTokenizer;
119use rerank::encoder_endpoints;
120use response_cache::ResponseCache;
121use sampling_knobs::SamplingKnobs;
122
123#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
132 Gguf(GgufModel),
133 Kimi(KimiModel),
134 Mla(MlaModel),
135 Gemma4(Gemma4Model),
136 Glm52(Glm52Model),
137}
138
139pub(crate) struct GgufModel {
140 decoder: Arc<Decoder>,
141 tokenizer: Arc<ServerTokenizer>,
142 stop_tokens: StopTokens,
143 bos_id: Option<usize>,
144 is_synthetic: bool,
145 chat_template: chat_template::PromptTemplate,
146}
147
148pub(crate) struct KimiModel {
149 engine: KimiEngine,
150 tokenizer: KimiTokenizer,
151 stop_tokens: StopTokens,
152 chat_template: chat_template::PromptTemplate,
153}
154
155pub(crate) struct MlaModel {
156 engine: MlaEngine,
157 tokenizer: ServerTokenizer,
158 stop_tokens: StopTokens,
159 bos_id: Option<usize>,
160 name: String,
161 chat_template: chat_template::PromptTemplate,
162}
163
164pub(crate) struct Gemma4Model {
165 engine: Gemma4Engine,
166 tokenizer: ServerTokenizer,
167 stop_tokens: StopTokens,
168 bos_id: Option<usize>,
169 name: String,
170 chat_template: chat_template::PromptTemplate,
171}
172
173pub(crate) struct Glm52Model {
174 engine: frink_models::Glm52Engine,
175 tokenizer: ServerTokenizer,
176 stop_tokens: StopTokens,
177 bos_id: Option<usize>,
178 name: String,
179 chat_template: chat_template::PromptTemplate,
180}
181
182impl Model {
183 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
184 match self {
185 Model::Gguf(m) => m.chat_template.clone(),
186 Model::Kimi(m) => m.chat_template.clone(),
187 Model::Mla(m) => m.chat_template.clone(),
188 Model::Gemma4(m) => m.chat_template.clone(),
189 Model::Glm52(m) => m.chat_template.clone(),
190 }
191 }
192
193 fn is_synthetic(&self) -> bool {
198 match self {
199 Model::Gguf(m) => m.is_synthetic,
200 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
201 }
202 }
203
204 fn tokenizer_kind(&self) -> &'static str {
205 match self {
206 Model::Gguf(m) => m.tokenizer.kind(),
207 Model::Kimi(_) => "kimi-tiktoken-bpe",
208 Model::Mla(m) => m.tokenizer.kind(),
209 Model::Gemma4(m) => m.tokenizer.kind(),
210 Model::Glm52(m) => m.tokenizer.kind(),
211 }
212 }
213
214 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
218 match self {
219 Model::Gguf(m) => m.decoder.expert_store_stats(),
220 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
221 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
222 }
223 }
224
225 pub(crate) fn name(&self) -> &str {
226 match self {
227 Model::Gguf(m) => m.decoder.config.name,
228 Model::Kimi(_) => "kimi-k3",
229 Model::Mla(m) => m.name.as_str(),
230 Model::Gemma4(m) => m.name.as_str(),
231 Model::Glm52(m) => m.name.as_str(),
232 }
233 }
234
235 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
262 match self {
263 Model::Gguf(m) => m.tokenizer.encode(text, specials),
264 Model::Kimi(m) => m
265 .tokenizer
266 .encode(text, specials)
267 .into_iter()
268 .map(|id| id as usize)
269 .collect(),
270 Model::Mla(m) => m.tokenizer.encode(text, specials),
271 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
272 Model::Glm52(m) => m.tokenizer.encode(text, specials),
273 }
274 }
275
276 pub(crate) fn bos_id(&self) -> Option<usize> {
285 match self {
286 Model::Gguf(m) => m.bos_id,
287 Model::Kimi(_) => None,
288 Model::Mla(m) => m.bos_id,
289 Model::Gemma4(m) => m.bos_id,
290 Model::Glm52(m) => m.bos_id,
291 }
292 }
293
294 pub(crate) fn decode(&self, ids: &[usize]) -> String {
295 match self {
296 Model::Gguf(m) => m.tokenizer.decode(ids),
297 Model::Kimi(m) => {
298 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
299 m.tokenizer.decode(&ids32)
300 }
301 Model::Mla(m) => m.tokenizer.decode(ids),
302 Model::Gemma4(m) => m.tokenizer.decode(ids),
303 Model::Glm52(m) => m.tokenizer.decode(ids),
304 }
305 }
306
307 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
310 match self {
311 Model::Gguf(m) => {
312 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
313 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
314 }
315 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
316 }
317 }
318
319 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
325 match self {
326 Model::Gguf(m) => Some(&m.decoder),
327 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
328 }
329 }
330
331 pub(crate) fn vocab_size(&self) -> Option<usize> {
332 match self {
333 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
334 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
335 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
336 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
337 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
338 }
339 }
340
341 fn has_real_vocabulary(&self) -> bool {
348 match self {
349 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
350 Model::Kimi(_) => true,
351 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
352 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
353 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
354 }
355 }
356}
357
358impl frink_models::dry::DryVocab for Model {
364 fn n_tokens(&self) -> usize {
365 self.vocab_size().unwrap_or(0)
366 }
367
368 fn detokenize(&self, token: usize) -> String {
369 self.decode(&[token])
370 }
371
372 fn tokenize(&self, text: &str) -> Vec<usize> {
373 self.encode(text, SpecialTokens::AsText)
374 }
375}
376
377pub(crate) struct AppState {
378 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
392 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
411 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
415 pub(crate) tasks: Arc<tasks::TaskRegistry>,
417 pub(crate) cancels: Arc<cancel::CancelRegistry>,
421 pub(crate) stats: stats::Stats,
424 pub(crate) streams: resume::StreamRegistry,
427 pub(crate) model_dir: Option<PathBuf>,
429 response_cache: Mutex<ResponseCache>,
433 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
445 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
458 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
468 sessions: session::SessionStore,
474 requests_total: std::sync::atomic::AtomicU64,
475 request_errors_total: std::sync::atomic::AtomicU64,
476 started_at: std::time::Instant,
477 last_request_ms: std::sync::atomic::AtomicU64,
482 detection: Arc<health::Detection>,
484 mcp: Option<mcp::LoadedMcpConfig>,
486 pub(crate) continuous_batching_enabled: bool,
490 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
494 loading_model: Mutex<Option<String>>,
498 last_load_error: Mutex<Option<(String, String)>>,
502 pub(crate) serving: Mutex<crate::stats::ServingStats>,
507 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
512 pub(crate) footprint:
518 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
519 pub(crate) started_unix: u64,
529}
530
531pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
537
538impl AppState {
539 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
543 self.active
544 .read()
545 .unwrap_or_else(|p| p.into_inner())
546 .clone()
547 }
548
549 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
555 self.active().ok_or_else(|| {
556 (
557 StatusCode::SERVICE_UNAVAILABLE,
558 Json(serde_json::json!({"error": {
559 "message": "no model is loaded; POST /admin/models/load with an id from \
560 GET /admin/models",
561 "type": "model_not_loaded"
562 }})),
563 )
564 })
565 }
566
567 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
575 Ok(Arc::clone(self.require_active()?.generative()?))
576 }
577
578 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
586 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
587 std::mem::replace(&mut *guard, next)
588 }
589
590 fn mark_request_finished(&self) {
594 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
595 self.last_request_ms
596 .store(ms, std::sync::atomic::Ordering::Relaxed);
597 }
598
599 pub(crate) fn uptime(&self) -> Duration {
600 self.started_at.elapsed()
601 }
602
603 pub(crate) fn requests_total(&self) -> u64 {
604 self.requests_total
605 .load(std::sync::atomic::Ordering::Relaxed)
606 }
607
608 pub(crate) fn errors_total(&self) -> u64 {
609 self.request_errors_total
610 .load(std::sync::atomic::Ordering::Relaxed)
611 }
612
613 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
614 lock_cache(&self.response_cache).stats()
615 }
616
617 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
620 let last = self
621 .last_request_ms
622 .load(std::sync::atomic::Ordering::Relaxed);
623 (last > 0)
624 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
625 .map(|age| age.max(0.0))
626 }
627
628 pub(crate) fn loading_model_id(&self) -> Option<String> {
629 self.loading_model
630 .lock()
631 .unwrap_or_else(|p| p.into_inner())
632 .clone()
633 }
634
635 pub(crate) fn set_loading_model(&self, id: Option<String>) {
636 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
637 }
638
639 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
640 self.last_load_error
641 .lock()
642 .unwrap_or_else(|p| p.into_inner())
643 .clone()
644 }
645
646 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
647 *self
648 .last_load_error
649 .lock()
650 .unwrap_or_else(|p| p.into_inner()) = error;
651 }
652
653 pub(crate) fn active_model_name(&self) -> Option<String> {
662 self.active().map(|a| a.name().to_string())
663 }
664
665 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
675 self.embedding
676 .clone()
677 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
678 }
679
680 pub(crate) fn embedding_model_name(&self) -> Option<String> {
684 match self.embedding_model() {
685 Some(e) => Some(e.name().to_string()),
686 None => self.active_model_name(),
687 }
688 }
689
690 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
691 self.stats.record(stats::entry(record));
692 }
693}
694
695fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
700 cache
701 .lock()
702 .unwrap_or_else(|poisoned| poisoned.into_inner())
703}
704
705#[derive(Debug, Clone, Deserialize)]
706#[serde(untagged)]
707pub(crate) enum MessageContent {
708 Text(String),
709 Parts(Vec<ContentPart>),
710}
711
712#[derive(Debug, Clone, Deserialize)]
713struct ContentPart {
714 #[serde(rename = "type")]
715 kind: String,
716 #[serde(default)]
717 text: Option<String>,
718 #[serde(default)]
719 image_url: Option<serde_json::Value>,
720}
721
722impl MessageContent {
723 fn as_text(&self) -> String {
724 match self {
725 Self::Text(s) => s.clone(),
726 Self::Parts(parts) => parts
727 .iter()
728 .filter_map(|p| p.text.as_deref())
729 .collect::<Vec<_>>()
730 .join(""),
731 }
732 }
733
734 fn has_image(&self) -> bool {
735 match self {
736 Self::Text(_) => false,
737 Self::Parts(parts) => parts
738 .iter()
739 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
740 }
741 }
742}
743
744#[derive(Debug, Clone, Deserialize)]
745pub(crate) struct ChatMessage {
746 pub(crate) role: String,
747 #[serde(default)]
751 pub(crate) content: Option<MessageContent>,
752 #[serde(default)]
756 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
757 #[serde(default)]
762 #[allow(dead_code)]
763 pub(crate) tool_call_id: Option<String>,
764 #[serde(default, alias = "reasoning")]
780 pub(crate) reasoning_content: Option<String>,
781}
782
783impl ChatMessage {
784 fn rendered_content(&self) -> String {
791 let mut out = self
792 .content
793 .as_ref()
794 .map(MessageContent::as_text)
795 .unwrap_or_default();
796 if let Some(calls) = &self.tool_calls {
797 for call in calls {
798 out.push_str(&format!(
799 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
800 call.function.name, call.function.arguments
801 ));
802 }
803 }
804 out
805 }
806}
807
808#[derive(Debug, Clone, Deserialize)]
809pub(crate) struct ToolCallIn {
810 #[serde(default)]
811 #[allow(dead_code)]
812 id: String,
813 #[serde(rename = "type", default)]
814 #[allow(dead_code)]
815 kind: String,
816 function: ToolCallFunctionIn,
817}
818
819#[derive(Debug, Clone, Deserialize)]
820struct ToolCallFunctionIn {
821 name: String,
822 arguments: String,
827}
828
829#[derive(Debug, Clone, Deserialize)]
832struct ToolDef {
833 #[serde(rename = "type", default)]
834 #[allow(dead_code)]
835 kind: String,
836 function: ToolFunctionDef,
837}
838
839#[derive(Debug, Clone, Deserialize)]
840struct ToolFunctionDef {
841 name: String,
842 #[serde(default)]
843 description: Option<String>,
844 #[serde(default)]
845 parameters: Option<serde_json::Value>,
846}
847
848#[derive(Debug, Clone, Deserialize)]
859#[serde(untagged)]
860enum ToolChoice {
861 Mode(String),
862 Specific(serde_json::Value),
863}
864
865#[derive(Deserialize)]
868#[serde(untagged)]
869enum StopParam {
870 One(String),
871 Many(Vec<String>),
872}
873
874#[derive(Deserialize)]
875struct ChatCompletionRequest {
876 model: String,
877 messages: Vec<ChatMessage>,
878 #[serde(default = "default_max_tokens")]
879 max_tokens: usize,
880 #[serde(default)]
881 temperature: Option<f32>,
882 #[serde(default)]
883 top_p: Option<f32>,
884 #[serde(default)]
889 min_p: Option<f32>,
890 #[serde(default)]
891 top_k: Option<usize>,
892 #[serde(default)]
893 repetition_penalty: Option<f32>,
894 #[serde(flatten)]
898 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
899 #[serde(flatten)]
903 unimplemented: crate::unimplemented_fields::UnimplementedFields,
904 #[serde(default)]
905 seed: Option<u64>,
906 #[serde(default)]
907 stop: Option<StopParam>,
908 #[serde(default)]
909 stream: Option<bool>,
910 #[serde(default)]
923 stream_resumable: Option<bool>,
924 #[serde(default)]
935 ignore_eos: Option<bool>,
936 #[serde(default)]
937 tools: Vec<ToolDef>,
938 #[serde(default)]
939 tool_choice: Option<ToolChoice>,
940 #[serde(default)]
948 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
949 #[serde(default)]
957 reasoning_effort: Option<String>,
958 #[serde(default)]
962 thinking: Option<ThinkingSwitch>,
963 #[serde(default)]
968 session_id: Option<String>,
969 #[serde(default, deserialize_with = "continuation::deserialize")]
977 continue_final_message: continuation::ContinueFinalMessage,
978 #[serde(default, alias = "thinking_budget_tokens")]
986 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
987 #[serde(default)]
989 logprobs: Option<bool>,
990 #[serde(default)]
991 top_logprobs: Option<u32>,
992 #[serde(default)]
993 presence_penalty: Option<f32>,
994 #[serde(default)]
995 frequency_penalty: Option<f32>,
996 #[serde(default)]
997 response_format: Option<serde_json::Value>,
998 #[serde(default)]
1005 logit_bias: Option<serde_json::Value>,
1006 #[serde(default)]
1010 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1011 #[serde(default)]
1021 samplers: Option<serde_json::Value>,
1022 #[serde(default)]
1032 grammar: Option<String>,
1033}
1034
1035const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1047
1048#[derive(Debug, Clone, Deserialize)]
1050pub(crate) struct ThinkingSwitch {
1051 #[serde(rename = "type")]
1052 pub(crate) kind: String,
1053}
1054
1055const THINKING_KWARG_KEYS: [&str; 4] = [
1059 "enable_thinking",
1060 "thinking",
1061 "thinking_mode",
1062 "reasoning_effort",
1063];
1064
1065const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1069
1070fn default_max_tokens() -> usize {
1071 DEFAULT_CHAT_MAX_TOKENS
1072}
1073
1074impl ChatCompletionRequest {
1075 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1084 let mut knobs = SamplingKnobs {
1085 temperature: self.temperature,
1086 top_p: self.top_p,
1087 min_p: self.min_p,
1088 top_k: self.top_k,
1089 repetition_penalty: self.repetition_penalty,
1090 presence_penalty: self.presence_penalty,
1091 frequency_penalty: self.frequency_penalty,
1092 penalty_last_n: None,
1096 sampler_order: unsupported_sampling::parse_sampler_order(
1097 self.samplers.as_ref(),
1098 "/v1/chat/completions",
1099 )?,
1100 ..SamplingKnobs::default()
1101 };
1102 self.extra_samplers.apply(&mut knobs);
1103 Ok(knobs)
1104 }
1105
1106 fn sampling_params(
1107 &self,
1108 model: crate::sampling_knobs::SamplerModel<'_>,
1109 ) -> Result<SamplingParams, ApiError> {
1110 self.sampling_knobs()?.resolve(model).map_err(|e| {
1111 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1112 })
1113 }
1114
1115 fn stop_sequences(&self) -> Vec<String> {
1116 self.stop
1117 .as_ref()
1118 .map(|s| match s {
1119 StopParam::One(v) => vec![v.clone()],
1120 StopParam::Many(v) => v.clone(),
1121 })
1122 .unwrap_or_default()
1123 }
1124
1125 fn several_choices(&self) -> bool {
1134 self.unimplemented.n.is_some_and(|n| n > 1)
1135 }
1136
1137 fn tools_active(&self) -> bool {
1138 !self.tools.is_empty()
1139 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1140 }
1141
1142 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1154 match &self.tool_choice {
1155 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1156 Some(ToolChoice::Specific(value)) => {
1157 let name = value
1161 .get("function")
1162 .and_then(|f| f.get("name"))
1163 .or_else(|| value.get("name"))
1164 .and_then(|n| n.as_str());
1165 match name {
1166 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1167 None => Err(invalid_request(
1168 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1169 function.name",
1170 "tool_choice",
1171 )),
1172 }
1173 }
1174 _ => Ok(None),
1175 }
1176 }
1177
1178 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1180 self.tools
1181 .iter()
1182 .map(|t| tool_grammar::ToolSpec {
1183 name: &t.function.name,
1184 parameters: t.function.parameters.as_ref(),
1185 })
1186 .collect()
1187 }
1188
1189 fn thinking_direction(&self) -> Option<bool> {
1232 if let Some(switch) = &self.thinking {
1233 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1234 "disabled" => Some(false),
1235 "enabled" => Some(true),
1236 _ => None,
1239 };
1240 }
1241 let effort = self.reasoning_effort.as_ref()?;
1242 DISABLE_EFFORTS
1243 .contains(&effort.trim().to_ascii_lowercase().as_str())
1244 .then_some(false)
1245 }
1246
1247 fn resolve_template_kwargs(
1248 &self,
1249 template: &chat_template::PromptTemplate,
1250 ) -> serde_json::Map<String, serde_json::Value> {
1251 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1252 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1256
1257 if !caller_steered {
1258 match self.thinking_direction() {
1259 Some(false) => {
1260 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1261 kwargs.insert(k, v);
1262 }
1263 return kwargs;
1266 }
1267 Some(true) => {
1268 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1269 kwargs.insert(k, v);
1270 }
1271 }
1272 None => {}
1273 }
1274 if let Some(effort) = &self.reasoning_effort {
1275 kwargs
1276 .entry("reasoning_effort".to_string())
1277 .or_insert_with(|| serde_json::json!(effort));
1278 }
1279 }
1280
1281 let offered: Vec<serde_json::Value> = if self.tools_active() {
1282 self.tools.iter().map(chat_template::tool_json).collect()
1283 } else {
1284 Vec::new()
1285 };
1286 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1287 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1288 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1289 kwargs.entry(k).or_insert(v);
1290 }
1291 }
1292 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1293 crate::policy::effort::EffortMapping::Mapped(to) => {
1294 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1295 }
1296 crate::policy::effort::EffortMapping::Dropped => {
1297 tracing::debug!(
1298 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1299 enough, so its own default applies"
1300 );
1301 }
1302 crate::policy::effort::EffortMapping::Unchanged => {}
1303 }
1304 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1305 kwargs
1306 }
1307
1308 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1311 if self.max_tokens == 0 {
1318 return Err(invalid_request(
1319 "max_tokens must be at least 1",
1320 "max_tokens",
1321 ));
1322 }
1323 if let Some(switch) = &self.thinking {
1328 let kind = switch.kind.trim().to_ascii_lowercase();
1329 if kind != "enabled" && kind != "disabled" {
1330 return Err(invalid_request(
1331 "thinking.type must be \"enabled\" or \"disabled\"",
1332 "thinking.type",
1333 ));
1334 }
1335 }
1336 for msg in &self.messages {
1337 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1338 return Err(unsupported_feature(
1339 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1340 ));
1341 }
1342 }
1343 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1344 return Err(unsupported_feature(
1345 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1346 ));
1347 }
1348 self.unimplemented.refuse("/v1/chat/completions")?;
1352 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1353 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1357 let stated_grammar =
1372 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1373 if let Some(forced) = self.forced_tool_choice()? {
1381 if self.tools.is_empty() {
1382 return Err(invalid_request(
1383 "tool_choice forces a tool call, but no tools were offered",
1384 "tool_choice",
1385 ));
1386 }
1387 if let tool_grammar::Forced::Named(name) = forced {
1388 if !self.tools.iter().any(|t| t.function.name == name) {
1389 return Err(invalid_request(
1390 &format!(
1391 "tool_choice names {name:?}, which is not one of the tools offered"
1392 ),
1393 "tool_choice",
1394 ));
1395 }
1396 }
1397 if stated_grammar.is_some() {
1407 return Err(invalid_request(
1408 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1409 are two different constraints on the same generation; send one",
1410 "tool_choice",
1411 ));
1412 }
1413 if self.json_object_mode() {
1414 return Err(invalid_request(
1415 "a forced tool_choice cannot be combined with response_format json_object: \
1416 the tool-call markers are not JSON",
1417 "tool_choice",
1418 ));
1419 }
1420 }
1421 Ok(())
1422 }
1423
1424 fn effective_stop_sequences(&self) -> Vec<String> {
1431 let mut stop = self.stop_sequences();
1432 if self.tools_active() {
1433 stop.push("</tool_call>".to_string());
1434 }
1435 stop
1436 }
1437
1438 fn json_object_mode(&self) -> bool {
1439 self.response_format
1440 .as_ref()
1441 .and_then(|v| v.get("type"))
1442 .and_then(|v| v.as_str())
1443 == Some("json_object")
1444 }
1445}
1446
1447#[derive(Serialize)]
1448struct ChatCompletionChoice {
1449 index: usize,
1450 message: ChatCompletionResponseMessage,
1451 finish_reason: &'static str,
1452}
1453
1454#[derive(Serialize)]
1455struct ChatCompletionResponseMessage {
1456 role: &'static str,
1457 #[serde(skip_serializing_if = "Option::is_none")]
1458 content: Option<String>,
1459 #[serde(skip_serializing_if = "Option::is_none")]
1463 reasoning_content: Option<String>,
1464 #[serde(skip_serializing_if = "Option::is_none")]
1465 tool_calls: Option<Vec<ToolCallOut>>,
1466}
1467
1468#[derive(Serialize, Clone)]
1469struct ToolCallOut {
1470 id: String,
1471 #[serde(rename = "type")]
1472 kind: &'static str,
1473 function: ToolCallFunctionOut,
1474}
1475
1476#[derive(Serialize, Clone)]
1484struct ToolCallDelta {
1485 index: usize,
1486 #[serde(skip_serializing_if = "Option::is_none")]
1487 id: Option<String>,
1488 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1489 kind: Option<&'static str>,
1490 function: ToolCallFunctionDelta,
1491}
1492
1493#[derive(Serialize, Clone, Default)]
1494struct ToolCallFunctionDelta {
1495 #[serde(skip_serializing_if = "Option::is_none")]
1496 name: Option<String>,
1497 #[serde(skip_serializing_if = "Option::is_none")]
1500 arguments: Option<String>,
1501}
1502
1503impl ToolCallDelta {
1504 fn whole(index: usize, name: String, arguments: String) -> Self {
1506 ToolCallDelta {
1507 index,
1508 id: Some(format!("call_{index}")),
1509 kind: Some("function"),
1510 function: ToolCallFunctionDelta {
1511 name: Some(name),
1512 arguments: Some(arguments),
1513 },
1514 }
1515 }
1516
1517 fn opening(index: usize, name: String) -> Self {
1519 ToolCallDelta {
1520 index,
1521 id: Some(format!("call_{index}")),
1522 kind: Some("function"),
1523 function: ToolCallFunctionDelta {
1524 name: Some(name),
1525 arguments: Some(String::new()),
1526 },
1527 }
1528 }
1529
1530 fn arguments(index: usize, fragment: String) -> Self {
1532 ToolCallDelta {
1533 index,
1534 id: None,
1535 kind: None,
1536 function: ToolCallFunctionDelta {
1537 name: None,
1538 arguments: Some(fragment),
1539 },
1540 }
1541 }
1542}
1543
1544#[derive(Serialize, Clone)]
1545struct ToolCallFunctionOut {
1546 name: String,
1547 arguments: String,
1551}
1552
1553#[derive(Serialize)]
1554struct ChatCompletionResponse {
1555 id: String,
1556 request_id: String,
1562 object: &'static str,
1563 model: String,
1564 choices: Vec<ChatCompletionChoice>,
1565 usage: generate::Usage,
1570 frink_cache: &'static str,
1581}
1582
1583#[derive(Serialize)]
1584struct ChatCompletionChunkDelta {
1585 #[serde(skip_serializing_if = "Option::is_none")]
1586 role: Option<&'static str>,
1587 #[serde(skip_serializing_if = "Option::is_none")]
1588 content: Option<String>,
1589 #[serde(skip_serializing_if = "Option::is_none")]
1591 reasoning_content: Option<String>,
1592 #[serde(skip_serializing_if = "Option::is_none")]
1593 tool_calls: Option<Vec<ToolCallDelta>>,
1594}
1595
1596#[derive(Serialize)]
1597struct ChatCompletionChunkChoice {
1598 index: usize,
1599 delta: ChatCompletionChunkDelta,
1600 finish_reason: Option<&'static str>,
1601}
1602
1603#[derive(Serialize)]
1604struct ChatCompletionChunk {
1605 id: String,
1606 #[serde(skip_serializing_if = "Option::is_none")]
1613 request_id: Option<String>,
1614 object: &'static str,
1615 model: String,
1616 choices: Vec<ChatCompletionChunkChoice>,
1617 #[serde(skip_serializing_if = "Option::is_none")]
1620 usage: Option<generate::Usage>,
1621}
1622
1623async fn health(State(state): State<Arc<AppState>>) -> Response {
1629 let snapshot = state.detection.snapshot();
1630 let mut capabilities = snapshot.capabilities;
1631 let active = state.active();
1632
1633 capabilities.push(match active.as_deref() {
1636 None => frink_api::Capability::unavailable(
1641 frink_api::health::capability::REAL_WEIGHTS,
1642 frink_api::health::reason::MODEL_NOT_LOADED,
1643 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1644 ),
1645 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1646 frink_api::health::capability::REAL_WEIGHTS,
1647 frink_api::health::reason::MODEL_NOT_LOADED,
1648 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1649 checkpoint. Output from this model is noise.",
1650 ),
1651 Some(active) if active.encoder().is_some() => {
1661 let endpoints = active
1662 .encoder()
1663 .map(|e| encoder_endpoints(e))
1664 .unwrap_or_default();
1665 let served_by = match endpoints.is_empty() {
1666 true => "no endpoint in this build serves it".to_string(),
1667 false => format!("served by {}", endpoints.join(" and ")),
1668 };
1669 frink_api::Capability::available(
1670 frink_api::health::capability::REAL_WEIGHTS,
1671 format!(
1672 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1673 {served_by}; generation endpoints refuse it.",
1674 active.name(),
1675 ),
1676 )
1677 }
1678 Some(active) => frink_api::Capability::available(
1679 frink_api::health::capability::REAL_WEIGHTS,
1680 format!("Serving the real checkpoint '{}'.", active.name()),
1681 ),
1682 });
1683 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1684 frink_api::Capability::available(
1685 frink_api::health::capability::CONTINUOUS_BATCHING,
1686 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1687 "On by default on Metal. Concurrent requests share one batched decode worker."
1688 } else {
1689 "Concurrent requests share one batched decode step."
1690 },
1691 )
1692 } else if state.metal_private_decode_gate.is_some() {
1693 frink_api::Capability::unavailable(
1694 frink_api::health::capability::CONTINUOUS_BATCHING,
1695 frink_api::health::reason::DISABLED,
1696 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1697 )
1698 } else {
1699 frink_api::Capability::unavailable(
1700 frink_api::health::capability::CONTINUOUS_BATCHING,
1701 frink_api::health::reason::DISABLED,
1702 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1703 )
1704 });
1705
1706 let last_request_ms = state
1707 .last_request_ms
1708 .load(std::sync::atomic::Ordering::Relaxed);
1709 let uptime = state.started_at.elapsed();
1710 let health_state = if active.is_none() {
1718 frink_api::HealthState::Unavailable
1719 } else {
1720 snapshot.state
1721 };
1722 let body = frink_api::HealthResponse {
1723 state: health_state,
1724 reason: match health_state {
1725 frink_api::HealthState::Ready => None,
1726 frink_api::HealthState::Unavailable => {
1727 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1728 }
1729 frink_api::HealthState::Detecting => {
1730 Some(frink_api::health::reason::DETECTING.to_string())
1731 }
1732 },
1733 detail: match health_state {
1734 frink_api::HealthState::Ready => None,
1735 frink_api::HealthState::Unavailable => Some(
1736 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1737 .to_string(),
1738 ),
1739 frink_api::HealthState::Detecting => {
1740 Some("Probing available compute backends.".to_string())
1741 }
1742 },
1743 model: active
1744 .as_deref()
1745 .map(|active| frink_api::health::ModelSummary {
1746 id: active.name().to_string(),
1747 tokenizer: active.tokenizer_kind().to_string(),
1748 synthetic_weights: active.is_synthetic(),
1749 }),
1750 capabilities,
1751 version: env!("CARGO_PKG_VERSION").to_string(),
1752 pid: std::process::id(),
1753 uptime_seconds: uptime.as_secs_f64(),
1754 server_time_unix_ms: std::time::SystemTime::now()
1755 .duration_since(std::time::UNIX_EPOCH)
1756 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1757 .unwrap_or(0),
1758 last_request_age_seconds: (last_request_ms > 0)
1759 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1760 .map(|age| age.max(0.0)),
1761 };
1762
1763 let status =
1764 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1765 (status, Json(body)).into_response()
1766}
1767
1768async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1769 let Some(active) = state.active() else {
1773 return Json(serde_json::json!({ "object": "list", "data": [] }));
1774 };
1775 let mut model_entry = serde_json::json!({
1776 "id": active.name(),
1777 "object": "model",
1778 "frink_synthetic_weights": active.is_synthetic(),
1779 "frink_tokenizer": active.tokenizer_kind(),
1780 });
1781 if let Some(encoder) = active.encoder() {
1787 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1788 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1789 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1790 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1791 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1792 }
1793 if let Some(model) = active.generative_opt() {
1801 let parser_configured = active.reasoning_format().is_some();
1802 let gears = model.chat_template().think_gears(parser_configured);
1803 if !gears.is_empty() {
1804 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1805 if let Some(default) = &gears.default {
1806 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1807 }
1808 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1811 }
1812 }
1813 if let Some(mcp) = &state.mcp {
1814 model_entry["frink_mcp"] = mcp.models_metadata();
1815 }
1816 Json(serde_json::json!({
1817 "object": "list",
1818 "data": [model_entry]
1819 }))
1820}
1821
1822async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1836 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1837 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1838 let active = state.active();
1839 Json(serde_json::json!({
1840 "model": active.as_ref().map(|a| a.name()),
1841 "state": state
1842 .maintenance
1843 .lock()
1844 .unwrap_or_else(|p| p.into_inner())
1845 .state()
1846 .as_str(),
1847 "uptime_s": state.uptime().as_secs(),
1848 "throughput": {
1849 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1850 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1851 },
1852 "requests": {
1853 "active": state.cancels.live_count(),
1854 "completed": state.stats.recorded_total(),
1855 "p95_ms": state.stats.p95_duration_ms(),
1856 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1857 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1858 "completion_tokens_total": state.stats.tokens_generated_total(),
1859 },
1860 "pools": cache_admin::pool_gauges(&state),
1866 "memory": cache_admin::footprint_json(&state),
1869 }))
1870}
1871
1872#[derive(Deserialize)]
1873struct RequestsQuery {
1874 #[serde(default)]
1875 since: u64,
1876 #[serde(default = "default_requests_limit")]
1877 limit: usize,
1878}
1879
1880fn default_requests_limit() -> usize {
1881 stats::MAX_PAGE
1882}
1883
1884async fn recent_requests(
1892 State(state): State<Arc<AppState>>,
1893 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1894) -> Json<serde_json::Value> {
1895 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1896 Json(serde_json::json!({
1897 "requests": rows,
1898 "next_cursor": cursor,
1899 "missed": missed,
1900 "total": state.stats.recorded_total(),
1901 }))
1902}
1903
1904#[derive(Serialize)]
1905struct CombinedCacheStats {
1906 response_cache: response_cache::CacheStats,
1907 prefix_cache: Option<frink_models::PrefixCacheStats>,
1909}
1910
1911async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1912 Json(CombinedCacheStats {
1913 response_cache: lock_cache(&state.response_cache).stats(),
1914 prefix_cache: state
1915 .prefix_cache
1916 .as_ref()
1917 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1918 })
1919}
1920
1921async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1926 use std::sync::atomic::Ordering;
1927
1928 let cache_stats = lock_cache(&state.response_cache).stats();
1929 let active = state.active();
1930 let requests_total = state.requests_total.load(Ordering::Relaxed);
1931 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1932 let uptime = state.started_at.elapsed().as_secs_f64();
1933
1934 let body = format!(
1935 "# HELP frink_requests_total Total chat completion requests received.\n\
1936 # TYPE frink_requests_total counter\n\
1937 frink_requests_total {requests_total}\n\
1938 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1939 # TYPE frink_request_errors_total counter\n\
1940 frink_request_errors_total {errors_total}\n\
1941 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1942 # TYPE frink_cache_hits_total counter\n\
1943 frink_cache_hits_total {}\n\
1944 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1945 # TYPE frink_cache_misses_total counter\n\
1946 frink_cache_misses_total {}\n\
1947 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1948 # TYPE frink_cache_entries gauge\n\
1949 frink_cache_entries {}\n\
1950 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1951 # TYPE frink_synthetic_weights gauge\n\
1952 frink_synthetic_weights {}\n\
1953 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1954 # TYPE frink_uptime_seconds gauge\n\
1955 frink_uptime_seconds {uptime}\n",
1956 cache_stats.hits,
1957 cache_stats.misses,
1958 cache_stats.entries,
1959 active
1963 .as_ref()
1964 .map(|a| a.is_synthetic() as u8)
1965 .unwrap_or(0),
1966 );
1967
1968 let body = match active
1972 .as_ref()
1973 .and_then(|a| a.expert_store_stats())
1974 {
1975 Some(es) => format!(
1976 "{body}\
1977 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
1978 # TYPE frink_expert_cache_hits_total counter\n\
1979 frink_expert_cache_hits_total {}\n\
1980 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1981 # TYPE frink_expert_cache_misses_total counter\n\
1982 frink_expert_cache_misses_total {}\n\
1983 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
1984 # TYPE frink_expert_cache_evictions_total counter\n\
1985 frink_expert_cache_evictions_total {}\n\
1986 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1987 # TYPE frink_expert_cache_pass_throughs_total counter\n\
1988 frink_expert_cache_pass_throughs_total {}\n\
1989 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1990 # TYPE frink_expert_cache_bytes_read_total counter\n\
1991 frink_expert_cache_bytes_read_total {}\n\
1992 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1993 # TYPE frink_expert_cache_resident_bytes gauge\n\
1994 frink_expert_cache_resident_bytes {}\n",
1995 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1996 ),
1997 None => body,
1998 };
1999
2000 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2005 Some(batcher) => {
2006 let sched = batcher.stats();
2007 format!(
2008 "{body}\
2009 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2010 # TYPE frink_prefill_chunks_total counter\n\
2011 frink_prefill_chunks_total {}\n\
2012 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2013 # TYPE frink_prefill_tokens_total counter\n\
2014 frink_prefill_tokens_total {}\n\
2015 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2016 # TYPE frink_decode_steps_total counter\n\
2017 frink_decode_steps_total {}\n\
2018 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2019 # TYPE frink_scheduler_queue_depth gauge\n\
2020 frink_scheduler_queue_depth {}\n\
2021 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2022 # TYPE frink_scheduler_queue_rejected_total counter\n\
2023 frink_scheduler_queue_rejected_total {}\n\
2024 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2025 # TYPE frink_kv_blocks_total gauge\n\
2026 frink_kv_blocks_total {}\n\
2027 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2028 # TYPE frink_kv_blocks_free gauge\n\
2029 frink_kv_blocks_free {}\n\
2030 # HELP frink_kv_block_size Token positions per KV block.\n\
2031 # TYPE frink_kv_block_size gauge\n\
2032 frink_kv_block_size {}\n\
2033 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2034 # TYPE frink_kv_rejected_too_large_total counter\n\
2035 frink_kv_rejected_too_large_total {}\n\
2036 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2037 # TYPE frink_kv_rejected_context_length_total counter\n\
2038 frink_kv_rejected_context_length_total {}\n\
2039 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2040 # TYPE frink_scheduler_aborted_total counter\n\
2041 frink_scheduler_aborted_total {}\n\
2042 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2043 # TYPE frink_scheduler_max_seqs gauge\n\
2044 frink_scheduler_max_seqs {}\n\
2045 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2046 # TYPE frink_scheduler_prefill_chunk gauge\n\
2047 frink_scheduler_prefill_chunk {}\n",
2048 sched.prefill_chunks,
2049 sched.prefill_tokens,
2050 sched.decode_steps,
2051 sched.queue_depth,
2052 sched.queue_rejected,
2053 sched.kv_blocks_total,
2054 sched.kv_blocks_free,
2055 sched.kv_block_size,
2056 sched.kv_rejected_too_large,
2057 sched.kv_rejected_context_length,
2058 sched.aborted,
2059 sched.max_seqs,
2060 sched.prefill_chunk,
2061 )
2062 }
2063 None => body,
2064 };
2065
2066 (
2067 [(
2068 axum::http::header::CONTENT_TYPE,
2069 "text/plain; version=0.0.4",
2070 )],
2071 body,
2072 )
2073 .into_response()
2074}
2075
2076pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2077
2078pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2084 (
2085 StatusCode::BAD_REQUEST,
2086 Json(serde_json::json!({"error": {
2087 "message": message,
2088 "type": "invalid_request_error",
2089 "param": param,
2090 "code": null,
2091 }})),
2092 )
2093}
2094
2095pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2096 (
2097 StatusCode::NOT_IMPLEMENTED,
2098 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2099 )
2100}
2101
2102pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2103 let status = match e {
2104 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2105 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2109 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2114 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2120 StatusCode::SERVICE_UNAVAILABLE
2121 }
2122 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2126 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2130 };
2131 tracing::warn!("decode error: {e}");
2132 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2133 if let generate::DecodeError::KvBudgetExceeded {
2139 binding,
2140 estimated_bytes,
2141 limit_bytes,
2142 positions,
2143 positions_limit,
2144 ..
2145 } = &e
2146 {
2147 body["error"]["type"] = serde_json::json!("invalid_request_error");
2148 body["error"]["code"] = serde_json::json!(binding);
2149 body["error"]["binding"] = serde_json::json!(binding);
2150 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2151 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2152 body["error"]["positions"] = serde_json::json!(positions);
2153 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2154 }
2155 if let Some(secs) = e.retry_after_secs() {
2159 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2160 }
2161 (status, Json(body))
2162}
2163
2164pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2165 tracing::error!("generation task panicked: {e}");
2166 (
2167 StatusCode::INTERNAL_SERVER_ERROR,
2168 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2169 )
2170}
2171
2172#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2182 model: &Model,
2183 prompt: &str,
2184 params: &GenerationParams,
2185 kv_pool: Option<&generate::KvPoolConfig>,
2186 paged_kv: Option<&generate::PagedKvConfig>,
2187 prefix_cache: Option<&Mutex<PrefixCache>>,
2188 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2189 ceiling: Option<&budget::ContextCeiling>,
2190 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2191 mut emit: impl FnMut(&str),
2192 ) -> Result<(Vec<generate::GeneratedChoice>, generate::Usage), generate::DecodeError> {
2198 let synthetic = model.is_synthetic();
2199 let _lora_lease = lora::lease(model, params.lora.as_deref());
2204 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2205 let params = &{
2211 let mut resolved = params.clone();
2212 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2213 model.encode(text, SpecialTokens::Parse)
2214 });
2215 resolved.reasoning_budget = resolved
2221 .reasoning_budget
2222 .armed(resolved.reasoning, prompt, |text| {
2223 model.encode(text, SpecialTokens::Parse)
2224 })
2225 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2226 resolved
2227 };
2228 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2229 let _metal_private_guard =
2230 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2231 let (finishes, usage) = match model {
2232 Model::Gguf(m) => {
2233 if let Some(batcher) = continuous_batcher {
2234 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2235 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2236 let (finish, _generated_ids, text, usage) = if synthetic {
2237 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2238 } else {
2239 batcher.generate_streaming(
2240 tokens,
2241 params.clone(),
2242 m.stop_tokens.clone(),
2243 Some(|chunk: &str| {
2244 if !chunk.is_empty() {
2245 chunks[0].push(chunk.to_string());
2246 emit(chunk);
2247 }
2248 }),
2249 )?
2250 };
2251 if !text.is_empty() && chunks[0].is_empty() {
2252 chunks[0].push(text);
2253 }
2254 (vec![(finish, Vec::new())], usage)
2261 } else {
2262 generate::generate(
2263 &m.decoder,
2264 m.tokenizer.as_ref(),
2265 &m.stop_tokens,
2266 m.bos_id,
2267 prompt,
2268 params,
2269 kv_pool,
2270 paged_kv,
2271 prefix_cache,
2272 ceiling,
2273 |choice, chunk| {
2274 chunks[choice].push(chunk.to_string());
2275 if !synthetic && choice == 0 {
2279 emit(chunk);
2280 }
2281 },
2282 )?
2283 }
2284 }
2285 Model::Kimi(m) => generate::generate_engine(
2286 &m.engine,
2287 &m.tokenizer,
2288 &m.stop_tokens,
2289 None,
2290 prompt,
2291 params,
2292 |chunk| {
2293 chunks[0].push(chunk.to_string());
2294 if !synthetic {
2295 emit(chunk);
2296 }
2297 },
2298 )?,
2299 Model::Mla(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::Gemma4(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 Model::Glm52(m) => generate::generate_engine(
2328 &m.engine,
2329 &m.tokenizer,
2330 &m.stop_tokens,
2331 m.bos_id,
2332 prompt,
2333 params,
2334 |chunk| {
2335 chunks[0].push(chunk.to_string());
2336 if !synthetic {
2337 emit(chunk);
2338 }
2339 },
2340 )?,
2341 };
2342
2343 let mut full = chunks[0].concat();
2344 if synthetic {
2345 full = format!(
2346 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2347 to serve a real model. Decoded ids -> {full:?}]"
2348 );
2349 emit(&full);
2350 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2351 emit(&full);
2352 }
2353
2354 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2360 let mut out: Vec<generate::GeneratedChoice> = finishes
2361 .into_iter()
2362 .zip(chunks.into_iter().map(|c| c.concat()))
2363 .map(|((finish, logprobs), text)| generate::GeneratedChoice {
2364 finish,
2365 text,
2366 logprobs,
2367 })
2368 .collect();
2369 if let Some(first) = out.first_mut() {
2370 if synthetic {
2378 first.logprobs.clear();
2379 }
2380 first.text = full;
2381 }
2382 Ok((out, usage))
2383}
2384
2385#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2390 model: &Model,
2391 prompt: &str,
2392 params: &GenerationParams,
2393 kv_pool: Option<&generate::KvPoolConfig>,
2394 paged_kv: Option<&generate::PagedKvConfig>,
2395 prefix_cache: Option<&Mutex<PrefixCache>>,
2396 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2397 ceiling: Option<&budget::ContextCeiling>,
2398 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2399 ) -> Result<(Vec<generate::GeneratedChoice>, generate::Usage), generate::DecodeError> {
2402 run_generation_emit(
2403 model,
2404 prompt,
2405 params,
2406 kv_pool,
2407 paged_kv,
2408 prefix_cache,
2409 continuous_batcher,
2410 ceiling,
2411 metal_private_decode_gate,
2412 |_| {},
2413 )
2414}
2415
2416pub(crate) fn prompt_from_messages(
2427 messages: &[ChatMessage],
2428 template: &chat_template::PromptTemplate,
2429 tools: &[ToolDef],
2430 extra: serde_json::Map<String, serde_json::Value>,
2431) -> Result<String, ApiError> {
2432 let rendered = if tools.is_empty() || template.handles_tools() {
2433 template.render(messages, tools, extra)
2434 } else {
2435 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2436 with_preamble.push(ChatMessage {
2437 role: "system".to_string(),
2438 content: Some(MessageContent::Text(tool_preamble(tools))),
2439 tool_calls: None,
2440 tool_call_id: None,
2441 reasoning_content: None,
2442 });
2443 with_preamble.extend_from_slice(messages);
2444 template.render(&with_preamble, &[], extra)
2445 };
2446 rendered.map_err(template_error_response)
2447}
2448
2449fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2454 (
2455 StatusCode::BAD_REQUEST,
2456 Json(serde_json::json!({
2457 "error": {
2458 "message": format!("chat template failed to render: {err}"),
2459 "type": "invalid_request_error",
2460 "param": "messages",
2461 "code": null,
2462 }
2463 })),
2464 )
2465}
2466
2467fn tool_preamble(tools: &[ToolDef]) -> String {
2480 let mut out = String::from(
2481 "You can call tools to help answer the user. To call a tool, respond with \
2482 EXACTLY one line in this format and nothing else:\n\
2483 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2484 object matching that tool's parameters>}}</tool_call>\n\n\
2485 Available tools:\n",
2486 );
2487 for t in tools {
2488 out.push_str(&format!(
2489 "- {}: {}\n parameters (JSON schema): {}\n",
2490 t.function.name,
2491 t.function.description.as_deref().unwrap_or(""),
2492 t.function
2493 .parameters
2494 .as_ref()
2495 .map(|v| v.to_string())
2496 .unwrap_or_else(|| "{}".to_string()),
2497 ));
2498 }
2499 out
2500}
2501
2502fn tool_call_deltas(
2512 events: Vec<crate::policy::parser::ToolCallEvent>,
2513 opened: &std::cell::Cell<usize>,
2514) -> (String, Vec<ToolCallDelta>) {
2515 let mut text = String::new();
2516 let mut deltas = Vec::new();
2517 for event in events {
2518 match event {
2519 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2520 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2521 opened.set(opened.get().max(index + 1));
2522 deltas.push(ToolCallDelta::opening(index, name));
2523 }
2524 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2525 if !fragment.is_empty() {
2526 deltas.push(ToolCallDelta::arguments(index, fragment));
2527 }
2528 }
2529 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2530 }
2531 }
2532 (text, deltas)
2533}
2534
2535fn build_response_message(
2546 text: String,
2547 tools: &[ToolDef],
2548 posture: output::OutputPosture,
2549 base_finish: &'static str,
2550) -> (ChatCompletionResponseMessage, &'static str) {
2551 let parsed = output::parse_output(&text, tools, posture);
2552 let calls: Vec<ToolCallOut> = parsed
2553 .calls
2554 .into_iter()
2555 .enumerate()
2556 .map(|(index, call)| ToolCallOut {
2557 id: format!("call_{index}"),
2558 kind: "function",
2559 function: ToolCallFunctionOut {
2560 name: call.name,
2561 arguments: call.arguments,
2562 },
2563 })
2564 .collect();
2565 if !calls.is_empty() {
2566 return (
2567 ChatCompletionResponseMessage {
2568 role: "assistant",
2569 content: None,
2570 reasoning_content: parsed.reasoning,
2571 tool_calls: Some(calls),
2572 },
2573 "tool_calls",
2574 );
2575 }
2576 (
2577 ChatCompletionResponseMessage {
2578 role: "assistant",
2579 content: Some(parsed.content),
2580 reasoning_content: parsed.reasoning,
2581 tool_calls: None,
2582 },
2583 base_finish,
2584 )
2585}
2586
2587fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2592 let mut history = match &req.session_id {
2593 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2594 None => req.messages.clone(),
2595 };
2596 if req.json_object_mode() {
2597 inject_json_object_system_hint(&mut history);
2598 }
2599 history
2600}
2601
2602fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2603 const HINT: &str =
2604 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2605 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2606 match &mut sys.content {
2607 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2608 s.push_str("\n\n");
2609 s.push_str(HINT);
2610 }
2611 None => {
2612 sys.content = Some(MessageContent::Text(HINT.to_string()));
2613 }
2614 _ => {}
2615 }
2616 } else {
2617 messages.insert(
2618 0,
2619 ChatMessage {
2620 role: "system".to_string(),
2621 content: Some(MessageContent::Text(HINT.to_string())),
2622 tool_calls: None,
2623 tool_call_id: None,
2624 reasoning_content: None,
2625 },
2626 );
2627 }
2628}
2629
2630async fn chat_completions(
2631 State(state): State<Arc<AppState>>,
2632 headers: axum::http::HeaderMap,
2633 Json(req): Json<ChatCompletionRequest>,
2634) -> Response {
2635 let attribution = attribution::Attribution::from_headers(&headers);
2636 state
2637 .requests_total
2638 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2639 let started = std::time::Instant::now();
2640
2641 let request_id = frink_api::next_request_id();
2645 let stream = req.stream.unwrap_or(false);
2646
2647 let refusal = cache_admin::check_admission(&state)
2653 .err()
2654 .or_else(|| req.validate_supported_fields().err());
2655 if let Some(err) = refusal {
2656 state
2657 .request_errors_total
2658 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2659 let response = err.into_response();
2660 state.record_request(stats::Record {
2661 request_id: &request_id,
2662 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2663 model: state.active_model_name(),
2664 status: response.status().as_u16(),
2665 stream,
2666 duration_ms: started.elapsed().as_millis() as u64,
2667 usage: None,
2668 attribution: &attribution,
2669 });
2670 return response;
2671 }
2672
2673 let response = if stream {
2674 chat_completions_stream(
2675 Arc::clone(&state),
2676 req,
2677 request_id.clone(),
2678 started,
2679 attribution.clone(),
2680 )
2681 .await
2682 .into_response()
2683 } else {
2684 chat_completions_full(
2685 Arc::clone(&state),
2686 req,
2687 request_id.clone(),
2688 started,
2689 attribution.clone(),
2690 )
2691 .await
2692 .into_response()
2693 };
2694
2695 if response.status().is_client_error() || response.status().is_server_error() {
2696 state
2697 .request_errors_total
2698 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2699 state.record_request(stats::Record {
2703 request_id: &request_id,
2704 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2705 model: state.active_model_name(),
2708 status: response.status().as_u16(),
2709 stream,
2710 duration_ms: started.elapsed().as_millis() as u64,
2711 usage: None,
2712 attribution: &attribution,
2713 });
2714 }
2715 state.mark_request_finished();
2716
2717 response
2718}
2719
2720async fn chat_completions_full(
2721 state: Arc<AppState>,
2722 req: ChatCompletionRequest,
2723 request_id: String,
2724 started: std::time::Instant,
2725 attribution: attribution::Attribution,
2726) -> Result<Json<ChatCompletionResponse>, ApiError> {
2727 let tools_active = req.tools_active();
2728 let active = state.require_active()?;
2732 let history = resolve_history(&state, &req);
2733 let template = active.generative()?.chat_template();
2734 let kwargs = req.resolve_template_kwargs(&template);
2735 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2736 let mut params =
2744 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2745 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2746 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2747
2748 let (completion, cache_status) = if let Some(cached) = key
2749 .as_ref()
2750 .and_then(|key| lock_cache(&state.response_cache).get(key))
2751 {
2752 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2753 (cached, "hit")
2754 } else {
2755 let (choices, usage) = decode_task::buffered(
2756 decode_task::DecodeHandles::take(&state, &active)?,
2757 prompt.clone(),
2758 params,
2759 )
2760 .await?;
2761
2762 let completion = response_cache::CachedCompletion {
2763 choices: choices.into_iter().map(|c| (c.finish, c.text)).collect(),
2764 usage,
2765 };
2766 let cache_status = match key {
2772 Some(key) => match completion.clone().cacheable() {
2776 Some(cacheable) => {
2777 tracing::debug!("cache miss for key {}", key.digest());
2778 lock_cache(&state.response_cache).put(key, cacheable);
2779 "miss"
2780 }
2781 None => "skip",
2782 },
2783 None => "skip",
2784 };
2785 (completion, cache_status)
2786 };
2787 let content = completion.first_text().to_string();
2790
2791 if req.json_object_mode() {
2792 json_mode::validate_json_object_output(&content)?;
2793 }
2794
2795 if let Some(id) = &req.session_id {
2799 state.sessions.store_reply(
2800 id,
2801 ChatMessage {
2802 role: "assistant".to_string(),
2803 content: Some(MessageContent::Text(content.clone())),
2804 tool_calls: None,
2805 tool_call_id: None,
2806 reasoning_content: None,
2807 },
2808 );
2809 }
2810
2811 let posture = output::OutputPosture::resolve_full(
2816 active.reasoning_format(),
2817 active.tool_call_format(),
2818 &prompt,
2819 );
2820 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2821 let rendered: Vec<ChatCompletionChoice> = completion
2822 .choices
2823 .into_iter()
2824 .enumerate()
2825 .map(|(index, (finish, text))| {
2826 let (message, finish_reason) =
2827 build_response_message(text, tools, posture, finish.as_str());
2828 ChatCompletionChoice {
2829 index,
2830 message,
2831 finish_reason,
2832 }
2833 })
2834 .collect();
2835
2836 state.record_request(stats::Record {
2837 request_id: &request_id,
2838 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2839 model: Some(active.name().to_string()),
2842 status: 200,
2843 stream: false,
2844 duration_ms: started.elapsed().as_millis() as u64,
2845 usage: Some(&completion.usage),
2846 attribution: &attribution,
2847 });
2848
2849 Ok(Json(ChatCompletionResponse {
2850 id: request_id.clone(),
2851 request_id,
2852 object: "chat.completion",
2853 model: req.model,
2854 choices: rendered,
2855 usage: completion.usage,
2856 frink_cache: cache_status,
2857 }))
2858}
2859
2860async fn chat_completions_stream(
2861 state: Arc<AppState>,
2862 req: ChatCompletionRequest,
2863 request_id: String,
2864 started: std::time::Instant,
2865 attribution: attribution::Attribution,
2866) -> Result<Response, ApiError> {
2867 if req.several_choices() {
2877 return Err(unsupported_feature(
2878 "`n` > 1 with `stream` is not implemented: the choices would arrive one after \
2879 another rather than interleaved by `choices[].index`. Send the request without \
2880 `stream`, which serves `n` on this route.",
2881 ));
2882 }
2883 let tools_active = req.tools_active();
2884 let active = state.require_active()?;
2888 let history = resolve_history(&state, &req);
2889 let template = active.generative()?.chat_template();
2890 let kwargs = req.resolve_template_kwargs(&template);
2891 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2892 let model_name = req.model.clone();
2893 let session_id = req.session_id.clone();
2894 let sessions = state.sessions.clone();
2895
2896 let model = Arc::clone(active.generative()?);
2897 let kv_pool = state.kv_pool.clone();
2898 let paged_kv = state.paged_kv.clone();
2899 let prefix_cache = state.prefix_cache.clone();
2900 let batcher = active.batcher.clone();
2901 let ceiling = active.ceiling.clone();
2902 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2903 let mut params =
2904 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2905 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2906 let stats_state = Arc::clone(&state);
2907 let served_model = active.name().to_string();
2910 let posture = output::OutputPosture::resolve_full(
2914 active.reasoning_format(),
2915 active.tool_call_format(),
2916 &prompt,
2917 );
2918 let offered_tools: Vec<ToolDef> = if tools_active {
2921 req.tools.clone()
2922 } else {
2923 Vec::new()
2924 };
2925
2926 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2931 params.cancel = Some(cancel_token.clone());
2932
2933 let overlap = true;
2942
2943 let slot = req
2947 .stream_resumable
2948 .unwrap_or(false)
2949 .then(|| state.streams.register(&request_id));
2950 let emitter = resume::Emitter::new(slot);
2951
2952 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2953 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2957 id: request_id.clone(),
2958 request_id: None,
2959 object: "chat.completion.chunk",
2960 model: model_name.clone(),
2961 choices: vec![ChatCompletionChunkChoice {
2962 index: 0,
2963 delta: ChatCompletionChunkDelta {
2964 role: None,
2965 content: None,
2966 reasoning_content: None,
2967 tool_calls: None,
2968 },
2969 finish_reason: None,
2970 }],
2971 usage: None,
2972 });
2973
2974 tokio::task::spawn_blocking(move || {
2975 let _cancel_guard = cancel_guard;
2978 let tx_chunks = tx.clone();
2979 let orphan_timeout = sse::orphan_timeout_from_env();
2984 let mut first = true;
2985 let head_request_id = request_id.clone();
2986 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
2996 Rc::new(RefCell::new(posture.reasoning_parser()));
2997 let emit_reasoning = Rc::clone(&stream_reasoning);
2998 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3002 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3003 );
3004 let emit_tools = Rc::clone(&stream_tools);
3005 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3009 let emit_streamed_calls = Rc::clone(&streamed_calls);
3010 let result = run_generation_emit(
3011 &model,
3012 &prompt,
3013 ¶ms,
3014 kv_pool.as_ref(),
3015 paged_kv.as_ref(),
3016 prefix_cache.as_deref(),
3017 batcher.as_ref(),
3018 ceiling.as_deref(),
3019 metal_private_decode_gate.as_deref(),
3020 |chunk| {
3021 if !overlap || chunk.is_empty() {
3022 return;
3023 }
3024 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3025 Some(parser) => {
3026 let delta = parser.push(chunk);
3027 (delta.reasoning, delta.content)
3028 }
3029 None => (String::new(), chunk.to_string()),
3030 };
3031 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3035 Some(parser) => {
3036 let (text, calls) =
3037 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3038 (text, calls)
3039 }
3040 None => (content, Vec::new()),
3041 };
3042 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3045 return;
3046 }
3047 let role = if first { Some("assistant") } else { None };
3048 let request_id = first.then(|| head_request_id.clone());
3049 first = false;
3050 let payload = ChatCompletionChunk {
3051 id: head_request_id.clone(),
3052 request_id,
3053 object: "chat.completion.chunk",
3054 model: model_name.clone(),
3055 choices: vec![ChatCompletionChunkChoice {
3056 index: 0,
3057 delta: ChatCompletionChunkDelta {
3058 role,
3059 content: (!content.is_empty()).then_some(content),
3060 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3061 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3062 },
3063 finish_reason: None,
3064 }],
3065 usage: None,
3066 };
3067 if let Err(why) =
3076 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3077 {
3078 if why == sse::SendFailure::Orphaned {
3079 tracing::warn!(
3080 "SSE stream {head_request_id} accepted nothing for the orphan \
3081 deadline; treating it as abandoned"
3082 );
3083 }
3084 if !emitter.is_resumable() {
3096 cancel_token.cancel();
3097 }
3098 }
3099 },
3100 );
3101
3102 let mut pending_request_id = first.then(|| request_id.clone());
3108
3109 match result {
3110 Ok((choices, usage)) => {
3113 let one = choices
3114 .into_iter()
3115 .next()
3116 .expect("a generation produces at least one choice");
3117 let (finish, full_text) = (one.finish, one.text);
3118 if let Some(id) = &session_id {
3119 sessions.store_reply(
3120 id,
3121 ChatMessage {
3122 role: "assistant".to_string(),
3123 content: Some(MessageContent::Text(full_text.clone())),
3124 tool_calls: None,
3125 tool_call_id: None,
3126 reasoning_content: None,
3127 },
3128 );
3129 }
3130 let mut streamed_finish: Option<&'static str> = None;
3136 if overlap {
3137 let tail = stream_reasoning
3138 .borrow_mut()
3139 .as_mut()
3140 .map(|parser| parser.flush())
3141 .unwrap_or_default();
3142 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3143 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3144 let mut events = parser.push(&content);
3145 events.extend(parser.finish());
3146 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3147 content = text;
3148 tool_calls = calls;
3149 }
3150 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3151 let payload = ChatCompletionChunk {
3152 id: request_id.clone(),
3153 request_id: pending_request_id.take(),
3154 object: "chat.completion.chunk",
3155 model: model_name.clone(),
3156 choices: vec![ChatCompletionChunkChoice {
3157 index: 0,
3158 delta: ChatCompletionChunkDelta {
3159 role: None,
3160 content: (!content.is_empty()).then_some(content),
3161 reasoning_content: (!tail.reasoning.is_empty())
3162 .then_some(tail.reasoning),
3163 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3164 },
3165 finish_reason: None,
3166 }],
3167 usage: None,
3168 };
3169 let _ =
3170 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3171 }
3172 if streamed_calls.get() > 0 {
3173 streamed_finish = Some("tool_calls");
3174 }
3175 } else {
3176 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3179 let tool_calls: Vec<ToolCallDelta> = parsed
3180 .calls
3181 .iter()
3182 .enumerate()
3183 .map(|(index, call)| {
3184 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3185 })
3186 .collect();
3187 if !tool_calls.is_empty() {
3188 streamed_finish = Some("tool_calls");
3189 }
3190 if !tool_calls.is_empty()
3191 || !parsed.content.is_empty()
3192 || parsed.reasoning.is_some()
3193 {
3194 let payload = ChatCompletionChunk {
3195 id: request_id.clone(),
3196 request_id: pending_request_id.take(),
3197 object: "chat.completion.chunk",
3198 model: model_name.clone(),
3199 choices: vec![ChatCompletionChunkChoice {
3200 index: 0,
3201 delta: ChatCompletionChunkDelta {
3202 role: Some("assistant"),
3203 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3204 .then(|| parsed.content.clone()),
3205 reasoning_content: parsed.reasoning.clone(),
3206 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3207 },
3208 finish_reason: None,
3209 }],
3210 usage: None,
3211 };
3212 let _ =
3213 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3214 }
3215 }
3216 let final_finish_reason = match streamed_finish {
3220 Some(reason) if finish.as_str() != "length" => reason,
3221 _ => finish.as_str(),
3222 };
3223 let final_payload = ChatCompletionChunk {
3224 id: request_id.clone(),
3225 request_id: pending_request_id.take(),
3226 object: "chat.completion.chunk",
3227 model: model_name,
3228 choices: vec![ChatCompletionChunkChoice {
3229 index: 0,
3230 delta: ChatCompletionChunkDelta {
3231 role: None,
3232 content: None,
3233 reasoning_content: None,
3234 tool_calls: None,
3235 },
3236 finish_reason: Some(final_finish_reason),
3237 }],
3238 usage: Some(usage.clone()),
3239 };
3240 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3241 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3242 stats_state.record_request(stats::Record {
3247 request_id: &request_id,
3248 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3249 model: Some(served_model.clone()),
3250 status: 200,
3251 stream: true,
3252 duration_ms: started.elapsed().as_millis() as u64,
3253 usage: Some(&usage),
3254 attribution: &attribution,
3255 });
3256 }
3257 Err(e) => {
3258 tracing::warn!("decode error on streamed request {request_id}: {e}");
3259 stats_state.record_request(stats::Record {
3265 request_id: &request_id,
3266 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3267 model: Some(served_model.clone()),
3268 status: 500,
3269 stream: true,
3270 duration_ms: started.elapsed().as_millis() as u64,
3271 usage: None,
3272 attribution: &attribution,
3273 });
3274 let payload = ChatCompletionChunk {
3275 id: request_id.clone(),
3276 request_id: pending_request_id.take(),
3277 object: "chat.completion.chunk",
3278 model: model_name,
3279 choices: vec![ChatCompletionChunkChoice {
3280 index: 0,
3281 delta: ChatCompletionChunkDelta {
3282 role: Some("assistant"),
3283 content: Some(format!("[error: {e}]")),
3284 reasoning_content: None,
3285 tool_calls: None,
3286 },
3287 finish_reason: Some("stop"),
3288 }],
3289 usage: None,
3290 };
3291 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3292 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3293 }
3294 }
3295 drop(emitter);
3299 });
3300
3301 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3302 Ok((
3325 [(
3326 axum::http::HeaderName::from_static("x-accel-buffering"),
3327 axum::http::HeaderValue::from_static("no"),
3328 )],
3329 Sse::new(stream),
3330 )
3331 .into_response())
3332}
3333
3334fn protected_routes() -> Router<Arc<AppState>> {
3365 use frink_api::routes;
3366
3367 Router::new()
3368 .route(routes::V1_MODELS, get(list_models))
3369 .route(routes::V1_RESPONSES, post(responses::responses))
3373 .route(
3374 &axum_path(routes::V1_RESPONSE),
3375 get(responses::responses_get),
3376 )
3377 .route(
3378 &axum_path(routes::V1_RESPONSE_CANCEL),
3379 post(responses::responses_cancel),
3380 )
3381 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3382 .route(routes::V1_STATS, get(serving_stats))
3383 .route(routes::V1_REQUESTS, get(recent_requests))
3384 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3385 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3386 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3387 .route(
3388 routes::LORA_ADAPTERS,
3389 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3390 )
3391 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3392 .route(routes::V1_CANCEL, post(cancel_generation))
3396 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3401 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3402 .route(routes::V1_MESSAGES, post(anthropic::messages))
3403 .route(
3404 routes::V1_MESSAGES_COUNT_TOKENS,
3405 post(anthropic::count_tokens),
3406 )
3407 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3408 .route(routes::COMPLETION, post(completion::completion))
3413 .route(routes::COMPLETIONS, post(completion::completion))
3414 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3415 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3416 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3422 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3423 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3424 .route(routes::V1_RERANK, post(rerank::rerank))
3428 .route(routes::RERANK, post(rerank::rerank))
3429 .route(routes::CACHE_STATS, get(cache_stats))
3430 .route(routes::METRICS, get(metrics))
3431 .route(routes::ADMIN_MODELS, get(admin::models))
3436 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3437 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3438 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3439 .route(routes::ADMIN_TASKS, get(admin::tasks))
3440 .route(&admin::cancel_route(), post(admin::cancel_task))
3441 .route(routes::ADMIN_STATS, get(admin::stats))
3442 .merge(conversations::router())
3446}
3447
3448fn axum_path(template: &str) -> String {
3449 let mut out = String::with_capacity(template.len());
3450 let mut rest = template;
3451 while let Some(open) = rest.find('{') {
3452 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3453 break;
3454 };
3455 out.push_str(&rest[..open]);
3456 out.push(':');
3457 out.push_str(&rest[open + 1..close]);
3458 rest = &rest[close + 1..];
3459 }
3460 out.push_str(rest);
3461 out
3462}
3463
3464async fn cancel_generation(
3473 State(state): State<Arc<AppState>>,
3474 Json(req): Json<frink_api::CancelGenerationRequest>,
3475) -> Response {
3476 let cancelled = state.cancels.cancel(&req.request_id);
3477 let status = if cancelled {
3478 StatusCode::OK
3479 } else {
3480 StatusCode::NOT_FOUND
3481 };
3482 let detail = if cancelled {
3483 "the generation was asked to stop; it ends at its next token".to_string()
3484 } else {
3485 "no generation with that request_id is running -- it has already \
3486 finished, was never issued, or was served by a path that does \
3487 not register for cancellation"
3488 .to_string()
3489 };
3490 (
3491 status,
3492 Json(frink_api::CancelGenerationResponse {
3493 request_id: req.request_id,
3494 cancelled,
3495 detail,
3496 }),
3497 )
3498 .into_response()
3499}
3500
3501type Activated = (
3505 Loaded,
3506 Option<serving::batch::ContinuousBatcher>,
3507 Option<Arc<budget::ContextCeiling>>,
3508);
3509
3510fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3528 let mut batcher = serving::batch::BatcherConfig::from_env();
3529 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3530 return batcher;
3533 }
3534 let Some(path) = path else {
3535 return batcher;
3536 };
3537 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3543 let Some((priced, gguf_ctx, source)) = priced else {
3544 return batcher;
3545 };
3546 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3547 tracing::warn!(
3551 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3552 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3553 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3554 admit on a number you choose.",
3555 priced.weights_bytes,
3556 priced.device_budget_bytes,
3557 );
3558 return batcher;
3559 };
3560 tracing::info!("{source}");
3561 tracing::info!("{}", derived.fit);
3562 let adopted = budget::apply_derived(&mut batcher, &derived);
3563 if adopted.max_context {
3564 tracing::info!(
3565 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3566 override with FRINK_CB_MAX_CONTEXT",
3567 derived.max_context
3568 );
3569 }
3570 if adopted.kv_blocks {
3571 tracing::info!(
3572 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3573 derived.kv_blocks,
3574 batcher.kv_block_size
3575 );
3576 }
3577 if let Some(narrowed) = adopted.max_context_narrowed {
3578 tracing::info!(
3579 "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",
3580 batcher.kv_blocks.unwrap_or_default(),
3581 batcher.kv_block_size
3582 );
3583 }
3584 batcher
3585}
3586
3587pub(crate) fn activate_loaded_model(
3596 loaded: model::LoadedModel,
3597 enable_continuous_batching: bool,
3598 path: Option<&str>,
3599 paged_kv: Option<&generate::PagedKvConfig>,
3600) -> Activated {
3601 match loaded {
3602 model::LoadedModel::Gguf(g) => {
3603 let decoder = Arc::new(g.decoder);
3604 let tokenizer = Arc::new(g.tokenizer);
3605 let config = price_batcher_config(path);
3606 let shape =
3612 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3613 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3614 let batcher = if enable_continuous_batching {
3615 tracing::info!(
3616 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3617 (stop sequences use the same pending-buffer trim as the private generate loop)"
3618 );
3619 let tok = Arc::clone(&tokenizer);
3620 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3621 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3622 Arc::clone(&decoder),
3623 decode,
3624 config,
3625 Arc::clone(&ceiling),
3626 paged_kv.cloned(),
3627 ))
3628 } else {
3629 None
3630 };
3631 (
3632 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3633 decoder,
3634 tokenizer,
3635 stop_tokens: g.stop_tokens,
3636 bos_id: g.bos_id,
3637 is_synthetic: g.is_synthetic,
3638 chat_template: g.chat_template,
3639 }))),
3640 batcher,
3641 Some(ceiling),
3642 )
3643 }
3644 model::LoadedModel::Kimi(k) => (
3645 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3646 engine: k.engine,
3647 tokenizer: k.tokenizer,
3648 stop_tokens: k.stop_tokens,
3649 chat_template: k.chat_template,
3650 }))),
3651 None,
3652 None,
3653 ),
3654 model::LoadedModel::Mla(m) => (
3655 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3656 engine: m.engine,
3657 tokenizer: m.tokenizer,
3658 stop_tokens: m.stop_tokens,
3659 bos_id: m.bos_id,
3660 name: m.name,
3661 chat_template: m.chat_template,
3662 }))),
3663 None,
3664 None,
3665 ),
3666 model::LoadedModel::Gemma4(m) => (
3667 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3668 engine: m.engine,
3669 tokenizer: m.tokenizer,
3670 stop_tokens: m.stop_tokens,
3671 bos_id: m.bos_id,
3672 name: m.name,
3673 chat_template: m.chat_template,
3674 }))),
3675 None,
3676 None,
3677 ),
3678 model::LoadedModel::Glm52(g) => (
3679 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3680 engine: g.engine,
3681 tokenizer: g.tokenizer,
3682 stop_tokens: g.stop_tokens,
3683 bos_id: g.bos_id,
3684 name: g.name,
3685 chat_template: g.chat_template,
3686 }))),
3687 None,
3688 None,
3689 ),
3690 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3695 }
3696}
3697
3698struct StartupModels {
3705 loaded: model::LoadedModel,
3706 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3707}
3708
3709fn continuous_batching_env() -> Option<bool> {
3710 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3711 .ok()
3712 .map(|v| v.trim().to_ascii_lowercase())
3713 .as_deref()
3714 {
3715 None => None,
3716 Some("1" | "true" | "yes" | "on") => Some(true),
3717 Some("0" | "false" | "no" | "off") => Some(false),
3718 _ => None,
3719 }
3720}
3721
3722fn metal_private_decode_active() -> bool {
3723 #[cfg(feature = "metal")]
3724 {
3725 BUILT_WITH_METAL
3726 && frink_metal::attn::metal_attn_enabled()
3727 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3728 }
3729 #[cfg(not(feature = "metal"))]
3730 {
3731 false
3732 }
3733}
3734
3735fn continuous_batching_compatible(
3736 loaded: &model::LoadedModel,
3737 kv_pool: &Option<generate::KvPoolConfig>,
3738 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3739 paged_kv: &Option<generate::PagedKvConfig>,
3740) -> bool {
3741 matches!(loaded, model::LoadedModel::Gguf(_))
3742 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3743}
3744
3745fn resolve_continuous_batching_enabled(
3746 loaded: &model::LoadedModel,
3747 kv_pool: &Option<generate::KvPoolConfig>,
3748 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3749 paged_kv: &Option<generate::PagedKvConfig>,
3750) -> bool {
3751 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3752 return false;
3753 }
3754 match continuous_batching_env() {
3755 Some(true) => true,
3756 Some(false) => false,
3757 None => metal_private_decode_active(),
3758 }
3759}
3760
3761fn acquire_metal_private_decode_gate(
3762 gate: Option<&std::sync::Mutex<()>>,
3763 used_batcher: bool,
3764) -> Option<std::sync::MutexGuard<'_, ()>> {
3765 if used_batcher {
3766 None
3767 } else {
3768 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3769 }
3770}
3771
3772fn build_app_state(
3773 models: StartupModels,
3774 kv_pool: Option<generate::KvPoolConfig>,
3775 paged_kv: Option<generate::PagedKvConfig>,
3776 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3777 enable_continuous_batching: bool,
3778 mcp: Option<mcp::LoadedMcpConfig>,
3779 detection: Arc<health::Detection>,
3780) -> AppState {
3781 let StartupModels { loaded, embedding } = models;
3782 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3783 let (loaded, batcher, ceiling) = activate_loaded_model(
3784 loaded,
3785 enable_continuous_batching,
3786 configured_path.as_deref(),
3787 paged_kv.as_ref(),
3788 );
3789 let id = startup_model_id();
3795 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3796 {
3797 None
3798 } else {
3799 tracing::info!(
3800 "Metal private-loop decode will serialize concurrent requests until \
3801 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3802 );
3803 Some(Arc::new(std::sync::Mutex::new(())))
3804 };
3805 AppState {
3806 embedding,
3807 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3808 id,
3809 loaded,
3810 batcher,
3811 ceiling,
3812 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3813 }))),
3814 paged_kv,
3815 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3816 tasks: Arc::new(tasks::TaskRegistry::new()),
3817 cancels: Arc::new(cancel::CancelRegistry::new()),
3818 stats: stats::Stats::new(),
3819 streams: resume::StreamRegistry::new(),
3820 model_dir: admin::model_dirs().into_iter().next(),
3821 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3822 kv_pool,
3823 prefix_cache,
3824 sessions: session::SessionStore::new(),
3825 requests_total: std::sync::atomic::AtomicU64::new(0),
3826 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3827 started_at: std::time::Instant::now(),
3828 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3829 detection,
3830 mcp,
3831 continuous_batching_enabled: enable_continuous_batching,
3832 metal_private_decode_gate,
3833 loading_model: Mutex::new(None),
3834 last_load_error: Mutex::new(None),
3835 serving: Mutex::new(crate::stats::ServingStats::default()),
3836 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3837 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3838 started_unix: unix_now(),
3839 }
3840}
3841
3842fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3850 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3851 return Ok(None);
3852 };
3853 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3854 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3855 tracing::info!(
3856 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3857 model.name(),
3858 model.architecture(),
3859 model.n_embd(),
3860 model.pooling_type().name(),
3861 model.n_ctx_train(),
3862 );
3863 Ok(Some(Arc::new(model)))
3864}
3865
3866fn unix_now() -> u64 {
3871 std::time::SystemTime::now()
3872 .duration_since(std::time::UNIX_EPOCH)
3873 .map(|d| d.as_secs())
3874 .unwrap_or(0)
3875}
3876
3877fn startup_model_id() -> Option<String> {
3882 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3883 let configured = std::fs::canonicalize(&configured).ok()?;
3884 admin::discover(&admin::model_dirs())
3885 .into_iter()
3886 .find(|d| {
3887 std::fs::canonicalize(&d.path)
3888 .map(|p| p == configured)
3889 .unwrap_or(false)
3890 })
3891 .map(|d| d.id)
3892}
3893
3894fn init_cpu_pool() {
3903 match frink_core::threads::init_cpu_pool() {
3904 Some(n) => eprintln!(
3905 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
3906 frink_core::threads::perf_core_count()
3907 ),
3908 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
3909 }
3910}
3911
3912fn announce_ready(addr: SocketAddr, scheme: &str) {
3925 use std::io::Write;
3926 let ready =
3927 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3928 let mut stdout = std::io::stdout().lock();
3929 let _ = writeln!(stdout, "{}", ready.to_line());
3930 let _ = stdout.flush();
3931}
3932
3933async fn shutdown_signal(exit_on_stdin_close: bool) {
3945 if !exit_on_stdin_close {
3946 std::future::pending::<()>().await;
3947 return;
3948 }
3949 let _ = tokio::task::spawn_blocking(|| {
3950 use std::io::Read;
3951 let mut sink = [0u8; 256];
3952 let mut stdin = std::io::stdin().lock();
3953 loop {
3954 match stdin.read(&mut sink) {
3955 Ok(0) => break,
3957 Ok(_) => continue,
3959 Err(e) => {
3960 tracing::warn!("stdin read failed ({e}); treating it as closed");
3961 break;
3962 }
3963 }
3964 }
3965 })
3966 .await;
3967 tracing::info!("stdin closed; shutting down");
3968}
3969
3970fn tokio_worker_threads() -> usize {
3975 std::env::var("FRINK_TOKIO_WORKERS")
3976 .ok()
3977 .and_then(|v| v.trim().parse::<usize>().ok())
3978 .filter(|n| *n > 0)
3979 .unwrap_or(2)
3980}
3981
3982fn install_ring_crypto_provider() {
3999 let _ = rustls::crypto::ring::default_provider().install_default();
4000}
4001
4002pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4008 if args.list_devices {
4009 frink_models::devices::print_available_devices();
4010 return Ok(());
4011 }
4012 apply_cli_overrides(&args)?;
4013
4014 let _instance = {
4018 use frink_core::instance::{register, InstancePolicy};
4019 let policy = if args.allow_multiple_instances {
4020 InstancePolicy::Multi
4021 } else {
4022 InstancePolicy::from_env_or(InstancePolicy::Single)
4023 };
4024 let model = std::env::var("FRINK_MODEL_PATH").ok();
4025 register(
4026 "server",
4027 model.as_deref(),
4028 frink_core::instance::current_backend(),
4029 policy,
4030 )
4031 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4032 };
4033
4034 let journal = journal::Journal::from_env();
4035 eprintln!(
4036 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4037 journal.path()
4038 );
4039 journal.append(&journal::Record::session_start(
4040 env!("CARGO_PKG_VERSION"),
4041 std::process::id(),
4042 ));
4043 journal::install_panic_hook(journal.clone());
4044
4045 let mcp_config_path = args.mcp_config.clone();
4046 let exit_on_stdin_close = args.exit_on_stdin_close
4047 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4048 .map(|v| v == "1")
4049 .unwrap_or(false);
4050
4051 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4055 init_cpu_pool();
4056
4057 let runtime = tokio::runtime::Builder::new_multi_thread()
4058 .worker_threads(tokio_worker_threads())
4059 .enable_all()
4060 .build()?;
4061 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4062
4063 let reason = match &result {
4064 Ok(()) => "normal".to_string(),
4065 Err(e) => e.to_string(),
4066 };
4067 journal.append(&journal::Record::session_exit(reason));
4068
4069 runtime.shutdown_background();
4074
4075 result
4076}
4077
4078async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4079 let _ = tracing_subscriber::fmt::try_init();
4085
4086 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4096 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4097 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4098 .map(|v| v == "1")
4099 .unwrap_or(false);
4100 if let Err(msg) =
4101 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4102 {
4103 anyhow::bail!(msg);
4104 }
4105
4106 let embedding_model = load_embedding_model()?;
4113
4114 let mut loaded = model::load()?;
4115 match &loaded {
4116 model::LoadedModel::Gguf(g) => tracing::info!(
4117 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4118 g.decoder.config.name,
4119 g.is_synthetic,
4120 g.tokenizer.kind()
4121 ),
4122 model::LoadedModel::Kimi(k) => tracing::info!(
4123 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4124 k.tokenizer.vocab_size()
4125 ),
4126 model::LoadedModel::Mla(m) => tracing::info!(
4127 "loaded MLA GGUF '{}' (tokenizer={})",
4128 m.name,
4129 m.tokenizer.kind()
4130 ),
4131 model::LoadedModel::Gemma4(m) => tracing::info!(
4132 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4133 m.name,
4134 m.tokenizer.kind()
4135 ),
4136 model::LoadedModel::Glm52(g) => tracing::info!(
4137 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4138 g.name,
4139 g.tokenizer.kind()
4140 ),
4141 model::LoadedModel::Encoder(_) => {}
4144 }
4145 let metal_default_moe_budget = {
4153 #[cfg(feature = "metal")]
4154 {
4155 frink_core::metal_dense_enabled()
4156 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4157 }
4158 #[cfg(not(feature = "metal"))]
4159 {
4160 false
4161 }
4162 };
4163 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4164 let budget: u64 = budget_str
4165 .parse()
4166 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4167 match &mut loaded {
4168 model::LoadedModel::Gguf(g) => {
4169 tracing::info!(
4170 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4171 (CUDA and/or Metal matvecs when built with the matching feature)"
4172 );
4173 g.decoder.gpu_vram_budget_bytes = Some(budget);
4174 }
4175 model::LoadedModel::Kimi(_) => {
4176 tracing::warn!(
4177 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4178 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4179 );
4180 }
4181 model::LoadedModel::Mla(_) => {
4182 tracing::warn!(
4183 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4184 FFN path only today; ignoring expert VRAM budget"
4185 );
4186 }
4187 model::LoadedModel::Gemma4(_) => {
4188 tracing::warn!(
4189 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4190 ignoring expert VRAM budget"
4191 );
4192 }
4193 model::LoadedModel::Glm52(_) => {
4194 tracing::warn!(
4195 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4196 GPU expert placement not wired yet; ignoring"
4197 );
4198 }
4199 model::LoadedModel::Encoder(_) => {
4200 tracing::warn!(
4201 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4202 it has no routed experts to place; ignoring"
4203 );
4204 }
4205 }
4206 } else if metal_default_moe_budget {
4207 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4211 if let model::LoadedModel::Gguf(g) = &mut loaded {
4212 tracing::info!(
4213 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4214 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4215 );
4216 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4217 }
4218 }
4219 #[cfg(feature = "cuda")]
4220 {
4221 if frink_core::cuda_dense_enabled() {
4222 tracing::info!(
4223 "CUDA dense matvec enabled for WeightMatrix::apply \
4224 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4225 );
4226 } else {
4227 tracing::info!(
4228 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4229 );
4230 }
4231 }
4232 #[cfg(feature = "metal")]
4233 {
4234 if frink_core::metal_dense_enabled() {
4235 tracing::info!(
4236 "Metal dense matvec enabled for WeightMatrix::apply \
4237 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4238 );
4239 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4240 Some("1") | Some("true") | Some("on") | Some("attn") => {
4241 tracing::info!(
4242 "Metal fused attention requested (FRINK_METAL_ATTN): \
4243 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4244 );
4245 }
4246 _ => {}
4247 }
4248 tracing::info!(
4249 "Metal greedy GPU argmax: temperature<=0 folds \
4250 final_norm+lm_head+argmax into the dense stack"
4251 );
4252 } else {
4253 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4254 }
4255 }
4256 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4267 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4268 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4269 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4270 panic!(
4271 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4272 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4273 );
4274 }
4275 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4276 (Ok(blocks), Ok(block_size), Err(_)) => {
4277 let total_blocks: usize = blocks
4278 .parse()
4279 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4280 let block_size: usize = block_size
4281 .parse()
4282 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4283 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4289 .ok()
4290 .map(|v| {
4291 v.parse()
4292 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4293 })
4294 .unwrap_or(0);
4295 tracing::info!(
4296 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4297 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4298 );
4299 Some(generate::KvPoolConfig {
4300 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4301 queue_wait: Duration::from_millis(queue_wait_ms),
4302 })
4303 }
4304 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4305 let block_size: usize = block_size
4306 .parse()
4307 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4308 let budget: u64 = byte_budget
4309 .parse()
4310 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4311 let cfg = match &loaded {
4312 model::LoadedModel::Gguf(g) => &g.decoder.config,
4313 model::LoadedModel::Kimi(_)
4314 | model::LoadedModel::Mla(_)
4315 | model::LoadedModel::Gemma4(_)
4316 | model::LoadedModel::Glm52(_)
4317 | model::LoadedModel::Encoder(_) => {
4318 panic!(
4319 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4320 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4321 );
4322 }
4323 };
4324 let bytes_per_block = block_size
4325 * cfg.kv_heads_all_layers()
4326 * (cfg.head_dim + cfg.v_head_dim())
4327 * std::mem::size_of::<f32>();
4328 assert!(
4329 bytes_per_block > 0,
4330 "derived KV block byte size must be positive (check model config and block size)"
4331 );
4332 let total_blocks = (budget as usize / bytes_per_block).max(1);
4333 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4334 .ok()
4335 .map(|v| {
4336 v.parse()
4337 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4338 })
4339 .unwrap_or(0);
4340 tracing::info!(
4341 "KV cache block pool enabled from byte budget: {budget} bytes / \
4342 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4343 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4344 cfg.n_layers
4345 );
4346 Some(generate::KvPoolConfig {
4347 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4348 queue_wait: Duration::from_millis(queue_wait_ms),
4349 })
4350 }
4351 (Err(_), Err(_), Err(_)) => None,
4352 (Err(_), Ok(_), Err(_)) => panic!(
4353 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4354 (or unset all three to disable KV cache pooling)"
4355 ),
4356 (Ok(_), Ok(_), Ok(_)) => {
4357 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4358 }
4359 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4360 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4361 set together (or neither, to disable KV cache pooling)"
4362 ),
4363 };
4364 let paged_kv = match (
4370 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4371 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4372 ) {
4373 (Ok(blocks), Ok(block_size)) => {
4374 assert!(
4375 kv_pool.is_none(),
4376 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4377 mutually exclusive: both bound the same KV memory, by different means. \
4378 Set one."
4379 );
4380 let blocks_per_layer: usize = blocks
4394 .parse()
4395 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4396 let block_size: usize = block_size
4397 .parse()
4398 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4399 let gguf = match &loaded {
4400 model::LoadedModel::Gguf(g) => g,
4401 _ => panic!(
4402 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4403 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4404 ),
4405 };
4406 let cfg = &gguf.decoder.config;
4407 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4408 .ok()
4409 .map(|v| {
4410 v.parse()
4411 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4412 })
4413 .unwrap_or(0);
4414 tracing::info!(
4415 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4416 layer across {} layers, shared by all concurrent requests, \
4417 {queue_wait_ms}ms admission queue wait",
4418 cfg.n_layers
4419 );
4420 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4425 block_size,
4426 ))));
4427 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4436 crate::policy::parser::ToolCallFormat::infer(
4437 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4438 )
4439 .opener(),
4440 |text| {
4441 gguf.tokenizer
4442 .encode(text, SpecialTokens::Parse)
4443 .into_iter()
4444 .map(|t| t as u32)
4445 .collect()
4446 },
4447 );
4448 if let Some(id) = anchor_token {
4449 tracing::info!(
4450 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4451 window stops short of where its next turn rejoins"
4452 );
4453 }
4454 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4455 .ok()
4456 .map(|v| {
4457 v.parse()
4458 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4459 })
4460 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4461 if let Some(window) = cfg.uniform_sliding_window() {
4462 tracing::info!(
4463 "Paged KV window slide enabled: every layer slides by {window} every \
4464 {slide_interval} decode steps, so a request holds its prompt and a \
4465 window rather than its whole context"
4466 );
4467 } else if cfg.kv_block_window().is_some() {
4468 tracing::info!(
4469 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4470 and a page group holds one block in every layer"
4471 );
4472 }
4473 Some(generate::PagedKvConfig {
4474 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4477 queue_wait: Duration::from_millis(queue_wait_ms),
4478 radix,
4479 anchor_token,
4480 slide_interval,
4481 })
4482 }
4483 (Err(_), Err(_)) => None,
4484 _ => panic!(
4485 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4486 (or neither, to disable paged KV)"
4487 ),
4488 };
4489 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4496 let max_entries: usize = v
4497 .parse()
4498 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4499 if kv_pool.is_some() {
4500 tracing::warn!(
4501 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4502 caching will never be consulted while a KV pool is configured"
4503 );
4504 }
4505 assert!(
4512 paged_kv.is_none(),
4513 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4514 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4515 produce, so the cache could never hit. Set one."
4516 );
4517 tracing::info!(
4518 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4519 all requests"
4520 );
4521 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4522 });
4523 if matches!(
4524 loaded,
4525 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4526 ) && (kv_pool.is_some() || prefix_cache.is_some())
4527 {
4528 tracing::warn!(
4529 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4530 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4531 frink_models::engine's module docs"
4532 );
4533 }
4534 let enable_cb =
4535 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4536 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4537 tracing::info!(
4538 "continuous batching enabled by default on Metal for safe parallel serving \
4539 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4540 );
4541 }
4542 if continuous_batching_env() == Some(true)
4543 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4544 && (kv_pool.is_some() || prefix_cache.is_some())
4545 {
4546 tracing::warn!(
4547 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4548 (those modes keep the private generate path)"
4549 );
4550 }
4551 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4552 if let Ok(chunk) = n.parse::<usize>() {
4553 if chunk > 0 {
4554 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4555 }
4556 }
4557 }
4558 if matches!(
4559 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4560 Some("1")
4561 ) {
4562 tracing::warn!(
4563 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4564 (minimal spill; full layer offload still planned)"
4565 );
4566 }
4567
4568 let mcp = match mcp_config_path {
4569 Some(path) => {
4570 let loaded = mcp::load_mcp_config(&path)?;
4571 tracing::info!(
4572 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4573 loaded.path,
4574 loaded.servers.len()
4575 );
4576 Some(loaded)
4577 }
4578 None => None,
4579 };
4580
4581 let detection = health::Detection::spawn();
4585
4586 let state = Arc::new(build_app_state(
4587 StartupModels {
4588 loaded,
4589 embedding: embedding_model,
4590 },
4591 kv_pool,
4592 paged_kv,
4593 prefix_cache,
4594 enable_cb,
4595 mcp,
4596 detection,
4597 ));
4598
4599 use frink_api::routes;
4603
4604 let public = Router::new().route(routes::HEALTH, get(health));
4609
4610 let mut protected = protected_routes();
4611
4612 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4617 tracing::info!("API key auth enabled");
4618 let auth = limits::AuthConfig {
4619 api_key: Arc::new(key),
4620 };
4621 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4622 auth,
4623 limits::require_api_key,
4624 ));
4625 }
4626 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4627 let rpm: u32 = rpm
4628 .parse()
4629 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4630 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4631 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4632 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4633 limiter,
4634 limits::rate_limit,
4635 ));
4636 }
4637 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4645 let origins = security::parse_cors_origins(&spec)
4646 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4647 tracing::info!(
4648 "CORS enabled: {} allow-listed origin(s) ({})",
4649 origins.len(),
4650 spec
4651 );
4652 let cors = tower_http::cors::CorsLayer::new()
4653 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4654 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4655 .allow_headers([
4656 axum::http::header::CONTENT_TYPE,
4657 axum::http::header::AUTHORIZATION,
4658 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4664 axum::http::HeaderName::from_static("last-event-id"),
4668 ]);
4669 protected = protected.route_layer(cors);
4670 }
4671
4672 let app = public
4676 .merge(protected)
4677 .layer(axum::middleware::from_fn(limits::retry_after))
4678 .with_state(state);
4679
4680 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4686 install_ring_crypto_provider();
4687 match tls_paths {
4692 Some(paths) => {
4693 let config =
4694 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4695 .await
4696 .map_err(|e| {
4697 anyhow::anyhow!(
4698 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4699 paths.cert,
4700 paths.key
4701 )
4702 })?;
4703 let socket_addr: std::net::SocketAddr = addr
4704 .parse()
4705 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4706 let listener = std::net::TcpListener::bind(socket_addr)?;
4707 listener.set_nonblocking(true)?;
4715 let bound = listener.local_addr()?;
4716 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4717 announce_ready(bound, "https");
4718
4719 let handle = axum_server::Handle::new();
4720 let shutdown_handle = handle.clone();
4721 tokio::spawn(async move {
4722 shutdown_signal(exit_on_stdin_close).await;
4723 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4724 });
4725 axum_server::from_tcp_rustls(listener, config)?
4726 .handle(handle)
4727 .serve(app.into_make_service())
4728 .await?;
4729 }
4730 None => {
4731 let listener = tokio::net::TcpListener::bind(&addr).await?;
4732 let bound = listener.local_addr()?;
4733 tracing::info!("frink-server listening on {bound}");
4734 announce_ready(bound, "http");
4735 axum::serve(listener, app)
4736 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4737 .await?;
4738 }
4739 }
4740 Ok(())
4741}
4742
4743#[cfg(test)]
4744pub(crate) mod tests {
4745 use super::*;
4746 use frink_models::config::test_dense_fixture;
4747
4748 #[test]
4749 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4750 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4751 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4752 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4753 assert_eq!(parsed.port, 51999);
4754 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4755 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4757 }
4758
4759 fn test_model() -> Model {
4760 let cfg = test_dense_fixture();
4764 Model::Gguf(GgufModel {
4765 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4766 tokenizer: Arc::new(ServerTokenizer::Byte),
4767 stop_tokens: StopTokens::default(),
4768 bos_id: None,
4769 is_synthetic: true,
4770 chat_template: chat_template::PromptTemplate::plain(),
4771 })
4772 }
4773
4774 fn greedy_params(max_tokens: usize) -> GenerationParams {
4775 GenerationParams {
4776 wants_logprobs: false,
4777 n: 1,
4778 reasoning: None,
4779 max_tokens,
4780 sampling: SamplingParams::default(),
4781 seed: 1,
4782 stop: Vec::new(),
4783 stop_token_ids: Vec::new(),
4784 json_object: false,
4785 grammar: None,
4786 cancel: None,
4787 ignore_eos: false,
4788 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4789 lora: None,
4790 }
4791 }
4792
4793 fn test_model_full_byte_vocab() -> Model {
4797 test_model_full_byte_vocab_with_eos(None)
4798 }
4799
4800 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4807 let mut cfg = test_dense_fixture();
4808 cfg.vocab_size = 256;
4809 Model::Gguf(GgufModel {
4810 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4811 tokenizer: Arc::new(ServerTokenizer::Byte),
4812 stop_tokens: StopTokens::from_eos(eos),
4813 bos_id: None,
4814 is_synthetic: true,
4815 chat_template: chat_template::PromptTemplate::plain(),
4816 })
4817 }
4818
4819 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4823 AppState {
4824 embedding: None,
4825 paged_kv: None,
4826 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4827 id: None,
4828 loaded: Loaded::Generative(Arc::new(model)),
4829 batcher: None,
4830 ceiling: None,
4831 checkpoint_path: None,
4832 }))),
4833 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4834 tasks: Arc::new(tasks::TaskRegistry::new()),
4835 cancels: Arc::new(cancel::CancelRegistry::new()),
4836 stats: stats::Stats::new(),
4837 streams: resume::StreamRegistry::new(),
4838 model_dir: None,
4839 response_cache: Mutex::new(response_cache),
4840 kv_pool: None,
4841 prefix_cache: None,
4842 sessions: session::SessionStore::new(),
4843 requests_total: std::sync::atomic::AtomicU64::new(0),
4844 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4845 started_at: std::time::Instant::now(),
4846 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4847 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4848 mcp: None,
4849 continuous_batching_enabled: false,
4850 metal_private_decode_gate: None,
4851 loading_model: Mutex::new(None),
4852 last_load_error: Mutex::new(None),
4853 serving: Mutex::new(crate::stats::ServingStats::default()),
4854 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4855 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4856 started_unix: unix_now(),
4857 }
4858 }
4859
4860 pub(crate) fn test_app() -> Router {
4869 test_app_with_state(Arc::new(test_state(
4870 test_model_full_byte_vocab(),
4871 ResponseCache::new(1000, Duration::from_secs(3600)),
4872 )))
4873 }
4874
4875 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4878 Router::new()
4888 .route(frink_api::routes::HEALTH, get(health))
4889 .merge(protected_routes())
4890 .with_state(state)
4891 }
4892
4893 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4894 let mut cfg = test_dense_fixture();
4895 cfg.name = name;
4896 cfg.vocab_size = vocab_size;
4897 Model::Gguf(GgufModel {
4898 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4899 tokenizer: Arc::new(ServerTokenizer::Byte),
4900 stop_tokens: StopTokens::default(),
4901 bos_id: None,
4902 is_synthetic: true,
4903 chat_template: chat_template::PromptTemplate::plain(),
4904 })
4905 }
4906
4907 fn model_with_template(name: &'static str, source: &str) -> Model {
4911 let mut cfg = test_dense_fixture();
4912 cfg.name = name;
4913 cfg.vocab_size = 256;
4914 Model::Gguf(GgufModel {
4915 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4916 tokenizer: Arc::new(ServerTokenizer::Byte),
4917 stop_tokens: StopTokens::default(),
4918 bos_id: None,
4919 is_synthetic: true,
4920 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4921 Some(source),
4922 Some("qwen3"),
4923 false,
4924 true,
4925 None,
4926 None,
4927 ),
4928 })
4929 }
4930
4931 #[tokio::test]
4940 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4941 let strict = "{% if messages | length > 1 %}\
4944 {{ raise_exception('this template takes one turn') }}\
4945 {% endif %}{{ messages[0].content }}";
4946 let state = Arc::new(test_state(
4947 model_with_template("strict", strict),
4948 ResponseCache::new(4, Duration::from_secs(60)),
4949 ));
4950 let app = test_app_with_state(state);
4951
4952 let (status, body) = post_json_uri(
4953 &app,
4954 "/v1/chat/completions",
4955 serde_json::json!({
4956 "model": "strict",
4957 "stream": true,
4958 "messages": [
4959 {"role": "user", "content": "one"},
4960 {"role": "user", "content": "two"},
4961 ],
4962 }),
4963 )
4964 .await;
4965 assert_eq!(status, StatusCode::BAD_REQUEST);
4966 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4967 assert!(
4968 body["error"]["message"]
4969 .as_str()
4970 .unwrap()
4971 .contains("one turn"),
4972 "the template's own message must reach the caller: {body}"
4973 );
4974
4975 let (status, _) = post_json_uri(
4977 &app,
4978 "/v1/chat/completions",
4979 serde_json::json!({
4980 "model": "strict",
4981 "stream": true,
4982 "max_tokens": 1,
4983 "messages": [{"role": "user", "content": "one"}],
4984 }),
4985 )
4986 .await;
4987 assert_eq!(status, StatusCode::OK);
4988 }
4989
4990 #[tokio::test]
4992 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4993 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4994 {% if reasoning_effort %}\
4995 {% if reasoning_effort not in ['low','medium','high'] %}\
4996 {{ raise_exception('bad effort') }}\
4997 {% endif %}[{{ reasoning_effort }}]\
4998 {% endif %}{{ messages[0].content }}";
4999 let state = Arc::new(test_state(
5000 model_with_template("thinker", reasoning),
5001 ResponseCache::new(4, Duration::from_secs(60)),
5002 ));
5003 let app = test_app_with_state(state);
5004 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5005 assert_eq!(status, StatusCode::OK);
5006 let entry = &models["data"][0];
5007 assert_eq!(
5008 entry["supported_reasoning_efforts"],
5009 serde_json::json!(["off", "low", "medium", "high"])
5010 );
5011 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5012 }
5013
5014 #[tokio::test]
5019 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5020 let app = test_app();
5021 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5022 let entry = &models["data"][0];
5023 assert!(entry.get("supported_reasoning_efforts").is_none());
5024 assert!(entry.get("default_reasoning_effort").is_none());
5025 }
5026
5027 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5028 Arc::new(ActiveModel {
5029 id: Some(name.to_string()),
5030 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5031 batcher: None,
5032 ceiling: None,
5033 checkpoint_path: None,
5034 })
5035 .tap_into(state)
5036 }
5037
5038 trait TapInto {
5040 fn tap_into(self, state: &AppState) -> Self;
5041 }
5042 impl TapInto for Arc<ActiveModel> {
5043 fn tap_into(self, state: &AppState) -> Self {
5044 state.swap_active(Some(Arc::clone(&self)));
5045 self
5046 }
5047 }
5048
5049 #[test]
5055 fn an_in_flight_request_keeps_the_model_it_started_on() {
5056 let state = test_state(
5057 named_test_model("model-a", 256),
5058 ResponseCache::new(4, Duration::from_secs(60)),
5059 );
5060
5061 let in_flight = state.active().expect("a model is loaded");
5064 assert_eq!(in_flight.name(), "model-a");
5065
5066 active_model(&state, "model-b");
5067
5068 assert_eq!(state.active().unwrap().name(), "model-b");
5070 assert_eq!(in_flight.name(), "model-a");
5072 let (choices, _usage) = run_generation(
5073 in_flight.generative().unwrap(),
5074 "hi",
5075 &greedy_params(3),
5076 None,
5077 None,
5078 None,
5079 None,
5080 None,
5081 None,
5082 )
5083 .expect("the old model must still decode after being swapped out");
5084 assert!(matches!(
5085 choices[0].finish,
5086 FinishReason::Length | FinishReason::Stop
5087 ));
5088 }
5089
5090 #[test]
5095 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5096 let state = test_state(
5097 named_test_model("model-a", 256),
5098 ResponseCache::new(4, Duration::from_secs(60)),
5099 );
5100 let in_flight = state.active().expect("a model is loaded");
5101 let weights = Arc::clone(in_flight.generative().unwrap());
5102 assert!(Arc::strong_count(&weights) >= 2);
5103
5104 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5105 id: Some("model-b".to_string()),
5106 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5107 batcher: None,
5108 ceiling: None,
5109 checkpoint_path: None,
5110 })));
5111 drop(previous);
5112 assert!(Arc::strong_count(&weights) >= 2);
5114 drop(in_flight);
5115 assert_eq!(Arc::strong_count(&weights), 1);
5116 }
5117
5118 #[tokio::test]
5122 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5123 let state = Arc::new(test_state(
5124 named_test_model("model-a", 256),
5125 ResponseCache::new(4, Duration::from_secs(60)),
5126 ));
5127 let app = test_app_with_state(Arc::clone(&state));
5128
5129 let (status, body) = post_json_uri(
5130 &app,
5131 frink_api::routes::ADMIN_MODELS_UNLOAD,
5132 serde_json::json!({}),
5133 )
5134 .await;
5135 assert_eq!(status, StatusCode::OK);
5136 assert_eq!(body["ok"], true);
5137 assert!(body["active"].is_null());
5138 assert!(state.active().is_none());
5139
5140 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5141 assert_eq!(status, StatusCode::OK);
5142 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5143 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5144
5145 let (status, body) = post_json_uri(
5146 &app,
5147 "/v1/chat/completions",
5148 serde_json::json!({
5149 "model": "x",
5150 "messages": [{"role": "user", "content": "hi"}]
5151 }),
5152 )
5153 .await;
5154 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5155 assert_eq!(body["error"]["type"], "model_not_loaded");
5156 }
5157
5158 #[tokio::test]
5162 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5163 let state = Arc::new(test_state(
5164 named_test_model("model-a", 256),
5165 ResponseCache::new(4, Duration::from_secs(60)),
5166 ));
5167 let app = test_app_with_state(Arc::clone(&state));
5168 state.swap_active(None);
5169
5170 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5171 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5174 assert_eq!(body["state"], "unavailable");
5175 assert_eq!(body["reason"], "model_not_loaded");
5176 assert!(body["model"].is_null());
5177 let real_weights = body["capabilities"]
5178 .as_array()
5179 .unwrap()
5180 .iter()
5181 .find(|c| c["id"] == "real_weights")
5182 .cloned()
5183 .expect("real_weights is always reported");
5184 assert_eq!(real_weights["available"], false);
5185 assert_eq!(real_weights["reason"], "model_not_loaded");
5186 }
5187
5188 #[tokio::test]
5192 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5193 let app = test_app();
5194
5195 let (status, completion) = post_json_uri(
5196 &app,
5197 "/v1/chat/completions",
5198 serde_json::json!({
5199 "model": "x",
5200 "messages": [{"role": "user", "content": "hi"}],
5201 "max_tokens": 4
5202 }),
5203 )
5204 .await;
5205 assert_eq!(status, StatusCode::OK);
5206 let request_id = completion["request_id"].as_str().unwrap().to_string();
5207
5208 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5209 assert_eq!(status, StatusCode::OK);
5210 let recent = stats["recent"].as_array().unwrap();
5211 assert_eq!(recent.len(), 1);
5212 let row = &recent[0];
5213 assert_eq!(row["request_id"], request_id);
5214 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5215 assert_eq!(row["status"], 200);
5216 assert_eq!(row["stream"], false);
5217 assert!(row["duration_ms"].is_number());
5220 assert!(row["decode_ms"].is_number());
5221 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5222 assert_eq!(
5223 stats["tokens_prompt_total"].as_u64().unwrap(),
5224 row["prompt_tokens"].as_u64().unwrap()
5225 );
5226 }
5227
5228 #[tokio::test]
5232 async fn a_rejected_request_is_recorded_too() {
5233 let state = Arc::new(test_state(
5234 named_test_model("model-a", 256),
5235 ResponseCache::new(4, Duration::from_secs(60)),
5236 ));
5237 let app = test_app_with_state(Arc::clone(&state));
5238 state.swap_active(None);
5239
5240 let (status, _) = post_json_uri(
5241 &app,
5242 "/v1/chat/completions",
5243 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5244 )
5245 .await;
5246 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5247
5248 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5249 let recent = stats["recent"].as_array().unwrap();
5250 assert_eq!(recent.len(), 1);
5251 assert_eq!(recent[0]["status"], 503);
5252 assert_eq!(recent[0]["completion_tokens"], 0);
5253 assert!(recent[0]["decode_ms"].is_null());
5254 assert_eq!(stats["errors_total"], 1);
5255 }
5256
5257 async fn post_json_with_headers(
5261 app: &Router,
5262 uri: &str,
5263 body: serde_json::Value,
5264 headers: &[(&str, &str)],
5265 ) -> (StatusCode, serde_json::Value) {
5266 use http_body_util::BodyExt;
5267 use tower::ServiceExt;
5268
5269 let mut builder = axum::http::Request::builder()
5270 .method("POST")
5271 .uri(uri)
5272 .header("content-type", "application/json");
5273 for (name, value) in headers {
5274 builder = builder.header(*name, *value);
5275 }
5276 let response = app
5277 .clone()
5278 .oneshot(
5279 builder
5280 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5281 .unwrap(),
5282 )
5283 .await
5284 .unwrap();
5285 let status = response.status();
5286 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5287 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5288 (status, json)
5289 }
5290
5291 #[tokio::test]
5295 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5296 let app = test_app();
5297
5298 let (status, _) = post_json_uri(
5299 &app,
5300 frink_api::routes::V1_TOKENIZE,
5301 serde_json::json!({"prompt": "hello"}),
5302 )
5303 .await;
5304 assert_eq!(status, StatusCode::OK);
5305 let (status, _) = post_json_uri(
5306 &app,
5307 frink_api::routes::V1_DETOKENIZE,
5308 serde_json::json!({"tokens": [104, 105]}),
5309 )
5310 .await;
5311 assert_eq!(status, StatusCode::OK);
5312 let (status, _) = post_json_uri(
5313 &app,
5314 frink_api::routes::V1_EMBEDDINGS,
5315 serde_json::json!({"input": "hello"}),
5316 )
5317 .await;
5318 assert_eq!(status, StatusCode::OK);
5319
5320 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5321 let routes: Vec<&str> = stats["recent"]
5322 .as_array()
5323 .unwrap()
5324 .iter()
5325 .map(|row| row["route"].as_str().unwrap())
5326 .collect();
5327 for expected in [
5328 frink_api::routes::V1_TOKENIZE,
5329 frink_api::routes::V1_DETOKENIZE,
5330 frink_api::routes::V1_EMBEDDINGS,
5331 ] {
5332 assert!(
5333 routes.contains(&expected),
5334 "{expected} is missing: {routes:?}"
5335 );
5336 }
5337
5338 let row = |route: &str| {
5339 stats["recent"]
5340 .as_array()
5341 .unwrap()
5342 .iter()
5343 .find(|r| r["route"] == route)
5344 .cloned()
5345 .unwrap()
5346 };
5347 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5351 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5352 assert!(embed["decode_ms"].is_null());
5353 assert_eq!(embed["completion_tokens"], 0);
5354 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5358 assert_eq!(
5359 stats["tokens_prompt_total"].as_u64().unwrap(),
5360 embed["prompt_tokens"].as_u64().unwrap(),
5361 "only the forward pass counted"
5362 );
5363 }
5364
5365 fn streaming_test_app() -> Router {
5370 let mut cfg = test_dense_fixture();
5371 cfg.vocab_size = 256;
5372 let model = Model::Gguf(GgufModel {
5373 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5374 tokenizer: Arc::new(ServerTokenizer::Byte),
5375 stop_tokens: StopTokens::default(),
5376 bos_id: None,
5377 is_synthetic: false,
5378 chat_template: chat_template::PromptTemplate::plain(),
5379 });
5380 test_app_with_state(Arc::new(test_state(
5381 model,
5382 ResponseCache::new(1000, Duration::from_secs(3600)),
5383 )))
5384 }
5385
5386 #[tokio::test]
5394 async fn chat_serves_several_choices_and_refuses_the_streaming_pair() {
5395 let app = test_app();
5396 let body = |n: u32, stream: bool| {
5397 serde_json::json!({
5398 "model": "x",
5399 "messages": [{"role": "user", "content": "hi"}],
5400 "max_tokens": 4,
5401 "temperature": 1.0,
5402 "n": n,
5403 "stream": stream
5404 })
5405 };
5406
5407 let (status, one) =
5408 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5409 assert_eq!(status, StatusCode::OK, "{one}");
5410
5411 let (status, three) =
5412 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5413 assert_eq!(status, StatusCode::OK, "{three}");
5414 let choices = three["choices"].as_array().expect("an array");
5415 assert_eq!(choices.len(), 3, "{three}");
5416 for (i, c) in choices.iter().enumerate() {
5417 assert_eq!(c["index"], i);
5418 assert!(c["message"]["role"].is_string(), "{c}");
5419 assert!(c["finish_reason"].is_string(), "{c}");
5420 }
5421 assert_eq!(
5423 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5424 "n = 3 billed the prompt more than once"
5425 );
5426
5427 let (status, refused) =
5430 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, true)).await;
5431 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{refused}");
5432 let message = refused["error"]["message"].as_str().unwrap_or_default();
5433 assert!(
5434 message.contains('n') && message.contains("stream"),
5435 "{refused}"
5436 );
5437 }
5438
5439 #[tokio::test]
5449 async fn every_route_refuses_the_same_unimplemented_fields() {
5450 let app = test_app();
5451 let fields = [
5452 ("n", serde_json::json!(3)),
5453 ("best_of", serde_json::json!(2)),
5454 ("prompt_logprobs", serde_json::json!(1)),
5455 ("echo", serde_json::json!(true)),
5456 ("use_beam_search", serde_json::json!(true)),
5457 ("truncate_prompt_tokens", serde_json::json!(8)),
5458 ("prompt_embeds", serde_json::json!("AA==")),
5459 ("allowed_token_ids", serde_json::json!([1, 2])),
5460 ("bad_words", serde_json::json!(["x"])),
5461 ("skip_special_tokens", serde_json::json!(false)),
5462 ("return_tokens_as_token_ids", serde_json::json!(true)),
5463 ];
5464 for (field, value) in fields {
5465 for (uri, base) in [
5466 (
5467 frink_api::routes::V1_CHAT_COMPLETIONS,
5468 serde_json::json!({
5469 "model": "x",
5470 "messages": [{"role": "user", "content": "hi"}],
5471 "max_tokens": 2
5472 }),
5473 ),
5474 (
5475 frink_api::routes::V1_COMPLETIONS,
5476 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
5477 ),
5478 (
5479 frink_api::routes::COMPLETION,
5480 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5481 ),
5482 ] {
5483 let mut body = base;
5484 body[field] = value.clone();
5485 if field == "n"
5490 && (uri == frink_api::routes::V1_COMPLETIONS
5491 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
5492 {
5493 let (status, answer) = post_json_uri(&app, uri, body).await;
5494 assert_eq!(
5495 status,
5496 StatusCode::OK,
5497 "{uri} refused a served `n`: {answer}"
5498 );
5499 assert_eq!(
5500 answer["choices"].as_array().map(Vec::len),
5501 Some(3),
5502 "{answer}"
5503 );
5504 continue;
5505 }
5506 let (status, answer) = post_json_uri(&app, uri, body).await;
5507 assert_eq!(
5508 status,
5509 StatusCode::NOT_IMPLEMENTED,
5510 "{uri} served `{field}` instead of refusing it: {answer}"
5511 );
5512 assert!(
5513 answer["error"]["message"]
5514 .as_str()
5515 .is_some_and(|m| m.contains(field)),
5516 "{uri} refused `{field}` without naming it: {answer}"
5517 );
5518 }
5519 }
5520 }
5521
5522 #[tokio::test]
5523 async fn the_native_completion_wire_is_not_the_openai_one() {
5524 let app = test_app();
5525
5526 let (status, native) = post_json_uri(
5527 &app,
5528 frink_api::routes::COMPLETION,
5529 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5530 )
5531 .await;
5532 assert_eq!(status, StatusCode::OK, "{native}");
5533 assert!(native["content"].is_string(), "{native}");
5534 assert_eq!(native["stop"], true);
5535 assert_eq!(native["stop_type"], "limit");
5536 assert_eq!(native["stopping_word"], "");
5537 assert_eq!(native["truncated"], false);
5538 assert_eq!(native["id_slot"], -1);
5539 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5540 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5541 assert!(
5542 native.get("choices").is_none(),
5543 "the native shape has no `choices`: {native}"
5544 );
5545
5546 let (status, openai) = post_json_uri(
5547 &app,
5548 frink_api::routes::V1_COMPLETIONS,
5549 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5550 )
5551 .await;
5552 assert_eq!(status, StatusCode::OK);
5553 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5554 assert!(
5555 openai.get("content").is_none(),
5556 "the OpenAI shape has no top-level `content`: {openai}"
5557 );
5558 }
5559
5560 #[tokio::test]
5564 async fn both_native_spellings_reach_the_same_handler() {
5565 let app = test_app();
5566 for route in [
5567 frink_api::routes::COMPLETION,
5568 frink_api::routes::COMPLETIONS,
5569 ] {
5570 let (status, body) = post_json_uri(
5571 &app,
5572 route,
5573 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5574 )
5575 .await;
5576 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5577 assert_eq!(body["stop"], true, "{route}");
5578 assert!(body["content"].is_string(), "{route}");
5579 }
5580
5581 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5584 let routes: Vec<&str> = stats["recent"]
5585 .as_array()
5586 .unwrap()
5587 .iter()
5588 .map(|row| row["route"].as_str().unwrap())
5589 .collect();
5590 assert!(
5591 routes.contains(&frink_api::routes::COMPLETION),
5592 "{routes:?}"
5593 );
5594 assert!(
5595 routes.contains(&frink_api::routes::COMPLETIONS),
5596 "{routes:?}"
5597 );
5598 }
5599
5600 #[tokio::test]
5606 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5607 let app = streaming_test_app();
5608 let raw = post_sse_raw_uri(
5609 &app,
5610 frink_api::routes::COMPLETION,
5611 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5612 )
5613 .await;
5614
5615 assert!(
5616 !raw.contains("[DONE]"),
5617 "llama.cpp's native stream has no sentinel: {raw}"
5618 );
5619 let frames: Vec<serde_json::Value> = raw
5620 .lines()
5621 .filter_map(|line| line.strip_prefix("data: "))
5622 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5623 .collect();
5624 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5625
5626 let (last, partials) = frames.split_last().unwrap();
5627 assert_eq!(last["stop"], true, "the last frame closes the stream");
5628 assert!(last["timings"].is_object(), "{last}");
5629 assert!(last["stop_type"].is_string(), "{last}");
5630 for partial in partials {
5631 assert_eq!(partial["stop"], false, "{partial}");
5632 assert!(partial["content"].is_string(), "{partial}");
5633 assert!(partial.get("timings").is_none(), "{partial}");
5637 assert!(partial.get("generation_settings").is_none(), "{partial}");
5638 }
5639 let streamed: String = partials
5642 .iter()
5643 .filter_map(|p| p["content"].as_str())
5644 .collect();
5645 assert_eq!(last["content"].as_str().unwrap(), streamed);
5646 }
5647
5648 #[tokio::test]
5653 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5654 let app = test_app();
5655 for body in [
5656 serde_json::json!({"prompt": "hi"}),
5657 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5658 ] {
5659 let (status, refusal) =
5660 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5661 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5662 assert!(
5663 refusal["error"]["message"]
5664 .as_str()
5665 .unwrap()
5666 .contains("n_predict"),
5667 "{refusal}"
5668 );
5669 }
5670 let (status, _) = post_json_uri(
5673 &app,
5674 frink_api::routes::COMPLETION,
5675 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5676 )
5677 .await;
5678 assert_eq!(status, StatusCode::OK);
5679 }
5680
5681 #[tokio::test]
5690 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5691 let app = streaming_test_app();
5692 let ask = |stop: serde_json::Value| {
5693 let app = app.clone();
5694 async move {
5695 post_json_uri(
5696 &app,
5697 frink_api::routes::COMPLETION,
5698 serde_json::json!({
5699 "prompt": "hi",
5700 "n_predict": 64,
5701 "ignore_eos": true,
5702 "stop": stop,
5703 }),
5704 )
5705 .await
5706 .1
5707 }
5708 };
5709
5710 let baseline = ask(serde_json::json!([])).await;
5711 assert_eq!(baseline["stop_type"], "limit");
5712 assert_eq!(baseline["stopping_word"], "");
5713 let text = baseline["content"].as_str().unwrap().to_string();
5714 let sentinel: String = text.chars().skip(1).take(2).collect();
5720 assert_eq!(
5721 sentinel.chars().count(),
5722 2,
5723 "the fixture must produce enough output to cut: {text:?}"
5724 );
5725 let cut = text.find(&sentinel).expect("it came out of this text");
5726
5727 let stopped = ask(serde_json::json!([sentinel])).await;
5728 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5729 assert_eq!(stopped["stopping_word"], sentinel);
5730 assert_eq!(
5731 stopped["content"].as_str().unwrap(),
5732 &text[..cut],
5733 "the answer must be cut at the sentinel, not run past it"
5734 );
5735 }
5736
5737 #[tokio::test]
5743 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5744 let app = test_app();
5745
5746 let (v1_status, v1) = post_json_uri(
5747 &app,
5748 frink_api::routes::V1_TOKENIZE,
5749 serde_json::json!({"prompt": "hello"}),
5750 )
5751 .await;
5752 let (alias_status, alias) = post_json_uri(
5753 &app,
5754 frink_api::routes::TOKENIZE,
5755 serde_json::json!({"content": "hello"}),
5756 )
5757 .await;
5758 assert_eq!(v1_status, StatusCode::OK);
5759 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5760 assert_eq!(v1["tokens"], alias["tokens"]);
5761 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5762
5763 let (status, both_ways) = post_json_uri(
5766 &app,
5767 frink_api::routes::TOKENIZE,
5768 serde_json::json!({"prompt": "hello"}),
5769 )
5770 .await;
5771 assert_eq!(status, StatusCode::OK);
5772 assert_eq!(both_ways["tokens"], v1["tokens"]);
5773 }
5774
5775 #[tokio::test]
5780 async fn detokenize_answers_under_both_dialects_keys() {
5781 let app = test_app();
5782 for route in [
5783 frink_api::routes::DETOKENIZE,
5784 frink_api::routes::V1_DETOKENIZE,
5785 ] {
5786 let (status, body) =
5787 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5788 assert_eq!(status, StatusCode::OK, "{route}");
5789 assert_eq!(body["text"], "hi", "{route}");
5790 assert_eq!(body["content"], body["text"], "{route}");
5791 }
5792 }
5793
5794 #[tokio::test]
5798 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5799 let app = test_app();
5800 let (status, _) = post_json_uri(
5801 &app,
5802 frink_api::routes::TOKENIZE,
5803 serde_json::json!({"content": "hello"}),
5804 )
5805 .await;
5806 assert_eq!(status, StatusCode::OK);
5807
5808 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5809 let routes: Vec<&str> = stats["recent"]
5810 .as_array()
5811 .unwrap()
5812 .iter()
5813 .map(|row| row["route"].as_str().unwrap())
5814 .collect();
5815 assert!(
5816 routes.contains(&frink_api::routes::TOKENIZE),
5817 "the alias must be its own row: {routes:?}"
5818 );
5819 assert!(
5820 !routes.contains(&frink_api::routes::V1_TOKENIZE),
5821 "nothing called /v1/tokenize: {routes:?}"
5822 );
5823 }
5824
5825 #[tokio::test]
5830 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5831 let mut cfg = test_dense_fixture();
5832 cfg.vocab_size = 256;
5833 let model = Model::Gguf(GgufModel {
5834 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5835 tokenizer: Arc::new(ServerTokenizer::Byte),
5836 stop_tokens: StopTokens::default(),
5837 bos_id: Some(7),
5838 is_synthetic: true,
5839 chat_template: chat_template::PromptTemplate::plain(),
5840 });
5841 let app = test_app_with_state(Arc::new(test_state(
5842 model,
5843 ResponseCache::new(1000, Duration::from_secs(3600)),
5844 )));
5845
5846 let (_, plain) = post_json_uri(
5847 &app,
5848 frink_api::routes::TOKENIZE,
5849 serde_json::json!({"content": "hi"}),
5850 )
5851 .await;
5852 let (_, special) = post_json_uri(
5853 &app,
5854 frink_api::routes::TOKENIZE,
5855 serde_json::json!({"content": "hi", "add_special": true}),
5856 )
5857 .await;
5858
5859 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5860 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5861 assert_eq!(special["count"], 3);
5862 }
5863
5864 #[tokio::test]
5867 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5868 let app = test_app();
5869 let (status, _) = post_json_uri(
5870 &app,
5871 frink_api::routes::V1_EMBEDDINGS,
5872 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5873 )
5874 .await;
5875 assert_eq!(status, StatusCode::BAD_REQUEST);
5876
5877 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5878 let recent = stats["recent"].as_array().unwrap();
5879 assert_eq!(recent.len(), 1);
5880 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
5881 assert_eq!(recent[0]["status"], 400);
5882 assert_eq!(
5883 recent[0]["prompt_tokens"], 0,
5884 "a rejected call embedded nothing"
5885 );
5886 }
5887
5888 #[tokio::test]
5891 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5892 let app = test_app();
5893 let key = "sk-monitor-secret";
5894 let (status, _) = post_json_with_headers(
5895 &app,
5896 "/v1/chat/completions",
5897 serde_json::json!({
5898 "model": "x",
5899 "messages": [{"role": "user", "content": "hi"}],
5900 "max_tokens": 2
5901 }),
5902 &[
5903 ("authorization", &format!("Bearer {key}")),
5904 ("x-frink-client", "frink-studio"),
5905 ],
5906 )
5907 .await;
5908 assert_eq!(status, StatusCode::OK);
5909
5910 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5911 let row = stats["recent"].as_array().unwrap()[0].clone();
5912 let fingerprint = row["via_api_key"]
5913 .as_str()
5914 .expect("the row names the key that served it")
5915 .to_string();
5916 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5917 assert!(!fingerprint.contains(key));
5918 assert!(
5919 !serde_json::to_string(&stats).unwrap().contains(key),
5920 "the stats payload must not carry the key in any form"
5921 );
5922 assert_eq!(row["client"], "frink-studio");
5923 }
5924
5925 #[tokio::test]
5928 async fn different_keys_are_different_callers_and_no_key_is_null() {
5929 let app = test_app();
5930 let body = serde_json::json!({
5931 "model": "x",
5932 "messages": [{"role": "user", "content": "hi"}],
5933 "max_tokens": 1
5934 });
5935 for headers in [
5936 vec![("authorization", "Bearer key-one")],
5937 vec![("authorization", "Bearer key-two")],
5938 vec![],
5939 ] {
5940 let (status, _) =
5941 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5942 assert_eq!(status, StatusCode::OK);
5943 }
5944
5945 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5946 let recent = stats["recent"].as_array().unwrap();
5947 assert_eq!(recent.len(), 3);
5948 let one = recent[0]["via_api_key"].as_str().unwrap();
5949 let two = recent[1]["via_api_key"].as_str().unwrap();
5950 assert_ne!(one, two, "two keys must not collapse into one caller");
5951 assert!(
5952 recent[2]["via_api_key"].is_null(),
5953 "an unauthenticated call is null, not a fingerprint of nothing"
5954 );
5955 assert!(recent[2]["client"].is_null());
5956 }
5957
5958 #[tokio::test]
5963 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5964 let state = Arc::new(test_state(
5965 named_test_model("really-loaded", 256),
5966 ResponseCache::new(4, Duration::from_secs(60)),
5967 ));
5968 let app = test_app_with_state(Arc::clone(&state));
5969
5970 let (status, _) = post_json_uri(
5971 &app,
5972 "/v1/chat/completions",
5973 serde_json::json!({
5974 "model": "gpt-4-turbo-that-is-not-here",
5975 "messages": [{"role": "user", "content": "hi"}],
5976 "max_tokens": 2
5977 }),
5978 )
5979 .await;
5980 assert_eq!(status, StatusCode::OK);
5981
5982 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5983 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5984
5985 state.swap_active(None);
5988 let (status, _) = post_json_uri(
5989 &app,
5990 "/v1/chat/completions",
5991 serde_json::json!({
5992 "model": "gpt-4-turbo-that-is-not-here",
5993 "messages": [{"role": "user", "content": "hi"}]
5994 }),
5995 )
5996 .await;
5997 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5998 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5999 let recent = stats["recent"].as_array().unwrap();
6000 assert!(recent[recent.len() - 1]["model"].is_null());
6001 }
6002
6003 #[tokio::test]
6007 async fn a_streamed_row_names_the_model_it_decoded_against() {
6008 let state = Arc::new(test_state(
6009 named_test_model("model-before", 256),
6010 ResponseCache::new(4, Duration::from_secs(60)),
6011 ));
6012 let app = test_app_with_state(Arc::clone(&state));
6013 let _ = post_sse_raw(&app, resumable_request()).await;
6014 active_model(&state, "model-after");
6016
6017 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6018 assert_eq!(stats["recent"][0]["model"], "model-before");
6019 }
6020
6021 #[tokio::test]
6024 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6025 let app = test_app();
6026 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6027 assert_eq!(status, StatusCode::OK);
6028 assert!(
6029 stats["queue_depth"].is_null(),
6030 "without continuous batching nothing queues, so there is nothing to measure"
6031 );
6032 assert!(stats["queue_rejected_total"].is_null());
6033 assert_eq!(
6034 stats["generating_now"], 0,
6035 "work in progress is measured and really is zero here"
6036 );
6037 }
6038
6039 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6044 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
6045 }
6046
6047 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6051 use http_body_util::BodyExt;
6052 use tower::ServiceExt;
6053
6054 let response = app
6055 .clone()
6056 .oneshot(
6057 axum::http::Request::builder()
6058 .method("POST")
6059 .uri(uri)
6060 .header("content-type", "application/json")
6061 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6062 .unwrap(),
6063 )
6064 .await
6065 .unwrap();
6066 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6067 String::from_utf8(bytes.to_vec()).unwrap()
6068 }
6069
6070 async fn get_json_with_headers(
6071 app: &Router,
6072 uri: &str,
6073 headers: &[(&str, &str)],
6074 ) -> (StatusCode, serde_json::Value) {
6075 use http_body_util::BodyExt;
6076 use tower::ServiceExt;
6077
6078 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6079 for (name, value) in headers {
6080 builder = builder.header(*name, *value);
6081 }
6082 let response = app
6083 .clone()
6084 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6085 .await
6086 .unwrap();
6087 let status = response.status();
6088 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6089 (
6090 status,
6091 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6092 )
6093 }
6094
6095 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6096 body.lines()
6097 .filter_map(|line| line.strip_prefix(field))
6098 .map(str::trim)
6099 .collect()
6100 }
6101
6102 fn resumable_request() -> serde_json::Value {
6103 serde_json::json!({
6104 "model": "m",
6105 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6106 "max_tokens": 4,
6107 "temperature": 0,
6108 "stream": true,
6109 "stream_resumable": true,
6110 })
6111 }
6112
6113 #[tokio::test]
6118 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6119 let app = test_app();
6120 let body = post_sse_raw(&app, resumable_request()).await;
6121
6122 let request_id = body
6123 .lines()
6124 .find_map(|l| l.strip_prefix("data: "))
6125 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6126 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6127 .expect("the first chunk names the request");
6128
6129 let ids = sse_field(&body, "id:");
6130 let datas = sse_field(&body, "data:");
6131 assert_eq!(
6132 ids.len(),
6133 datas.len(),
6134 "every event carries an id, or a reconnect cannot name where it stopped"
6135 );
6136 for (i, id) in ids.iter().enumerate() {
6137 assert_eq!(*id, format!("{request_id}:{i}"));
6138 }
6139 let retries = sse_field(&body, "retry:");
6140 assert_eq!(
6141 retries.len(),
6142 1,
6143 "the reconnect delay is stated once, not on every event"
6144 );
6145 assert_eq!(retries[0], "1500");
6146 assert!(
6147 body.contains("data: [DONE]"),
6148 "the end of stream is still stated"
6149 );
6150 }
6151
6152 #[tokio::test]
6156 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6157 let app = test_app();
6158 let mut request = resumable_request();
6159 request["stream_resumable"] = serde_json::json!(false);
6160 let body = post_sse_raw(&app, request).await;
6161 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6162 assert!(
6163 sse_field(&body, "id:").is_empty(),
6164 "an id promises a replay this stream cannot serve"
6165 );
6166 assert!(sse_field(&body, "retry:").is_empty());
6167 }
6168
6169 #[tokio::test]
6173 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6174 let app = test_app();
6175 let body = post_sse_raw(&app, resumable_request()).await;
6176 let request_id = sse_field(&body, "id:")[0]
6177 .rsplit_once(':')
6178 .unwrap()
6179 .0
6180 .to_string();
6181 let streamed: Vec<String> = sse_field(&body, "data:")
6182 .iter()
6183 .map(|d| d.to_string())
6184 .collect();
6185
6186 let (status, polled) = get_json(
6187 &app,
6188 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6189 )
6190 .await;
6191 assert_eq!(status, StatusCode::OK);
6192 let events: Vec<String> = polled["events"]
6193 .as_array()
6194 .unwrap()
6195 .iter()
6196 .map(|e| e["data"].as_str().unwrap().to_string())
6197 .collect();
6198 assert_eq!(
6199 events, streamed,
6200 "the fallback must deliver the same answer, not a re-run of it"
6201 );
6202 assert_eq!(polled["request_id"], request_id);
6203 assert_eq!(
6204 polled["done"], false,
6205 "events were still being handed out, so the client must ask again"
6206 );
6207
6208 let next = polled["next_index"].as_u64().unwrap();
6211 let (_, drained) = get_json(
6212 &app,
6213 &format!(
6214 "{}?from={next}",
6215 frink_api::routes::v1_stream_poll(&request_id)
6216 ),
6217 )
6218 .await;
6219 assert_eq!(drained["done"], true);
6220 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6221 }
6222
6223 #[tokio::test]
6227 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6228 let app = test_app();
6229 let body = post_sse_raw(&app, resumable_request()).await;
6230 let ids = sse_field(&body, "id:");
6231 let datas: Vec<String> = sse_field(&body, "data:")
6232 .iter()
6233 .map(|d| d.to_string())
6234 .collect();
6235 assert!(
6236 ids.len() >= 3,
6237 "need a few events to resume into the middle"
6238 );
6239 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6240
6241 let (status, resumed) = get_json_with_headers(
6242 &app,
6243 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6244 &[],
6245 )
6246 .await;
6247 assert_eq!(status, StatusCode::OK);
6248 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6249
6250 let (_, tail) = get_json(
6252 &app,
6253 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6254 )
6255 .await;
6256 let tail_events: Vec<String> = tail["events"]
6257 .as_array()
6258 .unwrap()
6259 .iter()
6260 .map(|e| e["data"].as_str().unwrap().to_string())
6261 .collect();
6262 assert_eq!(tail_events, datas[2..].to_vec());
6263 }
6264
6265 #[tokio::test]
6268 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6269 use http_body_util::BodyExt;
6270 use tower::ServiceExt;
6271
6272 let app = test_app();
6273 let body = post_sse_raw(&app, resumable_request()).await;
6274 let ids = sse_field(&body, "id:");
6275 let datas: Vec<String> = sse_field(&body, "data:")
6276 .iter()
6277 .map(|d| d.to_string())
6278 .collect();
6279 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6280
6281 let response = app
6282 .clone()
6283 .oneshot(
6284 axum::http::Request::builder()
6285 .method("GET")
6286 .uri(frink_api::routes::v1_stream(&request_id))
6287 .header("last-event-id", format!("{request_id}:0"))
6288 .body(axum::body::Body::empty())
6289 .unwrap(),
6290 )
6291 .await
6292 .unwrap();
6293 assert_eq!(response.status(), StatusCode::OK);
6294 assert_eq!(
6295 response
6296 .headers()
6297 .get("x-accel-buffering")
6298 .and_then(|v| v.to_str().ok()),
6299 Some("no"),
6300 "the reconnect needs the same anti-buffering header as the stream"
6301 );
6302 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6303 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6304 assert_eq!(
6305 sse_field(&resumed, "data:")
6306 .iter()
6307 .map(|d| d.to_string())
6308 .collect::<Vec<_>>(),
6309 datas[1..].to_vec()
6310 );
6311 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6312 }
6313
6314 #[tokio::test]
6318 async fn a_last_event_id_from_another_stream_is_refused() {
6319 let app = test_app();
6320 let body = post_sse_raw(&app, resumable_request()).await;
6321 let request_id = sse_field(&body, "id:")[0]
6322 .rsplit_once(':')
6323 .unwrap()
6324 .0
6325 .to_string();
6326
6327 let (status, err) = get_json_with_headers(
6328 &app,
6329 &frink_api::routes::v1_stream(&request_id),
6330 &[("last-event-id", "chatcmpl-someone-else:3")],
6331 )
6332 .await;
6333 assert_eq!(status, StatusCode::BAD_REQUEST);
6334 assert_eq!(err["error"]["code"], "bad_last_event_id");
6335 }
6336
6337 #[tokio::test]
6341 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6342 let app = test_app();
6343 let mut request = resumable_request();
6344 request["stream_resumable"] = serde_json::json!(false);
6345 let body = post_sse_raw(&app, request).await;
6346 let request_id = body
6347 .lines()
6348 .find_map(|l| l.strip_prefix("data: "))
6349 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6350 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6351 .unwrap();
6352
6353 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6354 assert_eq!(status, StatusCode::NOT_FOUND);
6355 assert_eq!(err["error"]["code"], "stream_not_found");
6356 assert!(err["error"]["message"]
6357 .as_str()
6358 .unwrap()
6359 .contains("stream_resumable"));
6360 }
6361
6362 #[test]
6366 fn the_axum_stream_patterns_match_the_published_templates() {
6367 assert_eq!(
6368 axum_path(frink_api::routes::V1_STREAM),
6369 "/v1/stream/:request_id"
6370 );
6371 assert_eq!(
6372 axum_path(frink_api::routes::V1_STREAM_POLL),
6373 "/v1/stream/:request_id/poll"
6374 );
6375 assert_eq!(
6376 frink_api::routes::v1_stream("abc"),
6377 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6378 );
6379 }
6380
6381 #[test]
6391 fn no_published_template_reaches_the_router_with_its_braces() {
6392 for template in [
6393 frink_api::routes::V1_STREAM,
6394 frink_api::routes::V1_STREAM_POLL,
6395 frink_api::routes::V1_RESPONSE,
6396 frink_api::routes::V1_RESPONSE_CANCEL,
6397 frink_api::routes::ADMIN_TASK_CANCEL,
6398 ] {
6399 assert!(
6400 template.contains('{'),
6401 "{template} is in the template list but has no placeholder"
6402 );
6403 let mounted = axum_path(template);
6404 assert!(
6405 !mounted.contains('{') && !mounted.contains('}'),
6406 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6407 );
6408 assert!(
6409 mounted.contains(':'),
6410 "{template} lost its placeholder entirely and would match one path only"
6411 );
6412 }
6413 }
6414
6415 #[tokio::test]
6423 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6424 let app = test_app();
6425 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6426 assert_eq!(status, StatusCode::NOT_FOUND);
6427 assert!(
6428 !body.is_null(),
6429 "empty body means axum never matched the route, so the id was read as a literal segment"
6430 );
6431 }
6432
6433 #[tokio::test]
6436 async fn the_task_list_starts_empty_rather_than_absent() {
6437 let app = test_app();
6438 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6439 assert_eq!(status, StatusCode::OK);
6440 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6441 }
6442
6443 #[tokio::test]
6452 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6453 assert!(
6454 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6455 "this test asserts the unconfigured behaviour"
6456 );
6457 let app = test_app();
6458 let (status, body) = post_json_uri(
6459 &app,
6460 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6461 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6462 )
6463 .await;
6464 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6465 assert!(
6466 body["error"]["message"]
6467 .as_str()
6468 .unwrap()
6469 .contains("--slot-save-path"),
6470 "{body}"
6471 );
6472 }
6473
6474 pub(crate) async fn post_json_uri(
6475 app: &Router,
6476 uri: &str,
6477 body: serde_json::Value,
6478 ) -> (StatusCode, serde_json::Value) {
6479 use http_body_util::BodyExt;
6480 use tower::ServiceExt;
6481
6482 let response = app
6483 .clone()
6484 .oneshot(
6485 axum::http::Request::builder()
6486 .method("POST")
6487 .uri(uri)
6488 .header("content-type", "application/json")
6489 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6490 .unwrap(),
6491 )
6492 .await
6493 .unwrap();
6494 let status = response.status();
6495 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6496 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6497 (status, json)
6498 }
6499
6500 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6501 post_json_uri(app, "/v1/chat/completions", body).await.1
6502 }
6503
6504 #[tokio::test]
6513 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6514 let app = test_app();
6515 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6516 assert_eq!(status, StatusCode::OK);
6517
6518 let memory = &body["memory"];
6519 if memory.is_null() {
6520 return;
6523 }
6524 assert!(
6525 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6526 "a read that produced a zero is a broken read, not an idle \
6527 engine: {memory}"
6528 );
6529 assert!(
6530 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6531 "the quantity must travel with the number: {memory}"
6532 );
6533 }
6534
6535 #[tokio::test]
6541 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6542 let app = test_app();
6543 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6544 assert_eq!(status, StatusCode::OK);
6545 for pool in ["kv_pages", "window_slots", "state_slots"] {
6546 assert!(
6547 body["pools"][pool].is_null(),
6548 "{pool} must be null rather than a zero row: {}",
6549 body["pools"]
6550 );
6551 }
6552 }
6553
6554 #[tokio::test]
6562 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6563 use http_body_util::BodyExt;
6564 use tower::ServiceExt;
6565
6566 let app = test_app();
6567 let send = |body: serde_json::Value| {
6568 let app = app.clone();
6569 async move {
6570 app.oneshot(
6571 axum::http::Request::builder()
6572 .method("POST")
6573 .uri(frink_api::routes::V1_MESSAGES)
6574 .header("content-type", "application/json")
6575 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6576 .unwrap(),
6577 )
6578 .await
6579 .unwrap()
6580 }
6581 };
6582
6583 let ok = send(serde_json::json!({
6584 "model": "test",
6585 "max_tokens": 1,
6586 "messages": [{"role": "user", "content": "hi"}],
6587 }))
6588 .await;
6589 assert_eq!(ok.status(), StatusCode::OK);
6590 let id = ok
6591 .headers()
6592 .get("request-id")
6593 .expect("a served message names its id")
6594 .to_str()
6595 .unwrap()
6596 .to_string();
6597 assert!(!id.is_empty());
6598
6599 let bad = send(serde_json::json!({"model": "test"})).await;
6602 assert!(bad.status().is_client_error());
6603 let other = bad.headers().get("request-id").expect("errors too");
6604 assert_ne!(other.to_str().unwrap(), id);
6605 let _ = bad.into_body().collect().await.unwrap();
6606 }
6607
6608 #[tokio::test]
6616 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6617 let state = Arc::new(test_state(
6618 test_model_full_byte_vocab(),
6619 ResponseCache::new(1000, Duration::from_secs(3600)),
6620 ));
6621 let app = test_app_with_state(Arc::clone(&state));
6622 let body = serde_json::json!({
6623 "model": "test",
6624 "messages": [{"role": "user", "content": "hi"}],
6625 "max_tokens": 1,
6626 });
6627
6628 state
6629 .maintenance
6630 .lock()
6631 .unwrap()
6632 .begin_rebuild()
6633 .expect("a fresh server is serving, so the rebuild starts");
6634 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6635 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6636 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6637
6638 state.maintenance.lock().unwrap().finish_rebuild(true);
6639 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6640 assert_eq!(
6641 status,
6642 StatusCode::OK,
6643 "the gate reopens; a rebuild is not a latch"
6644 );
6645 }
6646
6647 #[tokio::test]
6652 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6653 let app = test_app();
6654 let (status, body) = post_json_uri(
6655 &app,
6656 frink_api::routes::V1_CANCEL,
6657 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6658 )
6659 .await;
6660 assert_eq!(status, StatusCode::NOT_FOUND);
6661 assert_eq!(body["cancelled"], serde_json::json!(false));
6662 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6663 assert!(
6664 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6665 "the verdict must carry a human reason: {body}"
6666 );
6667 }
6668
6669 #[tokio::test]
6673 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6674 let state = Arc::new(test_state(
6675 test_model_full_byte_vocab(),
6676 ResponseCache::new(1000, Duration::from_secs(3600)),
6677 ));
6678 let app = test_app_with_state(Arc::clone(&state));
6679 let (token, _guard) = state.cancels.register("chatcmpl-live");
6680
6681 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6682 assert_eq!(status, StatusCode::OK);
6683 assert_eq!(before["generating_now"], serde_json::json!(1));
6684
6685 let (status, body) = post_json_uri(
6686 &app,
6687 frink_api::routes::V1_CANCEL,
6688 serde_json::json!({ "request_id": "chatcmpl-live" }),
6689 )
6690 .await;
6691 assert_eq!(status, StatusCode::OK);
6692 assert_eq!(body["cancelled"], serde_json::json!(true));
6693 assert!(
6694 token.is_cancelled(),
6695 "the endpoint answered ok without setting the flag the decode loop reads"
6696 );
6697 }
6698
6699 #[tokio::test]
6700 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6701 let app = test_app();
6702 let (status, tok) =
6703 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6704 assert_eq!(status, StatusCode::OK);
6705 let tokens = tok["tokens"].as_array().unwrap();
6706 assert_eq!(tok["count"], tokens.len());
6707 assert!(!tokens.is_empty());
6708
6709 let (status, detok) = post_json_uri(
6710 &app,
6711 "/v1/detokenize",
6712 serde_json::json!({ "tokens": tokens }),
6713 )
6714 .await;
6715 assert_eq!(status, StatusCode::OK);
6716 assert_eq!(detok["text"], "Hi");
6717
6718 let (status, emb) = post_json_uri(
6719 &app,
6720 "/v1/embeddings",
6721 serde_json::json!({
6722 "input": "Hi",
6723 "embedding_type": "mean"
6724 }),
6725 )
6726 .await;
6727 assert_eq!(status, StatusCode::OK);
6728 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6729 assert!(!vec.is_empty());
6730 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6731 }
6732
6733 #[tokio::test]
6739 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6740 let app = test_app();
6741 let (status, body) = post_json_uri(
6742 &app,
6743 "/v1/embeddings",
6744 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6745 )
6746 .await;
6747 assert_eq!(status, StatusCode::BAD_REQUEST);
6748 let msg = body["error"]["message"].as_str().unwrap();
6749 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6750 }
6751
6752 #[tokio::test]
6757 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6758 async fn a_real_embedding_model_serves_v1_embeddings() {
6759 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6760 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6761 if !path.exists() {
6762 eprintln!("SKIP: {} not present", path.display());
6763 return;
6764 }
6765 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6766 let mut state = test_state(
6767 test_model_full_byte_vocab(),
6768 ResponseCache::new(1000, Duration::from_secs(3600)),
6769 );
6770 state.embedding = Some(Arc::new(encoder));
6771 let app = test_app_with_state(Arc::new(state));
6772
6773 let (status, body) = post_json_uri(
6774 &app,
6775 "/v1/embeddings",
6776 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6777 )
6778 .await;
6779 assert_eq!(status, StatusCode::OK, "{body}");
6780 assert_eq!(body["model"], "bge-small-en-v1.5");
6781 let data = body["data"].as_array().unwrap();
6782 assert_eq!(data.len(), 2);
6783 for (i, row) in data.iter().enumerate() {
6784 assert_eq!(row["index"], i);
6785 let v: Vec<f64> = row["embedding"]
6786 .as_array()
6787 .unwrap()
6788 .iter()
6789 .map(|x| x.as_f64().unwrap())
6790 .collect();
6791 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6792 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6793 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6794 }
6795 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6798
6799 let (status, mean) = post_json_uri(
6803 &app,
6804 "/v1/embeddings",
6805 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6806 )
6807 .await;
6808 assert_eq!(status, StatusCode::OK);
6809 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6810 }
6811
6812 #[tokio::test]
6824 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6825 async fn an_encoder_can_be_the_loaded_model() {
6826 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6827 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6828 if !path.exists() {
6829 eprintln!("SKIP: {} not present", path.display());
6830 return;
6831 }
6832
6833 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6837 assert!(
6838 matches!(loaded, model::LoadedModel::Encoder(_)),
6839 "an encoder-only GGUF reached a decoder loader"
6840 );
6841 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6842 assert!(
6843 matches!(loaded, Loaded::Encoder(_)),
6844 "the encoder did not stay an encoder through activation"
6845 );
6846 assert!(
6847 batcher.is_none() && ceiling.is_none(),
6848 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6849 );
6850
6851 let state = test_state(
6852 test_model_full_byte_vocab(),
6853 ResponseCache::new(1000, Duration::from_secs(3600)),
6854 );
6855 state.swap_active(Some(Arc::new(ActiveModel {
6856 id: None,
6857 loaded,
6858 batcher,
6859 ceiling,
6860 checkpoint_path: None,
6861 })));
6862 let app = test_app_with_state(Arc::new(state));
6863
6864 let (status, body) = post_json_uri(
6866 &app,
6867 "/v1/embeddings",
6868 serde_json::json!({ "input": "Hello world" }),
6869 )
6870 .await;
6871 assert_eq!(status, StatusCode::OK, "{body}");
6872 assert_eq!(body["model"], "bge-small-en-v1.5");
6873 let v = body["data"][0]["embedding"].as_array().unwrap();
6874 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6875
6876 let (status, body) = post_json_uri(
6878 &app,
6879 "/v1/chat/completions",
6880 serde_json::json!({
6881 "model": "bge-small-en-v1.5",
6882 "messages": [{"role": "user", "content": "hi"}],
6883 }),
6884 )
6885 .await;
6886 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6887 let msg = body["error"]["message"].as_str().unwrap();
6888 for fact in [
6889 "bge-small-en-v1.5",
6890 "bert",
6891 "embedding model",
6892 "/v1/embeddings",
6893 ] {
6894 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6895 }
6896
6897 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
6899 assert_eq!(status, StatusCode::OK);
6900 let entry = &models["data"][0];
6901 assert_eq!(entry["id"], "bge-small-en-v1.5");
6902 assert_eq!(entry["frink_model_kind"], "embedding");
6903 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
6904 assert_eq!(entry["frink_n_embd"], 384);
6905 assert_eq!(entry["frink_pooling"], "CLS");
6906 assert_eq!(
6907 entry["frink_endpoints"],
6908 serde_json::json!(["/v1/embeddings"])
6909 );
6910 assert!(entry.get("supported_reasoning_efforts").is_none());
6913
6914 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
6916 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6917 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6918 assert_eq!(health["model"]["synthetic_weights"], false);
6919 let weights = health["capabilities"]
6920 .as_array()
6921 .unwrap()
6922 .iter()
6923 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
6924 .expect("a real-weights capability row");
6925 let detail = weights["detail"].as_str().unwrap_or_default();
6926 assert!(detail.contains("ENCODER"), "{detail}");
6927 let (status, body) = post_json_uri(
6935 &app,
6936 frink_api::routes::V1_TOKENIZE,
6937 serde_json::json!({ "content": "hello world" }),
6938 )
6939 .await;
6940 assert_eq!(
6941 status,
6942 StatusCode::OK,
6943 "an encoder has a real tokenizer: {body}"
6944 );
6945 let tokens = body["tokens"].as_array().expect("tokens array").clone();
6946 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
6947
6948 let (status, body) = post_json_uri(
6949 &app,
6950 frink_api::routes::V1_DETOKENIZE,
6951 serde_json::json!({ "tokens": tokens }),
6952 )
6953 .await;
6954 assert_eq!(status, StatusCode::OK, "{body}");
6955 let round_tripped = body["content"].as_str().expect("content").to_string();
6956 assert!(
6957 round_tripped.contains("hello") && round_tripped.contains("world"),
6958 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
6959 );
6960
6961 let (status, _) = post_json_uri(
6964 &app,
6965 "/v1/completions",
6966 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
6967 )
6968 .await;
6969 assert_eq!(
6970 status,
6971 StatusCode::NOT_IMPLEMENTED,
6972 "tokenizing an encoder must not have opened a path to generating with one"
6973 );
6974 }
6975
6976 #[tokio::test]
6981 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6982 use http_body_util::BodyExt;
6983 use tower::ServiceExt;
6984
6985 let fixture = concat!(
6986 "../frink-models/tests/fixtures/",
6987 "frink_real_moe_test.gguf"
6988 );
6989 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6990 let decoder = Decoder::from_gguf_with_expert_cache(
6991 &fixture,
6992 frink_models::config::test_moe_fixture(),
6993 Some(1024 * 1024),
6994 )
6995 .expect("MoE fixture must load store-backed");
6996
6997 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
7001 decoder.forward_token(1, 0, &mut caches);
7002
7003 let model = Model::Gguf(GgufModel {
7004 decoder: Arc::new(decoder),
7005 tokenizer: Arc::new(ServerTokenizer::Byte),
7006 stop_tokens: StopTokens::default(),
7007 bos_id: None,
7008 is_synthetic: false,
7009 chat_template: chat_template::PromptTemplate::plain(),
7010 });
7011 let state = Arc::new(test_state(
7012 model,
7013 ResponseCache::new(16, Duration::from_secs(60)),
7014 ));
7015 let app = Router::new()
7016 .route("/metrics", axum::routing::get(metrics))
7017 .route("/v1/chat/completions", post(chat_completions))
7018 .with_state(state);
7019
7020 let fetch_metrics = |app: Router| async move {
7021 let resp = app
7022 .oneshot(
7023 axum::http::Request::builder()
7024 .method("GET")
7025 .uri("/metrics")
7026 .body(axum::body::Body::empty())
7027 .unwrap(),
7028 )
7029 .await
7030 .unwrap();
7031 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7032 String::from_utf8(bytes.to_vec()).unwrap()
7033 };
7034
7035 let after = fetch_metrics(app.clone()).await;
7036 assert!(
7037 after.contains("frink_expert_cache_misses_total"),
7038 "streaming model must expose expert-cache metrics: {after}"
7039 );
7040 let misses: u64 = after
7041 .lines()
7042 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
7043 .and_then(|l| l.split_whitespace().nth(1))
7044 .and_then(|v| v.parse().ok())
7045 .expect("misses metric line must parse");
7046 assert!(
7047 misses > 0,
7048 "decode must have read experts through the store: {after}"
7049 );
7050 }
7051
7052 fn weather_tool() -> serde_json::Value {
7053 serde_json::json!({
7054 "type": "function",
7055 "function": {
7056 "name": "get_weather",
7057 "description": "Get the current weather for a location.",
7058 "parameters": {
7059 "type": "object",
7060 "properties": {"location": {"type": "string"}},
7061 "required": ["location"]
7062 }
7063 }
7064 })
7065 }
7066
7067 fn weather_tool_def() -> ToolDef {
7068 ToolDef {
7069 kind: "function".to_string(),
7070 function: ToolFunctionDef {
7071 name: "get_weather".to_string(),
7072 description: Some("Get the current weather for a location.".to_string()),
7073 parameters: Some(serde_json::json!({
7074 "type": "object",
7075 "properties": {"location": {"type": "string"}},
7076 "required": ["location"]
7077 })),
7078 },
7079 }
7080 }
7081
7082 #[test]
7083 fn tool_preamble_mentions_every_tool_name_and_description() {
7084 let preamble = tool_preamble(&[weather_tool_def()]);
7085 assert!(preamble.contains("get_weather"));
7086 assert!(preamble.contains("Get the current weather for a location."));
7087 assert!(preamble.contains("<tool_call>"));
7088 assert!(preamble.contains("</tool_call>"));
7089 }
7090
7091 #[test]
7092 fn a_real_marker_becomes_a_structured_tool_call() {
7093 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7094 let (message, finish) = build_response_message(
7095 text.to_string(),
7096 &[weather_tool_def()],
7097 output::OutputPosture::for_model("test-model"),
7098 "stop",
7099 );
7100 assert_eq!(finish, "tool_calls");
7101 let calls = message.tool_calls.expect("must carry a tool call");
7102 assert_eq!(calls[0].function.name, "get_weather");
7103 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7104 assert_eq!(parsed["location"], "Paris");
7105 }
7106
7107 #[test]
7108 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7109 let (message, finish) = build_response_message(
7110 "just an answer".to_string(),
7111 &[weather_tool_def()],
7112 output::OutputPosture::for_model("test-model"),
7113 "stop",
7114 );
7115 assert_eq!(finish, "stop");
7116 assert!(message.tool_calls.is_none());
7117 assert_eq!(message.content.as_deref(), Some("just an answer"));
7118 }
7119
7120 #[test]
7123 fn a_malformed_payload_is_not_a_tool_call() {
7124 let (message, finish) = build_response_message(
7125 "<tool_call>not valid json at all</tool_call>".to_string(),
7126 &[weather_tool_def()],
7127 output::OutputPosture::for_model("test-model"),
7128 "stop",
7129 );
7130 assert_eq!(finish, "stop");
7131 assert!(message.tool_calls.is_none());
7132 }
7133
7134 #[test]
7137 fn a_tool_that_was_never_offered_is_not_returned() {
7138 let (message, finish) = build_response_message(
7139 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7140 &[weather_tool_def()],
7141 output::OutputPosture::for_model("test-model"),
7142 "stop",
7143 );
7144 assert_eq!(finish, "stop");
7145 assert!(message.tool_calls.is_none());
7146 }
7147
7148 #[test]
7150 fn marker_text_with_no_tools_offered_stays_content() {
7151 let (message, finish) = build_response_message(
7152 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7153 &[],
7154 output::OutputPosture::for_model("test-model"),
7155 "stop",
7156 );
7157 assert_eq!(finish, "stop");
7158 assert!(message.tool_calls.is_none());
7159 assert!(message.content.is_some());
7160 }
7161
7162 #[test]
7166 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7167 let opened = std::cell::Cell::new(0usize);
7168 let mut parser = crate::policy::parser::ToolCallParser::new(
7169 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7170 vec![
7171 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7172 "write_file",
7173 serde_json::json!({"type": "object", "properties": {
7174 "path": {"type": "string"},
7175 "contents": {"type": "string"}
7176 }}),
7177 ),
7178 ],
7179 );
7180 let wire = "<tool_call><function=write_file>\
7181 <parameter=path>\n/tmp/x\n</parameter>\
7182 <parameter=contents>\nhello world\n</parameter>\
7183 </function></tool_call>";
7184
7185 let mut deltas = Vec::new();
7186 let mut text = String::new();
7187 for piece in wire.as_bytes().chunks(7) {
7188 let chunk = String::from_utf8_lossy(piece).into_owned();
7189 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7190 text.push_str(&more_text);
7191 deltas.extend(more);
7192 }
7193 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7194 text.push_str(&more_text);
7195 deltas.extend(more);
7196
7197 assert_eq!(opened.get(), 1, "one call opened");
7198 assert!(text.is_empty(), "the markers are not content: {text:?}");
7199
7200 let first = &deltas[0];
7201 assert_eq!(first.index, 0);
7202 assert_eq!(first.id.as_deref(), Some("call_0"));
7203 assert_eq!(first.kind, Some("function"));
7204 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7205
7206 let joined: String = deltas
7209 .iter()
7210 .filter_map(|d| d.function.arguments.clone())
7211 .collect();
7212 let parsed: serde_json::Value =
7213 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7214 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7215 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7216 assert!(
7217 deltas.len() >= 3,
7218 "the arguments arrived in pieces, not whole: {}",
7219 deltas.len()
7220 );
7221 assert!(
7222 deltas[1..].iter().all(|d| d.function.name.is_none()),
7223 "only the opening delta carries identity"
7224 );
7225 }
7226
7227 #[test]
7229 fn text_around_a_streamed_call_is_still_content() {
7230 let opened = std::cell::Cell::new(0usize);
7231 let mut parser = crate::policy::parser::ToolCallParser::new(
7232 crate::policy::parser::ToolCallFormat::Qwen25,
7233 vec![crate::policy::parser::tool_call::ToolSchema::new(
7234 "get_weather",
7235 )],
7236 );
7237 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7238 \"arguments\": {}}</tool_call> done";
7239 let mut text = String::new();
7240 for piece in wire.as_bytes().chunks(5) {
7241 let chunk = String::from_utf8_lossy(piece).into_owned();
7242 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7243 text.push_str(&more);
7244 }
7245 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7246 text.push_str(&more);
7247
7248 assert_eq!(opened.get(), 1);
7249 assert!(text.starts_with("let me check. "), "{text:?}");
7250 assert!(text.ends_with(" done"), "{text:?}");
7251 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7252 }
7253
7254 #[test]
7257 fn a_reasoning_block_is_split_out_of_the_answer() {
7258 let (message, finish) = build_response_message(
7259 "<think>weighing it up</think>The answer is 4.".to_string(),
7260 &[],
7261 output::OutputPosture::for_model("Qwen3-8B"),
7262 "stop",
7263 );
7264 assert_eq!(finish, "stop");
7265 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7266 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7267 }
7268
7269 #[test]
7272 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7273 let (message, _) = build_response_message(
7274 "Use the <think> tag like this.".to_string(),
7275 &[],
7276 output::OutputPosture::for_model("llama-3.1-8b"),
7277 "stop",
7278 );
7279 assert_eq!(
7280 message.content.as_deref(),
7281 Some("Use the <think> tag like this.")
7282 );
7283 assert!(message.reasoning_content.is_none());
7284 }
7285
7286 #[tokio::test]
7294 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7295 let app = test_app();
7296 let body = serde_json::json!({
7297 "model": "m",
7298 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7299 "max_tokens": 4,
7300 "temperature": 0,
7301 });
7302 let resp = post_json(&app, body).await;
7303 let message = &resp["choices"][0]["message"];
7304 assert!(message["content"].is_string());
7305 assert!(message.get("tool_calls").is_none());
7306 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7307 assert_eq!(resp["usage"]["completion_tokens"], 4);
7308 assert_eq!(
7309 resp["usage"]["total_tokens"],
7310 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7311 );
7312 }
7313
7314 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7315 use http_body_util::BodyExt;
7316 use tower::ServiceExt;
7317
7318 let response = app
7319 .clone()
7320 .oneshot(
7321 axum::http::Request::builder()
7322 .method("GET")
7323 .uri(uri)
7324 .body(axum::body::Body::empty())
7325 .unwrap(),
7326 )
7327 .await
7328 .unwrap();
7329 let status = response.status();
7330 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7331 (status, serde_json::from_slice(&bytes).unwrap())
7332 }
7333
7334 #[tokio::test]
7335 async fn health_answers_a_capability_handshake_not_a_boolean() {
7336 let app = test_app();
7337 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7338 assert_eq!(status, StatusCode::OK);
7339
7340 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7341 assert_eq!(health.state, frink_api::HealthState::Ready);
7342 assert!(health.pid > 0);
7343 assert!(health.server_time_unix_ms > 0);
7344 assert_eq!(health.last_request_age_seconds, None);
7347
7348 for id in [
7351 frink_api::health::capability::CPU,
7352 frink_api::health::capability::METAL,
7353 frink_api::health::capability::CUDA,
7354 frink_api::health::capability::REAL_WEIGHTS,
7355 frink_api::health::capability::CONTINUOUS_BATCHING,
7356 ] {
7357 let cap = health
7358 .capability(id)
7359 .unwrap_or_else(|| panic!("{id} missing"));
7360 assert!(!cap.reason.is_empty(), "{cap:?}");
7361 assert!(!cap.detail.is_empty(), "{cap:?}");
7362 }
7363 let weights = health
7367 .capability(frink_api::health::capability::REAL_WEIGHTS)
7368 .unwrap();
7369 assert!(!weights.available);
7370 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7371 assert!(health.model.as_ref().unwrap().synthetic_weights);
7372 }
7373
7374 #[tokio::test]
7375 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7376 let app = test_app();
7377 let _ = post_json(
7378 &app,
7379 serde_json::json!({
7380 "model": "m",
7381 "messages": [{"role": "user", "content": "\u{1}"}],
7382 "max_tokens": 1,
7383 "temperature": 0,
7384 }),
7385 )
7386 .await;
7387 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7388 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7389 let age = health
7390 .last_request_age_seconds
7391 .expect("a served request is evidence of liveness");
7392 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7393 }
7394
7395 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7397 use http_body_util::BodyExt;
7398 use tower::ServiceExt;
7399
7400 let response = app
7401 .clone()
7402 .oneshot(
7403 axum::http::Request::builder()
7404 .method("POST")
7405 .uri("/v1/chat/completions")
7406 .header("content-type", "application/json")
7407 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7408 .unwrap(),
7409 )
7410 .await
7411 .unwrap();
7412 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7413 String::from_utf8(bytes.to_vec())
7414 .unwrap()
7415 .lines()
7416 .filter_map(|line| line.strip_prefix("data: "))
7417 .filter(|payload| *payload != "[DONE]")
7418 .map(|payload| serde_json::from_str(payload).unwrap())
7419 .collect()
7420 }
7421
7422 #[tokio::test]
7423 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7424 let app = test_app();
7425 let chunks = post_sse_chunks(
7426 &app,
7427 serde_json::json!({
7428 "model": "m",
7429 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7430 "max_tokens": 4,
7431 "temperature": 0,
7432 "stream": true,
7433 }),
7434 )
7435 .await;
7436
7437 assert!(!chunks.is_empty());
7438 let request_id = chunks[0]["request_id"]
7439 .as_str()
7440 .expect("the first chunk names the request")
7441 .to_string();
7442 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7443 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7446 assert!(
7447 chunk.get("request_id").is_none(),
7448 "chunk {i} repeats request_id"
7449 );
7450 }
7451 for chunk in &chunks {
7454 assert_eq!(chunk["id"], serde_json::json!(request_id));
7455 }
7456
7457 let other = post_sse_chunks(
7458 &app,
7459 serde_json::json!({
7460 "model": "m",
7461 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7462 "max_tokens": 4,
7463 "temperature": 0,
7464 "stream": true,
7465 }),
7466 )
7467 .await;
7468 assert_ne!(
7469 other[0]["request_id"].as_str().unwrap(),
7470 request_id,
7471 "two concurrent chats must not share an id"
7472 );
7473 }
7474
7475 #[tokio::test]
7476 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7477 let app = test_app();
7478 let resp = post_json(
7479 &app,
7480 serde_json::json!({
7481 "model": "m",
7482 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7483 "max_tokens": 2,
7484 "temperature": 0,
7485 }),
7486 )
7487 .await;
7488 assert_eq!(resp["id"], resp["request_id"]);
7489 assert!(resp["request_id"]
7490 .as_str()
7491 .unwrap()
7492 .starts_with("chatcmpl-"));
7493 }
7494
7495 #[tokio::test]
7498 async fn usage_carries_separate_prefill_and_decode_timings() {
7499 let app = test_app();
7500 let resp = post_json(
7501 &app,
7502 serde_json::json!({
7503 "model": "m",
7504 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7505 "max_tokens": 4,
7506 "temperature": 0,
7507 }),
7508 )
7509 .await;
7510 let usage = &resp["usage"];
7511 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7512 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7513 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7514 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7515 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7517 }
7518
7519 #[tokio::test]
7527 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7528 let app = test_app();
7529 let body = serde_json::json!({
7530 "model": "m",
7531 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7532 "max_tokens": 4,
7533 "temperature": 0,
7534 "tools": [weather_tool()],
7535 });
7536 let resp = post_json(&app, body).await;
7537 let message = &resp["choices"][0]["message"];
7538 assert!(
7539 message["content"].is_string(),
7540 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7541 );
7542 assert!(message.get("tool_calls").is_none());
7543 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7547 }
7548
7549 #[tokio::test]
7553 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7554 let app = test_app();
7555 let body = serde_json::json!({
7556 "model": "m",
7557 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7558 "max_tokens": 3,
7559 "temperature": 0,
7560 });
7561 let first = post_json(&app, body.clone()).await;
7562 assert_eq!(first["frink_cache"], "miss");
7563 let second = post_json(&app, body).await;
7564 assert_eq!(second["frink_cache"], "hit");
7565 assert_eq!(
7566 first["choices"][0]["message"]["content"],
7567 second["choices"][0]["message"]["content"]
7568 );
7569 assert_eq!(
7570 first["choices"][0]["finish_reason"],
7571 second["choices"][0]["finish_reason"]
7572 );
7573 assert_eq!(first["usage"], second["usage"]);
7574 assert_eq!(second["usage"]["completion_tokens"], 3);
7575 }
7576
7577 #[tokio::test]
7588 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7589 let app = test_app();
7590 let plain = serde_json::json!({
7591 "model": "m",
7592 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7593 "max_tokens": 3,
7594 "temperature": 0,
7595 });
7596
7597 let first = post_json(&app, plain.clone()).await;
7598 assert_eq!(first["frink_cache"], "miss");
7599 let unconstrained = first["choices"][0]["message"]["content"]
7600 .as_str()
7601 .expect("content")
7602 .to_string();
7603
7604 let mut constrained = plain.clone();
7605 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7606 let second = post_json(&app, constrained).await;
7607 assert_eq!(
7608 second["frink_cache"], "miss",
7609 "a grammar is part of the key, so this body has never been answered"
7610 );
7611 let constrained_answer = second["choices"][0]["message"]["content"]
7615 .as_str()
7616 .expect("content")
7617 .to_string();
7618 assert!(
7619 constrained_answer.contains("-> \"yes\"]"),
7620 "the grammar must have been compiled AND applied, not skipped \
7621 by a cache hit: {constrained_answer}"
7622 );
7623 assert_ne!(
7624 constrained_answer, unconstrained,
7625 "the constrained request was served the unconstrained answer"
7626 );
7627
7628 let third = post_json(&app, plain).await;
7632 assert_eq!(third["frink_cache"], "hit");
7633 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7634 }
7635
7636 #[tokio::test]
7658 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7659 let state = Arc::new(test_state(
7660 test_model_full_byte_vocab(),
7661 ResponseCache::new(1000, Duration::from_secs(3600)),
7662 ));
7663 let app = test_app_with_state(state.clone());
7664 let plain = serde_json::json!({
7665 "model": "m",
7666 "messages": [
7667 {"role": "system", "content": "Answer in JSON when it helps."},
7668 {"role": "user", "content": "\u{1}\u{2}"},
7669 ],
7670 "max_tokens": 3,
7671 "temperature": 0,
7672 });
7673
7674 let first = post_json(&app, plain.clone()).await;
7675 assert_eq!(first["frink_cache"], "miss");
7676 assert_eq!(state.cache_stats().entries, 1);
7677
7678 let mut as_json = plain.clone();
7679 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7680 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7681 assert_eq!(
7682 status,
7683 StatusCode::BAD_REQUEST,
7684 "the demo banner is not a JSON object, whoever generated it"
7685 );
7686 assert_eq!(
7687 state.cache_stats().hits,
7688 0,
7689 "a json_object request must not be answered from an entry the \
7690 JSON mask never produced"
7691 );
7692 assert_eq!(
7693 state.cache_stats().entries,
7694 2,
7695 "json_object must key its own entry, not reuse the unconstrained \
7696 one it happens to render the same prompt as"
7697 );
7698 }
7699
7700 #[tokio::test]
7712 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7713 let app = test_app_with_state(Arc::new(test_state(
7714 test_model_full_byte_vocab_with_eos(Some(0x77)),
7715 ResponseCache::new(1000, Duration::from_secs(3600)),
7716 )));
7717 let body = serde_json::json!({
7718 "model": "m",
7719 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7720 "max_tokens": 6,
7721 "temperature": 0,
7722 });
7723
7724 let stopped = post_json(&app, body.clone()).await;
7725 assert_eq!(stopped["frink_cache"], "miss");
7726 assert_eq!(
7727 stopped["choices"][0]["finish_reason"], "stop",
7728 "the fixture is only meaningful if the model's EOS really fires here"
7729 );
7730 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7731
7732 let mut ignoring = body.clone();
7733 ignoring["ignore_eos"] = serde_json::json!(true);
7734 let ran_on = post_json(&app, ignoring).await;
7735 assert_eq!(
7736 ran_on["frink_cache"], "miss",
7737 "ignore_eos is part of the key, so this body has never been answered"
7738 );
7739 assert_eq!(
7740 ran_on["usage"]["completion_tokens"], 6,
7741 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7742 );
7743 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7744 assert_ne!(
7745 ran_on["choices"][0]["message"]["content"],
7746 stopped["choices"][0]["message"]["content"]
7747 );
7748 }
7749
7750 #[tokio::test]
7756 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7757 let session_app = test_app();
7758 let manual_app = test_app();
7759
7760 let turn1 = post_json(
7762 &session_app,
7763 serde_json::json!({
7764 "model": "m",
7765 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7766 "session_id": "s1",
7767 "max_tokens": 5,
7768 "temperature": 0,
7769 }),
7770 )
7771 .await;
7772 let reply1 = turn1["choices"][0]["message"]["content"]
7773 .as_str()
7774 .unwrap()
7775 .to_string();
7776
7777 let manual_turn1 = post_json(
7782 &manual_app,
7783 serde_json::json!({
7784 "model": "m",
7785 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7786 "max_tokens": 5,
7787 "temperature": 0,
7788 }),
7789 )
7790 .await;
7791 assert_eq!(
7792 manual_turn1["choices"][0]["message"]["content"]
7793 .as_str()
7794 .unwrap(),
7795 reply1
7796 );
7797
7798 let turn2 = post_json(
7800 &session_app,
7801 serde_json::json!({
7802 "model": "m",
7803 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7804 "session_id": "s1",
7805 "max_tokens": 5,
7806 "temperature": 0,
7807 }),
7808 )
7809 .await;
7810 let reply2 = turn2["choices"][0]["message"]["content"]
7811 .as_str()
7812 .unwrap()
7813 .to_string();
7814
7815 let manual_turn2 = post_json(
7819 &manual_app,
7820 serde_json::json!({
7821 "model": "m",
7822 "messages": [
7823 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7824 {"role": "assistant", "content": reply1},
7825 {"role": "user", "content": "\u{4}\u{5}"},
7826 ],
7827 "max_tokens": 5,
7828 "temperature": 0,
7829 }),
7830 )
7831 .await;
7832 assert_eq!(
7833 manual_turn2["choices"][0]["message"]["content"]
7834 .as_str()
7835 .unwrap(),
7836 reply2,
7837 "resuming a session must produce identical output to manually resending the full history"
7838 );
7839 }
7840
7841 #[test]
7844 fn lock_cache_recovers_from_a_poisoned_mutex() {
7845 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7846
7847 let poison_cache = Arc::clone(&cache);
7848 let _ = std::thread::spawn(move || {
7849 let _guard = poison_cache.lock().unwrap();
7850 panic!("simulated panic while holding the lock");
7851 })
7852 .join();
7853
7854 let recovered = lock_cache(&cache);
7856 assert_eq!(recovered.stats().entries, 0);
7857 }
7858
7859 #[test]
7860 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7861 let mut req_body = serde_json::json!({
7862 "model": "m",
7863 "messages": [{"role": "user", "content": "hi"}],
7864 });
7865 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7866 assert!(
7867 req.is_cacheable(),
7868 "default (temperature 0) must be cacheable"
7869 );
7870
7871 req_body["temperature"] = serde_json::json!(0.8);
7872 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7873 assert!(
7874 !req.is_cacheable(),
7875 "unseeded sampling must never be cacheable"
7876 );
7877
7878 req_body["seed"] = serde_json::json!(42);
7879 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7880 assert!(
7881 req.is_cacheable(),
7882 "sampling with an explicit seed is deterministic and must be cacheable"
7883 );
7884 }
7885
7886 const GRADED: &str = "{% if reasoning_effort %}\
7890 {% if reasoning_effort not in ['low','medium','high'] %}\
7891 {{ raise_exception('unsupported effort') }}\
7892 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7893 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7894
7895 fn graded_template() -> chat_template::PromptTemplate {
7896 chat_template::PromptTemplate::from_gguf_metadata(
7897 Some(GRADED),
7898 Some("qwen3"),
7899 false,
7900 true,
7901 None,
7902 None,
7903 )
7904 }
7905
7906 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7907 serde_json::from_value(value).expect("request")
7908 }
7909
7910 #[test]
7917 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7918 let req = chat_request(serde_json::json!({
7919 "model": "m",
7920 "messages": [{"role": "user", "content": "hi"}],
7921 "grammar": "root ::= \"a\"+",
7922 }));
7923 req.validate_supported_fields()
7924 .expect("a valid grammar is a valid request");
7925 let params = req
7926 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7927 .expect("a valid grammar compiles at params time too");
7928 assert!(
7929 params.grammar.is_some(),
7930 "the grammar was dropped between the wire and the sampler"
7931 );
7932 assert!(
7933 params.needs_vocab_logits(),
7934 "a grammar request that may fold lm_head into a GPU argmax is \
7935 a grammar request served unconstrained"
7936 );
7937
7938 let plain = chat_request(serde_json::json!({
7939 "model": "m",
7940 "messages": [{"role": "user", "content": "hi"}],
7941 }));
7942 assert!(plain
7943 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7944 .unwrap()
7945 .grammar
7946 .is_none());
7947 }
7948
7949 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7950 chat_request(serde_json::json!({
7951 "model": "m",
7952 "messages": [{"role": "user", "content": "weather in Rome?"}],
7953 "tools": [weather_tool()],
7954 "tool_choice": tool_choice,
7955 }))
7956 }
7957
7958 #[test]
7962 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7963 for choice in [
7964 serde_json::json!("required"),
7965 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7966 ] {
7967 let req = tool_request(choice.clone());
7968 req.validate_supported_fields()
7969 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7970 let params = req
7971 .generation_params_for_template(
7972 &graded_template(),
7973 "Qwen3-8B",
7974 crate::sampling_knobs::SamplerModel::absent(),
7975 )
7976 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7977 let grammar = params
7978 .grammar
7979 .as_ref()
7980 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7981 assert!(
7982 grammar.is_awaiting_trigger(),
7983 "the model must be free to think before it calls"
7984 );
7985 assert!(
7986 !grammar.allows_eog(),
7987 "{choice} must not be able to end the turn without a call"
7988 );
7989 assert!(
7995 params.needs_vocab_logits(),
7996 "{choice} would let a backend return a token id instead of logits"
7997 );
7998 assert!(
7999 !generate::greedy_gpu_fold_allowed(¶ms),
8000 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8001 );
8002 }
8003 }
8004
8005 #[test]
8007 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8008 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8009 let req = tool_request(choice.clone());
8010 req.validate_supported_fields().expect("still supported");
8011 let params = match req.generation_params_for_template(
8012 &graded_template(),
8013 "Qwen3-8B",
8014 crate::sampling_knobs::SamplerModel::absent(),
8015 ) {
8016 Ok(p) => p,
8017 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
8018 };
8019 assert!(
8020 params.grammar.is_none(),
8021 "{choice} does not force a call and must not be constrained"
8022 );
8023 }
8024 }
8025
8026 #[test]
8029 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8030 let req = chat_request(serde_json::json!({
8032 "model": "m",
8033 "messages": [{"role": "user", "content": "hi"}],
8034 "tool_choice": "required",
8035 }));
8036 let (status, _) = req
8037 .validate_supported_fields()
8038 .expect_err("nothing to call");
8039 assert_eq!(status, StatusCode::BAD_REQUEST);
8040
8041 let req =
8043 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8044 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8045 assert_eq!(status, StatusCode::BAD_REQUEST);
8046 assert_eq!(body["error"]["param"], "tool_choice");
8047
8048 let req = tool_request(serde_json::json!({"type": "function"}));
8050 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8051 assert_eq!(status, StatusCode::BAD_REQUEST);
8052
8053 let req = chat_request(serde_json::json!({
8055 "model": "m",
8056 "messages": [{"role": "user", "content": "hi"}],
8057 "tools": [weather_tool()],
8058 "tool_choice": "required",
8059 "grammar": "root ::= \"a\"+",
8060 }));
8061 let (status, _) = req
8062 .validate_supported_fields()
8063 .expect_err("a grammar and a forced call are two constraints");
8064 assert_eq!(status, StatusCode::BAD_REQUEST);
8065
8066 let req = tool_request(serde_json::json!("required"));
8072 let (status, Json(body)) = match req.generation_params_for_template(
8073 &graded_template(),
8074 "muse-glimmer-8b",
8075 crate::sampling_knobs::SamplerModel::absent(),
8076 ) {
8077 Err(e) => e,
8078 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
8079 };
8080 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8081 assert!(
8082 body["error"]["message"]
8083 .as_str()
8084 .unwrap()
8085 .contains("muse_glimmer"),
8086 "{body}"
8087 );
8088
8089 let req = tool_request(serde_json::json!("required"));
8095 let params = req
8096 .generation_params_for_template(
8097 &graded_template(),
8098 "gemma-4-E2B-it",
8099 crate::sampling_knobs::SamplerModel::absent(),
8100 )
8101 .expect("a gemma4 forced tool_choice is served");
8102 assert!(
8103 params.grammar.is_some(),
8104 "a forced tool_choice must arrive as the generation's grammar"
8105 );
8106 }
8107
8108 #[test]
8111 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8112 let req = chat_request(serde_json::json!({
8113 "model": "m",
8114 "messages": [{"role": "user", "content": "hi"}],
8115 "grammar": "root ::= \"a",
8116 }));
8117 let (status, Json(body)) = req
8118 .validate_supported_fields()
8119 .expect_err("this does not parse");
8120 assert_eq!(status, StatusCode::BAD_REQUEST);
8121 assert_eq!(body["error"]["param"], "grammar");
8122 assert!(
8123 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8124 .is_err(),
8125 "and again at params time"
8126 );
8127 }
8128
8129 #[test]
8136 fn response_format_json_schema_becomes_the_requests_grammar() {
8137 let req = chat_request(serde_json::json!({
8138 "model": "m",
8139 "messages": [{"role": "user", "content": "hi"}],
8140 "response_format": {
8141 "type": "json_schema",
8142 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8143 },
8144 }));
8145 req.validate_supported_fields()
8146 .expect("a boolean schema converts");
8147 let params = req
8148 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8149 .expect("and compiles");
8150 let grammar = params.grammar.expect("the schema is the grammar");
8151 let mut g = (*grammar).clone();
8152 g.accept_token(0, b"true").expect("a boolean is accepted");
8153 assert!(g.allows_eog(), "and completes the parse");
8154 assert!(
8155 !params.json_object,
8156 "a schema is not the json_object character-class mask"
8157 );
8158 }
8159
8160 #[test]
8164 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8165 let req = chat_request(serde_json::json!({
8166 "model": "m",
8167 "messages": [{"role": "user", "content": "hi"}],
8168 "response_format": {
8169 "type": "json_schema",
8170 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8171 },
8172 }));
8173 let (status, Json(body)) = req
8174 .validate_supported_fields()
8175 .expect_err("minimum has no grammar in this port");
8176 assert_eq!(status, StatusCode::BAD_REQUEST);
8177 assert!(
8178 body["error"]["message"]
8179 .as_str()
8180 .expect("a message")
8181 .contains("minimum"),
8182 "the refusal must name the keyword: {body}"
8183 );
8184 assert!(
8185 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8186 .is_err(),
8187 "and again at params time"
8188 );
8189 }
8190
8191 #[test]
8197 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8198 let req = chat_request(serde_json::json!({
8199 "model": "m",
8200 "messages": [{"role": "user", "content": "hi"}],
8201 "tool_choice": "required",
8202 "tools": [{
8203 "type": "function",
8204 "function": {"name": "f", "parameters": {"type": "object"}},
8205 }],
8206 "response_format": {
8207 "type": "json_schema",
8208 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8209 },
8210 }));
8211 let (status, Json(body)) = req
8212 .validate_supported_fields()
8213 .expect_err("two constraints, one generation");
8214 assert_eq!(status, StatusCode::BAD_REQUEST);
8215 assert_eq!(body["error"]["param"], "tool_choice");
8216 }
8217
8218 #[test]
8221 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8222 let req = chat_request(serde_json::json!({
8223 "model": "m",
8224 "messages": [{"role": "user", "content": "hi"}],
8225 }));
8226 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8227 }
8228
8229 #[test]
8234 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8235 let asked = chat_request(serde_json::json!({
8236 "model": "m",
8237 "messages": [{"role": "user", "content": "hi"}],
8238 "min_p": 0.07,
8239 }));
8240 assert_eq!(
8241 asked
8242 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8243 .expect("knobs")
8244 .min_p,
8245 0.07
8246 );
8247
8248 let silent = chat_request(serde_json::json!({
8249 "model": "m",
8250 "messages": [{"role": "user", "content": "hi"}],
8251 }));
8252 assert_eq!(
8253 silent
8254 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8255 .expect("knobs")
8256 .min_p,
8257 0.0,
8258 "an unset min_p must be off, not llama.cpp's CLI default"
8259 );
8260 }
8261
8262 #[test]
8271 fn no_sampler_knob_is_missing_from_the_cache_key() {
8272 let base = serde_json::json!({
8273 "model": "m",
8274 "messages": [{"role": "user", "content": "hi"}],
8275 "seed": 1,
8276 });
8277 let key_for = |body: serde_json::Value| {
8278 let req = chat_request(body);
8279 let params = req
8280 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8281 .expect("params");
8282 req.cache_key("prompt", ¶ms)
8283 };
8284 let baseline = key_for(base.clone());
8285 for (knob, value) in [
8286 ("temperature", serde_json::json!(0.5)),
8287 ("top_p", serde_json::json!(0.9)),
8288 ("min_p", serde_json::json!(0.05)),
8289 ("top_k", serde_json::json!(40)),
8290 ("repetition_penalty", serde_json::json!(1.1)),
8291 ("presence_penalty", serde_json::json!(0.3)),
8292 ("frequency_penalty", serde_json::json!(0.3)),
8293 (
8294 "samplers",
8295 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8296 ),
8297 ] {
8298 let mut body = base.clone();
8299 body[knob] = value;
8300 assert_ne!(
8301 key_for(body),
8302 baseline,
8303 "`{knob}` is not in the cache key: two requests differing \
8304 only in it would share one cached answer"
8305 );
8306 }
8307 }
8308
8309 #[test]
8319 fn no_constraint_is_missing_from_the_cache_key() {
8320 let base = serde_json::json!({
8321 "model": "m",
8322 "messages": [{"role": "user", "content": "pick one"}],
8323 });
8324 let key_for = |body: serde_json::Value| {
8325 let req = chat_request(body);
8326 let params = req
8327 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8328 .expect("params");
8329 req.cache_key("prompt", ¶ms)
8330 };
8331 let baseline = key_for(base.clone());
8332 for (field, value) in [
8333 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8334 (
8335 "response_format",
8336 serde_json::json!({"type": "json_object"}),
8337 ),
8338 (
8339 "response_format",
8340 serde_json::json!({"type": "json_schema", "json_schema": {
8341 "name": "answer",
8342 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8343 }}),
8344 ),
8345 ("ignore_eos", serde_json::json!(true)),
8346 ("stop", serde_json::json!(["\n"])),
8347 ("max_tokens", serde_json::json!(7)),
8348 ] {
8349 let mut body = base.clone();
8350 body[field] = value.clone();
8351 assert_ne!(
8352 key_for(body),
8353 baseline,
8354 "`{field}: {value}` is not in the cache key: two requests \
8355 differing only in it would share one cached answer"
8356 );
8357 }
8358 }
8359
8360 #[test]
8364 fn an_explicit_zero_output_budget_is_a_client_error() {
8365 let req = chat_request(serde_json::json!({
8366 "model": "m",
8367 "messages": [{"role": "user", "content": "hi"}],
8368 "max_tokens": 0,
8369 }));
8370 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8371 assert_eq!(status, StatusCode::BAD_REQUEST);
8372 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8373 }
8374
8375 #[test]
8378 fn a_request_can_turn_thinking_off() {
8379 let template = graded_template();
8380 for body in [
8381 serde_json::json!({
8382 "model": "m",
8383 "messages": [{"role": "user", "content": "hi"}],
8384 "reasoning_effort": "none",
8385 }),
8386 serde_json::json!({
8387 "model": "m",
8388 "messages": [{"role": "user", "content": "hi"}],
8389 "thinking": {"type": "disabled"},
8390 }),
8391 ] {
8392 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8393 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8394 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8395 assert!(!kwargs.contains_key("reasoning_effort"));
8398 }
8399 }
8400
8401 #[test]
8404 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8405 let template = graded_template();
8406 let kwargs = chat_request(serde_json::json!({
8407 "model": "m",
8408 "messages": [{"role": "user", "content": "hi"}],
8409 "reasoning_effort": "high",
8410 "thinking": {"type": "disabled"},
8411 }))
8412 .resolve_template_kwargs(&template);
8413 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8414 assert!(!kwargs.contains_key("reasoning_effort"));
8415 }
8416
8417 #[test]
8420 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8421 let req = chat_request(serde_json::json!({
8422 "model": "m",
8423 "messages": [{"role": "user", "content": "hi"}],
8424 "thinking": {"type": "disable"},
8425 }));
8426 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8427 assert_eq!(status, StatusCode::BAD_REQUEST);
8428 }
8429
8430 #[test]
8433 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8434 let template = graded_template();
8435 let kwargs = chat_request(serde_json::json!({
8436 "model": "m",
8437 "messages": [{"role": "user", "content": "hi"}],
8438 "reasoning_effort": "none",
8439 "chat_template_kwargs": {"enable_thinking": true},
8440 }))
8441 .resolve_template_kwargs(&template);
8442 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8443 }
8444
8445 #[test]
8449 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8450 let template = graded_template();
8451 let req = chat_request(serde_json::json!({
8452 "model": "m",
8453 "messages": [{"role": "user", "content": "hi"}],
8454 "reasoning_effort": "minimal",
8455 }));
8456 let kwargs = req.resolve_template_kwargs(&template);
8457 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8458 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8459 assert!(prompt.starts_with("E:low|"), "{prompt}");
8460 }
8461
8462 #[test]
8466 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8467 let template = graded_template();
8468 let req = chat_request(serde_json::json!({
8469 "model": "m",
8470 "messages": [{"role": "user", "content": "hi"}],
8471 "chat_template_kwargs": {"reasoning_effort": "none"},
8472 }));
8473 let kwargs = req.resolve_template_kwargs(&template);
8474 assert!(!kwargs.contains_key("reasoning_effort"));
8475 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8476 assert_eq!(prompt, "hi");
8477 }
8478
8479 #[test]
8482 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8483 let template = graded_template();
8484 let req = chat_request(serde_json::json!({
8485 "model": "m",
8486 "messages": [{"role": "user", "content": "hi"}],
8487 "reasoning_effort": "low",
8488 "chat_template_kwargs": {"reasoning_effort": "high"},
8489 }));
8490 assert_eq!(
8491 req.resolve_template_kwargs(&template)["reasoning_effort"],
8492 serde_json::json!("high")
8493 );
8494 }
8495
8496 #[test]
8500 fn offering_tools_turns_thinking_on_by_itself() {
8501 let template = graded_template();
8502 let quiet = chat_request(serde_json::json!({
8503 "model": "m",
8504 "messages": [{"role": "user", "content": "hi"}],
8505 }));
8506 assert!(!quiet
8507 .resolve_template_kwargs(&template)
8508 .contains_key("enable_thinking"));
8509
8510 let with_tools = chat_request(serde_json::json!({
8511 "model": "m",
8512 "messages": [{"role": "user", "content": "hi"}],
8513 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8514 }));
8515 let kwargs = with_tools.resolve_template_kwargs(&template);
8516 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8517 let prompt =
8518 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8519 assert!(prompt.starts_with("THINK|"), "{prompt}");
8520 }
8521
8522 #[test]
8527 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8528 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8529 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8530 Some("qwen3"),
8531 false,
8532 true,
8533 None,
8534 None,
8535 );
8536 let req = chat_request(serde_json::json!({
8537 "model": "m",
8538 "messages": [{"role": "user", "content": "hi"}],
8539 "chat_template_kwargs": {"enable_thinking": true},
8540 }));
8541 let kwargs = req.resolve_template_kwargs(&opener);
8542 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8543 assert!(prompt.ends_with("<think>"), "{prompt}");
8544
8545 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8548 let (message, _) = build_response_message(
8549 "weighing it up</think>Paris.".to_string(),
8550 &[],
8551 posture,
8552 "stop",
8553 );
8554 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8555 assert_eq!(message.content.as_deref(), Some("Paris."));
8556
8557 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8560 let (message, _) = build_response_message(
8561 "weighing it up</think>Paris.".to_string(),
8562 &[],
8563 closed,
8564 "stop",
8565 );
8566 assert_eq!(message.reasoning_content, None);
8567 }
8568
8569 #[test]
8570 fn stop_param_accepts_both_single_string_and_array() {
8571 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8572 "model": "m",
8573 "messages": [{"role": "user", "content": "hi"}],
8574 "stop": "END",
8575 }))
8576 .unwrap();
8577 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8578
8579 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8580 "model": "m",
8581 "messages": [{"role": "user", "content": "hi"}],
8582 "stop": ["A", "B"],
8583 }))
8584 .unwrap();
8585 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8586 }
8587
8588 #[test]
8589 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8590 let model = test_model();
8591 let result = run_generation(
8592 &model,
8593 "hello",
8594 &greedy_params(4),
8595 None,
8596 None,
8597 None,
8598 None,
8599 None,
8600 None,
8601 );
8602 assert!(matches!(
8603 result,
8604 Err(generate::DecodeError::TokenOutOfVocab { .. })
8605 ));
8606 }
8607
8608 #[test]
8612 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8613 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8615 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8616
8617 let holder_pool = Arc::clone(&pool);
8618 let holder = std::thread::spawn(move || {
8619 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8620 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8622 drop(held);
8623 });
8624 std::thread::sleep(Duration::from_millis(15));
8625
8626 let config = generate::KvPoolConfig {
8627 pool,
8628 queue_wait: Duration::ZERO,
8629 };
8630 let result = run_generation(
8631 &model,
8632 &prompt,
8633 &greedy_params(4),
8634 Some(&config),
8635 None,
8636 None,
8637 None,
8638 None,
8639 None,
8640 );
8641 assert!(matches!(
8642 result,
8643 Err(generate::DecodeError::KvPoolExhausted)
8644 ));
8645
8646 let (status, _body) = decode_error_response(result.unwrap_err());
8647 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8648 holder.join().unwrap();
8649 }
8650
8651 #[test]
8659 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8660 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8662 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8664 let config = generate::KvPoolConfig {
8665 pool,
8666 queue_wait: Duration::ZERO,
8667 };
8668
8669 let result = run_generation(
8670 &model,
8671 &prompt,
8672 &greedy_params(4),
8673 Some(&config),
8674 None,
8675 None,
8676 None,
8677 None,
8678 None,
8679 );
8680 let err = result.expect_err("one block cannot hold two layers' caches");
8681 assert!(
8682 matches!(
8683 &err,
8684 generate::DecodeError::KvBudgetExceeded { binding, .. }
8685 if *binding == frink_models::Ceiling::DeviceMemory.code()
8686 ),
8687 "expected an immovable device-memory refusal, got {err:?}"
8688 );
8689 let (status, _body) = decode_error_response(err);
8690 assert_eq!(status, StatusCode::BAD_REQUEST);
8691 }
8692
8693 #[test]
8699 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8700 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8701 queued: 512,
8702 cap: 512,
8703 });
8704 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8705 assert_eq!(body["error"]["retry_after_seconds"], 1);
8706 let message = body["error"]["message"].as_str().expect("message");
8707 assert!(message.contains("512"), "{message}");
8708 }
8709
8710 #[test]
8711 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8712 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8713 token: 99,
8714 vocab_size: 32,
8715 });
8716 assert!(
8717 body["error"]["retry_after_seconds"].is_null(),
8718 "retrying a prompt this model cannot tokenize never helps"
8719 );
8720 }
8721
8722 #[test]
8723 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8724 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8725 token: 99,
8726 vocab_size: 32,
8727 });
8728 assert_eq!(status, StatusCode::BAD_REQUEST);
8729 }
8730
8731 #[test]
8732 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8733 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8735 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8736 let config = generate::KvPoolConfig {
8737 pool: pool.clone(),
8738 queue_wait: Duration::ZERO,
8739 };
8740
8741 let (choices, _usage) = run_generation(
8742 &model,
8743 &prompt,
8744 &greedy_params(4),
8745 Some(&config),
8746 None,
8747 None,
8748 None,
8749 None,
8750 None,
8751 )
8752 .unwrap();
8753 assert_eq!(choices[0].finish, FinishReason::Length);
8754 assert_eq!(
8755 pool.lock().unwrap().free_blocks(),
8756 2,
8757 "a completed request must return its blocks to the pool"
8758 );
8759 }
8760
8761 #[tokio::test]
8766 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8767 let model = Arc::new(test_model());
8768 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8769
8770 let mut handles = Vec::new();
8771 for _ in 0..8 {
8772 let model = Arc::clone(&model);
8773 let prompt = prompt.clone();
8774 handles.push(tokio::task::spawn_blocking(move || {
8775 run_generation(
8776 &model,
8777 &prompt,
8778 &greedy_params(6),
8779 None,
8780 None,
8781 None,
8782 None,
8783 None,
8784 None,
8785 )
8786 .unwrap()
8787 }));
8788 }
8789
8790 let mut results = Vec::new();
8791 for h in handles {
8792 results.push(h.await.unwrap());
8793 }
8794 for r in &results[1..] {
8799 assert_eq!(r.0, results[0].0, "choices must match");
8803 assert_eq!(
8804 r.1.prompt_tokens, results[0].1.prompt_tokens,
8805 "prompt token count must match"
8806 );
8807 assert_eq!(
8808 r.1.completion_tokens, results[0].1.completion_tokens,
8809 "completion token count must match"
8810 );
8811 }
8812 }
8813
8814 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8821 let mut header_entries = Vec::new();
8822 let mut data = Vec::new();
8823 for (name, shape, values) in tensors {
8824 let start = data.len();
8825 for v in values {
8826 data.extend_from_slice(&v.to_le_bytes());
8827 }
8828 let end = data.len();
8829 let shape_str = shape
8830 .iter()
8831 .map(|d| d.to_string())
8832 .collect::<Vec<_>>()
8833 .join(",");
8834 header_entries.push(format!(
8835 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8836 ));
8837 }
8838 let header = format!("{{{}}}", header_entries.join(","));
8839 let header_bytes = header.as_bytes();
8840 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8841 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8842 out.extend_from_slice(header_bytes);
8843 out.extend_from_slice(&data);
8844 out
8845 }
8846
8847 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8858 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8859 use frink_models::kimi_loader::KimiRealHparams;
8860 use frink_moe::{GatingFunction, MoeLayerConfig};
8861
8862 let hidden_dim = 8;
8863 let kda_num_heads = 2;
8864 let kda_head_dim = 3;
8865 let kda_proj = kda_num_heads * kda_head_dim;
8866 let conv_kernel = 4;
8867 let dense_intermediate = 5;
8868 let vocab_size = 256;
8872 let mla_num_heads = 1;
8873 let mla_q_lora_rank = 2;
8874 let mla_kv_lora_rank = 2;
8875 let mla_qk_nope_head_dim = 2;
8876 let mla_qk_rope_head_dim = 2;
8877 let mla_v_head_dim = 2;
8878
8879 let model_cfg = frink_models::ModelConfig {
8880 rope_layers: frink_models::rope_layers::RopeLayers::All,
8881 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
8882 name: "synthetic-kimi-server-test",
8883 n_layers: 1,
8884 n_mtp_blocks: 0,
8885 hidden_dim,
8886 n_heads: 1,
8887 n_kv_heads: 1,
8888 head_dim: 4,
8889 v_head_dim: None,
8890 vocab_size,
8891 rope_theta: 10000.0,
8892 rms_norm_eps: 1e-5,
8893 post_norm_eps: 1e-5,
8894 sliding_window: None,
8895 moe: MoeLayerConfig {
8896 expert_weights_scale: 1.0,
8897 routed_weight_before_ffn: false,
8898 n_experts: 1,
8899 n_experts_active: 1,
8900 n_shared_experts: 0,
8901 hidden_dim,
8902 expert_ffn_dim: 4,
8903 gating: GatingFunction::Sigmoid,
8904 norm_topk_prob: true,
8905 expert_group_count: None,
8906 expert_group_used_count: None,
8907 },
8908 n_dense_leading_layers: 1,
8913 moe_interleave_step: None,
8914 norm_function: frink_models::norm::NormFunction::Rms,
8915 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8916 kda_layers: vec![1],
8917 full_attn_layers: vec![],
8918 mla: MlaConfig {
8919 num_heads: mla_num_heads,
8920 q_lora_rank: mla_q_lora_rank,
8921 kv_lora_rank: mla_kv_lora_rank,
8922 qk_nope_head_dim: mla_qk_nope_head_dim,
8923 qk_rope_head_dim: mla_qk_rope_head_dim,
8924 v_head_dim: mla_v_head_dim,
8925 use_output_gate: true,
8926 rope: None,
8927 },
8928 kda: KdaConfig {
8929 num_heads: kda_num_heads,
8930 head_dim: kda_head_dim,
8931 short_conv_kernel_size: conv_kernel,
8932 gate_lower_bound: -5.0,
8933 use_full_rank_gate: true,
8934 },
8935 }),
8936 rope_freqs: None,
8937 rope_attn_factor: 1.0,
8938 rope_dim: None,
8939 rope_dim_swa: None,
8940 rope_freqs_long: None,
8941 rope_freqs_short: None,
8942 rope_orig_ctx: None,
8943 rope_layout: frink_models::config::RopeLayout::Neox,
8944 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
8945 swa_layers: frink_models::swa_layers::SwaLayers::All,
8946 attn_logit_softcap: None,
8947 final_logit_softcap: None,
8948 embedding_scale: None,
8949 residual_scale: None,
8950 normed_residual_scale: None,
8951 clamp_kqv: None,
8952 attn_temperature: None,
8953 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
8954 block_sub_norms: false,
8955 parallel_residual: false,
8956 learned_positions: false,
8957 attn_value_scale: None,
8958 alibi_max_bias: None,
8959 layer_loops: None,
8960 skip_stream: false,
8961 parallel_ssm: false,
8962 swa_chunked: false,
8963 weightless_qk_norm: false,
8964 logit_multiplier: None,
8965 attention_scale: None,
8966 rope_theta_swa: None,
8967 ffn_activation: frink_models::config::FfnActivation::Swiglu,
8968 best_effort_fields: &["synthetic test config, not a real preset"],
8969 };
8970 let hp = KimiRealHparams {
8971 hidden_dim,
8972 kda_num_heads,
8973 kda_head_dim,
8974 mla_num_heads,
8975 mla_q_lora_rank,
8976 mla_kv_lora_rank,
8977 mla_qk_nope_head_dim,
8978 mla_qk_rope_head_dim,
8979 mla_v_head_dim,
8980 dense_intermediate_dim: dense_intermediate,
8981 moe_hidden_dim: hidden_dim,
8982 moe_intermediate_dim: 4,
8983 n_experts: 1,
8984 num_shared_experts: 0,
8985 };
8986
8987 let prefix = "language_model.model.layers.0";
8991 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8992 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8993 name: String,
8994 shape: Vec<usize>,
8995 n: usize| {
8996 tensors.push((name, shape, vec![0.01f32; n]));
8997 };
8998 push(
8999 &mut tensors,
9000 format!("{prefix}.input_layernorm.weight"),
9001 vec![hidden_dim],
9002 hidden_dim,
9003 );
9004 push(
9005 &mut tensors,
9006 format!("{prefix}.post_attention_layernorm.weight"),
9007 vec![hidden_dim],
9008 hidden_dim,
9009 );
9010 push(
9011 &mut tensors,
9012 format!("{prefix}.self_attention_res_norm.weight"),
9013 vec![hidden_dim],
9014 hidden_dim,
9015 );
9016 push(
9017 &mut tensors,
9018 format!("{prefix}.self_attention_res_proj.weight"),
9019 vec![1, hidden_dim],
9020 hidden_dim,
9021 );
9022 push(
9023 &mut tensors,
9024 format!("{prefix}.mlp_res_norm.weight"),
9025 vec![hidden_dim],
9026 hidden_dim,
9027 );
9028 push(
9029 &mut tensors,
9030 format!("{prefix}.mlp_res_proj.weight"),
9031 vec![1, hidden_dim],
9032 hidden_dim,
9033 );
9034 push(
9035 &mut tensors,
9036 format!("{prefix}.self_attn.q_proj.weight"),
9037 vec![kda_proj, hidden_dim],
9038 kda_proj * hidden_dim,
9039 );
9040 push(
9041 &mut tensors,
9042 format!("{prefix}.self_attn.k_proj.weight"),
9043 vec![kda_proj, hidden_dim],
9044 kda_proj * hidden_dim,
9045 );
9046 push(
9047 &mut tensors,
9048 format!("{prefix}.self_attn.v_proj.weight"),
9049 vec![kda_proj, hidden_dim],
9050 kda_proj * hidden_dim,
9051 );
9052 push(
9053 &mut tensors,
9054 format!("{prefix}.self_attn.q_conv1d.weight"),
9055 vec![kda_proj, 1, conv_kernel],
9056 kda_proj * conv_kernel,
9057 );
9058 push(
9059 &mut tensors,
9060 format!("{prefix}.self_attn.k_conv1d.weight"),
9061 vec![kda_proj, 1, conv_kernel],
9062 kda_proj * conv_kernel,
9063 );
9064 push(
9065 &mut tensors,
9066 format!("{prefix}.self_attn.v_conv1d.weight"),
9067 vec![kda_proj, 1, conv_kernel],
9068 kda_proj * conv_kernel,
9069 );
9070 push(
9071 &mut tensors,
9072 format!("{prefix}.self_attn.A_log"),
9073 vec![kda_num_heads],
9074 kda_num_heads,
9075 );
9076 push(
9077 &mut tensors,
9078 format!("{prefix}.self_attn.f_a_proj.weight"),
9079 vec![kda_head_dim, hidden_dim],
9080 kda_head_dim * hidden_dim,
9081 );
9082 push(
9083 &mut tensors,
9084 format!("{prefix}.self_attn.f_b_proj.weight"),
9085 vec![kda_proj, kda_head_dim],
9086 kda_proj * kda_head_dim,
9087 );
9088 push(
9089 &mut tensors,
9090 format!("{prefix}.self_attn.dt_bias"),
9091 vec![kda_proj],
9092 kda_proj,
9093 );
9094 push(
9095 &mut tensors,
9096 format!("{prefix}.self_attn.b_proj.weight"),
9097 vec![kda_num_heads, hidden_dim],
9098 kda_num_heads * hidden_dim,
9099 );
9100 push(
9101 &mut tensors,
9102 format!("{prefix}.self_attn.g_proj.weight"),
9103 vec![kda_proj, hidden_dim],
9104 kda_proj * hidden_dim,
9105 );
9106 push(
9107 &mut tensors,
9108 format!("{prefix}.self_attn.o_norm.weight"),
9109 vec![kda_head_dim],
9110 kda_head_dim,
9111 );
9112 push(
9113 &mut tensors,
9114 format!("{prefix}.self_attn.o_proj.weight"),
9115 vec![hidden_dim, kda_proj],
9116 hidden_dim * kda_proj,
9117 );
9118 push(
9119 &mut tensors,
9120 format!("{prefix}.mlp.gate_proj.weight"),
9121 vec![dense_intermediate, hidden_dim],
9122 dense_intermediate * hidden_dim,
9123 );
9124 push(
9125 &mut tensors,
9126 format!("{prefix}.mlp.up_proj.weight"),
9127 vec![dense_intermediate, hidden_dim],
9128 dense_intermediate * hidden_dim,
9129 );
9130 push(
9131 &mut tensors,
9132 format!("{prefix}.mlp.down_proj.weight"),
9133 vec![hidden_dim, dense_intermediate],
9134 hidden_dim * dense_intermediate,
9135 );
9136 push(
9137 &mut tensors,
9138 "language_model.model.embed_tokens.weight".to_string(),
9139 vec![vocab_size, hidden_dim],
9140 vocab_size * hidden_dim,
9141 );
9142 push(
9143 &mut tensors,
9144 "language_model.lm_head.weight".to_string(),
9145 vec![vocab_size, hidden_dim],
9146 vocab_size * hidden_dim,
9147 );
9148 push(
9149 &mut tensors,
9150 "language_model.model.norm.weight".to_string(),
9151 vec![hidden_dim],
9152 hidden_dim,
9153 );
9154 push(
9155 &mut tensors,
9156 "language_model.model.output_attn_res_norm.weight".to_string(),
9157 vec![hidden_dim],
9158 hidden_dim,
9159 );
9160 push(
9161 &mut tensors,
9162 "language_model.model.output_attn_res_proj.weight".to_string(),
9163 vec![1, hidden_dim],
9164 hidden_dim,
9165 );
9166
9167 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9180 let dir = std::env::temp_dir().join(format!(
9181 "frink_server_kimi_e2e_test_{}_{}",
9182 std::process::id(),
9183 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9184 ));
9185 std::fs::create_dir_all(&dir).unwrap();
9186 let shard_bytes = write_safetensors_shard(&tensors);
9187 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9188 let map_entries: Vec<String> = tensors
9189 .iter()
9190 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9191 .collect();
9192 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9193 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9194
9195 use base64::Engine;
9199 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9200 .map(|b| {
9201 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9202 format!("{b64} {b}")
9203 })
9204 .collect();
9205 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9206
9207 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9208 .expect("must load the synthetic Kimi checkpoint end to end");
9209 std::fs::remove_dir_all(&dir).ok();
9210 loaded
9211 }
9212
9213 #[test]
9221 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9222 let loaded = build_synthetic_kimi_loaded();
9223 let state = build_app_state(
9224 StartupModels {
9225 loaded: model::LoadedModel::Kimi(loaded),
9226 embedding: None,
9227 },
9228 None,
9229 None,
9230 None,
9231 false,
9232 None,
9233 Arc::new(health::Detection::ready(health::probe_backends())),
9234 );
9235 let active = state.active().expect("a freshly built state has a model");
9236 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9237 assert!(!active.is_synthetic());
9238
9239 let (choices, _usage) = run_generation(
9240 active.generative().unwrap(),
9241 "hi",
9242 &greedy_params(5),
9243 None,
9244 None,
9245 None,
9246 None,
9247 None,
9248 None,
9249 )
9250 .expect("a real Kimi checkpoint must generate without error");
9251 assert!(matches!(
9252 choices[0].finish,
9253 FinishReason::Length | FinishReason::Stop
9254 ));
9255 }
9256
9257 #[test]
9270 fn a_grammar_constrains_the_engine_decode_path() {
9271 let loaded = build_synthetic_kimi_loaded();
9272 let state = build_app_state(
9273 StartupModels {
9274 loaded: model::LoadedModel::Kimi(loaded),
9275 embedding: None,
9276 },
9277 None,
9278 None,
9279 None,
9280 false,
9281 None,
9282 Arc::new(health::Detection::ready(health::probe_backends())),
9283 );
9284 let active = state.active().expect("a freshly built state has a model");
9285
9286 let run = |grammar: Option<&str>| {
9287 let mut params = greedy_params(6);
9288 params.grammar = grammar.map(|src| {
9289 Arc::new(
9290 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9291 .expect("test grammar parses"),
9292 )
9293 });
9294 run_generation(
9295 active.generative().unwrap(),
9296 "hi",
9297 ¶ms,
9298 None,
9299 None,
9300 None,
9301 None,
9302 None,
9303 None,
9304 )
9305 };
9306
9307 let (choices, _) = run(None).expect("the unconstrained run must serve");
9308 let unconstrained = choices[0].text.clone();
9309 assert!(
9310 unconstrained.chars().any(|c| c != 'a'),
9311 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9312 constrained run below would prove nothing"
9313 );
9314
9315 let (choices, _) =
9316 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9317 let one = choices.into_iter().next().unwrap();
9318 let (finish, constrained) = (one.finish, one.text);
9319 assert!(
9320 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9321 "the engine decode path served text its grammar forbids ({constrained:?}): \
9322 the constraint was dropped between `generate_engine` and the sampler"
9323 );
9324 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9325 }
9326
9327 #[test]
9337 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9338 let loaded = build_synthetic_kimi_loaded();
9339 let state = build_app_state(
9340 StartupModels {
9341 loaded: model::LoadedModel::Kimi(loaded),
9342 embedding: None,
9343 },
9344 None,
9345 None,
9346 None,
9347 false,
9348 None,
9349 Arc::new(health::Detection::ready(health::probe_backends())),
9350 );
9351
9352 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9353 let kv_pool_config = generate::KvPoolConfig {
9354 pool: pool.clone(),
9355 queue_wait: Duration::ZERO,
9356 };
9357 let pc = Mutex::new(PrefixCache::new(4));
9358
9359 run_generation(
9360 state
9361 .active()
9362 .expect("a freshly built state has a model")
9363 .generative()
9364 .unwrap(),
9365 "hi",
9366 &greedy_params(5),
9367 Some(&kv_pool_config),
9368 None,
9369 Some(&pc),
9370 None,
9371 None,
9372 None,
9373 )
9374 .expect("a real Kimi checkpoint must generate without error");
9375
9376 assert_eq!(
9377 pool.lock().unwrap().free_blocks(),
9378 4,
9379 "the KV pool must be completely untouched by a Kimi request"
9380 );
9381 let stats = pc.lock().unwrap().stats();
9382 assert_eq!(
9383 stats.hits + stats.misses,
9384 0,
9385 "the prefix cache must never be consulted for a Kimi request"
9386 );
9387 }
9388}