1mod admin;
36mod anthropic;
37mod attribution;
38mod best_of;
39mod budget;
40mod cache_admin;
41mod cache_salt;
42mod cancel;
43mod chat_params;
44mod chat_stream_choice;
45mod chat_template;
46mod choice_stream;
47mod cli;
48mod completion;
49mod continuation;
50mod conversations;
51mod decode_task;
52mod embeddings;
53mod generate;
54mod grammar_request;
55mod health;
56mod journal;
57mod json_mode;
58mod limits;
59mod loaded;
60mod logprobs;
61mod lora;
62mod mcp;
63mod model;
64mod openai_extra;
65mod output;
66mod policy;
67mod prefill_batch;
68mod reasoning_budget;
69mod reasoning_tokens;
70mod request_tail;
71mod rerank;
72mod response_cache;
73pub(crate) mod responses;
74mod resume;
75mod round_robin;
76mod sample_step;
77mod sampling_knobs;
78mod sampling_loop;
79mod security;
80mod serving;
81mod session;
82mod slots;
83mod sse;
84mod stats;
85mod stop;
86mod stream_events;
87mod tasks;
88mod token_mask;
89mod tool_grammar;
90mod unimplemented_fields;
91mod unsupported_sampling;
92mod utf8_stream;
93
94use std::cell::RefCell;
95use std::convert::Infallible;
96use std::net::SocketAddr;
97use std::path::PathBuf;
98use std::rc::Rc;
99use std::sync::{Arc, Mutex, MutexGuard};
100use std::time::Duration;
101
102use axum::{
103 extract::State,
104 http::StatusCode,
105 response::sse::{Event, Sse},
106 response::{IntoResponse, Response},
107 routing::{get, post},
108 Json, Router,
109};
110use serde::{Deserialize, Serialize};
111
112use cli::apply_cli_overrides;
113pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
114
115use frink_core::cache::KvBlockPool;
116use frink_models::kimi_tokenizer::KimiTokenizer;
117use frink_models::sampling::SamplingParams;
118use frink_models::tokenizer::{SpecialTokens, StopTokens};
119use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
120#[cfg(test)]
121use generate::FinishReason;
122use generate::GenerationParams;
123pub(crate) use loaded::{ActiveModel, Loaded, SleptModel};
124use model::ServerTokenizer;
125use rerank::encoder_endpoints;
126use response_cache::ResponseCache;
127use sampling_knobs::SamplingKnobs;
128
129#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
138 Gguf(GgufModel),
139 Kimi(KimiModel),
140 Mla(MlaModel),
141 Gemma4(Gemma4Model),
142 Glm52(Glm52Model),
143}
144
145pub(crate) struct GgufModel {
146 decoder: Arc<Decoder>,
147 tokenizer: Arc<ServerTokenizer>,
148 stop_tokens: StopTokens,
149 bos_id: Option<usize>,
150 is_synthetic: bool,
151 chat_template: chat_template::PromptTemplate,
152}
153
154pub(crate) struct KimiModel {
155 engine: KimiEngine,
156 tokenizer: KimiTokenizer,
157 stop_tokens: StopTokens,
158 chat_template: chat_template::PromptTemplate,
159}
160
161pub(crate) struct MlaModel {
162 engine: MlaEngine,
163 tokenizer: ServerTokenizer,
164 stop_tokens: StopTokens,
165 bos_id: Option<usize>,
166 name: String,
167 chat_template: chat_template::PromptTemplate,
168}
169
170pub(crate) struct Gemma4Model {
171 engine: Gemma4Engine,
172 tokenizer: ServerTokenizer,
173 stop_tokens: StopTokens,
174 bos_id: Option<usize>,
175 name: String,
176 chat_template: chat_template::PromptTemplate,
177}
178
179pub(crate) struct Glm52Model {
180 engine: frink_models::Glm52Engine,
181 tokenizer: ServerTokenizer,
182 stop_tokens: StopTokens,
183 bos_id: Option<usize>,
184 name: String,
185 chat_template: chat_template::PromptTemplate,
186}
187
188impl Model {
189 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
190 match self {
191 Model::Gguf(m) => m.chat_template.clone(),
192 Model::Kimi(m) => m.chat_template.clone(),
193 Model::Mla(m) => m.chat_template.clone(),
194 Model::Gemma4(m) => m.chat_template.clone(),
195 Model::Glm52(m) => m.chat_template.clone(),
196 }
197 }
198
199 fn is_synthetic(&self) -> bool {
204 match self {
205 Model::Gguf(m) => m.is_synthetic,
206 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
207 }
208 }
209
210 fn tokenizer_kind(&self) -> &'static str {
211 match self {
212 Model::Gguf(m) => m.tokenizer.kind(),
213 Model::Kimi(_) => "kimi-tiktoken-bpe",
214 Model::Mla(m) => m.tokenizer.kind(),
215 Model::Gemma4(m) => m.tokenizer.kind(),
216 Model::Glm52(m) => m.tokenizer.kind(),
217 }
218 }
219
220 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
224 match self {
225 Model::Gguf(m) => m.decoder.expert_store_stats(),
226 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
227 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
228 }
229 }
230
231 pub(crate) fn name(&self) -> &str {
232 match self {
233 Model::Gguf(m) => m.decoder.config.name,
234 Model::Kimi(_) => "kimi-k3",
235 Model::Mla(m) => m.name.as_str(),
236 Model::Gemma4(m) => m.name.as_str(),
237 Model::Glm52(m) => m.name.as_str(),
238 }
239 }
240
241 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
268 match self {
269 Model::Gguf(m) => m.tokenizer.encode(text, specials),
270 Model::Kimi(m) => m
271 .tokenizer
272 .encode(text, specials)
273 .into_iter()
274 .map(|id| id as usize)
275 .collect(),
276 Model::Mla(m) => m.tokenizer.encode(text, specials),
277 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
278 Model::Glm52(m) => m.tokenizer.encode(text, specials),
279 }
280 }
281
282 pub(crate) fn bos_id(&self) -> Option<usize> {
291 match self {
292 Model::Gguf(m) => m.bos_id,
293 Model::Kimi(_) => None,
294 Model::Mla(m) => m.bos_id,
295 Model::Gemma4(m) => m.bos_id,
296 Model::Glm52(m) => m.bos_id,
297 }
298 }
299
300 pub(crate) fn decode(&self, ids: &[usize]) -> String {
301 match self {
302 Model::Gguf(m) => m.tokenizer.decode(ids),
303 Model::Kimi(m) => {
304 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
305 m.tokenizer.decode(&ids32)
306 }
307 Model::Mla(m) => m.tokenizer.decode(ids),
308 Model::Gemma4(m) => m.tokenizer.decode(ids),
309 Model::Glm52(m) => m.tokenizer.decode(ids),
310 }
311 }
312
313 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
316 match self {
317 Model::Gguf(m) => {
318 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
319 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
320 }
321 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
322 }
323 }
324
325 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
331 match self {
332 Model::Gguf(m) => Some(&m.decoder),
333 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
334 }
335 }
336
337 pub(crate) fn vocab_size(&self) -> Option<usize> {
338 match self {
339 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
340 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
341 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
342 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
343 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
344 }
345 }
346
347 fn has_real_vocabulary(&self) -> bool {
354 match self {
355 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
356 Model::Kimi(_) => true,
357 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
358 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
359 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
360 }
361 }
362}
363
364impl frink_models::dry::DryVocab for Model {
370 fn n_tokens(&self) -> usize {
371 self.vocab_size().unwrap_or(0)
372 }
373
374 fn detokenize(&self, token: usize) -> String {
375 self.decode(&[token])
376 }
377
378 fn tokenize(&self, text: &str) -> Vec<usize> {
379 self.encode(text, SpecialTokens::AsText)
380 }
381}
382
383pub(crate) struct AppState {
384 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
398 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
417 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
421 pub(crate) slept: Mutex<Option<SleptModel>>,
431 pub(crate) tasks: Arc<tasks::TaskRegistry>,
433 pub(crate) cancels: Arc<cancel::CancelRegistry>,
437 pub(crate) stats: stats::Stats,
440 pub(crate) streams: resume::StreamRegistry,
443 pub(crate) model_dir: Option<PathBuf>,
445 response_cache: Mutex<ResponseCache>,
449 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
461 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
474 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
484 sessions: session::SessionStore,
490 requests_total: std::sync::atomic::AtomicU64,
491 request_errors_total: std::sync::atomic::AtomicU64,
492 started_at: std::time::Instant,
493 last_request_ms: std::sync::atomic::AtomicU64,
498 detection: Arc<health::Detection>,
500 mcp: Option<mcp::LoadedMcpConfig>,
502 pub(crate) continuous_batching_enabled: bool,
506 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
510 loading_model: Mutex<Option<String>>,
514 last_load_error: Mutex<Option<(String, String)>>,
518 pub(crate) serving: Mutex<crate::stats::ServingStats>,
523 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
528 pub(crate) footprint:
534 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
535 pub(crate) started_unix: u64,
545}
546
547pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
553
554impl AppState {
555 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
559 self.active
560 .read()
561 .unwrap_or_else(|p| p.into_inner())
562 .clone()
563 }
564
565 pub(crate) fn is_sleeping(&self) -> bool {
572 self.slept
573 .lock()
574 .unwrap_or_else(|p| p.into_inner())
575 .is_some()
576 }
577
578 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
579 if let Some(active) = self.active() {
580 return Ok(active);
581 }
582 if self.is_sleeping() {
586 return Err((
587 StatusCode::SERVICE_UNAVAILABLE,
588 Json(serde_json::json!({"error": {
589 "message": "this server is asleep; POST /wake_up to reload the model it put \
590 away",
591 "type": "server_sleeping"
592 }})),
593 ));
594 }
595 Err((
596 StatusCode::SERVICE_UNAVAILABLE,
597 Json(serde_json::json!({"error": {
598 "message": "no model is loaded; POST /admin/models/load with an id from \
599 GET /admin/models",
600 "type": "model_not_loaded"
601 }})),
602 ))
603 }
604
605 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
613 Ok(Arc::clone(self.require_active()?.generative()?))
614 }
615
616 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
624 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
625 std::mem::replace(&mut *guard, next)
626 }
627
628 fn mark_request_finished(&self) {
632 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
633 self.last_request_ms
634 .store(ms, std::sync::atomic::Ordering::Relaxed);
635 }
636
637 pub(crate) fn uptime(&self) -> Duration {
638 self.started_at.elapsed()
639 }
640
641 pub(crate) fn requests_total(&self) -> u64 {
642 self.requests_total
643 .load(std::sync::atomic::Ordering::Relaxed)
644 }
645
646 pub(crate) fn errors_total(&self) -> u64 {
647 self.request_errors_total
648 .load(std::sync::atomic::Ordering::Relaxed)
649 }
650
651 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
652 lock_cache(&self.response_cache).stats()
653 }
654
655 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
658 let last = self
659 .last_request_ms
660 .load(std::sync::atomic::Ordering::Relaxed);
661 (last > 0)
662 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
663 .map(|age| age.max(0.0))
664 }
665
666 pub(crate) fn loading_model_id(&self) -> Option<String> {
667 self.loading_model
668 .lock()
669 .unwrap_or_else(|p| p.into_inner())
670 .clone()
671 }
672
673 pub(crate) fn set_loading_model(&self, id: Option<String>) {
674 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
675 }
676
677 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
678 self.last_load_error
679 .lock()
680 .unwrap_or_else(|p| p.into_inner())
681 .clone()
682 }
683
684 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
685 *self
686 .last_load_error
687 .lock()
688 .unwrap_or_else(|p| p.into_inner()) = error;
689 }
690
691 pub(crate) fn active_model_name(&self) -> Option<String> {
700 self.active().map(|a| a.name().to_string())
701 }
702
703 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
713 self.embedding
714 .clone()
715 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
716 }
717
718 pub(crate) fn embedding_model_name(&self) -> Option<String> {
722 match self.embedding_model() {
723 Some(e) => Some(e.name().to_string()),
724 None => self.active_model_name(),
725 }
726 }
727
728 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
729 self.stats.record(stats::entry(record));
730 }
731}
732
733fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
738 cache
739 .lock()
740 .unwrap_or_else(|poisoned| poisoned.into_inner())
741}
742
743#[derive(Debug, Clone, Deserialize)]
744#[serde(untagged)]
745pub(crate) enum MessageContent {
746 Text(String),
747 Parts(Vec<ContentPart>),
748}
749
750#[derive(Debug, Clone, Deserialize)]
751struct ContentPart {
752 #[serde(rename = "type")]
753 kind: String,
754 #[serde(default)]
755 text: Option<String>,
756 #[serde(default)]
757 image_url: Option<serde_json::Value>,
758}
759
760impl MessageContent {
761 fn as_text(&self) -> String {
762 match self {
763 Self::Text(s) => s.clone(),
764 Self::Parts(parts) => parts
765 .iter()
766 .filter_map(|p| p.text.as_deref())
767 .collect::<Vec<_>>()
768 .join(""),
769 }
770 }
771
772 fn has_image(&self) -> bool {
773 match self {
774 Self::Text(_) => false,
775 Self::Parts(parts) => parts
776 .iter()
777 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
778 }
779 }
780}
781
782#[derive(Debug, Clone, Deserialize)]
783pub(crate) struct ChatMessage {
784 pub(crate) role: String,
785 #[serde(default)]
789 pub(crate) content: Option<MessageContent>,
790 #[serde(default)]
794 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
795 #[serde(default)]
800 #[allow(dead_code)]
801 pub(crate) tool_call_id: Option<String>,
802 #[serde(default, alias = "reasoning")]
818 pub(crate) reasoning_content: Option<String>,
819}
820
821impl ChatMessage {
822 fn rendered_content(&self) -> String {
829 let mut out = self
830 .content
831 .as_ref()
832 .map(MessageContent::as_text)
833 .unwrap_or_default();
834 if let Some(calls) = &self.tool_calls {
835 for call in calls {
836 out.push_str(&format!(
837 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
838 call.function.name, call.function.arguments
839 ));
840 }
841 }
842 out
843 }
844}
845
846#[derive(Debug, Clone, Deserialize)]
847pub(crate) struct ToolCallIn {
848 #[serde(default)]
849 #[allow(dead_code)]
850 id: String,
851 #[serde(rename = "type", default)]
852 #[allow(dead_code)]
853 kind: String,
854 function: ToolCallFunctionIn,
855}
856
857#[derive(Debug, Clone, Deserialize)]
858struct ToolCallFunctionIn {
859 name: String,
860 arguments: String,
865}
866
867#[derive(Debug, Clone, Deserialize)]
870struct ToolDef {
871 #[serde(rename = "type", default)]
872 #[allow(dead_code)]
873 kind: String,
874 function: ToolFunctionDef,
875}
876
877#[derive(Debug, Clone, Deserialize)]
878struct ToolFunctionDef {
879 name: String,
880 #[serde(default)]
881 description: Option<String>,
882 #[serde(default)]
883 parameters: Option<serde_json::Value>,
884}
885
886#[derive(Debug, Clone, Deserialize)]
897#[serde(untagged)]
898enum ToolChoice {
899 Mode(String),
900 Specific(serde_json::Value),
901}
902
903#[derive(Deserialize)]
906#[serde(untagged)]
907enum StopParam {
908 One(String),
909 Many(Vec<String>),
910}
911
912#[derive(Deserialize)]
913struct ChatCompletionRequest {
914 model: String,
915 messages: Vec<ChatMessage>,
916 #[serde(default = "default_max_tokens")]
917 max_tokens: usize,
918 #[serde(default)]
919 temperature: Option<f32>,
920 #[serde(default)]
921 top_p: Option<f32>,
922 #[serde(default)]
927 min_p: Option<f32>,
928 #[serde(default)]
929 top_k: Option<usize>,
930 #[serde(default)]
931 repetition_penalty: Option<f32>,
932 #[serde(flatten)]
936 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
937 #[serde(flatten)]
941 unimplemented: crate::unimplemented_fields::UnimplementedFields,
942 #[serde(default)]
943 seed: Option<u64>,
944 #[serde(default)]
945 stop: Option<StopParam>,
946 #[serde(default)]
947 stream: Option<bool>,
948 #[serde(default)]
961 stream_resumable: Option<bool>,
962 #[serde(default)]
973 ignore_eos: Option<bool>,
974 #[serde(default)]
975 tools: Vec<ToolDef>,
976 #[serde(default)]
977 tool_choice: Option<ToolChoice>,
978 #[serde(default)]
986 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
987 #[serde(default)]
995 reasoning_effort: Option<String>,
996 #[serde(default)]
1000 thinking: Option<ThinkingSwitch>,
1001 #[serde(default)]
1006 session_id: Option<String>,
1007 #[serde(default, deserialize_with = "continuation::deserialize")]
1015 continue_final_message: continuation::ContinueFinalMessage,
1016 #[serde(default, alias = "thinking_budget_tokens")]
1024 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
1025 #[serde(default)]
1027 logprobs: Option<bool>,
1028 #[serde(default)]
1029 top_logprobs: Option<u32>,
1030 #[serde(default)]
1031 presence_penalty: Option<f32>,
1032 #[serde(default)]
1033 frequency_penalty: Option<f32>,
1034 #[serde(default)]
1035 response_format: Option<serde_json::Value>,
1036 #[serde(default)]
1043 logit_bias: Option<serde_json::Value>,
1044 #[serde(default)]
1048 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1049 #[serde(default)]
1059 samplers: Option<serde_json::Value>,
1060 #[serde(default)]
1070 grammar: Option<String>,
1071}
1072
1073const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1085
1086#[derive(Debug, Clone, Deserialize)]
1088pub(crate) struct ThinkingSwitch {
1089 #[serde(rename = "type")]
1090 pub(crate) kind: String,
1091}
1092
1093const THINKING_KWARG_KEYS: [&str; 4] = [
1097 "enable_thinking",
1098 "thinking",
1099 "thinking_mode",
1100 "reasoning_effort",
1101];
1102
1103const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1107
1108fn default_max_tokens() -> usize {
1109 DEFAULT_CHAT_MAX_TOKENS
1110}
1111
1112impl ChatCompletionRequest {
1113 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1122 let mut knobs = SamplingKnobs {
1123 temperature: self.temperature,
1124 top_p: self.top_p,
1125 min_p: self.min_p,
1126 top_k: self.top_k,
1127 repetition_penalty: self.repetition_penalty,
1128 presence_penalty: self.presence_penalty,
1129 frequency_penalty: self.frequency_penalty,
1130 penalty_last_n: None,
1134 sampler_order: unsupported_sampling::parse_sampler_order(
1135 self.samplers.as_ref(),
1136 "/v1/chat/completions",
1137 )?,
1138 ..SamplingKnobs::default()
1139 };
1140 self.extra_samplers.apply(&mut knobs);
1141 Ok(knobs)
1142 }
1143
1144 fn sampling_params(
1145 &self,
1146 model: crate::sampling_knobs::SamplerModel<'_>,
1147 ) -> Result<SamplingParams, ApiError> {
1148 self.sampling_knobs()?.resolve(model).map_err(|e| {
1149 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1150 })
1151 }
1152
1153 fn stop_sequences(&self) -> Vec<String> {
1154 self.stop
1155 .as_ref()
1156 .map(|s| match s {
1157 StopParam::One(v) => vec![v.clone()],
1158 StopParam::Many(v) => v.clone(),
1159 })
1160 .unwrap_or_default()
1161 }
1162
1163 fn n_logprobs(&self) -> Result<Option<usize>, ApiError> {
1177 const MAX: u32 = 20;
1178 match (self.logprobs, self.top_logprobs) {
1179 (Some(true), Some(n)) if n > MAX => Err(invalid_request(
1180 &format!(
1181 "`top_logprobs` is {n}; this server reports at most {MAX} alternatives per \
1182 position, as upstream does"
1183 ),
1184 "top_logprobs",
1185 )),
1186 (Some(true), Some(n)) => Ok(Some(n as usize)),
1187 (Some(true), None) => Ok(Some(0)),
1191 (_, Some(_)) => Err(invalid_request(
1192 "`top_logprobs` requires `logprobs: true`",
1193 "top_logprobs",
1194 )),
1195 _ => Ok(None),
1196 }
1197 }
1198
1199 fn tools_active(&self) -> bool {
1200 !self.tools.is_empty()
1201 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1202 }
1203
1204 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1216 match &self.tool_choice {
1217 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1218 Some(ToolChoice::Specific(value)) => {
1219 let name = value
1223 .get("function")
1224 .and_then(|f| f.get("name"))
1225 .or_else(|| value.get("name"))
1226 .and_then(|n| n.as_str());
1227 match name {
1228 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1229 None => Err(invalid_request(
1230 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1231 function.name",
1232 "tool_choice",
1233 )),
1234 }
1235 }
1236 _ => Ok(None),
1237 }
1238 }
1239
1240 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1242 self.tools
1243 .iter()
1244 .map(|t| tool_grammar::ToolSpec {
1245 name: &t.function.name,
1246 parameters: t.function.parameters.as_ref(),
1247 })
1248 .collect()
1249 }
1250
1251 fn thinking_direction(&self) -> Option<bool> {
1294 if let Some(switch) = &self.thinking {
1295 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1296 "disabled" => Some(false),
1297 "enabled" => Some(true),
1298 _ => None,
1301 };
1302 }
1303 let effort = self.reasoning_effort.as_ref()?;
1304 DISABLE_EFFORTS
1305 .contains(&effort.trim().to_ascii_lowercase().as_str())
1306 .then_some(false)
1307 }
1308
1309 fn resolve_template_kwargs(
1310 &self,
1311 template: &chat_template::PromptTemplate,
1312 ) -> serde_json::Map<String, serde_json::Value> {
1313 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1314 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1318
1319 if !caller_steered {
1320 match self.thinking_direction() {
1321 Some(false) => {
1322 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1323 kwargs.insert(k, v);
1324 }
1325 return kwargs;
1328 }
1329 Some(true) => {
1330 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1331 kwargs.insert(k, v);
1332 }
1333 }
1334 None => {}
1335 }
1336 if let Some(effort) = &self.reasoning_effort {
1337 kwargs
1338 .entry("reasoning_effort".to_string())
1339 .or_insert_with(|| serde_json::json!(effort));
1340 }
1341 }
1342
1343 let offered: Vec<serde_json::Value> = if self.tools_active() {
1344 self.tools.iter().map(chat_template::tool_json).collect()
1345 } else {
1346 Vec::new()
1347 };
1348 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1349 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1350 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1351 kwargs.entry(k).or_insert(v);
1352 }
1353 }
1354 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1355 crate::policy::effort::EffortMapping::Mapped(to) => {
1356 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1357 }
1358 crate::policy::effort::EffortMapping::Dropped => {
1359 tracing::debug!(
1360 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1361 enough, so its own default applies"
1362 );
1363 }
1364 crate::policy::effort::EffortMapping::Unchanged => {}
1365 }
1366 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1367 kwargs
1368 }
1369
1370 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1373 if self.max_tokens == 0 {
1380 return Err(invalid_request(
1381 "max_tokens must be at least 1",
1382 "max_tokens",
1383 ));
1384 }
1385 if let Some(switch) = &self.thinking {
1390 let kind = switch.kind.trim().to_ascii_lowercase();
1391 if kind != "enabled" && kind != "disabled" {
1392 return Err(invalid_request(
1393 "thinking.type must be \"enabled\" or \"disabled\"",
1394 "thinking.type",
1395 ));
1396 }
1397 }
1398 for msg in &self.messages {
1399 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1400 return Err(unsupported_feature(
1401 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1402 ));
1403 }
1404 }
1405 self.n_logprobs()?;
1409 self.unimplemented.refuse("/v1/chat/completions")?;
1413 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1414 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1418 let stated_grammar =
1433 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1434 if let Some(forced) = self.forced_tool_choice()? {
1442 if self.tools.is_empty() {
1443 return Err(invalid_request(
1444 "tool_choice forces a tool call, but no tools were offered",
1445 "tool_choice",
1446 ));
1447 }
1448 if let tool_grammar::Forced::Named(name) = forced {
1449 if !self.tools.iter().any(|t| t.function.name == name) {
1450 return Err(invalid_request(
1451 &format!(
1452 "tool_choice names {name:?}, which is not one of the tools offered"
1453 ),
1454 "tool_choice",
1455 ));
1456 }
1457 }
1458 if stated_grammar.is_some() {
1468 return Err(invalid_request(
1469 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1470 are two different constraints on the same generation; send one",
1471 "tool_choice",
1472 ));
1473 }
1474 if self.json_object_mode() {
1475 return Err(invalid_request(
1476 "a forced tool_choice cannot be combined with response_format json_object: \
1477 the tool-call markers are not JSON",
1478 "tool_choice",
1479 ));
1480 }
1481 }
1482 Ok(())
1483 }
1484
1485 fn effective_stop_sequences(&self) -> Vec<String> {
1492 let mut stop = self.stop_sequences();
1493 if self.tools_active() {
1494 stop.push("</tool_call>".to_string());
1495 }
1496 stop
1497 }
1498
1499 fn json_object_mode(&self) -> bool {
1500 self.response_format
1501 .as_ref()
1502 .and_then(|v| v.get("type"))
1503 .and_then(|v| v.as_str())
1504 == Some("json_object")
1505 }
1506}
1507
1508#[derive(Serialize)]
1509struct ChatCompletionChoice {
1510 index: usize,
1511 message: ChatCompletionResponseMessage,
1512 finish_reason: &'static str,
1513 #[serde(skip_serializing_if = "Option::is_none")]
1518 logprobs: Option<serde_json::Value>,
1519}
1520
1521#[derive(Serialize)]
1522struct ChatCompletionResponseMessage {
1523 role: &'static str,
1524 #[serde(skip_serializing_if = "Option::is_none")]
1525 content: Option<String>,
1526 #[serde(skip_serializing_if = "Option::is_none")]
1530 reasoning_content: Option<String>,
1531 #[serde(skip_serializing_if = "Option::is_none")]
1532 tool_calls: Option<Vec<ToolCallOut>>,
1533}
1534
1535#[derive(Serialize, Clone)]
1536struct ToolCallOut {
1537 id: String,
1538 #[serde(rename = "type")]
1539 kind: &'static str,
1540 function: ToolCallFunctionOut,
1541}
1542
1543#[derive(Serialize, Clone)]
1551struct ToolCallDelta {
1552 index: usize,
1553 #[serde(skip_serializing_if = "Option::is_none")]
1554 id: Option<String>,
1555 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1556 kind: Option<&'static str>,
1557 function: ToolCallFunctionDelta,
1558}
1559
1560#[derive(Serialize, Clone, Default)]
1561struct ToolCallFunctionDelta {
1562 #[serde(skip_serializing_if = "Option::is_none")]
1563 name: Option<String>,
1564 #[serde(skip_serializing_if = "Option::is_none")]
1567 arguments: Option<String>,
1568}
1569
1570impl ToolCallDelta {
1571 fn whole(index: usize, name: String, arguments: String) -> Self {
1573 ToolCallDelta {
1574 index,
1575 id: Some(format!("call_{index}")),
1576 kind: Some("function"),
1577 function: ToolCallFunctionDelta {
1578 name: Some(name),
1579 arguments: Some(arguments),
1580 },
1581 }
1582 }
1583
1584 fn opening(index: usize, name: String) -> Self {
1586 ToolCallDelta {
1587 index,
1588 id: Some(format!("call_{index}")),
1589 kind: Some("function"),
1590 function: ToolCallFunctionDelta {
1591 name: Some(name),
1592 arguments: Some(String::new()),
1593 },
1594 }
1595 }
1596
1597 fn arguments(index: usize, fragment: String) -> Self {
1599 ToolCallDelta {
1600 index,
1601 id: None,
1602 kind: None,
1603 function: ToolCallFunctionDelta {
1604 name: None,
1605 arguments: Some(fragment),
1606 },
1607 }
1608 }
1609}
1610
1611#[derive(Serialize, Clone)]
1612struct ToolCallFunctionOut {
1613 name: String,
1614 arguments: String,
1618}
1619
1620#[derive(Serialize)]
1621struct ChatCompletionResponse {
1622 id: String,
1623 request_id: String,
1629 object: &'static str,
1630 model: String,
1631 choices: Vec<ChatCompletionChoice>,
1632 usage: generate::Usage,
1637 frink_cache: &'static str,
1648}
1649
1650#[derive(Serialize)]
1651struct ChatCompletionChunkDelta {
1652 #[serde(skip_serializing_if = "Option::is_none")]
1653 role: Option<&'static str>,
1654 #[serde(skip_serializing_if = "Option::is_none")]
1655 content: Option<String>,
1656 #[serde(skip_serializing_if = "Option::is_none")]
1658 reasoning_content: Option<String>,
1659 #[serde(skip_serializing_if = "Option::is_none")]
1660 tool_calls: Option<Vec<ToolCallDelta>>,
1661}
1662
1663#[derive(Serialize)]
1664struct ChatCompletionChunkChoice {
1665 index: usize,
1666 delta: ChatCompletionChunkDelta,
1667 finish_reason: Option<&'static str>,
1668}
1669
1670#[derive(Serialize)]
1671struct ChatCompletionChunk {
1672 id: String,
1673 #[serde(skip_serializing_if = "Option::is_none")]
1680 request_id: Option<String>,
1681 object: &'static str,
1682 model: String,
1683 choices: Vec<ChatCompletionChunkChoice>,
1684 #[serde(skip_serializing_if = "Option::is_none")]
1687 usage: Option<generate::Usage>,
1688}
1689
1690async fn health(State(state): State<Arc<AppState>>) -> Response {
1696 let snapshot = state.detection.snapshot();
1697 let mut capabilities = snapshot.capabilities;
1698 let active = state.active();
1699
1700 capabilities.push(match active.as_deref() {
1703 None => frink_api::Capability::unavailable(
1708 frink_api::health::capability::REAL_WEIGHTS,
1709 frink_api::health::reason::MODEL_NOT_LOADED,
1710 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1711 ),
1712 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1713 frink_api::health::capability::REAL_WEIGHTS,
1714 frink_api::health::reason::MODEL_NOT_LOADED,
1715 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1716 checkpoint. Output from this model is noise.",
1717 ),
1718 Some(active) if active.encoder().is_some() => {
1728 let endpoints = active
1729 .encoder()
1730 .map(|e| encoder_endpoints(e))
1731 .unwrap_or_default();
1732 let served_by = match endpoints.is_empty() {
1733 true => "no endpoint in this build serves it".to_string(),
1734 false => format!("served by {}", endpoints.join(" and ")),
1735 };
1736 frink_api::Capability::available(
1737 frink_api::health::capability::REAL_WEIGHTS,
1738 format!(
1739 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1740 {served_by}; generation endpoints refuse it.",
1741 active.name(),
1742 ),
1743 )
1744 }
1745 Some(active) => frink_api::Capability::available(
1746 frink_api::health::capability::REAL_WEIGHTS,
1747 format!("Serving the real checkpoint '{}'.", active.name()),
1748 ),
1749 });
1750 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1751 frink_api::Capability::available(
1752 frink_api::health::capability::CONTINUOUS_BATCHING,
1753 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1754 "On by default on Metal. Concurrent requests share one batched decode worker."
1755 } else {
1756 "Concurrent requests share one batched decode step."
1757 },
1758 )
1759 } else if state.metal_private_decode_gate.is_some() {
1760 frink_api::Capability::unavailable(
1761 frink_api::health::capability::CONTINUOUS_BATCHING,
1762 frink_api::health::reason::DISABLED,
1763 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1764 )
1765 } else {
1766 frink_api::Capability::unavailable(
1767 frink_api::health::capability::CONTINUOUS_BATCHING,
1768 frink_api::health::reason::DISABLED,
1769 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1770 )
1771 });
1772
1773 let last_request_ms = state
1774 .last_request_ms
1775 .load(std::sync::atomic::Ordering::Relaxed);
1776 let uptime = state.started_at.elapsed();
1777 let health_state = if active.is_none() {
1785 frink_api::HealthState::Unavailable
1786 } else {
1787 snapshot.state
1788 };
1789 let body = frink_api::HealthResponse {
1790 state: health_state,
1791 reason: match health_state {
1792 frink_api::HealthState::Ready => None,
1793 frink_api::HealthState::Unavailable => {
1794 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1795 }
1796 frink_api::HealthState::Detecting => {
1797 Some(frink_api::health::reason::DETECTING.to_string())
1798 }
1799 },
1800 detail: match health_state {
1801 frink_api::HealthState::Ready => None,
1802 frink_api::HealthState::Unavailable => Some(
1803 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1804 .to_string(),
1805 ),
1806 frink_api::HealthState::Detecting => {
1807 Some("Probing available compute backends.".to_string())
1808 }
1809 },
1810 model: active
1811 .as_deref()
1812 .map(|active| frink_api::health::ModelSummary {
1813 id: active.name().to_string(),
1814 tokenizer: active.tokenizer_kind().to_string(),
1815 synthetic_weights: active.is_synthetic(),
1816 }),
1817 capabilities,
1818 version: env!("CARGO_PKG_VERSION").to_string(),
1819 pid: std::process::id(),
1820 uptime_seconds: uptime.as_secs_f64(),
1821 server_time_unix_ms: std::time::SystemTime::now()
1822 .duration_since(std::time::UNIX_EPOCH)
1823 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1824 .unwrap_or(0),
1825 last_request_age_seconds: (last_request_ms > 0)
1826 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1827 .map(|age| age.max(0.0)),
1828 };
1829
1830 let status =
1831 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1832 (status, Json(body)).into_response()
1833}
1834
1835async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1836 let Some(active) = state.active() else {
1840 return Json(serde_json::json!({ "object": "list", "data": [] }));
1841 };
1842 let mut model_entry = serde_json::json!({
1843 "id": active.name(),
1844 "object": "model",
1845 "frink_synthetic_weights": active.is_synthetic(),
1846 "frink_tokenizer": active.tokenizer_kind(),
1847 });
1848 if let Some(encoder) = active.encoder() {
1854 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1855 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1856 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1857 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1858 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1859 }
1860 if let Some(model) = active.generative_opt() {
1868 let parser_configured = active.reasoning_format().is_some();
1869 let gears = model.chat_template().think_gears(parser_configured);
1870 if !gears.is_empty() {
1871 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1872 if let Some(default) = &gears.default {
1873 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1874 }
1875 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1878 }
1879 }
1880 if let Some(mcp) = &state.mcp {
1881 model_entry["frink_mcp"] = mcp.models_metadata();
1882 }
1883 Json(serde_json::json!({
1884 "object": "list",
1885 "data": [model_entry]
1886 }))
1887}
1888
1889async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1903 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1904 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1905 let active = state.active();
1906 Json(serde_json::json!({
1907 "model": active.as_ref().map(|a| a.name()),
1908 "state": state
1909 .maintenance
1910 .lock()
1911 .unwrap_or_else(|p| p.into_inner())
1912 .state()
1913 .as_str(),
1914 "uptime_s": state.uptime().as_secs(),
1915 "throughput": {
1916 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1917 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1918 },
1919 "requests": {
1920 "active": state.cancels.live_count(),
1921 "completed": state.stats.recorded_total(),
1922 "p95_ms": state.stats.p95_duration_ms(),
1923 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1924 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1925 "completion_tokens_total": state.stats.tokens_generated_total(),
1926 },
1927 "pools": cache_admin::pool_gauges(&state),
1933 "memory": cache_admin::footprint_json(&state),
1936 }))
1937}
1938
1939#[derive(Deserialize)]
1940struct RequestsQuery {
1941 #[serde(default)]
1942 since: u64,
1943 #[serde(default = "default_requests_limit")]
1944 limit: usize,
1945}
1946
1947fn default_requests_limit() -> usize {
1948 stats::MAX_PAGE
1949}
1950
1951async fn recent_requests(
1959 State(state): State<Arc<AppState>>,
1960 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1961) -> Json<serde_json::Value> {
1962 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1963 Json(serde_json::json!({
1964 "requests": rows,
1965 "next_cursor": cursor,
1966 "missed": missed,
1967 "total": state.stats.recorded_total(),
1968 }))
1969}
1970
1971#[derive(Serialize)]
1972struct CombinedCacheStats {
1973 response_cache: response_cache::CacheStats,
1974 prefix_cache: Option<frink_models::PrefixCacheStats>,
1976}
1977
1978async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1979 Json(CombinedCacheStats {
1980 response_cache: lock_cache(&state.response_cache).stats(),
1981 prefix_cache: state
1982 .prefix_cache
1983 .as_ref()
1984 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1985 })
1986}
1987
1988async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1993 use std::sync::atomic::Ordering;
1994
1995 let cache_stats = lock_cache(&state.response_cache).stats();
1996 let active = state.active();
1997 let requests_total = state.requests_total.load(Ordering::Relaxed);
1998 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1999 let uptime = state.started_at.elapsed().as_secs_f64();
2000
2001 let body = format!(
2002 "# HELP frink_requests_total Total chat completion requests received.\n\
2003 # TYPE frink_requests_total counter\n\
2004 frink_requests_total {requests_total}\n\
2005 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
2006 # TYPE frink_request_errors_total counter\n\
2007 frink_request_errors_total {errors_total}\n\
2008 # HELP frink_cache_hits_total Whole-response cache hits.\n\
2009 # TYPE frink_cache_hits_total counter\n\
2010 frink_cache_hits_total {}\n\
2011 # HELP frink_cache_misses_total Whole-response cache misses.\n\
2012 # TYPE frink_cache_misses_total counter\n\
2013 frink_cache_misses_total {}\n\
2014 # HELP frink_cache_entries Current whole-response cache entry count.\n\
2015 # TYPE frink_cache_entries gauge\n\
2016 frink_cache_entries {}\n\
2017 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2018 # TYPE frink_synthetic_weights gauge\n\
2019 frink_synthetic_weights {}\n\
2020 # HELP frink_uptime_seconds Seconds since this server process started.\n\
2021 # TYPE frink_uptime_seconds gauge\n\
2022 frink_uptime_seconds {uptime}\n",
2023 cache_stats.hits,
2024 cache_stats.misses,
2025 cache_stats.entries,
2026 active
2030 .as_ref()
2031 .map(|a| a.is_synthetic() as u8)
2032 .unwrap_or(0),
2033 );
2034
2035 let body = match active
2039 .as_ref()
2040 .and_then(|a| a.expert_store_stats())
2041 {
2042 Some(es) => format!(
2043 "{body}\
2044 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
2045 # TYPE frink_expert_cache_hits_total counter\n\
2046 frink_expert_cache_hits_total {}\n\
2047 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2048 # TYPE frink_expert_cache_misses_total counter\n\
2049 frink_expert_cache_misses_total {}\n\
2050 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
2051 # TYPE frink_expert_cache_evictions_total counter\n\
2052 frink_expert_cache_evictions_total {}\n\
2053 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2054 # TYPE frink_expert_cache_pass_throughs_total counter\n\
2055 frink_expert_cache_pass_throughs_total {}\n\
2056 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2057 # TYPE frink_expert_cache_bytes_read_total counter\n\
2058 frink_expert_cache_bytes_read_total {}\n\
2059 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2060 # TYPE frink_expert_cache_resident_bytes gauge\n\
2061 frink_expert_cache_resident_bytes {}\n",
2062 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2063 ),
2064 None => body,
2065 };
2066
2067 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2072 Some(batcher) => {
2073 let sched = batcher.stats();
2074 format!(
2075 "{body}\
2076 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2077 # TYPE frink_prefill_chunks_total counter\n\
2078 frink_prefill_chunks_total {}\n\
2079 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2080 # TYPE frink_prefill_tokens_total counter\n\
2081 frink_prefill_tokens_total {}\n\
2082 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2083 # TYPE frink_decode_steps_total counter\n\
2084 frink_decode_steps_total {}\n\
2085 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2086 # TYPE frink_scheduler_queue_depth gauge\n\
2087 frink_scheduler_queue_depth {}\n\
2088 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2089 # TYPE frink_scheduler_queue_rejected_total counter\n\
2090 frink_scheduler_queue_rejected_total {}\n\
2091 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2092 # TYPE frink_kv_blocks_total gauge\n\
2093 frink_kv_blocks_total {}\n\
2094 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2095 # TYPE frink_kv_blocks_free gauge\n\
2096 frink_kv_blocks_free {}\n\
2097 # HELP frink_kv_block_size Token positions per KV block.\n\
2098 # TYPE frink_kv_block_size gauge\n\
2099 frink_kv_block_size {}\n\
2100 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2101 # TYPE frink_kv_rejected_too_large_total counter\n\
2102 frink_kv_rejected_too_large_total {}\n\
2103 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2104 # TYPE frink_kv_rejected_context_length_total counter\n\
2105 frink_kv_rejected_context_length_total {}\n\
2106 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2107 # TYPE frink_scheduler_aborted_total counter\n\
2108 frink_scheduler_aborted_total {}\n\
2109 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2110 # TYPE frink_scheduler_max_seqs gauge\n\
2111 frink_scheduler_max_seqs {}\n\
2112 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2113 # TYPE frink_scheduler_prefill_chunk gauge\n\
2114 frink_scheduler_prefill_chunk {}\n",
2115 sched.prefill_chunks,
2116 sched.prefill_tokens,
2117 sched.decode_steps,
2118 sched.queue_depth,
2119 sched.queue_rejected,
2120 sched.kv_blocks_total,
2121 sched.kv_blocks_free,
2122 sched.kv_block_size,
2123 sched.kv_rejected_too_large,
2124 sched.kv_rejected_context_length,
2125 sched.aborted,
2126 sched.max_seqs,
2127 sched.prefill_chunk,
2128 )
2129 }
2130 None => body,
2131 };
2132
2133 (
2134 [(
2135 axum::http::header::CONTENT_TYPE,
2136 "text/plain; version=0.0.4",
2137 )],
2138 body,
2139 )
2140 .into_response()
2141}
2142
2143pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2144
2145pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2151 (
2152 StatusCode::BAD_REQUEST,
2153 Json(serde_json::json!({"error": {
2154 "message": message,
2155 "type": "invalid_request_error",
2156 "param": param,
2157 "code": null,
2158 }})),
2159 )
2160}
2161
2162pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2163 (
2164 StatusCode::NOT_IMPLEMENTED,
2165 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2166 )
2167}
2168
2169pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2170 let status = match e {
2171 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2172 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2176 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2181 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2187 StatusCode::SERVICE_UNAVAILABLE
2188 }
2189 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2193 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2197 };
2198 tracing::warn!("decode error: {e}");
2199 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2200 if let generate::DecodeError::KvBudgetExceeded {
2206 binding,
2207 estimated_bytes,
2208 limit_bytes,
2209 positions,
2210 positions_limit,
2211 ..
2212 } = &e
2213 {
2214 body["error"]["type"] = serde_json::json!("invalid_request_error");
2215 body["error"]["code"] = serde_json::json!(binding);
2216 body["error"]["binding"] = serde_json::json!(binding);
2217 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2218 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2219 body["error"]["positions"] = serde_json::json!(positions);
2220 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2221 }
2222 if let Some(secs) = e.retry_after_secs() {
2226 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2227 }
2228 (status, Json(body))
2229}
2230
2231pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2232 tracing::error!("generation task panicked: {e}");
2233 (
2234 StatusCode::INTERNAL_SERVER_ERROR,
2235 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2236 )
2237}
2238
2239#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2249 model: &Model,
2250 prompt: &str,
2251 params: &GenerationParams,
2252 kv_pool: Option<&generate::KvPoolConfig>,
2253 paged_kv: Option<&generate::PagedKvConfig>,
2254 prefix_cache: Option<&Mutex<PrefixCache>>,
2255 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2256 ceiling: Option<&budget::ContextCeiling>,
2257 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2258 mut emit: impl FnMut(usize, &str),
2263) -> Result<generate::Generated, generate::DecodeError> {
2264 let synthetic = model.is_synthetic();
2265 let _lora_lease = lora::lease(model, params.lora.as_deref());
2270 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2271 let params = &{
2277 let mut resolved = params.clone();
2278 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2279 model.encode(text, SpecialTokens::Parse)
2280 });
2281 resolved
2286 .token_mask
2287 .resolve(|text| model.encode(text, SpecialTokens::Parse));
2288 resolved.reasoning_budget = resolved
2294 .reasoning_budget
2295 .armed(resolved.reasoning, prompt, |text| {
2296 model.encode(text, SpecialTokens::Parse)
2297 })
2298 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2299 resolved
2300 };
2301 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2302 let _metal_private_guard =
2303 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2304 let (finishes, prompt_rows, prompt_ids, truncated_prompt, usage) = match model {
2305 Model::Gguf(m) => {
2306 if let Some(batcher) = continuous_batcher {
2307 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2308 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2309 let (finish, _generated_ids, text, usage) = if synthetic {
2310 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2311 } else {
2312 batcher.generate_streaming(
2313 tokens,
2314 params.clone(),
2315 m.stop_tokens.clone(),
2316 Some(|chunk: &str| {
2317 if !chunk.is_empty() {
2318 chunks[0].push(chunk.to_string());
2319 emit(0, chunk);
2320 }
2321 }),
2322 )?
2323 };
2324 if !text.is_empty() && chunks[0].is_empty() {
2325 chunks[0].push(text);
2326 }
2327 (
2340 vec![(finish, Vec::new())],
2341 Vec::new(),
2342 Vec::new(),
2343 None,
2344 usage,
2345 )
2346 } else {
2347 generate::generate(
2348 &m.decoder,
2349 m.tokenizer.as_ref(),
2350 &m.stop_tokens,
2351 m.bos_id,
2352 prompt,
2353 params,
2354 kv_pool,
2355 paged_kv,
2356 prefix_cache,
2357 ceiling,
2358 |choice, chunk| {
2359 chunks[choice].push(chunk.to_string());
2360 if !synthetic {
2364 emit(choice, chunk);
2365 }
2366 },
2367 )?
2368 }
2369 }
2370 Model::Kimi(m) => generate::generate_engine(
2371 &m.engine,
2372 &m.tokenizer,
2373 &m.stop_tokens,
2374 None,
2375 prompt,
2376 params,
2377 |chunk| {
2378 chunks[0].push(chunk.to_string());
2379 if !synthetic {
2380 emit(0, chunk);
2381 }
2382 },
2383 )?,
2384 Model::Mla(m) => generate::generate_engine(
2385 &m.engine,
2386 &m.tokenizer,
2387 &m.stop_tokens,
2388 m.bos_id,
2389 prompt,
2390 params,
2391 |chunk| {
2392 chunks[0].push(chunk.to_string());
2393 if !synthetic {
2394 emit(0, chunk);
2395 }
2396 },
2397 )?,
2398 Model::Gemma4(m) => generate::generate_engine(
2399 &m.engine,
2400 &m.tokenizer,
2401 &m.stop_tokens,
2402 m.bos_id,
2403 prompt,
2404 params,
2405 |chunk| {
2406 chunks[0].push(chunk.to_string());
2407 if !synthetic {
2408 emit(0, chunk);
2409 }
2410 },
2411 )?,
2412 Model::Glm52(m) => generate::generate_engine(
2413 &m.engine,
2414 &m.tokenizer,
2415 &m.stop_tokens,
2416 m.bos_id,
2417 prompt,
2418 params,
2419 |chunk| {
2420 chunks[0].push(chunk.to_string());
2421 if !synthetic {
2422 emit(0, chunk);
2423 }
2424 },
2425 )?,
2426 };
2427
2428 let mut full = chunks[0].concat();
2429 if synthetic {
2430 full = format!(
2431 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2432 to serve a real model. Decoded ids -> {full:?}]"
2433 );
2434 emit(0, &full);
2435 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2436 emit(0, &full);
2437 }
2438
2439 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2445 let mut out: Vec<generate::GeneratedChoice> = finishes
2446 .into_iter()
2447 .zip(chunks.into_iter().map(|c| c.concat()))
2448 .map(|((finish, logprobs), text)| generate::GeneratedChoice {
2449 finish,
2450 text,
2451 logprobs,
2452 })
2453 .collect();
2454 if let Some(first) = out.first_mut() {
2455 if synthetic {
2463 first.logprobs.clear();
2464 }
2465 first.text = full;
2466 }
2467 Ok(generate::Generated {
2468 choices: out,
2469 prompt_rows,
2470 prompt_ids,
2471 truncated_prompt,
2472 usage,
2473 })
2474}
2475
2476#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2481 model: &Model,
2482 prompt: &str,
2483 params: &GenerationParams,
2484 kv_pool: Option<&generate::KvPoolConfig>,
2485 paged_kv: Option<&generate::PagedKvConfig>,
2486 prefix_cache: Option<&Mutex<PrefixCache>>,
2487 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2488 ceiling: Option<&budget::ContextCeiling>,
2489 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2490 ) -> Result<generate::Generated, generate::DecodeError> {
2493 run_generation_emit(
2494 model,
2495 prompt,
2496 params,
2497 kv_pool,
2498 paged_kv,
2499 prefix_cache,
2500 continuous_batcher,
2501 ceiling,
2502 metal_private_decode_gate,
2503 |_, _| {},
2504 )
2505}
2506
2507pub(crate) fn prompt_from_messages(
2518 messages: &[ChatMessage],
2519 template: &chat_template::PromptTemplate,
2520 tools: &[ToolDef],
2521 extra: serde_json::Map<String, serde_json::Value>,
2522) -> Result<String, ApiError> {
2523 let rendered = if tools.is_empty() || template.handles_tools() {
2524 template.render(messages, tools, extra)
2525 } else {
2526 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2527 with_preamble.push(ChatMessage {
2528 role: "system".to_string(),
2529 content: Some(MessageContent::Text(tool_preamble(tools))),
2530 tool_calls: None,
2531 tool_call_id: None,
2532 reasoning_content: None,
2533 });
2534 with_preamble.extend_from_slice(messages);
2535 template.render(&with_preamble, &[], extra)
2536 };
2537 rendered.map_err(template_error_response)
2538}
2539
2540fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2545 (
2546 StatusCode::BAD_REQUEST,
2547 Json(serde_json::json!({
2548 "error": {
2549 "message": format!("chat template failed to render: {err}"),
2550 "type": "invalid_request_error",
2551 "param": "messages",
2552 "code": null,
2553 }
2554 })),
2555 )
2556}
2557
2558fn tool_preamble(tools: &[ToolDef]) -> String {
2571 let mut out = String::from(
2572 "You can call tools to help answer the user. To call a tool, respond with \
2573 EXACTLY one line in this format and nothing else:\n\
2574 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2575 object matching that tool's parameters>}}</tool_call>\n\n\
2576 Available tools:\n",
2577 );
2578 for t in tools {
2579 out.push_str(&format!(
2580 "- {}: {}\n parameters (JSON schema): {}\n",
2581 t.function.name,
2582 t.function.description.as_deref().unwrap_or(""),
2583 t.function
2584 .parameters
2585 .as_ref()
2586 .map(|v| v.to_string())
2587 .unwrap_or_else(|| "{}".to_string()),
2588 ));
2589 }
2590 out
2591}
2592
2593fn tool_call_deltas(
2603 events: Vec<crate::policy::parser::ToolCallEvent>,
2604 opened: &std::cell::Cell<usize>,
2605) -> (String, Vec<ToolCallDelta>) {
2606 let mut text = String::new();
2607 let mut deltas = Vec::new();
2608 for event in events {
2609 match event {
2610 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2611 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2612 opened.set(opened.get().max(index + 1));
2613 deltas.push(ToolCallDelta::opening(index, name));
2614 }
2615 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2616 if !fragment.is_empty() {
2617 deltas.push(ToolCallDelta::arguments(index, fragment));
2618 }
2619 }
2620 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2621 }
2622 }
2623 (text, deltas)
2624}
2625
2626fn build_response_message(
2637 text: String,
2638 tools: &[ToolDef],
2639 posture: output::OutputPosture,
2640 base_finish: &'static str,
2641) -> (ChatCompletionResponseMessage, &'static str) {
2642 let parsed = output::parse_output(&text, tools, posture);
2643 let calls: Vec<ToolCallOut> = parsed
2644 .calls
2645 .into_iter()
2646 .enumerate()
2647 .map(|(index, call)| ToolCallOut {
2648 id: format!("call_{index}"),
2649 kind: "function",
2650 function: ToolCallFunctionOut {
2651 name: call.name,
2652 arguments: call.arguments,
2653 },
2654 })
2655 .collect();
2656 if !calls.is_empty() {
2657 return (
2658 ChatCompletionResponseMessage {
2659 role: "assistant",
2660 content: None,
2661 reasoning_content: parsed.reasoning,
2662 tool_calls: Some(calls),
2663 },
2664 "tool_calls",
2665 );
2666 }
2667 (
2668 ChatCompletionResponseMessage {
2669 role: "assistant",
2670 content: Some(parsed.content),
2671 reasoning_content: parsed.reasoning,
2672 tool_calls: None,
2673 },
2674 base_finish,
2675 )
2676}
2677
2678fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2683 let mut history = match &req.session_id {
2684 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2685 None => req.messages.clone(),
2686 };
2687 if req.json_object_mode() {
2688 inject_json_object_system_hint(&mut history);
2689 }
2690 history
2691}
2692
2693fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2694 const HINT: &str =
2695 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2696 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2697 match &mut sys.content {
2698 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2699 s.push_str("\n\n");
2700 s.push_str(HINT);
2701 }
2702 None => {
2703 sys.content = Some(MessageContent::Text(HINT.to_string()));
2704 }
2705 _ => {}
2706 }
2707 } else {
2708 messages.insert(
2709 0,
2710 ChatMessage {
2711 role: "system".to_string(),
2712 content: Some(MessageContent::Text(HINT.to_string())),
2713 tool_calls: None,
2714 tool_call_id: None,
2715 reasoning_content: None,
2716 },
2717 );
2718 }
2719}
2720
2721async fn chat_completions(
2722 State(state): State<Arc<AppState>>,
2723 headers: axum::http::HeaderMap,
2724 Json(req): Json<ChatCompletionRequest>,
2725) -> Response {
2726 let attribution = attribution::Attribution::from_headers(&headers);
2727 state
2728 .requests_total
2729 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2730 let started = std::time::Instant::now();
2731
2732 let request_id = frink_api::next_request_id();
2736 let stream = req.stream.unwrap_or(false);
2737
2738 let refusal = cache_admin::check_admission(&state)
2744 .err()
2745 .or_else(|| req.validate_supported_fields().err());
2746 if let Some(err) = refusal {
2747 state
2748 .request_errors_total
2749 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2750 let response = err.into_response();
2751 state.record_request(stats::Record {
2752 request_id: &request_id,
2753 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2754 model: state.active_model_name(),
2755 status: response.status().as_u16(),
2756 stream,
2757 duration_ms: started.elapsed().as_millis() as u64,
2758 usage: None,
2759 attribution: &attribution,
2760 });
2761 return response;
2762 }
2763
2764 let response = if stream {
2765 chat_completions_stream(
2766 Arc::clone(&state),
2767 req,
2768 request_id.clone(),
2769 started,
2770 attribution.clone(),
2771 )
2772 .await
2773 .into_response()
2774 } else {
2775 chat_completions_full(
2776 Arc::clone(&state),
2777 req,
2778 request_id.clone(),
2779 started,
2780 attribution.clone(),
2781 )
2782 .await
2783 .into_response()
2784 };
2785
2786 if response.status().is_client_error() || response.status().is_server_error() {
2787 state
2788 .request_errors_total
2789 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2790 state.record_request(stats::Record {
2794 request_id: &request_id,
2795 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2796 model: state.active_model_name(),
2799 status: response.status().as_u16(),
2800 stream,
2801 duration_ms: started.elapsed().as_millis() as u64,
2802 usage: None,
2803 attribution: &attribution,
2804 });
2805 }
2806 state.mark_request_finished();
2807
2808 response
2809}
2810
2811async fn chat_completions_full(
2812 state: Arc<AppState>,
2813 req: ChatCompletionRequest,
2814 request_id: String,
2815 started: std::time::Instant,
2816 attribution: attribution::Attribution,
2817) -> Result<Json<ChatCompletionResponse>, ApiError> {
2818 let tools_active = req.tools_active();
2819 let active = state.require_active()?;
2823 let history = resolve_history(&state, &req);
2824 let template = active.generative()?.chat_template();
2825 let kwargs = req.resolve_template_kwargs(&template);
2826 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2827 let mut params =
2835 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2836 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2837 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2838
2839 let mut generated_logprobs: Vec<crate::sampling_loop::PerTokenProbs> = Vec::new();
2843 let n_logprobs = req.n_logprobs()?;
2846 let decode_any = |id: usize| active.decode_any(&[id]);
2848 let (completion, cache_status) = if let Some(cached) = key
2849 .as_ref()
2850 .and_then(|key| lock_cache(&state.response_cache).get(key))
2851 {
2852 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2853 (cached, "hit")
2854 } else {
2855 let produced = decode_task::buffered(
2856 decode_task::DecodeHandles::take(&state, &active)?,
2857 prompt.clone(),
2858 params,
2859 )
2860 .await?;
2861 let usage = produced.usage;
2862 let choices = produced.choices;
2863
2864 generated_logprobs = choices.iter().map(|c| c.logprobs.clone()).collect();
2869 let completion = response_cache::CachedCompletion {
2870 choices: choices.into_iter().map(|c| (c.finish, c.text)).collect(),
2871 usage,
2872 };
2873 let cache_status = match key {
2879 Some(key) => match completion.clone().cacheable() {
2883 Some(cacheable) => {
2884 tracing::debug!("cache miss for key {}", key.digest());
2885 lock_cache(&state.response_cache).put(key, cacheable);
2886 "miss"
2887 }
2888 None => "skip",
2889 },
2890 None => "skip",
2891 };
2892 (completion, cache_status)
2893 };
2894 let content = completion.first_text().to_string();
2897
2898 if req.json_object_mode() {
2899 json_mode::validate_json_object_output(&content)?;
2900 }
2901
2902 if let Some(id) = &req.session_id {
2906 state.sessions.store_reply(
2907 id,
2908 ChatMessage {
2909 role: "assistant".to_string(),
2910 content: Some(MessageContent::Text(content.clone())),
2911 tool_calls: None,
2912 tool_call_id: None,
2913 reasoning_content: None,
2914 },
2915 );
2916 }
2917
2918 let posture = output::OutputPosture::resolve_full(
2923 active.reasoning_format(),
2924 active.tool_call_format(),
2925 &prompt,
2926 );
2927 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2928 let wanted = req.unimplemented.n.unwrap_or(1).max(1) as usize;
2932 let ranked: Vec<(generate::FinishReason, String)> = if completion.choices.len() > wanted {
2933 let scored: Vec<crate::generate::GeneratedChoice> = completion
2934 .choices
2935 .into_iter()
2936 .zip(
2937 generated_logprobs
2938 .iter()
2939 .cloned()
2940 .chain(std::iter::repeat(Vec::new())),
2941 )
2942 .map(
2943 |((finish, text), logprobs)| crate::generate::GeneratedChoice {
2944 finish,
2945 text,
2946 logprobs,
2947 },
2948 )
2949 .collect();
2950 let best = crate::best_of::take_best(scored, wanted);
2951 generated_logprobs = best.iter().map(|c| c.logprobs.clone()).collect();
2952 best.into_iter().map(|c| (c.finish, c.text)).collect()
2953 } else {
2954 completion.choices
2955 };
2956 let render_piece =
2962 crate::logprobs::piece_renderer(req.unimplemented.tokens_as_ids(), &decode_any);
2963 let rendered: Vec<ChatCompletionChoice> = ranked
2964 .into_iter()
2965 .enumerate()
2966 .map(|(index, (finish, text))| {
2967 let (message, finish_reason) =
2968 build_response_message(text, tools, posture, finish.as_str());
2969 ChatCompletionChoice {
2970 index,
2971 message,
2972 finish_reason,
2973 logprobs: n_logprobs.map(|k| {
2974 crate::logprobs::render_chat(
2975 generated_logprobs.get(index).unwrap_or(&Vec::new()),
2976 Some(k),
2977 render_piece.as_ref(),
2978 )
2979 }),
2980 }
2981 })
2982 .collect();
2983
2984 state.record_request(stats::Record {
2985 request_id: &request_id,
2986 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2987 model: Some(active.name().to_string()),
2990 status: 200,
2991 stream: false,
2992 duration_ms: started.elapsed().as_millis() as u64,
2993 usage: Some(&completion.usage),
2994 attribution: &attribution,
2995 });
2996
2997 Ok(Json(ChatCompletionResponse {
2998 id: request_id.clone(),
2999 request_id,
3000 object: "chat.completion",
3001 model: req.model,
3002 choices: rendered,
3003 usage: completion.usage,
3004 frink_cache: cache_status,
3005 }))
3006}
3007
3008async fn chat_completions_stream(
3009 state: Arc<AppState>,
3010 req: ChatCompletionRequest,
3011 request_id: String,
3012 started: std::time::Instant,
3013 attribution: attribution::Attribution,
3014) -> Result<Response, ApiError> {
3015 let tools_active = req.tools_active();
3025 let active = state.require_active()?;
3029 let history = resolve_history(&state, &req);
3030 let template = active.generative()?.chat_template();
3031 let kwargs = req.resolve_template_kwargs(&template);
3032 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
3033 let model_name = req.model.clone();
3034 let session_id = req.session_id.clone();
3035 let sessions = state.sessions.clone();
3036
3037 let model = Arc::clone(active.generative()?);
3038 let kv_pool = state.kv_pool.clone();
3039 let paged_kv = state.paged_kv.clone();
3040 let prefix_cache = state.prefix_cache.clone();
3041 let batcher = active.batcher.clone();
3042 let ceiling = active.ceiling.clone();
3043 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3044 let mut params =
3045 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
3046 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
3047 params.interleave_choices = params.n > 1;
3053 let stats_state = Arc::clone(&state);
3054 let served_model = active.name().to_string();
3057 let posture = output::OutputPosture::resolve_full(
3061 active.reasoning_format(),
3062 active.tool_call_format(),
3063 &prompt,
3064 );
3065 let offered_tools: Vec<ToolDef> = if tools_active {
3068 req.tools.clone()
3069 } else {
3070 Vec::new()
3071 };
3072
3073 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3078 params.cancel = Some(cancel_token.clone());
3079
3080 let overlap = true;
3089
3090 let slot = req
3094 .stream_resumable
3095 .unwrap_or(false)
3096 .then(|| state.streams.register(&request_id));
3097 let emitter = resume::Emitter::new(slot);
3098
3099 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3100 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3104 id: request_id.clone(),
3105 request_id: None,
3106 object: "chat.completion.chunk",
3107 model: model_name.clone(),
3108 choices: vec![ChatCompletionChunkChoice {
3109 index: 0,
3110 delta: ChatCompletionChunkDelta {
3111 role: None,
3112 content: None,
3113 reasoning_content: None,
3114 tool_calls: None,
3115 },
3116 finish_reason: None,
3117 }],
3118 usage: None,
3119 });
3120
3121 tokio::task::spawn_blocking(move || {
3122 let _cancel_guard = cancel_guard;
3125 let tx_chunks = tx.clone();
3126 let orphan_timeout = sse::orphan_timeout_from_env();
3131 let head_request_id = request_id.clone();
3132 let announced = std::cell::Cell::new(false);
3136 let emitters: Rc<RefCell<Vec<crate::chat_stream_choice::ChoiceEmitter>>> =
3142 Rc::new(RefCell::new(
3143 (0..params.n.max(1))
3144 .map(|_| {
3145 crate::chat_stream_choice::ChoiceEmitter::new(
3146 posture.reasoning_parser(),
3147 tools_active.then(|| posture.tool_call_parser(&offered_tools)),
3148 )
3149 })
3150 .collect(),
3151 ));
3152 let emit_choices = Rc::clone(&emitters);
3153 let result = run_generation_emit(
3154 &model,
3155 &prompt,
3156 ¶ms,
3157 kv_pool.as_ref(),
3158 paged_kv.as_ref(),
3159 prefix_cache.as_deref(),
3160 batcher.as_ref(),
3161 ceiling.as_deref(),
3162 metal_private_decode_gate.as_deref(),
3163 |choice, chunk| {
3164 if !overlap || chunk.is_empty() {
3165 return;
3166 }
3167 let mut held = emit_choices.borrow_mut();
3168 let Some(emitter_state) = held.get_mut(choice) else {
3169 return;
3170 };
3171 let delta = emitter_state.push(chunk);
3172 if delta.is_empty() {
3173 return;
3174 }
3175 let request_id = (!announced.get()).then(|| {
3178 announced.set(true);
3179 head_request_id.clone()
3180 });
3181 let wire = delta.into_choice(choice, emitter_state.start());
3182 drop(held);
3183 let payload = ChatCompletionChunk {
3184 id: head_request_id.clone(),
3185 request_id,
3186 object: "chat.completion.chunk",
3187 model: model_name.clone(),
3188 choices: vec![wire],
3189 usage: None,
3190 };
3191 if let Err(why) =
3200 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3201 {
3202 if why == sse::SendFailure::Orphaned {
3203 tracing::warn!(
3204 "SSE stream {head_request_id} accepted nothing for the orphan \
3205 deadline; treating it as abandoned"
3206 );
3207 }
3208 if !emitter.is_resumable() {
3220 cancel_token.cancel();
3221 }
3222 }
3223 },
3224 );
3225
3226 let mut pending_request_id = (!announced.get()).then(|| request_id.clone());
3232
3233 match result {
3234 Ok(generated) => {
3235 let usage = generated.usage;
3236 let produced: Vec<(generate::FinishReason, String)> = generated
3237 .choices
3238 .into_iter()
3239 .map(|c| (c.finish, c.text))
3240 .collect();
3241 assert!(
3242 !produced.is_empty(),
3243 "a generation produces at least one choice"
3244 );
3245 if let Some(id) = &session_id {
3250 sessions.store_reply(
3251 id,
3252 ChatMessage {
3253 role: "assistant".to_string(),
3254 content: Some(MessageContent::Text(produced[0].1.clone())),
3255 tool_calls: None,
3256 tool_call_id: None,
3257 reasoning_content: None,
3258 },
3259 );
3260 }
3261 for (index, (finish, full_text)) in produced.iter().enumerate() {
3262 let (finish, full_text) = (finish.clone(), full_text.as_str());
3263 let mut streamed_finish: Option<&'static str> = None;
3269 if overlap {
3270 let (tail, first, opened) = {
3271 let mut held = emitters.borrow_mut();
3272 let state = &mut held[index];
3273 let tail = state.flush();
3274 (tail, state.start(), state.opened_calls())
3275 };
3276 if !tail.is_empty() {
3277 let payload = ChatCompletionChunk {
3278 id: request_id.clone(),
3279 request_id: pending_request_id.take(),
3280 object: "chat.completion.chunk",
3281 model: model_name.clone(),
3282 choices: vec![tail.into_choice(index, first)],
3283 usage: None,
3284 };
3285 let _ = sse::send_or_orphan(
3286 &tx,
3287 Ok(emitter.event(&payload)),
3288 orphan_timeout,
3289 );
3290 }
3291 if opened > 0 {
3292 streamed_finish = Some("tool_calls");
3293 }
3294 } else {
3295 let parsed = output::parse_output(full_text, &offered_tools, posture);
3298 let tool_calls: Vec<ToolCallDelta> = parsed
3299 .calls
3300 .iter()
3301 .enumerate()
3302 .map(|(index, call)| {
3303 ToolCallDelta::whole(
3304 index,
3305 call.name.clone(),
3306 call.arguments.clone(),
3307 )
3308 })
3309 .collect();
3310 if !tool_calls.is_empty() {
3311 streamed_finish = Some("tool_calls");
3312 }
3313 if !tool_calls.is_empty()
3314 || !parsed.content.is_empty()
3315 || parsed.reasoning.is_some()
3316 {
3317 let payload = ChatCompletionChunk {
3318 id: request_id.clone(),
3319 request_id: pending_request_id.take(),
3320 object: "chat.completion.chunk",
3321 model: model_name.clone(),
3322 choices: vec![ChatCompletionChunkChoice {
3323 index,
3324 delta: ChatCompletionChunkDelta {
3325 role: Some("assistant"),
3326 content: (!parsed.content.is_empty()
3327 && tool_calls.is_empty())
3328 .then(|| parsed.content.clone()),
3329 reasoning_content: parsed.reasoning.clone(),
3330 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3331 },
3332 finish_reason: None,
3333 }],
3334 usage: None,
3335 };
3336 let _ = sse::send_or_orphan(
3337 &tx,
3338 Ok(emitter.event(&payload)),
3339 orphan_timeout,
3340 );
3341 }
3342 }
3343 let final_finish_reason = match streamed_finish {
3347 Some(reason) if finish.as_str() != "length" => reason,
3348 _ => finish.as_str(),
3349 };
3350 let last = index + 1 == produced.len();
3354 let final_payload = ChatCompletionChunk {
3355 id: request_id.clone(),
3356 request_id: pending_request_id.take(),
3357 object: "chat.completion.chunk",
3358 model: model_name.clone(),
3359 choices: vec![ChatCompletionChunkChoice {
3360 index,
3361 delta: ChatCompletionChunkDelta {
3362 role: None,
3363 content: None,
3364 reasoning_content: None,
3365 tool_calls: None,
3366 },
3367 finish_reason: Some(final_finish_reason),
3368 }],
3369 usage: last.then(|| usage.clone()),
3370 };
3371 let _ =
3372 sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3373 }
3374 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3375 stats_state.record_request(stats::Record {
3380 request_id: &request_id,
3381 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3382 model: Some(served_model.clone()),
3383 status: 200,
3384 stream: true,
3385 duration_ms: started.elapsed().as_millis() as u64,
3386 usage: Some(&usage),
3387 attribution: &attribution,
3388 });
3389 }
3390 Err(e) => {
3391 tracing::warn!("decode error on streamed request {request_id}: {e}");
3392 stats_state.record_request(stats::Record {
3398 request_id: &request_id,
3399 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3400 model: Some(served_model.clone()),
3401 status: 500,
3402 stream: true,
3403 duration_ms: started.elapsed().as_millis() as u64,
3404 usage: None,
3405 attribution: &attribution,
3406 });
3407 let payload = ChatCompletionChunk {
3408 id: request_id.clone(),
3409 request_id: pending_request_id.take(),
3410 object: "chat.completion.chunk",
3411 model: model_name,
3412 choices: vec![ChatCompletionChunkChoice {
3413 index: 0,
3414 delta: ChatCompletionChunkDelta {
3415 role: Some("assistant"),
3416 content: Some(format!("[error: {e}]")),
3417 reasoning_content: None,
3418 tool_calls: None,
3419 },
3420 finish_reason: Some("stop"),
3421 }],
3422 usage: None,
3423 };
3424 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3425 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3426 }
3427 }
3428 drop(emitter);
3432 });
3433
3434 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3435 Ok((
3458 [(
3459 axum::http::HeaderName::from_static("x-accel-buffering"),
3460 axum::http::HeaderValue::from_static("no"),
3461 )],
3462 Sse::new(stream),
3463 )
3464 .into_response())
3465}
3466
3467fn protected_routes() -> Router<Arc<AppState>> {
3498 use frink_api::routes;
3499
3500 Router::new()
3501 .route(routes::V1_MODELS, get(list_models))
3502 .route(routes::V1_RESPONSES, post(responses::responses))
3506 .route(
3507 &axum_path(routes::V1_RESPONSE),
3508 get(responses::responses_get),
3509 )
3510 .route(
3511 &axum_path(routes::V1_RESPONSE_CANCEL),
3512 post(responses::responses_cancel),
3513 )
3514 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3515 .route(routes::V1_STATS, get(serving_stats))
3516 .route(routes::V1_REQUESTS, get(recent_requests))
3517 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3518 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3519 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3520 .route(
3521 routes::LORA_ADAPTERS,
3522 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3523 )
3524 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3525 .route(routes::V1_CANCEL, post(cancel_generation))
3529 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3534 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3535 .route(routes::V1_MESSAGES, post(anthropic::messages))
3536 .route(
3537 routes::V1_MESSAGES_COUNT_TOKENS,
3538 post(anthropic::count_tokens),
3539 )
3540 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3541 .route(routes::COMPLETION, post(completion::completion))
3546 .route(routes::COMPLETIONS, post(completion::completion))
3547 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3548 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3549 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3555 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3556 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3557 .route(routes::V1_RERANK, post(rerank::rerank))
3561 .route(routes::RERANK, post(rerank::rerank))
3562 .route(routes::CACHE_STATS, get(cache_stats))
3563 .route(routes::METRICS, get(metrics))
3564 .route(routes::ADMIN_MODELS, get(admin::models))
3569 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3570 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3571 .route(routes::SLEEP, post(admin::sleep))
3575 .route(routes::WAKE_UP, post(admin::wake_up))
3576 .route(routes::IS_SLEEPING, get(admin::is_sleeping))
3577 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3578 .route(routes::ADMIN_TASKS, get(admin::tasks))
3579 .route(&admin::cancel_route(), post(admin::cancel_task))
3580 .route(routes::ADMIN_STATS, get(admin::stats))
3581 .merge(conversations::router())
3585}
3586
3587fn axum_path(template: &str) -> String {
3588 let mut out = String::with_capacity(template.len());
3589 let mut rest = template;
3590 while let Some(open) = rest.find('{') {
3591 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3592 break;
3593 };
3594 out.push_str(&rest[..open]);
3595 out.push(':');
3596 out.push_str(&rest[open + 1..close]);
3597 rest = &rest[close + 1..];
3598 }
3599 out.push_str(rest);
3600 out
3601}
3602
3603async fn cancel_generation(
3612 State(state): State<Arc<AppState>>,
3613 Json(req): Json<frink_api::CancelGenerationRequest>,
3614) -> Response {
3615 let cancelled = state.cancels.cancel(&req.request_id);
3616 let status = if cancelled {
3617 StatusCode::OK
3618 } else {
3619 StatusCode::NOT_FOUND
3620 };
3621 let detail = if cancelled {
3622 "the generation was asked to stop; it ends at its next token".to_string()
3623 } else {
3624 "no generation with that request_id is running -- it has already \
3625 finished, was never issued, or was served by a path that does \
3626 not register for cancellation"
3627 .to_string()
3628 };
3629 (
3630 status,
3631 Json(frink_api::CancelGenerationResponse {
3632 request_id: req.request_id,
3633 cancelled,
3634 detail,
3635 }),
3636 )
3637 .into_response()
3638}
3639
3640type Activated = (
3644 Loaded,
3645 Option<serving::batch::ContinuousBatcher>,
3646 Option<Arc<budget::ContextCeiling>>,
3647);
3648
3649fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3667 let mut batcher = serving::batch::BatcherConfig::from_env();
3668 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3669 return batcher;
3672 }
3673 let Some(path) = path else {
3674 return batcher;
3675 };
3676 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3682 let Some((priced, gguf_ctx, source)) = priced else {
3683 return batcher;
3684 };
3685 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3686 tracing::warn!(
3690 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3691 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3692 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3693 admit on a number you choose.",
3694 priced.weights_bytes,
3695 priced.device_budget_bytes,
3696 );
3697 return batcher;
3698 };
3699 tracing::info!("{source}");
3700 tracing::info!("{}", derived.fit);
3701 let adopted = budget::apply_derived(&mut batcher, &derived);
3702 if adopted.max_context {
3703 tracing::info!(
3704 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3705 override with FRINK_CB_MAX_CONTEXT",
3706 derived.max_context
3707 );
3708 }
3709 if adopted.kv_blocks {
3710 tracing::info!(
3711 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3712 derived.kv_blocks,
3713 batcher.kv_block_size
3714 );
3715 }
3716 if let Some(narrowed) = adopted.max_context_narrowed {
3717 tracing::info!(
3718 "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",
3719 batcher.kv_blocks.unwrap_or_default(),
3720 batcher.kv_block_size
3721 );
3722 }
3723 batcher
3724}
3725
3726pub(crate) fn activate_loaded_model(
3735 loaded: model::LoadedModel,
3736 enable_continuous_batching: bool,
3737 path: Option<&str>,
3738 paged_kv: Option<&generate::PagedKvConfig>,
3739) -> Activated {
3740 match loaded {
3741 model::LoadedModel::Gguf(g) => {
3742 let decoder = Arc::new(g.decoder);
3743 let tokenizer = Arc::new(g.tokenizer);
3744 let config = price_batcher_config(path);
3745 let shape =
3751 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3752 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3753 let batcher = if enable_continuous_batching {
3754 tracing::info!(
3755 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3756 (stop sequences use the same pending-buffer trim as the private generate loop)"
3757 );
3758 let tok = Arc::clone(&tokenizer);
3759 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3760 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3761 Arc::clone(&decoder),
3762 decode,
3763 config,
3764 Arc::clone(&ceiling),
3765 paged_kv.cloned(),
3766 ))
3767 } else {
3768 None
3769 };
3770 (
3771 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3772 decoder,
3773 tokenizer,
3774 stop_tokens: g.stop_tokens,
3775 bos_id: g.bos_id,
3776 is_synthetic: g.is_synthetic,
3777 chat_template: g.chat_template,
3778 }))),
3779 batcher,
3780 Some(ceiling),
3781 )
3782 }
3783 model::LoadedModel::Kimi(k) => (
3784 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3785 engine: k.engine,
3786 tokenizer: k.tokenizer,
3787 stop_tokens: k.stop_tokens,
3788 chat_template: k.chat_template,
3789 }))),
3790 None,
3791 None,
3792 ),
3793 model::LoadedModel::Mla(m) => (
3794 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3795 engine: m.engine,
3796 tokenizer: m.tokenizer,
3797 stop_tokens: m.stop_tokens,
3798 bos_id: m.bos_id,
3799 name: m.name,
3800 chat_template: m.chat_template,
3801 }))),
3802 None,
3803 None,
3804 ),
3805 model::LoadedModel::Gemma4(m) => (
3806 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3807 engine: m.engine,
3808 tokenizer: m.tokenizer,
3809 stop_tokens: m.stop_tokens,
3810 bos_id: m.bos_id,
3811 name: m.name,
3812 chat_template: m.chat_template,
3813 }))),
3814 None,
3815 None,
3816 ),
3817 model::LoadedModel::Glm52(g) => (
3818 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3819 engine: g.engine,
3820 tokenizer: g.tokenizer,
3821 stop_tokens: g.stop_tokens,
3822 bos_id: g.bos_id,
3823 name: g.name,
3824 chat_template: g.chat_template,
3825 }))),
3826 None,
3827 None,
3828 ),
3829 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3834 }
3835}
3836
3837struct StartupModels {
3844 loaded: model::LoadedModel,
3845 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3846}
3847
3848fn continuous_batching_env() -> Option<bool> {
3849 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3850 .ok()
3851 .map(|v| v.trim().to_ascii_lowercase())
3852 .as_deref()
3853 {
3854 None => None,
3855 Some("1" | "true" | "yes" | "on") => Some(true),
3856 Some("0" | "false" | "no" | "off") => Some(false),
3857 _ => None,
3858 }
3859}
3860
3861fn metal_private_decode_active() -> bool {
3862 #[cfg(feature = "metal")]
3863 {
3864 BUILT_WITH_METAL
3865 && frink_metal::attn::metal_attn_enabled()
3866 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3867 }
3868 #[cfg(not(feature = "metal"))]
3869 {
3870 false
3871 }
3872}
3873
3874fn continuous_batching_compatible(
3875 loaded: &model::LoadedModel,
3876 kv_pool: &Option<generate::KvPoolConfig>,
3877 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3878 paged_kv: &Option<generate::PagedKvConfig>,
3879) -> bool {
3880 matches!(loaded, model::LoadedModel::Gguf(_))
3881 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3882}
3883
3884fn resolve_continuous_batching_enabled(
3885 loaded: &model::LoadedModel,
3886 kv_pool: &Option<generate::KvPoolConfig>,
3887 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3888 paged_kv: &Option<generate::PagedKvConfig>,
3889) -> bool {
3890 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3891 return false;
3892 }
3893 match continuous_batching_env() {
3894 Some(true) => true,
3895 Some(false) => false,
3896 None => metal_private_decode_active(),
3897 }
3898}
3899
3900fn acquire_metal_private_decode_gate(
3901 gate: Option<&std::sync::Mutex<()>>,
3902 used_batcher: bool,
3903) -> Option<std::sync::MutexGuard<'_, ()>> {
3904 if used_batcher {
3905 None
3906 } else {
3907 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3908 }
3909}
3910
3911fn build_app_state(
3912 models: StartupModels,
3913 kv_pool: Option<generate::KvPoolConfig>,
3914 paged_kv: Option<generate::PagedKvConfig>,
3915 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3916 enable_continuous_batching: bool,
3917 mcp: Option<mcp::LoadedMcpConfig>,
3918 detection: Arc<health::Detection>,
3919) -> AppState {
3920 let StartupModels { loaded, embedding } = models;
3921 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3922 let (loaded, batcher, ceiling) = activate_loaded_model(
3923 loaded,
3924 enable_continuous_batching,
3925 configured_path.as_deref(),
3926 paged_kv.as_ref(),
3927 );
3928 let id = startup_model_id();
3934 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3935 {
3936 None
3937 } else {
3938 tracing::info!(
3939 "Metal private-loop decode will serialize concurrent requests until \
3940 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3941 );
3942 Some(Arc::new(std::sync::Mutex::new(())))
3943 };
3944 AppState {
3945 slept: Mutex::new(None),
3946 embedding,
3947 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3948 id,
3949 loaded,
3950 batcher,
3951 ceiling,
3952 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3953 }))),
3954 paged_kv,
3955 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3956 tasks: Arc::new(tasks::TaskRegistry::new()),
3957 cancels: Arc::new(cancel::CancelRegistry::new()),
3958 stats: stats::Stats::new(),
3959 streams: resume::StreamRegistry::new(),
3960 model_dir: admin::model_dirs().into_iter().next(),
3961 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3962 kv_pool,
3963 prefix_cache,
3964 sessions: session::SessionStore::new(),
3965 requests_total: std::sync::atomic::AtomicU64::new(0),
3966 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3967 started_at: std::time::Instant::now(),
3968 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3969 detection,
3970 mcp,
3971 continuous_batching_enabled: enable_continuous_batching,
3972 metal_private_decode_gate,
3973 loading_model: Mutex::new(None),
3974 last_load_error: Mutex::new(None),
3975 serving: Mutex::new(crate::stats::ServingStats::default()),
3976 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3977 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3978 started_unix: unix_now(),
3979 }
3980}
3981
3982fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3990 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3991 return Ok(None);
3992 };
3993 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3994 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3995 tracing::info!(
3996 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3997 model.name(),
3998 model.architecture(),
3999 model.n_embd(),
4000 model.pooling_type().name(),
4001 model.n_ctx_train(),
4002 );
4003 Ok(Some(Arc::new(model)))
4004}
4005
4006fn unix_now() -> u64 {
4011 std::time::SystemTime::now()
4012 .duration_since(std::time::UNIX_EPOCH)
4013 .map(|d| d.as_secs())
4014 .unwrap_or(0)
4015}
4016
4017fn startup_model_id() -> Option<String> {
4022 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
4023 let configured = std::fs::canonicalize(&configured).ok()?;
4024 admin::discover(&admin::model_dirs())
4025 .into_iter()
4026 .find(|d| {
4027 std::fs::canonicalize(&d.path)
4028 .map(|p| p == configured)
4029 .unwrap_or(false)
4030 })
4031 .map(|d| d.id)
4032}
4033
4034fn init_cpu_pool() {
4043 match frink_core::threads::init_cpu_pool() {
4044 Some(n) => eprintln!(
4045 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
4046 frink_core::threads::perf_core_count()
4047 ),
4048 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
4049 }
4050}
4051
4052fn announce_ready(addr: SocketAddr, scheme: &str) {
4065 use std::io::Write;
4066 let ready =
4067 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
4068 let mut stdout = std::io::stdout().lock();
4069 let _ = writeln!(stdout, "{}", ready.to_line());
4070 let _ = stdout.flush();
4071}
4072
4073async fn shutdown_signal(exit_on_stdin_close: bool) {
4085 if !exit_on_stdin_close {
4086 std::future::pending::<()>().await;
4087 return;
4088 }
4089 let _ = tokio::task::spawn_blocking(|| {
4090 use std::io::Read;
4091 let mut sink = [0u8; 256];
4092 let mut stdin = std::io::stdin().lock();
4093 loop {
4094 match stdin.read(&mut sink) {
4095 Ok(0) => break,
4097 Ok(_) => continue,
4099 Err(e) => {
4100 tracing::warn!("stdin read failed ({e}); treating it as closed");
4101 break;
4102 }
4103 }
4104 }
4105 })
4106 .await;
4107 tracing::info!("stdin closed; shutting down");
4108}
4109
4110fn tokio_worker_threads() -> usize {
4115 std::env::var("FRINK_TOKIO_WORKERS")
4116 .ok()
4117 .and_then(|v| v.trim().parse::<usize>().ok())
4118 .filter(|n| *n > 0)
4119 .unwrap_or(2)
4120}
4121
4122fn install_ring_crypto_provider() {
4139 let _ = rustls::crypto::ring::default_provider().install_default();
4140}
4141
4142pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4148 if args.list_devices {
4149 frink_models::devices::print_available_devices();
4150 return Ok(());
4151 }
4152 apply_cli_overrides(&args)?;
4153
4154 let _instance = {
4158 use frink_core::instance::{register, InstancePolicy};
4159 let policy = if args.allow_multiple_instances {
4160 InstancePolicy::Multi
4161 } else {
4162 InstancePolicy::from_env_or(InstancePolicy::Single)
4163 };
4164 let model = std::env::var("FRINK_MODEL_PATH").ok();
4165 register(
4166 "server",
4167 model.as_deref(),
4168 frink_core::instance::current_backend(),
4169 policy,
4170 )
4171 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4172 };
4173
4174 let journal = journal::Journal::from_env();
4175 eprintln!(
4176 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4177 journal.path()
4178 );
4179 journal.append(&journal::Record::session_start(
4180 env!("CARGO_PKG_VERSION"),
4181 std::process::id(),
4182 ));
4183 journal::install_panic_hook(journal.clone());
4184
4185 let mcp_config_path = args.mcp_config.clone();
4186 let exit_on_stdin_close = args.exit_on_stdin_close
4187 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4188 .map(|v| v == "1")
4189 .unwrap_or(false);
4190
4191 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4195 init_cpu_pool();
4196
4197 let runtime = tokio::runtime::Builder::new_multi_thread()
4198 .worker_threads(tokio_worker_threads())
4199 .enable_all()
4200 .build()?;
4201 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4202
4203 let reason = match &result {
4204 Ok(()) => "normal".to_string(),
4205 Err(e) => e.to_string(),
4206 };
4207 journal.append(&journal::Record::session_exit(reason));
4208
4209 runtime.shutdown_background();
4214
4215 result
4216}
4217
4218async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4219 let _ = tracing_subscriber::fmt::try_init();
4225
4226 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4236 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4237 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4238 .map(|v| v == "1")
4239 .unwrap_or(false);
4240 if let Err(msg) =
4241 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4242 {
4243 anyhow::bail!(msg);
4244 }
4245
4246 let embedding_model = load_embedding_model()?;
4253
4254 let mut loaded = model::load()?;
4255 match &loaded {
4256 model::LoadedModel::Gguf(g) => tracing::info!(
4257 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4258 g.decoder.config.name,
4259 g.is_synthetic,
4260 g.tokenizer.kind()
4261 ),
4262 model::LoadedModel::Kimi(k) => tracing::info!(
4263 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4264 k.tokenizer.vocab_size()
4265 ),
4266 model::LoadedModel::Mla(m) => tracing::info!(
4267 "loaded MLA GGUF '{}' (tokenizer={})",
4268 m.name,
4269 m.tokenizer.kind()
4270 ),
4271 model::LoadedModel::Gemma4(m) => tracing::info!(
4272 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4273 m.name,
4274 m.tokenizer.kind()
4275 ),
4276 model::LoadedModel::Glm52(g) => tracing::info!(
4277 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4278 g.name,
4279 g.tokenizer.kind()
4280 ),
4281 model::LoadedModel::Encoder(_) => {}
4284 }
4285 let metal_default_moe_budget = {
4293 #[cfg(feature = "metal")]
4294 {
4295 frink_core::metal_dense_enabled()
4296 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4297 }
4298 #[cfg(not(feature = "metal"))]
4299 {
4300 false
4301 }
4302 };
4303 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4304 let budget: u64 = budget_str
4305 .parse()
4306 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4307 match &mut loaded {
4308 model::LoadedModel::Gguf(g) => {
4309 tracing::info!(
4310 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4311 (CUDA and/or Metal matvecs when built with the matching feature)"
4312 );
4313 g.decoder.gpu_vram_budget_bytes = Some(budget);
4314 }
4315 model::LoadedModel::Kimi(_) => {
4316 tracing::warn!(
4317 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4318 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4319 );
4320 }
4321 model::LoadedModel::Mla(_) => {
4322 tracing::warn!(
4323 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4324 FFN path only today; ignoring expert VRAM budget"
4325 );
4326 }
4327 model::LoadedModel::Gemma4(_) => {
4328 tracing::warn!(
4329 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4330 ignoring expert VRAM budget"
4331 );
4332 }
4333 model::LoadedModel::Glm52(_) => {
4334 tracing::warn!(
4335 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4336 GPU expert placement not wired yet; ignoring"
4337 );
4338 }
4339 model::LoadedModel::Encoder(_) => {
4340 tracing::warn!(
4341 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4342 it has no routed experts to place; ignoring"
4343 );
4344 }
4345 }
4346 } else if metal_default_moe_budget {
4347 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4351 if let model::LoadedModel::Gguf(g) = &mut loaded {
4352 tracing::info!(
4353 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4354 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4355 );
4356 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4357 }
4358 }
4359 #[cfg(feature = "cuda")]
4360 {
4361 if frink_core::cuda_dense_enabled() {
4362 tracing::info!(
4363 "CUDA dense matvec enabled for WeightMatrix::apply \
4364 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4365 );
4366 } else {
4367 tracing::info!(
4368 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4369 );
4370 }
4371 }
4372 #[cfg(feature = "metal")]
4373 {
4374 if frink_core::metal_dense_enabled() {
4375 tracing::info!(
4376 "Metal dense matvec enabled for WeightMatrix::apply \
4377 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4378 );
4379 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4380 Some("1") | Some("true") | Some("on") | Some("attn") => {
4381 tracing::info!(
4382 "Metal fused attention requested (FRINK_METAL_ATTN): \
4383 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4384 );
4385 }
4386 _ => {}
4387 }
4388 tracing::info!(
4389 "Metal greedy GPU argmax: temperature<=0 folds \
4390 final_norm+lm_head+argmax into the dense stack"
4391 );
4392 } else {
4393 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4394 }
4395 }
4396 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4407 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4408 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4409 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4410 panic!(
4411 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4412 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4413 );
4414 }
4415 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4416 (Ok(blocks), Ok(block_size), Err(_)) => {
4417 let total_blocks: usize = blocks
4418 .parse()
4419 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4420 let block_size: usize = block_size
4421 .parse()
4422 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4423 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4429 .ok()
4430 .map(|v| {
4431 v.parse()
4432 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4433 })
4434 .unwrap_or(0);
4435 tracing::info!(
4436 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4437 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4438 );
4439 Some(generate::KvPoolConfig {
4440 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4441 queue_wait: Duration::from_millis(queue_wait_ms),
4442 })
4443 }
4444 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4445 let block_size: usize = block_size
4446 .parse()
4447 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4448 let budget: u64 = byte_budget
4449 .parse()
4450 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4451 let cfg = match &loaded {
4452 model::LoadedModel::Gguf(g) => &g.decoder.config,
4453 model::LoadedModel::Kimi(_)
4454 | model::LoadedModel::Mla(_)
4455 | model::LoadedModel::Gemma4(_)
4456 | model::LoadedModel::Glm52(_)
4457 | model::LoadedModel::Encoder(_) => {
4458 panic!(
4459 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4460 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4461 );
4462 }
4463 };
4464 let bytes_per_block = block_size
4465 * cfg.kv_heads_all_layers()
4466 * (cfg.head_dim + cfg.v_head_dim())
4467 * std::mem::size_of::<f32>();
4468 assert!(
4469 bytes_per_block > 0,
4470 "derived KV block byte size must be positive (check model config and block size)"
4471 );
4472 let total_blocks = (budget as usize / bytes_per_block).max(1);
4473 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4474 .ok()
4475 .map(|v| {
4476 v.parse()
4477 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4478 })
4479 .unwrap_or(0);
4480 tracing::info!(
4481 "KV cache block pool enabled from byte budget: {budget} bytes / \
4482 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4483 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4484 cfg.n_layers
4485 );
4486 Some(generate::KvPoolConfig {
4487 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4488 queue_wait: Duration::from_millis(queue_wait_ms),
4489 })
4490 }
4491 (Err(_), Err(_), Err(_)) => None,
4492 (Err(_), Ok(_), Err(_)) => panic!(
4493 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4494 (or unset all three to disable KV cache pooling)"
4495 ),
4496 (Ok(_), Ok(_), Ok(_)) => {
4497 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4498 }
4499 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4500 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4501 set together (or neither, to disable KV cache pooling)"
4502 ),
4503 };
4504 let paged_kv = match (
4510 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4511 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4512 ) {
4513 (Ok(blocks), Ok(block_size)) => {
4514 assert!(
4515 kv_pool.is_none(),
4516 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4517 mutually exclusive: both bound the same KV memory, by different means. \
4518 Set one."
4519 );
4520 let blocks_per_layer: usize = blocks
4534 .parse()
4535 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4536 let block_size: usize = block_size
4537 .parse()
4538 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4539 let gguf = match &loaded {
4540 model::LoadedModel::Gguf(g) => g,
4541 _ => panic!(
4542 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4543 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4544 ),
4545 };
4546 let cfg = &gguf.decoder.config;
4547 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4548 .ok()
4549 .map(|v| {
4550 v.parse()
4551 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4552 })
4553 .unwrap_or(0);
4554 tracing::info!(
4555 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4556 layer across {} layers, shared by all concurrent requests, \
4557 {queue_wait_ms}ms admission queue wait",
4558 cfg.n_layers
4559 );
4560 let radix = Some(Arc::new(Mutex::new(
4565 crate::policy::radix::SaltedRadix::new(block_size),
4566 )));
4567 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4576 crate::policy::parser::ToolCallFormat::infer(
4577 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4578 )
4579 .opener(),
4580 |text| {
4581 gguf.tokenizer
4582 .encode(text, SpecialTokens::Parse)
4583 .into_iter()
4584 .map(|t| t as u32)
4585 .collect()
4586 },
4587 );
4588 if let Some(id) = anchor_token {
4589 tracing::info!(
4590 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4591 window stops short of where its next turn rejoins"
4592 );
4593 }
4594 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4595 .ok()
4596 .map(|v| {
4597 v.parse()
4598 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4599 })
4600 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4601 if let Some(window) = cfg.uniform_sliding_window() {
4602 tracing::info!(
4603 "Paged KV window slide enabled: every layer slides by {window} every \
4604 {slide_interval} decode steps, so a request holds its prompt and a \
4605 window rather than its whole context"
4606 );
4607 } else if cfg.kv_block_window().is_some() {
4608 tracing::info!(
4609 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4610 and a page group holds one block in every layer"
4611 );
4612 }
4613 Some(generate::PagedKvConfig {
4614 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4617 queue_wait: Duration::from_millis(queue_wait_ms),
4618 radix,
4619 anchor_token,
4620 slide_interval,
4621 })
4622 }
4623 (Err(_), Err(_)) => None,
4624 _ => panic!(
4625 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4626 (or neither, to disable paged KV)"
4627 ),
4628 };
4629 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4636 let max_entries: usize = v
4637 .parse()
4638 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4639 if kv_pool.is_some() {
4640 tracing::warn!(
4641 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4642 caching will never be consulted while a KV pool is configured"
4643 );
4644 }
4645 assert!(
4652 paged_kv.is_none(),
4653 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4654 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4655 produce, so the cache could never hit. Set one."
4656 );
4657 tracing::info!(
4658 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4659 all requests"
4660 );
4661 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4662 });
4663 if matches!(
4664 loaded,
4665 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4666 ) && (kv_pool.is_some() || prefix_cache.is_some())
4667 {
4668 tracing::warn!(
4669 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4670 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4671 frink_models::engine's module docs"
4672 );
4673 }
4674 let enable_cb =
4675 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4676 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4677 tracing::info!(
4678 "continuous batching enabled by default on Metal for safe parallel serving \
4679 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4680 );
4681 }
4682 if continuous_batching_env() == Some(true)
4683 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4684 && (kv_pool.is_some() || prefix_cache.is_some())
4685 {
4686 tracing::warn!(
4687 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4688 (those modes keep the private generate path)"
4689 );
4690 }
4691 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4692 if let Ok(chunk) = n.parse::<usize>() {
4693 if chunk > 0 {
4694 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4695 }
4696 }
4697 }
4698 if matches!(
4699 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4700 Some("1")
4701 ) {
4702 tracing::warn!(
4703 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4704 (minimal spill; full layer offload still planned)"
4705 );
4706 }
4707
4708 let mcp = match mcp_config_path {
4709 Some(path) => {
4710 let loaded = mcp::load_mcp_config(&path)?;
4711 tracing::info!(
4712 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4713 loaded.path,
4714 loaded.servers.len()
4715 );
4716 Some(loaded)
4717 }
4718 None => None,
4719 };
4720
4721 let detection = health::Detection::spawn();
4725
4726 let state = Arc::new(build_app_state(
4727 StartupModels {
4728 loaded,
4729 embedding: embedding_model,
4730 },
4731 kv_pool,
4732 paged_kv,
4733 prefix_cache,
4734 enable_cb,
4735 mcp,
4736 detection,
4737 ));
4738
4739 use frink_api::routes;
4743
4744 let public = Router::new().route(routes::HEALTH, get(health));
4749
4750 let mut protected = protected_routes();
4751
4752 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4757 tracing::info!("API key auth enabled");
4758 let auth = limits::AuthConfig {
4759 api_key: Arc::new(key),
4760 };
4761 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4762 auth,
4763 limits::require_api_key,
4764 ));
4765 }
4766 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4767 let rpm: u32 = rpm
4768 .parse()
4769 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4770 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4771 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4772 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4773 limiter,
4774 limits::rate_limit,
4775 ));
4776 }
4777 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4785 let origins = security::parse_cors_origins(&spec)
4786 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4787 tracing::info!(
4788 "CORS enabled: {} allow-listed origin(s) ({})",
4789 origins.len(),
4790 spec
4791 );
4792 let cors = tower_http::cors::CorsLayer::new()
4793 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4794 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4795 .allow_headers([
4796 axum::http::header::CONTENT_TYPE,
4797 axum::http::header::AUTHORIZATION,
4798 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4804 axum::http::HeaderName::from_static("last-event-id"),
4808 ]);
4809 protected = protected.route_layer(cors);
4810 }
4811
4812 let app = public
4816 .merge(protected)
4817 .layer(axum::middleware::from_fn(limits::retry_after))
4818 .with_state(state);
4819
4820 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4826 install_ring_crypto_provider();
4827 match tls_paths {
4832 Some(paths) => {
4833 let config =
4834 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4835 .await
4836 .map_err(|e| {
4837 anyhow::anyhow!(
4838 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4839 paths.cert,
4840 paths.key
4841 )
4842 })?;
4843 let socket_addr: std::net::SocketAddr = addr
4844 .parse()
4845 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4846 let listener = std::net::TcpListener::bind(socket_addr)?;
4847 listener.set_nonblocking(true)?;
4855 let bound = listener.local_addr()?;
4856 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4857 announce_ready(bound, "https");
4858
4859 let handle = axum_server::Handle::new();
4860 let shutdown_handle = handle.clone();
4861 tokio::spawn(async move {
4862 shutdown_signal(exit_on_stdin_close).await;
4863 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4864 });
4865 axum_server::from_tcp_rustls(listener, config)?
4866 .handle(handle)
4867 .serve(app.into_make_service())
4868 .await?;
4869 }
4870 None => {
4871 let listener = tokio::net::TcpListener::bind(&addr).await?;
4872 let bound = listener.local_addr()?;
4873 tracing::info!("frink-server listening on {bound}");
4874 announce_ready(bound, "http");
4875 axum::serve(listener, app)
4876 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4877 .await?;
4878 }
4879 }
4880 Ok(())
4881}
4882
4883#[cfg(test)]
4884pub(crate) mod tests {
4885 use super::*;
4886 use frink_models::config::test_dense_fixture;
4887
4888 #[test]
4889 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4890 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4891 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4892 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4893 assert_eq!(parsed.port, 51999);
4894 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4895 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4897 }
4898
4899 fn test_model() -> Model {
4900 let cfg = test_dense_fixture();
4904 Model::Gguf(GgufModel {
4905 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4906 tokenizer: Arc::new(ServerTokenizer::Byte),
4907 stop_tokens: StopTokens::default(),
4908 bos_id: None,
4909 is_synthetic: true,
4910 chat_template: chat_template::PromptTemplate::plain(),
4911 })
4912 }
4913
4914 fn greedy_params(max_tokens: usize) -> GenerationParams {
4915 GenerationParams {
4916 cache_salt: None,
4917 prompt_logprobs: None,
4918 wants_logprobs: false,
4919 n: 1,
4920 interleave_choices: false,
4921 keep_special_tokens: false,
4922 truncate_prompt_tokens: None,
4923 token_mask: crate::token_mask::TokenMask::default(),
4924 reasoning: None,
4925 max_tokens,
4926 sampling: SamplingParams::default(),
4927 seed: 1,
4928 stop: Vec::new(),
4929 stop_token_ids: Vec::new(),
4930 json_object: false,
4931 grammar: None,
4932 cancel: None,
4933 ignore_eos: false,
4934 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4935 lora: None,
4936 }
4937 }
4938
4939 fn test_model_full_byte_vocab() -> Model {
4943 test_model_full_byte_vocab_with_eos(None)
4944 }
4945
4946 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4953 test_byte_model(eos, true)
4954 }
4955
4956 fn test_byte_model(eos: Option<usize>, synthetic: bool) -> Model {
4963 let mut cfg = test_dense_fixture();
4964 cfg.vocab_size = 256;
4965 Model::Gguf(GgufModel {
4966 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4967 tokenizer: Arc::new(ServerTokenizer::Byte),
4968 stop_tokens: StopTokens::from_eos(eos),
4969 bos_id: None,
4970 is_synthetic: synthetic,
4971 chat_template: chat_template::PromptTemplate::plain(),
4972 })
4973 }
4974
4975 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4979 test_state_at(model, response_cache, None)
4980 }
4981
4982 pub(crate) fn test_state_at(
4986 model: Model,
4987 response_cache: ResponseCache,
4988 checkpoint_path: Option<std::path::PathBuf>,
4989 ) -> AppState {
4990 AppState {
4991 slept: Mutex::new(None),
4992 embedding: None,
4993 paged_kv: None,
4994 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4995 id: None,
4996 loaded: Loaded::Generative(Arc::new(model)),
4997 batcher: None,
4998 ceiling: None,
4999 checkpoint_path,
5000 }))),
5001 load_in_progress: std::sync::atomic::AtomicBool::new(false),
5002 tasks: Arc::new(tasks::TaskRegistry::new()),
5003 cancels: Arc::new(cancel::CancelRegistry::new()),
5004 stats: stats::Stats::new(),
5005 streams: resume::StreamRegistry::new(),
5006 model_dir: None,
5007 response_cache: Mutex::new(response_cache),
5008 kv_pool: None,
5009 prefix_cache: None,
5010 sessions: session::SessionStore::new(),
5011 requests_total: std::sync::atomic::AtomicU64::new(0),
5012 request_errors_total: std::sync::atomic::AtomicU64::new(0),
5013 started_at: std::time::Instant::now(),
5014 last_request_ms: std::sync::atomic::AtomicU64::new(0),
5015 detection: Arc::new(health::Detection::ready(health::probe_backends())),
5016 mcp: None,
5017 continuous_batching_enabled: false,
5018 metal_private_decode_gate: None,
5019 loading_model: Mutex::new(None),
5020 last_load_error: Mutex::new(None),
5021 serving: Mutex::new(crate::stats::ServingStats::default()),
5022 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
5023 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
5024 started_unix: unix_now(),
5025 }
5026 }
5027
5028 pub(crate) fn test_app() -> Router {
5037 test_app_with_state(Arc::new(test_state(
5038 test_model_full_byte_vocab(),
5039 ResponseCache::new(1000, Duration::from_secs(3600)),
5040 )))
5041 }
5042
5043 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
5046 Router::new()
5056 .route(frink_api::routes::HEALTH, get(health))
5057 .merge(protected_routes())
5058 .with_state(state)
5059 }
5060
5061 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
5062 let mut cfg = test_dense_fixture();
5063 cfg.name = name;
5064 cfg.vocab_size = vocab_size;
5065 Model::Gguf(GgufModel {
5066 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5067 tokenizer: Arc::new(ServerTokenizer::Byte),
5068 stop_tokens: StopTokens::default(),
5069 bos_id: None,
5070 is_synthetic: true,
5071 chat_template: chat_template::PromptTemplate::plain(),
5072 })
5073 }
5074
5075 fn model_with_template(name: &'static str, source: &str) -> Model {
5079 let mut cfg = test_dense_fixture();
5080 cfg.name = name;
5081 cfg.vocab_size = 256;
5082 Model::Gguf(GgufModel {
5083 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5084 tokenizer: Arc::new(ServerTokenizer::Byte),
5085 stop_tokens: StopTokens::default(),
5086 bos_id: None,
5087 is_synthetic: true,
5088 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5089 Some(source),
5090 Some("qwen3"),
5091 false,
5092 true,
5093 None,
5094 None,
5095 ),
5096 })
5097 }
5098
5099 #[tokio::test]
5108 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5109 let strict = "{% if messages | length > 1 %}\
5112 {{ raise_exception('this template takes one turn') }}\
5113 {% endif %}{{ messages[0].content }}";
5114 let state = Arc::new(test_state(
5115 model_with_template("strict", strict),
5116 ResponseCache::new(4, Duration::from_secs(60)),
5117 ));
5118 let app = test_app_with_state(state);
5119
5120 let (status, body) = post_json_uri(
5121 &app,
5122 "/v1/chat/completions",
5123 serde_json::json!({
5124 "model": "strict",
5125 "stream": true,
5126 "messages": [
5127 {"role": "user", "content": "one"},
5128 {"role": "user", "content": "two"},
5129 ],
5130 }),
5131 )
5132 .await;
5133 assert_eq!(status, StatusCode::BAD_REQUEST);
5134 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5135 assert!(
5136 body["error"]["message"]
5137 .as_str()
5138 .unwrap()
5139 .contains("one turn"),
5140 "the template's own message must reach the caller: {body}"
5141 );
5142
5143 let (status, _) = post_json_uri(
5145 &app,
5146 "/v1/chat/completions",
5147 serde_json::json!({
5148 "model": "strict",
5149 "stream": true,
5150 "max_tokens": 1,
5151 "messages": [{"role": "user", "content": "one"}],
5152 }),
5153 )
5154 .await;
5155 assert_eq!(status, StatusCode::OK);
5156 }
5157
5158 #[tokio::test]
5160 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5161 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5162 {% if reasoning_effort %}\
5163 {% if reasoning_effort not in ['low','medium','high'] %}\
5164 {{ raise_exception('bad effort') }}\
5165 {% endif %}[{{ reasoning_effort }}]\
5166 {% endif %}{{ messages[0].content }}";
5167 let state = Arc::new(test_state(
5168 model_with_template("thinker", reasoning),
5169 ResponseCache::new(4, Duration::from_secs(60)),
5170 ));
5171 let app = test_app_with_state(state);
5172 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5173 assert_eq!(status, StatusCode::OK);
5174 let entry = &models["data"][0];
5175 assert_eq!(
5176 entry["supported_reasoning_efforts"],
5177 serde_json::json!(["off", "low", "medium", "high"])
5178 );
5179 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5180 }
5181
5182 #[tokio::test]
5187 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5188 let app = test_app();
5189 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5190 let entry = &models["data"][0];
5191 assert!(entry.get("supported_reasoning_efforts").is_none());
5192 assert!(entry.get("default_reasoning_effort").is_none());
5193 }
5194
5195 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5196 Arc::new(ActiveModel {
5197 id: Some(name.to_string()),
5198 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5199 batcher: None,
5200 ceiling: None,
5201 checkpoint_path: None,
5202 })
5203 .tap_into(state)
5204 }
5205
5206 trait TapInto {
5208 fn tap_into(self, state: &AppState) -> Self;
5209 }
5210 impl TapInto for Arc<ActiveModel> {
5211 fn tap_into(self, state: &AppState) -> Self {
5212 state.swap_active(Some(Arc::clone(&self)));
5213 self
5214 }
5215 }
5216
5217 #[test]
5223 fn an_in_flight_request_keeps_the_model_it_started_on() {
5224 let state = test_state(
5225 named_test_model("model-a", 256),
5226 ResponseCache::new(4, Duration::from_secs(60)),
5227 );
5228
5229 let in_flight = state.active().expect("a model is loaded");
5232 assert_eq!(in_flight.name(), "model-a");
5233
5234 active_model(&state, "model-b");
5235
5236 assert_eq!(state.active().unwrap().name(), "model-b");
5238 assert_eq!(in_flight.name(), "model-a");
5240 let produced = run_generation(
5241 in_flight.generative().unwrap(),
5242 "hi",
5243 &greedy_params(3),
5244 None,
5245 None,
5246 None,
5247 None,
5248 None,
5249 None,
5250 )
5251 .expect("the old model must still decode after being swapped out");
5252 assert!(matches!(
5253 produced.choices[0].finish,
5254 FinishReason::Length | FinishReason::Stop
5255 ));
5256 }
5257
5258 #[test]
5263 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5264 let state = test_state(
5265 named_test_model("model-a", 256),
5266 ResponseCache::new(4, Duration::from_secs(60)),
5267 );
5268 let in_flight = state.active().expect("a model is loaded");
5269 let weights = Arc::clone(in_flight.generative().unwrap());
5270 assert!(Arc::strong_count(&weights) >= 2);
5271
5272 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5273 id: Some("model-b".to_string()),
5274 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5275 batcher: None,
5276 ceiling: None,
5277 checkpoint_path: None,
5278 })));
5279 drop(previous);
5280 assert!(Arc::strong_count(&weights) >= 2);
5282 drop(in_flight);
5283 assert_eq!(Arc::strong_count(&weights), 1);
5284 }
5285
5286 #[tokio::test]
5290 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5291 let state = Arc::new(test_state(
5292 named_test_model("model-a", 256),
5293 ResponseCache::new(4, Duration::from_secs(60)),
5294 ));
5295 let app = test_app_with_state(Arc::clone(&state));
5296
5297 let (status, body) = post_json_uri(
5298 &app,
5299 frink_api::routes::ADMIN_MODELS_UNLOAD,
5300 serde_json::json!({}),
5301 )
5302 .await;
5303 assert_eq!(status, StatusCode::OK);
5304 assert_eq!(body["ok"], true);
5305 assert!(body["active"].is_null());
5306 assert!(state.active().is_none());
5307
5308 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5309 assert_eq!(status, StatusCode::OK);
5310 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5311 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5312
5313 let (status, body) = post_json_uri(
5314 &app,
5315 "/v1/chat/completions",
5316 serde_json::json!({
5317 "model": "x",
5318 "messages": [{"role": "user", "content": "hi"}]
5319 }),
5320 )
5321 .await;
5322 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5323 assert_eq!(body["error"]["type"], "model_not_loaded");
5324 }
5325
5326 #[tokio::test]
5330 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5331 let state = Arc::new(test_state(
5332 named_test_model("model-a", 256),
5333 ResponseCache::new(4, Duration::from_secs(60)),
5334 ));
5335 let app = test_app_with_state(Arc::clone(&state));
5336 state.swap_active(None);
5337
5338 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5339 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5342 assert_eq!(body["state"], "unavailable");
5343 assert_eq!(body["reason"], "model_not_loaded");
5344 assert!(body["model"].is_null());
5345 let real_weights = body["capabilities"]
5346 .as_array()
5347 .unwrap()
5348 .iter()
5349 .find(|c| c["id"] == "real_weights")
5350 .cloned()
5351 .expect("real_weights is always reported");
5352 assert_eq!(real_weights["available"], false);
5353 assert_eq!(real_weights["reason"], "model_not_loaded");
5354 }
5355
5356 #[tokio::test]
5360 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5361 let app = test_app();
5362
5363 let (status, completion) = post_json_uri(
5364 &app,
5365 "/v1/chat/completions",
5366 serde_json::json!({
5367 "model": "x",
5368 "messages": [{"role": "user", "content": "hi"}],
5369 "max_tokens": 4
5370 }),
5371 )
5372 .await;
5373 assert_eq!(status, StatusCode::OK);
5374 let request_id = completion["request_id"].as_str().unwrap().to_string();
5375
5376 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5377 assert_eq!(status, StatusCode::OK);
5378 let recent = stats["recent"].as_array().unwrap();
5379 assert_eq!(recent.len(), 1);
5380 let row = &recent[0];
5381 assert_eq!(row["request_id"], request_id);
5382 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5383 assert_eq!(row["status"], 200);
5384 assert_eq!(row["stream"], false);
5385 assert!(row["duration_ms"].is_number());
5388 assert!(row["decode_ms"].is_number());
5389 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5390 assert_eq!(
5391 stats["tokens_prompt_total"].as_u64().unwrap(),
5392 row["prompt_tokens"].as_u64().unwrap()
5393 );
5394 }
5395
5396 #[tokio::test]
5400 async fn a_rejected_request_is_recorded_too() {
5401 let state = Arc::new(test_state(
5402 named_test_model("model-a", 256),
5403 ResponseCache::new(4, Duration::from_secs(60)),
5404 ));
5405 let app = test_app_with_state(Arc::clone(&state));
5406 state.swap_active(None);
5407
5408 let (status, _) = post_json_uri(
5409 &app,
5410 "/v1/chat/completions",
5411 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5412 )
5413 .await;
5414 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5415
5416 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5417 let recent = stats["recent"].as_array().unwrap();
5418 assert_eq!(recent.len(), 1);
5419 assert_eq!(recent[0]["status"], 503);
5420 assert_eq!(recent[0]["completion_tokens"], 0);
5421 assert!(recent[0]["decode_ms"].is_null());
5422 assert_eq!(stats["errors_total"], 1);
5423 }
5424
5425 async fn post_json_with_headers(
5429 app: &Router,
5430 uri: &str,
5431 body: serde_json::Value,
5432 headers: &[(&str, &str)],
5433 ) -> (StatusCode, serde_json::Value) {
5434 use http_body_util::BodyExt;
5435 use tower::ServiceExt;
5436
5437 let mut builder = axum::http::Request::builder()
5438 .method("POST")
5439 .uri(uri)
5440 .header("content-type", "application/json");
5441 for (name, value) in headers {
5442 builder = builder.header(*name, *value);
5443 }
5444 let response = app
5445 .clone()
5446 .oneshot(
5447 builder
5448 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5449 .unwrap(),
5450 )
5451 .await
5452 .unwrap();
5453 let status = response.status();
5454 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5455 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5456 (status, json)
5457 }
5458
5459 #[tokio::test]
5463 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5464 let app = test_app();
5465
5466 let (status, _) = post_json_uri(
5467 &app,
5468 frink_api::routes::V1_TOKENIZE,
5469 serde_json::json!({"prompt": "hello"}),
5470 )
5471 .await;
5472 assert_eq!(status, StatusCode::OK);
5473 let (status, _) = post_json_uri(
5474 &app,
5475 frink_api::routes::V1_DETOKENIZE,
5476 serde_json::json!({"tokens": [104, 105]}),
5477 )
5478 .await;
5479 assert_eq!(status, StatusCode::OK);
5480 let (status, _) = post_json_uri(
5481 &app,
5482 frink_api::routes::V1_EMBEDDINGS,
5483 serde_json::json!({"input": "hello"}),
5484 )
5485 .await;
5486 assert_eq!(status, StatusCode::OK);
5487
5488 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5489 let routes: Vec<&str> = stats["recent"]
5490 .as_array()
5491 .unwrap()
5492 .iter()
5493 .map(|row| row["route"].as_str().unwrap())
5494 .collect();
5495 for expected in [
5496 frink_api::routes::V1_TOKENIZE,
5497 frink_api::routes::V1_DETOKENIZE,
5498 frink_api::routes::V1_EMBEDDINGS,
5499 ] {
5500 assert!(
5501 routes.contains(&expected),
5502 "{expected} is missing: {routes:?}"
5503 );
5504 }
5505
5506 let row = |route: &str| {
5507 stats["recent"]
5508 .as_array()
5509 .unwrap()
5510 .iter()
5511 .find(|r| r["route"] == route)
5512 .cloned()
5513 .unwrap()
5514 };
5515 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5519 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5520 assert!(embed["decode_ms"].is_null());
5521 assert_eq!(embed["completion_tokens"], 0);
5522 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5526 assert_eq!(
5527 stats["tokens_prompt_total"].as_u64().unwrap(),
5528 embed["prompt_tokens"].as_u64().unwrap(),
5529 "only the forward pass counted"
5530 );
5531 }
5532
5533 fn streaming_test_app() -> Router {
5538 let mut cfg = test_dense_fixture();
5539 cfg.vocab_size = 256;
5540 let model = Model::Gguf(GgufModel {
5541 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5542 tokenizer: Arc::new(ServerTokenizer::Byte),
5543 stop_tokens: StopTokens::default(),
5544 bos_id: None,
5545 is_synthetic: false,
5546 chat_template: chat_template::PromptTemplate::plain(),
5547 });
5548 test_app_with_state(Arc::new(test_state(
5549 model,
5550 ResponseCache::new(1000, Duration::from_secs(3600)),
5551 )))
5552 }
5553
5554 #[tokio::test]
5563 async fn chat_logprobs_are_rendered_and_are_never_served_from_cache() {
5564 let app = test_app();
5565 let body = |logprobs: Option<(bool, Option<u32>)>| {
5566 let mut b = serde_json::json!({
5567 "model": "x",
5568 "messages": [{"role": "user", "content": "hi"}],
5569 "max_tokens": 4
5570 });
5571 if let Some((on, top)) = logprobs {
5572 b["logprobs"] = serde_json::json!(on);
5573 if let Some(n) = top {
5574 b["top_logprobs"] = serde_json::json!(n);
5575 }
5576 }
5577 b
5578 };
5579
5580 let (status, plain) =
5582 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(None)).await;
5583 assert_eq!(status, StatusCode::OK, "{plain}");
5584 assert!(plain["choices"][0]["logprobs"].is_null(), "{plain}");
5585
5586 for attempt in 0..2 {
5590 let (status, with) = post_json_uri(
5591 &app,
5592 frink_api::routes::V1_CHAT_COMPLETIONS,
5593 body(Some((true, Some(2)))),
5594 )
5595 .await;
5596 assert_eq!(status, StatusCode::OK, "{with}");
5597 assert_ne!(
5598 with["frink_cache"], "hit",
5599 "attempt {attempt} replayed a cached answer for a logprobs request: {with}"
5600 );
5601 let lp = &with["choices"][0]["logprobs"];
5602 assert!(lp.is_object(), "attempt {attempt}: {with}");
5603 let content = lp["content"].as_array().expect("content");
5604 assert!(lp["tokens"].is_null(), "completions shape leaked: {lp}");
5606 for entry in content {
5607 assert!(entry["token"].is_string(), "{entry}");
5608 assert!(entry["bytes"].is_array(), "{entry}");
5609 let v = entry["logprob"].as_f64().expect("a real number");
5610 assert!(v <= 0.0 && v.is_finite(), "{entry}");
5611 let top = entry["top_logprobs"].as_array().expect("top_logprobs");
5612 assert!(top.len() <= 2, "asked for 2, got {}", top.len());
5613 }
5614 }
5615 }
5616
5617 #[tokio::test]
5623 async fn the_chat_logprobs_pair_is_validated() {
5624 let app = test_app();
5625 for (extra, why) in [
5626 (serde_json::json!({"top_logprobs": 3}), "without logprobs"),
5627 (
5628 serde_json::json!({"logprobs": true, "top_logprobs": 21}),
5629 "above the cap",
5630 ),
5631 ] {
5632 let mut body = serde_json::json!({
5633 "model": "x",
5634 "messages": [{"role": "user", "content": "hi"}],
5635 "max_tokens": 2
5636 });
5637 for (k, v) in extra.as_object().unwrap() {
5638 body[k] = v.clone();
5639 }
5640 let (status, answer) =
5641 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await;
5642 assert_eq!(status, StatusCode::BAD_REQUEST, "{why}: {answer}");
5643 assert!(
5644 answer["error"]["message"]
5645 .as_str()
5646 .is_some_and(|m| m.contains("top_logprobs")),
5647 "{why}: {answer}"
5648 );
5649 }
5650 }
5651
5652 #[tokio::test]
5661 async fn sleep_refuses_a_model_it_could_not_bring_back() {
5662 let app = test_app();
5663 let (status, answer) =
5664 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5665 assert_eq!(status, StatusCode::CONFLICT, "{answer}");
5666 assert_eq!(answer["error"]["type"], "not_reloadable", "{answer}");
5667 let (_, still) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5670 assert_eq!(still["is_sleeping"], false, "{still}");
5671 let (status, _) = post_json_uri(
5672 &app,
5673 frink_api::routes::V1_CHAT_COMPLETIONS,
5674 serde_json::json!({
5675 "model": "x",
5676 "messages": [{"role": "user", "content": "hi"}],
5677 "max_tokens": 2
5678 }),
5679 )
5680 .await;
5681 assert_eq!(status, StatusCode::OK, "a refused sleep unloaded the model");
5682 }
5683
5684 #[tokio::test]
5693 async fn sleep_remembers_what_unload_forgets() {
5694 let state = Arc::new(test_state_at(
5697 test_model_full_byte_vocab(),
5698 ResponseCache::new(1000, Duration::from_secs(3600)),
5699 Some(std::path::PathBuf::from("/nonexistent/fixture.gguf")),
5700 ));
5701 let app = test_app_with_state(Arc::clone(&state));
5702 let ask = || {
5703 let app = app.clone();
5704 async move {
5705 post_json_uri(
5706 &app,
5707 frink_api::routes::V1_CHAT_COMPLETIONS,
5708 serde_json::json!({
5709 "model": "x",
5710 "messages": [{"role": "user", "content": "hi"}],
5711 "max_tokens": 2
5712 }),
5713 )
5714 .await
5715 }
5716 };
5717
5718 let (status, _) = ask().await;
5719 assert_eq!(status, StatusCode::OK, "the fixture server serves");
5720 let (_, awake) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5721 assert_eq!(awake["is_sleeping"], false, "{awake}");
5722
5723 let (status, slept) =
5724 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5725 assert_eq!(status, StatusCode::OK, "{slept}");
5726 assert_eq!(slept["is_sleeping"], true, "{slept}");
5727 let (_, now) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5728 assert_eq!(now["is_sleeping"], true, "{now}");
5729
5730 let (status, refused) = ask().await;
5733 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE, "{refused}");
5734 assert_eq!(
5735 refused["error"]["type"], "server_sleeping",
5736 "an asleep server reported itself as empty: {refused}"
5737 );
5738
5739 let (status, again) =
5741 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5742 assert_eq!(status, StatusCode::OK, "{again}");
5743 assert_eq!(again["is_sleeping"], true, "{again}");
5744 }
5745
5746 #[tokio::test]
5750 async fn waking_a_server_that_is_awake_is_refused() {
5751 let app = test_app();
5752 let (status, answer) =
5753 post_json_uri(&app, frink_api::routes::WAKE_UP, serde_json::json!({})).await;
5754 assert_eq!(status, StatusCode::CONFLICT, "{answer}");
5755 assert_eq!(answer["error"]["type"], "not_sleeping", "{answer}");
5756 }
5757
5758 #[tokio::test]
5765 async fn a_salt_keeps_one_callers_cached_answer_from_another() {
5766 let app = test_app();
5767 let body = |salt: Option<&str>| {
5768 let mut b = serde_json::json!({
5769 "model": "x",
5770 "messages": [{"role": "user", "content": "the same prompt"}],
5771 "max_tokens": 4,
5772 "seed": 1
5773 });
5774 if let Some(s) = salt {
5775 b["cache_salt"] = serde_json::json!(s);
5776 }
5777 b
5778 };
5779 let post = |b: serde_json::Value| {
5780 let app = app.clone();
5781 async move { post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, b).await }
5782 };
5783
5784 let (status, _) = post(body(Some("tenant-a"))).await;
5786 assert_eq!(status, StatusCode::OK);
5787 let (_, again) = post(body(Some("tenant-a"))).await;
5788 assert_eq!(
5789 again["frink_cache"], "hit",
5790 "the owner did not get its own entry back: {again}"
5791 );
5792
5793 let (_, other) = post(body(Some("tenant-b"))).await;
5795 assert_ne!(
5796 other["frink_cache"], "hit",
5797 "a different caller was served tenant-a's answer: {other}"
5798 );
5799
5800 let (_, shared) = post(body(None)).await;
5802 assert_ne!(
5803 shared["frink_cache"], "hit",
5804 "an unsalted request was served a salted answer: {shared}"
5805 );
5806 }
5807
5808 #[tokio::test]
5811 async fn chat_serves_several_choices_from_one_prefill() {
5812 let app = test_app();
5813 let body = |n: u32, stream: bool| {
5814 serde_json::json!({
5815 "model": "x",
5816 "messages": [{"role": "user", "content": "hi"}],
5817 "max_tokens": 4,
5818 "temperature": 1.0,
5819 "n": n,
5820 "stream": stream
5821 })
5822 };
5823
5824 let (status, one) =
5825 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5826 assert_eq!(status, StatusCode::OK, "{one}");
5827
5828 let (status, three) =
5829 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5830 assert_eq!(status, StatusCode::OK, "{three}");
5831 let choices = three["choices"].as_array().expect("an array");
5832 assert_eq!(choices.len(), 3, "{three}");
5833 for (i, c) in choices.iter().enumerate() {
5834 assert_eq!(c["index"], i);
5835 assert!(c["message"]["role"].is_string(), "{c}");
5836 assert!(c["finish_reason"].is_string(), "{c}");
5837 }
5838 assert_eq!(
5840 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5841 "n = 3 billed the prompt more than once"
5842 );
5843 }
5844
5845 #[tokio::test]
5857 async fn a_streamed_n_interleaves_its_choices() {
5858 let app = streaming_test_app();
5859 let raw = post_sse_raw(
5860 &app,
5861 serde_json::json!({
5862 "model": "x",
5863 "messages": [{"role": "user", "content": "hi"}],
5864 "max_tokens": 6,
5865 "temperature": 1.0,
5866 "n": 3,
5867 "stream": true
5868 }),
5869 )
5870 .await;
5871
5872 let mut order: Vec<usize> = Vec::new();
5874 let mut finished: Vec<usize> = Vec::new();
5875 let mut usage_blocks = 0usize;
5876 for line in raw.lines() {
5877 let Some(rest) = line.strip_prefix("data: ") else {
5878 continue;
5879 };
5880 if rest.trim() == "[DONE]" {
5881 continue;
5882 }
5883 let v: serde_json::Value = serde_json::from_str(rest).expect(rest);
5884 if v.get("usage").is_some_and(|u| !u.is_null()) {
5885 usage_blocks += 1;
5886 }
5887 let Some(choice) = v["choices"].as_array().and_then(|c| c.first()) else {
5888 continue;
5889 };
5890 let index = choice["index"].as_u64().expect("an index") as usize;
5891 if choice["finish_reason"].is_string() {
5892 finished.push(index);
5893 continue;
5894 }
5895 order.push(index);
5896 }
5897
5898 assert_eq!(
5899 finished,
5900 vec![0, 1, 2],
5901 "one terminal chunk per choice, in index order: {raw}"
5902 );
5903 assert_eq!(usage_blocks, 1, "the usage block is the request's: {raw}");
5904 assert!(
5905 order.contains(&0) && order.contains(&2),
5906 "not every choice streamed: {order:?}"
5907 );
5908 let last_of_zero = order
5909 .iter()
5910 .rposition(|i| *i == 0)
5911 .expect("choice 0 streamed");
5912 let first_of_two = order
5913 .iter()
5914 .position(|i| *i == 2)
5915 .expect("choice 2 streamed");
5916 assert!(
5917 first_of_two < last_of_zero,
5918 "the choices arrived one after another rather than interleaved: {order:?}"
5919 );
5920 }
5921
5922 #[tokio::test]
5934 async fn skip_special_tokens_false_keeps_the_end_marker() {
5935 let probe = test_app_with_state(Arc::new(test_state(
5942 test_byte_model(None, false),
5943 ResponseCache::new(1000, Duration::from_secs(3600)),
5944 )));
5945 let (_, seen) = post_json_uri(
5946 &probe,
5947 frink_api::routes::V1_COMPLETIONS,
5948 serde_json::json!({
5949 "model": "x",
5950 "prompt": "\u{1}\u{2}",
5951 "max_tokens": 6,
5952 "temperature": 0,
5953 "logprobs": 1,
5954 "return_tokens_as_token_ids": true
5955 }),
5956 )
5957 .await;
5958 let eos: usize = seen["choices"][0]["logprobs"]["tokens"][1]
5959 .as_str()
5960 .and_then(|s| s.strip_prefix("token_id:"))
5961 .and_then(|s| s.parse().ok())
5962 .expect("a second generated token");
5963
5964 let app = test_app_with_state(Arc::new(test_state(
5965 test_byte_model(Some(eos), false),
5966 ResponseCache::new(1000, Duration::from_secs(3600)),
5967 )));
5968 let ask = |skip: bool| {
5969 serde_json::json!({
5970 "model": "x",
5971 "prompt": "\u{1}\u{2}",
5972 "max_tokens": 6,
5973 "temperature": 0,
5974 "skip_special_tokens": skip
5975 })
5976 };
5977
5978 let (status, kept) =
5979 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(false)).await;
5980 assert_eq!(status, StatusCode::OK, "{kept}");
5981 let (status, skipped) =
5982 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(true)).await;
5983 assert_eq!(status, StatusCode::OK, "{skipped}");
5984
5985 assert_eq!(
5988 kept["choices"][0]["finish_reason"], "stop",
5989 "the model did not end the turn: {kept}"
5990 );
5991 assert_eq!(
5992 skipped["choices"][0]["finish_reason"], "stop",
5993 "keeping the marker must not change WHEN it stops: {skipped}"
5994 );
5995
5996 let with = kept["choices"][0]["text"].as_str().expect("text");
5997 let without = skipped["choices"][0]["text"].as_str().expect("text");
5998 let marker = char::from(eos as u8);
6000 assert!(
6001 with.ends_with(marker),
6002 "the end marker was dropped: {with:?}"
6003 );
6004 assert!(
6005 !without.ends_with(marker),
6006 "the default must still skip it: {without:?}"
6007 );
6008 assert_eq!(
6009 with.len(),
6010 without.len() + marker.len_utf8(),
6011 "the two answers differ by more than the marker"
6012 );
6013 assert_eq!(
6016 kept["usage"]["completion_tokens"].as_u64().unwrap(),
6017 skipped["usage"]["completion_tokens"].as_u64().unwrap() + 1,
6018 "the kept marker was not counted"
6019 );
6020 }
6021
6022 #[tokio::test]
6028 async fn return_tokens_as_token_ids_renames_reported_tokens_only() {
6029 let app = streaming_test_app();
6030 let ask = |as_ids: bool| {
6031 serde_json::json!({
6032 "model": "x",
6033 "prompt": "hi",
6034 "max_tokens": 4,
6035 "temperature": 0,
6036 "logprobs": 2,
6037 "return_tokens_as_token_ids": as_ids
6038 })
6039 };
6040
6041 let (status, plain) =
6042 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(false)).await;
6043 assert_eq!(status, StatusCode::OK, "{plain}");
6044 let (status, by_id) =
6045 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(true)).await;
6046 assert_eq!(status, StatusCode::OK, "{by_id}");
6047
6048 let tokens = by_id["choices"][0]["logprobs"]["tokens"]
6049 .as_array()
6050 .expect("tokens");
6051 assert!(!tokens.is_empty(), "nothing was reported: {by_id}");
6052 for t in tokens {
6053 let s = t.as_str().expect("a piece");
6054 assert!(
6055 s.starts_with("token_id:") && s["token_id:".len()..].parse::<usize>().is_ok(),
6056 "reported as text rather than by id: {s:?}"
6057 );
6058 }
6059 let top = &by_id["choices"][0]["logprobs"]["top_logprobs"][0];
6063 for key in top.as_object().expect("a map").keys() {
6064 assert!(key.starts_with("token_id:"), "{key:?}");
6065 }
6066 assert_eq!(
6068 by_id["choices"][0]["text"], plain["choices"][0]["text"],
6069 "the completion's text changed, which the field does not do"
6070 );
6071 assert!(
6072 !plain["choices"][0]["logprobs"]["tokens"][0]
6073 .as_str()
6074 .unwrap_or_default()
6075 .starts_with("token_id:"),
6076 "the default already reported ids, so this proved nothing"
6077 );
6078 }
6079
6080 #[tokio::test]
6089 async fn echo_returns_the_prompt_with_offsets_that_index_it() {
6090 let app = streaming_test_app();
6091 let prompt = "hello";
6092 let (status, body) = post_json_uri(
6093 &app,
6094 frink_api::routes::V1_COMPLETIONS,
6095 serde_json::json!({
6096 "model": "x",
6097 "prompt": prompt,
6098 "max_tokens": 6,
6099 "temperature": 0,
6100 "echo": true,
6101 "logprobs": 2
6102 }),
6103 )
6104 .await;
6105 assert_eq!(status, StatusCode::OK, "{body}");
6106
6107 let text = body["choices"][0]["text"].as_str().expect("text");
6108 assert!(
6109 text.starts_with(prompt),
6110 "the prompt was not echoed: {text:?}"
6111 );
6112 assert!(
6113 text.len() > prompt.len(),
6114 "nothing was generated after the echo: {text:?}"
6115 );
6116
6117 let lp = &body["choices"][0]["logprobs"];
6118 let tokens = lp["tokens"].as_array().expect("tokens");
6119 let offsets = lp["text_offset"].as_array().expect("text_offset");
6120 let scores = lp["token_logprobs"].as_array().expect("token_logprobs");
6121 assert_eq!(tokens.len(), offsets.len());
6122 assert_eq!(tokens.len(), scores.len());
6123 assert!(
6124 tokens.len() > 6,
6125 "the arrays cover only the completion: {}",
6126 tokens.len()
6127 );
6128 assert!(scores[0].is_null(), "{lp}");
6130 for (i, (tok, off)) in tokens.iter().zip(offsets).enumerate() {
6132 let (piece, at) = (
6133 tok.as_str().expect("a piece"),
6134 off.as_u64().unwrap() as usize,
6135 );
6136 assert!(
6137 text[at..].starts_with(piece),
6138 "entry {i}: offset {at} does not start {piece:?} in {text:?}"
6139 );
6140 }
6141 }
6142
6143 #[tokio::test]
6153 async fn truncate_prompt_tokens_keeps_the_last_k_and_echo_reports_them() {
6154 let app = streaming_test_app();
6155 let prompt = "abcdefghij";
6156 let ask = |k: Option<u32>| {
6157 let mut b = serde_json::json!({
6158 "model": "x",
6159 "prompt": prompt,
6160 "max_tokens": 2,
6161 "temperature": 0,
6162 "echo": true
6163 });
6164 if let Some(k) = k {
6165 b["truncate_prompt_tokens"] = serde_json::json!(k);
6166 }
6167 b
6168 };
6169
6170 let (status, full) =
6171 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(None)).await;
6172 assert_eq!(status, StatusCode::OK, "{full}");
6173 let full_prompt_tokens = full["usage"]["prompt_tokens"].as_u64().expect("usage");
6174 assert!(full_prompt_tokens > 4, "the prompt is too short to cut");
6175
6176 let (status, cut) =
6177 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(Some(4))).await;
6178 assert_eq!(status, StatusCode::OK, "{cut}");
6179 assert_eq!(
6180 cut["usage"]["prompt_tokens"].as_u64(),
6181 Some(4),
6182 "the prompt was not truncated: {cut}"
6183 );
6184 let text = cut["choices"][0]["text"].as_str().expect("text");
6186 assert!(
6187 text.starts_with("ghij"),
6188 "echo reported a prompt the model never saw: {text:?}"
6189 );
6190 assert!(
6191 !text.starts_with(prompt),
6192 "the full prompt was echoed after a truncation: {text:?}"
6193 );
6194 }
6195
6196 #[tokio::test]
6200 async fn a_truncation_below_one_is_a_bad_request() {
6201 let app = streaming_test_app();
6202 for k in [0i64, -1] {
6203 let (status, body) = post_json_uri(
6204 &app,
6205 frink_api::routes::V1_COMPLETIONS,
6206 serde_json::json!({
6207 "model": "x",
6208 "prompt": "hi",
6209 "max_tokens": 2,
6210 "truncate_prompt_tokens": k
6211 }),
6212 )
6213 .await;
6214 assert_eq!(status, StatusCode::BAD_REQUEST, "k = {k}: {body}");
6215 }
6216 }
6217
6218 #[tokio::test]
6226 async fn allowed_token_ids_restricts_the_draw() {
6227 let app = streaming_test_app();
6228 let body = |allowed: Option<serde_json::Value>| {
6229 let mut b = serde_json::json!({
6230 "model": "x",
6231 "prompt": "hi",
6232 "max_tokens": 16,
6233 "temperature": 1.0,
6234 "seed": 3
6235 });
6236 if let Some(ids) = allowed {
6237 b["allowed_token_ids"] = ids;
6238 }
6239 b
6240 };
6241
6242 let (status, free) =
6245 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, body(None)).await;
6246 assert_eq!(status, StatusCode::OK, "{free}");
6247 let free_text = free["choices"][0]["text"].as_str().unwrap_or_default();
6248
6249 let (status, restricted) = post_json_uri(
6250 &app,
6251 frink_api::routes::V1_COMPLETIONS,
6252 body(Some(serde_json::json!([65, 66]))),
6254 )
6255 .await;
6256 assert_eq!(status, StatusCode::OK, "{restricted}");
6257 let text = restricted["choices"][0]["text"]
6258 .as_str()
6259 .unwrap_or_default();
6260 assert!(!text.is_empty(), "nothing was generated: {restricted}");
6261 assert!(
6262 text.chars().all(|c| c == 'A' || c == 'B'),
6263 "a token outside `allowed_token_ids` was drawn: {text:?}"
6264 );
6265 assert!(
6268 !free_text.chars().all(|c| c == 'A' || c == 'B'),
6269 "the unrestricted answer was already inside the allowed set: {free_text:?}"
6270 );
6271 }
6272
6273 #[tokio::test]
6280 async fn an_empty_allowed_token_ids_is_a_bad_request() {
6281 let app = streaming_test_app();
6282 let (status, body) = post_json_uri(
6283 &app,
6284 frink_api::routes::V1_COMPLETIONS,
6285 serde_json::json!({
6286 "model": "x",
6287 "prompt": "hi",
6288 "max_tokens": 4,
6289 "allowed_token_ids": []
6290 }),
6291 )
6292 .await;
6293 assert_eq!(status, StatusCode::BAD_REQUEST, "{body}");
6294 assert!(
6295 body["error"]["message"]
6296 .as_str()
6297 .unwrap_or_default()
6298 .contains("allowed_token_ids"),
6299 "{body}"
6300 );
6301 }
6302
6303 #[tokio::test]
6309 async fn bad_words_removes_a_token_without_ending_the_generation() {
6310 let app = streaming_test_app();
6311 let ask = |bad: Option<serde_json::Value>| {
6312 let mut b = serde_json::json!({
6313 "model": "x",
6314 "prompt": "hi",
6315 "max_tokens": 24,
6316 "temperature": 1.0,
6317 "seed": 11
6318 });
6319 if let Some(words) = bad {
6320 b["bad_words"] = words;
6321 }
6322 b
6323 };
6324
6325 let (status, free) =
6326 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(None)).await;
6327 assert_eq!(status, StatusCode::OK, "{free}");
6328 let free_text = free["choices"][0]["text"]
6329 .as_str()
6330 .unwrap_or_default()
6331 .to_string();
6332 let target = free_text
6335 .chars()
6336 .find(|c| c.is_ascii() && !c.is_control())
6337 .expect("the model produced some ascii");
6338
6339 let (status, steered) = post_json_uri(
6340 &app,
6341 frink_api::routes::V1_COMPLETIONS,
6342 ask(Some(serde_json::json!([target.to_string()]))),
6343 )
6344 .await;
6345 assert_eq!(status, StatusCode::OK, "{steered}");
6346 let text = steered["choices"][0]["text"].as_str().unwrap_or_default();
6347 assert!(
6348 !text.contains(target),
6349 "the forbidden {target:?} came back anyway: {text:?}"
6350 );
6351 assert_eq!(
6354 steered["usage"]["completion_tokens"], free["usage"]["completion_tokens"],
6355 "the generation ended early, so `bad_words` acted like `stop`: {steered}"
6356 );
6357 }
6358
6359 #[tokio::test]
6369 async fn every_route_refuses_the_same_unimplemented_fields() {
6370 let app = test_app();
6371 let fields = [
6372 ("n", serde_json::json!(3)),
6373 ("best_of", serde_json::json!(2)),
6374 ("prompt_logprobs", serde_json::json!(1)),
6375 ("echo", serde_json::json!(true)),
6376 ("use_beam_search", serde_json::json!(true)),
6377 ("truncate_prompt_tokens", serde_json::json!(8)),
6378 ("prompt_embeds", serde_json::json!("AA==")),
6379 ("skip_special_tokens", serde_json::json!(false)),
6380 ("return_tokens_as_token_ids", serde_json::json!(true)),
6381 ];
6382 for (field, value) in fields {
6383 for (uri, base) in [
6384 (
6385 frink_api::routes::V1_CHAT_COMPLETIONS,
6386 serde_json::json!({
6387 "model": "x",
6388 "messages": [{"role": "user", "content": "hi"}],
6389 "max_tokens": 2
6390 }),
6391 ),
6392 (
6393 frink_api::routes::V1_COMPLETIONS,
6394 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
6395 ),
6396 (
6397 frink_api::routes::COMPLETION,
6398 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6399 ),
6400 ] {
6401 let mut body = base;
6402 body[field] = value.clone();
6403 if field == "prompt_logprobs" && uri == frink_api::routes::V1_COMPLETIONS {
6410 let (status, answer) = post_json_uri(&app, uri, body).await;
6411 assert_eq!(status, StatusCode::OK, "{uri} refused it: {answer}");
6412 assert!(
6413 answer["prompt_logprobs"].is_array(),
6414 "served without the field: {answer}"
6415 );
6416 continue;
6417 }
6418 if field == "echo" && uri == frink_api::routes::V1_COMPLETIONS {
6422 let (status, answer) = post_json_uri(&app, uri, body).await;
6423 assert_eq!(status, StatusCode::OK, "{uri} refused `echo`: {answer}");
6424 assert!(
6425 answer["choices"][0]["text"]
6426 .as_str()
6427 .unwrap_or_default()
6428 .starts_with("hi"),
6429 "served without echoing the prompt: {answer}"
6430 );
6431 continue;
6432 }
6433 if field == "skip_special_tokens" || field == "return_tokens_as_token_ids" {
6437 let (status, answer) = post_json_uri(&app, uri, body).await;
6438 assert_eq!(status, StatusCode::OK, "{uri} refused `{field}`: {answer}");
6439 continue;
6440 }
6441 if field == "truncate_prompt_tokens" {
6444 let (status, answer) = post_json_uri(&app, uri, body).await;
6445 assert_eq!(
6446 status,
6447 StatusCode::OK,
6448 "{uri} refused `truncate_prompt_tokens`: {answer}"
6449 );
6450 continue;
6451 }
6452 if (field == "n" || field == "best_of")
6453 && (uri == frink_api::routes::V1_COMPLETIONS
6454 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
6455 {
6456 let (status, answer) = post_json_uri(&app, uri, body).await;
6457 assert_eq!(
6458 status,
6459 StatusCode::OK,
6460 "{uri} refused a served `{field}`: {answer}"
6461 );
6462 let want = if field == "n" { 3 } else { 1 };
6466 assert_eq!(
6467 answer["choices"].as_array().map(Vec::len),
6468 Some(want),
6469 "{field}: {answer}"
6470 );
6471 continue;
6472 }
6473 let (status, answer) = post_json_uri(&app, uri, body).await;
6474 assert_eq!(
6475 status,
6476 StatusCode::NOT_IMPLEMENTED,
6477 "{uri} served `{field}` instead of refusing it: {answer}"
6478 );
6479 assert!(
6480 answer["error"]["message"]
6481 .as_str()
6482 .is_some_and(|m| m.contains(field)),
6483 "{uri} refused `{field}` without naming it: {answer}"
6484 );
6485 }
6486 }
6487 }
6488
6489 #[tokio::test]
6490 async fn the_native_completion_wire_is_not_the_openai_one() {
6491 let app = test_app();
6492
6493 let (status, native) = post_json_uri(
6494 &app,
6495 frink_api::routes::COMPLETION,
6496 serde_json::json!({"prompt": "hi", "n_predict": 4}),
6497 )
6498 .await;
6499 assert_eq!(status, StatusCode::OK, "{native}");
6500 assert!(native["content"].is_string(), "{native}");
6501 assert_eq!(native["stop"], true);
6502 assert_eq!(native["stop_type"], "limit");
6503 assert_eq!(native["stopping_word"], "");
6504 assert_eq!(native["truncated"], false);
6505 assert_eq!(native["id_slot"], -1);
6506 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
6507 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
6508 assert!(
6509 native.get("choices").is_none(),
6510 "the native shape has no `choices`: {native}"
6511 );
6512
6513 let (status, openai) = post_json_uri(
6514 &app,
6515 frink_api::routes::V1_COMPLETIONS,
6516 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
6517 )
6518 .await;
6519 assert_eq!(status, StatusCode::OK);
6520 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
6521 assert!(
6522 openai.get("content").is_none(),
6523 "the OpenAI shape has no top-level `content`: {openai}"
6524 );
6525 }
6526
6527 #[tokio::test]
6531 async fn both_native_spellings_reach_the_same_handler() {
6532 let app = test_app();
6533 for route in [
6534 frink_api::routes::COMPLETION,
6535 frink_api::routes::COMPLETIONS,
6536 ] {
6537 let (status, body) = post_json_uri(
6538 &app,
6539 route,
6540 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
6541 )
6542 .await;
6543 assert_eq!(status, StatusCode::OK, "{route}: {body}");
6544 assert_eq!(body["stop"], true, "{route}");
6545 assert!(body["content"].is_string(), "{route}");
6546 }
6547
6548 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6551 let routes: Vec<&str> = stats["recent"]
6552 .as_array()
6553 .unwrap()
6554 .iter()
6555 .map(|row| row["route"].as_str().unwrap())
6556 .collect();
6557 assert!(
6558 routes.contains(&frink_api::routes::COMPLETION),
6559 "{routes:?}"
6560 );
6561 assert!(
6562 routes.contains(&frink_api::routes::COMPLETIONS),
6563 "{routes:?}"
6564 );
6565 }
6566
6567 #[tokio::test]
6573 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
6574 let app = streaming_test_app();
6575 let raw = post_sse_raw_uri(
6576 &app,
6577 frink_api::routes::COMPLETION,
6578 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
6579 )
6580 .await;
6581
6582 assert!(
6583 !raw.contains("[DONE]"),
6584 "llama.cpp's native stream has no sentinel: {raw}"
6585 );
6586 let frames: Vec<serde_json::Value> = raw
6587 .lines()
6588 .filter_map(|line| line.strip_prefix("data: "))
6589 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
6590 .collect();
6591 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
6592
6593 let (last, partials) = frames.split_last().unwrap();
6594 assert_eq!(last["stop"], true, "the last frame closes the stream");
6595 assert!(last["timings"].is_object(), "{last}");
6596 assert!(last["stop_type"].is_string(), "{last}");
6597 for partial in partials {
6598 assert_eq!(partial["stop"], false, "{partial}");
6599 assert!(partial["content"].is_string(), "{partial}");
6600 assert!(partial.get("timings").is_none(), "{partial}");
6604 assert!(partial.get("generation_settings").is_none(), "{partial}");
6605 }
6606 let streamed: String = partials
6609 .iter()
6610 .filter_map(|p| p["content"].as_str())
6611 .collect();
6612 assert_eq!(last["content"].as_str().unwrap(), streamed);
6613 }
6614
6615 #[tokio::test]
6620 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
6621 let app = test_app();
6622 for body in [
6623 serde_json::json!({"prompt": "hi"}),
6624 serde_json::json!({"prompt": "hi", "n_predict": -1}),
6625 ] {
6626 let (status, refusal) =
6627 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
6628 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
6629 assert!(
6630 refusal["error"]["message"]
6631 .as_str()
6632 .unwrap()
6633 .contains("n_predict"),
6634 "{refusal}"
6635 );
6636 }
6637 let (status, _) = post_json_uri(
6640 &app,
6641 frink_api::routes::COMPLETION,
6642 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6643 )
6644 .await;
6645 assert_eq!(status, StatusCode::OK);
6646 }
6647
6648 #[tokio::test]
6657 async fn a_stop_string_halts_the_answer_and_is_named_back() {
6658 let app = streaming_test_app();
6659 let ask = |stop: serde_json::Value| {
6660 let app = app.clone();
6661 async move {
6662 post_json_uri(
6663 &app,
6664 frink_api::routes::COMPLETION,
6665 serde_json::json!({
6666 "prompt": "hi",
6667 "n_predict": 64,
6668 "ignore_eos": true,
6669 "stop": stop,
6670 }),
6671 )
6672 .await
6673 .1
6674 }
6675 };
6676
6677 let baseline = ask(serde_json::json!([])).await;
6678 assert_eq!(baseline["stop_type"], "limit");
6679 assert_eq!(baseline["stopping_word"], "");
6680 let text = baseline["content"].as_str().unwrap().to_string();
6681 let sentinel: String = text.chars().skip(1).take(2).collect();
6687 assert_eq!(
6688 sentinel.chars().count(),
6689 2,
6690 "the fixture must produce enough output to cut: {text:?}"
6691 );
6692 let cut = text.find(&sentinel).expect("it came out of this text");
6693
6694 let stopped = ask(serde_json::json!([sentinel])).await;
6695 assert_eq!(stopped["stop_type"], "word", "{stopped}");
6696 assert_eq!(stopped["stopping_word"], sentinel);
6697 assert_eq!(
6698 stopped["content"].as_str().unwrap(),
6699 &text[..cut],
6700 "the answer must be cut at the sentinel, not run past it"
6701 );
6702 }
6703
6704 #[tokio::test]
6710 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
6711 let app = test_app();
6712
6713 let (v1_status, v1) = post_json_uri(
6714 &app,
6715 frink_api::routes::V1_TOKENIZE,
6716 serde_json::json!({"prompt": "hello"}),
6717 )
6718 .await;
6719 let (alias_status, alias) = post_json_uri(
6720 &app,
6721 frink_api::routes::TOKENIZE,
6722 serde_json::json!({"content": "hello"}),
6723 )
6724 .await;
6725 assert_eq!(v1_status, StatusCode::OK);
6726 assert_eq!(alias_status, StatusCode::OK, "{alias}");
6727 assert_eq!(v1["tokens"], alias["tokens"]);
6728 assert!(!alias["tokens"].as_array().unwrap().is_empty());
6729
6730 let (status, both_ways) = post_json_uri(
6733 &app,
6734 frink_api::routes::TOKENIZE,
6735 serde_json::json!({"prompt": "hello"}),
6736 )
6737 .await;
6738 assert_eq!(status, StatusCode::OK);
6739 assert_eq!(both_ways["tokens"], v1["tokens"]);
6740 }
6741
6742 #[tokio::test]
6747 async fn detokenize_answers_under_both_dialects_keys() {
6748 let app = test_app();
6749 for route in [
6750 frink_api::routes::DETOKENIZE,
6751 frink_api::routes::V1_DETOKENIZE,
6752 ] {
6753 let (status, body) =
6754 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
6755 assert_eq!(status, StatusCode::OK, "{route}");
6756 assert_eq!(body["text"], "hi", "{route}");
6757 assert_eq!(body["content"], body["text"], "{route}");
6758 }
6759 }
6760
6761 #[tokio::test]
6765 async fn the_alias_is_recorded_under_the_path_the_client_called() {
6766 let app = test_app();
6767 let (status, _) = post_json_uri(
6768 &app,
6769 frink_api::routes::TOKENIZE,
6770 serde_json::json!({"content": "hello"}),
6771 )
6772 .await;
6773 assert_eq!(status, StatusCode::OK);
6774
6775 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6776 let routes: Vec<&str> = stats["recent"]
6777 .as_array()
6778 .unwrap()
6779 .iter()
6780 .map(|row| row["route"].as_str().unwrap())
6781 .collect();
6782 assert!(
6783 routes.contains(&frink_api::routes::TOKENIZE),
6784 "the alias must be its own row: {routes:?}"
6785 );
6786 assert!(
6787 !routes.contains(&frink_api::routes::V1_TOKENIZE),
6788 "nothing called /v1/tokenize: {routes:?}"
6789 );
6790 }
6791
6792 #[tokio::test]
6797 async fn add_special_prepends_the_same_bos_the_decoder_would() {
6798 let mut cfg = test_dense_fixture();
6799 cfg.vocab_size = 256;
6800 let model = Model::Gguf(GgufModel {
6801 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
6802 tokenizer: Arc::new(ServerTokenizer::Byte),
6803 stop_tokens: StopTokens::default(),
6804 bos_id: Some(7),
6805 is_synthetic: true,
6806 chat_template: chat_template::PromptTemplate::plain(),
6807 });
6808 let app = test_app_with_state(Arc::new(test_state(
6809 model,
6810 ResponseCache::new(1000, Duration::from_secs(3600)),
6811 )));
6812
6813 let (_, plain) = post_json_uri(
6814 &app,
6815 frink_api::routes::TOKENIZE,
6816 serde_json::json!({"content": "hi"}),
6817 )
6818 .await;
6819 let (_, special) = post_json_uri(
6820 &app,
6821 frink_api::routes::TOKENIZE,
6822 serde_json::json!({"content": "hi", "add_special": true}),
6823 )
6824 .await;
6825
6826 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6827 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6828 assert_eq!(special["count"], 3);
6829 }
6830
6831 #[tokio::test]
6834 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6835 let app = test_app();
6836 let (status, _) = post_json_uri(
6837 &app,
6838 frink_api::routes::V1_EMBEDDINGS,
6839 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6840 )
6841 .await;
6842 assert_eq!(status, StatusCode::BAD_REQUEST);
6843
6844 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6845 let recent = stats["recent"].as_array().unwrap();
6846 assert_eq!(recent.len(), 1);
6847 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
6848 assert_eq!(recent[0]["status"], 400);
6849 assert_eq!(
6850 recent[0]["prompt_tokens"], 0,
6851 "a rejected call embedded nothing"
6852 );
6853 }
6854
6855 #[tokio::test]
6858 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6859 let app = test_app();
6860 let key = "sk-monitor-secret";
6861 let (status, _) = post_json_with_headers(
6862 &app,
6863 "/v1/chat/completions",
6864 serde_json::json!({
6865 "model": "x",
6866 "messages": [{"role": "user", "content": "hi"}],
6867 "max_tokens": 2
6868 }),
6869 &[
6870 ("authorization", &format!("Bearer {key}")),
6871 ("x-frink-client", "frink-studio"),
6872 ],
6873 )
6874 .await;
6875 assert_eq!(status, StatusCode::OK);
6876
6877 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6878 let row = stats["recent"].as_array().unwrap()[0].clone();
6879 let fingerprint = row["via_api_key"]
6880 .as_str()
6881 .expect("the row names the key that served it")
6882 .to_string();
6883 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6884 assert!(!fingerprint.contains(key));
6885 assert!(
6886 !serde_json::to_string(&stats).unwrap().contains(key),
6887 "the stats payload must not carry the key in any form"
6888 );
6889 assert_eq!(row["client"], "frink-studio");
6890 }
6891
6892 #[tokio::test]
6895 async fn different_keys_are_different_callers_and_no_key_is_null() {
6896 let app = test_app();
6897 let body = serde_json::json!({
6898 "model": "x",
6899 "messages": [{"role": "user", "content": "hi"}],
6900 "max_tokens": 1
6901 });
6902 for headers in [
6903 vec![("authorization", "Bearer key-one")],
6904 vec![("authorization", "Bearer key-two")],
6905 vec![],
6906 ] {
6907 let (status, _) =
6908 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6909 assert_eq!(status, StatusCode::OK);
6910 }
6911
6912 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6913 let recent = stats["recent"].as_array().unwrap();
6914 assert_eq!(recent.len(), 3);
6915 let one = recent[0]["via_api_key"].as_str().unwrap();
6916 let two = recent[1]["via_api_key"].as_str().unwrap();
6917 assert_ne!(one, two, "two keys must not collapse into one caller");
6918 assert!(
6919 recent[2]["via_api_key"].is_null(),
6920 "an unauthenticated call is null, not a fingerprint of nothing"
6921 );
6922 assert!(recent[2]["client"].is_null());
6923 }
6924
6925 #[tokio::test]
6930 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6931 let state = Arc::new(test_state(
6932 named_test_model("really-loaded", 256),
6933 ResponseCache::new(4, Duration::from_secs(60)),
6934 ));
6935 let app = test_app_with_state(Arc::clone(&state));
6936
6937 let (status, _) = post_json_uri(
6938 &app,
6939 "/v1/chat/completions",
6940 serde_json::json!({
6941 "model": "gpt-4-turbo-that-is-not-here",
6942 "messages": [{"role": "user", "content": "hi"}],
6943 "max_tokens": 2
6944 }),
6945 )
6946 .await;
6947 assert_eq!(status, StatusCode::OK);
6948
6949 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6950 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6951
6952 state.swap_active(None);
6955 let (status, _) = post_json_uri(
6956 &app,
6957 "/v1/chat/completions",
6958 serde_json::json!({
6959 "model": "gpt-4-turbo-that-is-not-here",
6960 "messages": [{"role": "user", "content": "hi"}]
6961 }),
6962 )
6963 .await;
6964 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6965 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6966 let recent = stats["recent"].as_array().unwrap();
6967 assert!(recent[recent.len() - 1]["model"].is_null());
6968 }
6969
6970 #[tokio::test]
6974 async fn a_streamed_row_names_the_model_it_decoded_against() {
6975 let state = Arc::new(test_state(
6976 named_test_model("model-before", 256),
6977 ResponseCache::new(4, Duration::from_secs(60)),
6978 ));
6979 let app = test_app_with_state(Arc::clone(&state));
6980 let _ = post_sse_raw(&app, resumable_request()).await;
6981 active_model(&state, "model-after");
6983
6984 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6985 assert_eq!(stats["recent"][0]["model"], "model-before");
6986 }
6987
6988 #[tokio::test]
6991 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6992 let app = test_app();
6993 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6994 assert_eq!(status, StatusCode::OK);
6995 assert!(
6996 stats["queue_depth"].is_null(),
6997 "without continuous batching nothing queues, so there is nothing to measure"
6998 );
6999 assert!(stats["queue_rejected_total"].is_null());
7000 assert_eq!(
7001 stats["generating_now"], 0,
7002 "work in progress is measured and really is zero here"
7003 );
7004 }
7005
7006 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
7011 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
7012 }
7013
7014 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
7018 use http_body_util::BodyExt;
7019 use tower::ServiceExt;
7020
7021 let response = app
7022 .clone()
7023 .oneshot(
7024 axum::http::Request::builder()
7025 .method("POST")
7026 .uri(uri)
7027 .header("content-type", "application/json")
7028 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7029 .unwrap(),
7030 )
7031 .await
7032 .unwrap();
7033 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7034 String::from_utf8(bytes.to_vec()).unwrap()
7035 }
7036
7037 async fn get_json_with_headers(
7038 app: &Router,
7039 uri: &str,
7040 headers: &[(&str, &str)],
7041 ) -> (StatusCode, serde_json::Value) {
7042 use http_body_util::BodyExt;
7043 use tower::ServiceExt;
7044
7045 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
7046 for (name, value) in headers {
7047 builder = builder.header(*name, *value);
7048 }
7049 let response = app
7050 .clone()
7051 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
7052 .await
7053 .unwrap();
7054 let status = response.status();
7055 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7056 (
7057 status,
7058 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
7059 )
7060 }
7061
7062 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
7063 body.lines()
7064 .filter_map(|line| line.strip_prefix(field))
7065 .map(str::trim)
7066 .collect()
7067 }
7068
7069 fn resumable_request() -> serde_json::Value {
7070 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 "stream": true,
7076 "stream_resumable": true,
7077 })
7078 }
7079
7080 #[tokio::test]
7085 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
7086 let app = test_app();
7087 let body = post_sse_raw(&app, resumable_request()).await;
7088
7089 let request_id = body
7090 .lines()
7091 .find_map(|l| l.strip_prefix("data: "))
7092 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
7093 .and_then(|v| v["request_id"].as_str().map(str::to_string))
7094 .expect("the first chunk names the request");
7095
7096 let ids = sse_field(&body, "id:");
7097 let datas = sse_field(&body, "data:");
7098 assert_eq!(
7099 ids.len(),
7100 datas.len(),
7101 "every event carries an id, or a reconnect cannot name where it stopped"
7102 );
7103 for (i, id) in ids.iter().enumerate() {
7104 assert_eq!(*id, format!("{request_id}:{i}"));
7105 }
7106 let retries = sse_field(&body, "retry:");
7107 assert_eq!(
7108 retries.len(),
7109 1,
7110 "the reconnect delay is stated once, not on every event"
7111 );
7112 assert_eq!(retries[0], "1500");
7113 assert!(
7114 body.contains("data: [DONE]"),
7115 "the end of stream is still stated"
7116 );
7117 }
7118
7119 #[tokio::test]
7123 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
7124 let app = test_app();
7125 let mut request = resumable_request();
7126 request["stream_resumable"] = serde_json::json!(false);
7127 let body = post_sse_raw(&app, request).await;
7128 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
7129 assert!(
7130 sse_field(&body, "id:").is_empty(),
7131 "an id promises a replay this stream cannot serve"
7132 );
7133 assert!(sse_field(&body, "retry:").is_empty());
7134 }
7135
7136 #[tokio::test]
7140 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
7141 let app = test_app();
7142 let body = post_sse_raw(&app, resumable_request()).await;
7143 let request_id = sse_field(&body, "id:")[0]
7144 .rsplit_once(':')
7145 .unwrap()
7146 .0
7147 .to_string();
7148 let streamed: Vec<String> = sse_field(&body, "data:")
7149 .iter()
7150 .map(|d| d.to_string())
7151 .collect();
7152
7153 let (status, polled) = get_json(
7154 &app,
7155 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
7156 )
7157 .await;
7158 assert_eq!(status, StatusCode::OK);
7159 let events: Vec<String> = polled["events"]
7160 .as_array()
7161 .unwrap()
7162 .iter()
7163 .map(|e| e["data"].as_str().unwrap().to_string())
7164 .collect();
7165 assert_eq!(
7166 events, streamed,
7167 "the fallback must deliver the same answer, not a re-run of it"
7168 );
7169 assert_eq!(polled["request_id"], request_id);
7170 assert_eq!(
7171 polled["done"], false,
7172 "events were still being handed out, so the client must ask again"
7173 );
7174
7175 let next = polled["next_index"].as_u64().unwrap();
7178 let (_, drained) = get_json(
7179 &app,
7180 &format!(
7181 "{}?from={next}",
7182 frink_api::routes::v1_stream_poll(&request_id)
7183 ),
7184 )
7185 .await;
7186 assert_eq!(drained["done"], true);
7187 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
7188 }
7189
7190 #[tokio::test]
7194 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
7195 let app = test_app();
7196 let body = post_sse_raw(&app, resumable_request()).await;
7197 let ids = sse_field(&body, "id:");
7198 let datas: Vec<String> = sse_field(&body, "data:")
7199 .iter()
7200 .map(|d| d.to_string())
7201 .collect();
7202 assert!(
7203 ids.len() >= 3,
7204 "need a few events to resume into the middle"
7205 );
7206 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
7207
7208 let (status, resumed) = get_json_with_headers(
7209 &app,
7210 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
7211 &[],
7212 )
7213 .await;
7214 assert_eq!(status, StatusCode::OK);
7215 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
7216
7217 let (_, tail) = get_json(
7219 &app,
7220 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
7221 )
7222 .await;
7223 let tail_events: Vec<String> = tail["events"]
7224 .as_array()
7225 .unwrap()
7226 .iter()
7227 .map(|e| e["data"].as_str().unwrap().to_string())
7228 .collect();
7229 assert_eq!(tail_events, datas[2..].to_vec());
7230 }
7231
7232 #[tokio::test]
7235 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
7236 use http_body_util::BodyExt;
7237 use tower::ServiceExt;
7238
7239 let app = test_app();
7240 let body = post_sse_raw(&app, resumable_request()).await;
7241 let ids = sse_field(&body, "id:");
7242 let datas: Vec<String> = sse_field(&body, "data:")
7243 .iter()
7244 .map(|d| d.to_string())
7245 .collect();
7246 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
7247
7248 let response = app
7249 .clone()
7250 .oneshot(
7251 axum::http::Request::builder()
7252 .method("GET")
7253 .uri(frink_api::routes::v1_stream(&request_id))
7254 .header("last-event-id", format!("{request_id}:0"))
7255 .body(axum::body::Body::empty())
7256 .unwrap(),
7257 )
7258 .await
7259 .unwrap();
7260 assert_eq!(response.status(), StatusCode::OK);
7261 assert_eq!(
7262 response
7263 .headers()
7264 .get("x-accel-buffering")
7265 .and_then(|v| v.to_str().ok()),
7266 Some("no"),
7267 "the reconnect needs the same anti-buffering header as the stream"
7268 );
7269 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7270 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
7271 assert_eq!(
7272 sse_field(&resumed, "data:")
7273 .iter()
7274 .map(|d| d.to_string())
7275 .collect::<Vec<_>>(),
7276 datas[1..].to_vec()
7277 );
7278 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
7279 }
7280
7281 #[tokio::test]
7285 async fn a_last_event_id_from_another_stream_is_refused() {
7286 let app = test_app();
7287 let body = post_sse_raw(&app, resumable_request()).await;
7288 let request_id = sse_field(&body, "id:")[0]
7289 .rsplit_once(':')
7290 .unwrap()
7291 .0
7292 .to_string();
7293
7294 let (status, err) = get_json_with_headers(
7295 &app,
7296 &frink_api::routes::v1_stream(&request_id),
7297 &[("last-event-id", "chatcmpl-someone-else:3")],
7298 )
7299 .await;
7300 assert_eq!(status, StatusCode::BAD_REQUEST);
7301 assert_eq!(err["error"]["code"], "bad_last_event_id");
7302 }
7303
7304 #[tokio::test]
7308 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
7309 let app = test_app();
7310 let mut request = resumable_request();
7311 request["stream_resumable"] = serde_json::json!(false);
7312 let body = post_sse_raw(&app, request).await;
7313 let request_id = body
7314 .lines()
7315 .find_map(|l| l.strip_prefix("data: "))
7316 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
7317 .and_then(|v| v["request_id"].as_str().map(str::to_string))
7318 .unwrap();
7319
7320 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
7321 assert_eq!(status, StatusCode::NOT_FOUND);
7322 assert_eq!(err["error"]["code"], "stream_not_found");
7323 assert!(err["error"]["message"]
7324 .as_str()
7325 .unwrap()
7326 .contains("stream_resumable"));
7327 }
7328
7329 #[test]
7333 fn the_axum_stream_patterns_match_the_published_templates() {
7334 assert_eq!(
7335 axum_path(frink_api::routes::V1_STREAM),
7336 "/v1/stream/:request_id"
7337 );
7338 assert_eq!(
7339 axum_path(frink_api::routes::V1_STREAM_POLL),
7340 "/v1/stream/:request_id/poll"
7341 );
7342 assert_eq!(
7343 frink_api::routes::v1_stream("abc"),
7344 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
7345 );
7346 }
7347
7348 #[test]
7358 fn no_published_template_reaches_the_router_with_its_braces() {
7359 for template in [
7360 frink_api::routes::V1_STREAM,
7361 frink_api::routes::V1_STREAM_POLL,
7362 frink_api::routes::V1_RESPONSE,
7363 frink_api::routes::V1_RESPONSE_CANCEL,
7364 frink_api::routes::ADMIN_TASK_CANCEL,
7365 ] {
7366 assert!(
7367 template.contains('{'),
7368 "{template} is in the template list but has no placeholder"
7369 );
7370 let mounted = axum_path(template);
7371 assert!(
7372 !mounted.contains('{') && !mounted.contains('}'),
7373 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
7374 );
7375 assert!(
7376 mounted.contains(':'),
7377 "{template} lost its placeholder entirely and would match one path only"
7378 );
7379 }
7380 }
7381
7382 #[tokio::test]
7390 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
7391 let app = test_app();
7392 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
7393 assert_eq!(status, StatusCode::NOT_FOUND);
7394 assert!(
7395 !body.is_null(),
7396 "empty body means axum never matched the route, so the id was read as a literal segment"
7397 );
7398 }
7399
7400 #[tokio::test]
7403 async fn the_task_list_starts_empty_rather_than_absent() {
7404 let app = test_app();
7405 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
7406 assert_eq!(status, StatusCode::OK);
7407 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
7408 }
7409
7410 #[tokio::test]
7419 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
7420 assert!(
7421 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
7422 "this test asserts the unconfigured behaviour"
7423 );
7424 let app = test_app();
7425 let (status, body) = post_json_uri(
7426 &app,
7427 &format!("{}?action=save", frink_api::routes::slots_id(0)),
7428 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
7429 )
7430 .await;
7431 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7432 assert!(
7433 body["error"]["message"]
7434 .as_str()
7435 .unwrap()
7436 .contains("--slot-save-path"),
7437 "{body}"
7438 );
7439 }
7440
7441 pub(crate) async fn post_json_uri(
7442 app: &Router,
7443 uri: &str,
7444 body: serde_json::Value,
7445 ) -> (StatusCode, serde_json::Value) {
7446 use http_body_util::BodyExt;
7447 use tower::ServiceExt;
7448
7449 let response = app
7450 .clone()
7451 .oneshot(
7452 axum::http::Request::builder()
7453 .method("POST")
7454 .uri(uri)
7455 .header("content-type", "application/json")
7456 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7457 .unwrap(),
7458 )
7459 .await
7460 .unwrap();
7461 let status = response.status();
7462 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7463 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
7464 (status, json)
7465 }
7466
7467 pub(crate) async fn get_json_uri(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7470 use http_body_util::BodyExt;
7471 use tower::ServiceExt;
7472
7473 let response = app
7474 .clone()
7475 .oneshot(
7476 axum::http::Request::builder()
7477 .method("GET")
7478 .uri(uri)
7479 .body(axum::body::Body::empty())
7480 .unwrap(),
7481 )
7482 .await
7483 .unwrap();
7484 let status = response.status();
7485 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7486 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
7487 (status, json)
7488 }
7489
7490 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
7491 post_json_uri(app, "/v1/chat/completions", body).await.1
7492 }
7493
7494 #[tokio::test]
7503 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
7504 let app = test_app();
7505 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
7506 assert_eq!(status, StatusCode::OK);
7507
7508 let memory = &body["memory"];
7509 if memory.is_null() {
7510 return;
7513 }
7514 assert!(
7515 memory["bytes"].as_u64().is_some_and(|b| b > 0),
7516 "a read that produced a zero is a broken read, not an idle \
7517 engine: {memory}"
7518 );
7519 assert!(
7520 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
7521 "the quantity must travel with the number: {memory}"
7522 );
7523 }
7524
7525 #[tokio::test]
7531 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
7532 let app = test_app();
7533 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
7534 assert_eq!(status, StatusCode::OK);
7535 for pool in ["kv_pages", "window_slots", "state_slots"] {
7536 assert!(
7537 body["pools"][pool].is_null(),
7538 "{pool} must be null rather than a zero row: {}",
7539 body["pools"]
7540 );
7541 }
7542 }
7543
7544 #[tokio::test]
7552 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
7553 use http_body_util::BodyExt;
7554 use tower::ServiceExt;
7555
7556 let app = test_app();
7557 let send = |body: serde_json::Value| {
7558 let app = app.clone();
7559 async move {
7560 app.oneshot(
7561 axum::http::Request::builder()
7562 .method("POST")
7563 .uri(frink_api::routes::V1_MESSAGES)
7564 .header("content-type", "application/json")
7565 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7566 .unwrap(),
7567 )
7568 .await
7569 .unwrap()
7570 }
7571 };
7572
7573 let ok = send(serde_json::json!({
7574 "model": "test",
7575 "max_tokens": 1,
7576 "messages": [{"role": "user", "content": "hi"}],
7577 }))
7578 .await;
7579 assert_eq!(ok.status(), StatusCode::OK);
7580 let id = ok
7581 .headers()
7582 .get("request-id")
7583 .expect("a served message names its id")
7584 .to_str()
7585 .unwrap()
7586 .to_string();
7587 assert!(!id.is_empty());
7588
7589 let bad = send(serde_json::json!({"model": "test"})).await;
7592 assert!(bad.status().is_client_error());
7593 let other = bad.headers().get("request-id").expect("errors too");
7594 assert_ne!(other.to_str().unwrap(), id);
7595 let _ = bad.into_body().collect().await.unwrap();
7596 }
7597
7598 #[tokio::test]
7606 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
7607 let state = Arc::new(test_state(
7608 test_model_full_byte_vocab(),
7609 ResponseCache::new(1000, Duration::from_secs(3600)),
7610 ));
7611 let app = test_app_with_state(Arc::clone(&state));
7612 let body = serde_json::json!({
7613 "model": "test",
7614 "messages": [{"role": "user", "content": "hi"}],
7615 "max_tokens": 1,
7616 });
7617
7618 state
7619 .maintenance
7620 .lock()
7621 .unwrap()
7622 .begin_rebuild()
7623 .expect("a fresh server is serving, so the rebuild starts");
7624 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
7625 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7626 assert_eq!(refused["error"]["type"], "cache_rebuilding");
7627
7628 state.maintenance.lock().unwrap().finish_rebuild(true);
7629 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
7630 assert_eq!(
7631 status,
7632 StatusCode::OK,
7633 "the gate reopens; a rebuild is not a latch"
7634 );
7635 }
7636
7637 #[tokio::test]
7642 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
7643 let app = test_app();
7644 let (status, body) = post_json_uri(
7645 &app,
7646 frink_api::routes::V1_CANCEL,
7647 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
7648 )
7649 .await;
7650 assert_eq!(status, StatusCode::NOT_FOUND);
7651 assert_eq!(body["cancelled"], serde_json::json!(false));
7652 assert_eq!(body["request_id"], "chatcmpl-never-issued");
7653 assert!(
7654 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
7655 "the verdict must carry a human reason: {body}"
7656 );
7657 }
7658
7659 #[tokio::test]
7663 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
7664 let state = Arc::new(test_state(
7665 test_model_full_byte_vocab(),
7666 ResponseCache::new(1000, Duration::from_secs(3600)),
7667 ));
7668 let app = test_app_with_state(Arc::clone(&state));
7669 let (token, _guard) = state.cancels.register("chatcmpl-live");
7670
7671 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7672 assert_eq!(status, StatusCode::OK);
7673 assert_eq!(before["generating_now"], serde_json::json!(1));
7674
7675 let (status, body) = post_json_uri(
7676 &app,
7677 frink_api::routes::V1_CANCEL,
7678 serde_json::json!({ "request_id": "chatcmpl-live" }),
7679 )
7680 .await;
7681 assert_eq!(status, StatusCode::OK);
7682 assert_eq!(body["cancelled"], serde_json::json!(true));
7683 assert!(
7684 token.is_cancelled(),
7685 "the endpoint answered ok without setting the flag the decode loop reads"
7686 );
7687 }
7688
7689 #[tokio::test]
7690 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
7691 let app = test_app();
7692 let (status, tok) =
7693 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
7694 assert_eq!(status, StatusCode::OK);
7695 let tokens = tok["tokens"].as_array().unwrap();
7696 assert_eq!(tok["count"], tokens.len());
7697 assert!(!tokens.is_empty());
7698
7699 let (status, detok) = post_json_uri(
7700 &app,
7701 "/v1/detokenize",
7702 serde_json::json!({ "tokens": tokens }),
7703 )
7704 .await;
7705 assert_eq!(status, StatusCode::OK);
7706 assert_eq!(detok["text"], "Hi");
7707
7708 let (status, emb) = post_json_uri(
7709 &app,
7710 "/v1/embeddings",
7711 serde_json::json!({
7712 "input": "Hi",
7713 "embedding_type": "mean"
7714 }),
7715 )
7716 .await;
7717 assert_eq!(status, StatusCode::OK);
7718 let vec = emb["data"][0]["embedding"].as_array().unwrap();
7719 assert!(!vec.is_empty());
7720 assert!(vec.iter().all(|v| v.as_f64().is_some()));
7721 }
7722
7723 #[tokio::test]
7729 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
7730 let app = test_app();
7731 let (status, body) = post_json_uri(
7732 &app,
7733 "/v1/embeddings",
7734 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
7735 )
7736 .await;
7737 assert_eq!(status, StatusCode::BAD_REQUEST);
7738 let msg = body["error"]["message"].as_str().unwrap();
7739 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
7740 }
7741
7742 #[tokio::test]
7747 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7748 async fn a_real_embedding_model_serves_v1_embeddings() {
7749 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7750 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7751 if !path.exists() {
7752 eprintln!("SKIP: {} not present", path.display());
7753 return;
7754 }
7755 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
7756 let mut state = test_state(
7757 test_model_full_byte_vocab(),
7758 ResponseCache::new(1000, Duration::from_secs(3600)),
7759 );
7760 state.embedding = Some(Arc::new(encoder));
7761 let app = test_app_with_state(Arc::new(state));
7762
7763 let (status, body) = post_json_uri(
7764 &app,
7765 "/v1/embeddings",
7766 serde_json::json!({ "input": ["Hello world", "a second input"] }),
7767 )
7768 .await;
7769 assert_eq!(status, StatusCode::OK, "{body}");
7770 assert_eq!(body["model"], "bge-small-en-v1.5");
7771 let data = body["data"].as_array().unwrap();
7772 assert_eq!(data.len(), 2);
7773 for (i, row) in data.iter().enumerate() {
7774 assert_eq!(row["index"], i);
7775 let v: Vec<f64> = row["embedding"]
7776 .as_array()
7777 .unwrap()
7778 .iter()
7779 .map(|x| x.as_f64().unwrap())
7780 .collect();
7781 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
7782 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
7783 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
7784 }
7785 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
7788
7789 let (status, mean) = post_json_uri(
7793 &app,
7794 "/v1/embeddings",
7795 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
7796 )
7797 .await;
7798 assert_eq!(status, StatusCode::OK);
7799 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
7800 }
7801
7802 #[tokio::test]
7814 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7815 async fn an_encoder_can_be_the_loaded_model() {
7816 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7817 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7818 if !path.exists() {
7819 eprintln!("SKIP: {} not present", path.display());
7820 return;
7821 }
7822
7823 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
7827 assert!(
7828 matches!(loaded, model::LoadedModel::Encoder(_)),
7829 "an encoder-only GGUF reached a decoder loader"
7830 );
7831 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
7832 assert!(
7833 matches!(loaded, Loaded::Encoder(_)),
7834 "the encoder did not stay an encoder through activation"
7835 );
7836 assert!(
7837 batcher.is_none() && ceiling.is_none(),
7838 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7839 );
7840
7841 let state = test_state(
7842 test_model_full_byte_vocab(),
7843 ResponseCache::new(1000, Duration::from_secs(3600)),
7844 );
7845 state.swap_active(Some(Arc::new(ActiveModel {
7846 id: None,
7847 loaded,
7848 batcher,
7849 ceiling,
7850 checkpoint_path: None,
7851 })));
7852 let app = test_app_with_state(Arc::new(state));
7853
7854 let (status, body) = post_json_uri(
7856 &app,
7857 "/v1/embeddings",
7858 serde_json::json!({ "input": "Hello world" }),
7859 )
7860 .await;
7861 assert_eq!(status, StatusCode::OK, "{body}");
7862 assert_eq!(body["model"], "bge-small-en-v1.5");
7863 let v = body["data"][0]["embedding"].as_array().unwrap();
7864 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
7865
7866 let (status, body) = post_json_uri(
7868 &app,
7869 "/v1/chat/completions",
7870 serde_json::json!({
7871 "model": "bge-small-en-v1.5",
7872 "messages": [{"role": "user", "content": "hi"}],
7873 }),
7874 )
7875 .await;
7876 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
7877 let msg = body["error"]["message"].as_str().unwrap();
7878 for fact in [
7879 "bge-small-en-v1.5",
7880 "bert",
7881 "embedding model",
7882 "/v1/embeddings",
7883 ] {
7884 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
7885 }
7886
7887 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
7889 assert_eq!(status, StatusCode::OK);
7890 let entry = &models["data"][0];
7891 assert_eq!(entry["id"], "bge-small-en-v1.5");
7892 assert_eq!(entry["frink_model_kind"], "embedding");
7893 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
7894 assert_eq!(entry["frink_n_embd"], 384);
7895 assert_eq!(entry["frink_pooling"], "CLS");
7896 assert_eq!(
7897 entry["frink_endpoints"],
7898 serde_json::json!(["/v1/embeddings"])
7899 );
7900 assert!(entry.get("supported_reasoning_efforts").is_none());
7903
7904 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
7906 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
7907 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
7908 assert_eq!(health["model"]["synthetic_weights"], false);
7909 let weights = health["capabilities"]
7910 .as_array()
7911 .unwrap()
7912 .iter()
7913 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
7914 .expect("a real-weights capability row");
7915 let detail = weights["detail"].as_str().unwrap_or_default();
7916 assert!(detail.contains("ENCODER"), "{detail}");
7917 let (status, body) = post_json_uri(
7925 &app,
7926 frink_api::routes::V1_TOKENIZE,
7927 serde_json::json!({ "content": "hello world" }),
7928 )
7929 .await;
7930 assert_eq!(
7931 status,
7932 StatusCode::OK,
7933 "an encoder has a real tokenizer: {body}"
7934 );
7935 let tokens = body["tokens"].as_array().expect("tokens array").clone();
7936 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
7937
7938 let (status, body) = post_json_uri(
7939 &app,
7940 frink_api::routes::V1_DETOKENIZE,
7941 serde_json::json!({ "tokens": tokens }),
7942 )
7943 .await;
7944 assert_eq!(status, StatusCode::OK, "{body}");
7945 let round_tripped = body["content"].as_str().expect("content").to_string();
7946 assert!(
7947 round_tripped.contains("hello") && round_tripped.contains("world"),
7948 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
7949 );
7950
7951 let (status, _) = post_json_uri(
7954 &app,
7955 "/v1/completions",
7956 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
7957 )
7958 .await;
7959 assert_eq!(
7960 status,
7961 StatusCode::NOT_IMPLEMENTED,
7962 "tokenizing an encoder must not have opened a path to generating with one"
7963 );
7964 }
7965
7966 #[tokio::test]
7971 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7972 use http_body_util::BodyExt;
7973 use tower::ServiceExt;
7974
7975 let fixture = concat!(
7976 "../frink-models/tests/fixtures/",
7977 "frink_real_moe_test.gguf"
7978 );
7979 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7980 let decoder = Decoder::from_gguf_with_expert_cache(
7981 &fixture,
7982 frink_models::config::test_moe_fixture(),
7983 Some(1024 * 1024),
7984 )
7985 .expect("MoE fixture must load store-backed");
7986
7987 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
7991 decoder.forward_token(1, 0, &mut caches);
7992
7993 let model = Model::Gguf(GgufModel {
7994 decoder: Arc::new(decoder),
7995 tokenizer: Arc::new(ServerTokenizer::Byte),
7996 stop_tokens: StopTokens::default(),
7997 bos_id: None,
7998 is_synthetic: false,
7999 chat_template: chat_template::PromptTemplate::plain(),
8000 });
8001 let state = Arc::new(test_state(
8002 model,
8003 ResponseCache::new(16, Duration::from_secs(60)),
8004 ));
8005 let app = Router::new()
8006 .route("/metrics", axum::routing::get(metrics))
8007 .route("/v1/chat/completions", post(chat_completions))
8008 .with_state(state);
8009
8010 let fetch_metrics = |app: Router| async move {
8011 let resp = app
8012 .oneshot(
8013 axum::http::Request::builder()
8014 .method("GET")
8015 .uri("/metrics")
8016 .body(axum::body::Body::empty())
8017 .unwrap(),
8018 )
8019 .await
8020 .unwrap();
8021 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
8022 String::from_utf8(bytes.to_vec()).unwrap()
8023 };
8024
8025 let after = fetch_metrics(app.clone()).await;
8026 assert!(
8027 after.contains("frink_expert_cache_misses_total"),
8028 "streaming model must expose expert-cache metrics: {after}"
8029 );
8030 let misses: u64 = after
8031 .lines()
8032 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
8033 .and_then(|l| l.split_whitespace().nth(1))
8034 .and_then(|v| v.parse().ok())
8035 .expect("misses metric line must parse");
8036 assert!(
8037 misses > 0,
8038 "decode must have read experts through the store: {after}"
8039 );
8040 }
8041
8042 fn weather_tool() -> serde_json::Value {
8043 serde_json::json!({
8044 "type": "function",
8045 "function": {
8046 "name": "get_weather",
8047 "description": "Get the current weather for a location.",
8048 "parameters": {
8049 "type": "object",
8050 "properties": {"location": {"type": "string"}},
8051 "required": ["location"]
8052 }
8053 }
8054 })
8055 }
8056
8057 fn weather_tool_def() -> ToolDef {
8058 ToolDef {
8059 kind: "function".to_string(),
8060 function: ToolFunctionDef {
8061 name: "get_weather".to_string(),
8062 description: Some("Get the current weather for a location.".to_string()),
8063 parameters: Some(serde_json::json!({
8064 "type": "object",
8065 "properties": {"location": {"type": "string"}},
8066 "required": ["location"]
8067 })),
8068 },
8069 }
8070 }
8071
8072 #[test]
8073 fn tool_preamble_mentions_every_tool_name_and_description() {
8074 let preamble = tool_preamble(&[weather_tool_def()]);
8075 assert!(preamble.contains("get_weather"));
8076 assert!(preamble.contains("Get the current weather for a location."));
8077 assert!(preamble.contains("<tool_call>"));
8078 assert!(preamble.contains("</tool_call>"));
8079 }
8080
8081 #[test]
8082 fn a_real_marker_becomes_a_structured_tool_call() {
8083 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
8084 let (message, finish) = build_response_message(
8085 text.to_string(),
8086 &[weather_tool_def()],
8087 output::OutputPosture::for_model("test-model"),
8088 "stop",
8089 );
8090 assert_eq!(finish, "tool_calls");
8091 let calls = message.tool_calls.expect("must carry a tool call");
8092 assert_eq!(calls[0].function.name, "get_weather");
8093 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
8094 assert_eq!(parsed["location"], "Paris");
8095 }
8096
8097 #[test]
8098 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
8099 let (message, finish) = build_response_message(
8100 "just an answer".to_string(),
8101 &[weather_tool_def()],
8102 output::OutputPosture::for_model("test-model"),
8103 "stop",
8104 );
8105 assert_eq!(finish, "stop");
8106 assert!(message.tool_calls.is_none());
8107 assert_eq!(message.content.as_deref(), Some("just an answer"));
8108 }
8109
8110 #[test]
8113 fn a_malformed_payload_is_not_a_tool_call() {
8114 let (message, finish) = build_response_message(
8115 "<tool_call>not valid json at all</tool_call>".to_string(),
8116 &[weather_tool_def()],
8117 output::OutputPosture::for_model("test-model"),
8118 "stop",
8119 );
8120 assert_eq!(finish, "stop");
8121 assert!(message.tool_calls.is_none());
8122 }
8123
8124 #[test]
8127 fn a_tool_that_was_never_offered_is_not_returned() {
8128 let (message, finish) = build_response_message(
8129 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
8130 &[weather_tool_def()],
8131 output::OutputPosture::for_model("test-model"),
8132 "stop",
8133 );
8134 assert_eq!(finish, "stop");
8135 assert!(message.tool_calls.is_none());
8136 }
8137
8138 #[test]
8140 fn marker_text_with_no_tools_offered_stays_content() {
8141 let (message, finish) = build_response_message(
8142 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
8143 &[],
8144 output::OutputPosture::for_model("test-model"),
8145 "stop",
8146 );
8147 assert_eq!(finish, "stop");
8148 assert!(message.tool_calls.is_none());
8149 assert!(message.content.is_some());
8150 }
8151
8152 #[test]
8156 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
8157 let opened = std::cell::Cell::new(0usize);
8158 let mut parser = crate::policy::parser::ToolCallParser::new(
8159 crate::policy::parser::ToolCallFormat::Qwen3Coder,
8160 vec![
8161 crate::policy::parser::tool_call::ToolSchema::with_parameters(
8162 "write_file",
8163 serde_json::json!({"type": "object", "properties": {
8164 "path": {"type": "string"},
8165 "contents": {"type": "string"}
8166 }}),
8167 ),
8168 ],
8169 );
8170 let wire = "<tool_call><function=write_file>\
8171 <parameter=path>\n/tmp/x\n</parameter>\
8172 <parameter=contents>\nhello world\n</parameter>\
8173 </function></tool_call>";
8174
8175 let mut deltas = Vec::new();
8176 let mut text = String::new();
8177 for piece in wire.as_bytes().chunks(7) {
8178 let chunk = String::from_utf8_lossy(piece).into_owned();
8179 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
8180 text.push_str(&more_text);
8181 deltas.extend(more);
8182 }
8183 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
8184 text.push_str(&more_text);
8185 deltas.extend(more);
8186
8187 assert_eq!(opened.get(), 1, "one call opened");
8188 assert!(text.is_empty(), "the markers are not content: {text:?}");
8189
8190 let first = &deltas[0];
8191 assert_eq!(first.index, 0);
8192 assert_eq!(first.id.as_deref(), Some("call_0"));
8193 assert_eq!(first.kind, Some("function"));
8194 assert_eq!(first.function.name.as_deref(), Some("write_file"));
8195
8196 let joined: String = deltas
8199 .iter()
8200 .filter_map(|d| d.function.arguments.clone())
8201 .collect();
8202 let parsed: serde_json::Value =
8203 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
8204 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
8205 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
8206 assert!(
8207 deltas.len() >= 3,
8208 "the arguments arrived in pieces, not whole: {}",
8209 deltas.len()
8210 );
8211 assert!(
8212 deltas[1..].iter().all(|d| d.function.name.is_none()),
8213 "only the opening delta carries identity"
8214 );
8215 }
8216
8217 #[test]
8219 fn text_around_a_streamed_call_is_still_content() {
8220 let opened = std::cell::Cell::new(0usize);
8221 let mut parser = crate::policy::parser::ToolCallParser::new(
8222 crate::policy::parser::ToolCallFormat::Qwen25,
8223 vec![crate::policy::parser::tool_call::ToolSchema::new(
8224 "get_weather",
8225 )],
8226 );
8227 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
8228 \"arguments\": {}}</tool_call> done";
8229 let mut text = String::new();
8230 for piece in wire.as_bytes().chunks(5) {
8231 let chunk = String::from_utf8_lossy(piece).into_owned();
8232 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
8233 text.push_str(&more);
8234 }
8235 let (more, _) = tool_call_deltas(parser.finish(), &opened);
8236 text.push_str(&more);
8237
8238 assert_eq!(opened.get(), 1);
8239 assert!(text.starts_with("let me check. "), "{text:?}");
8240 assert!(text.ends_with(" done"), "{text:?}");
8241 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
8242 }
8243
8244 #[test]
8247 fn a_reasoning_block_is_split_out_of_the_answer() {
8248 let (message, finish) = build_response_message(
8249 "<think>weighing it up</think>The answer is 4.".to_string(),
8250 &[],
8251 output::OutputPosture::for_model("Qwen3-8B"),
8252 "stop",
8253 );
8254 assert_eq!(finish, "stop");
8255 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
8256 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8257 }
8258
8259 #[test]
8262 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
8263 let (message, _) = build_response_message(
8264 "Use the <think> tag like this.".to_string(),
8265 &[],
8266 output::OutputPosture::for_model("llama-3.1-8b"),
8267 "stop",
8268 );
8269 assert_eq!(
8270 message.content.as_deref(),
8271 Some("Use the <think> tag like this.")
8272 );
8273 assert!(message.reasoning_content.is_none());
8274 }
8275
8276 #[tokio::test]
8284 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
8285 let app = test_app();
8286 let body = serde_json::json!({
8287 "model": "m",
8288 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8289 "max_tokens": 4,
8290 "temperature": 0,
8291 });
8292 let resp = post_json(&app, body).await;
8293 let message = &resp["choices"][0]["message"];
8294 assert!(message["content"].is_string());
8295 assert!(message.get("tool_calls").is_none());
8296 assert_eq!(resp["choices"][0]["finish_reason"], "length");
8297 assert_eq!(resp["usage"]["completion_tokens"], 4);
8298 assert_eq!(
8299 resp["usage"]["total_tokens"],
8300 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
8301 );
8302 }
8303
8304 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
8305 use http_body_util::BodyExt;
8306 use tower::ServiceExt;
8307
8308 let response = app
8309 .clone()
8310 .oneshot(
8311 axum::http::Request::builder()
8312 .method("GET")
8313 .uri(uri)
8314 .body(axum::body::Body::empty())
8315 .unwrap(),
8316 )
8317 .await
8318 .unwrap();
8319 let status = response.status();
8320 let bytes = response.into_body().collect().await.unwrap().to_bytes();
8321 (status, serde_json::from_slice(&bytes).unwrap())
8322 }
8323
8324 #[tokio::test]
8325 async fn health_answers_a_capability_handshake_not_a_boolean() {
8326 let app = test_app();
8327 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
8328 assert_eq!(status, StatusCode::OK);
8329
8330 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
8331 assert_eq!(health.state, frink_api::HealthState::Ready);
8332 assert!(health.pid > 0);
8333 assert!(health.server_time_unix_ms > 0);
8334 assert_eq!(health.last_request_age_seconds, None);
8337
8338 for id in [
8341 frink_api::health::capability::CPU,
8342 frink_api::health::capability::METAL,
8343 frink_api::health::capability::CUDA,
8344 frink_api::health::capability::REAL_WEIGHTS,
8345 frink_api::health::capability::CONTINUOUS_BATCHING,
8346 ] {
8347 let cap = health
8348 .capability(id)
8349 .unwrap_or_else(|| panic!("{id} missing"));
8350 assert!(!cap.reason.is_empty(), "{cap:?}");
8351 assert!(!cap.detail.is_empty(), "{cap:?}");
8352 }
8353 let weights = health
8357 .capability(frink_api::health::capability::REAL_WEIGHTS)
8358 .unwrap();
8359 assert!(!weights.available);
8360 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
8361 assert!(health.model.as_ref().unwrap().synthetic_weights);
8362 }
8363
8364 #[tokio::test]
8365 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
8366 let app = test_app();
8367 let _ = post_json(
8368 &app,
8369 serde_json::json!({
8370 "model": "m",
8371 "messages": [{"role": "user", "content": "\u{1}"}],
8372 "max_tokens": 1,
8373 "temperature": 0,
8374 }),
8375 )
8376 .await;
8377 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
8378 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
8379 let age = health
8380 .last_request_age_seconds
8381 .expect("a served request is evidence of liveness");
8382 assert!((0.0..5.0).contains(&age), "implausible age {age}");
8383 }
8384
8385 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
8387 use http_body_util::BodyExt;
8388 use tower::ServiceExt;
8389
8390 let response = app
8391 .clone()
8392 .oneshot(
8393 axum::http::Request::builder()
8394 .method("POST")
8395 .uri("/v1/chat/completions")
8396 .header("content-type", "application/json")
8397 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
8398 .unwrap(),
8399 )
8400 .await
8401 .unwrap();
8402 let bytes = response.into_body().collect().await.unwrap().to_bytes();
8403 String::from_utf8(bytes.to_vec())
8404 .unwrap()
8405 .lines()
8406 .filter_map(|line| line.strip_prefix("data: "))
8407 .filter(|payload| *payload != "[DONE]")
8408 .map(|payload| serde_json::from_str(payload).unwrap())
8409 .collect()
8410 }
8411
8412 #[tokio::test]
8413 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
8414 let app = test_app();
8415 let chunks = post_sse_chunks(
8416 &app,
8417 serde_json::json!({
8418 "model": "m",
8419 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8420 "max_tokens": 4,
8421 "temperature": 0,
8422 "stream": true,
8423 }),
8424 )
8425 .await;
8426
8427 assert!(!chunks.is_empty());
8428 let request_id = chunks[0]["request_id"]
8429 .as_str()
8430 .expect("the first chunk names the request")
8431 .to_string();
8432 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
8433 for (i, chunk) in chunks.iter().enumerate().skip(1) {
8436 assert!(
8437 chunk.get("request_id").is_none(),
8438 "chunk {i} repeats request_id"
8439 );
8440 }
8441 for chunk in &chunks {
8444 assert_eq!(chunk["id"], serde_json::json!(request_id));
8445 }
8446
8447 let other = post_sse_chunks(
8448 &app,
8449 serde_json::json!({
8450 "model": "m",
8451 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8452 "max_tokens": 4,
8453 "temperature": 0,
8454 "stream": true,
8455 }),
8456 )
8457 .await;
8458 assert_ne!(
8459 other[0]["request_id"].as_str().unwrap(),
8460 request_id,
8461 "two concurrent chats must not share an id"
8462 );
8463 }
8464
8465 #[tokio::test]
8466 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
8467 let app = test_app();
8468 let resp = post_json(
8469 &app,
8470 serde_json::json!({
8471 "model": "m",
8472 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8473 "max_tokens": 2,
8474 "temperature": 0,
8475 }),
8476 )
8477 .await;
8478 assert_eq!(resp["id"], resp["request_id"]);
8479 assert!(resp["request_id"]
8480 .as_str()
8481 .unwrap()
8482 .starts_with("chatcmpl-"));
8483 }
8484
8485 #[tokio::test]
8488 async fn usage_carries_separate_prefill_and_decode_timings() {
8489 let app = test_app();
8490 let resp = post_json(
8491 &app,
8492 serde_json::json!({
8493 "model": "m",
8494 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8495 "max_tokens": 4,
8496 "temperature": 0,
8497 }),
8498 )
8499 .await;
8500 let usage = &resp["usage"];
8501 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
8502 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
8503 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
8504 assert!(usage["predicted_per_second"].is_number(), "{usage}");
8505 assert!(usage.get("cached_tokens").is_none(), "{usage}");
8507 }
8508
8509 #[tokio::test]
8517 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
8518 let app = test_app();
8519 let body = serde_json::json!({
8520 "model": "m",
8521 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8522 "max_tokens": 4,
8523 "temperature": 0,
8524 "tools": [weather_tool()],
8525 });
8526 let resp = post_json(&app, body).await;
8527 let message = &resp["choices"][0]["message"];
8528 assert!(
8529 message["content"].is_string(),
8530 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
8531 );
8532 assert!(message.get("tool_calls").is_none());
8533 assert_eq!(resp["choices"][0]["finish_reason"], "length");
8537 }
8538
8539 #[tokio::test]
8543 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
8544 let app = test_app();
8545 let body = serde_json::json!({
8546 "model": "m",
8547 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8548 "max_tokens": 3,
8549 "temperature": 0,
8550 });
8551 let first = post_json(&app, body.clone()).await;
8552 assert_eq!(first["frink_cache"], "miss");
8553 let second = post_json(&app, body).await;
8554 assert_eq!(second["frink_cache"], "hit");
8555 assert_eq!(
8556 first["choices"][0]["message"]["content"],
8557 second["choices"][0]["message"]["content"]
8558 );
8559 assert_eq!(
8560 first["choices"][0]["finish_reason"],
8561 second["choices"][0]["finish_reason"]
8562 );
8563 assert_eq!(first["usage"], second["usage"]);
8564 assert_eq!(second["usage"]["completion_tokens"], 3);
8565 }
8566
8567 #[tokio::test]
8578 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
8579 let app = test_app();
8580 let plain = serde_json::json!({
8581 "model": "m",
8582 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8583 "max_tokens": 3,
8584 "temperature": 0,
8585 });
8586
8587 let first = post_json(&app, plain.clone()).await;
8588 assert_eq!(first["frink_cache"], "miss");
8589 let unconstrained = first["choices"][0]["message"]["content"]
8590 .as_str()
8591 .expect("content")
8592 .to_string();
8593
8594 let mut constrained = plain.clone();
8595 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
8596 let second = post_json(&app, constrained).await;
8597 assert_eq!(
8598 second["frink_cache"], "miss",
8599 "a grammar is part of the key, so this body has never been answered"
8600 );
8601 let constrained_answer = second["choices"][0]["message"]["content"]
8605 .as_str()
8606 .expect("content")
8607 .to_string();
8608 assert!(
8609 constrained_answer.contains("-> \"yes\"]"),
8610 "the grammar must have been compiled AND applied, not skipped \
8611 by a cache hit: {constrained_answer}"
8612 );
8613 assert_ne!(
8614 constrained_answer, unconstrained,
8615 "the constrained request was served the unconstrained answer"
8616 );
8617
8618 let third = post_json(&app, plain).await;
8622 assert_eq!(third["frink_cache"], "hit");
8623 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
8624 }
8625
8626 #[tokio::test]
8648 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
8649 let state = Arc::new(test_state(
8650 test_model_full_byte_vocab(),
8651 ResponseCache::new(1000, Duration::from_secs(3600)),
8652 ));
8653 let app = test_app_with_state(state.clone());
8654 let plain = serde_json::json!({
8655 "model": "m",
8656 "messages": [
8657 {"role": "system", "content": "Answer in JSON when it helps."},
8658 {"role": "user", "content": "\u{1}\u{2}"},
8659 ],
8660 "max_tokens": 3,
8661 "temperature": 0,
8662 });
8663
8664 let first = post_json(&app, plain.clone()).await;
8665 assert_eq!(first["frink_cache"], "miss");
8666 assert_eq!(state.cache_stats().entries, 1);
8667
8668 let mut as_json = plain.clone();
8669 as_json["response_format"] = serde_json::json!({"type": "json_object"});
8670 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
8671 assert_eq!(
8672 status,
8673 StatusCode::BAD_REQUEST,
8674 "the demo banner is not a JSON object, whoever generated it"
8675 );
8676 assert_eq!(
8677 state.cache_stats().hits,
8678 0,
8679 "a json_object request must not be answered from an entry the \
8680 JSON mask never produced"
8681 );
8682 assert_eq!(
8683 state.cache_stats().entries,
8684 2,
8685 "json_object must key its own entry, not reuse the unconstrained \
8686 one it happens to render the same prompt as"
8687 );
8688 }
8689
8690 #[tokio::test]
8702 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
8703 let app = test_app_with_state(Arc::new(test_state(
8704 test_model_full_byte_vocab_with_eos(Some(0x77)),
8705 ResponseCache::new(1000, Duration::from_secs(3600)),
8706 )));
8707 let body = serde_json::json!({
8708 "model": "m",
8709 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8710 "max_tokens": 6,
8711 "temperature": 0,
8712 });
8713
8714 let stopped = post_json(&app, body.clone()).await;
8715 assert_eq!(stopped["frink_cache"], "miss");
8716 assert_eq!(
8717 stopped["choices"][0]["finish_reason"], "stop",
8718 "the fixture is only meaningful if the model's EOS really fires here"
8719 );
8720 assert_eq!(stopped["usage"]["completion_tokens"], 1);
8721
8722 let mut ignoring = body.clone();
8723 ignoring["ignore_eos"] = serde_json::json!(true);
8724 let ran_on = post_json(&app, ignoring).await;
8725 assert_eq!(
8726 ran_on["frink_cache"], "miss",
8727 "ignore_eos is part of the key, so this body has never been answered"
8728 );
8729 assert_eq!(
8730 ran_on["usage"]["completion_tokens"], 6,
8731 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
8732 );
8733 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
8734 assert_ne!(
8735 ran_on["choices"][0]["message"]["content"],
8736 stopped["choices"][0]["message"]["content"]
8737 );
8738 }
8739
8740 #[tokio::test]
8746 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
8747 let session_app = test_app();
8748 let manual_app = test_app();
8749
8750 let turn1 = post_json(
8752 &session_app,
8753 serde_json::json!({
8754 "model": "m",
8755 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8756 "session_id": "s1",
8757 "max_tokens": 5,
8758 "temperature": 0,
8759 }),
8760 )
8761 .await;
8762 let reply1 = turn1["choices"][0]["message"]["content"]
8763 .as_str()
8764 .unwrap()
8765 .to_string();
8766
8767 let manual_turn1 = post_json(
8772 &manual_app,
8773 serde_json::json!({
8774 "model": "m",
8775 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8776 "max_tokens": 5,
8777 "temperature": 0,
8778 }),
8779 )
8780 .await;
8781 assert_eq!(
8782 manual_turn1["choices"][0]["message"]["content"]
8783 .as_str()
8784 .unwrap(),
8785 reply1
8786 );
8787
8788 let turn2 = post_json(
8790 &session_app,
8791 serde_json::json!({
8792 "model": "m",
8793 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
8794 "session_id": "s1",
8795 "max_tokens": 5,
8796 "temperature": 0,
8797 }),
8798 )
8799 .await;
8800 let reply2 = turn2["choices"][0]["message"]["content"]
8801 .as_str()
8802 .unwrap()
8803 .to_string();
8804
8805 let manual_turn2 = post_json(
8809 &manual_app,
8810 serde_json::json!({
8811 "model": "m",
8812 "messages": [
8813 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
8814 {"role": "assistant", "content": reply1},
8815 {"role": "user", "content": "\u{4}\u{5}"},
8816 ],
8817 "max_tokens": 5,
8818 "temperature": 0,
8819 }),
8820 )
8821 .await;
8822 assert_eq!(
8823 manual_turn2["choices"][0]["message"]["content"]
8824 .as_str()
8825 .unwrap(),
8826 reply2,
8827 "resuming a session must produce identical output to manually resending the full history"
8828 );
8829 }
8830
8831 #[test]
8834 fn lock_cache_recovers_from_a_poisoned_mutex() {
8835 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
8836
8837 let poison_cache = Arc::clone(&cache);
8838 let _ = std::thread::spawn(move || {
8839 let _guard = poison_cache.lock().unwrap();
8840 panic!("simulated panic while holding the lock");
8841 })
8842 .join();
8843
8844 let recovered = lock_cache(&cache);
8846 assert_eq!(recovered.stats().entries, 0);
8847 }
8848
8849 #[test]
8850 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8851 let mut req_body = serde_json::json!({
8852 "model": "m",
8853 "messages": [{"role": "user", "content": "hi"}],
8854 });
8855 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8856 assert!(
8857 req.is_cacheable(),
8858 "default (temperature 0) must be cacheable"
8859 );
8860
8861 req_body["temperature"] = serde_json::json!(0.8);
8862 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8863 assert!(
8864 !req.is_cacheable(),
8865 "unseeded sampling must never be cacheable"
8866 );
8867
8868 req_body["seed"] = serde_json::json!(42);
8869 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
8870 assert!(
8871 req.is_cacheable(),
8872 "sampling with an explicit seed is deterministic and must be cacheable"
8873 );
8874 }
8875
8876 const GRADED: &str = "{% if reasoning_effort %}\
8880 {% if reasoning_effort not in ['low','medium','high'] %}\
8881 {{ raise_exception('unsupported effort') }}\
8882 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
8883 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
8884
8885 fn graded_template() -> chat_template::PromptTemplate {
8886 chat_template::PromptTemplate::from_gguf_metadata(
8887 Some(GRADED),
8888 Some("qwen3"),
8889 false,
8890 true,
8891 None,
8892 None,
8893 )
8894 }
8895
8896 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
8897 serde_json::from_value(value).expect("request")
8898 }
8899
8900 #[test]
8907 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
8908 let req = chat_request(serde_json::json!({
8909 "model": "m",
8910 "messages": [{"role": "user", "content": "hi"}],
8911 "grammar": "root ::= \"a\"+",
8912 }));
8913 req.validate_supported_fields()
8914 .expect("a valid grammar is a valid request");
8915 let params = req
8916 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8917 .expect("a valid grammar compiles at params time too");
8918 assert!(
8919 params.grammar.is_some(),
8920 "the grammar was dropped between the wire and the sampler"
8921 );
8922 assert!(
8923 params.needs_vocab_logits(),
8924 "a grammar request that may fold lm_head into a GPU argmax is \
8925 a grammar request served unconstrained"
8926 );
8927
8928 let plain = chat_request(serde_json::json!({
8929 "model": "m",
8930 "messages": [{"role": "user", "content": "hi"}],
8931 }));
8932 assert!(plain
8933 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8934 .unwrap()
8935 .grammar
8936 .is_none());
8937 }
8938
8939 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
8940 chat_request(serde_json::json!({
8941 "model": "m",
8942 "messages": [{"role": "user", "content": "weather in Rome?"}],
8943 "tools": [weather_tool()],
8944 "tool_choice": tool_choice,
8945 }))
8946 }
8947
8948 #[test]
8952 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
8953 for choice in [
8954 serde_json::json!("required"),
8955 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
8956 ] {
8957 let req = tool_request(choice.clone());
8958 req.validate_supported_fields()
8959 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
8960 let params = req
8961 .generation_params_for_template(
8962 &graded_template(),
8963 "Qwen3-8B",
8964 crate::sampling_knobs::SamplerModel::absent(),
8965 )
8966 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
8967 let grammar = params
8968 .grammar
8969 .as_ref()
8970 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
8971 assert!(
8972 grammar.is_awaiting_trigger(),
8973 "the model must be free to think before it calls"
8974 );
8975 assert!(
8976 !grammar.allows_eog(),
8977 "{choice} must not be able to end the turn without a call"
8978 );
8979 assert!(
8985 params.needs_vocab_logits(),
8986 "{choice} would let a backend return a token id instead of logits"
8987 );
8988 assert!(
8989 !generate::greedy_gpu_fold_allowed(¶ms),
8990 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8991 );
8992 }
8993 }
8994
8995 #[test]
8997 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8998 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8999 let req = tool_request(choice.clone());
9000 req.validate_supported_fields().expect("still supported");
9001 let params = match req.generation_params_for_template(
9002 &graded_template(),
9003 "Qwen3-8B",
9004 crate::sampling_knobs::SamplerModel::absent(),
9005 ) {
9006 Ok(p) => p,
9007 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
9008 };
9009 assert!(
9010 params.grammar.is_none(),
9011 "{choice} does not force a call and must not be constrained"
9012 );
9013 }
9014 }
9015
9016 #[test]
9019 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
9020 let req = chat_request(serde_json::json!({
9022 "model": "m",
9023 "messages": [{"role": "user", "content": "hi"}],
9024 "tool_choice": "required",
9025 }));
9026 let (status, _) = req
9027 .validate_supported_fields()
9028 .expect_err("nothing to call");
9029 assert_eq!(status, StatusCode::BAD_REQUEST);
9030
9031 let req =
9033 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
9034 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
9035 assert_eq!(status, StatusCode::BAD_REQUEST);
9036 assert_eq!(body["error"]["param"], "tool_choice");
9037
9038 let req = tool_request(serde_json::json!({"type": "function"}));
9040 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
9041 assert_eq!(status, StatusCode::BAD_REQUEST);
9042
9043 let req = chat_request(serde_json::json!({
9045 "model": "m",
9046 "messages": [{"role": "user", "content": "hi"}],
9047 "tools": [weather_tool()],
9048 "tool_choice": "required",
9049 "grammar": "root ::= \"a\"+",
9050 }));
9051 let (status, _) = req
9052 .validate_supported_fields()
9053 .expect_err("a grammar and a forced call are two constraints");
9054 assert_eq!(status, StatusCode::BAD_REQUEST);
9055
9056 let req = tool_request(serde_json::json!("required"));
9062 let (status, Json(body)) = match req.generation_params_for_template(
9063 &graded_template(),
9064 "muse-glimmer-8b",
9065 crate::sampling_knobs::SamplerModel::absent(),
9066 ) {
9067 Err(e) => e,
9068 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
9069 };
9070 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
9071 assert!(
9072 body["error"]["message"]
9073 .as_str()
9074 .unwrap()
9075 .contains("muse_glimmer"),
9076 "{body}"
9077 );
9078
9079 let req = tool_request(serde_json::json!("required"));
9085 let params = req
9086 .generation_params_for_template(
9087 &graded_template(),
9088 "gemma-4-E2B-it",
9089 crate::sampling_knobs::SamplerModel::absent(),
9090 )
9091 .expect("a gemma4 forced tool_choice is served");
9092 assert!(
9093 params.grammar.is_some(),
9094 "a forced tool_choice must arrive as the generation's grammar"
9095 );
9096 }
9097
9098 #[test]
9101 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
9102 let req = chat_request(serde_json::json!({
9103 "model": "m",
9104 "messages": [{"role": "user", "content": "hi"}],
9105 "grammar": "root ::= \"a",
9106 }));
9107 let (status, Json(body)) = req
9108 .validate_supported_fields()
9109 .expect_err("this does not parse");
9110 assert_eq!(status, StatusCode::BAD_REQUEST);
9111 assert_eq!(body["error"]["param"], "grammar");
9112 assert!(
9113 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
9114 .is_err(),
9115 "and again at params time"
9116 );
9117 }
9118
9119 #[test]
9126 fn response_format_json_schema_becomes_the_requests_grammar() {
9127 let req = chat_request(serde_json::json!({
9128 "model": "m",
9129 "messages": [{"role": "user", "content": "hi"}],
9130 "response_format": {
9131 "type": "json_schema",
9132 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
9133 },
9134 }));
9135 req.validate_supported_fields()
9136 .expect("a boolean schema converts");
9137 let params = req
9138 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9139 .expect("and compiles");
9140 let grammar = params.grammar.expect("the schema is the grammar");
9141 let mut g = (*grammar).clone();
9142 g.accept_token(0, b"true").expect("a boolean is accepted");
9143 assert!(g.allows_eog(), "and completes the parse");
9144 assert!(
9145 !params.json_object,
9146 "a schema is not the json_object character-class mask"
9147 );
9148 }
9149
9150 #[test]
9154 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
9155 let req = chat_request(serde_json::json!({
9156 "model": "m",
9157 "messages": [{"role": "user", "content": "hi"}],
9158 "response_format": {
9159 "type": "json_schema",
9160 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
9161 },
9162 }));
9163 let (status, Json(body)) = req
9164 .validate_supported_fields()
9165 .expect_err("minimum has no grammar in this port");
9166 assert_eq!(status, StatusCode::BAD_REQUEST);
9167 assert!(
9168 body["error"]["message"]
9169 .as_str()
9170 .expect("a message")
9171 .contains("minimum"),
9172 "the refusal must name the keyword: {body}"
9173 );
9174 assert!(
9175 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
9176 .is_err(),
9177 "and again at params time"
9178 );
9179 }
9180
9181 #[test]
9187 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
9188 let req = chat_request(serde_json::json!({
9189 "model": "m",
9190 "messages": [{"role": "user", "content": "hi"}],
9191 "tool_choice": "required",
9192 "tools": [{
9193 "type": "function",
9194 "function": {"name": "f", "parameters": {"type": "object"}},
9195 }],
9196 "response_format": {
9197 "type": "json_schema",
9198 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
9199 },
9200 }));
9201 let (status, Json(body)) = req
9202 .validate_supported_fields()
9203 .expect_err("two constraints, one generation");
9204 assert_eq!(status, StatusCode::BAD_REQUEST);
9205 assert_eq!(body["error"]["param"], "tool_choice");
9206 }
9207
9208 #[test]
9211 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
9212 let req = chat_request(serde_json::json!({
9213 "model": "m",
9214 "messages": [{"role": "user", "content": "hi"}],
9215 }));
9216 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
9217 }
9218
9219 #[test]
9224 fn min_p_reaches_the_sampler_from_the_chat_wire() {
9225 let asked = chat_request(serde_json::json!({
9226 "model": "m",
9227 "messages": [{"role": "user", "content": "hi"}],
9228 "min_p": 0.07,
9229 }));
9230 assert_eq!(
9231 asked
9232 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
9233 .expect("knobs")
9234 .min_p,
9235 0.07
9236 );
9237
9238 let silent = chat_request(serde_json::json!({
9239 "model": "m",
9240 "messages": [{"role": "user", "content": "hi"}],
9241 }));
9242 assert_eq!(
9243 silent
9244 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
9245 .expect("knobs")
9246 .min_p,
9247 0.0,
9248 "an unset min_p must be off, not llama.cpp's CLI default"
9249 );
9250 }
9251
9252 #[test]
9261 fn no_sampler_knob_is_missing_from_the_cache_key() {
9262 let base = serde_json::json!({
9263 "model": "m",
9264 "messages": [{"role": "user", "content": "hi"}],
9265 "seed": 1,
9266 });
9267 let key_for = |body: serde_json::Value| {
9268 let req = chat_request(body);
9269 let params = req
9270 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9271 .expect("params");
9272 req.cache_key("prompt", ¶ms)
9273 };
9274 let baseline = key_for(base.clone());
9275 for (knob, value) in [
9276 ("temperature", serde_json::json!(0.5)),
9277 ("top_p", serde_json::json!(0.9)),
9278 ("min_p", serde_json::json!(0.05)),
9279 ("top_k", serde_json::json!(40)),
9280 ("repetition_penalty", serde_json::json!(1.1)),
9281 ("presence_penalty", serde_json::json!(0.3)),
9282 ("frequency_penalty", serde_json::json!(0.3)),
9283 (
9284 "samplers",
9285 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
9286 ),
9287 ] {
9288 let mut body = base.clone();
9289 body[knob] = value;
9290 assert_ne!(
9291 key_for(body),
9292 baseline,
9293 "`{knob}` is not in the cache key: two requests differing \
9294 only in it would share one cached answer"
9295 );
9296 }
9297 }
9298
9299 #[test]
9309 fn no_constraint_is_missing_from_the_cache_key() {
9310 let base = serde_json::json!({
9311 "model": "m",
9312 "messages": [{"role": "user", "content": "pick one"}],
9313 });
9314 let key_for = |body: serde_json::Value| {
9315 let req = chat_request(body);
9316 let params = req
9317 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9318 .expect("params");
9319 req.cache_key("prompt", ¶ms)
9320 };
9321 let baseline = key_for(base.clone());
9322 for (field, value) in [
9323 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
9324 (
9325 "response_format",
9326 serde_json::json!({"type": "json_object"}),
9327 ),
9328 (
9329 "response_format",
9330 serde_json::json!({"type": "json_schema", "json_schema": {
9331 "name": "answer",
9332 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
9333 }}),
9334 ),
9335 ("ignore_eos", serde_json::json!(true)),
9336 ("stop", serde_json::json!(["\n"])),
9337 ("max_tokens", serde_json::json!(7)),
9338 ] {
9339 let mut body = base.clone();
9340 body[field] = value.clone();
9341 assert_ne!(
9342 key_for(body),
9343 baseline,
9344 "`{field}: {value}` is not in the cache key: two requests \
9345 differing only in it would share one cached answer"
9346 );
9347 }
9348 }
9349
9350 #[test]
9354 fn an_explicit_zero_output_budget_is_a_client_error() {
9355 let req = chat_request(serde_json::json!({
9356 "model": "m",
9357 "messages": [{"role": "user", "content": "hi"}],
9358 "max_tokens": 0,
9359 }));
9360 let (status, body) = req.validate_supported_fields().expect_err("rejected");
9361 assert_eq!(status, StatusCode::BAD_REQUEST);
9362 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
9363 }
9364
9365 #[test]
9368 fn a_request_can_turn_thinking_off() {
9369 let template = graded_template();
9370 for body in [
9371 serde_json::json!({
9372 "model": "m",
9373 "messages": [{"role": "user", "content": "hi"}],
9374 "reasoning_effort": "none",
9375 }),
9376 serde_json::json!({
9377 "model": "m",
9378 "messages": [{"role": "user", "content": "hi"}],
9379 "thinking": {"type": "disabled"},
9380 }),
9381 ] {
9382 let kwargs = chat_request(body).resolve_template_kwargs(&template);
9383 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
9384 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
9385 assert!(!kwargs.contains_key("reasoning_effort"));
9388 }
9389 }
9390
9391 #[test]
9394 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
9395 let template = graded_template();
9396 let kwargs = chat_request(serde_json::json!({
9397 "model": "m",
9398 "messages": [{"role": "user", "content": "hi"}],
9399 "reasoning_effort": "high",
9400 "thinking": {"type": "disabled"},
9401 }))
9402 .resolve_template_kwargs(&template);
9403 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
9404 assert!(!kwargs.contains_key("reasoning_effort"));
9405 }
9406
9407 #[test]
9410 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
9411 let req = chat_request(serde_json::json!({
9412 "model": "m",
9413 "messages": [{"role": "user", "content": "hi"}],
9414 "thinking": {"type": "disable"},
9415 }));
9416 let (status, _) = req.validate_supported_fields().expect_err("rejected");
9417 assert_eq!(status, StatusCode::BAD_REQUEST);
9418 }
9419
9420 #[test]
9423 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
9424 let template = graded_template();
9425 let kwargs = chat_request(serde_json::json!({
9426 "model": "m",
9427 "messages": [{"role": "user", "content": "hi"}],
9428 "reasoning_effort": "none",
9429 "chat_template_kwargs": {"enable_thinking": true},
9430 }))
9431 .resolve_template_kwargs(&template);
9432 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
9433 }
9434
9435 #[test]
9439 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
9440 let template = graded_template();
9441 let req = chat_request(serde_json::json!({
9442 "model": "m",
9443 "messages": [{"role": "user", "content": "hi"}],
9444 "reasoning_effort": "minimal",
9445 }));
9446 let kwargs = req.resolve_template_kwargs(&template);
9447 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
9448 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
9449 assert!(prompt.starts_with("E:low|"), "{prompt}");
9450 }
9451
9452 #[test]
9456 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
9457 let template = graded_template();
9458 let req = chat_request(serde_json::json!({
9459 "model": "m",
9460 "messages": [{"role": "user", "content": "hi"}],
9461 "chat_template_kwargs": {"reasoning_effort": "none"},
9462 }));
9463 let kwargs = req.resolve_template_kwargs(&template);
9464 assert!(!kwargs.contains_key("reasoning_effort"));
9465 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
9466 assert_eq!(prompt, "hi");
9467 }
9468
9469 #[test]
9472 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
9473 let template = graded_template();
9474 let req = chat_request(serde_json::json!({
9475 "model": "m",
9476 "messages": [{"role": "user", "content": "hi"}],
9477 "reasoning_effort": "low",
9478 "chat_template_kwargs": {"reasoning_effort": "high"},
9479 }));
9480 assert_eq!(
9481 req.resolve_template_kwargs(&template)["reasoning_effort"],
9482 serde_json::json!("high")
9483 );
9484 }
9485
9486 #[test]
9490 fn offering_tools_turns_thinking_on_by_itself() {
9491 let template = graded_template();
9492 let quiet = chat_request(serde_json::json!({
9493 "model": "m",
9494 "messages": [{"role": "user", "content": "hi"}],
9495 }));
9496 assert!(!quiet
9497 .resolve_template_kwargs(&template)
9498 .contains_key("enable_thinking"));
9499
9500 let with_tools = chat_request(serde_json::json!({
9501 "model": "m",
9502 "messages": [{"role": "user", "content": "hi"}],
9503 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
9504 }));
9505 let kwargs = with_tools.resolve_template_kwargs(&template);
9506 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
9507 let prompt =
9508 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
9509 assert!(prompt.starts_with("THINK|"), "{prompt}");
9510 }
9511
9512 #[test]
9517 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
9518 let opener = chat_template::PromptTemplate::from_gguf_metadata(
9519 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
9520 Some("qwen3"),
9521 false,
9522 true,
9523 None,
9524 None,
9525 );
9526 let req = chat_request(serde_json::json!({
9527 "model": "m",
9528 "messages": [{"role": "user", "content": "hi"}],
9529 "chat_template_kwargs": {"enable_thinking": true},
9530 }));
9531 let kwargs = req.resolve_template_kwargs(&opener);
9532 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
9533 assert!(prompt.ends_with("<think>"), "{prompt}");
9534
9535 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
9538 let (message, _) = build_response_message(
9539 "weighing it up</think>Paris.".to_string(),
9540 &[],
9541 posture,
9542 "stop",
9543 );
9544 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
9545 assert_eq!(message.content.as_deref(), Some("Paris."));
9546
9547 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
9550 let (message, _) = build_response_message(
9551 "weighing it up</think>Paris.".to_string(),
9552 &[],
9553 closed,
9554 "stop",
9555 );
9556 assert_eq!(message.reasoning_content, None);
9557 }
9558
9559 #[test]
9560 fn stop_param_accepts_both_single_string_and_array() {
9561 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
9562 "model": "m",
9563 "messages": [{"role": "user", "content": "hi"}],
9564 "stop": "END",
9565 }))
9566 .unwrap();
9567 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
9568
9569 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
9570 "model": "m",
9571 "messages": [{"role": "user", "content": "hi"}],
9572 "stop": ["A", "B"],
9573 }))
9574 .unwrap();
9575 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
9576 }
9577
9578 #[test]
9579 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
9580 let model = test_model();
9581 let result = run_generation(
9582 &model,
9583 "hello",
9584 &greedy_params(4),
9585 None,
9586 None,
9587 None,
9588 None,
9589 None,
9590 None,
9591 );
9592 assert!(matches!(
9593 result,
9594 Err(generate::DecodeError::TokenOutOfVocab { .. })
9595 ));
9596 }
9597
9598 #[test]
9602 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
9603 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9605 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
9606
9607 let holder_pool = Arc::clone(&pool);
9608 let holder = std::thread::spawn(move || {
9609 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
9610 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
9612 drop(held);
9613 });
9614 std::thread::sleep(Duration::from_millis(15));
9615
9616 let config = generate::KvPoolConfig {
9617 pool,
9618 queue_wait: Duration::ZERO,
9619 };
9620 let result = run_generation(
9621 &model,
9622 &prompt,
9623 &greedy_params(4),
9624 Some(&config),
9625 None,
9626 None,
9627 None,
9628 None,
9629 None,
9630 );
9631 assert!(matches!(
9632 result,
9633 Err(generate::DecodeError::KvPoolExhausted)
9634 ));
9635
9636 let (status, _body) = decode_error_response(result.unwrap_err());
9637 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
9638 holder.join().unwrap();
9639 }
9640
9641 #[test]
9649 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
9650 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9652 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
9654 let config = generate::KvPoolConfig {
9655 pool,
9656 queue_wait: Duration::ZERO,
9657 };
9658
9659 let result = run_generation(
9660 &model,
9661 &prompt,
9662 &greedy_params(4),
9663 Some(&config),
9664 None,
9665 None,
9666 None,
9667 None,
9668 None,
9669 );
9670 let err = result.expect_err("one block cannot hold two layers' caches");
9671 assert!(
9672 matches!(
9673 &err,
9674 generate::DecodeError::KvBudgetExceeded { binding, .. }
9675 if *binding == frink_models::Ceiling::DeviceMemory.code()
9676 ),
9677 "expected an immovable device-memory refusal, got {err:?}"
9678 );
9679 let (status, _body) = decode_error_response(err);
9680 assert_eq!(status, StatusCode::BAD_REQUEST);
9681 }
9682
9683 #[test]
9689 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
9690 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
9691 queued: 512,
9692 cap: 512,
9693 });
9694 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
9695 assert_eq!(body["error"]["retry_after_seconds"], 1);
9696 let message = body["error"]["message"].as_str().expect("message");
9697 assert!(message.contains("512"), "{message}");
9698 }
9699
9700 #[test]
9701 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
9702 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
9703 token: 99,
9704 vocab_size: 32,
9705 });
9706 assert!(
9707 body["error"]["retry_after_seconds"].is_null(),
9708 "retrying a prompt this model cannot tokenize never helps"
9709 );
9710 }
9711
9712 #[test]
9713 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
9714 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
9715 token: 99,
9716 vocab_size: 32,
9717 });
9718 assert_eq!(status, StatusCode::BAD_REQUEST);
9719 }
9720
9721 #[test]
9722 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
9723 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9725 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
9726 let config = generate::KvPoolConfig {
9727 pool: pool.clone(),
9728 queue_wait: Duration::ZERO,
9729 };
9730
9731 let produced = run_generation(
9732 &model,
9733 &prompt,
9734 &greedy_params(4),
9735 Some(&config),
9736 None,
9737 None,
9738 None,
9739 None,
9740 None,
9741 )
9742 .unwrap();
9743 assert_eq!(produced.choices[0].finish, FinishReason::Length);
9744 assert_eq!(
9745 pool.lock().unwrap().free_blocks(),
9746 2,
9747 "a completed request must return its blocks to the pool"
9748 );
9749 }
9750
9751 #[tokio::test]
9756 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
9757 let model = Arc::new(test_model());
9758 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9759
9760 let mut handles = Vec::new();
9761 for _ in 0..8 {
9762 let model = Arc::clone(&model);
9763 let prompt = prompt.clone();
9764 handles.push(tokio::task::spawn_blocking(move || {
9765 run_generation(
9766 &model,
9767 &prompt,
9768 &greedy_params(6),
9769 None,
9770 None,
9771 None,
9772 None,
9773 None,
9774 None,
9775 )
9776 .unwrap()
9777 }));
9778 }
9779
9780 let mut results = Vec::new();
9781 for h in handles {
9782 results.push(h.await.unwrap());
9783 }
9784 for r in &results[1..] {
9789 assert_eq!(r.choices, results[0].choices, "choices must match");
9793 assert_eq!(
9794 r.usage.prompt_tokens, results[0].usage.prompt_tokens,
9795 "prompt token count must match"
9796 );
9797 assert_eq!(
9798 r.usage.completion_tokens, results[0].usage.completion_tokens,
9799 "completion token count must match"
9800 );
9801 }
9802 }
9803
9804 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
9811 let mut header_entries = Vec::new();
9812 let mut data = Vec::new();
9813 for (name, shape, values) in tensors {
9814 let start = data.len();
9815 for v in values {
9816 data.extend_from_slice(&v.to_le_bytes());
9817 }
9818 let end = data.len();
9819 let shape_str = shape
9820 .iter()
9821 .map(|d| d.to_string())
9822 .collect::<Vec<_>>()
9823 .join(",");
9824 header_entries.push(format!(
9825 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
9826 ));
9827 }
9828 let header = format!("{{{}}}", header_entries.join(","));
9829 let header_bytes = header.as_bytes();
9830 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
9831 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
9832 out.extend_from_slice(header_bytes);
9833 out.extend_from_slice(&data);
9834 out
9835 }
9836
9837 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9848 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9849 use frink_models::kimi_loader::KimiRealHparams;
9850 use frink_moe::{GatingFunction, MoeLayerConfig};
9851
9852 let hidden_dim = 8;
9853 let kda_num_heads = 2;
9854 let kda_head_dim = 3;
9855 let kda_proj = kda_num_heads * kda_head_dim;
9856 let conv_kernel = 4;
9857 let dense_intermediate = 5;
9858 let vocab_size = 256;
9862 let mla_num_heads = 1;
9863 let mla_q_lora_rank = 2;
9864 let mla_kv_lora_rank = 2;
9865 let mla_qk_nope_head_dim = 2;
9866 let mla_qk_rope_head_dim = 2;
9867 let mla_v_head_dim = 2;
9868
9869 let model_cfg = frink_models::ModelConfig {
9870 rope_layers: frink_models::rope_layers::RopeLayers::All,
9871 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
9872 name: "synthetic-kimi-server-test",
9873 n_layers: 1,
9874 n_mtp_blocks: 0,
9875 hidden_dim,
9876 n_heads: 1,
9877 n_kv_heads: 1,
9878 head_dim: 4,
9879 v_head_dim: None,
9880 vocab_size,
9881 rope_theta: 10000.0,
9882 rms_norm_eps: 1e-5,
9883 post_norm_eps: 1e-5,
9884 sliding_window: None,
9885 moe: MoeLayerConfig {
9886 expert_weights_scale: 1.0,
9887 routed_weight_before_ffn: false,
9888 n_experts: 1,
9889 n_experts_active: 1,
9890 n_shared_experts: 0,
9891 hidden_dim,
9892 expert_ffn_dim: 4,
9893 gating: GatingFunction::Sigmoid,
9894 norm_topk_prob: true,
9895 expert_group_count: None,
9896 expert_group_used_count: None,
9897 },
9898 n_dense_leading_layers: 1,
9903 moe_interleave_step: None,
9904 norm_function: frink_models::norm::NormFunction::Rms,
9905 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
9906 kda_layers: vec![1],
9907 full_attn_layers: vec![],
9908 mla: MlaConfig {
9909 num_heads: mla_num_heads,
9910 q_lora_rank: mla_q_lora_rank,
9911 kv_lora_rank: mla_kv_lora_rank,
9912 qk_nope_head_dim: mla_qk_nope_head_dim,
9913 qk_rope_head_dim: mla_qk_rope_head_dim,
9914 v_head_dim: mla_v_head_dim,
9915 use_output_gate: true,
9916 rope: None,
9917 },
9918 kda: KdaConfig {
9919 num_heads: kda_num_heads,
9920 head_dim: kda_head_dim,
9921 short_conv_kernel_size: conv_kernel,
9922 gate_lower_bound: -5.0,
9923 use_full_rank_gate: true,
9924 },
9925 }),
9926 rope_freqs: None,
9927 rope_attn_factor: 1.0,
9928 rope_dim: None,
9929 rope_dim_swa: None,
9930 rope_freqs_long: None,
9931 rope_freqs_short: None,
9932 rope_orig_ctx: None,
9933 rope_layout: frink_models::config::RopeLayout::Neox,
9934 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
9935 swa_layers: frink_models::swa_layers::SwaLayers::All,
9936 attn_logit_softcap: None,
9937 final_logit_softcap: None,
9938 embedding_scale: None,
9939 residual_scale: None,
9940 normed_residual_scale: None,
9941 clamp_kqv: None,
9942 attn_temperature: None,
9943 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
9944 block_sub_norms: false,
9945 parallel_residual: false,
9946 learned_positions: false,
9947 attn_value_scale: None,
9948 alibi_max_bias: None,
9949 layer_loops: None,
9950 skip_stream: false,
9951 parallel_ssm: false,
9952 swa_chunked: false,
9953 weightless_qk_norm: false,
9954 logit_multiplier: None,
9955 attention_scale: None,
9956 rope_theta_swa: None,
9957 ffn_activation: frink_models::config::FfnActivation::Swiglu,
9958 best_effort_fields: &["synthetic test config, not a real preset"],
9959 };
9960 let hp = KimiRealHparams {
9961 hidden_dim,
9962 kda_num_heads,
9963 kda_head_dim,
9964 mla_num_heads,
9965 mla_q_lora_rank,
9966 mla_kv_lora_rank,
9967 mla_qk_nope_head_dim,
9968 mla_qk_rope_head_dim,
9969 mla_v_head_dim,
9970 dense_intermediate_dim: dense_intermediate,
9971 moe_hidden_dim: hidden_dim,
9972 moe_intermediate_dim: 4,
9973 n_experts: 1,
9974 num_shared_experts: 0,
9975 };
9976
9977 let prefix = "language_model.model.layers.0";
9981 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
9982 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
9983 name: String,
9984 shape: Vec<usize>,
9985 n: usize| {
9986 tensors.push((name, shape, vec![0.01f32; n]));
9987 };
9988 push(
9989 &mut tensors,
9990 format!("{prefix}.input_layernorm.weight"),
9991 vec![hidden_dim],
9992 hidden_dim,
9993 );
9994 push(
9995 &mut tensors,
9996 format!("{prefix}.post_attention_layernorm.weight"),
9997 vec![hidden_dim],
9998 hidden_dim,
9999 );
10000 push(
10001 &mut tensors,
10002 format!("{prefix}.self_attention_res_norm.weight"),
10003 vec![hidden_dim],
10004 hidden_dim,
10005 );
10006 push(
10007 &mut tensors,
10008 format!("{prefix}.self_attention_res_proj.weight"),
10009 vec![1, hidden_dim],
10010 hidden_dim,
10011 );
10012 push(
10013 &mut tensors,
10014 format!("{prefix}.mlp_res_norm.weight"),
10015 vec![hidden_dim],
10016 hidden_dim,
10017 );
10018 push(
10019 &mut tensors,
10020 format!("{prefix}.mlp_res_proj.weight"),
10021 vec![1, hidden_dim],
10022 hidden_dim,
10023 );
10024 push(
10025 &mut tensors,
10026 format!("{prefix}.self_attn.q_proj.weight"),
10027 vec![kda_proj, hidden_dim],
10028 kda_proj * hidden_dim,
10029 );
10030 push(
10031 &mut tensors,
10032 format!("{prefix}.self_attn.k_proj.weight"),
10033 vec![kda_proj, hidden_dim],
10034 kda_proj * hidden_dim,
10035 );
10036 push(
10037 &mut tensors,
10038 format!("{prefix}.self_attn.v_proj.weight"),
10039 vec![kda_proj, hidden_dim],
10040 kda_proj * hidden_dim,
10041 );
10042 push(
10043 &mut tensors,
10044 format!("{prefix}.self_attn.q_conv1d.weight"),
10045 vec![kda_proj, 1, conv_kernel],
10046 kda_proj * conv_kernel,
10047 );
10048 push(
10049 &mut tensors,
10050 format!("{prefix}.self_attn.k_conv1d.weight"),
10051 vec![kda_proj, 1, conv_kernel],
10052 kda_proj * conv_kernel,
10053 );
10054 push(
10055 &mut tensors,
10056 format!("{prefix}.self_attn.v_conv1d.weight"),
10057 vec![kda_proj, 1, conv_kernel],
10058 kda_proj * conv_kernel,
10059 );
10060 push(
10061 &mut tensors,
10062 format!("{prefix}.self_attn.A_log"),
10063 vec![kda_num_heads],
10064 kda_num_heads,
10065 );
10066 push(
10067 &mut tensors,
10068 format!("{prefix}.self_attn.f_a_proj.weight"),
10069 vec![kda_head_dim, hidden_dim],
10070 kda_head_dim * hidden_dim,
10071 );
10072 push(
10073 &mut tensors,
10074 format!("{prefix}.self_attn.f_b_proj.weight"),
10075 vec![kda_proj, kda_head_dim],
10076 kda_proj * kda_head_dim,
10077 );
10078 push(
10079 &mut tensors,
10080 format!("{prefix}.self_attn.dt_bias"),
10081 vec![kda_proj],
10082 kda_proj,
10083 );
10084 push(
10085 &mut tensors,
10086 format!("{prefix}.self_attn.b_proj.weight"),
10087 vec![kda_num_heads, hidden_dim],
10088 kda_num_heads * hidden_dim,
10089 );
10090 push(
10091 &mut tensors,
10092 format!("{prefix}.self_attn.g_proj.weight"),
10093 vec![kda_proj, hidden_dim],
10094 kda_proj * hidden_dim,
10095 );
10096 push(
10097 &mut tensors,
10098 format!("{prefix}.self_attn.o_norm.weight"),
10099 vec![kda_head_dim],
10100 kda_head_dim,
10101 );
10102 push(
10103 &mut tensors,
10104 format!("{prefix}.self_attn.o_proj.weight"),
10105 vec![hidden_dim, kda_proj],
10106 hidden_dim * kda_proj,
10107 );
10108 push(
10109 &mut tensors,
10110 format!("{prefix}.mlp.gate_proj.weight"),
10111 vec![dense_intermediate, hidden_dim],
10112 dense_intermediate * hidden_dim,
10113 );
10114 push(
10115 &mut tensors,
10116 format!("{prefix}.mlp.up_proj.weight"),
10117 vec![dense_intermediate, hidden_dim],
10118 dense_intermediate * hidden_dim,
10119 );
10120 push(
10121 &mut tensors,
10122 format!("{prefix}.mlp.down_proj.weight"),
10123 vec![hidden_dim, dense_intermediate],
10124 hidden_dim * dense_intermediate,
10125 );
10126 push(
10127 &mut tensors,
10128 "language_model.model.embed_tokens.weight".to_string(),
10129 vec![vocab_size, hidden_dim],
10130 vocab_size * hidden_dim,
10131 );
10132 push(
10133 &mut tensors,
10134 "language_model.lm_head.weight".to_string(),
10135 vec![vocab_size, hidden_dim],
10136 vocab_size * hidden_dim,
10137 );
10138 push(
10139 &mut tensors,
10140 "language_model.model.norm.weight".to_string(),
10141 vec![hidden_dim],
10142 hidden_dim,
10143 );
10144 push(
10145 &mut tensors,
10146 "language_model.model.output_attn_res_norm.weight".to_string(),
10147 vec![hidden_dim],
10148 hidden_dim,
10149 );
10150 push(
10151 &mut tensors,
10152 "language_model.model.output_attn_res_proj.weight".to_string(),
10153 vec![1, hidden_dim],
10154 hidden_dim,
10155 );
10156
10157 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
10170 let dir = std::env::temp_dir().join(format!(
10171 "frink_server_kimi_e2e_test_{}_{}",
10172 std::process::id(),
10173 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
10174 ));
10175 std::fs::create_dir_all(&dir).unwrap();
10176 let shard_bytes = write_safetensors_shard(&tensors);
10177 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
10178 let map_entries: Vec<String> = tensors
10179 .iter()
10180 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
10181 .collect();
10182 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
10183 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
10184
10185 use base64::Engine;
10189 let vocab_lines: Vec<String> = (0..vocab_size as u32)
10190 .map(|b| {
10191 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
10192 format!("{b64} {b}")
10193 })
10194 .collect();
10195 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
10196
10197 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
10198 .expect("must load the synthetic Kimi checkpoint end to end");
10199 std::fs::remove_dir_all(&dir).ok();
10200 loaded
10201 }
10202
10203 #[test]
10211 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
10212 let loaded = build_synthetic_kimi_loaded();
10213 let state = build_app_state(
10214 StartupModels {
10215 loaded: model::LoadedModel::Kimi(loaded),
10216 embedding: None,
10217 },
10218 None,
10219 None,
10220 None,
10221 false,
10222 None,
10223 Arc::new(health::Detection::ready(health::probe_backends())),
10224 );
10225 let active = state.active().expect("a freshly built state has a model");
10226 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
10227 assert!(!active.is_synthetic());
10228
10229 let produced = run_generation(
10230 active.generative().unwrap(),
10231 "hi",
10232 &greedy_params(5),
10233 None,
10234 None,
10235 None,
10236 None,
10237 None,
10238 None,
10239 )
10240 .expect("a real Kimi checkpoint must generate without error");
10241 assert!(matches!(
10242 produced.choices[0].finish,
10243 FinishReason::Length | FinishReason::Stop
10244 ));
10245 }
10246
10247 #[test]
10260 fn a_grammar_constrains_the_engine_decode_path() {
10261 let loaded = build_synthetic_kimi_loaded();
10262 let state = build_app_state(
10263 StartupModels {
10264 loaded: model::LoadedModel::Kimi(loaded),
10265 embedding: None,
10266 },
10267 None,
10268 None,
10269 None,
10270 false,
10271 None,
10272 Arc::new(health::Detection::ready(health::probe_backends())),
10273 );
10274 let active = state.active().expect("a freshly built state has a model");
10275
10276 let run = |grammar: Option<&str>| {
10277 let mut params = greedy_params(6);
10278 params.grammar = grammar.map(|src| {
10279 Arc::new(
10280 frink_models::grammar::Grammar::from_str_with_root(src, "root")
10281 .expect("test grammar parses"),
10282 )
10283 });
10284 run_generation(
10285 active.generative().unwrap(),
10286 "hi",
10287 ¶ms,
10288 None,
10289 None,
10290 None,
10291 None,
10292 None,
10293 None,
10294 )
10295 };
10296
10297 let produced = run(None).expect("the unconstrained run must serve");
10298 let unconstrained = produced.choices[0].text.clone();
10299 assert!(
10300 unconstrained.chars().any(|c| c != 'a'),
10301 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
10302 constrained run below would prove nothing"
10303 );
10304
10305 let produced =
10306 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
10307 let one = produced.choices.into_iter().next().unwrap();
10308 let (finish, constrained) = (one.finish, one.text);
10309 assert!(
10310 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
10311 "the engine decode path served text its grammar forbids ({constrained:?}): \
10312 the constraint was dropped between `generate_engine` and the sampler"
10313 );
10314 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
10315 }
10316
10317 #[test]
10327 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
10328 let loaded = build_synthetic_kimi_loaded();
10329 let state = build_app_state(
10330 StartupModels {
10331 loaded: model::LoadedModel::Kimi(loaded),
10332 embedding: None,
10333 },
10334 None,
10335 None,
10336 None,
10337 false,
10338 None,
10339 Arc::new(health::Detection::ready(health::probe_backends())),
10340 );
10341
10342 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
10343 let kv_pool_config = generate::KvPoolConfig {
10344 pool: pool.clone(),
10345 queue_wait: Duration::ZERO,
10346 };
10347 let pc = Mutex::new(PrefixCache::new(4));
10348
10349 run_generation(
10350 state
10351 .active()
10352 .expect("a freshly built state has a model")
10353 .generative()
10354 .unwrap(),
10355 "hi",
10356 &greedy_params(5),
10357 Some(&kv_pool_config),
10358 None,
10359 Some(&pc),
10360 None,
10361 None,
10362 None,
10363 )
10364 .expect("a real Kimi checkpoint must generate without error");
10365
10366 assert_eq!(
10367 pool.lock().unwrap().free_blocks(),
10368 4,
10369 "the KV pool must be completely untouched by a Kimi request"
10370 );
10371 let stats = pc.lock().unwrap().stats();
10372 assert_eq!(
10373 stats.hits + stats.misses,
10374 0,
10375 "the prefix cache must never be consulted for a Kimi request"
10376 );
10377 }
10378}