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, 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 (vec![(finish, Vec::new())], Vec::new(), Vec::new(), usage)
2337 } else {
2338 generate::generate(
2339 &m.decoder,
2340 m.tokenizer.as_ref(),
2341 &m.stop_tokens,
2342 m.bos_id,
2343 prompt,
2344 params,
2345 kv_pool,
2346 paged_kv,
2347 prefix_cache,
2348 ceiling,
2349 |choice, chunk| {
2350 chunks[choice].push(chunk.to_string());
2351 if !synthetic {
2355 emit(choice, chunk);
2356 }
2357 },
2358 )?
2359 }
2360 }
2361 Model::Kimi(m) => generate::generate_engine(
2362 &m.engine,
2363 &m.tokenizer,
2364 &m.stop_tokens,
2365 None,
2366 prompt,
2367 params,
2368 |chunk| {
2369 chunks[0].push(chunk.to_string());
2370 if !synthetic {
2371 emit(0, chunk);
2372 }
2373 },
2374 )?,
2375 Model::Mla(m) => generate::generate_engine(
2376 &m.engine,
2377 &m.tokenizer,
2378 &m.stop_tokens,
2379 m.bos_id,
2380 prompt,
2381 params,
2382 |chunk| {
2383 chunks[0].push(chunk.to_string());
2384 if !synthetic {
2385 emit(0, chunk);
2386 }
2387 },
2388 )?,
2389 Model::Gemma4(m) => generate::generate_engine(
2390 &m.engine,
2391 &m.tokenizer,
2392 &m.stop_tokens,
2393 m.bos_id,
2394 prompt,
2395 params,
2396 |chunk| {
2397 chunks[0].push(chunk.to_string());
2398 if !synthetic {
2399 emit(0, chunk);
2400 }
2401 },
2402 )?,
2403 Model::Glm52(m) => generate::generate_engine(
2404 &m.engine,
2405 &m.tokenizer,
2406 &m.stop_tokens,
2407 m.bos_id,
2408 prompt,
2409 params,
2410 |chunk| {
2411 chunks[0].push(chunk.to_string());
2412 if !synthetic {
2413 emit(0, chunk);
2414 }
2415 },
2416 )?,
2417 };
2418
2419 let mut full = chunks[0].concat();
2420 if synthetic {
2421 full = format!(
2422 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2423 to serve a real model. Decoded ids -> {full:?}]"
2424 );
2425 emit(0, &full);
2426 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2427 emit(0, &full);
2428 }
2429
2430 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2436 let mut out: Vec<generate::GeneratedChoice> = finishes
2437 .into_iter()
2438 .zip(chunks.into_iter().map(|c| c.concat()))
2439 .map(|((finish, logprobs), text)| generate::GeneratedChoice {
2440 finish,
2441 text,
2442 logprobs,
2443 })
2444 .collect();
2445 if let Some(first) = out.first_mut() {
2446 if synthetic {
2454 first.logprobs.clear();
2455 }
2456 first.text = full;
2457 }
2458 Ok(generate::Generated {
2459 choices: out,
2460 prompt_rows,
2461 prompt_ids,
2462 usage,
2463 })
2464}
2465
2466#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2471 model: &Model,
2472 prompt: &str,
2473 params: &GenerationParams,
2474 kv_pool: Option<&generate::KvPoolConfig>,
2475 paged_kv: Option<&generate::PagedKvConfig>,
2476 prefix_cache: Option<&Mutex<PrefixCache>>,
2477 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2478 ceiling: Option<&budget::ContextCeiling>,
2479 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2480 ) -> Result<generate::Generated, generate::DecodeError> {
2483 run_generation_emit(
2484 model,
2485 prompt,
2486 params,
2487 kv_pool,
2488 paged_kv,
2489 prefix_cache,
2490 continuous_batcher,
2491 ceiling,
2492 metal_private_decode_gate,
2493 |_, _| {},
2494 )
2495}
2496
2497pub(crate) fn prompt_from_messages(
2508 messages: &[ChatMessage],
2509 template: &chat_template::PromptTemplate,
2510 tools: &[ToolDef],
2511 extra: serde_json::Map<String, serde_json::Value>,
2512) -> Result<String, ApiError> {
2513 let rendered = if tools.is_empty() || template.handles_tools() {
2514 template.render(messages, tools, extra)
2515 } else {
2516 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2517 with_preamble.push(ChatMessage {
2518 role: "system".to_string(),
2519 content: Some(MessageContent::Text(tool_preamble(tools))),
2520 tool_calls: None,
2521 tool_call_id: None,
2522 reasoning_content: None,
2523 });
2524 with_preamble.extend_from_slice(messages);
2525 template.render(&with_preamble, &[], extra)
2526 };
2527 rendered.map_err(template_error_response)
2528}
2529
2530fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2535 (
2536 StatusCode::BAD_REQUEST,
2537 Json(serde_json::json!({
2538 "error": {
2539 "message": format!("chat template failed to render: {err}"),
2540 "type": "invalid_request_error",
2541 "param": "messages",
2542 "code": null,
2543 }
2544 })),
2545 )
2546}
2547
2548fn tool_preamble(tools: &[ToolDef]) -> String {
2561 let mut out = String::from(
2562 "You can call tools to help answer the user. To call a tool, respond with \
2563 EXACTLY one line in this format and nothing else:\n\
2564 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2565 object matching that tool's parameters>}}</tool_call>\n\n\
2566 Available tools:\n",
2567 );
2568 for t in tools {
2569 out.push_str(&format!(
2570 "- {}: {}\n parameters (JSON schema): {}\n",
2571 t.function.name,
2572 t.function.description.as_deref().unwrap_or(""),
2573 t.function
2574 .parameters
2575 .as_ref()
2576 .map(|v| v.to_string())
2577 .unwrap_or_else(|| "{}".to_string()),
2578 ));
2579 }
2580 out
2581}
2582
2583fn tool_call_deltas(
2593 events: Vec<crate::policy::parser::ToolCallEvent>,
2594 opened: &std::cell::Cell<usize>,
2595) -> (String, Vec<ToolCallDelta>) {
2596 let mut text = String::new();
2597 let mut deltas = Vec::new();
2598 for event in events {
2599 match event {
2600 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2601 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2602 opened.set(opened.get().max(index + 1));
2603 deltas.push(ToolCallDelta::opening(index, name));
2604 }
2605 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2606 if !fragment.is_empty() {
2607 deltas.push(ToolCallDelta::arguments(index, fragment));
2608 }
2609 }
2610 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2611 }
2612 }
2613 (text, deltas)
2614}
2615
2616fn build_response_message(
2627 text: String,
2628 tools: &[ToolDef],
2629 posture: output::OutputPosture,
2630 base_finish: &'static str,
2631) -> (ChatCompletionResponseMessage, &'static str) {
2632 let parsed = output::parse_output(&text, tools, posture);
2633 let calls: Vec<ToolCallOut> = parsed
2634 .calls
2635 .into_iter()
2636 .enumerate()
2637 .map(|(index, call)| ToolCallOut {
2638 id: format!("call_{index}"),
2639 kind: "function",
2640 function: ToolCallFunctionOut {
2641 name: call.name,
2642 arguments: call.arguments,
2643 },
2644 })
2645 .collect();
2646 if !calls.is_empty() {
2647 return (
2648 ChatCompletionResponseMessage {
2649 role: "assistant",
2650 content: None,
2651 reasoning_content: parsed.reasoning,
2652 tool_calls: Some(calls),
2653 },
2654 "tool_calls",
2655 );
2656 }
2657 (
2658 ChatCompletionResponseMessage {
2659 role: "assistant",
2660 content: Some(parsed.content),
2661 reasoning_content: parsed.reasoning,
2662 tool_calls: None,
2663 },
2664 base_finish,
2665 )
2666}
2667
2668fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2673 let mut history = match &req.session_id {
2674 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2675 None => req.messages.clone(),
2676 };
2677 if req.json_object_mode() {
2678 inject_json_object_system_hint(&mut history);
2679 }
2680 history
2681}
2682
2683fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2684 const HINT: &str =
2685 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2686 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2687 match &mut sys.content {
2688 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2689 s.push_str("\n\n");
2690 s.push_str(HINT);
2691 }
2692 None => {
2693 sys.content = Some(MessageContent::Text(HINT.to_string()));
2694 }
2695 _ => {}
2696 }
2697 } else {
2698 messages.insert(
2699 0,
2700 ChatMessage {
2701 role: "system".to_string(),
2702 content: Some(MessageContent::Text(HINT.to_string())),
2703 tool_calls: None,
2704 tool_call_id: None,
2705 reasoning_content: None,
2706 },
2707 );
2708 }
2709}
2710
2711async fn chat_completions(
2712 State(state): State<Arc<AppState>>,
2713 headers: axum::http::HeaderMap,
2714 Json(req): Json<ChatCompletionRequest>,
2715) -> Response {
2716 let attribution = attribution::Attribution::from_headers(&headers);
2717 state
2718 .requests_total
2719 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2720 let started = std::time::Instant::now();
2721
2722 let request_id = frink_api::next_request_id();
2726 let stream = req.stream.unwrap_or(false);
2727
2728 let refusal = cache_admin::check_admission(&state)
2734 .err()
2735 .or_else(|| req.validate_supported_fields().err());
2736 if let Some(err) = refusal {
2737 state
2738 .request_errors_total
2739 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2740 let response = err.into_response();
2741 state.record_request(stats::Record {
2742 request_id: &request_id,
2743 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2744 model: state.active_model_name(),
2745 status: response.status().as_u16(),
2746 stream,
2747 duration_ms: started.elapsed().as_millis() as u64,
2748 usage: None,
2749 attribution: &attribution,
2750 });
2751 return response;
2752 }
2753
2754 let response = if stream {
2755 chat_completions_stream(
2756 Arc::clone(&state),
2757 req,
2758 request_id.clone(),
2759 started,
2760 attribution.clone(),
2761 )
2762 .await
2763 .into_response()
2764 } else {
2765 chat_completions_full(
2766 Arc::clone(&state),
2767 req,
2768 request_id.clone(),
2769 started,
2770 attribution.clone(),
2771 )
2772 .await
2773 .into_response()
2774 };
2775
2776 if response.status().is_client_error() || response.status().is_server_error() {
2777 state
2778 .request_errors_total
2779 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2780 state.record_request(stats::Record {
2784 request_id: &request_id,
2785 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2786 model: state.active_model_name(),
2789 status: response.status().as_u16(),
2790 stream,
2791 duration_ms: started.elapsed().as_millis() as u64,
2792 usage: None,
2793 attribution: &attribution,
2794 });
2795 }
2796 state.mark_request_finished();
2797
2798 response
2799}
2800
2801async fn chat_completions_full(
2802 state: Arc<AppState>,
2803 req: ChatCompletionRequest,
2804 request_id: String,
2805 started: std::time::Instant,
2806 attribution: attribution::Attribution,
2807) -> Result<Json<ChatCompletionResponse>, ApiError> {
2808 let tools_active = req.tools_active();
2809 let active = state.require_active()?;
2813 let history = resolve_history(&state, &req);
2814 let template = active.generative()?.chat_template();
2815 let kwargs = req.resolve_template_kwargs(&template);
2816 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2817 let mut params =
2825 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2826 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2827 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2828
2829 let mut generated_logprobs: Vec<crate::sampling_loop::PerTokenProbs> = Vec::new();
2833 let n_logprobs = req.n_logprobs()?;
2836 let decode_piece = |id: usize| active.decode_any(&[id]);
2838 let (completion, cache_status) = if let Some(cached) = key
2839 .as_ref()
2840 .and_then(|key| lock_cache(&state.response_cache).get(key))
2841 {
2842 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2843 (cached, "hit")
2844 } else {
2845 let produced = decode_task::buffered(
2846 decode_task::DecodeHandles::take(&state, &active)?,
2847 prompt.clone(),
2848 params,
2849 )
2850 .await?;
2851 let usage = produced.usage;
2852 let choices = produced.choices;
2853
2854 generated_logprobs = choices.iter().map(|c| c.logprobs.clone()).collect();
2859 let completion = response_cache::CachedCompletion {
2860 choices: choices.into_iter().map(|c| (c.finish, c.text)).collect(),
2861 usage,
2862 };
2863 let cache_status = match key {
2869 Some(key) => match completion.clone().cacheable() {
2873 Some(cacheable) => {
2874 tracing::debug!("cache miss for key {}", key.digest());
2875 lock_cache(&state.response_cache).put(key, cacheable);
2876 "miss"
2877 }
2878 None => "skip",
2879 },
2880 None => "skip",
2881 };
2882 (completion, cache_status)
2883 };
2884 let content = completion.first_text().to_string();
2887
2888 if req.json_object_mode() {
2889 json_mode::validate_json_object_output(&content)?;
2890 }
2891
2892 if let Some(id) = &req.session_id {
2896 state.sessions.store_reply(
2897 id,
2898 ChatMessage {
2899 role: "assistant".to_string(),
2900 content: Some(MessageContent::Text(content.clone())),
2901 tool_calls: None,
2902 tool_call_id: None,
2903 reasoning_content: None,
2904 },
2905 );
2906 }
2907
2908 let posture = output::OutputPosture::resolve_full(
2913 active.reasoning_format(),
2914 active.tool_call_format(),
2915 &prompt,
2916 );
2917 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2918 let wanted = req.unimplemented.n.unwrap_or(1).max(1) as usize;
2922 let ranked: Vec<(generate::FinishReason, String)> = if completion.choices.len() > wanted {
2923 let scored: Vec<crate::generate::GeneratedChoice> = completion
2924 .choices
2925 .into_iter()
2926 .zip(
2927 generated_logprobs
2928 .iter()
2929 .cloned()
2930 .chain(std::iter::repeat(Vec::new())),
2931 )
2932 .map(
2933 |((finish, text), logprobs)| crate::generate::GeneratedChoice {
2934 finish,
2935 text,
2936 logprobs,
2937 },
2938 )
2939 .collect();
2940 let best = crate::best_of::take_best(scored, wanted);
2941 generated_logprobs = best.iter().map(|c| c.logprobs.clone()).collect();
2942 best.into_iter().map(|c| (c.finish, c.text)).collect()
2943 } else {
2944 completion.choices
2945 };
2946 let rendered: Vec<ChatCompletionChoice> = ranked
2947 .into_iter()
2948 .enumerate()
2949 .map(|(index, (finish, text))| {
2950 let (message, finish_reason) =
2951 build_response_message(text, tools, posture, finish.as_str());
2952 ChatCompletionChoice {
2953 index,
2954 message,
2955 finish_reason,
2956 logprobs: n_logprobs.map(|k| {
2957 crate::logprobs::render_chat(
2958 generated_logprobs.get(index).unwrap_or(&Vec::new()),
2959 Some(k),
2960 &decode_piece,
2961 )
2962 }),
2963 }
2964 })
2965 .collect();
2966
2967 state.record_request(stats::Record {
2968 request_id: &request_id,
2969 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2970 model: Some(active.name().to_string()),
2973 status: 200,
2974 stream: false,
2975 duration_ms: started.elapsed().as_millis() as u64,
2976 usage: Some(&completion.usage),
2977 attribution: &attribution,
2978 });
2979
2980 Ok(Json(ChatCompletionResponse {
2981 id: request_id.clone(),
2982 request_id,
2983 object: "chat.completion",
2984 model: req.model,
2985 choices: rendered,
2986 usage: completion.usage,
2987 frink_cache: cache_status,
2988 }))
2989}
2990
2991async fn chat_completions_stream(
2992 state: Arc<AppState>,
2993 req: ChatCompletionRequest,
2994 request_id: String,
2995 started: std::time::Instant,
2996 attribution: attribution::Attribution,
2997) -> Result<Response, ApiError> {
2998 let tools_active = req.tools_active();
3008 let active = state.require_active()?;
3012 let history = resolve_history(&state, &req);
3013 let template = active.generative()?.chat_template();
3014 let kwargs = req.resolve_template_kwargs(&template);
3015 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
3016 let model_name = req.model.clone();
3017 let session_id = req.session_id.clone();
3018 let sessions = state.sessions.clone();
3019
3020 let model = Arc::clone(active.generative()?);
3021 let kv_pool = state.kv_pool.clone();
3022 let paged_kv = state.paged_kv.clone();
3023 let prefix_cache = state.prefix_cache.clone();
3024 let batcher = active.batcher.clone();
3025 let ceiling = active.ceiling.clone();
3026 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3027 let mut params =
3028 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
3029 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
3030 params.interleave_choices = params.n > 1;
3036 let stats_state = Arc::clone(&state);
3037 let served_model = active.name().to_string();
3040 let posture = output::OutputPosture::resolve_full(
3044 active.reasoning_format(),
3045 active.tool_call_format(),
3046 &prompt,
3047 );
3048 let offered_tools: Vec<ToolDef> = if tools_active {
3051 req.tools.clone()
3052 } else {
3053 Vec::new()
3054 };
3055
3056 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3061 params.cancel = Some(cancel_token.clone());
3062
3063 let overlap = true;
3072
3073 let slot = req
3077 .stream_resumable
3078 .unwrap_or(false)
3079 .then(|| state.streams.register(&request_id));
3080 let emitter = resume::Emitter::new(slot);
3081
3082 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3083 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3087 id: request_id.clone(),
3088 request_id: None,
3089 object: "chat.completion.chunk",
3090 model: model_name.clone(),
3091 choices: vec![ChatCompletionChunkChoice {
3092 index: 0,
3093 delta: ChatCompletionChunkDelta {
3094 role: None,
3095 content: None,
3096 reasoning_content: None,
3097 tool_calls: None,
3098 },
3099 finish_reason: None,
3100 }],
3101 usage: None,
3102 });
3103
3104 tokio::task::spawn_blocking(move || {
3105 let _cancel_guard = cancel_guard;
3108 let tx_chunks = tx.clone();
3109 let orphan_timeout = sse::orphan_timeout_from_env();
3114 let head_request_id = request_id.clone();
3115 let announced = std::cell::Cell::new(false);
3119 let emitters: Rc<RefCell<Vec<crate::chat_stream_choice::ChoiceEmitter>>> =
3125 Rc::new(RefCell::new(
3126 (0..params.n.max(1))
3127 .map(|_| {
3128 crate::chat_stream_choice::ChoiceEmitter::new(
3129 posture.reasoning_parser(),
3130 tools_active.then(|| posture.tool_call_parser(&offered_tools)),
3131 )
3132 })
3133 .collect(),
3134 ));
3135 let emit_choices = Rc::clone(&emitters);
3136 let result = run_generation_emit(
3137 &model,
3138 &prompt,
3139 ¶ms,
3140 kv_pool.as_ref(),
3141 paged_kv.as_ref(),
3142 prefix_cache.as_deref(),
3143 batcher.as_ref(),
3144 ceiling.as_deref(),
3145 metal_private_decode_gate.as_deref(),
3146 |choice, chunk| {
3147 if !overlap || chunk.is_empty() {
3148 return;
3149 }
3150 let mut held = emit_choices.borrow_mut();
3151 let Some(emitter_state) = held.get_mut(choice) else {
3152 return;
3153 };
3154 let delta = emitter_state.push(chunk);
3155 if delta.is_empty() {
3156 return;
3157 }
3158 let request_id = (!announced.get()).then(|| {
3161 announced.set(true);
3162 head_request_id.clone()
3163 });
3164 let wire = delta.into_choice(choice, emitter_state.start());
3165 drop(held);
3166 let payload = ChatCompletionChunk {
3167 id: head_request_id.clone(),
3168 request_id,
3169 object: "chat.completion.chunk",
3170 model: model_name.clone(),
3171 choices: vec![wire],
3172 usage: None,
3173 };
3174 if let Err(why) =
3183 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3184 {
3185 if why == sse::SendFailure::Orphaned {
3186 tracing::warn!(
3187 "SSE stream {head_request_id} accepted nothing for the orphan \
3188 deadline; treating it as abandoned"
3189 );
3190 }
3191 if !emitter.is_resumable() {
3203 cancel_token.cancel();
3204 }
3205 }
3206 },
3207 );
3208
3209 let mut pending_request_id = (!announced.get()).then(|| request_id.clone());
3215
3216 match result {
3217 Ok(generated) => {
3218 let usage = generated.usage;
3219 let produced: Vec<(generate::FinishReason, String)> = generated
3220 .choices
3221 .into_iter()
3222 .map(|c| (c.finish, c.text))
3223 .collect();
3224 assert!(
3225 !produced.is_empty(),
3226 "a generation produces at least one choice"
3227 );
3228 if let Some(id) = &session_id {
3233 sessions.store_reply(
3234 id,
3235 ChatMessage {
3236 role: "assistant".to_string(),
3237 content: Some(MessageContent::Text(produced[0].1.clone())),
3238 tool_calls: None,
3239 tool_call_id: None,
3240 reasoning_content: None,
3241 },
3242 );
3243 }
3244 for (index, (finish, full_text)) in produced.iter().enumerate() {
3245 let (finish, full_text) = (finish.clone(), full_text.as_str());
3246 let mut streamed_finish: Option<&'static str> = None;
3252 if overlap {
3253 let (tail, first, opened) = {
3254 let mut held = emitters.borrow_mut();
3255 let state = &mut held[index];
3256 let tail = state.flush();
3257 (tail, state.start(), state.opened_calls())
3258 };
3259 if !tail.is_empty() {
3260 let payload = ChatCompletionChunk {
3261 id: request_id.clone(),
3262 request_id: pending_request_id.take(),
3263 object: "chat.completion.chunk",
3264 model: model_name.clone(),
3265 choices: vec![tail.into_choice(index, first)],
3266 usage: None,
3267 };
3268 let _ = sse::send_or_orphan(
3269 &tx,
3270 Ok(emitter.event(&payload)),
3271 orphan_timeout,
3272 );
3273 }
3274 if opened > 0 {
3275 streamed_finish = Some("tool_calls");
3276 }
3277 } else {
3278 let parsed = output::parse_output(full_text, &offered_tools, posture);
3281 let tool_calls: Vec<ToolCallDelta> = parsed
3282 .calls
3283 .iter()
3284 .enumerate()
3285 .map(|(index, call)| {
3286 ToolCallDelta::whole(
3287 index,
3288 call.name.clone(),
3289 call.arguments.clone(),
3290 )
3291 })
3292 .collect();
3293 if !tool_calls.is_empty() {
3294 streamed_finish = Some("tool_calls");
3295 }
3296 if !tool_calls.is_empty()
3297 || !parsed.content.is_empty()
3298 || parsed.reasoning.is_some()
3299 {
3300 let payload = ChatCompletionChunk {
3301 id: request_id.clone(),
3302 request_id: pending_request_id.take(),
3303 object: "chat.completion.chunk",
3304 model: model_name.clone(),
3305 choices: vec![ChatCompletionChunkChoice {
3306 index,
3307 delta: ChatCompletionChunkDelta {
3308 role: Some("assistant"),
3309 content: (!parsed.content.is_empty()
3310 && tool_calls.is_empty())
3311 .then(|| parsed.content.clone()),
3312 reasoning_content: parsed.reasoning.clone(),
3313 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3314 },
3315 finish_reason: None,
3316 }],
3317 usage: None,
3318 };
3319 let _ = sse::send_or_orphan(
3320 &tx,
3321 Ok(emitter.event(&payload)),
3322 orphan_timeout,
3323 );
3324 }
3325 }
3326 let final_finish_reason = match streamed_finish {
3330 Some(reason) if finish.as_str() != "length" => reason,
3331 _ => finish.as_str(),
3332 };
3333 let last = index + 1 == produced.len();
3337 let final_payload = ChatCompletionChunk {
3338 id: request_id.clone(),
3339 request_id: pending_request_id.take(),
3340 object: "chat.completion.chunk",
3341 model: model_name.clone(),
3342 choices: vec![ChatCompletionChunkChoice {
3343 index,
3344 delta: ChatCompletionChunkDelta {
3345 role: None,
3346 content: None,
3347 reasoning_content: None,
3348 tool_calls: None,
3349 },
3350 finish_reason: Some(final_finish_reason),
3351 }],
3352 usage: last.then(|| usage.clone()),
3353 };
3354 let _ =
3355 sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3356 }
3357 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3358 stats_state.record_request(stats::Record {
3363 request_id: &request_id,
3364 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3365 model: Some(served_model.clone()),
3366 status: 200,
3367 stream: true,
3368 duration_ms: started.elapsed().as_millis() as u64,
3369 usage: Some(&usage),
3370 attribution: &attribution,
3371 });
3372 }
3373 Err(e) => {
3374 tracing::warn!("decode error on streamed request {request_id}: {e}");
3375 stats_state.record_request(stats::Record {
3381 request_id: &request_id,
3382 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3383 model: Some(served_model.clone()),
3384 status: 500,
3385 stream: true,
3386 duration_ms: started.elapsed().as_millis() as u64,
3387 usage: None,
3388 attribution: &attribution,
3389 });
3390 let payload = ChatCompletionChunk {
3391 id: request_id.clone(),
3392 request_id: pending_request_id.take(),
3393 object: "chat.completion.chunk",
3394 model: model_name,
3395 choices: vec![ChatCompletionChunkChoice {
3396 index: 0,
3397 delta: ChatCompletionChunkDelta {
3398 role: Some("assistant"),
3399 content: Some(format!("[error: {e}]")),
3400 reasoning_content: None,
3401 tool_calls: None,
3402 },
3403 finish_reason: Some("stop"),
3404 }],
3405 usage: None,
3406 };
3407 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3408 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3409 }
3410 }
3411 drop(emitter);
3415 });
3416
3417 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3418 Ok((
3441 [(
3442 axum::http::HeaderName::from_static("x-accel-buffering"),
3443 axum::http::HeaderValue::from_static("no"),
3444 )],
3445 Sse::new(stream),
3446 )
3447 .into_response())
3448}
3449
3450fn protected_routes() -> Router<Arc<AppState>> {
3481 use frink_api::routes;
3482
3483 Router::new()
3484 .route(routes::V1_MODELS, get(list_models))
3485 .route(routes::V1_RESPONSES, post(responses::responses))
3489 .route(
3490 &axum_path(routes::V1_RESPONSE),
3491 get(responses::responses_get),
3492 )
3493 .route(
3494 &axum_path(routes::V1_RESPONSE_CANCEL),
3495 post(responses::responses_cancel),
3496 )
3497 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3498 .route(routes::V1_STATS, get(serving_stats))
3499 .route(routes::V1_REQUESTS, get(recent_requests))
3500 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3501 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3502 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3503 .route(
3504 routes::LORA_ADAPTERS,
3505 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3506 )
3507 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3508 .route(routes::V1_CANCEL, post(cancel_generation))
3512 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3517 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3518 .route(routes::V1_MESSAGES, post(anthropic::messages))
3519 .route(
3520 routes::V1_MESSAGES_COUNT_TOKENS,
3521 post(anthropic::count_tokens),
3522 )
3523 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3524 .route(routes::COMPLETION, post(completion::completion))
3529 .route(routes::COMPLETIONS, post(completion::completion))
3530 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3531 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3532 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3538 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3539 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3540 .route(routes::V1_RERANK, post(rerank::rerank))
3544 .route(routes::RERANK, post(rerank::rerank))
3545 .route(routes::CACHE_STATS, get(cache_stats))
3546 .route(routes::METRICS, get(metrics))
3547 .route(routes::ADMIN_MODELS, get(admin::models))
3552 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3553 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3554 .route(routes::SLEEP, post(admin::sleep))
3558 .route(routes::WAKE_UP, post(admin::wake_up))
3559 .route(routes::IS_SLEEPING, get(admin::is_sleeping))
3560 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3561 .route(routes::ADMIN_TASKS, get(admin::tasks))
3562 .route(&admin::cancel_route(), post(admin::cancel_task))
3563 .route(routes::ADMIN_STATS, get(admin::stats))
3564 .merge(conversations::router())
3568}
3569
3570fn axum_path(template: &str) -> String {
3571 let mut out = String::with_capacity(template.len());
3572 let mut rest = template;
3573 while let Some(open) = rest.find('{') {
3574 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3575 break;
3576 };
3577 out.push_str(&rest[..open]);
3578 out.push(':');
3579 out.push_str(&rest[open + 1..close]);
3580 rest = &rest[close + 1..];
3581 }
3582 out.push_str(rest);
3583 out
3584}
3585
3586async fn cancel_generation(
3595 State(state): State<Arc<AppState>>,
3596 Json(req): Json<frink_api::CancelGenerationRequest>,
3597) -> Response {
3598 let cancelled = state.cancels.cancel(&req.request_id);
3599 let status = if cancelled {
3600 StatusCode::OK
3601 } else {
3602 StatusCode::NOT_FOUND
3603 };
3604 let detail = if cancelled {
3605 "the generation was asked to stop; it ends at its next token".to_string()
3606 } else {
3607 "no generation with that request_id is running -- it has already \
3608 finished, was never issued, or was served by a path that does \
3609 not register for cancellation"
3610 .to_string()
3611 };
3612 (
3613 status,
3614 Json(frink_api::CancelGenerationResponse {
3615 request_id: req.request_id,
3616 cancelled,
3617 detail,
3618 }),
3619 )
3620 .into_response()
3621}
3622
3623type Activated = (
3627 Loaded,
3628 Option<serving::batch::ContinuousBatcher>,
3629 Option<Arc<budget::ContextCeiling>>,
3630);
3631
3632fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3650 let mut batcher = serving::batch::BatcherConfig::from_env();
3651 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3652 return batcher;
3655 }
3656 let Some(path) = path else {
3657 return batcher;
3658 };
3659 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3665 let Some((priced, gguf_ctx, source)) = priced else {
3666 return batcher;
3667 };
3668 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3669 tracing::warn!(
3673 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3674 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3675 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3676 admit on a number you choose.",
3677 priced.weights_bytes,
3678 priced.device_budget_bytes,
3679 );
3680 return batcher;
3681 };
3682 tracing::info!("{source}");
3683 tracing::info!("{}", derived.fit);
3684 let adopted = budget::apply_derived(&mut batcher, &derived);
3685 if adopted.max_context {
3686 tracing::info!(
3687 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3688 override with FRINK_CB_MAX_CONTEXT",
3689 derived.max_context
3690 );
3691 }
3692 if adopted.kv_blocks {
3693 tracing::info!(
3694 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3695 derived.kv_blocks,
3696 batcher.kv_block_size
3697 );
3698 }
3699 if let Some(narrowed) = adopted.max_context_narrowed {
3700 tracing::info!(
3701 "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",
3702 batcher.kv_blocks.unwrap_or_default(),
3703 batcher.kv_block_size
3704 );
3705 }
3706 batcher
3707}
3708
3709pub(crate) fn activate_loaded_model(
3718 loaded: model::LoadedModel,
3719 enable_continuous_batching: bool,
3720 path: Option<&str>,
3721 paged_kv: Option<&generate::PagedKvConfig>,
3722) -> Activated {
3723 match loaded {
3724 model::LoadedModel::Gguf(g) => {
3725 let decoder = Arc::new(g.decoder);
3726 let tokenizer = Arc::new(g.tokenizer);
3727 let config = price_batcher_config(path);
3728 let shape =
3734 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3735 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3736 let batcher = if enable_continuous_batching {
3737 tracing::info!(
3738 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3739 (stop sequences use the same pending-buffer trim as the private generate loop)"
3740 );
3741 let tok = Arc::clone(&tokenizer);
3742 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3743 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3744 Arc::clone(&decoder),
3745 decode,
3746 config,
3747 Arc::clone(&ceiling),
3748 paged_kv.cloned(),
3749 ))
3750 } else {
3751 None
3752 };
3753 (
3754 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3755 decoder,
3756 tokenizer,
3757 stop_tokens: g.stop_tokens,
3758 bos_id: g.bos_id,
3759 is_synthetic: g.is_synthetic,
3760 chat_template: g.chat_template,
3761 }))),
3762 batcher,
3763 Some(ceiling),
3764 )
3765 }
3766 model::LoadedModel::Kimi(k) => (
3767 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3768 engine: k.engine,
3769 tokenizer: k.tokenizer,
3770 stop_tokens: k.stop_tokens,
3771 chat_template: k.chat_template,
3772 }))),
3773 None,
3774 None,
3775 ),
3776 model::LoadedModel::Mla(m) => (
3777 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3778 engine: m.engine,
3779 tokenizer: m.tokenizer,
3780 stop_tokens: m.stop_tokens,
3781 bos_id: m.bos_id,
3782 name: m.name,
3783 chat_template: m.chat_template,
3784 }))),
3785 None,
3786 None,
3787 ),
3788 model::LoadedModel::Gemma4(m) => (
3789 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3790 engine: m.engine,
3791 tokenizer: m.tokenizer,
3792 stop_tokens: m.stop_tokens,
3793 bos_id: m.bos_id,
3794 name: m.name,
3795 chat_template: m.chat_template,
3796 }))),
3797 None,
3798 None,
3799 ),
3800 model::LoadedModel::Glm52(g) => (
3801 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3802 engine: g.engine,
3803 tokenizer: g.tokenizer,
3804 stop_tokens: g.stop_tokens,
3805 bos_id: g.bos_id,
3806 name: g.name,
3807 chat_template: g.chat_template,
3808 }))),
3809 None,
3810 None,
3811 ),
3812 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3817 }
3818}
3819
3820struct StartupModels {
3827 loaded: model::LoadedModel,
3828 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3829}
3830
3831fn continuous_batching_env() -> Option<bool> {
3832 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3833 .ok()
3834 .map(|v| v.trim().to_ascii_lowercase())
3835 .as_deref()
3836 {
3837 None => None,
3838 Some("1" | "true" | "yes" | "on") => Some(true),
3839 Some("0" | "false" | "no" | "off") => Some(false),
3840 _ => None,
3841 }
3842}
3843
3844fn metal_private_decode_active() -> bool {
3845 #[cfg(feature = "metal")]
3846 {
3847 BUILT_WITH_METAL
3848 && frink_metal::attn::metal_attn_enabled()
3849 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3850 }
3851 #[cfg(not(feature = "metal"))]
3852 {
3853 false
3854 }
3855}
3856
3857fn continuous_batching_compatible(
3858 loaded: &model::LoadedModel,
3859 kv_pool: &Option<generate::KvPoolConfig>,
3860 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3861 paged_kv: &Option<generate::PagedKvConfig>,
3862) -> bool {
3863 matches!(loaded, model::LoadedModel::Gguf(_))
3864 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3865}
3866
3867fn resolve_continuous_batching_enabled(
3868 loaded: &model::LoadedModel,
3869 kv_pool: &Option<generate::KvPoolConfig>,
3870 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3871 paged_kv: &Option<generate::PagedKvConfig>,
3872) -> bool {
3873 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3874 return false;
3875 }
3876 match continuous_batching_env() {
3877 Some(true) => true,
3878 Some(false) => false,
3879 None => metal_private_decode_active(),
3880 }
3881}
3882
3883fn acquire_metal_private_decode_gate(
3884 gate: Option<&std::sync::Mutex<()>>,
3885 used_batcher: bool,
3886) -> Option<std::sync::MutexGuard<'_, ()>> {
3887 if used_batcher {
3888 None
3889 } else {
3890 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3891 }
3892}
3893
3894fn build_app_state(
3895 models: StartupModels,
3896 kv_pool: Option<generate::KvPoolConfig>,
3897 paged_kv: Option<generate::PagedKvConfig>,
3898 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3899 enable_continuous_batching: bool,
3900 mcp: Option<mcp::LoadedMcpConfig>,
3901 detection: Arc<health::Detection>,
3902) -> AppState {
3903 let StartupModels { loaded, embedding } = models;
3904 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3905 let (loaded, batcher, ceiling) = activate_loaded_model(
3906 loaded,
3907 enable_continuous_batching,
3908 configured_path.as_deref(),
3909 paged_kv.as_ref(),
3910 );
3911 let id = startup_model_id();
3917 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3918 {
3919 None
3920 } else {
3921 tracing::info!(
3922 "Metal private-loop decode will serialize concurrent requests until \
3923 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3924 );
3925 Some(Arc::new(std::sync::Mutex::new(())))
3926 };
3927 AppState {
3928 slept: Mutex::new(None),
3929 embedding,
3930 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3931 id,
3932 loaded,
3933 batcher,
3934 ceiling,
3935 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3936 }))),
3937 paged_kv,
3938 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3939 tasks: Arc::new(tasks::TaskRegistry::new()),
3940 cancels: Arc::new(cancel::CancelRegistry::new()),
3941 stats: stats::Stats::new(),
3942 streams: resume::StreamRegistry::new(),
3943 model_dir: admin::model_dirs().into_iter().next(),
3944 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3945 kv_pool,
3946 prefix_cache,
3947 sessions: session::SessionStore::new(),
3948 requests_total: std::sync::atomic::AtomicU64::new(0),
3949 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3950 started_at: std::time::Instant::now(),
3951 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3952 detection,
3953 mcp,
3954 continuous_batching_enabled: enable_continuous_batching,
3955 metal_private_decode_gate,
3956 loading_model: Mutex::new(None),
3957 last_load_error: Mutex::new(None),
3958 serving: Mutex::new(crate::stats::ServingStats::default()),
3959 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3960 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3961 started_unix: unix_now(),
3962 }
3963}
3964
3965fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3973 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3974 return Ok(None);
3975 };
3976 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3977 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3978 tracing::info!(
3979 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3980 model.name(),
3981 model.architecture(),
3982 model.n_embd(),
3983 model.pooling_type().name(),
3984 model.n_ctx_train(),
3985 );
3986 Ok(Some(Arc::new(model)))
3987}
3988
3989fn unix_now() -> u64 {
3994 std::time::SystemTime::now()
3995 .duration_since(std::time::UNIX_EPOCH)
3996 .map(|d| d.as_secs())
3997 .unwrap_or(0)
3998}
3999
4000fn startup_model_id() -> Option<String> {
4005 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
4006 let configured = std::fs::canonicalize(&configured).ok()?;
4007 admin::discover(&admin::model_dirs())
4008 .into_iter()
4009 .find(|d| {
4010 std::fs::canonicalize(&d.path)
4011 .map(|p| p == configured)
4012 .unwrap_or(false)
4013 })
4014 .map(|d| d.id)
4015}
4016
4017fn init_cpu_pool() {
4026 match frink_core::threads::init_cpu_pool() {
4027 Some(n) => eprintln!(
4028 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
4029 frink_core::threads::perf_core_count()
4030 ),
4031 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
4032 }
4033}
4034
4035fn announce_ready(addr: SocketAddr, scheme: &str) {
4048 use std::io::Write;
4049 let ready =
4050 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
4051 let mut stdout = std::io::stdout().lock();
4052 let _ = writeln!(stdout, "{}", ready.to_line());
4053 let _ = stdout.flush();
4054}
4055
4056async fn shutdown_signal(exit_on_stdin_close: bool) {
4068 if !exit_on_stdin_close {
4069 std::future::pending::<()>().await;
4070 return;
4071 }
4072 let _ = tokio::task::spawn_blocking(|| {
4073 use std::io::Read;
4074 let mut sink = [0u8; 256];
4075 let mut stdin = std::io::stdin().lock();
4076 loop {
4077 match stdin.read(&mut sink) {
4078 Ok(0) => break,
4080 Ok(_) => continue,
4082 Err(e) => {
4083 tracing::warn!("stdin read failed ({e}); treating it as closed");
4084 break;
4085 }
4086 }
4087 }
4088 })
4089 .await;
4090 tracing::info!("stdin closed; shutting down");
4091}
4092
4093fn tokio_worker_threads() -> usize {
4098 std::env::var("FRINK_TOKIO_WORKERS")
4099 .ok()
4100 .and_then(|v| v.trim().parse::<usize>().ok())
4101 .filter(|n| *n > 0)
4102 .unwrap_or(2)
4103}
4104
4105fn install_ring_crypto_provider() {
4122 let _ = rustls::crypto::ring::default_provider().install_default();
4123}
4124
4125pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4131 if args.list_devices {
4132 frink_models::devices::print_available_devices();
4133 return Ok(());
4134 }
4135 apply_cli_overrides(&args)?;
4136
4137 let _instance = {
4141 use frink_core::instance::{register, InstancePolicy};
4142 let policy = if args.allow_multiple_instances {
4143 InstancePolicy::Multi
4144 } else {
4145 InstancePolicy::from_env_or(InstancePolicy::Single)
4146 };
4147 let model = std::env::var("FRINK_MODEL_PATH").ok();
4148 register(
4149 "server",
4150 model.as_deref(),
4151 frink_core::instance::current_backend(),
4152 policy,
4153 )
4154 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4155 };
4156
4157 let journal = journal::Journal::from_env();
4158 eprintln!(
4159 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4160 journal.path()
4161 );
4162 journal.append(&journal::Record::session_start(
4163 env!("CARGO_PKG_VERSION"),
4164 std::process::id(),
4165 ));
4166 journal::install_panic_hook(journal.clone());
4167
4168 let mcp_config_path = args.mcp_config.clone();
4169 let exit_on_stdin_close = args.exit_on_stdin_close
4170 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4171 .map(|v| v == "1")
4172 .unwrap_or(false);
4173
4174 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4178 init_cpu_pool();
4179
4180 let runtime = tokio::runtime::Builder::new_multi_thread()
4181 .worker_threads(tokio_worker_threads())
4182 .enable_all()
4183 .build()?;
4184 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4185
4186 let reason = match &result {
4187 Ok(()) => "normal".to_string(),
4188 Err(e) => e.to_string(),
4189 };
4190 journal.append(&journal::Record::session_exit(reason));
4191
4192 runtime.shutdown_background();
4197
4198 result
4199}
4200
4201async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4202 let _ = tracing_subscriber::fmt::try_init();
4208
4209 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4219 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4220 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4221 .map(|v| v == "1")
4222 .unwrap_or(false);
4223 if let Err(msg) =
4224 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4225 {
4226 anyhow::bail!(msg);
4227 }
4228
4229 let embedding_model = load_embedding_model()?;
4236
4237 let mut loaded = model::load()?;
4238 match &loaded {
4239 model::LoadedModel::Gguf(g) => tracing::info!(
4240 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4241 g.decoder.config.name,
4242 g.is_synthetic,
4243 g.tokenizer.kind()
4244 ),
4245 model::LoadedModel::Kimi(k) => tracing::info!(
4246 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4247 k.tokenizer.vocab_size()
4248 ),
4249 model::LoadedModel::Mla(m) => tracing::info!(
4250 "loaded MLA GGUF '{}' (tokenizer={})",
4251 m.name,
4252 m.tokenizer.kind()
4253 ),
4254 model::LoadedModel::Gemma4(m) => tracing::info!(
4255 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4256 m.name,
4257 m.tokenizer.kind()
4258 ),
4259 model::LoadedModel::Glm52(g) => tracing::info!(
4260 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4261 g.name,
4262 g.tokenizer.kind()
4263 ),
4264 model::LoadedModel::Encoder(_) => {}
4267 }
4268 let metal_default_moe_budget = {
4276 #[cfg(feature = "metal")]
4277 {
4278 frink_core::metal_dense_enabled()
4279 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4280 }
4281 #[cfg(not(feature = "metal"))]
4282 {
4283 false
4284 }
4285 };
4286 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4287 let budget: u64 = budget_str
4288 .parse()
4289 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4290 match &mut loaded {
4291 model::LoadedModel::Gguf(g) => {
4292 tracing::info!(
4293 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4294 (CUDA and/or Metal matvecs when built with the matching feature)"
4295 );
4296 g.decoder.gpu_vram_budget_bytes = Some(budget);
4297 }
4298 model::LoadedModel::Kimi(_) => {
4299 tracing::warn!(
4300 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4301 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4302 );
4303 }
4304 model::LoadedModel::Mla(_) => {
4305 tracing::warn!(
4306 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4307 FFN path only today; ignoring expert VRAM budget"
4308 );
4309 }
4310 model::LoadedModel::Gemma4(_) => {
4311 tracing::warn!(
4312 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4313 ignoring expert VRAM budget"
4314 );
4315 }
4316 model::LoadedModel::Glm52(_) => {
4317 tracing::warn!(
4318 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4319 GPU expert placement not wired yet; ignoring"
4320 );
4321 }
4322 model::LoadedModel::Encoder(_) => {
4323 tracing::warn!(
4324 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4325 it has no routed experts to place; ignoring"
4326 );
4327 }
4328 }
4329 } else if metal_default_moe_budget {
4330 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4334 if let model::LoadedModel::Gguf(g) = &mut loaded {
4335 tracing::info!(
4336 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4337 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4338 );
4339 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4340 }
4341 }
4342 #[cfg(feature = "cuda")]
4343 {
4344 if frink_core::cuda_dense_enabled() {
4345 tracing::info!(
4346 "CUDA dense matvec enabled for WeightMatrix::apply \
4347 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4348 );
4349 } else {
4350 tracing::info!(
4351 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4352 );
4353 }
4354 }
4355 #[cfg(feature = "metal")]
4356 {
4357 if frink_core::metal_dense_enabled() {
4358 tracing::info!(
4359 "Metal dense matvec enabled for WeightMatrix::apply \
4360 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4361 );
4362 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4363 Some("1") | Some("true") | Some("on") | Some("attn") => {
4364 tracing::info!(
4365 "Metal fused attention requested (FRINK_METAL_ATTN): \
4366 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4367 );
4368 }
4369 _ => {}
4370 }
4371 tracing::info!(
4372 "Metal greedy GPU argmax: temperature<=0 folds \
4373 final_norm+lm_head+argmax into the dense stack"
4374 );
4375 } else {
4376 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4377 }
4378 }
4379 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4390 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4391 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4392 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4393 panic!(
4394 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4395 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4396 );
4397 }
4398 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4399 (Ok(blocks), Ok(block_size), Err(_)) => {
4400 let total_blocks: usize = blocks
4401 .parse()
4402 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4403 let block_size: usize = block_size
4404 .parse()
4405 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4406 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4412 .ok()
4413 .map(|v| {
4414 v.parse()
4415 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4416 })
4417 .unwrap_or(0);
4418 tracing::info!(
4419 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4420 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4421 );
4422 Some(generate::KvPoolConfig {
4423 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4424 queue_wait: Duration::from_millis(queue_wait_ms),
4425 })
4426 }
4427 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4428 let block_size: usize = block_size
4429 .parse()
4430 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4431 let budget: u64 = byte_budget
4432 .parse()
4433 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4434 let cfg = match &loaded {
4435 model::LoadedModel::Gguf(g) => &g.decoder.config,
4436 model::LoadedModel::Kimi(_)
4437 | model::LoadedModel::Mla(_)
4438 | model::LoadedModel::Gemma4(_)
4439 | model::LoadedModel::Glm52(_)
4440 | model::LoadedModel::Encoder(_) => {
4441 panic!(
4442 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4443 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4444 );
4445 }
4446 };
4447 let bytes_per_block = block_size
4448 * cfg.kv_heads_all_layers()
4449 * (cfg.head_dim + cfg.v_head_dim())
4450 * std::mem::size_of::<f32>();
4451 assert!(
4452 bytes_per_block > 0,
4453 "derived KV block byte size must be positive (check model config and block size)"
4454 );
4455 let total_blocks = (budget as usize / bytes_per_block).max(1);
4456 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4457 .ok()
4458 .map(|v| {
4459 v.parse()
4460 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4461 })
4462 .unwrap_or(0);
4463 tracing::info!(
4464 "KV cache block pool enabled from byte budget: {budget} bytes / \
4465 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4466 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4467 cfg.n_layers
4468 );
4469 Some(generate::KvPoolConfig {
4470 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4471 queue_wait: Duration::from_millis(queue_wait_ms),
4472 })
4473 }
4474 (Err(_), Err(_), Err(_)) => None,
4475 (Err(_), Ok(_), Err(_)) => panic!(
4476 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4477 (or unset all three to disable KV cache pooling)"
4478 ),
4479 (Ok(_), Ok(_), Ok(_)) => {
4480 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4481 }
4482 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4483 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4484 set together (or neither, to disable KV cache pooling)"
4485 ),
4486 };
4487 let paged_kv = match (
4493 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4494 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4495 ) {
4496 (Ok(blocks), Ok(block_size)) => {
4497 assert!(
4498 kv_pool.is_none(),
4499 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4500 mutually exclusive: both bound the same KV memory, by different means. \
4501 Set one."
4502 );
4503 let blocks_per_layer: usize = blocks
4517 .parse()
4518 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4519 let block_size: usize = block_size
4520 .parse()
4521 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4522 let gguf = match &loaded {
4523 model::LoadedModel::Gguf(g) => g,
4524 _ => panic!(
4525 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4526 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4527 ),
4528 };
4529 let cfg = &gguf.decoder.config;
4530 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4531 .ok()
4532 .map(|v| {
4533 v.parse()
4534 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4535 })
4536 .unwrap_or(0);
4537 tracing::info!(
4538 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4539 layer across {} layers, shared by all concurrent requests, \
4540 {queue_wait_ms}ms admission queue wait",
4541 cfg.n_layers
4542 );
4543 let radix = Some(Arc::new(Mutex::new(
4548 crate::policy::radix::SaltedRadix::new(block_size),
4549 )));
4550 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4559 crate::policy::parser::ToolCallFormat::infer(
4560 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4561 )
4562 .opener(),
4563 |text| {
4564 gguf.tokenizer
4565 .encode(text, SpecialTokens::Parse)
4566 .into_iter()
4567 .map(|t| t as u32)
4568 .collect()
4569 },
4570 );
4571 if let Some(id) = anchor_token {
4572 tracing::info!(
4573 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4574 window stops short of where its next turn rejoins"
4575 );
4576 }
4577 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4578 .ok()
4579 .map(|v| {
4580 v.parse()
4581 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4582 })
4583 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4584 if let Some(window) = cfg.uniform_sliding_window() {
4585 tracing::info!(
4586 "Paged KV window slide enabled: every layer slides by {window} every \
4587 {slide_interval} decode steps, so a request holds its prompt and a \
4588 window rather than its whole context"
4589 );
4590 } else if cfg.kv_block_window().is_some() {
4591 tracing::info!(
4592 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4593 and a page group holds one block in every layer"
4594 );
4595 }
4596 Some(generate::PagedKvConfig {
4597 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4600 queue_wait: Duration::from_millis(queue_wait_ms),
4601 radix,
4602 anchor_token,
4603 slide_interval,
4604 })
4605 }
4606 (Err(_), Err(_)) => None,
4607 _ => panic!(
4608 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4609 (or neither, to disable paged KV)"
4610 ),
4611 };
4612 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4619 let max_entries: usize = v
4620 .parse()
4621 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4622 if kv_pool.is_some() {
4623 tracing::warn!(
4624 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4625 caching will never be consulted while a KV pool is configured"
4626 );
4627 }
4628 assert!(
4635 paged_kv.is_none(),
4636 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4637 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4638 produce, so the cache could never hit. Set one."
4639 );
4640 tracing::info!(
4641 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4642 all requests"
4643 );
4644 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4645 });
4646 if matches!(
4647 loaded,
4648 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4649 ) && (kv_pool.is_some() || prefix_cache.is_some())
4650 {
4651 tracing::warn!(
4652 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4653 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4654 frink_models::engine's module docs"
4655 );
4656 }
4657 let enable_cb =
4658 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4659 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4660 tracing::info!(
4661 "continuous batching enabled by default on Metal for safe parallel serving \
4662 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4663 );
4664 }
4665 if continuous_batching_env() == Some(true)
4666 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4667 && (kv_pool.is_some() || prefix_cache.is_some())
4668 {
4669 tracing::warn!(
4670 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4671 (those modes keep the private generate path)"
4672 );
4673 }
4674 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4675 if let Ok(chunk) = n.parse::<usize>() {
4676 if chunk > 0 {
4677 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4678 }
4679 }
4680 }
4681 if matches!(
4682 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4683 Some("1")
4684 ) {
4685 tracing::warn!(
4686 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4687 (minimal spill; full layer offload still planned)"
4688 );
4689 }
4690
4691 let mcp = match mcp_config_path {
4692 Some(path) => {
4693 let loaded = mcp::load_mcp_config(&path)?;
4694 tracing::info!(
4695 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4696 loaded.path,
4697 loaded.servers.len()
4698 );
4699 Some(loaded)
4700 }
4701 None => None,
4702 };
4703
4704 let detection = health::Detection::spawn();
4708
4709 let state = Arc::new(build_app_state(
4710 StartupModels {
4711 loaded,
4712 embedding: embedding_model,
4713 },
4714 kv_pool,
4715 paged_kv,
4716 prefix_cache,
4717 enable_cb,
4718 mcp,
4719 detection,
4720 ));
4721
4722 use frink_api::routes;
4726
4727 let public = Router::new().route(routes::HEALTH, get(health));
4732
4733 let mut protected = protected_routes();
4734
4735 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4740 tracing::info!("API key auth enabled");
4741 let auth = limits::AuthConfig {
4742 api_key: Arc::new(key),
4743 };
4744 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4745 auth,
4746 limits::require_api_key,
4747 ));
4748 }
4749 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4750 let rpm: u32 = rpm
4751 .parse()
4752 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4753 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4754 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4755 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4756 limiter,
4757 limits::rate_limit,
4758 ));
4759 }
4760 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4768 let origins = security::parse_cors_origins(&spec)
4769 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4770 tracing::info!(
4771 "CORS enabled: {} allow-listed origin(s) ({})",
4772 origins.len(),
4773 spec
4774 );
4775 let cors = tower_http::cors::CorsLayer::new()
4776 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4777 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4778 .allow_headers([
4779 axum::http::header::CONTENT_TYPE,
4780 axum::http::header::AUTHORIZATION,
4781 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4787 axum::http::HeaderName::from_static("last-event-id"),
4791 ]);
4792 protected = protected.route_layer(cors);
4793 }
4794
4795 let app = public
4799 .merge(protected)
4800 .layer(axum::middleware::from_fn(limits::retry_after))
4801 .with_state(state);
4802
4803 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4809 install_ring_crypto_provider();
4810 match tls_paths {
4815 Some(paths) => {
4816 let config =
4817 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4818 .await
4819 .map_err(|e| {
4820 anyhow::anyhow!(
4821 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4822 paths.cert,
4823 paths.key
4824 )
4825 })?;
4826 let socket_addr: std::net::SocketAddr = addr
4827 .parse()
4828 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4829 let listener = std::net::TcpListener::bind(socket_addr)?;
4830 listener.set_nonblocking(true)?;
4838 let bound = listener.local_addr()?;
4839 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4840 announce_ready(bound, "https");
4841
4842 let handle = axum_server::Handle::new();
4843 let shutdown_handle = handle.clone();
4844 tokio::spawn(async move {
4845 shutdown_signal(exit_on_stdin_close).await;
4846 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4847 });
4848 axum_server::from_tcp_rustls(listener, config)?
4849 .handle(handle)
4850 .serve(app.into_make_service())
4851 .await?;
4852 }
4853 None => {
4854 let listener = tokio::net::TcpListener::bind(&addr).await?;
4855 let bound = listener.local_addr()?;
4856 tracing::info!("frink-server listening on {bound}");
4857 announce_ready(bound, "http");
4858 axum::serve(listener, app)
4859 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4860 .await?;
4861 }
4862 }
4863 Ok(())
4864}
4865
4866#[cfg(test)]
4867pub(crate) mod tests {
4868 use super::*;
4869 use frink_models::config::test_dense_fixture;
4870
4871 #[test]
4872 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4873 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4874 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4875 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4876 assert_eq!(parsed.port, 51999);
4877 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4878 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4880 }
4881
4882 fn test_model() -> Model {
4883 let cfg = test_dense_fixture();
4887 Model::Gguf(GgufModel {
4888 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4889 tokenizer: Arc::new(ServerTokenizer::Byte),
4890 stop_tokens: StopTokens::default(),
4891 bos_id: None,
4892 is_synthetic: true,
4893 chat_template: chat_template::PromptTemplate::plain(),
4894 })
4895 }
4896
4897 fn greedy_params(max_tokens: usize) -> GenerationParams {
4898 GenerationParams {
4899 cache_salt: None,
4900 prompt_logprobs: None,
4901 wants_logprobs: false,
4902 n: 1,
4903 interleave_choices: false,
4904 token_mask: crate::token_mask::TokenMask::default(),
4905 reasoning: None,
4906 max_tokens,
4907 sampling: SamplingParams::default(),
4908 seed: 1,
4909 stop: Vec::new(),
4910 stop_token_ids: Vec::new(),
4911 json_object: false,
4912 grammar: None,
4913 cancel: None,
4914 ignore_eos: false,
4915 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4916 lora: None,
4917 }
4918 }
4919
4920 fn test_model_full_byte_vocab() -> Model {
4924 test_model_full_byte_vocab_with_eos(None)
4925 }
4926
4927 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4934 let mut cfg = test_dense_fixture();
4935 cfg.vocab_size = 256;
4936 Model::Gguf(GgufModel {
4937 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4938 tokenizer: Arc::new(ServerTokenizer::Byte),
4939 stop_tokens: StopTokens::from_eos(eos),
4940 bos_id: None,
4941 is_synthetic: true,
4942 chat_template: chat_template::PromptTemplate::plain(),
4943 })
4944 }
4945
4946 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4950 test_state_at(model, response_cache, None)
4951 }
4952
4953 pub(crate) fn test_state_at(
4957 model: Model,
4958 response_cache: ResponseCache,
4959 checkpoint_path: Option<std::path::PathBuf>,
4960 ) -> AppState {
4961 AppState {
4962 slept: Mutex::new(None),
4963 embedding: None,
4964 paged_kv: None,
4965 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4966 id: None,
4967 loaded: Loaded::Generative(Arc::new(model)),
4968 batcher: None,
4969 ceiling: None,
4970 checkpoint_path,
4971 }))),
4972 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4973 tasks: Arc::new(tasks::TaskRegistry::new()),
4974 cancels: Arc::new(cancel::CancelRegistry::new()),
4975 stats: stats::Stats::new(),
4976 streams: resume::StreamRegistry::new(),
4977 model_dir: None,
4978 response_cache: Mutex::new(response_cache),
4979 kv_pool: None,
4980 prefix_cache: None,
4981 sessions: session::SessionStore::new(),
4982 requests_total: std::sync::atomic::AtomicU64::new(0),
4983 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4984 started_at: std::time::Instant::now(),
4985 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4986 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4987 mcp: None,
4988 continuous_batching_enabled: false,
4989 metal_private_decode_gate: None,
4990 loading_model: Mutex::new(None),
4991 last_load_error: Mutex::new(None),
4992 serving: Mutex::new(crate::stats::ServingStats::default()),
4993 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4994 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4995 started_unix: unix_now(),
4996 }
4997 }
4998
4999 pub(crate) fn test_app() -> Router {
5008 test_app_with_state(Arc::new(test_state(
5009 test_model_full_byte_vocab(),
5010 ResponseCache::new(1000, Duration::from_secs(3600)),
5011 )))
5012 }
5013
5014 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
5017 Router::new()
5027 .route(frink_api::routes::HEALTH, get(health))
5028 .merge(protected_routes())
5029 .with_state(state)
5030 }
5031
5032 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
5033 let mut cfg = test_dense_fixture();
5034 cfg.name = name;
5035 cfg.vocab_size = vocab_size;
5036 Model::Gguf(GgufModel {
5037 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5038 tokenizer: Arc::new(ServerTokenizer::Byte),
5039 stop_tokens: StopTokens::default(),
5040 bos_id: None,
5041 is_synthetic: true,
5042 chat_template: chat_template::PromptTemplate::plain(),
5043 })
5044 }
5045
5046 fn model_with_template(name: &'static str, source: &str) -> Model {
5050 let mut cfg = test_dense_fixture();
5051 cfg.name = name;
5052 cfg.vocab_size = 256;
5053 Model::Gguf(GgufModel {
5054 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5055 tokenizer: Arc::new(ServerTokenizer::Byte),
5056 stop_tokens: StopTokens::default(),
5057 bos_id: None,
5058 is_synthetic: true,
5059 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5060 Some(source),
5061 Some("qwen3"),
5062 false,
5063 true,
5064 None,
5065 None,
5066 ),
5067 })
5068 }
5069
5070 #[tokio::test]
5079 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5080 let strict = "{% if messages | length > 1 %}\
5083 {{ raise_exception('this template takes one turn') }}\
5084 {% endif %}{{ messages[0].content }}";
5085 let state = Arc::new(test_state(
5086 model_with_template("strict", strict),
5087 ResponseCache::new(4, Duration::from_secs(60)),
5088 ));
5089 let app = test_app_with_state(state);
5090
5091 let (status, body) = post_json_uri(
5092 &app,
5093 "/v1/chat/completions",
5094 serde_json::json!({
5095 "model": "strict",
5096 "stream": true,
5097 "messages": [
5098 {"role": "user", "content": "one"},
5099 {"role": "user", "content": "two"},
5100 ],
5101 }),
5102 )
5103 .await;
5104 assert_eq!(status, StatusCode::BAD_REQUEST);
5105 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5106 assert!(
5107 body["error"]["message"]
5108 .as_str()
5109 .unwrap()
5110 .contains("one turn"),
5111 "the template's own message must reach the caller: {body}"
5112 );
5113
5114 let (status, _) = post_json_uri(
5116 &app,
5117 "/v1/chat/completions",
5118 serde_json::json!({
5119 "model": "strict",
5120 "stream": true,
5121 "max_tokens": 1,
5122 "messages": [{"role": "user", "content": "one"}],
5123 }),
5124 )
5125 .await;
5126 assert_eq!(status, StatusCode::OK);
5127 }
5128
5129 #[tokio::test]
5131 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5132 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5133 {% if reasoning_effort %}\
5134 {% if reasoning_effort not in ['low','medium','high'] %}\
5135 {{ raise_exception('bad effort') }}\
5136 {% endif %}[{{ reasoning_effort }}]\
5137 {% endif %}{{ messages[0].content }}";
5138 let state = Arc::new(test_state(
5139 model_with_template("thinker", reasoning),
5140 ResponseCache::new(4, Duration::from_secs(60)),
5141 ));
5142 let app = test_app_with_state(state);
5143 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5144 assert_eq!(status, StatusCode::OK);
5145 let entry = &models["data"][0];
5146 assert_eq!(
5147 entry["supported_reasoning_efforts"],
5148 serde_json::json!(["off", "low", "medium", "high"])
5149 );
5150 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5151 }
5152
5153 #[tokio::test]
5158 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5159 let app = test_app();
5160 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5161 let entry = &models["data"][0];
5162 assert!(entry.get("supported_reasoning_efforts").is_none());
5163 assert!(entry.get("default_reasoning_effort").is_none());
5164 }
5165
5166 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5167 Arc::new(ActiveModel {
5168 id: Some(name.to_string()),
5169 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5170 batcher: None,
5171 ceiling: None,
5172 checkpoint_path: None,
5173 })
5174 .tap_into(state)
5175 }
5176
5177 trait TapInto {
5179 fn tap_into(self, state: &AppState) -> Self;
5180 }
5181 impl TapInto for Arc<ActiveModel> {
5182 fn tap_into(self, state: &AppState) -> Self {
5183 state.swap_active(Some(Arc::clone(&self)));
5184 self
5185 }
5186 }
5187
5188 #[test]
5194 fn an_in_flight_request_keeps_the_model_it_started_on() {
5195 let state = test_state(
5196 named_test_model("model-a", 256),
5197 ResponseCache::new(4, Duration::from_secs(60)),
5198 );
5199
5200 let in_flight = state.active().expect("a model is loaded");
5203 assert_eq!(in_flight.name(), "model-a");
5204
5205 active_model(&state, "model-b");
5206
5207 assert_eq!(state.active().unwrap().name(), "model-b");
5209 assert_eq!(in_flight.name(), "model-a");
5211 let produced = run_generation(
5212 in_flight.generative().unwrap(),
5213 "hi",
5214 &greedy_params(3),
5215 None,
5216 None,
5217 None,
5218 None,
5219 None,
5220 None,
5221 )
5222 .expect("the old model must still decode after being swapped out");
5223 assert!(matches!(
5224 produced.choices[0].finish,
5225 FinishReason::Length | FinishReason::Stop
5226 ));
5227 }
5228
5229 #[test]
5234 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5235 let state = test_state(
5236 named_test_model("model-a", 256),
5237 ResponseCache::new(4, Duration::from_secs(60)),
5238 );
5239 let in_flight = state.active().expect("a model is loaded");
5240 let weights = Arc::clone(in_flight.generative().unwrap());
5241 assert!(Arc::strong_count(&weights) >= 2);
5242
5243 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5244 id: Some("model-b".to_string()),
5245 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5246 batcher: None,
5247 ceiling: None,
5248 checkpoint_path: None,
5249 })));
5250 drop(previous);
5251 assert!(Arc::strong_count(&weights) >= 2);
5253 drop(in_flight);
5254 assert_eq!(Arc::strong_count(&weights), 1);
5255 }
5256
5257 #[tokio::test]
5261 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5262 let state = Arc::new(test_state(
5263 named_test_model("model-a", 256),
5264 ResponseCache::new(4, Duration::from_secs(60)),
5265 ));
5266 let app = test_app_with_state(Arc::clone(&state));
5267
5268 let (status, body) = post_json_uri(
5269 &app,
5270 frink_api::routes::ADMIN_MODELS_UNLOAD,
5271 serde_json::json!({}),
5272 )
5273 .await;
5274 assert_eq!(status, StatusCode::OK);
5275 assert_eq!(body["ok"], true);
5276 assert!(body["active"].is_null());
5277 assert!(state.active().is_none());
5278
5279 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5280 assert_eq!(status, StatusCode::OK);
5281 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5282 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5283
5284 let (status, body) = post_json_uri(
5285 &app,
5286 "/v1/chat/completions",
5287 serde_json::json!({
5288 "model": "x",
5289 "messages": [{"role": "user", "content": "hi"}]
5290 }),
5291 )
5292 .await;
5293 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5294 assert_eq!(body["error"]["type"], "model_not_loaded");
5295 }
5296
5297 #[tokio::test]
5301 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5302 let state = Arc::new(test_state(
5303 named_test_model("model-a", 256),
5304 ResponseCache::new(4, Duration::from_secs(60)),
5305 ));
5306 let app = test_app_with_state(Arc::clone(&state));
5307 state.swap_active(None);
5308
5309 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5310 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5313 assert_eq!(body["state"], "unavailable");
5314 assert_eq!(body["reason"], "model_not_loaded");
5315 assert!(body["model"].is_null());
5316 let real_weights = body["capabilities"]
5317 .as_array()
5318 .unwrap()
5319 .iter()
5320 .find(|c| c["id"] == "real_weights")
5321 .cloned()
5322 .expect("real_weights is always reported");
5323 assert_eq!(real_weights["available"], false);
5324 assert_eq!(real_weights["reason"], "model_not_loaded");
5325 }
5326
5327 #[tokio::test]
5331 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5332 let app = test_app();
5333
5334 let (status, completion) = post_json_uri(
5335 &app,
5336 "/v1/chat/completions",
5337 serde_json::json!({
5338 "model": "x",
5339 "messages": [{"role": "user", "content": "hi"}],
5340 "max_tokens": 4
5341 }),
5342 )
5343 .await;
5344 assert_eq!(status, StatusCode::OK);
5345 let request_id = completion["request_id"].as_str().unwrap().to_string();
5346
5347 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5348 assert_eq!(status, StatusCode::OK);
5349 let recent = stats["recent"].as_array().unwrap();
5350 assert_eq!(recent.len(), 1);
5351 let row = &recent[0];
5352 assert_eq!(row["request_id"], request_id);
5353 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5354 assert_eq!(row["status"], 200);
5355 assert_eq!(row["stream"], false);
5356 assert!(row["duration_ms"].is_number());
5359 assert!(row["decode_ms"].is_number());
5360 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5361 assert_eq!(
5362 stats["tokens_prompt_total"].as_u64().unwrap(),
5363 row["prompt_tokens"].as_u64().unwrap()
5364 );
5365 }
5366
5367 #[tokio::test]
5371 async fn a_rejected_request_is_recorded_too() {
5372 let state = Arc::new(test_state(
5373 named_test_model("model-a", 256),
5374 ResponseCache::new(4, Duration::from_secs(60)),
5375 ));
5376 let app = test_app_with_state(Arc::clone(&state));
5377 state.swap_active(None);
5378
5379 let (status, _) = post_json_uri(
5380 &app,
5381 "/v1/chat/completions",
5382 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5383 )
5384 .await;
5385 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5386
5387 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5388 let recent = stats["recent"].as_array().unwrap();
5389 assert_eq!(recent.len(), 1);
5390 assert_eq!(recent[0]["status"], 503);
5391 assert_eq!(recent[0]["completion_tokens"], 0);
5392 assert!(recent[0]["decode_ms"].is_null());
5393 assert_eq!(stats["errors_total"], 1);
5394 }
5395
5396 async fn post_json_with_headers(
5400 app: &Router,
5401 uri: &str,
5402 body: serde_json::Value,
5403 headers: &[(&str, &str)],
5404 ) -> (StatusCode, serde_json::Value) {
5405 use http_body_util::BodyExt;
5406 use tower::ServiceExt;
5407
5408 let mut builder = axum::http::Request::builder()
5409 .method("POST")
5410 .uri(uri)
5411 .header("content-type", "application/json");
5412 for (name, value) in headers {
5413 builder = builder.header(*name, *value);
5414 }
5415 let response = app
5416 .clone()
5417 .oneshot(
5418 builder
5419 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5420 .unwrap(),
5421 )
5422 .await
5423 .unwrap();
5424 let status = response.status();
5425 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5426 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5427 (status, json)
5428 }
5429
5430 #[tokio::test]
5434 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5435 let app = test_app();
5436
5437 let (status, _) = post_json_uri(
5438 &app,
5439 frink_api::routes::V1_TOKENIZE,
5440 serde_json::json!({"prompt": "hello"}),
5441 )
5442 .await;
5443 assert_eq!(status, StatusCode::OK);
5444 let (status, _) = post_json_uri(
5445 &app,
5446 frink_api::routes::V1_DETOKENIZE,
5447 serde_json::json!({"tokens": [104, 105]}),
5448 )
5449 .await;
5450 assert_eq!(status, StatusCode::OK);
5451 let (status, _) = post_json_uri(
5452 &app,
5453 frink_api::routes::V1_EMBEDDINGS,
5454 serde_json::json!({"input": "hello"}),
5455 )
5456 .await;
5457 assert_eq!(status, StatusCode::OK);
5458
5459 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5460 let routes: Vec<&str> = stats["recent"]
5461 .as_array()
5462 .unwrap()
5463 .iter()
5464 .map(|row| row["route"].as_str().unwrap())
5465 .collect();
5466 for expected in [
5467 frink_api::routes::V1_TOKENIZE,
5468 frink_api::routes::V1_DETOKENIZE,
5469 frink_api::routes::V1_EMBEDDINGS,
5470 ] {
5471 assert!(
5472 routes.contains(&expected),
5473 "{expected} is missing: {routes:?}"
5474 );
5475 }
5476
5477 let row = |route: &str| {
5478 stats["recent"]
5479 .as_array()
5480 .unwrap()
5481 .iter()
5482 .find(|r| r["route"] == route)
5483 .cloned()
5484 .unwrap()
5485 };
5486 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5490 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5491 assert!(embed["decode_ms"].is_null());
5492 assert_eq!(embed["completion_tokens"], 0);
5493 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5497 assert_eq!(
5498 stats["tokens_prompt_total"].as_u64().unwrap(),
5499 embed["prompt_tokens"].as_u64().unwrap(),
5500 "only the forward pass counted"
5501 );
5502 }
5503
5504 fn streaming_test_app() -> Router {
5509 let mut cfg = test_dense_fixture();
5510 cfg.vocab_size = 256;
5511 let model = Model::Gguf(GgufModel {
5512 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5513 tokenizer: Arc::new(ServerTokenizer::Byte),
5514 stop_tokens: StopTokens::default(),
5515 bos_id: None,
5516 is_synthetic: false,
5517 chat_template: chat_template::PromptTemplate::plain(),
5518 });
5519 test_app_with_state(Arc::new(test_state(
5520 model,
5521 ResponseCache::new(1000, Duration::from_secs(3600)),
5522 )))
5523 }
5524
5525 #[tokio::test]
5534 async fn chat_logprobs_are_rendered_and_are_never_served_from_cache() {
5535 let app = test_app();
5536 let body = |logprobs: Option<(bool, Option<u32>)>| {
5537 let mut b = serde_json::json!({
5538 "model": "x",
5539 "messages": [{"role": "user", "content": "hi"}],
5540 "max_tokens": 4
5541 });
5542 if let Some((on, top)) = logprobs {
5543 b["logprobs"] = serde_json::json!(on);
5544 if let Some(n) = top {
5545 b["top_logprobs"] = serde_json::json!(n);
5546 }
5547 }
5548 b
5549 };
5550
5551 let (status, plain) =
5553 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(None)).await;
5554 assert_eq!(status, StatusCode::OK, "{plain}");
5555 assert!(plain["choices"][0]["logprobs"].is_null(), "{plain}");
5556
5557 for attempt in 0..2 {
5561 let (status, with) = post_json_uri(
5562 &app,
5563 frink_api::routes::V1_CHAT_COMPLETIONS,
5564 body(Some((true, Some(2)))),
5565 )
5566 .await;
5567 assert_eq!(status, StatusCode::OK, "{with}");
5568 assert_ne!(
5569 with["frink_cache"], "hit",
5570 "attempt {attempt} replayed a cached answer for a logprobs request: {with}"
5571 );
5572 let lp = &with["choices"][0]["logprobs"];
5573 assert!(lp.is_object(), "attempt {attempt}: {with}");
5574 let content = lp["content"].as_array().expect("content");
5575 assert!(lp["tokens"].is_null(), "completions shape leaked: {lp}");
5577 for entry in content {
5578 assert!(entry["token"].is_string(), "{entry}");
5579 assert!(entry["bytes"].is_array(), "{entry}");
5580 let v = entry["logprob"].as_f64().expect("a real number");
5581 assert!(v <= 0.0 && v.is_finite(), "{entry}");
5582 let top = entry["top_logprobs"].as_array().expect("top_logprobs");
5583 assert!(top.len() <= 2, "asked for 2, got {}", top.len());
5584 }
5585 }
5586 }
5587
5588 #[tokio::test]
5594 async fn the_chat_logprobs_pair_is_validated() {
5595 let app = test_app();
5596 for (extra, why) in [
5597 (serde_json::json!({"top_logprobs": 3}), "without logprobs"),
5598 (
5599 serde_json::json!({"logprobs": true, "top_logprobs": 21}),
5600 "above the cap",
5601 ),
5602 ] {
5603 let mut body = serde_json::json!({
5604 "model": "x",
5605 "messages": [{"role": "user", "content": "hi"}],
5606 "max_tokens": 2
5607 });
5608 for (k, v) in extra.as_object().unwrap() {
5609 body[k] = v.clone();
5610 }
5611 let (status, answer) =
5612 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await;
5613 assert_eq!(status, StatusCode::BAD_REQUEST, "{why}: {answer}");
5614 assert!(
5615 answer["error"]["message"]
5616 .as_str()
5617 .is_some_and(|m| m.contains("top_logprobs")),
5618 "{why}: {answer}"
5619 );
5620 }
5621 }
5622
5623 #[tokio::test]
5632 async fn sleep_refuses_a_model_it_could_not_bring_back() {
5633 let app = test_app();
5634 let (status, answer) =
5635 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5636 assert_eq!(status, StatusCode::CONFLICT, "{answer}");
5637 assert_eq!(answer["error"]["type"], "not_reloadable", "{answer}");
5638 let (_, still) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5641 assert_eq!(still["is_sleeping"], false, "{still}");
5642 let (status, _) = post_json_uri(
5643 &app,
5644 frink_api::routes::V1_CHAT_COMPLETIONS,
5645 serde_json::json!({
5646 "model": "x",
5647 "messages": [{"role": "user", "content": "hi"}],
5648 "max_tokens": 2
5649 }),
5650 )
5651 .await;
5652 assert_eq!(status, StatusCode::OK, "a refused sleep unloaded the model");
5653 }
5654
5655 #[tokio::test]
5664 async fn sleep_remembers_what_unload_forgets() {
5665 let state = Arc::new(test_state_at(
5668 test_model_full_byte_vocab(),
5669 ResponseCache::new(1000, Duration::from_secs(3600)),
5670 Some(std::path::PathBuf::from("/nonexistent/fixture.gguf")),
5671 ));
5672 let app = test_app_with_state(Arc::clone(&state));
5673 let ask = || {
5674 let app = app.clone();
5675 async move {
5676 post_json_uri(
5677 &app,
5678 frink_api::routes::V1_CHAT_COMPLETIONS,
5679 serde_json::json!({
5680 "model": "x",
5681 "messages": [{"role": "user", "content": "hi"}],
5682 "max_tokens": 2
5683 }),
5684 )
5685 .await
5686 }
5687 };
5688
5689 let (status, _) = ask().await;
5690 assert_eq!(status, StatusCode::OK, "the fixture server serves");
5691 let (_, awake) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5692 assert_eq!(awake["is_sleeping"], false, "{awake}");
5693
5694 let (status, slept) =
5695 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5696 assert_eq!(status, StatusCode::OK, "{slept}");
5697 assert_eq!(slept["is_sleeping"], true, "{slept}");
5698 let (_, now) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5699 assert_eq!(now["is_sleeping"], true, "{now}");
5700
5701 let (status, refused) = ask().await;
5704 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE, "{refused}");
5705 assert_eq!(
5706 refused["error"]["type"], "server_sleeping",
5707 "an asleep server reported itself as empty: {refused}"
5708 );
5709
5710 let (status, again) =
5712 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5713 assert_eq!(status, StatusCode::OK, "{again}");
5714 assert_eq!(again["is_sleeping"], true, "{again}");
5715 }
5716
5717 #[tokio::test]
5721 async fn waking_a_server_that_is_awake_is_refused() {
5722 let app = test_app();
5723 let (status, answer) =
5724 post_json_uri(&app, frink_api::routes::WAKE_UP, serde_json::json!({})).await;
5725 assert_eq!(status, StatusCode::CONFLICT, "{answer}");
5726 assert_eq!(answer["error"]["type"], "not_sleeping", "{answer}");
5727 }
5728
5729 #[tokio::test]
5736 async fn a_salt_keeps_one_callers_cached_answer_from_another() {
5737 let app = test_app();
5738 let body = |salt: Option<&str>| {
5739 let mut b = serde_json::json!({
5740 "model": "x",
5741 "messages": [{"role": "user", "content": "the same prompt"}],
5742 "max_tokens": 4,
5743 "seed": 1
5744 });
5745 if let Some(s) = salt {
5746 b["cache_salt"] = serde_json::json!(s);
5747 }
5748 b
5749 };
5750 let post = |b: serde_json::Value| {
5751 let app = app.clone();
5752 async move { post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, b).await }
5753 };
5754
5755 let (status, _) = post(body(Some("tenant-a"))).await;
5757 assert_eq!(status, StatusCode::OK);
5758 let (_, again) = post(body(Some("tenant-a"))).await;
5759 assert_eq!(
5760 again["frink_cache"], "hit",
5761 "the owner did not get its own entry back: {again}"
5762 );
5763
5764 let (_, other) = post(body(Some("tenant-b"))).await;
5766 assert_ne!(
5767 other["frink_cache"], "hit",
5768 "a different caller was served tenant-a's answer: {other}"
5769 );
5770
5771 let (_, shared) = post(body(None)).await;
5773 assert_ne!(
5774 shared["frink_cache"], "hit",
5775 "an unsalted request was served a salted answer: {shared}"
5776 );
5777 }
5778
5779 #[tokio::test]
5782 async fn chat_serves_several_choices_from_one_prefill() {
5783 let app = test_app();
5784 let body = |n: u32, stream: bool| {
5785 serde_json::json!({
5786 "model": "x",
5787 "messages": [{"role": "user", "content": "hi"}],
5788 "max_tokens": 4,
5789 "temperature": 1.0,
5790 "n": n,
5791 "stream": stream
5792 })
5793 };
5794
5795 let (status, one) =
5796 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5797 assert_eq!(status, StatusCode::OK, "{one}");
5798
5799 let (status, three) =
5800 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5801 assert_eq!(status, StatusCode::OK, "{three}");
5802 let choices = three["choices"].as_array().expect("an array");
5803 assert_eq!(choices.len(), 3, "{three}");
5804 for (i, c) in choices.iter().enumerate() {
5805 assert_eq!(c["index"], i);
5806 assert!(c["message"]["role"].is_string(), "{c}");
5807 assert!(c["finish_reason"].is_string(), "{c}");
5808 }
5809 assert_eq!(
5811 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5812 "n = 3 billed the prompt more than once"
5813 );
5814 }
5815
5816 #[tokio::test]
5828 async fn a_streamed_n_interleaves_its_choices() {
5829 let app = streaming_test_app();
5830 let raw = post_sse_raw(
5831 &app,
5832 serde_json::json!({
5833 "model": "x",
5834 "messages": [{"role": "user", "content": "hi"}],
5835 "max_tokens": 6,
5836 "temperature": 1.0,
5837 "n": 3,
5838 "stream": true
5839 }),
5840 )
5841 .await;
5842
5843 let mut order: Vec<usize> = Vec::new();
5845 let mut finished: Vec<usize> = Vec::new();
5846 let mut usage_blocks = 0usize;
5847 for line in raw.lines() {
5848 let Some(rest) = line.strip_prefix("data: ") else {
5849 continue;
5850 };
5851 if rest.trim() == "[DONE]" {
5852 continue;
5853 }
5854 let v: serde_json::Value = serde_json::from_str(rest).expect(rest);
5855 if v.get("usage").is_some_and(|u| !u.is_null()) {
5856 usage_blocks += 1;
5857 }
5858 let Some(choice) = v["choices"].as_array().and_then(|c| c.first()) else {
5859 continue;
5860 };
5861 let index = choice["index"].as_u64().expect("an index") as usize;
5862 if choice["finish_reason"].is_string() {
5863 finished.push(index);
5864 continue;
5865 }
5866 order.push(index);
5867 }
5868
5869 assert_eq!(
5870 finished,
5871 vec![0, 1, 2],
5872 "one terminal chunk per choice, in index order: {raw}"
5873 );
5874 assert_eq!(usage_blocks, 1, "the usage block is the request's: {raw}");
5875 assert!(
5876 order.contains(&0) && order.contains(&2),
5877 "not every choice streamed: {order:?}"
5878 );
5879 let last_of_zero = order
5880 .iter()
5881 .rposition(|i| *i == 0)
5882 .expect("choice 0 streamed");
5883 let first_of_two = order
5884 .iter()
5885 .position(|i| *i == 2)
5886 .expect("choice 2 streamed");
5887 assert!(
5888 first_of_two < last_of_zero,
5889 "the choices arrived one after another rather than interleaved: {order:?}"
5890 );
5891 }
5892
5893 #[tokio::test]
5901 async fn allowed_token_ids_restricts_the_draw() {
5902 let app = streaming_test_app();
5903 let body = |allowed: Option<serde_json::Value>| {
5904 let mut b = serde_json::json!({
5905 "model": "x",
5906 "prompt": "hi",
5907 "max_tokens": 16,
5908 "temperature": 1.0,
5909 "seed": 3
5910 });
5911 if let Some(ids) = allowed {
5912 b["allowed_token_ids"] = ids;
5913 }
5914 b
5915 };
5916
5917 let (status, free) =
5920 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, body(None)).await;
5921 assert_eq!(status, StatusCode::OK, "{free}");
5922 let free_text = free["choices"][0]["text"].as_str().unwrap_or_default();
5923
5924 let (status, restricted) = post_json_uri(
5925 &app,
5926 frink_api::routes::V1_COMPLETIONS,
5927 body(Some(serde_json::json!([65, 66]))),
5929 )
5930 .await;
5931 assert_eq!(status, StatusCode::OK, "{restricted}");
5932 let text = restricted["choices"][0]["text"]
5933 .as_str()
5934 .unwrap_or_default();
5935 assert!(!text.is_empty(), "nothing was generated: {restricted}");
5936 assert!(
5937 text.chars().all(|c| c == 'A' || c == 'B'),
5938 "a token outside `allowed_token_ids` was drawn: {text:?}"
5939 );
5940 assert!(
5943 !free_text.chars().all(|c| c == 'A' || c == 'B'),
5944 "the unrestricted answer was already inside the allowed set: {free_text:?}"
5945 );
5946 }
5947
5948 #[tokio::test]
5955 async fn an_empty_allowed_token_ids_is_a_bad_request() {
5956 let app = streaming_test_app();
5957 let (status, body) = post_json_uri(
5958 &app,
5959 frink_api::routes::V1_COMPLETIONS,
5960 serde_json::json!({
5961 "model": "x",
5962 "prompt": "hi",
5963 "max_tokens": 4,
5964 "allowed_token_ids": []
5965 }),
5966 )
5967 .await;
5968 assert_eq!(status, StatusCode::BAD_REQUEST, "{body}");
5969 assert!(
5970 body["error"]["message"]
5971 .as_str()
5972 .unwrap_or_default()
5973 .contains("allowed_token_ids"),
5974 "{body}"
5975 );
5976 }
5977
5978 #[tokio::test]
5984 async fn bad_words_removes_a_token_without_ending_the_generation() {
5985 let app = streaming_test_app();
5986 let ask = |bad: Option<serde_json::Value>| {
5987 let mut b = serde_json::json!({
5988 "model": "x",
5989 "prompt": "hi",
5990 "max_tokens": 24,
5991 "temperature": 1.0,
5992 "seed": 11
5993 });
5994 if let Some(words) = bad {
5995 b["bad_words"] = words;
5996 }
5997 b
5998 };
5999
6000 let (status, free) =
6001 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(None)).await;
6002 assert_eq!(status, StatusCode::OK, "{free}");
6003 let free_text = free["choices"][0]["text"]
6004 .as_str()
6005 .unwrap_or_default()
6006 .to_string();
6007 let target = free_text
6010 .chars()
6011 .find(|c| c.is_ascii() && !c.is_control())
6012 .expect("the model produced some ascii");
6013
6014 let (status, steered) = post_json_uri(
6015 &app,
6016 frink_api::routes::V1_COMPLETIONS,
6017 ask(Some(serde_json::json!([target.to_string()]))),
6018 )
6019 .await;
6020 assert_eq!(status, StatusCode::OK, "{steered}");
6021 let text = steered["choices"][0]["text"].as_str().unwrap_or_default();
6022 assert!(
6023 !text.contains(target),
6024 "the forbidden {target:?} came back anyway: {text:?}"
6025 );
6026 assert_eq!(
6029 steered["usage"]["completion_tokens"], free["usage"]["completion_tokens"],
6030 "the generation ended early, so `bad_words` acted like `stop`: {steered}"
6031 );
6032 }
6033
6034 #[tokio::test]
6044 async fn every_route_refuses_the_same_unimplemented_fields() {
6045 let app = test_app();
6046 let fields = [
6047 ("n", serde_json::json!(3)),
6048 ("best_of", serde_json::json!(2)),
6049 ("prompt_logprobs", serde_json::json!(1)),
6050 ("echo", serde_json::json!(true)),
6051 ("use_beam_search", serde_json::json!(true)),
6052 ("truncate_prompt_tokens", serde_json::json!(8)),
6053 ("prompt_embeds", serde_json::json!("AA==")),
6054 ("skip_special_tokens", serde_json::json!(false)),
6055 ("return_tokens_as_token_ids", serde_json::json!(true)),
6056 ];
6057 for (field, value) in fields {
6058 for (uri, base) in [
6059 (
6060 frink_api::routes::V1_CHAT_COMPLETIONS,
6061 serde_json::json!({
6062 "model": "x",
6063 "messages": [{"role": "user", "content": "hi"}],
6064 "max_tokens": 2
6065 }),
6066 ),
6067 (
6068 frink_api::routes::V1_COMPLETIONS,
6069 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
6070 ),
6071 (
6072 frink_api::routes::COMPLETION,
6073 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6074 ),
6075 ] {
6076 let mut body = base;
6077 body[field] = value.clone();
6078 if field == "prompt_logprobs" && uri == frink_api::routes::V1_COMPLETIONS {
6085 let (status, answer) = post_json_uri(&app, uri, body).await;
6086 assert_eq!(status, StatusCode::OK, "{uri} refused it: {answer}");
6087 assert!(
6088 answer["prompt_logprobs"].is_array(),
6089 "served without the field: {answer}"
6090 );
6091 continue;
6092 }
6093 if (field == "n" || field == "best_of")
6094 && (uri == frink_api::routes::V1_COMPLETIONS
6095 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
6096 {
6097 let (status, answer) = post_json_uri(&app, uri, body).await;
6098 assert_eq!(
6099 status,
6100 StatusCode::OK,
6101 "{uri} refused a served `{field}`: {answer}"
6102 );
6103 let want = if field == "n" { 3 } else { 1 };
6107 assert_eq!(
6108 answer["choices"].as_array().map(Vec::len),
6109 Some(want),
6110 "{field}: {answer}"
6111 );
6112 continue;
6113 }
6114 let (status, answer) = post_json_uri(&app, uri, body).await;
6115 assert_eq!(
6116 status,
6117 StatusCode::NOT_IMPLEMENTED,
6118 "{uri} served `{field}` instead of refusing it: {answer}"
6119 );
6120 assert!(
6121 answer["error"]["message"]
6122 .as_str()
6123 .is_some_and(|m| m.contains(field)),
6124 "{uri} refused `{field}` without naming it: {answer}"
6125 );
6126 }
6127 }
6128 }
6129
6130 #[tokio::test]
6131 async fn the_native_completion_wire_is_not_the_openai_one() {
6132 let app = test_app();
6133
6134 let (status, native) = post_json_uri(
6135 &app,
6136 frink_api::routes::COMPLETION,
6137 serde_json::json!({"prompt": "hi", "n_predict": 4}),
6138 )
6139 .await;
6140 assert_eq!(status, StatusCode::OK, "{native}");
6141 assert!(native["content"].is_string(), "{native}");
6142 assert_eq!(native["stop"], true);
6143 assert_eq!(native["stop_type"], "limit");
6144 assert_eq!(native["stopping_word"], "");
6145 assert_eq!(native["truncated"], false);
6146 assert_eq!(native["id_slot"], -1);
6147 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
6148 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
6149 assert!(
6150 native.get("choices").is_none(),
6151 "the native shape has no `choices`: {native}"
6152 );
6153
6154 let (status, openai) = post_json_uri(
6155 &app,
6156 frink_api::routes::V1_COMPLETIONS,
6157 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
6158 )
6159 .await;
6160 assert_eq!(status, StatusCode::OK);
6161 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
6162 assert!(
6163 openai.get("content").is_none(),
6164 "the OpenAI shape has no top-level `content`: {openai}"
6165 );
6166 }
6167
6168 #[tokio::test]
6172 async fn both_native_spellings_reach_the_same_handler() {
6173 let app = test_app();
6174 for route in [
6175 frink_api::routes::COMPLETION,
6176 frink_api::routes::COMPLETIONS,
6177 ] {
6178 let (status, body) = post_json_uri(
6179 &app,
6180 route,
6181 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
6182 )
6183 .await;
6184 assert_eq!(status, StatusCode::OK, "{route}: {body}");
6185 assert_eq!(body["stop"], true, "{route}");
6186 assert!(body["content"].is_string(), "{route}");
6187 }
6188
6189 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6192 let routes: Vec<&str> = stats["recent"]
6193 .as_array()
6194 .unwrap()
6195 .iter()
6196 .map(|row| row["route"].as_str().unwrap())
6197 .collect();
6198 assert!(
6199 routes.contains(&frink_api::routes::COMPLETION),
6200 "{routes:?}"
6201 );
6202 assert!(
6203 routes.contains(&frink_api::routes::COMPLETIONS),
6204 "{routes:?}"
6205 );
6206 }
6207
6208 #[tokio::test]
6214 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
6215 let app = streaming_test_app();
6216 let raw = post_sse_raw_uri(
6217 &app,
6218 frink_api::routes::COMPLETION,
6219 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
6220 )
6221 .await;
6222
6223 assert!(
6224 !raw.contains("[DONE]"),
6225 "llama.cpp's native stream has no sentinel: {raw}"
6226 );
6227 let frames: Vec<serde_json::Value> = raw
6228 .lines()
6229 .filter_map(|line| line.strip_prefix("data: "))
6230 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
6231 .collect();
6232 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
6233
6234 let (last, partials) = frames.split_last().unwrap();
6235 assert_eq!(last["stop"], true, "the last frame closes the stream");
6236 assert!(last["timings"].is_object(), "{last}");
6237 assert!(last["stop_type"].is_string(), "{last}");
6238 for partial in partials {
6239 assert_eq!(partial["stop"], false, "{partial}");
6240 assert!(partial["content"].is_string(), "{partial}");
6241 assert!(partial.get("timings").is_none(), "{partial}");
6245 assert!(partial.get("generation_settings").is_none(), "{partial}");
6246 }
6247 let streamed: String = partials
6250 .iter()
6251 .filter_map(|p| p["content"].as_str())
6252 .collect();
6253 assert_eq!(last["content"].as_str().unwrap(), streamed);
6254 }
6255
6256 #[tokio::test]
6261 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
6262 let app = test_app();
6263 for body in [
6264 serde_json::json!({"prompt": "hi"}),
6265 serde_json::json!({"prompt": "hi", "n_predict": -1}),
6266 ] {
6267 let (status, refusal) =
6268 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
6269 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
6270 assert!(
6271 refusal["error"]["message"]
6272 .as_str()
6273 .unwrap()
6274 .contains("n_predict"),
6275 "{refusal}"
6276 );
6277 }
6278 let (status, _) = post_json_uri(
6281 &app,
6282 frink_api::routes::COMPLETION,
6283 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6284 )
6285 .await;
6286 assert_eq!(status, StatusCode::OK);
6287 }
6288
6289 #[tokio::test]
6298 async fn a_stop_string_halts_the_answer_and_is_named_back() {
6299 let app = streaming_test_app();
6300 let ask = |stop: serde_json::Value| {
6301 let app = app.clone();
6302 async move {
6303 post_json_uri(
6304 &app,
6305 frink_api::routes::COMPLETION,
6306 serde_json::json!({
6307 "prompt": "hi",
6308 "n_predict": 64,
6309 "ignore_eos": true,
6310 "stop": stop,
6311 }),
6312 )
6313 .await
6314 .1
6315 }
6316 };
6317
6318 let baseline = ask(serde_json::json!([])).await;
6319 assert_eq!(baseline["stop_type"], "limit");
6320 assert_eq!(baseline["stopping_word"], "");
6321 let text = baseline["content"].as_str().unwrap().to_string();
6322 let sentinel: String = text.chars().skip(1).take(2).collect();
6328 assert_eq!(
6329 sentinel.chars().count(),
6330 2,
6331 "the fixture must produce enough output to cut: {text:?}"
6332 );
6333 let cut = text.find(&sentinel).expect("it came out of this text");
6334
6335 let stopped = ask(serde_json::json!([sentinel])).await;
6336 assert_eq!(stopped["stop_type"], "word", "{stopped}");
6337 assert_eq!(stopped["stopping_word"], sentinel);
6338 assert_eq!(
6339 stopped["content"].as_str().unwrap(),
6340 &text[..cut],
6341 "the answer must be cut at the sentinel, not run past it"
6342 );
6343 }
6344
6345 #[tokio::test]
6351 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
6352 let app = test_app();
6353
6354 let (v1_status, v1) = post_json_uri(
6355 &app,
6356 frink_api::routes::V1_TOKENIZE,
6357 serde_json::json!({"prompt": "hello"}),
6358 )
6359 .await;
6360 let (alias_status, alias) = post_json_uri(
6361 &app,
6362 frink_api::routes::TOKENIZE,
6363 serde_json::json!({"content": "hello"}),
6364 )
6365 .await;
6366 assert_eq!(v1_status, StatusCode::OK);
6367 assert_eq!(alias_status, StatusCode::OK, "{alias}");
6368 assert_eq!(v1["tokens"], alias["tokens"]);
6369 assert!(!alias["tokens"].as_array().unwrap().is_empty());
6370
6371 let (status, both_ways) = post_json_uri(
6374 &app,
6375 frink_api::routes::TOKENIZE,
6376 serde_json::json!({"prompt": "hello"}),
6377 )
6378 .await;
6379 assert_eq!(status, StatusCode::OK);
6380 assert_eq!(both_ways["tokens"], v1["tokens"]);
6381 }
6382
6383 #[tokio::test]
6388 async fn detokenize_answers_under_both_dialects_keys() {
6389 let app = test_app();
6390 for route in [
6391 frink_api::routes::DETOKENIZE,
6392 frink_api::routes::V1_DETOKENIZE,
6393 ] {
6394 let (status, body) =
6395 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
6396 assert_eq!(status, StatusCode::OK, "{route}");
6397 assert_eq!(body["text"], "hi", "{route}");
6398 assert_eq!(body["content"], body["text"], "{route}");
6399 }
6400 }
6401
6402 #[tokio::test]
6406 async fn the_alias_is_recorded_under_the_path_the_client_called() {
6407 let app = test_app();
6408 let (status, _) = post_json_uri(
6409 &app,
6410 frink_api::routes::TOKENIZE,
6411 serde_json::json!({"content": "hello"}),
6412 )
6413 .await;
6414 assert_eq!(status, StatusCode::OK);
6415
6416 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6417 let routes: Vec<&str> = stats["recent"]
6418 .as_array()
6419 .unwrap()
6420 .iter()
6421 .map(|row| row["route"].as_str().unwrap())
6422 .collect();
6423 assert!(
6424 routes.contains(&frink_api::routes::TOKENIZE),
6425 "the alias must be its own row: {routes:?}"
6426 );
6427 assert!(
6428 !routes.contains(&frink_api::routes::V1_TOKENIZE),
6429 "nothing called /v1/tokenize: {routes:?}"
6430 );
6431 }
6432
6433 #[tokio::test]
6438 async fn add_special_prepends_the_same_bos_the_decoder_would() {
6439 let mut cfg = test_dense_fixture();
6440 cfg.vocab_size = 256;
6441 let model = Model::Gguf(GgufModel {
6442 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
6443 tokenizer: Arc::new(ServerTokenizer::Byte),
6444 stop_tokens: StopTokens::default(),
6445 bos_id: Some(7),
6446 is_synthetic: true,
6447 chat_template: chat_template::PromptTemplate::plain(),
6448 });
6449 let app = test_app_with_state(Arc::new(test_state(
6450 model,
6451 ResponseCache::new(1000, Duration::from_secs(3600)),
6452 )));
6453
6454 let (_, plain) = post_json_uri(
6455 &app,
6456 frink_api::routes::TOKENIZE,
6457 serde_json::json!({"content": "hi"}),
6458 )
6459 .await;
6460 let (_, special) = post_json_uri(
6461 &app,
6462 frink_api::routes::TOKENIZE,
6463 serde_json::json!({"content": "hi", "add_special": true}),
6464 )
6465 .await;
6466
6467 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6468 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6469 assert_eq!(special["count"], 3);
6470 }
6471
6472 #[tokio::test]
6475 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6476 let app = test_app();
6477 let (status, _) = post_json_uri(
6478 &app,
6479 frink_api::routes::V1_EMBEDDINGS,
6480 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6481 )
6482 .await;
6483 assert_eq!(status, StatusCode::BAD_REQUEST);
6484
6485 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6486 let recent = stats["recent"].as_array().unwrap();
6487 assert_eq!(recent.len(), 1);
6488 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
6489 assert_eq!(recent[0]["status"], 400);
6490 assert_eq!(
6491 recent[0]["prompt_tokens"], 0,
6492 "a rejected call embedded nothing"
6493 );
6494 }
6495
6496 #[tokio::test]
6499 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6500 let app = test_app();
6501 let key = "sk-monitor-secret";
6502 let (status, _) = post_json_with_headers(
6503 &app,
6504 "/v1/chat/completions",
6505 serde_json::json!({
6506 "model": "x",
6507 "messages": [{"role": "user", "content": "hi"}],
6508 "max_tokens": 2
6509 }),
6510 &[
6511 ("authorization", &format!("Bearer {key}")),
6512 ("x-frink-client", "frink-studio"),
6513 ],
6514 )
6515 .await;
6516 assert_eq!(status, StatusCode::OK);
6517
6518 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6519 let row = stats["recent"].as_array().unwrap()[0].clone();
6520 let fingerprint = row["via_api_key"]
6521 .as_str()
6522 .expect("the row names the key that served it")
6523 .to_string();
6524 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6525 assert!(!fingerprint.contains(key));
6526 assert!(
6527 !serde_json::to_string(&stats).unwrap().contains(key),
6528 "the stats payload must not carry the key in any form"
6529 );
6530 assert_eq!(row["client"], "frink-studio");
6531 }
6532
6533 #[tokio::test]
6536 async fn different_keys_are_different_callers_and_no_key_is_null() {
6537 let app = test_app();
6538 let body = serde_json::json!({
6539 "model": "x",
6540 "messages": [{"role": "user", "content": "hi"}],
6541 "max_tokens": 1
6542 });
6543 for headers in [
6544 vec![("authorization", "Bearer key-one")],
6545 vec![("authorization", "Bearer key-two")],
6546 vec![],
6547 ] {
6548 let (status, _) =
6549 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6550 assert_eq!(status, StatusCode::OK);
6551 }
6552
6553 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6554 let recent = stats["recent"].as_array().unwrap();
6555 assert_eq!(recent.len(), 3);
6556 let one = recent[0]["via_api_key"].as_str().unwrap();
6557 let two = recent[1]["via_api_key"].as_str().unwrap();
6558 assert_ne!(one, two, "two keys must not collapse into one caller");
6559 assert!(
6560 recent[2]["via_api_key"].is_null(),
6561 "an unauthenticated call is null, not a fingerprint of nothing"
6562 );
6563 assert!(recent[2]["client"].is_null());
6564 }
6565
6566 #[tokio::test]
6571 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6572 let state = Arc::new(test_state(
6573 named_test_model("really-loaded", 256),
6574 ResponseCache::new(4, Duration::from_secs(60)),
6575 ));
6576 let app = test_app_with_state(Arc::clone(&state));
6577
6578 let (status, _) = post_json_uri(
6579 &app,
6580 "/v1/chat/completions",
6581 serde_json::json!({
6582 "model": "gpt-4-turbo-that-is-not-here",
6583 "messages": [{"role": "user", "content": "hi"}],
6584 "max_tokens": 2
6585 }),
6586 )
6587 .await;
6588 assert_eq!(status, StatusCode::OK);
6589
6590 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6591 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6592
6593 state.swap_active(None);
6596 let (status, _) = post_json_uri(
6597 &app,
6598 "/v1/chat/completions",
6599 serde_json::json!({
6600 "model": "gpt-4-turbo-that-is-not-here",
6601 "messages": [{"role": "user", "content": "hi"}]
6602 }),
6603 )
6604 .await;
6605 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6606 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6607 let recent = stats["recent"].as_array().unwrap();
6608 assert!(recent[recent.len() - 1]["model"].is_null());
6609 }
6610
6611 #[tokio::test]
6615 async fn a_streamed_row_names_the_model_it_decoded_against() {
6616 let state = Arc::new(test_state(
6617 named_test_model("model-before", 256),
6618 ResponseCache::new(4, Duration::from_secs(60)),
6619 ));
6620 let app = test_app_with_state(Arc::clone(&state));
6621 let _ = post_sse_raw(&app, resumable_request()).await;
6622 active_model(&state, "model-after");
6624
6625 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6626 assert_eq!(stats["recent"][0]["model"], "model-before");
6627 }
6628
6629 #[tokio::test]
6632 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6633 let app = test_app();
6634 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6635 assert_eq!(status, StatusCode::OK);
6636 assert!(
6637 stats["queue_depth"].is_null(),
6638 "without continuous batching nothing queues, so there is nothing to measure"
6639 );
6640 assert!(stats["queue_rejected_total"].is_null());
6641 assert_eq!(
6642 stats["generating_now"], 0,
6643 "work in progress is measured and really is zero here"
6644 );
6645 }
6646
6647 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6652 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
6653 }
6654
6655 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6659 use http_body_util::BodyExt;
6660 use tower::ServiceExt;
6661
6662 let response = app
6663 .clone()
6664 .oneshot(
6665 axum::http::Request::builder()
6666 .method("POST")
6667 .uri(uri)
6668 .header("content-type", "application/json")
6669 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6670 .unwrap(),
6671 )
6672 .await
6673 .unwrap();
6674 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6675 String::from_utf8(bytes.to_vec()).unwrap()
6676 }
6677
6678 async fn get_json_with_headers(
6679 app: &Router,
6680 uri: &str,
6681 headers: &[(&str, &str)],
6682 ) -> (StatusCode, serde_json::Value) {
6683 use http_body_util::BodyExt;
6684 use tower::ServiceExt;
6685
6686 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6687 for (name, value) in headers {
6688 builder = builder.header(*name, *value);
6689 }
6690 let response = app
6691 .clone()
6692 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6693 .await
6694 .unwrap();
6695 let status = response.status();
6696 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6697 (
6698 status,
6699 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6700 )
6701 }
6702
6703 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6704 body.lines()
6705 .filter_map(|line| line.strip_prefix(field))
6706 .map(str::trim)
6707 .collect()
6708 }
6709
6710 fn resumable_request() -> serde_json::Value {
6711 serde_json::json!({
6712 "model": "m",
6713 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6714 "max_tokens": 4,
6715 "temperature": 0,
6716 "stream": true,
6717 "stream_resumable": true,
6718 })
6719 }
6720
6721 #[tokio::test]
6726 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6727 let app = test_app();
6728 let body = post_sse_raw(&app, resumable_request()).await;
6729
6730 let request_id = body
6731 .lines()
6732 .find_map(|l| l.strip_prefix("data: "))
6733 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6734 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6735 .expect("the first chunk names the request");
6736
6737 let ids = sse_field(&body, "id:");
6738 let datas = sse_field(&body, "data:");
6739 assert_eq!(
6740 ids.len(),
6741 datas.len(),
6742 "every event carries an id, or a reconnect cannot name where it stopped"
6743 );
6744 for (i, id) in ids.iter().enumerate() {
6745 assert_eq!(*id, format!("{request_id}:{i}"));
6746 }
6747 let retries = sse_field(&body, "retry:");
6748 assert_eq!(
6749 retries.len(),
6750 1,
6751 "the reconnect delay is stated once, not on every event"
6752 );
6753 assert_eq!(retries[0], "1500");
6754 assert!(
6755 body.contains("data: [DONE]"),
6756 "the end of stream is still stated"
6757 );
6758 }
6759
6760 #[tokio::test]
6764 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6765 let app = test_app();
6766 let mut request = resumable_request();
6767 request["stream_resumable"] = serde_json::json!(false);
6768 let body = post_sse_raw(&app, request).await;
6769 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6770 assert!(
6771 sse_field(&body, "id:").is_empty(),
6772 "an id promises a replay this stream cannot serve"
6773 );
6774 assert!(sse_field(&body, "retry:").is_empty());
6775 }
6776
6777 #[tokio::test]
6781 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6782 let app = test_app();
6783 let body = post_sse_raw(&app, resumable_request()).await;
6784 let request_id = sse_field(&body, "id:")[0]
6785 .rsplit_once(':')
6786 .unwrap()
6787 .0
6788 .to_string();
6789 let streamed: Vec<String> = sse_field(&body, "data:")
6790 .iter()
6791 .map(|d| d.to_string())
6792 .collect();
6793
6794 let (status, polled) = get_json(
6795 &app,
6796 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6797 )
6798 .await;
6799 assert_eq!(status, StatusCode::OK);
6800 let events: Vec<String> = polled["events"]
6801 .as_array()
6802 .unwrap()
6803 .iter()
6804 .map(|e| e["data"].as_str().unwrap().to_string())
6805 .collect();
6806 assert_eq!(
6807 events, streamed,
6808 "the fallback must deliver the same answer, not a re-run of it"
6809 );
6810 assert_eq!(polled["request_id"], request_id);
6811 assert_eq!(
6812 polled["done"], false,
6813 "events were still being handed out, so the client must ask again"
6814 );
6815
6816 let next = polled["next_index"].as_u64().unwrap();
6819 let (_, drained) = get_json(
6820 &app,
6821 &format!(
6822 "{}?from={next}",
6823 frink_api::routes::v1_stream_poll(&request_id)
6824 ),
6825 )
6826 .await;
6827 assert_eq!(drained["done"], true);
6828 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6829 }
6830
6831 #[tokio::test]
6835 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6836 let app = test_app();
6837 let body = post_sse_raw(&app, resumable_request()).await;
6838 let ids = sse_field(&body, "id:");
6839 let datas: Vec<String> = sse_field(&body, "data:")
6840 .iter()
6841 .map(|d| d.to_string())
6842 .collect();
6843 assert!(
6844 ids.len() >= 3,
6845 "need a few events to resume into the middle"
6846 );
6847 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6848
6849 let (status, resumed) = get_json_with_headers(
6850 &app,
6851 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6852 &[],
6853 )
6854 .await;
6855 assert_eq!(status, StatusCode::OK);
6856 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6857
6858 let (_, tail) = get_json(
6860 &app,
6861 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6862 )
6863 .await;
6864 let tail_events: Vec<String> = tail["events"]
6865 .as_array()
6866 .unwrap()
6867 .iter()
6868 .map(|e| e["data"].as_str().unwrap().to_string())
6869 .collect();
6870 assert_eq!(tail_events, datas[2..].to_vec());
6871 }
6872
6873 #[tokio::test]
6876 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6877 use http_body_util::BodyExt;
6878 use tower::ServiceExt;
6879
6880 let app = test_app();
6881 let body = post_sse_raw(&app, resumable_request()).await;
6882 let ids = sse_field(&body, "id:");
6883 let datas: Vec<String> = sse_field(&body, "data:")
6884 .iter()
6885 .map(|d| d.to_string())
6886 .collect();
6887 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6888
6889 let response = app
6890 .clone()
6891 .oneshot(
6892 axum::http::Request::builder()
6893 .method("GET")
6894 .uri(frink_api::routes::v1_stream(&request_id))
6895 .header("last-event-id", format!("{request_id}:0"))
6896 .body(axum::body::Body::empty())
6897 .unwrap(),
6898 )
6899 .await
6900 .unwrap();
6901 assert_eq!(response.status(), StatusCode::OK);
6902 assert_eq!(
6903 response
6904 .headers()
6905 .get("x-accel-buffering")
6906 .and_then(|v| v.to_str().ok()),
6907 Some("no"),
6908 "the reconnect needs the same anti-buffering header as the stream"
6909 );
6910 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6911 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6912 assert_eq!(
6913 sse_field(&resumed, "data:")
6914 .iter()
6915 .map(|d| d.to_string())
6916 .collect::<Vec<_>>(),
6917 datas[1..].to_vec()
6918 );
6919 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6920 }
6921
6922 #[tokio::test]
6926 async fn a_last_event_id_from_another_stream_is_refused() {
6927 let app = test_app();
6928 let body = post_sse_raw(&app, resumable_request()).await;
6929 let request_id = sse_field(&body, "id:")[0]
6930 .rsplit_once(':')
6931 .unwrap()
6932 .0
6933 .to_string();
6934
6935 let (status, err) = get_json_with_headers(
6936 &app,
6937 &frink_api::routes::v1_stream(&request_id),
6938 &[("last-event-id", "chatcmpl-someone-else:3")],
6939 )
6940 .await;
6941 assert_eq!(status, StatusCode::BAD_REQUEST);
6942 assert_eq!(err["error"]["code"], "bad_last_event_id");
6943 }
6944
6945 #[tokio::test]
6949 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6950 let app = test_app();
6951 let mut request = resumable_request();
6952 request["stream_resumable"] = serde_json::json!(false);
6953 let body = post_sse_raw(&app, request).await;
6954 let request_id = body
6955 .lines()
6956 .find_map(|l| l.strip_prefix("data: "))
6957 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6958 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6959 .unwrap();
6960
6961 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6962 assert_eq!(status, StatusCode::NOT_FOUND);
6963 assert_eq!(err["error"]["code"], "stream_not_found");
6964 assert!(err["error"]["message"]
6965 .as_str()
6966 .unwrap()
6967 .contains("stream_resumable"));
6968 }
6969
6970 #[test]
6974 fn the_axum_stream_patterns_match_the_published_templates() {
6975 assert_eq!(
6976 axum_path(frink_api::routes::V1_STREAM),
6977 "/v1/stream/:request_id"
6978 );
6979 assert_eq!(
6980 axum_path(frink_api::routes::V1_STREAM_POLL),
6981 "/v1/stream/:request_id/poll"
6982 );
6983 assert_eq!(
6984 frink_api::routes::v1_stream("abc"),
6985 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6986 );
6987 }
6988
6989 #[test]
6999 fn no_published_template_reaches_the_router_with_its_braces() {
7000 for template in [
7001 frink_api::routes::V1_STREAM,
7002 frink_api::routes::V1_STREAM_POLL,
7003 frink_api::routes::V1_RESPONSE,
7004 frink_api::routes::V1_RESPONSE_CANCEL,
7005 frink_api::routes::ADMIN_TASK_CANCEL,
7006 ] {
7007 assert!(
7008 template.contains('{'),
7009 "{template} is in the template list but has no placeholder"
7010 );
7011 let mounted = axum_path(template);
7012 assert!(
7013 !mounted.contains('{') && !mounted.contains('}'),
7014 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
7015 );
7016 assert!(
7017 mounted.contains(':'),
7018 "{template} lost its placeholder entirely and would match one path only"
7019 );
7020 }
7021 }
7022
7023 #[tokio::test]
7031 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
7032 let app = test_app();
7033 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
7034 assert_eq!(status, StatusCode::NOT_FOUND);
7035 assert!(
7036 !body.is_null(),
7037 "empty body means axum never matched the route, so the id was read as a literal segment"
7038 );
7039 }
7040
7041 #[tokio::test]
7044 async fn the_task_list_starts_empty_rather_than_absent() {
7045 let app = test_app();
7046 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
7047 assert_eq!(status, StatusCode::OK);
7048 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
7049 }
7050
7051 #[tokio::test]
7060 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
7061 assert!(
7062 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
7063 "this test asserts the unconfigured behaviour"
7064 );
7065 let app = test_app();
7066 let (status, body) = post_json_uri(
7067 &app,
7068 &format!("{}?action=save", frink_api::routes::slots_id(0)),
7069 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
7070 )
7071 .await;
7072 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7073 assert!(
7074 body["error"]["message"]
7075 .as_str()
7076 .unwrap()
7077 .contains("--slot-save-path"),
7078 "{body}"
7079 );
7080 }
7081
7082 pub(crate) async fn post_json_uri(
7083 app: &Router,
7084 uri: &str,
7085 body: serde_json::Value,
7086 ) -> (StatusCode, serde_json::Value) {
7087 use http_body_util::BodyExt;
7088 use tower::ServiceExt;
7089
7090 let response = app
7091 .clone()
7092 .oneshot(
7093 axum::http::Request::builder()
7094 .method("POST")
7095 .uri(uri)
7096 .header("content-type", "application/json")
7097 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7098 .unwrap(),
7099 )
7100 .await
7101 .unwrap();
7102 let status = response.status();
7103 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7104 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
7105 (status, json)
7106 }
7107
7108 pub(crate) async fn get_json_uri(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7111 use http_body_util::BodyExt;
7112 use tower::ServiceExt;
7113
7114 let response = app
7115 .clone()
7116 .oneshot(
7117 axum::http::Request::builder()
7118 .method("GET")
7119 .uri(uri)
7120 .body(axum::body::Body::empty())
7121 .unwrap(),
7122 )
7123 .await
7124 .unwrap();
7125 let status = response.status();
7126 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7127 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
7128 (status, json)
7129 }
7130
7131 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
7132 post_json_uri(app, "/v1/chat/completions", body).await.1
7133 }
7134
7135 #[tokio::test]
7144 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
7145 let app = test_app();
7146 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
7147 assert_eq!(status, StatusCode::OK);
7148
7149 let memory = &body["memory"];
7150 if memory.is_null() {
7151 return;
7154 }
7155 assert!(
7156 memory["bytes"].as_u64().is_some_and(|b| b > 0),
7157 "a read that produced a zero is a broken read, not an idle \
7158 engine: {memory}"
7159 );
7160 assert!(
7161 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
7162 "the quantity must travel with the number: {memory}"
7163 );
7164 }
7165
7166 #[tokio::test]
7172 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
7173 let app = test_app();
7174 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
7175 assert_eq!(status, StatusCode::OK);
7176 for pool in ["kv_pages", "window_slots", "state_slots"] {
7177 assert!(
7178 body["pools"][pool].is_null(),
7179 "{pool} must be null rather than a zero row: {}",
7180 body["pools"]
7181 );
7182 }
7183 }
7184
7185 #[tokio::test]
7193 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
7194 use http_body_util::BodyExt;
7195 use tower::ServiceExt;
7196
7197 let app = test_app();
7198 let send = |body: serde_json::Value| {
7199 let app = app.clone();
7200 async move {
7201 app.oneshot(
7202 axum::http::Request::builder()
7203 .method("POST")
7204 .uri(frink_api::routes::V1_MESSAGES)
7205 .header("content-type", "application/json")
7206 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7207 .unwrap(),
7208 )
7209 .await
7210 .unwrap()
7211 }
7212 };
7213
7214 let ok = send(serde_json::json!({
7215 "model": "test",
7216 "max_tokens": 1,
7217 "messages": [{"role": "user", "content": "hi"}],
7218 }))
7219 .await;
7220 assert_eq!(ok.status(), StatusCode::OK);
7221 let id = ok
7222 .headers()
7223 .get("request-id")
7224 .expect("a served message names its id")
7225 .to_str()
7226 .unwrap()
7227 .to_string();
7228 assert!(!id.is_empty());
7229
7230 let bad = send(serde_json::json!({"model": "test"})).await;
7233 assert!(bad.status().is_client_error());
7234 let other = bad.headers().get("request-id").expect("errors too");
7235 assert_ne!(other.to_str().unwrap(), id);
7236 let _ = bad.into_body().collect().await.unwrap();
7237 }
7238
7239 #[tokio::test]
7247 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
7248 let state = Arc::new(test_state(
7249 test_model_full_byte_vocab(),
7250 ResponseCache::new(1000, Duration::from_secs(3600)),
7251 ));
7252 let app = test_app_with_state(Arc::clone(&state));
7253 let body = serde_json::json!({
7254 "model": "test",
7255 "messages": [{"role": "user", "content": "hi"}],
7256 "max_tokens": 1,
7257 });
7258
7259 state
7260 .maintenance
7261 .lock()
7262 .unwrap()
7263 .begin_rebuild()
7264 .expect("a fresh server is serving, so the rebuild starts");
7265 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
7266 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7267 assert_eq!(refused["error"]["type"], "cache_rebuilding");
7268
7269 state.maintenance.lock().unwrap().finish_rebuild(true);
7270 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
7271 assert_eq!(
7272 status,
7273 StatusCode::OK,
7274 "the gate reopens; a rebuild is not a latch"
7275 );
7276 }
7277
7278 #[tokio::test]
7283 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
7284 let app = test_app();
7285 let (status, body) = post_json_uri(
7286 &app,
7287 frink_api::routes::V1_CANCEL,
7288 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
7289 )
7290 .await;
7291 assert_eq!(status, StatusCode::NOT_FOUND);
7292 assert_eq!(body["cancelled"], serde_json::json!(false));
7293 assert_eq!(body["request_id"], "chatcmpl-never-issued");
7294 assert!(
7295 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
7296 "the verdict must carry a human reason: {body}"
7297 );
7298 }
7299
7300 #[tokio::test]
7304 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
7305 let state = Arc::new(test_state(
7306 test_model_full_byte_vocab(),
7307 ResponseCache::new(1000, Duration::from_secs(3600)),
7308 ));
7309 let app = test_app_with_state(Arc::clone(&state));
7310 let (token, _guard) = state.cancels.register("chatcmpl-live");
7311
7312 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7313 assert_eq!(status, StatusCode::OK);
7314 assert_eq!(before["generating_now"], serde_json::json!(1));
7315
7316 let (status, body) = post_json_uri(
7317 &app,
7318 frink_api::routes::V1_CANCEL,
7319 serde_json::json!({ "request_id": "chatcmpl-live" }),
7320 )
7321 .await;
7322 assert_eq!(status, StatusCode::OK);
7323 assert_eq!(body["cancelled"], serde_json::json!(true));
7324 assert!(
7325 token.is_cancelled(),
7326 "the endpoint answered ok without setting the flag the decode loop reads"
7327 );
7328 }
7329
7330 #[tokio::test]
7331 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
7332 let app = test_app();
7333 let (status, tok) =
7334 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
7335 assert_eq!(status, StatusCode::OK);
7336 let tokens = tok["tokens"].as_array().unwrap();
7337 assert_eq!(tok["count"], tokens.len());
7338 assert!(!tokens.is_empty());
7339
7340 let (status, detok) = post_json_uri(
7341 &app,
7342 "/v1/detokenize",
7343 serde_json::json!({ "tokens": tokens }),
7344 )
7345 .await;
7346 assert_eq!(status, StatusCode::OK);
7347 assert_eq!(detok["text"], "Hi");
7348
7349 let (status, emb) = post_json_uri(
7350 &app,
7351 "/v1/embeddings",
7352 serde_json::json!({
7353 "input": "Hi",
7354 "embedding_type": "mean"
7355 }),
7356 )
7357 .await;
7358 assert_eq!(status, StatusCode::OK);
7359 let vec = emb["data"][0]["embedding"].as_array().unwrap();
7360 assert!(!vec.is_empty());
7361 assert!(vec.iter().all(|v| v.as_f64().is_some()));
7362 }
7363
7364 #[tokio::test]
7370 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
7371 let app = test_app();
7372 let (status, body) = post_json_uri(
7373 &app,
7374 "/v1/embeddings",
7375 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
7376 )
7377 .await;
7378 assert_eq!(status, StatusCode::BAD_REQUEST);
7379 let msg = body["error"]["message"].as_str().unwrap();
7380 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
7381 }
7382
7383 #[tokio::test]
7388 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7389 async fn a_real_embedding_model_serves_v1_embeddings() {
7390 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7391 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7392 if !path.exists() {
7393 eprintln!("SKIP: {} not present", path.display());
7394 return;
7395 }
7396 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
7397 let mut state = test_state(
7398 test_model_full_byte_vocab(),
7399 ResponseCache::new(1000, Duration::from_secs(3600)),
7400 );
7401 state.embedding = Some(Arc::new(encoder));
7402 let app = test_app_with_state(Arc::new(state));
7403
7404 let (status, body) = post_json_uri(
7405 &app,
7406 "/v1/embeddings",
7407 serde_json::json!({ "input": ["Hello world", "a second input"] }),
7408 )
7409 .await;
7410 assert_eq!(status, StatusCode::OK, "{body}");
7411 assert_eq!(body["model"], "bge-small-en-v1.5");
7412 let data = body["data"].as_array().unwrap();
7413 assert_eq!(data.len(), 2);
7414 for (i, row) in data.iter().enumerate() {
7415 assert_eq!(row["index"], i);
7416 let v: Vec<f64> = row["embedding"]
7417 .as_array()
7418 .unwrap()
7419 .iter()
7420 .map(|x| x.as_f64().unwrap())
7421 .collect();
7422 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
7423 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
7424 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
7425 }
7426 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
7429
7430 let (status, mean) = post_json_uri(
7434 &app,
7435 "/v1/embeddings",
7436 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
7437 )
7438 .await;
7439 assert_eq!(status, StatusCode::OK);
7440 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
7441 }
7442
7443 #[tokio::test]
7455 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7456 async fn an_encoder_can_be_the_loaded_model() {
7457 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7458 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7459 if !path.exists() {
7460 eprintln!("SKIP: {} not present", path.display());
7461 return;
7462 }
7463
7464 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
7468 assert!(
7469 matches!(loaded, model::LoadedModel::Encoder(_)),
7470 "an encoder-only GGUF reached a decoder loader"
7471 );
7472 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
7473 assert!(
7474 matches!(loaded, Loaded::Encoder(_)),
7475 "the encoder did not stay an encoder through activation"
7476 );
7477 assert!(
7478 batcher.is_none() && ceiling.is_none(),
7479 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7480 );
7481
7482 let state = test_state(
7483 test_model_full_byte_vocab(),
7484 ResponseCache::new(1000, Duration::from_secs(3600)),
7485 );
7486 state.swap_active(Some(Arc::new(ActiveModel {
7487 id: None,
7488 loaded,
7489 batcher,
7490 ceiling,
7491 checkpoint_path: None,
7492 })));
7493 let app = test_app_with_state(Arc::new(state));
7494
7495 let (status, body) = post_json_uri(
7497 &app,
7498 "/v1/embeddings",
7499 serde_json::json!({ "input": "Hello world" }),
7500 )
7501 .await;
7502 assert_eq!(status, StatusCode::OK, "{body}");
7503 assert_eq!(body["model"], "bge-small-en-v1.5");
7504 let v = body["data"][0]["embedding"].as_array().unwrap();
7505 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
7506
7507 let (status, body) = post_json_uri(
7509 &app,
7510 "/v1/chat/completions",
7511 serde_json::json!({
7512 "model": "bge-small-en-v1.5",
7513 "messages": [{"role": "user", "content": "hi"}],
7514 }),
7515 )
7516 .await;
7517 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
7518 let msg = body["error"]["message"].as_str().unwrap();
7519 for fact in [
7520 "bge-small-en-v1.5",
7521 "bert",
7522 "embedding model",
7523 "/v1/embeddings",
7524 ] {
7525 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
7526 }
7527
7528 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
7530 assert_eq!(status, StatusCode::OK);
7531 let entry = &models["data"][0];
7532 assert_eq!(entry["id"], "bge-small-en-v1.5");
7533 assert_eq!(entry["frink_model_kind"], "embedding");
7534 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
7535 assert_eq!(entry["frink_n_embd"], 384);
7536 assert_eq!(entry["frink_pooling"], "CLS");
7537 assert_eq!(
7538 entry["frink_endpoints"],
7539 serde_json::json!(["/v1/embeddings"])
7540 );
7541 assert!(entry.get("supported_reasoning_efforts").is_none());
7544
7545 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
7547 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
7548 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
7549 assert_eq!(health["model"]["synthetic_weights"], false);
7550 let weights = health["capabilities"]
7551 .as_array()
7552 .unwrap()
7553 .iter()
7554 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
7555 .expect("a real-weights capability row");
7556 let detail = weights["detail"].as_str().unwrap_or_default();
7557 assert!(detail.contains("ENCODER"), "{detail}");
7558 let (status, body) = post_json_uri(
7566 &app,
7567 frink_api::routes::V1_TOKENIZE,
7568 serde_json::json!({ "content": "hello world" }),
7569 )
7570 .await;
7571 assert_eq!(
7572 status,
7573 StatusCode::OK,
7574 "an encoder has a real tokenizer: {body}"
7575 );
7576 let tokens = body["tokens"].as_array().expect("tokens array").clone();
7577 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
7578
7579 let (status, body) = post_json_uri(
7580 &app,
7581 frink_api::routes::V1_DETOKENIZE,
7582 serde_json::json!({ "tokens": tokens }),
7583 )
7584 .await;
7585 assert_eq!(status, StatusCode::OK, "{body}");
7586 let round_tripped = body["content"].as_str().expect("content").to_string();
7587 assert!(
7588 round_tripped.contains("hello") && round_tripped.contains("world"),
7589 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
7590 );
7591
7592 let (status, _) = post_json_uri(
7595 &app,
7596 "/v1/completions",
7597 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
7598 )
7599 .await;
7600 assert_eq!(
7601 status,
7602 StatusCode::NOT_IMPLEMENTED,
7603 "tokenizing an encoder must not have opened a path to generating with one"
7604 );
7605 }
7606
7607 #[tokio::test]
7612 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7613 use http_body_util::BodyExt;
7614 use tower::ServiceExt;
7615
7616 let fixture = concat!(
7617 "../frink-models/tests/fixtures/",
7618 "frink_real_moe_test.gguf"
7619 );
7620 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7621 let decoder = Decoder::from_gguf_with_expert_cache(
7622 &fixture,
7623 frink_models::config::test_moe_fixture(),
7624 Some(1024 * 1024),
7625 )
7626 .expect("MoE fixture must load store-backed");
7627
7628 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
7632 decoder.forward_token(1, 0, &mut caches);
7633
7634 let model = Model::Gguf(GgufModel {
7635 decoder: Arc::new(decoder),
7636 tokenizer: Arc::new(ServerTokenizer::Byte),
7637 stop_tokens: StopTokens::default(),
7638 bos_id: None,
7639 is_synthetic: false,
7640 chat_template: chat_template::PromptTemplate::plain(),
7641 });
7642 let state = Arc::new(test_state(
7643 model,
7644 ResponseCache::new(16, Duration::from_secs(60)),
7645 ));
7646 let app = Router::new()
7647 .route("/metrics", axum::routing::get(metrics))
7648 .route("/v1/chat/completions", post(chat_completions))
7649 .with_state(state);
7650
7651 let fetch_metrics = |app: Router| async move {
7652 let resp = app
7653 .oneshot(
7654 axum::http::Request::builder()
7655 .method("GET")
7656 .uri("/metrics")
7657 .body(axum::body::Body::empty())
7658 .unwrap(),
7659 )
7660 .await
7661 .unwrap();
7662 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7663 String::from_utf8(bytes.to_vec()).unwrap()
7664 };
7665
7666 let after = fetch_metrics(app.clone()).await;
7667 assert!(
7668 after.contains("frink_expert_cache_misses_total"),
7669 "streaming model must expose expert-cache metrics: {after}"
7670 );
7671 let misses: u64 = after
7672 .lines()
7673 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
7674 .and_then(|l| l.split_whitespace().nth(1))
7675 .and_then(|v| v.parse().ok())
7676 .expect("misses metric line must parse");
7677 assert!(
7678 misses > 0,
7679 "decode must have read experts through the store: {after}"
7680 );
7681 }
7682
7683 fn weather_tool() -> serde_json::Value {
7684 serde_json::json!({
7685 "type": "function",
7686 "function": {
7687 "name": "get_weather",
7688 "description": "Get the current weather for a location.",
7689 "parameters": {
7690 "type": "object",
7691 "properties": {"location": {"type": "string"}},
7692 "required": ["location"]
7693 }
7694 }
7695 })
7696 }
7697
7698 fn weather_tool_def() -> ToolDef {
7699 ToolDef {
7700 kind: "function".to_string(),
7701 function: ToolFunctionDef {
7702 name: "get_weather".to_string(),
7703 description: Some("Get the current weather for a location.".to_string()),
7704 parameters: Some(serde_json::json!({
7705 "type": "object",
7706 "properties": {"location": {"type": "string"}},
7707 "required": ["location"]
7708 })),
7709 },
7710 }
7711 }
7712
7713 #[test]
7714 fn tool_preamble_mentions_every_tool_name_and_description() {
7715 let preamble = tool_preamble(&[weather_tool_def()]);
7716 assert!(preamble.contains("get_weather"));
7717 assert!(preamble.contains("Get the current weather for a location."));
7718 assert!(preamble.contains("<tool_call>"));
7719 assert!(preamble.contains("</tool_call>"));
7720 }
7721
7722 #[test]
7723 fn a_real_marker_becomes_a_structured_tool_call() {
7724 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7725 let (message, finish) = build_response_message(
7726 text.to_string(),
7727 &[weather_tool_def()],
7728 output::OutputPosture::for_model("test-model"),
7729 "stop",
7730 );
7731 assert_eq!(finish, "tool_calls");
7732 let calls = message.tool_calls.expect("must carry a tool call");
7733 assert_eq!(calls[0].function.name, "get_weather");
7734 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7735 assert_eq!(parsed["location"], "Paris");
7736 }
7737
7738 #[test]
7739 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7740 let (message, finish) = build_response_message(
7741 "just an answer".to_string(),
7742 &[weather_tool_def()],
7743 output::OutputPosture::for_model("test-model"),
7744 "stop",
7745 );
7746 assert_eq!(finish, "stop");
7747 assert!(message.tool_calls.is_none());
7748 assert_eq!(message.content.as_deref(), Some("just an answer"));
7749 }
7750
7751 #[test]
7754 fn a_malformed_payload_is_not_a_tool_call() {
7755 let (message, finish) = build_response_message(
7756 "<tool_call>not valid json at all</tool_call>".to_string(),
7757 &[weather_tool_def()],
7758 output::OutputPosture::for_model("test-model"),
7759 "stop",
7760 );
7761 assert_eq!(finish, "stop");
7762 assert!(message.tool_calls.is_none());
7763 }
7764
7765 #[test]
7768 fn a_tool_that_was_never_offered_is_not_returned() {
7769 let (message, finish) = build_response_message(
7770 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7771 &[weather_tool_def()],
7772 output::OutputPosture::for_model("test-model"),
7773 "stop",
7774 );
7775 assert_eq!(finish, "stop");
7776 assert!(message.tool_calls.is_none());
7777 }
7778
7779 #[test]
7781 fn marker_text_with_no_tools_offered_stays_content() {
7782 let (message, finish) = build_response_message(
7783 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7784 &[],
7785 output::OutputPosture::for_model("test-model"),
7786 "stop",
7787 );
7788 assert_eq!(finish, "stop");
7789 assert!(message.tool_calls.is_none());
7790 assert!(message.content.is_some());
7791 }
7792
7793 #[test]
7797 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7798 let opened = std::cell::Cell::new(0usize);
7799 let mut parser = crate::policy::parser::ToolCallParser::new(
7800 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7801 vec![
7802 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7803 "write_file",
7804 serde_json::json!({"type": "object", "properties": {
7805 "path": {"type": "string"},
7806 "contents": {"type": "string"}
7807 }}),
7808 ),
7809 ],
7810 );
7811 let wire = "<tool_call><function=write_file>\
7812 <parameter=path>\n/tmp/x\n</parameter>\
7813 <parameter=contents>\nhello world\n</parameter>\
7814 </function></tool_call>";
7815
7816 let mut deltas = Vec::new();
7817 let mut text = String::new();
7818 for piece in wire.as_bytes().chunks(7) {
7819 let chunk = String::from_utf8_lossy(piece).into_owned();
7820 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7821 text.push_str(&more_text);
7822 deltas.extend(more);
7823 }
7824 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7825 text.push_str(&more_text);
7826 deltas.extend(more);
7827
7828 assert_eq!(opened.get(), 1, "one call opened");
7829 assert!(text.is_empty(), "the markers are not content: {text:?}");
7830
7831 let first = &deltas[0];
7832 assert_eq!(first.index, 0);
7833 assert_eq!(first.id.as_deref(), Some("call_0"));
7834 assert_eq!(first.kind, Some("function"));
7835 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7836
7837 let joined: String = deltas
7840 .iter()
7841 .filter_map(|d| d.function.arguments.clone())
7842 .collect();
7843 let parsed: serde_json::Value =
7844 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7845 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7846 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7847 assert!(
7848 deltas.len() >= 3,
7849 "the arguments arrived in pieces, not whole: {}",
7850 deltas.len()
7851 );
7852 assert!(
7853 deltas[1..].iter().all(|d| d.function.name.is_none()),
7854 "only the opening delta carries identity"
7855 );
7856 }
7857
7858 #[test]
7860 fn text_around_a_streamed_call_is_still_content() {
7861 let opened = std::cell::Cell::new(0usize);
7862 let mut parser = crate::policy::parser::ToolCallParser::new(
7863 crate::policy::parser::ToolCallFormat::Qwen25,
7864 vec![crate::policy::parser::tool_call::ToolSchema::new(
7865 "get_weather",
7866 )],
7867 );
7868 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7869 \"arguments\": {}}</tool_call> done";
7870 let mut text = String::new();
7871 for piece in wire.as_bytes().chunks(5) {
7872 let chunk = String::from_utf8_lossy(piece).into_owned();
7873 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7874 text.push_str(&more);
7875 }
7876 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7877 text.push_str(&more);
7878
7879 assert_eq!(opened.get(), 1);
7880 assert!(text.starts_with("let me check. "), "{text:?}");
7881 assert!(text.ends_with(" done"), "{text:?}");
7882 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7883 }
7884
7885 #[test]
7888 fn a_reasoning_block_is_split_out_of_the_answer() {
7889 let (message, finish) = build_response_message(
7890 "<think>weighing it up</think>The answer is 4.".to_string(),
7891 &[],
7892 output::OutputPosture::for_model("Qwen3-8B"),
7893 "stop",
7894 );
7895 assert_eq!(finish, "stop");
7896 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7897 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7898 }
7899
7900 #[test]
7903 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7904 let (message, _) = build_response_message(
7905 "Use the <think> tag like this.".to_string(),
7906 &[],
7907 output::OutputPosture::for_model("llama-3.1-8b"),
7908 "stop",
7909 );
7910 assert_eq!(
7911 message.content.as_deref(),
7912 Some("Use the <think> tag like this.")
7913 );
7914 assert!(message.reasoning_content.is_none());
7915 }
7916
7917 #[tokio::test]
7925 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7926 let app = test_app();
7927 let body = serde_json::json!({
7928 "model": "m",
7929 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7930 "max_tokens": 4,
7931 "temperature": 0,
7932 });
7933 let resp = post_json(&app, body).await;
7934 let message = &resp["choices"][0]["message"];
7935 assert!(message["content"].is_string());
7936 assert!(message.get("tool_calls").is_none());
7937 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7938 assert_eq!(resp["usage"]["completion_tokens"], 4);
7939 assert_eq!(
7940 resp["usage"]["total_tokens"],
7941 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7942 );
7943 }
7944
7945 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7946 use http_body_util::BodyExt;
7947 use tower::ServiceExt;
7948
7949 let response = app
7950 .clone()
7951 .oneshot(
7952 axum::http::Request::builder()
7953 .method("GET")
7954 .uri(uri)
7955 .body(axum::body::Body::empty())
7956 .unwrap(),
7957 )
7958 .await
7959 .unwrap();
7960 let status = response.status();
7961 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7962 (status, serde_json::from_slice(&bytes).unwrap())
7963 }
7964
7965 #[tokio::test]
7966 async fn health_answers_a_capability_handshake_not_a_boolean() {
7967 let app = test_app();
7968 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7969 assert_eq!(status, StatusCode::OK);
7970
7971 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7972 assert_eq!(health.state, frink_api::HealthState::Ready);
7973 assert!(health.pid > 0);
7974 assert!(health.server_time_unix_ms > 0);
7975 assert_eq!(health.last_request_age_seconds, None);
7978
7979 for id in [
7982 frink_api::health::capability::CPU,
7983 frink_api::health::capability::METAL,
7984 frink_api::health::capability::CUDA,
7985 frink_api::health::capability::REAL_WEIGHTS,
7986 frink_api::health::capability::CONTINUOUS_BATCHING,
7987 ] {
7988 let cap = health
7989 .capability(id)
7990 .unwrap_or_else(|| panic!("{id} missing"));
7991 assert!(!cap.reason.is_empty(), "{cap:?}");
7992 assert!(!cap.detail.is_empty(), "{cap:?}");
7993 }
7994 let weights = health
7998 .capability(frink_api::health::capability::REAL_WEIGHTS)
7999 .unwrap();
8000 assert!(!weights.available);
8001 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
8002 assert!(health.model.as_ref().unwrap().synthetic_weights);
8003 }
8004
8005 #[tokio::test]
8006 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
8007 let app = test_app();
8008 let _ = post_json(
8009 &app,
8010 serde_json::json!({
8011 "model": "m",
8012 "messages": [{"role": "user", "content": "\u{1}"}],
8013 "max_tokens": 1,
8014 "temperature": 0,
8015 }),
8016 )
8017 .await;
8018 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
8019 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
8020 let age = health
8021 .last_request_age_seconds
8022 .expect("a served request is evidence of liveness");
8023 assert!((0.0..5.0).contains(&age), "implausible age {age}");
8024 }
8025
8026 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
8028 use http_body_util::BodyExt;
8029 use tower::ServiceExt;
8030
8031 let response = app
8032 .clone()
8033 .oneshot(
8034 axum::http::Request::builder()
8035 .method("POST")
8036 .uri("/v1/chat/completions")
8037 .header("content-type", "application/json")
8038 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
8039 .unwrap(),
8040 )
8041 .await
8042 .unwrap();
8043 let bytes = response.into_body().collect().await.unwrap().to_bytes();
8044 String::from_utf8(bytes.to_vec())
8045 .unwrap()
8046 .lines()
8047 .filter_map(|line| line.strip_prefix("data: "))
8048 .filter(|payload| *payload != "[DONE]")
8049 .map(|payload| serde_json::from_str(payload).unwrap())
8050 .collect()
8051 }
8052
8053 #[tokio::test]
8054 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
8055 let app = test_app();
8056 let chunks = post_sse_chunks(
8057 &app,
8058 serde_json::json!({
8059 "model": "m",
8060 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8061 "max_tokens": 4,
8062 "temperature": 0,
8063 "stream": true,
8064 }),
8065 )
8066 .await;
8067
8068 assert!(!chunks.is_empty());
8069 let request_id = chunks[0]["request_id"]
8070 .as_str()
8071 .expect("the first chunk names the request")
8072 .to_string();
8073 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
8074 for (i, chunk) in chunks.iter().enumerate().skip(1) {
8077 assert!(
8078 chunk.get("request_id").is_none(),
8079 "chunk {i} repeats request_id"
8080 );
8081 }
8082 for chunk in &chunks {
8085 assert_eq!(chunk["id"], serde_json::json!(request_id));
8086 }
8087
8088 let other = post_sse_chunks(
8089 &app,
8090 serde_json::json!({
8091 "model": "m",
8092 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8093 "max_tokens": 4,
8094 "temperature": 0,
8095 "stream": true,
8096 }),
8097 )
8098 .await;
8099 assert_ne!(
8100 other[0]["request_id"].as_str().unwrap(),
8101 request_id,
8102 "two concurrent chats must not share an id"
8103 );
8104 }
8105
8106 #[tokio::test]
8107 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
8108 let app = test_app();
8109 let resp = post_json(
8110 &app,
8111 serde_json::json!({
8112 "model": "m",
8113 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8114 "max_tokens": 2,
8115 "temperature": 0,
8116 }),
8117 )
8118 .await;
8119 assert_eq!(resp["id"], resp["request_id"]);
8120 assert!(resp["request_id"]
8121 .as_str()
8122 .unwrap()
8123 .starts_with("chatcmpl-"));
8124 }
8125
8126 #[tokio::test]
8129 async fn usage_carries_separate_prefill_and_decode_timings() {
8130 let app = test_app();
8131 let resp = post_json(
8132 &app,
8133 serde_json::json!({
8134 "model": "m",
8135 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8136 "max_tokens": 4,
8137 "temperature": 0,
8138 }),
8139 )
8140 .await;
8141 let usage = &resp["usage"];
8142 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
8143 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
8144 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
8145 assert!(usage["predicted_per_second"].is_number(), "{usage}");
8146 assert!(usage.get("cached_tokens").is_none(), "{usage}");
8148 }
8149
8150 #[tokio::test]
8158 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
8159 let app = test_app();
8160 let body = serde_json::json!({
8161 "model": "m",
8162 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8163 "max_tokens": 4,
8164 "temperature": 0,
8165 "tools": [weather_tool()],
8166 });
8167 let resp = post_json(&app, body).await;
8168 let message = &resp["choices"][0]["message"];
8169 assert!(
8170 message["content"].is_string(),
8171 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
8172 );
8173 assert!(message.get("tool_calls").is_none());
8174 assert_eq!(resp["choices"][0]["finish_reason"], "length");
8178 }
8179
8180 #[tokio::test]
8184 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
8185 let app = test_app();
8186 let body = serde_json::json!({
8187 "model": "m",
8188 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8189 "max_tokens": 3,
8190 "temperature": 0,
8191 });
8192 let first = post_json(&app, body.clone()).await;
8193 assert_eq!(first["frink_cache"], "miss");
8194 let second = post_json(&app, body).await;
8195 assert_eq!(second["frink_cache"], "hit");
8196 assert_eq!(
8197 first["choices"][0]["message"]["content"],
8198 second["choices"][0]["message"]["content"]
8199 );
8200 assert_eq!(
8201 first["choices"][0]["finish_reason"],
8202 second["choices"][0]["finish_reason"]
8203 );
8204 assert_eq!(first["usage"], second["usage"]);
8205 assert_eq!(second["usage"]["completion_tokens"], 3);
8206 }
8207
8208 #[tokio::test]
8219 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
8220 let app = test_app();
8221 let plain = serde_json::json!({
8222 "model": "m",
8223 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8224 "max_tokens": 3,
8225 "temperature": 0,
8226 });
8227
8228 let first = post_json(&app, plain.clone()).await;
8229 assert_eq!(first["frink_cache"], "miss");
8230 let unconstrained = first["choices"][0]["message"]["content"]
8231 .as_str()
8232 .expect("content")
8233 .to_string();
8234
8235 let mut constrained = plain.clone();
8236 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
8237 let second = post_json(&app, constrained).await;
8238 assert_eq!(
8239 second["frink_cache"], "miss",
8240 "a grammar is part of the key, so this body has never been answered"
8241 );
8242 let constrained_answer = second["choices"][0]["message"]["content"]
8246 .as_str()
8247 .expect("content")
8248 .to_string();
8249 assert!(
8250 constrained_answer.contains("-> \"yes\"]"),
8251 "the grammar must have been compiled AND applied, not skipped \
8252 by a cache hit: {constrained_answer}"
8253 );
8254 assert_ne!(
8255 constrained_answer, unconstrained,
8256 "the constrained request was served the unconstrained answer"
8257 );
8258
8259 let third = post_json(&app, plain).await;
8263 assert_eq!(third["frink_cache"], "hit");
8264 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
8265 }
8266
8267 #[tokio::test]
8289 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
8290 let state = Arc::new(test_state(
8291 test_model_full_byte_vocab(),
8292 ResponseCache::new(1000, Duration::from_secs(3600)),
8293 ));
8294 let app = test_app_with_state(state.clone());
8295 let plain = serde_json::json!({
8296 "model": "m",
8297 "messages": [
8298 {"role": "system", "content": "Answer in JSON when it helps."},
8299 {"role": "user", "content": "\u{1}\u{2}"},
8300 ],
8301 "max_tokens": 3,
8302 "temperature": 0,
8303 });
8304
8305 let first = post_json(&app, plain.clone()).await;
8306 assert_eq!(first["frink_cache"], "miss");
8307 assert_eq!(state.cache_stats().entries, 1);
8308
8309 let mut as_json = plain.clone();
8310 as_json["response_format"] = serde_json::json!({"type": "json_object"});
8311 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
8312 assert_eq!(
8313 status,
8314 StatusCode::BAD_REQUEST,
8315 "the demo banner is not a JSON object, whoever generated it"
8316 );
8317 assert_eq!(
8318 state.cache_stats().hits,
8319 0,
8320 "a json_object request must not be answered from an entry the \
8321 JSON mask never produced"
8322 );
8323 assert_eq!(
8324 state.cache_stats().entries,
8325 2,
8326 "json_object must key its own entry, not reuse the unconstrained \
8327 one it happens to render the same prompt as"
8328 );
8329 }
8330
8331 #[tokio::test]
8343 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
8344 let app = test_app_with_state(Arc::new(test_state(
8345 test_model_full_byte_vocab_with_eos(Some(0x77)),
8346 ResponseCache::new(1000, Duration::from_secs(3600)),
8347 )));
8348 let body = serde_json::json!({
8349 "model": "m",
8350 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8351 "max_tokens": 6,
8352 "temperature": 0,
8353 });
8354
8355 let stopped = post_json(&app, body.clone()).await;
8356 assert_eq!(stopped["frink_cache"], "miss");
8357 assert_eq!(
8358 stopped["choices"][0]["finish_reason"], "stop",
8359 "the fixture is only meaningful if the model's EOS really fires here"
8360 );
8361 assert_eq!(stopped["usage"]["completion_tokens"], 1);
8362
8363 let mut ignoring = body.clone();
8364 ignoring["ignore_eos"] = serde_json::json!(true);
8365 let ran_on = post_json(&app, ignoring).await;
8366 assert_eq!(
8367 ran_on["frink_cache"], "miss",
8368 "ignore_eos is part of the key, so this body has never been answered"
8369 );
8370 assert_eq!(
8371 ran_on["usage"]["completion_tokens"], 6,
8372 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
8373 );
8374 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
8375 assert_ne!(
8376 ran_on["choices"][0]["message"]["content"],
8377 stopped["choices"][0]["message"]["content"]
8378 );
8379 }
8380
8381 #[tokio::test]
8387 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
8388 let session_app = test_app();
8389 let manual_app = test_app();
8390
8391 let turn1 = post_json(
8393 &session_app,
8394 serde_json::json!({
8395 "model": "m",
8396 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8397 "session_id": "s1",
8398 "max_tokens": 5,
8399 "temperature": 0,
8400 }),
8401 )
8402 .await;
8403 let reply1 = turn1["choices"][0]["message"]["content"]
8404 .as_str()
8405 .unwrap()
8406 .to_string();
8407
8408 let manual_turn1 = post_json(
8413 &manual_app,
8414 serde_json::json!({
8415 "model": "m",
8416 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8417 "max_tokens": 5,
8418 "temperature": 0,
8419 }),
8420 )
8421 .await;
8422 assert_eq!(
8423 manual_turn1["choices"][0]["message"]["content"]
8424 .as_str()
8425 .unwrap(),
8426 reply1
8427 );
8428
8429 let turn2 = post_json(
8431 &session_app,
8432 serde_json::json!({
8433 "model": "m",
8434 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
8435 "session_id": "s1",
8436 "max_tokens": 5,
8437 "temperature": 0,
8438 }),
8439 )
8440 .await;
8441 let reply2 = turn2["choices"][0]["message"]["content"]
8442 .as_str()
8443 .unwrap()
8444 .to_string();
8445
8446 let manual_turn2 = post_json(
8450 &manual_app,
8451 serde_json::json!({
8452 "model": "m",
8453 "messages": [
8454 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
8455 {"role": "assistant", "content": reply1},
8456 {"role": "user", "content": "\u{4}\u{5}"},
8457 ],
8458 "max_tokens": 5,
8459 "temperature": 0,
8460 }),
8461 )
8462 .await;
8463 assert_eq!(
8464 manual_turn2["choices"][0]["message"]["content"]
8465 .as_str()
8466 .unwrap(),
8467 reply2,
8468 "resuming a session must produce identical output to manually resending the full history"
8469 );
8470 }
8471
8472 #[test]
8475 fn lock_cache_recovers_from_a_poisoned_mutex() {
8476 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
8477
8478 let poison_cache = Arc::clone(&cache);
8479 let _ = std::thread::spawn(move || {
8480 let _guard = poison_cache.lock().unwrap();
8481 panic!("simulated panic while holding the lock");
8482 })
8483 .join();
8484
8485 let recovered = lock_cache(&cache);
8487 assert_eq!(recovered.stats().entries, 0);
8488 }
8489
8490 #[test]
8491 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8492 let mut req_body = serde_json::json!({
8493 "model": "m",
8494 "messages": [{"role": "user", "content": "hi"}],
8495 });
8496 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8497 assert!(
8498 req.is_cacheable(),
8499 "default (temperature 0) must be cacheable"
8500 );
8501
8502 req_body["temperature"] = serde_json::json!(0.8);
8503 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8504 assert!(
8505 !req.is_cacheable(),
8506 "unseeded sampling must never be cacheable"
8507 );
8508
8509 req_body["seed"] = serde_json::json!(42);
8510 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
8511 assert!(
8512 req.is_cacheable(),
8513 "sampling with an explicit seed is deterministic and must be cacheable"
8514 );
8515 }
8516
8517 const GRADED: &str = "{% if reasoning_effort %}\
8521 {% if reasoning_effort not in ['low','medium','high'] %}\
8522 {{ raise_exception('unsupported effort') }}\
8523 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
8524 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
8525
8526 fn graded_template() -> chat_template::PromptTemplate {
8527 chat_template::PromptTemplate::from_gguf_metadata(
8528 Some(GRADED),
8529 Some("qwen3"),
8530 false,
8531 true,
8532 None,
8533 None,
8534 )
8535 }
8536
8537 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
8538 serde_json::from_value(value).expect("request")
8539 }
8540
8541 #[test]
8548 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
8549 let req = chat_request(serde_json::json!({
8550 "model": "m",
8551 "messages": [{"role": "user", "content": "hi"}],
8552 "grammar": "root ::= \"a\"+",
8553 }));
8554 req.validate_supported_fields()
8555 .expect("a valid grammar is a valid request");
8556 let params = req
8557 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8558 .expect("a valid grammar compiles at params time too");
8559 assert!(
8560 params.grammar.is_some(),
8561 "the grammar was dropped between the wire and the sampler"
8562 );
8563 assert!(
8564 params.needs_vocab_logits(),
8565 "a grammar request that may fold lm_head into a GPU argmax is \
8566 a grammar request served unconstrained"
8567 );
8568
8569 let plain = chat_request(serde_json::json!({
8570 "model": "m",
8571 "messages": [{"role": "user", "content": "hi"}],
8572 }));
8573 assert!(plain
8574 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8575 .unwrap()
8576 .grammar
8577 .is_none());
8578 }
8579
8580 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
8581 chat_request(serde_json::json!({
8582 "model": "m",
8583 "messages": [{"role": "user", "content": "weather in Rome?"}],
8584 "tools": [weather_tool()],
8585 "tool_choice": tool_choice,
8586 }))
8587 }
8588
8589 #[test]
8593 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
8594 for choice in [
8595 serde_json::json!("required"),
8596 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
8597 ] {
8598 let req = tool_request(choice.clone());
8599 req.validate_supported_fields()
8600 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
8601 let params = req
8602 .generation_params_for_template(
8603 &graded_template(),
8604 "Qwen3-8B",
8605 crate::sampling_knobs::SamplerModel::absent(),
8606 )
8607 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
8608 let grammar = params
8609 .grammar
8610 .as_ref()
8611 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
8612 assert!(
8613 grammar.is_awaiting_trigger(),
8614 "the model must be free to think before it calls"
8615 );
8616 assert!(
8617 !grammar.allows_eog(),
8618 "{choice} must not be able to end the turn without a call"
8619 );
8620 assert!(
8626 params.needs_vocab_logits(),
8627 "{choice} would let a backend return a token id instead of logits"
8628 );
8629 assert!(
8630 !generate::greedy_gpu_fold_allowed(¶ms),
8631 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8632 );
8633 }
8634 }
8635
8636 #[test]
8638 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8639 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8640 let req = tool_request(choice.clone());
8641 req.validate_supported_fields().expect("still supported");
8642 let params = match req.generation_params_for_template(
8643 &graded_template(),
8644 "Qwen3-8B",
8645 crate::sampling_knobs::SamplerModel::absent(),
8646 ) {
8647 Ok(p) => p,
8648 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
8649 };
8650 assert!(
8651 params.grammar.is_none(),
8652 "{choice} does not force a call and must not be constrained"
8653 );
8654 }
8655 }
8656
8657 #[test]
8660 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8661 let req = chat_request(serde_json::json!({
8663 "model": "m",
8664 "messages": [{"role": "user", "content": "hi"}],
8665 "tool_choice": "required",
8666 }));
8667 let (status, _) = req
8668 .validate_supported_fields()
8669 .expect_err("nothing to call");
8670 assert_eq!(status, StatusCode::BAD_REQUEST);
8671
8672 let req =
8674 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8675 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8676 assert_eq!(status, StatusCode::BAD_REQUEST);
8677 assert_eq!(body["error"]["param"], "tool_choice");
8678
8679 let req = tool_request(serde_json::json!({"type": "function"}));
8681 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8682 assert_eq!(status, StatusCode::BAD_REQUEST);
8683
8684 let req = chat_request(serde_json::json!({
8686 "model": "m",
8687 "messages": [{"role": "user", "content": "hi"}],
8688 "tools": [weather_tool()],
8689 "tool_choice": "required",
8690 "grammar": "root ::= \"a\"+",
8691 }));
8692 let (status, _) = req
8693 .validate_supported_fields()
8694 .expect_err("a grammar and a forced call are two constraints");
8695 assert_eq!(status, StatusCode::BAD_REQUEST);
8696
8697 let req = tool_request(serde_json::json!("required"));
8703 let (status, Json(body)) = match req.generation_params_for_template(
8704 &graded_template(),
8705 "muse-glimmer-8b",
8706 crate::sampling_knobs::SamplerModel::absent(),
8707 ) {
8708 Err(e) => e,
8709 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
8710 };
8711 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8712 assert!(
8713 body["error"]["message"]
8714 .as_str()
8715 .unwrap()
8716 .contains("muse_glimmer"),
8717 "{body}"
8718 );
8719
8720 let req = tool_request(serde_json::json!("required"));
8726 let params = req
8727 .generation_params_for_template(
8728 &graded_template(),
8729 "gemma-4-E2B-it",
8730 crate::sampling_knobs::SamplerModel::absent(),
8731 )
8732 .expect("a gemma4 forced tool_choice is served");
8733 assert!(
8734 params.grammar.is_some(),
8735 "a forced tool_choice must arrive as the generation's grammar"
8736 );
8737 }
8738
8739 #[test]
8742 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8743 let req = chat_request(serde_json::json!({
8744 "model": "m",
8745 "messages": [{"role": "user", "content": "hi"}],
8746 "grammar": "root ::= \"a",
8747 }));
8748 let (status, Json(body)) = req
8749 .validate_supported_fields()
8750 .expect_err("this does not parse");
8751 assert_eq!(status, StatusCode::BAD_REQUEST);
8752 assert_eq!(body["error"]["param"], "grammar");
8753 assert!(
8754 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8755 .is_err(),
8756 "and again at params time"
8757 );
8758 }
8759
8760 #[test]
8767 fn response_format_json_schema_becomes_the_requests_grammar() {
8768 let req = chat_request(serde_json::json!({
8769 "model": "m",
8770 "messages": [{"role": "user", "content": "hi"}],
8771 "response_format": {
8772 "type": "json_schema",
8773 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8774 },
8775 }));
8776 req.validate_supported_fields()
8777 .expect("a boolean schema converts");
8778 let params = req
8779 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8780 .expect("and compiles");
8781 let grammar = params.grammar.expect("the schema is the grammar");
8782 let mut g = (*grammar).clone();
8783 g.accept_token(0, b"true").expect("a boolean is accepted");
8784 assert!(g.allows_eog(), "and completes the parse");
8785 assert!(
8786 !params.json_object,
8787 "a schema is not the json_object character-class mask"
8788 );
8789 }
8790
8791 #[test]
8795 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8796 let req = chat_request(serde_json::json!({
8797 "model": "m",
8798 "messages": [{"role": "user", "content": "hi"}],
8799 "response_format": {
8800 "type": "json_schema",
8801 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8802 },
8803 }));
8804 let (status, Json(body)) = req
8805 .validate_supported_fields()
8806 .expect_err("minimum has no grammar in this port");
8807 assert_eq!(status, StatusCode::BAD_REQUEST);
8808 assert!(
8809 body["error"]["message"]
8810 .as_str()
8811 .expect("a message")
8812 .contains("minimum"),
8813 "the refusal must name the keyword: {body}"
8814 );
8815 assert!(
8816 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8817 .is_err(),
8818 "and again at params time"
8819 );
8820 }
8821
8822 #[test]
8828 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8829 let req = chat_request(serde_json::json!({
8830 "model": "m",
8831 "messages": [{"role": "user", "content": "hi"}],
8832 "tool_choice": "required",
8833 "tools": [{
8834 "type": "function",
8835 "function": {"name": "f", "parameters": {"type": "object"}},
8836 }],
8837 "response_format": {
8838 "type": "json_schema",
8839 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8840 },
8841 }));
8842 let (status, Json(body)) = req
8843 .validate_supported_fields()
8844 .expect_err("two constraints, one generation");
8845 assert_eq!(status, StatusCode::BAD_REQUEST);
8846 assert_eq!(body["error"]["param"], "tool_choice");
8847 }
8848
8849 #[test]
8852 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8853 let req = chat_request(serde_json::json!({
8854 "model": "m",
8855 "messages": [{"role": "user", "content": "hi"}],
8856 }));
8857 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8858 }
8859
8860 #[test]
8865 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8866 let asked = chat_request(serde_json::json!({
8867 "model": "m",
8868 "messages": [{"role": "user", "content": "hi"}],
8869 "min_p": 0.07,
8870 }));
8871 assert_eq!(
8872 asked
8873 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8874 .expect("knobs")
8875 .min_p,
8876 0.07
8877 );
8878
8879 let silent = chat_request(serde_json::json!({
8880 "model": "m",
8881 "messages": [{"role": "user", "content": "hi"}],
8882 }));
8883 assert_eq!(
8884 silent
8885 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8886 .expect("knobs")
8887 .min_p,
8888 0.0,
8889 "an unset min_p must be off, not llama.cpp's CLI default"
8890 );
8891 }
8892
8893 #[test]
8902 fn no_sampler_knob_is_missing_from_the_cache_key() {
8903 let base = serde_json::json!({
8904 "model": "m",
8905 "messages": [{"role": "user", "content": "hi"}],
8906 "seed": 1,
8907 });
8908 let key_for = |body: serde_json::Value| {
8909 let req = chat_request(body);
8910 let params = req
8911 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8912 .expect("params");
8913 req.cache_key("prompt", ¶ms)
8914 };
8915 let baseline = key_for(base.clone());
8916 for (knob, value) in [
8917 ("temperature", serde_json::json!(0.5)),
8918 ("top_p", serde_json::json!(0.9)),
8919 ("min_p", serde_json::json!(0.05)),
8920 ("top_k", serde_json::json!(40)),
8921 ("repetition_penalty", serde_json::json!(1.1)),
8922 ("presence_penalty", serde_json::json!(0.3)),
8923 ("frequency_penalty", serde_json::json!(0.3)),
8924 (
8925 "samplers",
8926 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8927 ),
8928 ] {
8929 let mut body = base.clone();
8930 body[knob] = value;
8931 assert_ne!(
8932 key_for(body),
8933 baseline,
8934 "`{knob}` is not in the cache key: two requests differing \
8935 only in it would share one cached answer"
8936 );
8937 }
8938 }
8939
8940 #[test]
8950 fn no_constraint_is_missing_from_the_cache_key() {
8951 let base = serde_json::json!({
8952 "model": "m",
8953 "messages": [{"role": "user", "content": "pick one"}],
8954 });
8955 let key_for = |body: serde_json::Value| {
8956 let req = chat_request(body);
8957 let params = req
8958 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8959 .expect("params");
8960 req.cache_key("prompt", ¶ms)
8961 };
8962 let baseline = key_for(base.clone());
8963 for (field, value) in [
8964 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8965 (
8966 "response_format",
8967 serde_json::json!({"type": "json_object"}),
8968 ),
8969 (
8970 "response_format",
8971 serde_json::json!({"type": "json_schema", "json_schema": {
8972 "name": "answer",
8973 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8974 }}),
8975 ),
8976 ("ignore_eos", serde_json::json!(true)),
8977 ("stop", serde_json::json!(["\n"])),
8978 ("max_tokens", serde_json::json!(7)),
8979 ] {
8980 let mut body = base.clone();
8981 body[field] = value.clone();
8982 assert_ne!(
8983 key_for(body),
8984 baseline,
8985 "`{field}: {value}` is not in the cache key: two requests \
8986 differing only in it would share one cached answer"
8987 );
8988 }
8989 }
8990
8991 #[test]
8995 fn an_explicit_zero_output_budget_is_a_client_error() {
8996 let req = chat_request(serde_json::json!({
8997 "model": "m",
8998 "messages": [{"role": "user", "content": "hi"}],
8999 "max_tokens": 0,
9000 }));
9001 let (status, body) = req.validate_supported_fields().expect_err("rejected");
9002 assert_eq!(status, StatusCode::BAD_REQUEST);
9003 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
9004 }
9005
9006 #[test]
9009 fn a_request_can_turn_thinking_off() {
9010 let template = graded_template();
9011 for body in [
9012 serde_json::json!({
9013 "model": "m",
9014 "messages": [{"role": "user", "content": "hi"}],
9015 "reasoning_effort": "none",
9016 }),
9017 serde_json::json!({
9018 "model": "m",
9019 "messages": [{"role": "user", "content": "hi"}],
9020 "thinking": {"type": "disabled"},
9021 }),
9022 ] {
9023 let kwargs = chat_request(body).resolve_template_kwargs(&template);
9024 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
9025 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
9026 assert!(!kwargs.contains_key("reasoning_effort"));
9029 }
9030 }
9031
9032 #[test]
9035 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
9036 let template = graded_template();
9037 let kwargs = chat_request(serde_json::json!({
9038 "model": "m",
9039 "messages": [{"role": "user", "content": "hi"}],
9040 "reasoning_effort": "high",
9041 "thinking": {"type": "disabled"},
9042 }))
9043 .resolve_template_kwargs(&template);
9044 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
9045 assert!(!kwargs.contains_key("reasoning_effort"));
9046 }
9047
9048 #[test]
9051 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
9052 let req = chat_request(serde_json::json!({
9053 "model": "m",
9054 "messages": [{"role": "user", "content": "hi"}],
9055 "thinking": {"type": "disable"},
9056 }));
9057 let (status, _) = req.validate_supported_fields().expect_err("rejected");
9058 assert_eq!(status, StatusCode::BAD_REQUEST);
9059 }
9060
9061 #[test]
9064 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
9065 let template = graded_template();
9066 let kwargs = chat_request(serde_json::json!({
9067 "model": "m",
9068 "messages": [{"role": "user", "content": "hi"}],
9069 "reasoning_effort": "none",
9070 "chat_template_kwargs": {"enable_thinking": true},
9071 }))
9072 .resolve_template_kwargs(&template);
9073 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
9074 }
9075
9076 #[test]
9080 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
9081 let template = graded_template();
9082 let req = chat_request(serde_json::json!({
9083 "model": "m",
9084 "messages": [{"role": "user", "content": "hi"}],
9085 "reasoning_effort": "minimal",
9086 }));
9087 let kwargs = req.resolve_template_kwargs(&template);
9088 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
9089 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
9090 assert!(prompt.starts_with("E:low|"), "{prompt}");
9091 }
9092
9093 #[test]
9097 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
9098 let template = graded_template();
9099 let req = chat_request(serde_json::json!({
9100 "model": "m",
9101 "messages": [{"role": "user", "content": "hi"}],
9102 "chat_template_kwargs": {"reasoning_effort": "none"},
9103 }));
9104 let kwargs = req.resolve_template_kwargs(&template);
9105 assert!(!kwargs.contains_key("reasoning_effort"));
9106 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
9107 assert_eq!(prompt, "hi");
9108 }
9109
9110 #[test]
9113 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
9114 let template = graded_template();
9115 let req = chat_request(serde_json::json!({
9116 "model": "m",
9117 "messages": [{"role": "user", "content": "hi"}],
9118 "reasoning_effort": "low",
9119 "chat_template_kwargs": {"reasoning_effort": "high"},
9120 }));
9121 assert_eq!(
9122 req.resolve_template_kwargs(&template)["reasoning_effort"],
9123 serde_json::json!("high")
9124 );
9125 }
9126
9127 #[test]
9131 fn offering_tools_turns_thinking_on_by_itself() {
9132 let template = graded_template();
9133 let quiet = chat_request(serde_json::json!({
9134 "model": "m",
9135 "messages": [{"role": "user", "content": "hi"}],
9136 }));
9137 assert!(!quiet
9138 .resolve_template_kwargs(&template)
9139 .contains_key("enable_thinking"));
9140
9141 let with_tools = chat_request(serde_json::json!({
9142 "model": "m",
9143 "messages": [{"role": "user", "content": "hi"}],
9144 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
9145 }));
9146 let kwargs = with_tools.resolve_template_kwargs(&template);
9147 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
9148 let prompt =
9149 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
9150 assert!(prompt.starts_with("THINK|"), "{prompt}");
9151 }
9152
9153 #[test]
9158 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
9159 let opener = chat_template::PromptTemplate::from_gguf_metadata(
9160 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
9161 Some("qwen3"),
9162 false,
9163 true,
9164 None,
9165 None,
9166 );
9167 let req = chat_request(serde_json::json!({
9168 "model": "m",
9169 "messages": [{"role": "user", "content": "hi"}],
9170 "chat_template_kwargs": {"enable_thinking": true},
9171 }));
9172 let kwargs = req.resolve_template_kwargs(&opener);
9173 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
9174 assert!(prompt.ends_with("<think>"), "{prompt}");
9175
9176 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
9179 let (message, _) = build_response_message(
9180 "weighing it up</think>Paris.".to_string(),
9181 &[],
9182 posture,
9183 "stop",
9184 );
9185 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
9186 assert_eq!(message.content.as_deref(), Some("Paris."));
9187
9188 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
9191 let (message, _) = build_response_message(
9192 "weighing it up</think>Paris.".to_string(),
9193 &[],
9194 closed,
9195 "stop",
9196 );
9197 assert_eq!(message.reasoning_content, None);
9198 }
9199
9200 #[test]
9201 fn stop_param_accepts_both_single_string_and_array() {
9202 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
9203 "model": "m",
9204 "messages": [{"role": "user", "content": "hi"}],
9205 "stop": "END",
9206 }))
9207 .unwrap();
9208 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
9209
9210 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
9211 "model": "m",
9212 "messages": [{"role": "user", "content": "hi"}],
9213 "stop": ["A", "B"],
9214 }))
9215 .unwrap();
9216 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
9217 }
9218
9219 #[test]
9220 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
9221 let model = test_model();
9222 let result = run_generation(
9223 &model,
9224 "hello",
9225 &greedy_params(4),
9226 None,
9227 None,
9228 None,
9229 None,
9230 None,
9231 None,
9232 );
9233 assert!(matches!(
9234 result,
9235 Err(generate::DecodeError::TokenOutOfVocab { .. })
9236 ));
9237 }
9238
9239 #[test]
9243 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
9244 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9246 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
9247
9248 let holder_pool = Arc::clone(&pool);
9249 let holder = std::thread::spawn(move || {
9250 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
9251 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
9253 drop(held);
9254 });
9255 std::thread::sleep(Duration::from_millis(15));
9256
9257 let config = generate::KvPoolConfig {
9258 pool,
9259 queue_wait: Duration::ZERO,
9260 };
9261 let result = run_generation(
9262 &model,
9263 &prompt,
9264 &greedy_params(4),
9265 Some(&config),
9266 None,
9267 None,
9268 None,
9269 None,
9270 None,
9271 );
9272 assert!(matches!(
9273 result,
9274 Err(generate::DecodeError::KvPoolExhausted)
9275 ));
9276
9277 let (status, _body) = decode_error_response(result.unwrap_err());
9278 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
9279 holder.join().unwrap();
9280 }
9281
9282 #[test]
9290 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
9291 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9293 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
9295 let config = generate::KvPoolConfig {
9296 pool,
9297 queue_wait: Duration::ZERO,
9298 };
9299
9300 let result = run_generation(
9301 &model,
9302 &prompt,
9303 &greedy_params(4),
9304 Some(&config),
9305 None,
9306 None,
9307 None,
9308 None,
9309 None,
9310 );
9311 let err = result.expect_err("one block cannot hold two layers' caches");
9312 assert!(
9313 matches!(
9314 &err,
9315 generate::DecodeError::KvBudgetExceeded { binding, .. }
9316 if *binding == frink_models::Ceiling::DeviceMemory.code()
9317 ),
9318 "expected an immovable device-memory refusal, got {err:?}"
9319 );
9320 let (status, _body) = decode_error_response(err);
9321 assert_eq!(status, StatusCode::BAD_REQUEST);
9322 }
9323
9324 #[test]
9330 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
9331 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
9332 queued: 512,
9333 cap: 512,
9334 });
9335 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
9336 assert_eq!(body["error"]["retry_after_seconds"], 1);
9337 let message = body["error"]["message"].as_str().expect("message");
9338 assert!(message.contains("512"), "{message}");
9339 }
9340
9341 #[test]
9342 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
9343 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
9344 token: 99,
9345 vocab_size: 32,
9346 });
9347 assert!(
9348 body["error"]["retry_after_seconds"].is_null(),
9349 "retrying a prompt this model cannot tokenize never helps"
9350 );
9351 }
9352
9353 #[test]
9354 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
9355 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
9356 token: 99,
9357 vocab_size: 32,
9358 });
9359 assert_eq!(status, StatusCode::BAD_REQUEST);
9360 }
9361
9362 #[test]
9363 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
9364 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9366 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
9367 let config = generate::KvPoolConfig {
9368 pool: pool.clone(),
9369 queue_wait: Duration::ZERO,
9370 };
9371
9372 let produced = run_generation(
9373 &model,
9374 &prompt,
9375 &greedy_params(4),
9376 Some(&config),
9377 None,
9378 None,
9379 None,
9380 None,
9381 None,
9382 )
9383 .unwrap();
9384 assert_eq!(produced.choices[0].finish, FinishReason::Length);
9385 assert_eq!(
9386 pool.lock().unwrap().free_blocks(),
9387 2,
9388 "a completed request must return its blocks to the pool"
9389 );
9390 }
9391
9392 #[tokio::test]
9397 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
9398 let model = Arc::new(test_model());
9399 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9400
9401 let mut handles = Vec::new();
9402 for _ in 0..8 {
9403 let model = Arc::clone(&model);
9404 let prompt = prompt.clone();
9405 handles.push(tokio::task::spawn_blocking(move || {
9406 run_generation(
9407 &model,
9408 &prompt,
9409 &greedy_params(6),
9410 None,
9411 None,
9412 None,
9413 None,
9414 None,
9415 None,
9416 )
9417 .unwrap()
9418 }));
9419 }
9420
9421 let mut results = Vec::new();
9422 for h in handles {
9423 results.push(h.await.unwrap());
9424 }
9425 for r in &results[1..] {
9430 assert_eq!(r.choices, results[0].choices, "choices must match");
9434 assert_eq!(
9435 r.usage.prompt_tokens, results[0].usage.prompt_tokens,
9436 "prompt token count must match"
9437 );
9438 assert_eq!(
9439 r.usage.completion_tokens, results[0].usage.completion_tokens,
9440 "completion token count must match"
9441 );
9442 }
9443 }
9444
9445 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
9452 let mut header_entries = Vec::new();
9453 let mut data = Vec::new();
9454 for (name, shape, values) in tensors {
9455 let start = data.len();
9456 for v in values {
9457 data.extend_from_slice(&v.to_le_bytes());
9458 }
9459 let end = data.len();
9460 let shape_str = shape
9461 .iter()
9462 .map(|d| d.to_string())
9463 .collect::<Vec<_>>()
9464 .join(",");
9465 header_entries.push(format!(
9466 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
9467 ));
9468 }
9469 let header = format!("{{{}}}", header_entries.join(","));
9470 let header_bytes = header.as_bytes();
9471 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
9472 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
9473 out.extend_from_slice(header_bytes);
9474 out.extend_from_slice(&data);
9475 out
9476 }
9477
9478 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9489 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9490 use frink_models::kimi_loader::KimiRealHparams;
9491 use frink_moe::{GatingFunction, MoeLayerConfig};
9492
9493 let hidden_dim = 8;
9494 let kda_num_heads = 2;
9495 let kda_head_dim = 3;
9496 let kda_proj = kda_num_heads * kda_head_dim;
9497 let conv_kernel = 4;
9498 let dense_intermediate = 5;
9499 let vocab_size = 256;
9503 let mla_num_heads = 1;
9504 let mla_q_lora_rank = 2;
9505 let mla_kv_lora_rank = 2;
9506 let mla_qk_nope_head_dim = 2;
9507 let mla_qk_rope_head_dim = 2;
9508 let mla_v_head_dim = 2;
9509
9510 let model_cfg = frink_models::ModelConfig {
9511 rope_layers: frink_models::rope_layers::RopeLayers::All,
9512 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
9513 name: "synthetic-kimi-server-test",
9514 n_layers: 1,
9515 n_mtp_blocks: 0,
9516 hidden_dim,
9517 n_heads: 1,
9518 n_kv_heads: 1,
9519 head_dim: 4,
9520 v_head_dim: None,
9521 vocab_size,
9522 rope_theta: 10000.0,
9523 rms_norm_eps: 1e-5,
9524 post_norm_eps: 1e-5,
9525 sliding_window: None,
9526 moe: MoeLayerConfig {
9527 expert_weights_scale: 1.0,
9528 routed_weight_before_ffn: false,
9529 n_experts: 1,
9530 n_experts_active: 1,
9531 n_shared_experts: 0,
9532 hidden_dim,
9533 expert_ffn_dim: 4,
9534 gating: GatingFunction::Sigmoid,
9535 norm_topk_prob: true,
9536 expert_group_count: None,
9537 expert_group_used_count: None,
9538 },
9539 n_dense_leading_layers: 1,
9544 moe_interleave_step: None,
9545 norm_function: frink_models::norm::NormFunction::Rms,
9546 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
9547 kda_layers: vec![1],
9548 full_attn_layers: vec![],
9549 mla: MlaConfig {
9550 num_heads: mla_num_heads,
9551 q_lora_rank: mla_q_lora_rank,
9552 kv_lora_rank: mla_kv_lora_rank,
9553 qk_nope_head_dim: mla_qk_nope_head_dim,
9554 qk_rope_head_dim: mla_qk_rope_head_dim,
9555 v_head_dim: mla_v_head_dim,
9556 use_output_gate: true,
9557 rope: None,
9558 },
9559 kda: KdaConfig {
9560 num_heads: kda_num_heads,
9561 head_dim: kda_head_dim,
9562 short_conv_kernel_size: conv_kernel,
9563 gate_lower_bound: -5.0,
9564 use_full_rank_gate: true,
9565 },
9566 }),
9567 rope_freqs: None,
9568 rope_attn_factor: 1.0,
9569 rope_dim: None,
9570 rope_dim_swa: None,
9571 rope_freqs_long: None,
9572 rope_freqs_short: None,
9573 rope_orig_ctx: None,
9574 rope_layout: frink_models::config::RopeLayout::Neox,
9575 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
9576 swa_layers: frink_models::swa_layers::SwaLayers::All,
9577 attn_logit_softcap: None,
9578 final_logit_softcap: None,
9579 embedding_scale: None,
9580 residual_scale: None,
9581 normed_residual_scale: None,
9582 clamp_kqv: None,
9583 attn_temperature: None,
9584 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
9585 block_sub_norms: false,
9586 parallel_residual: false,
9587 learned_positions: false,
9588 attn_value_scale: None,
9589 alibi_max_bias: None,
9590 layer_loops: None,
9591 skip_stream: false,
9592 parallel_ssm: false,
9593 swa_chunked: false,
9594 weightless_qk_norm: false,
9595 logit_multiplier: None,
9596 attention_scale: None,
9597 rope_theta_swa: None,
9598 ffn_activation: frink_models::config::FfnActivation::Swiglu,
9599 best_effort_fields: &["synthetic test config, not a real preset"],
9600 };
9601 let hp = KimiRealHparams {
9602 hidden_dim,
9603 kda_num_heads,
9604 kda_head_dim,
9605 mla_num_heads,
9606 mla_q_lora_rank,
9607 mla_kv_lora_rank,
9608 mla_qk_nope_head_dim,
9609 mla_qk_rope_head_dim,
9610 mla_v_head_dim,
9611 dense_intermediate_dim: dense_intermediate,
9612 moe_hidden_dim: hidden_dim,
9613 moe_intermediate_dim: 4,
9614 n_experts: 1,
9615 num_shared_experts: 0,
9616 };
9617
9618 let prefix = "language_model.model.layers.0";
9622 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
9623 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
9624 name: String,
9625 shape: Vec<usize>,
9626 n: usize| {
9627 tensors.push((name, shape, vec![0.01f32; n]));
9628 };
9629 push(
9630 &mut tensors,
9631 format!("{prefix}.input_layernorm.weight"),
9632 vec![hidden_dim],
9633 hidden_dim,
9634 );
9635 push(
9636 &mut tensors,
9637 format!("{prefix}.post_attention_layernorm.weight"),
9638 vec![hidden_dim],
9639 hidden_dim,
9640 );
9641 push(
9642 &mut tensors,
9643 format!("{prefix}.self_attention_res_norm.weight"),
9644 vec![hidden_dim],
9645 hidden_dim,
9646 );
9647 push(
9648 &mut tensors,
9649 format!("{prefix}.self_attention_res_proj.weight"),
9650 vec![1, hidden_dim],
9651 hidden_dim,
9652 );
9653 push(
9654 &mut tensors,
9655 format!("{prefix}.mlp_res_norm.weight"),
9656 vec![hidden_dim],
9657 hidden_dim,
9658 );
9659 push(
9660 &mut tensors,
9661 format!("{prefix}.mlp_res_proj.weight"),
9662 vec![1, hidden_dim],
9663 hidden_dim,
9664 );
9665 push(
9666 &mut tensors,
9667 format!("{prefix}.self_attn.q_proj.weight"),
9668 vec![kda_proj, hidden_dim],
9669 kda_proj * hidden_dim,
9670 );
9671 push(
9672 &mut tensors,
9673 format!("{prefix}.self_attn.k_proj.weight"),
9674 vec![kda_proj, hidden_dim],
9675 kda_proj * hidden_dim,
9676 );
9677 push(
9678 &mut tensors,
9679 format!("{prefix}.self_attn.v_proj.weight"),
9680 vec![kda_proj, hidden_dim],
9681 kda_proj * hidden_dim,
9682 );
9683 push(
9684 &mut tensors,
9685 format!("{prefix}.self_attn.q_conv1d.weight"),
9686 vec![kda_proj, 1, conv_kernel],
9687 kda_proj * conv_kernel,
9688 );
9689 push(
9690 &mut tensors,
9691 format!("{prefix}.self_attn.k_conv1d.weight"),
9692 vec![kda_proj, 1, conv_kernel],
9693 kda_proj * conv_kernel,
9694 );
9695 push(
9696 &mut tensors,
9697 format!("{prefix}.self_attn.v_conv1d.weight"),
9698 vec![kda_proj, 1, conv_kernel],
9699 kda_proj * conv_kernel,
9700 );
9701 push(
9702 &mut tensors,
9703 format!("{prefix}.self_attn.A_log"),
9704 vec![kda_num_heads],
9705 kda_num_heads,
9706 );
9707 push(
9708 &mut tensors,
9709 format!("{prefix}.self_attn.f_a_proj.weight"),
9710 vec![kda_head_dim, hidden_dim],
9711 kda_head_dim * hidden_dim,
9712 );
9713 push(
9714 &mut tensors,
9715 format!("{prefix}.self_attn.f_b_proj.weight"),
9716 vec![kda_proj, kda_head_dim],
9717 kda_proj * kda_head_dim,
9718 );
9719 push(
9720 &mut tensors,
9721 format!("{prefix}.self_attn.dt_bias"),
9722 vec![kda_proj],
9723 kda_proj,
9724 );
9725 push(
9726 &mut tensors,
9727 format!("{prefix}.self_attn.b_proj.weight"),
9728 vec![kda_num_heads, hidden_dim],
9729 kda_num_heads * hidden_dim,
9730 );
9731 push(
9732 &mut tensors,
9733 format!("{prefix}.self_attn.g_proj.weight"),
9734 vec![kda_proj, hidden_dim],
9735 kda_proj * hidden_dim,
9736 );
9737 push(
9738 &mut tensors,
9739 format!("{prefix}.self_attn.o_norm.weight"),
9740 vec![kda_head_dim],
9741 kda_head_dim,
9742 );
9743 push(
9744 &mut tensors,
9745 format!("{prefix}.self_attn.o_proj.weight"),
9746 vec![hidden_dim, kda_proj],
9747 hidden_dim * kda_proj,
9748 );
9749 push(
9750 &mut tensors,
9751 format!("{prefix}.mlp.gate_proj.weight"),
9752 vec![dense_intermediate, hidden_dim],
9753 dense_intermediate * hidden_dim,
9754 );
9755 push(
9756 &mut tensors,
9757 format!("{prefix}.mlp.up_proj.weight"),
9758 vec![dense_intermediate, hidden_dim],
9759 dense_intermediate * hidden_dim,
9760 );
9761 push(
9762 &mut tensors,
9763 format!("{prefix}.mlp.down_proj.weight"),
9764 vec![hidden_dim, dense_intermediate],
9765 hidden_dim * dense_intermediate,
9766 );
9767 push(
9768 &mut tensors,
9769 "language_model.model.embed_tokens.weight".to_string(),
9770 vec![vocab_size, hidden_dim],
9771 vocab_size * hidden_dim,
9772 );
9773 push(
9774 &mut tensors,
9775 "language_model.lm_head.weight".to_string(),
9776 vec![vocab_size, hidden_dim],
9777 vocab_size * hidden_dim,
9778 );
9779 push(
9780 &mut tensors,
9781 "language_model.model.norm.weight".to_string(),
9782 vec![hidden_dim],
9783 hidden_dim,
9784 );
9785 push(
9786 &mut tensors,
9787 "language_model.model.output_attn_res_norm.weight".to_string(),
9788 vec![hidden_dim],
9789 hidden_dim,
9790 );
9791 push(
9792 &mut tensors,
9793 "language_model.model.output_attn_res_proj.weight".to_string(),
9794 vec![1, hidden_dim],
9795 hidden_dim,
9796 );
9797
9798 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9811 let dir = std::env::temp_dir().join(format!(
9812 "frink_server_kimi_e2e_test_{}_{}",
9813 std::process::id(),
9814 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9815 ));
9816 std::fs::create_dir_all(&dir).unwrap();
9817 let shard_bytes = write_safetensors_shard(&tensors);
9818 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9819 let map_entries: Vec<String> = tensors
9820 .iter()
9821 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9822 .collect();
9823 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9824 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9825
9826 use base64::Engine;
9830 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9831 .map(|b| {
9832 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9833 format!("{b64} {b}")
9834 })
9835 .collect();
9836 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9837
9838 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9839 .expect("must load the synthetic Kimi checkpoint end to end");
9840 std::fs::remove_dir_all(&dir).ok();
9841 loaded
9842 }
9843
9844 #[test]
9852 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9853 let loaded = build_synthetic_kimi_loaded();
9854 let state = build_app_state(
9855 StartupModels {
9856 loaded: model::LoadedModel::Kimi(loaded),
9857 embedding: None,
9858 },
9859 None,
9860 None,
9861 None,
9862 false,
9863 None,
9864 Arc::new(health::Detection::ready(health::probe_backends())),
9865 );
9866 let active = state.active().expect("a freshly built state has a model");
9867 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9868 assert!(!active.is_synthetic());
9869
9870 let produced = run_generation(
9871 active.generative().unwrap(),
9872 "hi",
9873 &greedy_params(5),
9874 None,
9875 None,
9876 None,
9877 None,
9878 None,
9879 None,
9880 )
9881 .expect("a real Kimi checkpoint must generate without error");
9882 assert!(matches!(
9883 produced.choices[0].finish,
9884 FinishReason::Length | FinishReason::Stop
9885 ));
9886 }
9887
9888 #[test]
9901 fn a_grammar_constrains_the_engine_decode_path() {
9902 let loaded = build_synthetic_kimi_loaded();
9903 let state = build_app_state(
9904 StartupModels {
9905 loaded: model::LoadedModel::Kimi(loaded),
9906 embedding: None,
9907 },
9908 None,
9909 None,
9910 None,
9911 false,
9912 None,
9913 Arc::new(health::Detection::ready(health::probe_backends())),
9914 );
9915 let active = state.active().expect("a freshly built state has a model");
9916
9917 let run = |grammar: Option<&str>| {
9918 let mut params = greedy_params(6);
9919 params.grammar = grammar.map(|src| {
9920 Arc::new(
9921 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9922 .expect("test grammar parses"),
9923 )
9924 });
9925 run_generation(
9926 active.generative().unwrap(),
9927 "hi",
9928 ¶ms,
9929 None,
9930 None,
9931 None,
9932 None,
9933 None,
9934 None,
9935 )
9936 };
9937
9938 let produced = run(None).expect("the unconstrained run must serve");
9939 let unconstrained = produced.choices[0].text.clone();
9940 assert!(
9941 unconstrained.chars().any(|c| c != 'a'),
9942 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9943 constrained run below would prove nothing"
9944 );
9945
9946 let produced =
9947 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9948 let one = produced.choices.into_iter().next().unwrap();
9949 let (finish, constrained) = (one.finish, one.text);
9950 assert!(
9951 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9952 "the engine decode path served text its grammar forbids ({constrained:?}): \
9953 the constraint was dropped between `generate_engine` and the sampler"
9954 );
9955 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9956 }
9957
9958 #[test]
9968 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9969 let loaded = build_synthetic_kimi_loaded();
9970 let state = build_app_state(
9971 StartupModels {
9972 loaded: model::LoadedModel::Kimi(loaded),
9973 embedding: None,
9974 },
9975 None,
9976 None,
9977 None,
9978 false,
9979 None,
9980 Arc::new(health::Detection::ready(health::probe_backends())),
9981 );
9982
9983 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9984 let kv_pool_config = generate::KvPoolConfig {
9985 pool: pool.clone(),
9986 queue_wait: Duration::ZERO,
9987 };
9988 let pc = Mutex::new(PrefixCache::new(4));
9989
9990 run_generation(
9991 state
9992 .active()
9993 .expect("a freshly built state has a model")
9994 .generative()
9995 .unwrap(),
9996 "hi",
9997 &greedy_params(5),
9998 Some(&kv_pool_config),
9999 None,
10000 Some(&pc),
10001 None,
10002 None,
10003 None,
10004 )
10005 .expect("a real Kimi checkpoint must generate without error");
10006
10007 assert_eq!(
10008 pool.lock().unwrap().free_blocks(),
10009 4,
10010 "the KV pool must be completely untouched by a Kimi request"
10011 );
10012 let stats = pc.lock().unwrap().stats();
10013 assert_eq!(
10014 stats.hits + stats.misses,
10015 0,
10016 "the prefix cache must never be consulted for a Kimi request"
10017 );
10018 }
10019}