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