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