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 logit_bias;
61mod logprobs;
62mod lora;
63mod mcp;
64mod model;
65mod openai_extra;
66mod output;
67mod policy;
68mod prefill_batch;
69mod reasoning_budget;
70mod reasoning_tokens;
71mod request_tail;
72mod rerank;
73mod response_cache;
74pub(crate) mod responses;
75mod resume;
76mod round_robin;
77mod sample_step;
78mod sampling_knobs;
79mod sampling_loop;
80mod security;
81mod serving;
82mod session;
83mod slots;
84mod sse;
85mod stats;
86mod stop;
87mod stream_events;
88mod tasks;
89mod token_mask;
90mod tool_grammar;
91mod unimplemented_fields;
92mod unsupported_sampling;
93mod utf8_stream;
94
95use std::cell::RefCell;
96use std::convert::Infallible;
97use std::net::SocketAddr;
98use std::path::PathBuf;
99use std::rc::Rc;
100use std::sync::{Arc, Mutex, MutexGuard};
101use std::time::Duration;
102
103use axum::{
104 extract::State,
105 http::StatusCode,
106 response::sse::{Event, Sse},
107 response::{IntoResponse, Response},
108 routing::{get, post},
109 Json, Router,
110};
111use serde::{Deserialize, Serialize};
112
113use cli::apply_cli_overrides;
114pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
115
116use frink_core::cache::KvBlockPool;
117use frink_models::kimi_tokenizer::KimiTokenizer;
118use frink_models::sampling::SamplingParams;
119use frink_models::tokenizer::{SpecialTokens, StopTokens};
120use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
121#[cfg(test)]
122use generate::FinishReason;
123use generate::GenerationParams;
124pub(crate) use loaded::{ActiveModel, Loaded, SleptModel};
125use model::ServerTokenizer;
126use rerank::encoder_endpoints;
127use response_cache::ResponseCache;
128use sampling_knobs::SamplingKnobs;
129
130#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
139 Gguf(GgufModel),
140 Kimi(KimiModel),
141 Mla(MlaModel),
142 Gemma4(Gemma4Model),
143 Glm52(Glm52Model),
144}
145
146pub(crate) struct GgufModel {
147 decoder: Arc<Decoder>,
148 tokenizer: Arc<ServerTokenizer>,
149 stop_tokens: StopTokens,
150 bos_id: Option<usize>,
151 is_synthetic: bool,
152 chat_template: chat_template::PromptTemplate,
153}
154
155pub(crate) struct KimiModel {
156 engine: KimiEngine,
157 tokenizer: KimiTokenizer,
158 stop_tokens: StopTokens,
159 chat_template: chat_template::PromptTemplate,
160}
161
162pub(crate) struct MlaModel {
163 engine: MlaEngine,
164 tokenizer: ServerTokenizer,
165 stop_tokens: StopTokens,
166 bos_id: Option<usize>,
167 name: String,
168 chat_template: chat_template::PromptTemplate,
169}
170
171pub(crate) struct Gemma4Model {
172 engine: Gemma4Engine,
173 tokenizer: ServerTokenizer,
174 stop_tokens: StopTokens,
175 bos_id: Option<usize>,
176 name: String,
177 chat_template: chat_template::PromptTemplate,
178}
179
180pub(crate) struct Glm52Model {
181 engine: frink_models::Glm52Engine,
182 tokenizer: ServerTokenizer,
183 stop_tokens: StopTokens,
184 bos_id: Option<usize>,
185 name: String,
186 chat_template: chat_template::PromptTemplate,
187}
188
189impl Model {
190 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
191 match self {
192 Model::Gguf(m) => m.chat_template.clone(),
193 Model::Kimi(m) => m.chat_template.clone(),
194 Model::Mla(m) => m.chat_template.clone(),
195 Model::Gemma4(m) => m.chat_template.clone(),
196 Model::Glm52(m) => m.chat_template.clone(),
197 }
198 }
199
200 fn is_synthetic(&self) -> bool {
205 match self {
206 Model::Gguf(m) => m.is_synthetic,
207 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
208 }
209 }
210
211 fn tokenizer_kind(&self) -> &'static str {
212 match self {
213 Model::Gguf(m) => m.tokenizer.kind(),
214 Model::Kimi(_) => "kimi-tiktoken-bpe",
215 Model::Mla(m) => m.tokenizer.kind(),
216 Model::Gemma4(m) => m.tokenizer.kind(),
217 Model::Glm52(m) => m.tokenizer.kind(),
218 }
219 }
220
221 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
225 match self {
226 Model::Gguf(m) => m.decoder.expert_store_stats(),
227 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
228 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
229 }
230 }
231
232 pub(crate) fn name(&self) -> &str {
233 match self {
234 Model::Gguf(m) => m.decoder.config.name,
235 Model::Kimi(_) => "kimi-k3",
236 Model::Mla(m) => m.name.as_str(),
237 Model::Gemma4(m) => m.name.as_str(),
238 Model::Glm52(m) => m.name.as_str(),
239 }
240 }
241
242 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
269 match self {
270 Model::Gguf(m) => m.tokenizer.encode(text, specials),
271 Model::Kimi(m) => m
272 .tokenizer
273 .encode(text, specials)
274 .into_iter()
275 .map(|id| id as usize)
276 .collect(),
277 Model::Mla(m) => m.tokenizer.encode(text, specials),
278 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
279 Model::Glm52(m) => m.tokenizer.encode(text, specials),
280 }
281 }
282
283 pub(crate) fn bos_id(&self) -> Option<usize> {
292 match self {
293 Model::Gguf(m) => m.bos_id,
294 Model::Kimi(_) => None,
295 Model::Mla(m) => m.bos_id,
296 Model::Gemma4(m) => m.bos_id,
297 Model::Glm52(m) => m.bos_id,
298 }
299 }
300
301 pub(crate) fn decode(&self, ids: &[usize]) -> String {
302 match self {
303 Model::Gguf(m) => m.tokenizer.decode(ids),
304 Model::Kimi(m) => {
305 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
306 m.tokenizer.decode(&ids32)
307 }
308 Model::Mla(m) => m.tokenizer.decode(ids),
309 Model::Gemma4(m) => m.tokenizer.decode(ids),
310 Model::Glm52(m) => m.tokenizer.decode(ids),
311 }
312 }
313
314 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
317 match self {
318 Model::Gguf(m) => {
319 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
320 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
321 }
322 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
323 }
324 }
325
326 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
332 match self {
333 Model::Gguf(m) => Some(&m.decoder),
334 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
335 }
336 }
337
338 pub(crate) fn vocab_size(&self) -> Option<usize> {
339 match self {
340 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
341 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
342 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
343 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
344 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
345 }
346 }
347
348 fn has_real_vocabulary(&self) -> bool {
355 match self {
356 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
357 Model::Kimi(_) => true,
358 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
359 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
360 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
361 }
362 }
363}
364
365impl frink_models::dry::DryVocab for Model {
371 fn n_tokens(&self) -> usize {
372 self.vocab_size().unwrap_or(0)
373 }
374
375 fn detokenize(&self, token: usize) -> String {
376 self.decode(&[token])
377 }
378
379 fn tokenize(&self, text: &str) -> Vec<usize> {
380 self.encode(text, SpecialTokens::AsText)
381 }
382}
383
384pub(crate) struct AppState {
385 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
399 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
418 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
422 pub(crate) slept: Mutex<Option<SleptModel>>,
432 pub(crate) tasks: Arc<tasks::TaskRegistry>,
434 pub(crate) cancels: Arc<cancel::CancelRegistry>,
438 pub(crate) stats: stats::Stats,
441 pub(crate) streams: resume::StreamRegistry,
444 pub(crate) model_dir: Option<PathBuf>,
446 response_cache: Mutex<ResponseCache>,
450 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
462 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
475 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
485 sessions: session::SessionStore,
491 requests_total: std::sync::atomic::AtomicU64,
492 request_errors_total: std::sync::atomic::AtomicU64,
493 started_at: std::time::Instant,
494 last_request_ms: std::sync::atomic::AtomicU64,
499 detection: Arc<health::Detection>,
501 mcp: Option<mcp::LoadedMcpConfig>,
503 pub(crate) continuous_batching_enabled: bool,
507 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
511 loading_model: Mutex<Option<String>>,
515 last_load_error: Mutex<Option<(String, String)>>,
519 pub(crate) serving: Mutex<crate::stats::ServingStats>,
524 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
529 pub(crate) footprint:
535 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
536 pub(crate) started_unix: u64,
546}
547
548pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
554
555impl AppState {
556 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
560 self.active
561 .read()
562 .unwrap_or_else(|p| p.into_inner())
563 .clone()
564 }
565
566 pub(crate) fn is_sleeping(&self) -> bool {
573 self.slept
574 .lock()
575 .unwrap_or_else(|p| p.into_inner())
576 .is_some()
577 }
578
579 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
580 if let Some(active) = self.active() {
581 return Ok(active);
582 }
583 if self.is_sleeping() {
587 return Err((
588 StatusCode::SERVICE_UNAVAILABLE,
589 Json(serde_json::json!({"error": {
590 "message": "this server is asleep; POST /wake_up to reload the model it put \
591 away",
592 "type": "server_sleeping"
593 }})),
594 ));
595 }
596 Err((
597 StatusCode::SERVICE_UNAVAILABLE,
598 Json(serde_json::json!({"error": {
599 "message": "no model is loaded; POST /admin/models/load with an id from \
600 GET /admin/models",
601 "type": "model_not_loaded"
602 }})),
603 ))
604 }
605
606 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
614 Ok(Arc::clone(self.require_active()?.generative()?))
615 }
616
617 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
625 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
626 std::mem::replace(&mut *guard, next)
627 }
628
629 fn mark_request_finished(&self) {
633 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
634 self.last_request_ms
635 .store(ms, std::sync::atomic::Ordering::Relaxed);
636 }
637
638 pub(crate) fn uptime(&self) -> Duration {
639 self.started_at.elapsed()
640 }
641
642 pub(crate) fn requests_total(&self) -> u64 {
643 self.requests_total
644 .load(std::sync::atomic::Ordering::Relaxed)
645 }
646
647 pub(crate) fn errors_total(&self) -> u64 {
648 self.request_errors_total
649 .load(std::sync::atomic::Ordering::Relaxed)
650 }
651
652 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
653 lock_cache(&self.response_cache).stats()
654 }
655
656 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
659 let last = self
660 .last_request_ms
661 .load(std::sync::atomic::Ordering::Relaxed);
662 (last > 0)
663 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
664 .map(|age| age.max(0.0))
665 }
666
667 pub(crate) fn loading_model_id(&self) -> Option<String> {
668 self.loading_model
669 .lock()
670 .unwrap_or_else(|p| p.into_inner())
671 .clone()
672 }
673
674 pub(crate) fn set_loading_model(&self, id: Option<String>) {
675 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
676 }
677
678 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
679 self.last_load_error
680 .lock()
681 .unwrap_or_else(|p| p.into_inner())
682 .clone()
683 }
684
685 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
686 *self
687 .last_load_error
688 .lock()
689 .unwrap_or_else(|p| p.into_inner()) = error;
690 }
691
692 pub(crate) fn active_model_name(&self) -> Option<String> {
701 self.active().map(|a| a.name().to_string())
702 }
703
704 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
714 self.embedding
715 .clone()
716 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
717 }
718
719 pub(crate) fn embedding_model_name(&self) -> Option<String> {
723 match self.embedding_model() {
724 Some(e) => Some(e.name().to_string()),
725 None => self.active_model_name(),
726 }
727 }
728
729 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
730 self.stats.record(stats::entry(record));
731 }
732}
733
734fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
739 cache
740 .lock()
741 .unwrap_or_else(|poisoned| poisoned.into_inner())
742}
743
744#[derive(Debug, Clone, Deserialize)]
745#[serde(untagged)]
746pub(crate) enum MessageContent {
747 Text(String),
748 Parts(Vec<ContentPart>),
749}
750
751#[derive(Debug, Clone, Deserialize)]
752struct ContentPart {
753 #[serde(rename = "type")]
754 kind: String,
755 #[serde(default)]
756 text: Option<String>,
757 #[serde(default)]
758 image_url: Option<serde_json::Value>,
759}
760
761impl MessageContent {
762 fn as_text(&self) -> String {
763 match self {
764 Self::Text(s) => s.clone(),
765 Self::Parts(parts) => parts
766 .iter()
767 .filter_map(|p| p.text.as_deref())
768 .collect::<Vec<_>>()
769 .join(""),
770 }
771 }
772
773 fn has_image(&self) -> bool {
774 match self {
775 Self::Text(_) => false,
776 Self::Parts(parts) => parts
777 .iter()
778 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
779 }
780 }
781}
782
783#[derive(Debug, Clone, Deserialize)]
784pub(crate) struct ChatMessage {
785 pub(crate) role: String,
786 #[serde(default)]
790 pub(crate) content: Option<MessageContent>,
791 #[serde(default)]
795 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
796 #[serde(default)]
801 #[allow(dead_code)]
802 pub(crate) tool_call_id: Option<String>,
803 #[serde(default, alias = "reasoning")]
819 pub(crate) reasoning_content: Option<String>,
820}
821
822impl ChatMessage {
823 fn rendered_content(&self) -> String {
830 let mut out = self
831 .content
832 .as_ref()
833 .map(MessageContent::as_text)
834 .unwrap_or_default();
835 if let Some(calls) = &self.tool_calls {
836 for call in calls {
837 out.push_str(&format!(
838 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
839 call.function.name, call.function.arguments
840 ));
841 }
842 }
843 out
844 }
845}
846
847#[derive(Debug, Clone, Deserialize)]
848pub(crate) struct ToolCallIn {
849 #[serde(default)]
850 #[allow(dead_code)]
851 id: String,
852 #[serde(rename = "type", default)]
853 #[allow(dead_code)]
854 kind: String,
855 function: ToolCallFunctionIn,
856}
857
858#[derive(Debug, Clone, Deserialize)]
859struct ToolCallFunctionIn {
860 name: String,
861 arguments: String,
866}
867
868#[derive(Debug, Clone, Deserialize)]
871struct ToolDef {
872 #[serde(rename = "type", default)]
873 #[allow(dead_code)]
874 kind: String,
875 function: ToolFunctionDef,
876}
877
878#[derive(Debug, Clone, Deserialize)]
879struct ToolFunctionDef {
880 name: String,
881 #[serde(default)]
882 description: Option<String>,
883 #[serde(default)]
884 parameters: Option<serde_json::Value>,
885}
886
887#[derive(Debug, Clone, Deserialize)]
898#[serde(untagged)]
899enum ToolChoice {
900 Mode(String),
901 Specific(serde_json::Value),
902}
903
904#[derive(Deserialize)]
907#[serde(untagged)]
908enum StopParam {
909 One(String),
910 Many(Vec<String>),
911}
912
913#[derive(Deserialize)]
914struct ChatCompletionRequest {
915 model: String,
916 messages: Vec<ChatMessage>,
917 #[serde(default = "default_max_tokens")]
918 max_tokens: usize,
919 #[serde(default)]
920 temperature: Option<f32>,
921 #[serde(default)]
922 top_p: Option<f32>,
923 #[serde(default)]
928 min_p: Option<f32>,
929 #[serde(default)]
930 top_k: Option<usize>,
931 #[serde(default)]
932 repetition_penalty: Option<f32>,
933 #[serde(flatten)]
937 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
938 #[serde(flatten)]
942 unimplemented: crate::unimplemented_fields::UnimplementedFields,
943 #[serde(default)]
944 seed: Option<u64>,
945 #[serde(default)]
946 stop: Option<StopParam>,
947 #[serde(default)]
948 stream: Option<bool>,
949 #[serde(default)]
962 stream_resumable: Option<bool>,
963 #[serde(default)]
974 ignore_eos: Option<bool>,
975 #[serde(default)]
976 tools: Vec<ToolDef>,
977 #[serde(default)]
978 tool_choice: Option<ToolChoice>,
979 #[serde(default)]
987 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
988 #[serde(default)]
996 reasoning_effort: Option<String>,
997 #[serde(default)]
1001 thinking: Option<ThinkingSwitch>,
1002 #[serde(default)]
1007 session_id: Option<String>,
1008 #[serde(default, deserialize_with = "continuation::deserialize")]
1016 continue_final_message: continuation::ContinueFinalMessage,
1017 #[serde(default, alias = "thinking_budget_tokens")]
1025 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
1026 #[serde(default)]
1028 logprobs: Option<bool>,
1029 #[serde(default)]
1030 top_logprobs: Option<u32>,
1031 #[serde(default)]
1032 presence_penalty: Option<f32>,
1033 #[serde(default)]
1034 frequency_penalty: Option<f32>,
1035 #[serde(default)]
1036 response_format: Option<serde_json::Value>,
1037 #[serde(default)]
1044 logit_bias: Option<serde_json::Value>,
1045 #[serde(default)]
1049 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1050 #[serde(default)]
1060 samplers: Option<serde_json::Value>,
1061 #[serde(default)]
1071 grammar: Option<String>,
1072}
1073
1074const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1086
1087#[derive(Debug, Clone, Deserialize)]
1089pub(crate) struct ThinkingSwitch {
1090 #[serde(rename = "type")]
1091 pub(crate) kind: String,
1092}
1093
1094const THINKING_KWARG_KEYS: [&str; 4] = [
1098 "enable_thinking",
1099 "thinking",
1100 "thinking_mode",
1101 "reasoning_effort",
1102];
1103
1104const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1108
1109fn default_max_tokens() -> usize {
1110 DEFAULT_CHAT_MAX_TOKENS
1111}
1112
1113impl ChatCompletionRequest {
1114 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1123 let mut knobs = SamplingKnobs {
1124 temperature: self.temperature,
1125 top_p: self.top_p,
1126 min_p: self.min_p,
1127 top_k: self.top_k,
1128 repetition_penalty: self.repetition_penalty,
1129 presence_penalty: self.presence_penalty,
1130 frequency_penalty: self.frequency_penalty,
1131 penalty_last_n: None,
1135 sampler_order: unsupported_sampling::parse_sampler_order(
1136 self.samplers.as_ref(),
1137 "/v1/chat/completions",
1138 )?,
1139 ..SamplingKnobs::default()
1140 };
1141 self.extra_samplers.apply(&mut knobs);
1142 Ok(knobs)
1143 }
1144
1145 fn sampling_params(
1146 &self,
1147 model: crate::sampling_knobs::SamplerModel<'_>,
1148 ) -> Result<SamplingParams, ApiError> {
1149 self.sampling_knobs()?.resolve(model).map_err(|e| {
1150 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1151 })
1152 }
1153
1154 fn stop_sequences(&self) -> Vec<String> {
1155 self.stop
1156 .as_ref()
1157 .map(|s| match s {
1158 StopParam::One(v) => vec![v.clone()],
1159 StopParam::Many(v) => v.clone(),
1160 })
1161 .unwrap_or_default()
1162 }
1163
1164 fn n_logprobs(&self) -> Result<Option<usize>, ApiError> {
1178 const MAX: u32 = 20;
1179 match (self.logprobs, self.top_logprobs) {
1180 (Some(true), Some(n)) if n > MAX => Err(invalid_request(
1181 &format!(
1182 "`top_logprobs` is {n}; this server reports at most {MAX} alternatives per \
1183 position, as upstream does"
1184 ),
1185 "top_logprobs",
1186 )),
1187 (Some(true), Some(n)) => Ok(Some(n as usize)),
1188 (Some(true), None) => Ok(Some(0)),
1192 (_, Some(_)) => Err(invalid_request(
1193 "`top_logprobs` requires `logprobs: true`",
1194 "top_logprobs",
1195 )),
1196 _ => Ok(None),
1197 }
1198 }
1199
1200 fn tools_active(&self) -> bool {
1201 !self.tools.is_empty()
1202 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1203 }
1204
1205 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1217 match &self.tool_choice {
1218 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1219 Some(ToolChoice::Specific(value)) => {
1220 let name = value
1224 .get("function")
1225 .and_then(|f| f.get("name"))
1226 .or_else(|| value.get("name"))
1227 .and_then(|n| n.as_str());
1228 match name {
1229 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1230 None => Err(invalid_request(
1231 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1232 function.name",
1233 "tool_choice",
1234 )),
1235 }
1236 }
1237 _ => Ok(None),
1238 }
1239 }
1240
1241 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1243 self.tools
1244 .iter()
1245 .map(|t| tool_grammar::ToolSpec {
1246 name: &t.function.name,
1247 parameters: t.function.parameters.as_ref(),
1248 })
1249 .collect()
1250 }
1251
1252 fn thinking_direction(&self) -> Option<bool> {
1295 if let Some(switch) = &self.thinking {
1296 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1297 "disabled" => Some(false),
1298 "enabled" => Some(true),
1299 _ => None,
1302 };
1303 }
1304 let effort = self.reasoning_effort.as_ref()?;
1305 DISABLE_EFFORTS
1306 .contains(&effort.trim().to_ascii_lowercase().as_str())
1307 .then_some(false)
1308 }
1309
1310 fn resolve_template_kwargs(
1311 &self,
1312 template: &chat_template::PromptTemplate,
1313 ) -> serde_json::Map<String, serde_json::Value> {
1314 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1315 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1319
1320 if !caller_steered {
1321 match self.thinking_direction() {
1322 Some(false) => {
1323 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1324 kwargs.insert(k, v);
1325 }
1326 return kwargs;
1329 }
1330 Some(true) => {
1331 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1332 kwargs.insert(k, v);
1333 }
1334 }
1335 None => {}
1336 }
1337 if let Some(effort) = &self.reasoning_effort {
1338 kwargs
1339 .entry("reasoning_effort".to_string())
1340 .or_insert_with(|| serde_json::json!(effort));
1341 }
1342 }
1343
1344 let offered: Vec<serde_json::Value> = if self.tools_active() {
1345 self.tools.iter().map(chat_template::tool_json).collect()
1346 } else {
1347 Vec::new()
1348 };
1349 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1350 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1351 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1352 kwargs.entry(k).or_insert(v);
1353 }
1354 }
1355 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1356 crate::policy::effort::EffortMapping::Mapped(to) => {
1357 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1358 }
1359 crate::policy::effort::EffortMapping::Dropped => {
1360 tracing::debug!(
1361 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1362 enough, so its own default applies"
1363 );
1364 }
1365 crate::policy::effort::EffortMapping::Unchanged => {}
1366 }
1367 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1368 kwargs
1369 }
1370
1371 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1374 if self.max_tokens == 0 {
1381 return Err(invalid_request(
1382 "max_tokens must be at least 1",
1383 "max_tokens",
1384 ));
1385 }
1386 if let Some(switch) = &self.thinking {
1391 let kind = switch.kind.trim().to_ascii_lowercase();
1392 if kind != "enabled" && kind != "disabled" {
1393 return Err(invalid_request(
1394 "thinking.type must be \"enabled\" or \"disabled\"",
1395 "thinking.type",
1396 ));
1397 }
1398 }
1399 for msg in &self.messages {
1400 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1401 return Err(unsupported_feature(
1402 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1403 ));
1404 }
1405 }
1406 self.n_logprobs()?;
1410 self.unimplemented.refuse("/v1/chat/completions")?;
1414 logit_bias::LogitBias::parse(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1418 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1422 let stated_grammar =
1437 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1438 if let Some(forced) = self.forced_tool_choice()? {
1446 if self.tools.is_empty() {
1447 return Err(invalid_request(
1448 "tool_choice forces a tool call, but no tools were offered",
1449 "tool_choice",
1450 ));
1451 }
1452 if let tool_grammar::Forced::Named(name) = forced {
1453 if !self.tools.iter().any(|t| t.function.name == name) {
1454 return Err(invalid_request(
1455 &format!(
1456 "tool_choice names {name:?}, which is not one of the tools offered"
1457 ),
1458 "tool_choice",
1459 ));
1460 }
1461 }
1462 if stated_grammar.is_some() {
1472 return Err(invalid_request(
1473 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1474 are two different constraints on the same generation; send one",
1475 "tool_choice",
1476 ));
1477 }
1478 if self.json_object_mode() {
1479 return Err(invalid_request(
1480 "a forced tool_choice cannot be combined with response_format json_object: \
1481 the tool-call markers are not JSON",
1482 "tool_choice",
1483 ));
1484 }
1485 }
1486 Ok(())
1487 }
1488
1489 fn effective_stop_sequences(&self) -> Vec<String> {
1496 let mut stop = self.stop_sequences();
1497 if self.tools_active() {
1498 stop.push("</tool_call>".to_string());
1499 }
1500 stop
1501 }
1502
1503 fn json_object_mode(&self) -> bool {
1504 self.response_format
1505 .as_ref()
1506 .and_then(|v| v.get("type"))
1507 .and_then(|v| v.as_str())
1508 == Some("json_object")
1509 }
1510}
1511
1512#[derive(Serialize)]
1513struct ChatCompletionChoice {
1514 index: usize,
1515 message: ChatCompletionResponseMessage,
1516 finish_reason: &'static str,
1517 #[serde(skip_serializing_if = "Option::is_none")]
1522 logprobs: Option<serde_json::Value>,
1523}
1524
1525#[derive(Serialize)]
1526struct ChatCompletionResponseMessage {
1527 role: &'static str,
1528 #[serde(skip_serializing_if = "Option::is_none")]
1529 content: Option<String>,
1530 #[serde(skip_serializing_if = "Option::is_none")]
1534 reasoning_content: Option<String>,
1535 #[serde(skip_serializing_if = "Option::is_none")]
1536 tool_calls: Option<Vec<ToolCallOut>>,
1537}
1538
1539#[derive(Serialize, Clone)]
1540struct ToolCallOut {
1541 id: String,
1542 #[serde(rename = "type")]
1543 kind: &'static str,
1544 function: ToolCallFunctionOut,
1545}
1546
1547#[derive(Serialize, Clone)]
1555struct ToolCallDelta {
1556 index: usize,
1557 #[serde(skip_serializing_if = "Option::is_none")]
1558 id: Option<String>,
1559 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1560 kind: Option<&'static str>,
1561 function: ToolCallFunctionDelta,
1562}
1563
1564#[derive(Serialize, Clone, Default)]
1565struct ToolCallFunctionDelta {
1566 #[serde(skip_serializing_if = "Option::is_none")]
1567 name: Option<String>,
1568 #[serde(skip_serializing_if = "Option::is_none")]
1571 arguments: Option<String>,
1572}
1573
1574impl ToolCallDelta {
1575 fn whole(index: usize, name: String, arguments: String) -> Self {
1577 ToolCallDelta {
1578 index,
1579 id: Some(format!("call_{index}")),
1580 kind: Some("function"),
1581 function: ToolCallFunctionDelta {
1582 name: Some(name),
1583 arguments: Some(arguments),
1584 },
1585 }
1586 }
1587
1588 fn opening(index: usize, name: String) -> Self {
1590 ToolCallDelta {
1591 index,
1592 id: Some(format!("call_{index}")),
1593 kind: Some("function"),
1594 function: ToolCallFunctionDelta {
1595 name: Some(name),
1596 arguments: Some(String::new()),
1597 },
1598 }
1599 }
1600
1601 fn arguments(index: usize, fragment: String) -> Self {
1603 ToolCallDelta {
1604 index,
1605 id: None,
1606 kind: None,
1607 function: ToolCallFunctionDelta {
1608 name: None,
1609 arguments: Some(fragment),
1610 },
1611 }
1612 }
1613}
1614
1615#[derive(Serialize, Clone)]
1616struct ToolCallFunctionOut {
1617 name: String,
1618 arguments: String,
1622}
1623
1624#[derive(Serialize)]
1625struct ChatCompletionResponse {
1626 id: String,
1627 request_id: String,
1633 object: &'static str,
1634 model: String,
1635 choices: Vec<ChatCompletionChoice>,
1636 usage: generate::Usage,
1641 frink_cache: &'static str,
1652}
1653
1654#[derive(Serialize)]
1655struct ChatCompletionChunkDelta {
1656 #[serde(skip_serializing_if = "Option::is_none")]
1657 role: Option<&'static str>,
1658 #[serde(skip_serializing_if = "Option::is_none")]
1659 content: Option<String>,
1660 #[serde(skip_serializing_if = "Option::is_none")]
1662 reasoning_content: Option<String>,
1663 #[serde(skip_serializing_if = "Option::is_none")]
1664 tool_calls: Option<Vec<ToolCallDelta>>,
1665}
1666
1667#[derive(Serialize)]
1668struct ChatCompletionChunkChoice {
1669 index: usize,
1670 delta: ChatCompletionChunkDelta,
1671 finish_reason: Option<&'static str>,
1672}
1673
1674#[derive(Serialize)]
1675struct ChatCompletionChunk {
1676 id: String,
1677 #[serde(skip_serializing_if = "Option::is_none")]
1684 request_id: Option<String>,
1685 object: &'static str,
1686 model: String,
1687 choices: Vec<ChatCompletionChunkChoice>,
1688 #[serde(skip_serializing_if = "Option::is_none")]
1691 usage: Option<generate::Usage>,
1692}
1693
1694async fn health(State(state): State<Arc<AppState>>) -> Response {
1700 let snapshot = state.detection.snapshot();
1701 let mut capabilities = snapshot.capabilities;
1702 let active = state.active();
1703
1704 capabilities.push(match active.as_deref() {
1707 None => frink_api::Capability::unavailable(
1712 frink_api::health::capability::REAL_WEIGHTS,
1713 frink_api::health::reason::MODEL_NOT_LOADED,
1714 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1715 ),
1716 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1717 frink_api::health::capability::REAL_WEIGHTS,
1718 frink_api::health::reason::MODEL_NOT_LOADED,
1719 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1720 checkpoint. Output from this model is noise.",
1721 ),
1722 Some(active) if active.encoder().is_some() => {
1732 let endpoints = active
1733 .encoder()
1734 .map(|e| encoder_endpoints(e))
1735 .unwrap_or_default();
1736 let served_by = match endpoints.is_empty() {
1737 true => "no endpoint in this build serves it".to_string(),
1738 false => format!("served by {}", endpoints.join(" and ")),
1739 };
1740 frink_api::Capability::available(
1741 frink_api::health::capability::REAL_WEIGHTS,
1742 format!(
1743 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1744 {served_by}; generation endpoints refuse it.",
1745 active.name(),
1746 ),
1747 )
1748 }
1749 Some(active) => frink_api::Capability::available(
1750 frink_api::health::capability::REAL_WEIGHTS,
1751 format!("Serving the real checkpoint '{}'.", active.name()),
1752 ),
1753 });
1754 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1755 frink_api::Capability::available(
1756 frink_api::health::capability::CONTINUOUS_BATCHING,
1757 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1758 "On by default on Metal. Concurrent requests share one batched decode worker."
1759 } else {
1760 "Concurrent requests share one batched decode step."
1761 },
1762 )
1763 } else if state.metal_private_decode_gate.is_some() {
1764 frink_api::Capability::unavailable(
1765 frink_api::health::capability::CONTINUOUS_BATCHING,
1766 frink_api::health::reason::DISABLED,
1767 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1768 )
1769 } else {
1770 frink_api::Capability::unavailable(
1771 frink_api::health::capability::CONTINUOUS_BATCHING,
1772 frink_api::health::reason::DISABLED,
1773 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1774 )
1775 });
1776
1777 let last_request_ms = state
1778 .last_request_ms
1779 .load(std::sync::atomic::Ordering::Relaxed);
1780 let uptime = state.started_at.elapsed();
1781 let health_state = if active.is_none() {
1789 frink_api::HealthState::Unavailable
1790 } else {
1791 snapshot.state
1792 };
1793 let body = frink_api::HealthResponse {
1794 state: health_state,
1795 reason: match health_state {
1796 frink_api::HealthState::Ready => None,
1797 frink_api::HealthState::Unavailable => {
1798 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1799 }
1800 frink_api::HealthState::Detecting => {
1801 Some(frink_api::health::reason::DETECTING.to_string())
1802 }
1803 },
1804 detail: match health_state {
1805 frink_api::HealthState::Ready => None,
1806 frink_api::HealthState::Unavailable => Some(
1807 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1808 .to_string(),
1809 ),
1810 frink_api::HealthState::Detecting => {
1811 Some("Probing available compute backends.".to_string())
1812 }
1813 },
1814 model: active
1815 .as_deref()
1816 .map(|active| frink_api::health::ModelSummary {
1817 id: active.name().to_string(),
1818 tokenizer: active.tokenizer_kind().to_string(),
1819 synthetic_weights: active.is_synthetic(),
1820 }),
1821 capabilities,
1822 version: env!("CARGO_PKG_VERSION").to_string(),
1823 pid: std::process::id(),
1824 uptime_seconds: uptime.as_secs_f64(),
1825 server_time_unix_ms: std::time::SystemTime::now()
1826 .duration_since(std::time::UNIX_EPOCH)
1827 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1828 .unwrap_or(0),
1829 last_request_age_seconds: (last_request_ms > 0)
1830 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1831 .map(|age| age.max(0.0)),
1832 };
1833
1834 let status =
1835 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1836 (status, Json(body)).into_response()
1837}
1838
1839async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1840 let Some(active) = state.active() else {
1844 return Json(serde_json::json!({ "object": "list", "data": [] }));
1845 };
1846 let mut model_entry = serde_json::json!({
1847 "id": active.name(),
1848 "object": "model",
1849 "frink_synthetic_weights": active.is_synthetic(),
1850 "frink_tokenizer": active.tokenizer_kind(),
1851 });
1852 if let Some(encoder) = active.encoder() {
1858 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1859 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1860 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1861 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1862 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1863 }
1864 if let Some(model) = active.generative_opt() {
1872 let parser_configured = active.reasoning_format().is_some();
1873 let gears = model.chat_template().think_gears(parser_configured);
1874 if !gears.is_empty() {
1875 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1876 if let Some(default) = &gears.default {
1877 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1878 }
1879 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1882 }
1883 }
1884 if let Some(mcp) = &state.mcp {
1885 model_entry["frink_mcp"] = mcp.models_metadata();
1886 }
1887 Json(serde_json::json!({
1888 "object": "list",
1889 "data": [model_entry]
1890 }))
1891}
1892
1893async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1907 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1908 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1909 let active = state.active();
1910 Json(serde_json::json!({
1911 "model": active.as_ref().map(|a| a.name()),
1912 "state": state
1913 .maintenance
1914 .lock()
1915 .unwrap_or_else(|p| p.into_inner())
1916 .state()
1917 .as_str(),
1918 "uptime_s": state.uptime().as_secs(),
1919 "throughput": {
1920 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1921 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1922 },
1923 "requests": {
1924 "active": state.cancels.live_count(),
1925 "completed": state.stats.recorded_total(),
1926 "p95_ms": state.stats.p95_duration_ms(),
1927 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1928 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1929 "completion_tokens_total": state.stats.tokens_generated_total(),
1930 },
1931 "pools": cache_admin::pool_gauges(&state),
1937 "memory": cache_admin::footprint_json(&state),
1940 }))
1941}
1942
1943#[derive(Deserialize)]
1944struct RequestsQuery {
1945 #[serde(default)]
1946 since: u64,
1947 #[serde(default = "default_requests_limit")]
1948 limit: usize,
1949}
1950
1951fn default_requests_limit() -> usize {
1952 stats::MAX_PAGE
1953}
1954
1955async fn recent_requests(
1963 State(state): State<Arc<AppState>>,
1964 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1965) -> Json<serde_json::Value> {
1966 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1967 Json(serde_json::json!({
1968 "requests": rows,
1969 "next_cursor": cursor,
1970 "missed": missed,
1971 "total": state.stats.recorded_total(),
1972 }))
1973}
1974
1975#[derive(Serialize)]
1976struct CombinedCacheStats {
1977 response_cache: response_cache::CacheStats,
1978 prefix_cache: Option<frink_models::PrefixCacheStats>,
1980}
1981
1982async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1983 Json(CombinedCacheStats {
1984 response_cache: lock_cache(&state.response_cache).stats(),
1985 prefix_cache: state
1986 .prefix_cache
1987 .as_ref()
1988 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1989 })
1990}
1991
1992async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1997 use std::sync::atomic::Ordering;
1998
1999 let cache_stats = lock_cache(&state.response_cache).stats();
2000 let active = state.active();
2001 let requests_total = state.requests_total.load(Ordering::Relaxed);
2002 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
2003 let uptime = state.started_at.elapsed().as_secs_f64();
2004
2005 let body = format!(
2006 "# HELP frink_requests_total Total chat completion requests received.\n\
2007 # TYPE frink_requests_total counter\n\
2008 frink_requests_total {requests_total}\n\
2009 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
2010 # TYPE frink_request_errors_total counter\n\
2011 frink_request_errors_total {errors_total}\n\
2012 # HELP frink_cache_hits_total Whole-response cache hits.\n\
2013 # TYPE frink_cache_hits_total counter\n\
2014 frink_cache_hits_total {}\n\
2015 # HELP frink_cache_misses_total Whole-response cache misses.\n\
2016 # TYPE frink_cache_misses_total counter\n\
2017 frink_cache_misses_total {}\n\
2018 # HELP frink_cache_entries Current whole-response cache entry count.\n\
2019 # TYPE frink_cache_entries gauge\n\
2020 frink_cache_entries {}\n\
2021 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
2022 # TYPE frink_synthetic_weights gauge\n\
2023 frink_synthetic_weights {}\n\
2024 # HELP frink_uptime_seconds Seconds since this server process started.\n\
2025 # TYPE frink_uptime_seconds gauge\n\
2026 frink_uptime_seconds {uptime}\n",
2027 cache_stats.hits,
2028 cache_stats.misses,
2029 cache_stats.entries,
2030 active
2034 .as_ref()
2035 .map(|a| a.is_synthetic() as u8)
2036 .unwrap_or(0),
2037 );
2038
2039 let body = match active
2043 .as_ref()
2044 .and_then(|a| a.expert_store_stats())
2045 {
2046 Some(es) => format!(
2047 "{body}\
2048 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
2049 # TYPE frink_expert_cache_hits_total counter\n\
2050 frink_expert_cache_hits_total {}\n\
2051 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2052 # TYPE frink_expert_cache_misses_total counter\n\
2053 frink_expert_cache_misses_total {}\n\
2054 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
2055 # TYPE frink_expert_cache_evictions_total counter\n\
2056 frink_expert_cache_evictions_total {}\n\
2057 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2058 # TYPE frink_expert_cache_pass_throughs_total counter\n\
2059 frink_expert_cache_pass_throughs_total {}\n\
2060 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2061 # TYPE frink_expert_cache_bytes_read_total counter\n\
2062 frink_expert_cache_bytes_read_total {}\n\
2063 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2064 # TYPE frink_expert_cache_resident_bytes gauge\n\
2065 frink_expert_cache_resident_bytes {}\n",
2066 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2067 ),
2068 None => body,
2069 };
2070
2071 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2076 Some(batcher) => {
2077 let sched = batcher.stats();
2078 format!(
2079 "{body}\
2080 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2081 # TYPE frink_prefill_chunks_total counter\n\
2082 frink_prefill_chunks_total {}\n\
2083 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2084 # TYPE frink_prefill_tokens_total counter\n\
2085 frink_prefill_tokens_total {}\n\
2086 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2087 # TYPE frink_decode_steps_total counter\n\
2088 frink_decode_steps_total {}\n\
2089 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2090 # TYPE frink_scheduler_queue_depth gauge\n\
2091 frink_scheduler_queue_depth {}\n\
2092 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2093 # TYPE frink_scheduler_queue_rejected_total counter\n\
2094 frink_scheduler_queue_rejected_total {}\n\
2095 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2096 # TYPE frink_kv_blocks_total gauge\n\
2097 frink_kv_blocks_total {}\n\
2098 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2099 # TYPE frink_kv_blocks_free gauge\n\
2100 frink_kv_blocks_free {}\n\
2101 # HELP frink_kv_block_size Token positions per KV block.\n\
2102 # TYPE frink_kv_block_size gauge\n\
2103 frink_kv_block_size {}\n\
2104 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2105 # TYPE frink_kv_rejected_too_large_total counter\n\
2106 frink_kv_rejected_too_large_total {}\n\
2107 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2108 # TYPE frink_kv_rejected_context_length_total counter\n\
2109 frink_kv_rejected_context_length_total {}\n\
2110 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2111 # TYPE frink_scheduler_aborted_total counter\n\
2112 frink_scheduler_aborted_total {}\n\
2113 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2114 # TYPE frink_scheduler_max_seqs gauge\n\
2115 frink_scheduler_max_seqs {}\n\
2116 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2117 # TYPE frink_scheduler_prefill_chunk gauge\n\
2118 frink_scheduler_prefill_chunk {}\n",
2119 sched.prefill_chunks,
2120 sched.prefill_tokens,
2121 sched.decode_steps,
2122 sched.queue_depth,
2123 sched.queue_rejected,
2124 sched.kv_blocks_total,
2125 sched.kv_blocks_free,
2126 sched.kv_block_size,
2127 sched.kv_rejected_too_large,
2128 sched.kv_rejected_context_length,
2129 sched.aborted,
2130 sched.max_seqs,
2131 sched.prefill_chunk,
2132 )
2133 }
2134 None => body,
2135 };
2136
2137 (
2138 [(
2139 axum::http::header::CONTENT_TYPE,
2140 "text/plain; version=0.0.4",
2141 )],
2142 body,
2143 )
2144 .into_response()
2145}
2146
2147pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2148
2149pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2155 (
2156 StatusCode::BAD_REQUEST,
2157 Json(serde_json::json!({"error": {
2158 "message": message,
2159 "type": "invalid_request_error",
2160 "param": param,
2161 "code": null,
2162 }})),
2163 )
2164}
2165
2166pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2167 (
2168 StatusCode::NOT_IMPLEMENTED,
2169 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2170 )
2171}
2172
2173pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2174 let status = match e {
2175 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2176 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2180 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2185 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2191 StatusCode::SERVICE_UNAVAILABLE
2192 }
2193 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2197 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2201 };
2202 tracing::warn!("decode error: {e}");
2203 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2204 if let generate::DecodeError::KvBudgetExceeded {
2210 binding,
2211 estimated_bytes,
2212 limit_bytes,
2213 positions,
2214 positions_limit,
2215 ..
2216 } = &e
2217 {
2218 body["error"]["type"] = serde_json::json!("invalid_request_error");
2219 body["error"]["code"] = serde_json::json!(binding);
2220 body["error"]["binding"] = serde_json::json!(binding);
2221 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2222 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2223 body["error"]["positions"] = serde_json::json!(positions);
2224 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2225 }
2226 if let Some(secs) = e.retry_after_secs() {
2230 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2231 }
2232 (status, Json(body))
2233}
2234
2235pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2236 tracing::error!("generation task panicked: {e}");
2237 (
2238 StatusCode::INTERNAL_SERVER_ERROR,
2239 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2240 )
2241}
2242
2243#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2253 model: &Model,
2254 prompt: &str,
2255 params: &GenerationParams,
2256 kv_pool: Option<&generate::KvPoolConfig>,
2257 paged_kv: Option<&generate::PagedKvConfig>,
2258 prefix_cache: Option<&Mutex<PrefixCache>>,
2259 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2260 ceiling: Option<&budget::ContextCeiling>,
2261 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2262 mut emit: impl FnMut(usize, &str),
2267) -> Result<generate::Generated, generate::DecodeError> {
2268 let synthetic = model.is_synthetic();
2269 let _lora_lease = lora::lease(model, params.lora.as_deref());
2274 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2275 let params = &{
2281 let mut resolved = params.clone();
2282 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2283 model.encode(text, SpecialTokens::Parse)
2284 });
2285 resolved
2290 .token_mask
2291 .resolve(|text| model.encode(text, SpecialTokens::Parse));
2292 resolved.reasoning_budget = resolved
2298 .reasoning_budget
2299 .armed(resolved.reasoning, prompt, |text| {
2300 model.encode(text, SpecialTokens::Parse)
2301 })
2302 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2303 resolved
2304 };
2305 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2306 let _metal_private_guard =
2307 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2308 let (finishes, prompt_rows, prompt_ids, truncated_prompt, usage) = match model {
2309 Model::Gguf(m) => {
2310 if let Some(batcher) = continuous_batcher {
2311 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2312 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2313 let (finish, _generated_ids, text, usage) = if synthetic {
2314 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2315 } else {
2316 batcher.generate_streaming(
2317 tokens,
2318 params.clone(),
2319 m.stop_tokens.clone(),
2320 Some(|chunk: &str| {
2321 if !chunk.is_empty() {
2322 chunks[0].push(chunk.to_string());
2323 emit(0, chunk);
2324 }
2325 }),
2326 )?
2327 };
2328 if !text.is_empty() && chunks[0].is_empty() {
2329 chunks[0].push(text);
2330 }
2331 (
2344 vec![(finish, Vec::new())],
2345 Vec::new(),
2346 Vec::new(),
2347 None,
2348 usage,
2349 )
2350 } else {
2351 generate::generate(
2352 &m.decoder,
2353 m.tokenizer.as_ref(),
2354 &m.stop_tokens,
2355 m.bos_id,
2356 prompt,
2357 params,
2358 kv_pool,
2359 paged_kv,
2360 prefix_cache,
2361 ceiling,
2362 |choice, chunk| {
2363 chunks[choice].push(chunk.to_string());
2364 if !synthetic {
2368 emit(choice, chunk);
2369 }
2370 },
2371 )?
2372 }
2373 }
2374 Model::Kimi(m) => generate::generate_engine(
2375 &m.engine,
2376 &m.tokenizer,
2377 &m.stop_tokens,
2378 None,
2379 prompt,
2380 params,
2381 |chunk| {
2382 chunks[0].push(chunk.to_string());
2383 if !synthetic {
2384 emit(0, chunk);
2385 }
2386 },
2387 )?,
2388 Model::Mla(m) => generate::generate_engine(
2389 &m.engine,
2390 &m.tokenizer,
2391 &m.stop_tokens,
2392 m.bos_id,
2393 prompt,
2394 params,
2395 |chunk| {
2396 chunks[0].push(chunk.to_string());
2397 if !synthetic {
2398 emit(0, chunk);
2399 }
2400 },
2401 )?,
2402 Model::Gemma4(m) => generate::generate_engine(
2403 &m.engine,
2404 &m.tokenizer,
2405 &m.stop_tokens,
2406 m.bos_id,
2407 prompt,
2408 params,
2409 |chunk| {
2410 chunks[0].push(chunk.to_string());
2411 if !synthetic {
2412 emit(0, chunk);
2413 }
2414 },
2415 )?,
2416 Model::Glm52(m) => generate::generate_engine(
2417 &m.engine,
2418 &m.tokenizer,
2419 &m.stop_tokens,
2420 m.bos_id,
2421 prompt,
2422 params,
2423 |chunk| {
2424 chunks[0].push(chunk.to_string());
2425 if !synthetic {
2426 emit(0, chunk);
2427 }
2428 },
2429 )?,
2430 };
2431
2432 let mut full = chunks[0].concat();
2433 if synthetic {
2434 full = format!(
2435 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2436 to serve a real model. Decoded ids -> {full:?}]"
2437 );
2438 emit(0, &full);
2439 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2440 emit(0, &full);
2441 }
2442
2443 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2449 let mut out: Vec<generate::GeneratedChoice> = finishes
2450 .into_iter()
2451 .zip(chunks.into_iter().map(|c| c.concat()))
2452 .map(|((finish, logprobs), text)| generate::GeneratedChoice {
2453 finish,
2454 text,
2455 logprobs,
2456 })
2457 .collect();
2458 if let Some(first) = out.first_mut() {
2459 if synthetic {
2467 first.logprobs.clear();
2468 }
2469 first.text = full;
2470 }
2471 Ok(generate::Generated {
2472 choices: out,
2473 prompt_rows,
2474 prompt_ids,
2475 truncated_prompt,
2476 usage,
2477 })
2478}
2479
2480#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2485 model: &Model,
2486 prompt: &str,
2487 params: &GenerationParams,
2488 kv_pool: Option<&generate::KvPoolConfig>,
2489 paged_kv: Option<&generate::PagedKvConfig>,
2490 prefix_cache: Option<&Mutex<PrefixCache>>,
2491 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2492 ceiling: Option<&budget::ContextCeiling>,
2493 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2494 ) -> Result<generate::Generated, generate::DecodeError> {
2497 run_generation_emit(
2498 model,
2499 prompt,
2500 params,
2501 kv_pool,
2502 paged_kv,
2503 prefix_cache,
2504 continuous_batcher,
2505 ceiling,
2506 metal_private_decode_gate,
2507 |_, _| {},
2508 )
2509}
2510
2511pub(crate) fn prompt_from_messages(
2522 messages: &[ChatMessage],
2523 template: &chat_template::PromptTemplate,
2524 tools: &[ToolDef],
2525 extra: serde_json::Map<String, serde_json::Value>,
2526) -> Result<String, ApiError> {
2527 let rendered = if tools.is_empty() || template.handles_tools() {
2528 template.render(messages, tools, extra)
2529 } else {
2530 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2531 with_preamble.push(ChatMessage {
2532 role: "system".to_string(),
2533 content: Some(MessageContent::Text(tool_preamble(tools))),
2534 tool_calls: None,
2535 tool_call_id: None,
2536 reasoning_content: None,
2537 });
2538 with_preamble.extend_from_slice(messages);
2539 template.render(&with_preamble, &[], extra)
2540 };
2541 rendered.map_err(template_error_response)
2542}
2543
2544fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2549 (
2550 StatusCode::BAD_REQUEST,
2551 Json(serde_json::json!({
2552 "error": {
2553 "message": format!("chat template failed to render: {err}"),
2554 "type": "invalid_request_error",
2555 "param": "messages",
2556 "code": null,
2557 }
2558 })),
2559 )
2560}
2561
2562fn tool_preamble(tools: &[ToolDef]) -> String {
2575 let mut out = String::from(
2576 "You can call tools to help answer the user. To call a tool, respond with \
2577 EXACTLY one line in this format and nothing else:\n\
2578 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2579 object matching that tool's parameters>}}</tool_call>\n\n\
2580 Available tools:\n",
2581 );
2582 for t in tools {
2583 out.push_str(&format!(
2584 "- {}: {}\n parameters (JSON schema): {}\n",
2585 t.function.name,
2586 t.function.description.as_deref().unwrap_or(""),
2587 t.function
2588 .parameters
2589 .as_ref()
2590 .map(|v| v.to_string())
2591 .unwrap_or_else(|| "{}".to_string()),
2592 ));
2593 }
2594 out
2595}
2596
2597fn tool_call_deltas(
2607 events: Vec<crate::policy::parser::ToolCallEvent>,
2608 opened: &std::cell::Cell<usize>,
2609) -> (String, Vec<ToolCallDelta>) {
2610 let mut text = String::new();
2611 let mut deltas = Vec::new();
2612 for event in events {
2613 match event {
2614 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2615 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2616 opened.set(opened.get().max(index + 1));
2617 deltas.push(ToolCallDelta::opening(index, name));
2618 }
2619 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2620 if !fragment.is_empty() {
2621 deltas.push(ToolCallDelta::arguments(index, fragment));
2622 }
2623 }
2624 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2625 }
2626 }
2627 (text, deltas)
2628}
2629
2630fn build_response_message(
2641 text: String,
2642 tools: &[ToolDef],
2643 posture: output::OutputPosture,
2644 base_finish: &'static str,
2645) -> (ChatCompletionResponseMessage, &'static str) {
2646 let parsed = output::parse_output(&text, tools, posture);
2647 let calls: Vec<ToolCallOut> = parsed
2648 .calls
2649 .into_iter()
2650 .enumerate()
2651 .map(|(index, call)| ToolCallOut {
2652 id: format!("call_{index}"),
2653 kind: "function",
2654 function: ToolCallFunctionOut {
2655 name: call.name,
2656 arguments: call.arguments,
2657 },
2658 })
2659 .collect();
2660 if !calls.is_empty() {
2661 return (
2662 ChatCompletionResponseMessage {
2663 role: "assistant",
2664 content: None,
2665 reasoning_content: parsed.reasoning,
2666 tool_calls: Some(calls),
2667 },
2668 "tool_calls",
2669 );
2670 }
2671 (
2672 ChatCompletionResponseMessage {
2673 role: "assistant",
2674 content: Some(parsed.content),
2675 reasoning_content: parsed.reasoning,
2676 tool_calls: None,
2677 },
2678 base_finish,
2679 )
2680}
2681
2682fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2687 let mut history = match &req.session_id {
2688 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2689 None => req.messages.clone(),
2690 };
2691 if req.json_object_mode() {
2692 inject_json_object_system_hint(&mut history);
2693 }
2694 history
2695}
2696
2697fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2698 const HINT: &str =
2699 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2700 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2701 match &mut sys.content {
2702 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2703 s.push_str("\n\n");
2704 s.push_str(HINT);
2705 }
2706 None => {
2707 sys.content = Some(MessageContent::Text(HINT.to_string()));
2708 }
2709 _ => {}
2710 }
2711 } else {
2712 messages.insert(
2713 0,
2714 ChatMessage {
2715 role: "system".to_string(),
2716 content: Some(MessageContent::Text(HINT.to_string())),
2717 tool_calls: None,
2718 tool_call_id: None,
2719 reasoning_content: None,
2720 },
2721 );
2722 }
2723}
2724
2725async fn chat_completions(
2726 State(state): State<Arc<AppState>>,
2727 headers: axum::http::HeaderMap,
2728 Json(req): Json<ChatCompletionRequest>,
2729) -> Response {
2730 let attribution = attribution::Attribution::from_headers(&headers);
2731 state
2732 .requests_total
2733 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2734 let started = std::time::Instant::now();
2735
2736 let request_id = frink_api::next_request_id();
2740 let stream = req.stream.unwrap_or(false);
2741
2742 let refusal = cache_admin::check_admission(&state)
2748 .err()
2749 .or_else(|| req.validate_supported_fields().err());
2750 if let Some(err) = refusal {
2751 state
2752 .request_errors_total
2753 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2754 let response = err.into_response();
2755 state.record_request(stats::Record {
2756 request_id: &request_id,
2757 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2758 model: state.active_model_name(),
2759 status: response.status().as_u16(),
2760 stream,
2761 duration_ms: started.elapsed().as_millis() as u64,
2762 usage: None,
2763 attribution: &attribution,
2764 });
2765 return response;
2766 }
2767
2768 let response = if stream {
2769 chat_completions_stream(
2770 Arc::clone(&state),
2771 req,
2772 request_id.clone(),
2773 started,
2774 attribution.clone(),
2775 )
2776 .await
2777 .into_response()
2778 } else {
2779 chat_completions_full(
2780 Arc::clone(&state),
2781 req,
2782 request_id.clone(),
2783 started,
2784 attribution.clone(),
2785 )
2786 .await
2787 .into_response()
2788 };
2789
2790 if response.status().is_client_error() || response.status().is_server_error() {
2791 state
2792 .request_errors_total
2793 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2794 state.record_request(stats::Record {
2798 request_id: &request_id,
2799 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2800 model: state.active_model_name(),
2803 status: response.status().as_u16(),
2804 stream,
2805 duration_ms: started.elapsed().as_millis() as u64,
2806 usage: None,
2807 attribution: &attribution,
2808 });
2809 }
2810 state.mark_request_finished();
2811
2812 response
2813}
2814
2815async fn chat_completions_full(
2816 state: Arc<AppState>,
2817 req: ChatCompletionRequest,
2818 request_id: String,
2819 started: std::time::Instant,
2820 attribution: attribution::Attribution,
2821) -> Result<Json<ChatCompletionResponse>, ApiError> {
2822 let tools_active = req.tools_active();
2823 let active = state.require_active()?;
2827 let history = resolve_history(&state, &req);
2828 let template = active.generative()?.chat_template();
2829 let kwargs = req.resolve_template_kwargs(&template);
2830 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2831 let mut params =
2839 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2840 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2841 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2842
2843 let mut generated_logprobs: Vec<crate::sampling_loop::PerTokenProbs> = Vec::new();
2847 let n_logprobs = req.n_logprobs()?;
2850 let decode_any = |id: usize| active.decode_any(&[id]);
2852 let (completion, cache_status) = if let Some(cached) = key
2853 .as_ref()
2854 .and_then(|key| lock_cache(&state.response_cache).get(key))
2855 {
2856 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2857 (cached, "hit")
2858 } else {
2859 let produced = decode_task::buffered(
2860 decode_task::DecodeHandles::take(&state, &active)?,
2861 prompt.clone(),
2862 params,
2863 )
2864 .await?;
2865 let usage = produced.usage;
2866 let choices = produced.choices;
2867
2868 generated_logprobs = choices.iter().map(|c| c.logprobs.clone()).collect();
2873 let completion = response_cache::CachedCompletion {
2874 choices: choices.into_iter().map(|c| (c.finish, c.text)).collect(),
2875 usage,
2876 };
2877 let cache_status = match key {
2883 Some(key) => match completion.clone().cacheable() {
2887 Some(cacheable) => {
2888 tracing::debug!("cache miss for key {}", key.digest());
2889 lock_cache(&state.response_cache).put(key, cacheable);
2890 "miss"
2891 }
2892 None => "skip",
2893 },
2894 None => "skip",
2895 };
2896 (completion, cache_status)
2897 };
2898 let content = completion.first_text().to_string();
2901
2902 if req.json_object_mode() {
2903 json_mode::validate_json_object_output(&content)?;
2904 }
2905
2906 if let Some(id) = &req.session_id {
2910 state.sessions.store_reply(
2911 id,
2912 ChatMessage {
2913 role: "assistant".to_string(),
2914 content: Some(MessageContent::Text(content.clone())),
2915 tool_calls: None,
2916 tool_call_id: None,
2917 reasoning_content: None,
2918 },
2919 );
2920 }
2921
2922 let posture = output::OutputPosture::resolve_full(
2927 active.reasoning_format(),
2928 active.tool_call_format(),
2929 &prompt,
2930 );
2931 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2932 let wanted = req.unimplemented.n.unwrap_or(1).max(1) as usize;
2936 let ranked: Vec<(generate::FinishReason, String)> = if completion.choices.len() > wanted {
2937 let scored: Vec<crate::generate::GeneratedChoice> = completion
2938 .choices
2939 .into_iter()
2940 .zip(
2941 generated_logprobs
2942 .iter()
2943 .cloned()
2944 .chain(std::iter::repeat(Vec::new())),
2945 )
2946 .map(
2947 |((finish, text), logprobs)| crate::generate::GeneratedChoice {
2948 finish,
2949 text,
2950 logprobs,
2951 },
2952 )
2953 .collect();
2954 let best = crate::best_of::take_best(scored, wanted);
2955 generated_logprobs = best.iter().map(|c| c.logprobs.clone()).collect();
2956 best.into_iter().map(|c| (c.finish, c.text)).collect()
2957 } else {
2958 completion.choices
2959 };
2960 let render_piece =
2966 crate::logprobs::piece_renderer(req.unimplemented.tokens_as_ids(), &decode_any);
2967 let rendered: Vec<ChatCompletionChoice> = ranked
2968 .into_iter()
2969 .enumerate()
2970 .map(|(index, (finish, text))| {
2971 let (message, finish_reason) =
2972 build_response_message(text, tools, posture, finish.as_str());
2973 ChatCompletionChoice {
2974 index,
2975 message,
2976 finish_reason,
2977 logprobs: n_logprobs.map(|k| {
2978 crate::logprobs::render_chat(
2979 generated_logprobs.get(index).unwrap_or(&Vec::new()),
2980 Some(k),
2981 render_piece.as_ref(),
2982 )
2983 }),
2984 }
2985 })
2986 .collect();
2987
2988 state.record_request(stats::Record {
2989 request_id: &request_id,
2990 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2991 model: Some(active.name().to_string()),
2994 status: 200,
2995 stream: false,
2996 duration_ms: started.elapsed().as_millis() as u64,
2997 usage: Some(&completion.usage),
2998 attribution: &attribution,
2999 });
3000
3001 Ok(Json(ChatCompletionResponse {
3002 id: request_id.clone(),
3003 request_id,
3004 object: "chat.completion",
3005 model: req.model,
3006 choices: rendered,
3007 usage: completion.usage,
3008 frink_cache: cache_status,
3009 }))
3010}
3011
3012async fn chat_completions_stream(
3013 state: Arc<AppState>,
3014 req: ChatCompletionRequest,
3015 request_id: String,
3016 started: std::time::Instant,
3017 attribution: attribution::Attribution,
3018) -> Result<Response, ApiError> {
3019 let tools_active = req.tools_active();
3029 let active = state.require_active()?;
3033 let history = resolve_history(&state, &req);
3034 let template = active.generative()?.chat_template();
3035 let kwargs = req.resolve_template_kwargs(&template);
3036 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
3037 let model_name = req.model.clone();
3038 let session_id = req.session_id.clone();
3039 let sessions = state.sessions.clone();
3040
3041 let model = Arc::clone(active.generative()?);
3042 let kv_pool = state.kv_pool.clone();
3043 let paged_kv = state.paged_kv.clone();
3044 let prefix_cache = state.prefix_cache.clone();
3045 let batcher = active.batcher.clone();
3046 let ceiling = active.ceiling.clone();
3047 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3048 let mut params =
3049 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
3050 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
3051 params.interleave_choices = params.n > 1;
3057 let stats_state = Arc::clone(&state);
3058 let served_model = active.name().to_string();
3061 let posture = output::OutputPosture::resolve_full(
3065 active.reasoning_format(),
3066 active.tool_call_format(),
3067 &prompt,
3068 );
3069 let offered_tools: Vec<ToolDef> = if tools_active {
3072 req.tools.clone()
3073 } else {
3074 Vec::new()
3075 };
3076
3077 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3082 params.cancel = Some(cancel_token.clone());
3083
3084 let overlap = true;
3093
3094 let slot = req
3098 .stream_resumable
3099 .unwrap_or(false)
3100 .then(|| state.streams.register(&request_id));
3101 let emitter = resume::Emitter::new(slot);
3102
3103 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3104 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3108 id: request_id.clone(),
3109 request_id: None,
3110 object: "chat.completion.chunk",
3111 model: model_name.clone(),
3112 choices: vec![ChatCompletionChunkChoice {
3113 index: 0,
3114 delta: ChatCompletionChunkDelta {
3115 role: None,
3116 content: None,
3117 reasoning_content: None,
3118 tool_calls: None,
3119 },
3120 finish_reason: None,
3121 }],
3122 usage: None,
3123 });
3124
3125 tokio::task::spawn_blocking(move || {
3126 let _cancel_guard = cancel_guard;
3129 let tx_chunks = tx.clone();
3130 let orphan_timeout = sse::orphan_timeout_from_env();
3135 let head_request_id = request_id.clone();
3136 let announced = std::cell::Cell::new(false);
3140 let emitters: Rc<RefCell<Vec<crate::chat_stream_choice::ChoiceEmitter>>> =
3146 Rc::new(RefCell::new(
3147 (0..params.n.max(1))
3148 .map(|_| {
3149 crate::chat_stream_choice::ChoiceEmitter::new(
3150 posture.reasoning_parser(),
3151 tools_active.then(|| posture.tool_call_parser(&offered_tools)),
3152 )
3153 })
3154 .collect(),
3155 ));
3156 let emit_choices = Rc::clone(&emitters);
3157 let result = run_generation_emit(
3158 &model,
3159 &prompt,
3160 ¶ms,
3161 kv_pool.as_ref(),
3162 paged_kv.as_ref(),
3163 prefix_cache.as_deref(),
3164 batcher.as_ref(),
3165 ceiling.as_deref(),
3166 metal_private_decode_gate.as_deref(),
3167 |choice, chunk| {
3168 if !overlap || chunk.is_empty() {
3169 return;
3170 }
3171 let mut held = emit_choices.borrow_mut();
3172 let Some(emitter_state) = held.get_mut(choice) else {
3173 return;
3174 };
3175 let delta = emitter_state.push(chunk);
3176 if delta.is_empty() {
3177 return;
3178 }
3179 let request_id = (!announced.get()).then(|| {
3182 announced.set(true);
3183 head_request_id.clone()
3184 });
3185 let wire = delta.into_choice(choice, emitter_state.start());
3186 drop(held);
3187 let payload = ChatCompletionChunk {
3188 id: head_request_id.clone(),
3189 request_id,
3190 object: "chat.completion.chunk",
3191 model: model_name.clone(),
3192 choices: vec![wire],
3193 usage: None,
3194 };
3195 if let Err(why) =
3204 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3205 {
3206 if why == sse::SendFailure::Orphaned {
3207 tracing::warn!(
3208 "SSE stream {head_request_id} accepted nothing for the orphan \
3209 deadline; treating it as abandoned"
3210 );
3211 }
3212 if !emitter.is_resumable() {
3224 cancel_token.cancel();
3225 }
3226 }
3227 },
3228 );
3229
3230 let mut pending_request_id = (!announced.get()).then(|| request_id.clone());
3236
3237 match result {
3238 Ok(generated) => {
3239 let usage = generated.usage;
3240 let produced: Vec<(generate::FinishReason, String)> = generated
3241 .choices
3242 .into_iter()
3243 .map(|c| (c.finish, c.text))
3244 .collect();
3245 assert!(
3246 !produced.is_empty(),
3247 "a generation produces at least one choice"
3248 );
3249 if let Some(id) = &session_id {
3254 sessions.store_reply(
3255 id,
3256 ChatMessage {
3257 role: "assistant".to_string(),
3258 content: Some(MessageContent::Text(produced[0].1.clone())),
3259 tool_calls: None,
3260 tool_call_id: None,
3261 reasoning_content: None,
3262 },
3263 );
3264 }
3265 for (index, (finish, full_text)) in produced.iter().enumerate() {
3266 let (finish, full_text) = (finish.clone(), full_text.as_str());
3267 let mut streamed_finish: Option<&'static str> = None;
3273 if overlap {
3274 let (tail, first, opened) = {
3275 let mut held = emitters.borrow_mut();
3276 let state = &mut held[index];
3277 let tail = state.flush();
3278 (tail, state.start(), state.opened_calls())
3279 };
3280 if !tail.is_empty() {
3281 let payload = ChatCompletionChunk {
3282 id: request_id.clone(),
3283 request_id: pending_request_id.take(),
3284 object: "chat.completion.chunk",
3285 model: model_name.clone(),
3286 choices: vec![tail.into_choice(index, first)],
3287 usage: None,
3288 };
3289 let _ = sse::send_or_orphan(
3290 &tx,
3291 Ok(emitter.event(&payload)),
3292 orphan_timeout,
3293 );
3294 }
3295 if opened > 0 {
3296 streamed_finish = Some("tool_calls");
3297 }
3298 } else {
3299 let parsed = output::parse_output(full_text, &offered_tools, posture);
3302 let tool_calls: Vec<ToolCallDelta> = parsed
3303 .calls
3304 .iter()
3305 .enumerate()
3306 .map(|(index, call)| {
3307 ToolCallDelta::whole(
3308 index,
3309 call.name.clone(),
3310 call.arguments.clone(),
3311 )
3312 })
3313 .collect();
3314 if !tool_calls.is_empty() {
3315 streamed_finish = Some("tool_calls");
3316 }
3317 if !tool_calls.is_empty()
3318 || !parsed.content.is_empty()
3319 || parsed.reasoning.is_some()
3320 {
3321 let payload = ChatCompletionChunk {
3322 id: request_id.clone(),
3323 request_id: pending_request_id.take(),
3324 object: "chat.completion.chunk",
3325 model: model_name.clone(),
3326 choices: vec![ChatCompletionChunkChoice {
3327 index,
3328 delta: ChatCompletionChunkDelta {
3329 role: Some("assistant"),
3330 content: (!parsed.content.is_empty()
3331 && tool_calls.is_empty())
3332 .then(|| parsed.content.clone()),
3333 reasoning_content: parsed.reasoning.clone(),
3334 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3335 },
3336 finish_reason: None,
3337 }],
3338 usage: None,
3339 };
3340 let _ = sse::send_or_orphan(
3341 &tx,
3342 Ok(emitter.event(&payload)),
3343 orphan_timeout,
3344 );
3345 }
3346 }
3347 let final_finish_reason = match streamed_finish {
3351 Some(reason) if finish.as_str() != "length" => reason,
3352 _ => finish.as_str(),
3353 };
3354 let last = index + 1 == produced.len();
3358 let final_payload = ChatCompletionChunk {
3359 id: request_id.clone(),
3360 request_id: pending_request_id.take(),
3361 object: "chat.completion.chunk",
3362 model: model_name.clone(),
3363 choices: vec![ChatCompletionChunkChoice {
3364 index,
3365 delta: ChatCompletionChunkDelta {
3366 role: None,
3367 content: None,
3368 reasoning_content: None,
3369 tool_calls: None,
3370 },
3371 finish_reason: Some(final_finish_reason),
3372 }],
3373 usage: last.then(|| usage.clone()),
3374 };
3375 let _ =
3376 sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3377 }
3378 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3379 stats_state.record_request(stats::Record {
3384 request_id: &request_id,
3385 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3386 model: Some(served_model.clone()),
3387 status: 200,
3388 stream: true,
3389 duration_ms: started.elapsed().as_millis() as u64,
3390 usage: Some(&usage),
3391 attribution: &attribution,
3392 });
3393 }
3394 Err(e) => {
3395 tracing::warn!("decode error on streamed request {request_id}: {e}");
3396 stats_state.record_request(stats::Record {
3402 request_id: &request_id,
3403 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3404 model: Some(served_model.clone()),
3405 status: 500,
3406 stream: true,
3407 duration_ms: started.elapsed().as_millis() as u64,
3408 usage: None,
3409 attribution: &attribution,
3410 });
3411 let payload = ChatCompletionChunk {
3412 id: request_id.clone(),
3413 request_id: pending_request_id.take(),
3414 object: "chat.completion.chunk",
3415 model: model_name,
3416 choices: vec![ChatCompletionChunkChoice {
3417 index: 0,
3418 delta: ChatCompletionChunkDelta {
3419 role: Some("assistant"),
3420 content: Some(format!("[error: {e}]")),
3421 reasoning_content: None,
3422 tool_calls: None,
3423 },
3424 finish_reason: Some("stop"),
3425 }],
3426 usage: None,
3427 };
3428 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3429 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3430 }
3431 }
3432 drop(emitter);
3436 });
3437
3438 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3439 Ok((
3462 [(
3463 axum::http::HeaderName::from_static("x-accel-buffering"),
3464 axum::http::HeaderValue::from_static("no"),
3465 )],
3466 Sse::new(stream),
3467 )
3468 .into_response())
3469}
3470
3471fn protected_routes() -> Router<Arc<AppState>> {
3502 use frink_api::routes;
3503
3504 Router::new()
3505 .route(routes::V1_MODELS, get(list_models))
3506 .route(routes::V1_RESPONSES, post(responses::responses))
3510 .route(
3511 &axum_path(routes::V1_RESPONSE),
3512 get(responses::responses_get),
3513 )
3514 .route(
3515 &axum_path(routes::V1_RESPONSE_CANCEL),
3516 post(responses::responses_cancel),
3517 )
3518 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3519 .route(routes::V1_STATS, get(serving_stats))
3520 .route(routes::V1_REQUESTS, get(recent_requests))
3521 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3522 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3523 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3524 .route(
3525 routes::LORA_ADAPTERS,
3526 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3527 )
3528 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3529 .route(routes::V1_CANCEL, post(cancel_generation))
3533 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3538 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3539 .route(routes::V1_MESSAGES, post(anthropic::messages))
3540 .route(
3541 routes::V1_MESSAGES_COUNT_TOKENS,
3542 post(anthropic::count_tokens),
3543 )
3544 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3545 .route(routes::COMPLETION, post(completion::completion))
3550 .route(routes::COMPLETIONS, post(completion::completion))
3551 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3552 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3553 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3559 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3560 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3561 .route(routes::V1_RERANK, post(rerank::rerank))
3565 .route(routes::RERANK, post(rerank::rerank))
3566 .route(routes::CACHE_STATS, get(cache_stats))
3567 .route(routes::METRICS, get(metrics))
3568 .route(routes::ADMIN_MODELS, get(admin::models))
3573 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3574 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3575 .route(routes::SLEEP, post(admin::sleep))
3579 .route(routes::WAKE_UP, post(admin::wake_up))
3580 .route(routes::IS_SLEEPING, get(admin::is_sleeping))
3581 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3582 .route(routes::ADMIN_TASKS, get(admin::tasks))
3583 .route(&admin::cancel_route(), post(admin::cancel_task))
3584 .route(routes::ADMIN_STATS, get(admin::stats))
3585 .merge(conversations::router())
3589}
3590
3591fn axum_path(template: &str) -> String {
3592 let mut out = String::with_capacity(template.len());
3593 let mut rest = template;
3594 while let Some(open) = rest.find('{') {
3595 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3596 break;
3597 };
3598 out.push_str(&rest[..open]);
3599 out.push(':');
3600 out.push_str(&rest[open + 1..close]);
3601 rest = &rest[close + 1..];
3602 }
3603 out.push_str(rest);
3604 out
3605}
3606
3607async fn cancel_generation(
3616 State(state): State<Arc<AppState>>,
3617 Json(req): Json<frink_api::CancelGenerationRequest>,
3618) -> Response {
3619 let cancelled = state.cancels.cancel(&req.request_id);
3620 let status = if cancelled {
3621 StatusCode::OK
3622 } else {
3623 StatusCode::NOT_FOUND
3624 };
3625 let detail = if cancelled {
3626 "the generation was asked to stop; it ends at its next token".to_string()
3627 } else {
3628 "no generation with that request_id is running -- it has already \
3629 finished, was never issued, or was served by a path that does \
3630 not register for cancellation"
3631 .to_string()
3632 };
3633 (
3634 status,
3635 Json(frink_api::CancelGenerationResponse {
3636 request_id: req.request_id,
3637 cancelled,
3638 detail,
3639 }),
3640 )
3641 .into_response()
3642}
3643
3644type Activated = (
3648 Loaded,
3649 Option<serving::batch::ContinuousBatcher>,
3650 Option<Arc<budget::ContextCeiling>>,
3651);
3652
3653fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3671 let mut batcher = serving::batch::BatcherConfig::from_env();
3672 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3673 return batcher;
3676 }
3677 let Some(path) = path else {
3678 return batcher;
3679 };
3680 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3686 let Some((priced, gguf_ctx, source)) = priced else {
3687 return batcher;
3688 };
3689 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3690 tracing::warn!(
3694 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3695 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3696 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3697 admit on a number you choose.",
3698 priced.weights_bytes,
3699 priced.device_budget_bytes,
3700 );
3701 return batcher;
3702 };
3703 tracing::info!("{source}");
3704 tracing::info!("{}", derived.fit);
3705 let adopted = budget::apply_derived(&mut batcher, &derived);
3706 if adopted.max_context {
3707 tracing::info!(
3708 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3709 override with FRINK_CB_MAX_CONTEXT",
3710 derived.max_context
3711 );
3712 }
3713 if adopted.kv_blocks {
3714 tracing::info!(
3715 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3716 derived.kv_blocks,
3717 batcher.kv_block_size
3718 );
3719 }
3720 if let Some(narrowed) = adopted.max_context_narrowed {
3721 tracing::info!(
3722 "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",
3723 batcher.kv_blocks.unwrap_or_default(),
3724 batcher.kv_block_size
3725 );
3726 }
3727 batcher
3728}
3729
3730pub(crate) fn activate_loaded_model(
3739 loaded: model::LoadedModel,
3740 enable_continuous_batching: bool,
3741 path: Option<&str>,
3742 paged_kv: Option<&generate::PagedKvConfig>,
3743) -> Activated {
3744 match loaded {
3745 model::LoadedModel::Gguf(g) => {
3746 let decoder = Arc::new(g.decoder);
3747 let tokenizer = Arc::new(g.tokenizer);
3748 let config = price_batcher_config(path);
3749 let shape =
3755 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3756 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3757 let batcher = if enable_continuous_batching {
3758 tracing::info!(
3759 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3760 (stop sequences use the same pending-buffer trim as the private generate loop)"
3761 );
3762 let tok = Arc::clone(&tokenizer);
3763 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3764 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3765 Arc::clone(&decoder),
3766 decode,
3767 config,
3768 Arc::clone(&ceiling),
3769 paged_kv.cloned(),
3770 ))
3771 } else {
3772 None
3773 };
3774 (
3775 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3776 decoder,
3777 tokenizer,
3778 stop_tokens: g.stop_tokens,
3779 bos_id: g.bos_id,
3780 is_synthetic: g.is_synthetic,
3781 chat_template: g.chat_template,
3782 }))),
3783 batcher,
3784 Some(ceiling),
3785 )
3786 }
3787 model::LoadedModel::Kimi(k) => (
3788 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3789 engine: k.engine,
3790 tokenizer: k.tokenizer,
3791 stop_tokens: k.stop_tokens,
3792 chat_template: k.chat_template,
3793 }))),
3794 None,
3795 None,
3796 ),
3797 model::LoadedModel::Mla(m) => (
3798 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3799 engine: m.engine,
3800 tokenizer: m.tokenizer,
3801 stop_tokens: m.stop_tokens,
3802 bos_id: m.bos_id,
3803 name: m.name,
3804 chat_template: m.chat_template,
3805 }))),
3806 None,
3807 None,
3808 ),
3809 model::LoadedModel::Gemma4(m) => (
3810 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3811 engine: m.engine,
3812 tokenizer: m.tokenizer,
3813 stop_tokens: m.stop_tokens,
3814 bos_id: m.bos_id,
3815 name: m.name,
3816 chat_template: m.chat_template,
3817 }))),
3818 None,
3819 None,
3820 ),
3821 model::LoadedModel::Glm52(g) => (
3822 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3823 engine: g.engine,
3824 tokenizer: g.tokenizer,
3825 stop_tokens: g.stop_tokens,
3826 bos_id: g.bos_id,
3827 name: g.name,
3828 chat_template: g.chat_template,
3829 }))),
3830 None,
3831 None,
3832 ),
3833 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3838 }
3839}
3840
3841struct StartupModels {
3848 loaded: model::LoadedModel,
3849 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3850}
3851
3852fn continuous_batching_env() -> Option<bool> {
3853 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3854 .ok()
3855 .map(|v| v.trim().to_ascii_lowercase())
3856 .as_deref()
3857 {
3858 None => None,
3859 Some("1" | "true" | "yes" | "on") => Some(true),
3860 Some("0" | "false" | "no" | "off") => Some(false),
3861 _ => None,
3862 }
3863}
3864
3865fn metal_private_decode_active() -> bool {
3866 #[cfg(feature = "metal")]
3867 {
3868 BUILT_WITH_METAL
3869 && frink_metal::attn::metal_attn_enabled()
3870 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3871 }
3872 #[cfg(not(feature = "metal"))]
3873 {
3874 false
3875 }
3876}
3877
3878fn continuous_batching_compatible(
3879 loaded: &model::LoadedModel,
3880 kv_pool: &Option<generate::KvPoolConfig>,
3881 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3882 paged_kv: &Option<generate::PagedKvConfig>,
3883) -> bool {
3884 matches!(loaded, model::LoadedModel::Gguf(_))
3885 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3886}
3887
3888fn resolve_continuous_batching_enabled(
3889 loaded: &model::LoadedModel,
3890 kv_pool: &Option<generate::KvPoolConfig>,
3891 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3892 paged_kv: &Option<generate::PagedKvConfig>,
3893) -> bool {
3894 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3895 return false;
3896 }
3897 match continuous_batching_env() {
3898 Some(true) => true,
3899 Some(false) => false,
3900 None => metal_private_decode_active(),
3901 }
3902}
3903
3904fn acquire_metal_private_decode_gate(
3905 gate: Option<&std::sync::Mutex<()>>,
3906 used_batcher: bool,
3907) -> Option<std::sync::MutexGuard<'_, ()>> {
3908 if used_batcher {
3909 None
3910 } else {
3911 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3912 }
3913}
3914
3915fn build_app_state(
3916 models: StartupModels,
3917 kv_pool: Option<generate::KvPoolConfig>,
3918 paged_kv: Option<generate::PagedKvConfig>,
3919 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3920 enable_continuous_batching: bool,
3921 mcp: Option<mcp::LoadedMcpConfig>,
3922 detection: Arc<health::Detection>,
3923) -> AppState {
3924 let StartupModels { loaded, embedding } = models;
3925 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3926 let (loaded, batcher, ceiling) = activate_loaded_model(
3927 loaded,
3928 enable_continuous_batching,
3929 configured_path.as_deref(),
3930 paged_kv.as_ref(),
3931 );
3932 let id = startup_model_id();
3938 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3939 {
3940 None
3941 } else {
3942 tracing::info!(
3943 "Metal private-loop decode will serialize concurrent requests until \
3944 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3945 );
3946 Some(Arc::new(std::sync::Mutex::new(())))
3947 };
3948 AppState {
3949 slept: Mutex::new(None),
3950 embedding,
3951 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3952 id,
3953 loaded,
3954 batcher,
3955 ceiling,
3956 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3957 }))),
3958 paged_kv,
3959 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3960 tasks: Arc::new(tasks::TaskRegistry::new()),
3961 cancels: Arc::new(cancel::CancelRegistry::new()),
3962 stats: stats::Stats::new(),
3963 streams: resume::StreamRegistry::new(),
3964 model_dir: admin::model_dirs().into_iter().next(),
3965 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3966 kv_pool,
3967 prefix_cache,
3968 sessions: session::SessionStore::new(),
3969 requests_total: std::sync::atomic::AtomicU64::new(0),
3970 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3971 started_at: std::time::Instant::now(),
3972 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3973 detection,
3974 mcp,
3975 continuous_batching_enabled: enable_continuous_batching,
3976 metal_private_decode_gate,
3977 loading_model: Mutex::new(None),
3978 last_load_error: Mutex::new(None),
3979 serving: Mutex::new(crate::stats::ServingStats::default()),
3980 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3981 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3982 started_unix: unix_now(),
3983 }
3984}
3985
3986fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3994 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3995 return Ok(None);
3996 };
3997 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3998 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3999 tracing::info!(
4000 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
4001 model.name(),
4002 model.architecture(),
4003 model.n_embd(),
4004 model.pooling_type().name(),
4005 model.n_ctx_train(),
4006 );
4007 Ok(Some(Arc::new(model)))
4008}
4009
4010fn unix_now() -> u64 {
4015 std::time::SystemTime::now()
4016 .duration_since(std::time::UNIX_EPOCH)
4017 .map(|d| d.as_secs())
4018 .unwrap_or(0)
4019}
4020
4021fn startup_model_id() -> Option<String> {
4026 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
4027 let configured = std::fs::canonicalize(&configured).ok()?;
4028 admin::discover(&admin::model_dirs())
4029 .into_iter()
4030 .find(|d| {
4031 std::fs::canonicalize(&d.path)
4032 .map(|p| p == configured)
4033 .unwrap_or(false)
4034 })
4035 .map(|d| d.id)
4036}
4037
4038fn init_cpu_pool() {
4047 match frink_core::threads::init_cpu_pool() {
4048 Some(n) => eprintln!(
4049 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
4050 frink_core::threads::perf_core_count()
4051 ),
4052 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
4053 }
4054}
4055
4056fn announce_ready(addr: SocketAddr, scheme: &str) {
4069 use std::io::Write;
4070 let ready =
4071 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
4072 let mut stdout = std::io::stdout().lock();
4073 let _ = writeln!(stdout, "{}", ready.to_line());
4074 let _ = stdout.flush();
4075}
4076
4077async fn shutdown_signal(exit_on_stdin_close: bool) {
4089 if !exit_on_stdin_close {
4090 std::future::pending::<()>().await;
4091 return;
4092 }
4093 let _ = tokio::task::spawn_blocking(|| {
4094 use std::io::Read;
4095 let mut sink = [0u8; 256];
4096 let mut stdin = std::io::stdin().lock();
4097 loop {
4098 match stdin.read(&mut sink) {
4099 Ok(0) => break,
4101 Ok(_) => continue,
4103 Err(e) => {
4104 tracing::warn!("stdin read failed ({e}); treating it as closed");
4105 break;
4106 }
4107 }
4108 }
4109 })
4110 .await;
4111 tracing::info!("stdin closed; shutting down");
4112}
4113
4114fn tokio_worker_threads() -> usize {
4119 std::env::var("FRINK_TOKIO_WORKERS")
4120 .ok()
4121 .and_then(|v| v.trim().parse::<usize>().ok())
4122 .filter(|n| *n > 0)
4123 .unwrap_or(2)
4124}
4125
4126fn install_ring_crypto_provider() {
4143 let _ = rustls::crypto::ring::default_provider().install_default();
4144}
4145
4146pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4152 if args.list_devices {
4153 frink_models::devices::print_available_devices();
4154 return Ok(());
4155 }
4156 apply_cli_overrides(&args)?;
4157
4158 let _instance = {
4162 use frink_core::instance::{register, InstancePolicy};
4163 let policy = if args.allow_multiple_instances {
4164 InstancePolicy::Multi
4165 } else {
4166 InstancePolicy::from_env_or(InstancePolicy::Single)
4167 };
4168 let model = std::env::var("FRINK_MODEL_PATH").ok();
4169 register(
4170 "server",
4171 model.as_deref(),
4172 frink_core::instance::current_backend(),
4173 policy,
4174 )
4175 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4176 };
4177
4178 let journal = journal::Journal::from_env();
4179 eprintln!(
4180 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4181 journal.path()
4182 );
4183 journal.append(&journal::Record::session_start(
4184 env!("CARGO_PKG_VERSION"),
4185 std::process::id(),
4186 ));
4187 journal::install_panic_hook(journal.clone());
4188
4189 let mcp_config_path = args.mcp_config.clone();
4190 let exit_on_stdin_close = args.exit_on_stdin_close
4191 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4192 .map(|v| v == "1")
4193 .unwrap_or(false);
4194
4195 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4199 init_cpu_pool();
4200
4201 let runtime = tokio::runtime::Builder::new_multi_thread()
4202 .worker_threads(tokio_worker_threads())
4203 .enable_all()
4204 .build()?;
4205 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4206
4207 let reason = match &result {
4208 Ok(()) => "normal".to_string(),
4209 Err(e) => e.to_string(),
4210 };
4211 journal.append(&journal::Record::session_exit(reason));
4212
4213 runtime.shutdown_background();
4218
4219 result
4220}
4221
4222async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4223 let _ = tracing_subscriber::fmt::try_init();
4229
4230 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4240 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4241 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4242 .map(|v| v == "1")
4243 .unwrap_or(false);
4244 if let Err(msg) =
4245 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4246 {
4247 anyhow::bail!(msg);
4248 }
4249
4250 let embedding_model = load_embedding_model()?;
4257
4258 let mut loaded = model::load()?;
4259 match &loaded {
4260 model::LoadedModel::Gguf(g) => tracing::info!(
4261 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4262 g.decoder.config.name,
4263 g.is_synthetic,
4264 g.tokenizer.kind()
4265 ),
4266 model::LoadedModel::Kimi(k) => tracing::info!(
4267 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4268 k.tokenizer.vocab_size()
4269 ),
4270 model::LoadedModel::Mla(m) => tracing::info!(
4271 "loaded MLA GGUF '{}' (tokenizer={})",
4272 m.name,
4273 m.tokenizer.kind()
4274 ),
4275 model::LoadedModel::Gemma4(m) => tracing::info!(
4276 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4277 m.name,
4278 m.tokenizer.kind()
4279 ),
4280 model::LoadedModel::Glm52(g) => tracing::info!(
4281 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4282 g.name,
4283 g.tokenizer.kind()
4284 ),
4285 model::LoadedModel::Encoder(_) => {}
4288 }
4289 let metal_default_moe_budget = {
4297 #[cfg(feature = "metal")]
4298 {
4299 frink_core::metal_dense_enabled()
4300 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4301 }
4302 #[cfg(not(feature = "metal"))]
4303 {
4304 false
4305 }
4306 };
4307 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4308 let budget: u64 = budget_str
4309 .parse()
4310 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4311 match &mut loaded {
4312 model::LoadedModel::Gguf(g) => {
4313 tracing::info!(
4314 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4315 (CUDA and/or Metal matvecs when built with the matching feature)"
4316 );
4317 g.decoder.gpu_vram_budget_bytes = Some(budget);
4318 }
4319 model::LoadedModel::Kimi(_) => {
4320 tracing::warn!(
4321 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4322 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4323 );
4324 }
4325 model::LoadedModel::Mla(_) => {
4326 tracing::warn!(
4327 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4328 FFN path only today; ignoring expert VRAM budget"
4329 );
4330 }
4331 model::LoadedModel::Gemma4(_) => {
4332 tracing::warn!(
4333 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4334 ignoring expert VRAM budget"
4335 );
4336 }
4337 model::LoadedModel::Glm52(_) => {
4338 tracing::warn!(
4339 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4340 GPU expert placement not wired yet; ignoring"
4341 );
4342 }
4343 model::LoadedModel::Encoder(_) => {
4344 tracing::warn!(
4345 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4346 it has no routed experts to place; ignoring"
4347 );
4348 }
4349 }
4350 } else if metal_default_moe_budget {
4351 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4355 if let model::LoadedModel::Gguf(g) = &mut loaded {
4356 tracing::info!(
4357 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4358 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4359 );
4360 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4361 }
4362 }
4363 #[cfg(feature = "cuda")]
4364 {
4365 if frink_core::cuda_dense_enabled() {
4366 tracing::info!(
4367 "CUDA dense matvec enabled for WeightMatrix::apply \
4368 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4369 );
4370 } else {
4371 tracing::info!(
4372 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4373 );
4374 }
4375 }
4376 #[cfg(feature = "metal")]
4377 {
4378 if frink_core::metal_dense_enabled() {
4379 tracing::info!(
4380 "Metal dense matvec enabled for WeightMatrix::apply \
4381 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4382 );
4383 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4384 Some("1") | Some("true") | Some("on") | Some("attn") => {
4385 tracing::info!(
4386 "Metal fused attention requested (FRINK_METAL_ATTN): \
4387 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4388 );
4389 }
4390 _ => {}
4391 }
4392 tracing::info!(
4393 "Metal greedy GPU argmax: temperature<=0 folds \
4394 final_norm+lm_head+argmax into the dense stack"
4395 );
4396 } else {
4397 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4398 }
4399 }
4400 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4411 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4412 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4413 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4414 panic!(
4415 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4416 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4417 );
4418 }
4419 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4420 (Ok(blocks), Ok(block_size), Err(_)) => {
4421 let total_blocks: usize = blocks
4422 .parse()
4423 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4424 let block_size: usize = block_size
4425 .parse()
4426 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4427 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4433 .ok()
4434 .map(|v| {
4435 v.parse()
4436 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4437 })
4438 .unwrap_or(0);
4439 tracing::info!(
4440 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4441 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4442 );
4443 Some(generate::KvPoolConfig {
4444 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4445 queue_wait: Duration::from_millis(queue_wait_ms),
4446 })
4447 }
4448 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4449 let block_size: usize = block_size
4450 .parse()
4451 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4452 let budget: u64 = byte_budget
4453 .parse()
4454 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4455 let cfg = match &loaded {
4456 model::LoadedModel::Gguf(g) => &g.decoder.config,
4457 model::LoadedModel::Kimi(_)
4458 | model::LoadedModel::Mla(_)
4459 | model::LoadedModel::Gemma4(_)
4460 | model::LoadedModel::Glm52(_)
4461 | model::LoadedModel::Encoder(_) => {
4462 panic!(
4463 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4464 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4465 );
4466 }
4467 };
4468 let bytes_per_block = block_size
4469 * cfg.kv_heads_all_layers()
4470 * (cfg.head_dim + cfg.v_head_dim())
4471 * std::mem::size_of::<f32>();
4472 assert!(
4473 bytes_per_block > 0,
4474 "derived KV block byte size must be positive (check model config and block size)"
4475 );
4476 let total_blocks = (budget as usize / bytes_per_block).max(1);
4477 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4478 .ok()
4479 .map(|v| {
4480 v.parse()
4481 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4482 })
4483 .unwrap_or(0);
4484 tracing::info!(
4485 "KV cache block pool enabled from byte budget: {budget} bytes / \
4486 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4487 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4488 cfg.n_layers
4489 );
4490 Some(generate::KvPoolConfig {
4491 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4492 queue_wait: Duration::from_millis(queue_wait_ms),
4493 })
4494 }
4495 (Err(_), Err(_), Err(_)) => None,
4496 (Err(_), Ok(_), Err(_)) => panic!(
4497 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4498 (or unset all three to disable KV cache pooling)"
4499 ),
4500 (Ok(_), Ok(_), Ok(_)) => {
4501 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4502 }
4503 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4504 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4505 set together (or neither, to disable KV cache pooling)"
4506 ),
4507 };
4508 let paged_kv = match (
4514 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4515 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4516 ) {
4517 (Ok(blocks), Ok(block_size)) => {
4518 assert!(
4519 kv_pool.is_none(),
4520 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4521 mutually exclusive: both bound the same KV memory, by different means. \
4522 Set one."
4523 );
4524 let blocks_per_layer: usize = blocks
4538 .parse()
4539 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4540 let block_size: usize = block_size
4541 .parse()
4542 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4543 let gguf = match &loaded {
4544 model::LoadedModel::Gguf(g) => g,
4545 _ => panic!(
4546 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4547 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4548 ),
4549 };
4550 let cfg = &gguf.decoder.config;
4551 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4552 .ok()
4553 .map(|v| {
4554 v.parse()
4555 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4556 })
4557 .unwrap_or(0);
4558 tracing::info!(
4559 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4560 layer across {} layers, shared by all concurrent requests, \
4561 {queue_wait_ms}ms admission queue wait",
4562 cfg.n_layers
4563 );
4564 let radix = Some(Arc::new(Mutex::new(
4569 crate::policy::radix::SaltedRadix::new(block_size),
4570 )));
4571 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4580 crate::policy::parser::ToolCallFormat::infer(
4581 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4582 )
4583 .opener(),
4584 |text| {
4585 gguf.tokenizer
4586 .encode(text, SpecialTokens::Parse)
4587 .into_iter()
4588 .map(|t| t as u32)
4589 .collect()
4590 },
4591 );
4592 if let Some(id) = anchor_token {
4593 tracing::info!(
4594 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4595 window stops short of where its next turn rejoins"
4596 );
4597 }
4598 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4599 .ok()
4600 .map(|v| {
4601 v.parse()
4602 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4603 })
4604 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4605 if let Some(window) = cfg.uniform_sliding_window() {
4606 tracing::info!(
4607 "Paged KV window slide enabled: every layer slides by {window} every \
4608 {slide_interval} decode steps, so a request holds its prompt and a \
4609 window rather than its whole context"
4610 );
4611 } else if cfg.kv_block_window().is_some() {
4612 tracing::info!(
4613 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4614 and a page group holds one block in every layer"
4615 );
4616 }
4617 Some(generate::PagedKvConfig {
4618 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4621 queue_wait: Duration::from_millis(queue_wait_ms),
4622 radix,
4623 anchor_token,
4624 slide_interval,
4625 })
4626 }
4627 (Err(_), Err(_)) => None,
4628 _ => panic!(
4629 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4630 (or neither, to disable paged KV)"
4631 ),
4632 };
4633 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4640 let max_entries: usize = v
4641 .parse()
4642 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4643 if kv_pool.is_some() {
4644 tracing::warn!(
4645 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4646 caching will never be consulted while a KV pool is configured"
4647 );
4648 }
4649 assert!(
4656 paged_kv.is_none(),
4657 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4658 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4659 produce, so the cache could never hit. Set one."
4660 );
4661 tracing::info!(
4662 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4663 all requests"
4664 );
4665 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4666 });
4667 if matches!(
4668 loaded,
4669 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4670 ) && (kv_pool.is_some() || prefix_cache.is_some())
4671 {
4672 tracing::warn!(
4673 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4674 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4675 frink_models::engine's module docs"
4676 );
4677 }
4678 let enable_cb =
4679 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4680 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4681 tracing::info!(
4682 "continuous batching enabled by default on Metal for safe parallel serving \
4683 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4684 );
4685 }
4686 if continuous_batching_env() == Some(true)
4687 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4688 && (kv_pool.is_some() || prefix_cache.is_some())
4689 {
4690 tracing::warn!(
4691 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4692 (those modes keep the private generate path)"
4693 );
4694 }
4695 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4696 if let Ok(chunk) = n.parse::<usize>() {
4697 if chunk > 0 {
4698 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4699 }
4700 }
4701 }
4702 if matches!(
4703 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4704 Some("1")
4705 ) {
4706 tracing::warn!(
4707 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4708 (minimal spill; full layer offload still planned)"
4709 );
4710 }
4711
4712 let mcp = match mcp_config_path {
4713 Some(path) => {
4714 let loaded = mcp::load_mcp_config(&path)?;
4715 tracing::info!(
4716 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4717 loaded.path,
4718 loaded.servers.len()
4719 );
4720 Some(loaded)
4721 }
4722 None => None,
4723 };
4724
4725 let detection = health::Detection::spawn();
4729
4730 let state = Arc::new(build_app_state(
4731 StartupModels {
4732 loaded,
4733 embedding: embedding_model,
4734 },
4735 kv_pool,
4736 paged_kv,
4737 prefix_cache,
4738 enable_cb,
4739 mcp,
4740 detection,
4741 ));
4742
4743 use frink_api::routes;
4747
4748 let public = Router::new().route(routes::HEALTH, get(health));
4753
4754 let mut protected = protected_routes();
4755
4756 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4761 tracing::info!("API key auth enabled");
4762 let auth = limits::AuthConfig {
4763 api_key: Arc::new(key),
4764 };
4765 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4766 auth,
4767 limits::require_api_key,
4768 ));
4769 }
4770 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4771 let rpm: u32 = rpm
4772 .parse()
4773 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4774 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4775 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4776 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4777 limiter,
4778 limits::rate_limit,
4779 ));
4780 }
4781 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4789 let origins = security::parse_cors_origins(&spec)
4790 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4791 tracing::info!(
4792 "CORS enabled: {} allow-listed origin(s) ({})",
4793 origins.len(),
4794 spec
4795 );
4796 let cors = tower_http::cors::CorsLayer::new()
4797 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4798 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4799 .allow_headers([
4800 axum::http::header::CONTENT_TYPE,
4801 axum::http::header::AUTHORIZATION,
4802 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4808 axum::http::HeaderName::from_static("last-event-id"),
4812 ]);
4813 protected = protected.route_layer(cors);
4814 }
4815
4816 let app = public
4820 .merge(protected)
4821 .layer(axum::middleware::from_fn(limits::retry_after))
4822 .with_state(state);
4823
4824 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4830 install_ring_crypto_provider();
4831 match tls_paths {
4836 Some(paths) => {
4837 let config =
4838 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4839 .await
4840 .map_err(|e| {
4841 anyhow::anyhow!(
4842 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4843 paths.cert,
4844 paths.key
4845 )
4846 })?;
4847 let socket_addr: std::net::SocketAddr = addr
4848 .parse()
4849 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4850 let listener = std::net::TcpListener::bind(socket_addr)?;
4851 listener.set_nonblocking(true)?;
4859 let bound = listener.local_addr()?;
4860 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4861 announce_ready(bound, "https");
4862
4863 let handle = axum_server::Handle::new();
4864 let shutdown_handle = handle.clone();
4865 tokio::spawn(async move {
4866 shutdown_signal(exit_on_stdin_close).await;
4867 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4868 });
4869 axum_server::from_tcp_rustls(listener, config)?
4870 .handle(handle)
4871 .serve(app.into_make_service())
4872 .await?;
4873 }
4874 None => {
4875 let listener = tokio::net::TcpListener::bind(&addr).await?;
4876 let bound = listener.local_addr()?;
4877 tracing::info!("frink-server listening on {bound}");
4878 announce_ready(bound, "http");
4879 axum::serve(listener, app)
4880 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4881 .await?;
4882 }
4883 }
4884 Ok(())
4885}
4886
4887#[cfg(test)]
4888pub(crate) mod tests {
4889 use super::*;
4890 use frink_models::config::test_dense_fixture;
4891
4892 #[test]
4893 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4894 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4895 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4896 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4897 assert_eq!(parsed.port, 51999);
4898 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4899 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4901 }
4902
4903 fn test_model() -> Model {
4904 let cfg = test_dense_fixture();
4908 Model::Gguf(GgufModel {
4909 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4910 tokenizer: Arc::new(ServerTokenizer::Byte),
4911 stop_tokens: StopTokens::default(),
4912 bos_id: None,
4913 is_synthetic: true,
4914 chat_template: chat_template::PromptTemplate::plain(),
4915 })
4916 }
4917
4918 fn greedy_params(max_tokens: usize) -> GenerationParams {
4919 GenerationParams {
4920 cache_salt: None,
4921 prompt_logprobs: None,
4922 wants_logprobs: false,
4923 n: 1,
4924 interleave_choices: false,
4925 logit_bias: crate::logit_bias::LogitBias::default(),
4926 keep_special_tokens: false,
4927 truncate_prompt_tokens: None,
4928 token_mask: crate::token_mask::TokenMask::default(),
4929 reasoning: None,
4930 max_tokens,
4931 sampling: SamplingParams::default(),
4932 seed: 1,
4933 stop: Vec::new(),
4934 stop_token_ids: Vec::new(),
4935 json_object: false,
4936 grammar: None,
4937 cancel: None,
4938 ignore_eos: false,
4939 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4940 lora: None,
4941 }
4942 }
4943
4944 fn test_model_full_byte_vocab() -> Model {
4948 test_model_full_byte_vocab_with_eos(None)
4949 }
4950
4951 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4958 test_byte_model(eos, true)
4959 }
4960
4961 fn test_byte_model(eos: Option<usize>, synthetic: bool) -> Model {
4968 let mut cfg = test_dense_fixture();
4969 cfg.vocab_size = 256;
4970 Model::Gguf(GgufModel {
4971 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4972 tokenizer: Arc::new(ServerTokenizer::Byte),
4973 stop_tokens: StopTokens::from_eos(eos),
4974 bos_id: None,
4975 is_synthetic: synthetic,
4976 chat_template: chat_template::PromptTemplate::plain(),
4977 })
4978 }
4979
4980 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4984 test_state_at(model, response_cache, None)
4985 }
4986
4987 pub(crate) fn test_state_at(
4991 model: Model,
4992 response_cache: ResponseCache,
4993 checkpoint_path: Option<std::path::PathBuf>,
4994 ) -> AppState {
4995 AppState {
4996 slept: Mutex::new(None),
4997 embedding: None,
4998 paged_kv: None,
4999 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
5000 id: None,
5001 loaded: Loaded::Generative(Arc::new(model)),
5002 batcher: None,
5003 ceiling: None,
5004 checkpoint_path,
5005 }))),
5006 load_in_progress: std::sync::atomic::AtomicBool::new(false),
5007 tasks: Arc::new(tasks::TaskRegistry::new()),
5008 cancels: Arc::new(cancel::CancelRegistry::new()),
5009 stats: stats::Stats::new(),
5010 streams: resume::StreamRegistry::new(),
5011 model_dir: None,
5012 response_cache: Mutex::new(response_cache),
5013 kv_pool: None,
5014 prefix_cache: None,
5015 sessions: session::SessionStore::new(),
5016 requests_total: std::sync::atomic::AtomicU64::new(0),
5017 request_errors_total: std::sync::atomic::AtomicU64::new(0),
5018 started_at: std::time::Instant::now(),
5019 last_request_ms: std::sync::atomic::AtomicU64::new(0),
5020 detection: Arc::new(health::Detection::ready(health::probe_backends())),
5021 mcp: None,
5022 continuous_batching_enabled: false,
5023 metal_private_decode_gate: None,
5024 loading_model: Mutex::new(None),
5025 last_load_error: Mutex::new(None),
5026 serving: Mutex::new(crate::stats::ServingStats::default()),
5027 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
5028 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
5029 started_unix: unix_now(),
5030 }
5031 }
5032
5033 pub(crate) fn test_app() -> Router {
5042 test_app_with_state(Arc::new(test_state(
5043 test_model_full_byte_vocab(),
5044 ResponseCache::new(1000, Duration::from_secs(3600)),
5045 )))
5046 }
5047
5048 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
5051 Router::new()
5061 .route(frink_api::routes::HEALTH, get(health))
5062 .merge(protected_routes())
5063 .with_state(state)
5064 }
5065
5066 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
5067 let mut cfg = test_dense_fixture();
5068 cfg.name = name;
5069 cfg.vocab_size = vocab_size;
5070 Model::Gguf(GgufModel {
5071 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5072 tokenizer: Arc::new(ServerTokenizer::Byte),
5073 stop_tokens: StopTokens::default(),
5074 bos_id: None,
5075 is_synthetic: true,
5076 chat_template: chat_template::PromptTemplate::plain(),
5077 })
5078 }
5079
5080 fn model_with_template(name: &'static str, source: &str) -> Model {
5084 let mut cfg = test_dense_fixture();
5085 cfg.name = name;
5086 cfg.vocab_size = 256;
5087 Model::Gguf(GgufModel {
5088 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5089 tokenizer: Arc::new(ServerTokenizer::Byte),
5090 stop_tokens: StopTokens::default(),
5091 bos_id: None,
5092 is_synthetic: true,
5093 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5094 Some(source),
5095 Some("qwen3"),
5096 false,
5097 true,
5098 None,
5099 None,
5100 ),
5101 })
5102 }
5103
5104 #[tokio::test]
5113 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5114 let strict = "{% if messages | length > 1 %}\
5117 {{ raise_exception('this template takes one turn') }}\
5118 {% endif %}{{ messages[0].content }}";
5119 let state = Arc::new(test_state(
5120 model_with_template("strict", strict),
5121 ResponseCache::new(4, Duration::from_secs(60)),
5122 ));
5123 let app = test_app_with_state(state);
5124
5125 let (status, body) = post_json_uri(
5126 &app,
5127 "/v1/chat/completions",
5128 serde_json::json!({
5129 "model": "strict",
5130 "stream": true,
5131 "messages": [
5132 {"role": "user", "content": "one"},
5133 {"role": "user", "content": "two"},
5134 ],
5135 }),
5136 )
5137 .await;
5138 assert_eq!(status, StatusCode::BAD_REQUEST);
5139 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5140 assert!(
5141 body["error"]["message"]
5142 .as_str()
5143 .unwrap()
5144 .contains("one turn"),
5145 "the template's own message must reach the caller: {body}"
5146 );
5147
5148 let (status, _) = post_json_uri(
5150 &app,
5151 "/v1/chat/completions",
5152 serde_json::json!({
5153 "model": "strict",
5154 "stream": true,
5155 "max_tokens": 1,
5156 "messages": [{"role": "user", "content": "one"}],
5157 }),
5158 )
5159 .await;
5160 assert_eq!(status, StatusCode::OK);
5161 }
5162
5163 #[tokio::test]
5165 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5166 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5167 {% if reasoning_effort %}\
5168 {% if reasoning_effort not in ['low','medium','high'] %}\
5169 {{ raise_exception('bad effort') }}\
5170 {% endif %}[{{ reasoning_effort }}]\
5171 {% endif %}{{ messages[0].content }}";
5172 let state = Arc::new(test_state(
5173 model_with_template("thinker", reasoning),
5174 ResponseCache::new(4, Duration::from_secs(60)),
5175 ));
5176 let app = test_app_with_state(state);
5177 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5178 assert_eq!(status, StatusCode::OK);
5179 let entry = &models["data"][0];
5180 assert_eq!(
5181 entry["supported_reasoning_efforts"],
5182 serde_json::json!(["off", "low", "medium", "high"])
5183 );
5184 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5185 }
5186
5187 #[tokio::test]
5192 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5193 let app = test_app();
5194 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5195 let entry = &models["data"][0];
5196 assert!(entry.get("supported_reasoning_efforts").is_none());
5197 assert!(entry.get("default_reasoning_effort").is_none());
5198 }
5199
5200 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5201 Arc::new(ActiveModel {
5202 id: Some(name.to_string()),
5203 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5204 batcher: None,
5205 ceiling: None,
5206 checkpoint_path: None,
5207 })
5208 .tap_into(state)
5209 }
5210
5211 trait TapInto {
5213 fn tap_into(self, state: &AppState) -> Self;
5214 }
5215 impl TapInto for Arc<ActiveModel> {
5216 fn tap_into(self, state: &AppState) -> Self {
5217 state.swap_active(Some(Arc::clone(&self)));
5218 self
5219 }
5220 }
5221
5222 #[test]
5228 fn an_in_flight_request_keeps_the_model_it_started_on() {
5229 let state = test_state(
5230 named_test_model("model-a", 256),
5231 ResponseCache::new(4, Duration::from_secs(60)),
5232 );
5233
5234 let in_flight = state.active().expect("a model is loaded");
5237 assert_eq!(in_flight.name(), "model-a");
5238
5239 active_model(&state, "model-b");
5240
5241 assert_eq!(state.active().unwrap().name(), "model-b");
5243 assert_eq!(in_flight.name(), "model-a");
5245 let produced = run_generation(
5246 in_flight.generative().unwrap(),
5247 "hi",
5248 &greedy_params(3),
5249 None,
5250 None,
5251 None,
5252 None,
5253 None,
5254 None,
5255 )
5256 .expect("the old model must still decode after being swapped out");
5257 assert!(matches!(
5258 produced.choices[0].finish,
5259 FinishReason::Length | FinishReason::Stop
5260 ));
5261 }
5262
5263 #[test]
5268 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5269 let state = test_state(
5270 named_test_model("model-a", 256),
5271 ResponseCache::new(4, Duration::from_secs(60)),
5272 );
5273 let in_flight = state.active().expect("a model is loaded");
5274 let weights = Arc::clone(in_flight.generative().unwrap());
5275 assert!(Arc::strong_count(&weights) >= 2);
5276
5277 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5278 id: Some("model-b".to_string()),
5279 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5280 batcher: None,
5281 ceiling: None,
5282 checkpoint_path: None,
5283 })));
5284 drop(previous);
5285 assert!(Arc::strong_count(&weights) >= 2);
5287 drop(in_flight);
5288 assert_eq!(Arc::strong_count(&weights), 1);
5289 }
5290
5291 #[tokio::test]
5295 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5296 let state = Arc::new(test_state(
5297 named_test_model("model-a", 256),
5298 ResponseCache::new(4, Duration::from_secs(60)),
5299 ));
5300 let app = test_app_with_state(Arc::clone(&state));
5301
5302 let (status, body) = post_json_uri(
5303 &app,
5304 frink_api::routes::ADMIN_MODELS_UNLOAD,
5305 serde_json::json!({}),
5306 )
5307 .await;
5308 assert_eq!(status, StatusCode::OK);
5309 assert_eq!(body["ok"], true);
5310 assert!(body["active"].is_null());
5311 assert!(state.active().is_none());
5312
5313 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5314 assert_eq!(status, StatusCode::OK);
5315 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5316 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5317
5318 let (status, body) = post_json_uri(
5319 &app,
5320 "/v1/chat/completions",
5321 serde_json::json!({
5322 "model": "x",
5323 "messages": [{"role": "user", "content": "hi"}]
5324 }),
5325 )
5326 .await;
5327 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5328 assert_eq!(body["error"]["type"], "model_not_loaded");
5329 }
5330
5331 #[tokio::test]
5335 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5336 let state = Arc::new(test_state(
5337 named_test_model("model-a", 256),
5338 ResponseCache::new(4, Duration::from_secs(60)),
5339 ));
5340 let app = test_app_with_state(Arc::clone(&state));
5341 state.swap_active(None);
5342
5343 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5344 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5347 assert_eq!(body["state"], "unavailable");
5348 assert_eq!(body["reason"], "model_not_loaded");
5349 assert!(body["model"].is_null());
5350 let real_weights = body["capabilities"]
5351 .as_array()
5352 .unwrap()
5353 .iter()
5354 .find(|c| c["id"] == "real_weights")
5355 .cloned()
5356 .expect("real_weights is always reported");
5357 assert_eq!(real_weights["available"], false);
5358 assert_eq!(real_weights["reason"], "model_not_loaded");
5359 }
5360
5361 #[tokio::test]
5365 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5366 let app = test_app();
5367
5368 let (status, completion) = post_json_uri(
5369 &app,
5370 "/v1/chat/completions",
5371 serde_json::json!({
5372 "model": "x",
5373 "messages": [{"role": "user", "content": "hi"}],
5374 "max_tokens": 4
5375 }),
5376 )
5377 .await;
5378 assert_eq!(status, StatusCode::OK);
5379 let request_id = completion["request_id"].as_str().unwrap().to_string();
5380
5381 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5382 assert_eq!(status, StatusCode::OK);
5383 let recent = stats["recent"].as_array().unwrap();
5384 assert_eq!(recent.len(), 1);
5385 let row = &recent[0];
5386 assert_eq!(row["request_id"], request_id);
5387 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5388 assert_eq!(row["status"], 200);
5389 assert_eq!(row["stream"], false);
5390 assert!(row["duration_ms"].is_number());
5393 assert!(row["decode_ms"].is_number());
5394 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5395 assert_eq!(
5396 stats["tokens_prompt_total"].as_u64().unwrap(),
5397 row["prompt_tokens"].as_u64().unwrap()
5398 );
5399 }
5400
5401 #[tokio::test]
5405 async fn a_rejected_request_is_recorded_too() {
5406 let state = Arc::new(test_state(
5407 named_test_model("model-a", 256),
5408 ResponseCache::new(4, Duration::from_secs(60)),
5409 ));
5410 let app = test_app_with_state(Arc::clone(&state));
5411 state.swap_active(None);
5412
5413 let (status, _) = post_json_uri(
5414 &app,
5415 "/v1/chat/completions",
5416 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5417 )
5418 .await;
5419 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5420
5421 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5422 let recent = stats["recent"].as_array().unwrap();
5423 assert_eq!(recent.len(), 1);
5424 assert_eq!(recent[0]["status"], 503);
5425 assert_eq!(recent[0]["completion_tokens"], 0);
5426 assert!(recent[0]["decode_ms"].is_null());
5427 assert_eq!(stats["errors_total"], 1);
5428 }
5429
5430 async fn post_json_with_headers(
5434 app: &Router,
5435 uri: &str,
5436 body: serde_json::Value,
5437 headers: &[(&str, &str)],
5438 ) -> (StatusCode, serde_json::Value) {
5439 use http_body_util::BodyExt;
5440 use tower::ServiceExt;
5441
5442 let mut builder = axum::http::Request::builder()
5443 .method("POST")
5444 .uri(uri)
5445 .header("content-type", "application/json");
5446 for (name, value) in headers {
5447 builder = builder.header(*name, *value);
5448 }
5449 let response = app
5450 .clone()
5451 .oneshot(
5452 builder
5453 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5454 .unwrap(),
5455 )
5456 .await
5457 .unwrap();
5458 let status = response.status();
5459 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5460 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5461 (status, json)
5462 }
5463
5464 #[tokio::test]
5468 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5469 let app = test_app();
5470
5471 let (status, _) = post_json_uri(
5472 &app,
5473 frink_api::routes::V1_TOKENIZE,
5474 serde_json::json!({"prompt": "hello"}),
5475 )
5476 .await;
5477 assert_eq!(status, StatusCode::OK);
5478 let (status, _) = post_json_uri(
5479 &app,
5480 frink_api::routes::V1_DETOKENIZE,
5481 serde_json::json!({"tokens": [104, 105]}),
5482 )
5483 .await;
5484 assert_eq!(status, StatusCode::OK);
5485 let (status, _) = post_json_uri(
5486 &app,
5487 frink_api::routes::V1_EMBEDDINGS,
5488 serde_json::json!({"input": "hello"}),
5489 )
5490 .await;
5491 assert_eq!(status, StatusCode::OK);
5492
5493 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5494 let routes: Vec<&str> = stats["recent"]
5495 .as_array()
5496 .unwrap()
5497 .iter()
5498 .map(|row| row["route"].as_str().unwrap())
5499 .collect();
5500 for expected in [
5501 frink_api::routes::V1_TOKENIZE,
5502 frink_api::routes::V1_DETOKENIZE,
5503 frink_api::routes::V1_EMBEDDINGS,
5504 ] {
5505 assert!(
5506 routes.contains(&expected),
5507 "{expected} is missing: {routes:?}"
5508 );
5509 }
5510
5511 let row = |route: &str| {
5512 stats["recent"]
5513 .as_array()
5514 .unwrap()
5515 .iter()
5516 .find(|r| r["route"] == route)
5517 .cloned()
5518 .unwrap()
5519 };
5520 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5524 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5525 assert!(embed["decode_ms"].is_null());
5526 assert_eq!(embed["completion_tokens"], 0);
5527 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5531 assert_eq!(
5532 stats["tokens_prompt_total"].as_u64().unwrap(),
5533 embed["prompt_tokens"].as_u64().unwrap(),
5534 "only the forward pass counted"
5535 );
5536 }
5537
5538 fn streaming_test_app() -> Router {
5543 let mut cfg = test_dense_fixture();
5544 cfg.vocab_size = 256;
5545 let model = Model::Gguf(GgufModel {
5546 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5547 tokenizer: Arc::new(ServerTokenizer::Byte),
5548 stop_tokens: StopTokens::default(),
5549 bos_id: None,
5550 is_synthetic: false,
5551 chat_template: chat_template::PromptTemplate::plain(),
5552 });
5553 test_app_with_state(Arc::new(test_state(
5554 model,
5555 ResponseCache::new(1000, Duration::from_secs(3600)),
5556 )))
5557 }
5558
5559 #[tokio::test]
5568 async fn chat_logprobs_are_rendered_and_are_never_served_from_cache() {
5569 let app = test_app();
5570 let body = |logprobs: Option<(bool, Option<u32>)>| {
5571 let mut b = serde_json::json!({
5572 "model": "x",
5573 "messages": [{"role": "user", "content": "hi"}],
5574 "max_tokens": 4
5575 });
5576 if let Some((on, top)) = logprobs {
5577 b["logprobs"] = serde_json::json!(on);
5578 if let Some(n) = top {
5579 b["top_logprobs"] = serde_json::json!(n);
5580 }
5581 }
5582 b
5583 };
5584
5585 let (status, plain) =
5587 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(None)).await;
5588 assert_eq!(status, StatusCode::OK, "{plain}");
5589 assert!(plain["choices"][0]["logprobs"].is_null(), "{plain}");
5590
5591 for attempt in 0..2 {
5595 let (status, with) = post_json_uri(
5596 &app,
5597 frink_api::routes::V1_CHAT_COMPLETIONS,
5598 body(Some((true, Some(2)))),
5599 )
5600 .await;
5601 assert_eq!(status, StatusCode::OK, "{with}");
5602 assert_ne!(
5603 with["frink_cache"], "hit",
5604 "attempt {attempt} replayed a cached answer for a logprobs request: {with}"
5605 );
5606 let lp = &with["choices"][0]["logprobs"];
5607 assert!(lp.is_object(), "attempt {attempt}: {with}");
5608 let content = lp["content"].as_array().expect("content");
5609 assert!(lp["tokens"].is_null(), "completions shape leaked: {lp}");
5611 for entry in content {
5612 assert!(entry["token"].is_string(), "{entry}");
5613 assert!(entry["bytes"].is_array(), "{entry}");
5614 let v = entry["logprob"].as_f64().expect("a real number");
5615 assert!(v <= 0.0 && v.is_finite(), "{entry}");
5616 let top = entry["top_logprobs"].as_array().expect("top_logprobs");
5617 assert!(top.len() <= 2, "asked for 2, got {}", top.len());
5618 }
5619 }
5620 }
5621
5622 #[tokio::test]
5628 async fn the_chat_logprobs_pair_is_validated() {
5629 let app = test_app();
5630 for (extra, why) in [
5631 (serde_json::json!({"top_logprobs": 3}), "without logprobs"),
5632 (
5633 serde_json::json!({"logprobs": true, "top_logprobs": 21}),
5634 "above the cap",
5635 ),
5636 ] {
5637 let mut body = serde_json::json!({
5638 "model": "x",
5639 "messages": [{"role": "user", "content": "hi"}],
5640 "max_tokens": 2
5641 });
5642 for (k, v) in extra.as_object().unwrap() {
5643 body[k] = v.clone();
5644 }
5645 let (status, answer) =
5646 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await;
5647 assert_eq!(status, StatusCode::BAD_REQUEST, "{why}: {answer}");
5648 assert!(
5649 answer["error"]["message"]
5650 .as_str()
5651 .is_some_and(|m| m.contains("top_logprobs")),
5652 "{why}: {answer}"
5653 );
5654 }
5655 }
5656
5657 #[tokio::test]
5666 async fn sleep_refuses_a_model_it_could_not_bring_back() {
5667 let app = test_app();
5668 let (status, answer) =
5669 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5670 assert_eq!(status, StatusCode::CONFLICT, "{answer}");
5671 assert_eq!(answer["error"]["type"], "not_reloadable", "{answer}");
5672 let (_, still) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5675 assert_eq!(still["is_sleeping"], false, "{still}");
5676 let (status, _) = 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 assert_eq!(status, StatusCode::OK, "a refused sleep unloaded the model");
5687 }
5688
5689 #[tokio::test]
5698 async fn sleep_remembers_what_unload_forgets() {
5699 let state = Arc::new(test_state_at(
5702 test_model_full_byte_vocab(),
5703 ResponseCache::new(1000, Duration::from_secs(3600)),
5704 Some(std::path::PathBuf::from("/nonexistent/fixture.gguf")),
5705 ));
5706 let app = test_app_with_state(Arc::clone(&state));
5707 let ask = || {
5708 let app = app.clone();
5709 async move {
5710 post_json_uri(
5711 &app,
5712 frink_api::routes::V1_CHAT_COMPLETIONS,
5713 serde_json::json!({
5714 "model": "x",
5715 "messages": [{"role": "user", "content": "hi"}],
5716 "max_tokens": 2
5717 }),
5718 )
5719 .await
5720 }
5721 };
5722
5723 let (status, _) = ask().await;
5724 assert_eq!(status, StatusCode::OK, "the fixture server serves");
5725 let (_, awake) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5726 assert_eq!(awake["is_sleeping"], false, "{awake}");
5727
5728 let (status, slept) =
5729 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5730 assert_eq!(status, StatusCode::OK, "{slept}");
5731 assert_eq!(slept["is_sleeping"], true, "{slept}");
5732 let (_, now) = get_json_uri(&app, frink_api::routes::IS_SLEEPING).await;
5733 assert_eq!(now["is_sleeping"], true, "{now}");
5734
5735 let (status, refused) = ask().await;
5738 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE, "{refused}");
5739 assert_eq!(
5740 refused["error"]["type"], "server_sleeping",
5741 "an asleep server reported itself as empty: {refused}"
5742 );
5743
5744 let (status, again) =
5746 post_json_uri(&app, frink_api::routes::SLEEP, serde_json::json!({})).await;
5747 assert_eq!(status, StatusCode::OK, "{again}");
5748 assert_eq!(again["is_sleeping"], true, "{again}");
5749 }
5750
5751 #[tokio::test]
5755 async fn waking_a_server_that_is_awake_is_refused() {
5756 let app = test_app();
5757 let (status, answer) =
5758 post_json_uri(&app, frink_api::routes::WAKE_UP, serde_json::json!({})).await;
5759 assert_eq!(status, StatusCode::CONFLICT, "{answer}");
5760 assert_eq!(answer["error"]["type"], "not_sleeping", "{answer}");
5761 }
5762
5763 #[tokio::test]
5770 async fn a_salt_keeps_one_callers_cached_answer_from_another() {
5771 let app = test_app();
5772 let body = |salt: Option<&str>| {
5773 let mut b = serde_json::json!({
5774 "model": "x",
5775 "messages": [{"role": "user", "content": "the same prompt"}],
5776 "max_tokens": 4,
5777 "seed": 1
5778 });
5779 if let Some(s) = salt {
5780 b["cache_salt"] = serde_json::json!(s);
5781 }
5782 b
5783 };
5784 let post = |b: serde_json::Value| {
5785 let app = app.clone();
5786 async move { post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, b).await }
5787 };
5788
5789 let (status, _) = post(body(Some("tenant-a"))).await;
5791 assert_eq!(status, StatusCode::OK);
5792 let (_, again) = post(body(Some("tenant-a"))).await;
5793 assert_eq!(
5794 again["frink_cache"], "hit",
5795 "the owner did not get its own entry back: {again}"
5796 );
5797
5798 let (_, other) = post(body(Some("tenant-b"))).await;
5800 assert_ne!(
5801 other["frink_cache"], "hit",
5802 "a different caller was served tenant-a's answer: {other}"
5803 );
5804
5805 let (_, shared) = post(body(None)).await;
5807 assert_ne!(
5808 shared["frink_cache"], "hit",
5809 "an unsalted request was served a salted answer: {shared}"
5810 );
5811 }
5812
5813 #[tokio::test]
5816 async fn chat_serves_several_choices_from_one_prefill() {
5817 let app = test_app();
5818 let body = |n: u32, stream: bool| {
5819 serde_json::json!({
5820 "model": "x",
5821 "messages": [{"role": "user", "content": "hi"}],
5822 "max_tokens": 4,
5823 "temperature": 1.0,
5824 "n": n,
5825 "stream": stream
5826 })
5827 };
5828
5829 let (status, one) =
5830 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5831 assert_eq!(status, StatusCode::OK, "{one}");
5832
5833 let (status, three) =
5834 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5835 assert_eq!(status, StatusCode::OK, "{three}");
5836 let choices = three["choices"].as_array().expect("an array");
5837 assert_eq!(choices.len(), 3, "{three}");
5838 for (i, c) in choices.iter().enumerate() {
5839 assert_eq!(c["index"], i);
5840 assert!(c["message"]["role"].is_string(), "{c}");
5841 assert!(c["finish_reason"].is_string(), "{c}");
5842 }
5843 assert_eq!(
5845 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5846 "n = 3 billed the prompt more than once"
5847 );
5848 }
5849
5850 #[tokio::test]
5862 async fn a_streamed_n_interleaves_its_choices() {
5863 let app = streaming_test_app();
5864 let raw = post_sse_raw(
5865 &app,
5866 serde_json::json!({
5867 "model": "x",
5868 "messages": [{"role": "user", "content": "hi"}],
5869 "max_tokens": 6,
5870 "temperature": 1.0,
5871 "n": 3,
5872 "stream": true
5873 }),
5874 )
5875 .await;
5876
5877 let mut order: Vec<usize> = Vec::new();
5879 let mut finished: Vec<usize> = Vec::new();
5880 let mut usage_blocks = 0usize;
5881 for line in raw.lines() {
5882 let Some(rest) = line.strip_prefix("data: ") else {
5883 continue;
5884 };
5885 if rest.trim() == "[DONE]" {
5886 continue;
5887 }
5888 let v: serde_json::Value = serde_json::from_str(rest).expect(rest);
5889 if v.get("usage").is_some_and(|u| !u.is_null()) {
5890 usage_blocks += 1;
5891 }
5892 let Some(choice) = v["choices"].as_array().and_then(|c| c.first()) else {
5893 continue;
5894 };
5895 let index = choice["index"].as_u64().expect("an index") as usize;
5896 if choice["finish_reason"].is_string() {
5897 finished.push(index);
5898 continue;
5899 }
5900 order.push(index);
5901 }
5902
5903 assert_eq!(
5904 finished,
5905 vec![0, 1, 2],
5906 "one terminal chunk per choice, in index order: {raw}"
5907 );
5908 assert_eq!(usage_blocks, 1, "the usage block is the request's: {raw}");
5909 assert!(
5910 order.contains(&0) && order.contains(&2),
5911 "not every choice streamed: {order:?}"
5912 );
5913 let last_of_zero = order
5914 .iter()
5915 .rposition(|i| *i == 0)
5916 .expect("choice 0 streamed");
5917 let first_of_two = order
5918 .iter()
5919 .position(|i| *i == 2)
5920 .expect("choice 2 streamed");
5921 assert!(
5922 first_of_two < last_of_zero,
5923 "the choices arrived one after another rather than interleaved: {order:?}"
5924 );
5925 }
5926
5927 #[tokio::test]
5934 async fn logit_bias_moves_the_draw() {
5935 let app = streaming_test_app();
5936 let ask = |bias: Option<serde_json::Value>| {
5937 let mut b = serde_json::json!({
5938 "model": "x",
5939 "prompt": "hi",
5940 "max_tokens": 8,
5941 "temperature": 0
5942 });
5943 if let Some(v) = bias {
5944 b["logit_bias"] = v;
5945 }
5946 b
5947 };
5948
5949 let (status, plain) =
5950 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(None)).await;
5951 assert_eq!(status, StatusCode::OK, "{plain}");
5952 let free = plain["choices"][0]["text"].as_str().unwrap_or_default();
5953
5954 let (status, biased) = post_json_uri(
5956 &app,
5957 frink_api::routes::V1_COMPLETIONS,
5958 ask(Some(serde_json::json!({"65": 100.0}))),
5959 )
5960 .await;
5961 assert_eq!(status, StatusCode::OK, "{biased}");
5962 let text = biased["choices"][0]["text"].as_str().expect("text");
5963 assert!(!text.is_empty(), "nothing was generated: {biased}");
5964 assert!(
5965 text.chars().all(|c| c == 'A'),
5966 "the bias did not reach the sampler: {text:?}"
5967 );
5968 assert!(
5969 !free.chars().all(|c| c == 'A'),
5970 "the unbiased answer was already all As, so this proved nothing: {free:?}"
5971 );
5972
5973 let (status, chat) = post_json_uri(
5976 &app,
5977 frink_api::routes::V1_CHAT_COMPLETIONS,
5978 serde_json::json!({
5979 "model": "x",
5980 "messages": [{"role": "user", "content": "hi"}],
5981 "max_tokens": 8,
5982 "temperature": 0,
5983 "logit_bias": {"65": 100.0}
5984 }),
5985 )
5986 .await;
5987 assert_eq!(status, StatusCode::OK, "{chat}");
5988 let content = chat["choices"][0]["message"]["content"]
5989 .as_str()
5990 .unwrap_or_default();
5991 assert!(
5992 !content.is_empty() && content.chars().all(|c| c == 'A'),
5993 "the chat wire ignored the bias: {content:?}"
5994 );
5995 }
5996
5997 #[tokio::test]
6005 async fn a_bias_cannot_beat_allowed_token_ids() {
6006 let app = streaming_test_app();
6007 let (status, body) = post_json_uri(
6008 &app,
6009 frink_api::routes::V1_COMPLETIONS,
6010 serde_json::json!({
6011 "model": "x",
6012 "prompt": "hi",
6013 "max_tokens": 8,
6014 "temperature": 0,
6015 "logit_bias": {"65": 100.0},
6017 "allowed_token_ids": [66, 67]
6018 }),
6019 )
6020 .await;
6021 assert_eq!(status, StatusCode::OK, "{body}");
6022 let text = body["choices"][0]["text"].as_str().expect("text");
6023 assert!(!text.is_empty(), "nothing was generated: {body}");
6024 assert!(
6025 !text.contains('A'),
6026 "a bias produced a token the constraint forbade: {text:?}"
6027 );
6028 assert!(
6029 text.chars().all(|c| c == 'B' || c == 'C'),
6030 "the allowed set was not honoured: {text:?}"
6031 );
6032 }
6033
6034 #[tokio::test]
6037 async fn a_bias_outside_the_range_is_a_bad_request() {
6038 let app = streaming_test_app();
6039 let (status, body) = post_json_uri(
6040 &app,
6041 frink_api::routes::V1_COMPLETIONS,
6042 serde_json::json!({
6043 "model": "x",
6044 "prompt": "hi",
6045 "max_tokens": 2,
6046 "logit_bias": {"65": 1000.0}
6047 }),
6048 )
6049 .await;
6050 assert_eq!(status, StatusCode::BAD_REQUEST, "{body}");
6051 assert!(
6052 body["error"]["message"]
6053 .as_str()
6054 .unwrap_or_default()
6055 .contains("logit_bias"),
6056 "{body}"
6057 );
6058 }
6059
6060 #[tokio::test]
6072 async fn skip_special_tokens_false_keeps_the_end_marker() {
6073 let probe = test_app_with_state(Arc::new(test_state(
6080 test_byte_model(None, false),
6081 ResponseCache::new(1000, Duration::from_secs(3600)),
6082 )));
6083 let (_, seen) = post_json_uri(
6084 &probe,
6085 frink_api::routes::V1_COMPLETIONS,
6086 serde_json::json!({
6087 "model": "x",
6088 "prompt": "\u{1}\u{2}",
6089 "max_tokens": 6,
6090 "temperature": 0,
6091 "logprobs": 1,
6092 "return_tokens_as_token_ids": true
6093 }),
6094 )
6095 .await;
6096 let eos: usize = seen["choices"][0]["logprobs"]["tokens"][1]
6097 .as_str()
6098 .and_then(|s| s.strip_prefix("token_id:"))
6099 .and_then(|s| s.parse().ok())
6100 .expect("a second generated token");
6101
6102 let app = test_app_with_state(Arc::new(test_state(
6103 test_byte_model(Some(eos), false),
6104 ResponseCache::new(1000, Duration::from_secs(3600)),
6105 )));
6106 let ask = |skip: bool| {
6107 serde_json::json!({
6108 "model": "x",
6109 "prompt": "\u{1}\u{2}",
6110 "max_tokens": 6,
6111 "temperature": 0,
6112 "skip_special_tokens": skip
6113 })
6114 };
6115
6116 let (status, kept) =
6117 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(false)).await;
6118 assert_eq!(status, StatusCode::OK, "{kept}");
6119 let (status, skipped) =
6120 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(true)).await;
6121 assert_eq!(status, StatusCode::OK, "{skipped}");
6122
6123 assert_eq!(
6126 kept["choices"][0]["finish_reason"], "stop",
6127 "the model did not end the turn: {kept}"
6128 );
6129 assert_eq!(
6130 skipped["choices"][0]["finish_reason"], "stop",
6131 "keeping the marker must not change WHEN it stops: {skipped}"
6132 );
6133
6134 let with = kept["choices"][0]["text"].as_str().expect("text");
6135 let without = skipped["choices"][0]["text"].as_str().expect("text");
6136 let marker = char::from(eos as u8);
6138 assert!(
6139 with.ends_with(marker),
6140 "the end marker was dropped: {with:?}"
6141 );
6142 assert!(
6143 !without.ends_with(marker),
6144 "the default must still skip it: {without:?}"
6145 );
6146 assert_eq!(
6147 with.len(),
6148 without.len() + marker.len_utf8(),
6149 "the two answers differ by more than the marker"
6150 );
6151 assert_eq!(
6154 kept["usage"]["completion_tokens"].as_u64().unwrap(),
6155 skipped["usage"]["completion_tokens"].as_u64().unwrap() + 1,
6156 "the kept marker was not counted"
6157 );
6158 }
6159
6160 #[tokio::test]
6166 async fn return_tokens_as_token_ids_renames_reported_tokens_only() {
6167 let app = streaming_test_app();
6168 let ask = |as_ids: bool| {
6169 serde_json::json!({
6170 "model": "x",
6171 "prompt": "hi",
6172 "max_tokens": 4,
6173 "temperature": 0,
6174 "logprobs": 2,
6175 "return_tokens_as_token_ids": as_ids
6176 })
6177 };
6178
6179 let (status, plain) =
6180 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(false)).await;
6181 assert_eq!(status, StatusCode::OK, "{plain}");
6182 let (status, by_id) =
6183 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(true)).await;
6184 assert_eq!(status, StatusCode::OK, "{by_id}");
6185
6186 let tokens = by_id["choices"][0]["logprobs"]["tokens"]
6187 .as_array()
6188 .expect("tokens");
6189 assert!(!tokens.is_empty(), "nothing was reported: {by_id}");
6190 for t in tokens {
6191 let s = t.as_str().expect("a piece");
6192 assert!(
6193 s.starts_with("token_id:") && s["token_id:".len()..].parse::<usize>().is_ok(),
6194 "reported as text rather than by id: {s:?}"
6195 );
6196 }
6197 let top = &by_id["choices"][0]["logprobs"]["top_logprobs"][0];
6201 for key in top.as_object().expect("a map").keys() {
6202 assert!(key.starts_with("token_id:"), "{key:?}");
6203 }
6204 assert_eq!(
6206 by_id["choices"][0]["text"], plain["choices"][0]["text"],
6207 "the completion's text changed, which the field does not do"
6208 );
6209 assert!(
6210 !plain["choices"][0]["logprobs"]["tokens"][0]
6211 .as_str()
6212 .unwrap_or_default()
6213 .starts_with("token_id:"),
6214 "the default already reported ids, so this proved nothing"
6215 );
6216 }
6217
6218 #[tokio::test]
6227 async fn echo_returns_the_prompt_with_offsets_that_index_it() {
6228 let app = streaming_test_app();
6229 let prompt = "hello";
6230 let (status, body) = post_json_uri(
6231 &app,
6232 frink_api::routes::V1_COMPLETIONS,
6233 serde_json::json!({
6234 "model": "x",
6235 "prompt": prompt,
6236 "max_tokens": 6,
6237 "temperature": 0,
6238 "echo": true,
6239 "logprobs": 2
6240 }),
6241 )
6242 .await;
6243 assert_eq!(status, StatusCode::OK, "{body}");
6244
6245 let text = body["choices"][0]["text"].as_str().expect("text");
6246 assert!(
6247 text.starts_with(prompt),
6248 "the prompt was not echoed: {text:?}"
6249 );
6250 assert!(
6251 text.len() > prompt.len(),
6252 "nothing was generated after the echo: {text:?}"
6253 );
6254
6255 let lp = &body["choices"][0]["logprobs"];
6256 let tokens = lp["tokens"].as_array().expect("tokens");
6257 let offsets = lp["text_offset"].as_array().expect("text_offset");
6258 let scores = lp["token_logprobs"].as_array().expect("token_logprobs");
6259 assert_eq!(tokens.len(), offsets.len());
6260 assert_eq!(tokens.len(), scores.len());
6261 assert!(
6262 tokens.len() > 6,
6263 "the arrays cover only the completion: {}",
6264 tokens.len()
6265 );
6266 assert!(scores[0].is_null(), "{lp}");
6268 for (i, (tok, off)) in tokens.iter().zip(offsets).enumerate() {
6270 let (piece, at) = (
6271 tok.as_str().expect("a piece"),
6272 off.as_u64().unwrap() as usize,
6273 );
6274 assert!(
6275 text[at..].starts_with(piece),
6276 "entry {i}: offset {at} does not start {piece:?} in {text:?}"
6277 );
6278 }
6279 }
6280
6281 #[tokio::test]
6291 async fn truncate_prompt_tokens_keeps_the_last_k_and_echo_reports_them() {
6292 let app = streaming_test_app();
6293 let prompt = "abcdefghij";
6294 let ask = |k: Option<u32>| {
6295 let mut b = serde_json::json!({
6296 "model": "x",
6297 "prompt": prompt,
6298 "max_tokens": 2,
6299 "temperature": 0,
6300 "echo": true
6301 });
6302 if let Some(k) = k {
6303 b["truncate_prompt_tokens"] = serde_json::json!(k);
6304 }
6305 b
6306 };
6307
6308 let (status, full) =
6309 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(None)).await;
6310 assert_eq!(status, StatusCode::OK, "{full}");
6311 let full_prompt_tokens = full["usage"]["prompt_tokens"].as_u64().expect("usage");
6312 assert!(full_prompt_tokens > 4, "the prompt is too short to cut");
6313
6314 let (status, cut) =
6315 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(Some(4))).await;
6316 assert_eq!(status, StatusCode::OK, "{cut}");
6317 assert_eq!(
6318 cut["usage"]["prompt_tokens"].as_u64(),
6319 Some(4),
6320 "the prompt was not truncated: {cut}"
6321 );
6322 let text = cut["choices"][0]["text"].as_str().expect("text");
6324 assert!(
6325 text.starts_with("ghij"),
6326 "echo reported a prompt the model never saw: {text:?}"
6327 );
6328 assert!(
6329 !text.starts_with(prompt),
6330 "the full prompt was echoed after a truncation: {text:?}"
6331 );
6332 }
6333
6334 #[tokio::test]
6338 async fn a_truncation_below_one_is_a_bad_request() {
6339 let app = streaming_test_app();
6340 for k in [0i64, -1] {
6341 let (status, body) = post_json_uri(
6342 &app,
6343 frink_api::routes::V1_COMPLETIONS,
6344 serde_json::json!({
6345 "model": "x",
6346 "prompt": "hi",
6347 "max_tokens": 2,
6348 "truncate_prompt_tokens": k
6349 }),
6350 )
6351 .await;
6352 assert_eq!(status, StatusCode::BAD_REQUEST, "k = {k}: {body}");
6353 }
6354 }
6355
6356 #[tokio::test]
6364 async fn allowed_token_ids_restricts_the_draw() {
6365 let app = streaming_test_app();
6366 let body = |allowed: Option<serde_json::Value>| {
6367 let mut b = serde_json::json!({
6368 "model": "x",
6369 "prompt": "hi",
6370 "max_tokens": 16,
6371 "temperature": 1.0,
6372 "seed": 3
6373 });
6374 if let Some(ids) = allowed {
6375 b["allowed_token_ids"] = ids;
6376 }
6377 b
6378 };
6379
6380 let (status, free) =
6383 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, body(None)).await;
6384 assert_eq!(status, StatusCode::OK, "{free}");
6385 let free_text = free["choices"][0]["text"].as_str().unwrap_or_default();
6386
6387 let (status, restricted) = post_json_uri(
6388 &app,
6389 frink_api::routes::V1_COMPLETIONS,
6390 body(Some(serde_json::json!([65, 66]))),
6392 )
6393 .await;
6394 assert_eq!(status, StatusCode::OK, "{restricted}");
6395 let text = restricted["choices"][0]["text"]
6396 .as_str()
6397 .unwrap_or_default();
6398 assert!(!text.is_empty(), "nothing was generated: {restricted}");
6399 assert!(
6400 text.chars().all(|c| c == 'A' || c == 'B'),
6401 "a token outside `allowed_token_ids` was drawn: {text:?}"
6402 );
6403 assert!(
6406 !free_text.chars().all(|c| c == 'A' || c == 'B'),
6407 "the unrestricted answer was already inside the allowed set: {free_text:?}"
6408 );
6409 }
6410
6411 #[tokio::test]
6418 async fn an_empty_allowed_token_ids_is_a_bad_request() {
6419 let app = streaming_test_app();
6420 let (status, body) = post_json_uri(
6421 &app,
6422 frink_api::routes::V1_COMPLETIONS,
6423 serde_json::json!({
6424 "model": "x",
6425 "prompt": "hi",
6426 "max_tokens": 4,
6427 "allowed_token_ids": []
6428 }),
6429 )
6430 .await;
6431 assert_eq!(status, StatusCode::BAD_REQUEST, "{body}");
6432 assert!(
6433 body["error"]["message"]
6434 .as_str()
6435 .unwrap_or_default()
6436 .contains("allowed_token_ids"),
6437 "{body}"
6438 );
6439 }
6440
6441 #[tokio::test]
6447 async fn bad_words_removes_a_token_without_ending_the_generation() {
6448 let app = streaming_test_app();
6449 let ask = |bad: Option<serde_json::Value>| {
6450 let mut b = serde_json::json!({
6451 "model": "x",
6452 "prompt": "hi",
6453 "max_tokens": 24,
6454 "temperature": 1.0,
6455 "seed": 11
6456 });
6457 if let Some(words) = bad {
6458 b["bad_words"] = words;
6459 }
6460 b
6461 };
6462
6463 let (status, free) =
6464 post_json_uri(&app, frink_api::routes::V1_COMPLETIONS, ask(None)).await;
6465 assert_eq!(status, StatusCode::OK, "{free}");
6466 let free_text = free["choices"][0]["text"]
6467 .as_str()
6468 .unwrap_or_default()
6469 .to_string();
6470 let target = free_text
6473 .chars()
6474 .find(|c| c.is_ascii() && !c.is_control())
6475 .expect("the model produced some ascii");
6476
6477 let (status, steered) = post_json_uri(
6478 &app,
6479 frink_api::routes::V1_COMPLETIONS,
6480 ask(Some(serde_json::json!([target.to_string()]))),
6481 )
6482 .await;
6483 assert_eq!(status, StatusCode::OK, "{steered}");
6484 let text = steered["choices"][0]["text"].as_str().unwrap_or_default();
6485 assert!(
6486 !text.contains(target),
6487 "the forbidden {target:?} came back anyway: {text:?}"
6488 );
6489 assert_eq!(
6492 steered["usage"]["completion_tokens"], free["usage"]["completion_tokens"],
6493 "the generation ended early, so `bad_words` acted like `stop`: {steered}"
6494 );
6495 }
6496
6497 #[tokio::test]
6507 async fn every_route_refuses_the_same_unimplemented_fields() {
6508 let app = test_app();
6509 let fields = [
6510 ("n", serde_json::json!(3)),
6511 ("best_of", serde_json::json!(2)),
6512 ("prompt_logprobs", serde_json::json!(1)),
6513 ("echo", serde_json::json!(true)),
6514 ("use_beam_search", serde_json::json!(true)),
6515 ("truncate_prompt_tokens", serde_json::json!(8)),
6516 ("prompt_embeds", serde_json::json!("AA==")),
6517 ("skip_special_tokens", serde_json::json!(false)),
6518 ("return_tokens_as_token_ids", serde_json::json!(true)),
6519 ];
6520 for (field, value) in fields {
6521 for (uri, base) in [
6522 (
6523 frink_api::routes::V1_CHAT_COMPLETIONS,
6524 serde_json::json!({
6525 "model": "x",
6526 "messages": [{"role": "user", "content": "hi"}],
6527 "max_tokens": 2
6528 }),
6529 ),
6530 (
6531 frink_api::routes::V1_COMPLETIONS,
6532 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
6533 ),
6534 (
6535 frink_api::routes::COMPLETION,
6536 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6537 ),
6538 ] {
6539 let mut body = base;
6540 body[field] = value.clone();
6541 if field == "prompt_logprobs" && uri == frink_api::routes::V1_COMPLETIONS {
6548 let (status, answer) = post_json_uri(&app, uri, body).await;
6549 assert_eq!(status, StatusCode::OK, "{uri} refused it: {answer}");
6550 assert!(
6551 answer["prompt_logprobs"].is_array(),
6552 "served without the field: {answer}"
6553 );
6554 continue;
6555 }
6556 if field == "echo" && uri == frink_api::routes::V1_COMPLETIONS {
6560 let (status, answer) = post_json_uri(&app, uri, body).await;
6561 assert_eq!(status, StatusCode::OK, "{uri} refused `echo`: {answer}");
6562 assert!(
6563 answer["choices"][0]["text"]
6564 .as_str()
6565 .unwrap_or_default()
6566 .starts_with("hi"),
6567 "served without echoing the prompt: {answer}"
6568 );
6569 continue;
6570 }
6571 if field == "skip_special_tokens" || field == "return_tokens_as_token_ids" {
6575 let (status, answer) = post_json_uri(&app, uri, body).await;
6576 assert_eq!(status, StatusCode::OK, "{uri} refused `{field}`: {answer}");
6577 continue;
6578 }
6579 if field == "truncate_prompt_tokens" {
6582 let (status, answer) = post_json_uri(&app, uri, body).await;
6583 assert_eq!(
6584 status,
6585 StatusCode::OK,
6586 "{uri} refused `truncate_prompt_tokens`: {answer}"
6587 );
6588 continue;
6589 }
6590 if (field == "n" || field == "best_of")
6591 && (uri == frink_api::routes::V1_COMPLETIONS
6592 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
6593 {
6594 let (status, answer) = post_json_uri(&app, uri, body).await;
6595 assert_eq!(
6596 status,
6597 StatusCode::OK,
6598 "{uri} refused a served `{field}`: {answer}"
6599 );
6600 let want = if field == "n" { 3 } else { 1 };
6604 assert_eq!(
6605 answer["choices"].as_array().map(Vec::len),
6606 Some(want),
6607 "{field}: {answer}"
6608 );
6609 continue;
6610 }
6611 let (status, answer) = post_json_uri(&app, uri, body).await;
6612 assert_eq!(
6613 status,
6614 StatusCode::NOT_IMPLEMENTED,
6615 "{uri} served `{field}` instead of refusing it: {answer}"
6616 );
6617 assert!(
6618 answer["error"]["message"]
6619 .as_str()
6620 .is_some_and(|m| m.contains(field)),
6621 "{uri} refused `{field}` without naming it: {answer}"
6622 );
6623 }
6624 }
6625 }
6626
6627 #[tokio::test]
6628 async fn the_native_completion_wire_is_not_the_openai_one() {
6629 let app = test_app();
6630
6631 let (status, native) = post_json_uri(
6632 &app,
6633 frink_api::routes::COMPLETION,
6634 serde_json::json!({"prompt": "hi", "n_predict": 4}),
6635 )
6636 .await;
6637 assert_eq!(status, StatusCode::OK, "{native}");
6638 assert!(native["content"].is_string(), "{native}");
6639 assert_eq!(native["stop"], true);
6640 assert_eq!(native["stop_type"], "limit");
6641 assert_eq!(native["stopping_word"], "");
6642 assert_eq!(native["truncated"], false);
6643 assert_eq!(native["id_slot"], -1);
6644 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
6645 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
6646 assert!(
6647 native.get("choices").is_none(),
6648 "the native shape has no `choices`: {native}"
6649 );
6650
6651 let (status, openai) = post_json_uri(
6652 &app,
6653 frink_api::routes::V1_COMPLETIONS,
6654 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
6655 )
6656 .await;
6657 assert_eq!(status, StatusCode::OK);
6658 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
6659 assert!(
6660 openai.get("content").is_none(),
6661 "the OpenAI shape has no top-level `content`: {openai}"
6662 );
6663 }
6664
6665 #[tokio::test]
6669 async fn both_native_spellings_reach_the_same_handler() {
6670 let app = test_app();
6671 for route in [
6672 frink_api::routes::COMPLETION,
6673 frink_api::routes::COMPLETIONS,
6674 ] {
6675 let (status, body) = post_json_uri(
6676 &app,
6677 route,
6678 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
6679 )
6680 .await;
6681 assert_eq!(status, StatusCode::OK, "{route}: {body}");
6682 assert_eq!(body["stop"], true, "{route}");
6683 assert!(body["content"].is_string(), "{route}");
6684 }
6685
6686 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6689 let routes: Vec<&str> = stats["recent"]
6690 .as_array()
6691 .unwrap()
6692 .iter()
6693 .map(|row| row["route"].as_str().unwrap())
6694 .collect();
6695 assert!(
6696 routes.contains(&frink_api::routes::COMPLETION),
6697 "{routes:?}"
6698 );
6699 assert!(
6700 routes.contains(&frink_api::routes::COMPLETIONS),
6701 "{routes:?}"
6702 );
6703 }
6704
6705 #[tokio::test]
6711 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
6712 let app = streaming_test_app();
6713 let raw = post_sse_raw_uri(
6714 &app,
6715 frink_api::routes::COMPLETION,
6716 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
6717 )
6718 .await;
6719
6720 assert!(
6721 !raw.contains("[DONE]"),
6722 "llama.cpp's native stream has no sentinel: {raw}"
6723 );
6724 let frames: Vec<serde_json::Value> = raw
6725 .lines()
6726 .filter_map(|line| line.strip_prefix("data: "))
6727 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
6728 .collect();
6729 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
6730
6731 let (last, partials) = frames.split_last().unwrap();
6732 assert_eq!(last["stop"], true, "the last frame closes the stream");
6733 assert!(last["timings"].is_object(), "{last}");
6734 assert!(last["stop_type"].is_string(), "{last}");
6735 for partial in partials {
6736 assert_eq!(partial["stop"], false, "{partial}");
6737 assert!(partial["content"].is_string(), "{partial}");
6738 assert!(partial.get("timings").is_none(), "{partial}");
6742 assert!(partial.get("generation_settings").is_none(), "{partial}");
6743 }
6744 let streamed: String = partials
6747 .iter()
6748 .filter_map(|p| p["content"].as_str())
6749 .collect();
6750 assert_eq!(last["content"].as_str().unwrap(), streamed);
6751 }
6752
6753 #[tokio::test]
6758 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
6759 let app = test_app();
6760 for body in [
6761 serde_json::json!({"prompt": "hi"}),
6762 serde_json::json!({"prompt": "hi", "n_predict": -1}),
6763 ] {
6764 let (status, refusal) =
6765 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
6766 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
6767 assert!(
6768 refusal["error"]["message"]
6769 .as_str()
6770 .unwrap()
6771 .contains("n_predict"),
6772 "{refusal}"
6773 );
6774 }
6775 let (status, _) = post_json_uri(
6778 &app,
6779 frink_api::routes::COMPLETION,
6780 serde_json::json!({"prompt": "hi", "n_predict": 2}),
6781 )
6782 .await;
6783 assert_eq!(status, StatusCode::OK);
6784 }
6785
6786 #[tokio::test]
6795 async fn a_stop_string_halts_the_answer_and_is_named_back() {
6796 let app = streaming_test_app();
6797 let ask = |stop: serde_json::Value| {
6798 let app = app.clone();
6799 async move {
6800 post_json_uri(
6801 &app,
6802 frink_api::routes::COMPLETION,
6803 serde_json::json!({
6804 "prompt": "hi",
6805 "n_predict": 64,
6806 "ignore_eos": true,
6807 "stop": stop,
6808 }),
6809 )
6810 .await
6811 .1
6812 }
6813 };
6814
6815 let baseline = ask(serde_json::json!([])).await;
6816 assert_eq!(baseline["stop_type"], "limit");
6817 assert_eq!(baseline["stopping_word"], "");
6818 let text = baseline["content"].as_str().unwrap().to_string();
6819 let sentinel: String = text.chars().skip(1).take(2).collect();
6825 assert_eq!(
6826 sentinel.chars().count(),
6827 2,
6828 "the fixture must produce enough output to cut: {text:?}"
6829 );
6830 let cut = text.find(&sentinel).expect("it came out of this text");
6831
6832 let stopped = ask(serde_json::json!([sentinel])).await;
6833 assert_eq!(stopped["stop_type"], "word", "{stopped}");
6834 assert_eq!(stopped["stopping_word"], sentinel);
6835 assert_eq!(
6836 stopped["content"].as_str().unwrap(),
6837 &text[..cut],
6838 "the answer must be cut at the sentinel, not run past it"
6839 );
6840 }
6841
6842 #[tokio::test]
6848 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
6849 let app = test_app();
6850
6851 let (v1_status, v1) = post_json_uri(
6852 &app,
6853 frink_api::routes::V1_TOKENIZE,
6854 serde_json::json!({"prompt": "hello"}),
6855 )
6856 .await;
6857 let (alias_status, alias) = post_json_uri(
6858 &app,
6859 frink_api::routes::TOKENIZE,
6860 serde_json::json!({"content": "hello"}),
6861 )
6862 .await;
6863 assert_eq!(v1_status, StatusCode::OK);
6864 assert_eq!(alias_status, StatusCode::OK, "{alias}");
6865 assert_eq!(v1["tokens"], alias["tokens"]);
6866 assert!(!alias["tokens"].as_array().unwrap().is_empty());
6867
6868 let (status, both_ways) = post_json_uri(
6871 &app,
6872 frink_api::routes::TOKENIZE,
6873 serde_json::json!({"prompt": "hello"}),
6874 )
6875 .await;
6876 assert_eq!(status, StatusCode::OK);
6877 assert_eq!(both_ways["tokens"], v1["tokens"]);
6878 }
6879
6880 #[tokio::test]
6885 async fn detokenize_answers_under_both_dialects_keys() {
6886 let app = test_app();
6887 for route in [
6888 frink_api::routes::DETOKENIZE,
6889 frink_api::routes::V1_DETOKENIZE,
6890 ] {
6891 let (status, body) =
6892 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
6893 assert_eq!(status, StatusCode::OK, "{route}");
6894 assert_eq!(body["text"], "hi", "{route}");
6895 assert_eq!(body["content"], body["text"], "{route}");
6896 }
6897 }
6898
6899 #[tokio::test]
6903 async fn the_alias_is_recorded_under_the_path_the_client_called() {
6904 let app = test_app();
6905 let (status, _) = post_json_uri(
6906 &app,
6907 frink_api::routes::TOKENIZE,
6908 serde_json::json!({"content": "hello"}),
6909 )
6910 .await;
6911 assert_eq!(status, StatusCode::OK);
6912
6913 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6914 let routes: Vec<&str> = stats["recent"]
6915 .as_array()
6916 .unwrap()
6917 .iter()
6918 .map(|row| row["route"].as_str().unwrap())
6919 .collect();
6920 assert!(
6921 routes.contains(&frink_api::routes::TOKENIZE),
6922 "the alias must be its own row: {routes:?}"
6923 );
6924 assert!(
6925 !routes.contains(&frink_api::routes::V1_TOKENIZE),
6926 "nothing called /v1/tokenize: {routes:?}"
6927 );
6928 }
6929
6930 #[tokio::test]
6935 async fn add_special_prepends_the_same_bos_the_decoder_would() {
6936 let mut cfg = test_dense_fixture();
6937 cfg.vocab_size = 256;
6938 let model = Model::Gguf(GgufModel {
6939 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
6940 tokenizer: Arc::new(ServerTokenizer::Byte),
6941 stop_tokens: StopTokens::default(),
6942 bos_id: Some(7),
6943 is_synthetic: true,
6944 chat_template: chat_template::PromptTemplate::plain(),
6945 });
6946 let app = test_app_with_state(Arc::new(test_state(
6947 model,
6948 ResponseCache::new(1000, Duration::from_secs(3600)),
6949 )));
6950
6951 let (_, plain) = post_json_uri(
6952 &app,
6953 frink_api::routes::TOKENIZE,
6954 serde_json::json!({"content": "hi"}),
6955 )
6956 .await;
6957 let (_, special) = post_json_uri(
6958 &app,
6959 frink_api::routes::TOKENIZE,
6960 serde_json::json!({"content": "hi", "add_special": true}),
6961 )
6962 .await;
6963
6964 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6965 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6966 assert_eq!(special["count"], 3);
6967 }
6968
6969 #[tokio::test]
6972 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6973 let app = test_app();
6974 let (status, _) = post_json_uri(
6975 &app,
6976 frink_api::routes::V1_EMBEDDINGS,
6977 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6978 )
6979 .await;
6980 assert_eq!(status, StatusCode::BAD_REQUEST);
6981
6982 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6983 let recent = stats["recent"].as_array().unwrap();
6984 assert_eq!(recent.len(), 1);
6985 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
6986 assert_eq!(recent[0]["status"], 400);
6987 assert_eq!(
6988 recent[0]["prompt_tokens"], 0,
6989 "a rejected call embedded nothing"
6990 );
6991 }
6992
6993 #[tokio::test]
6996 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6997 let app = test_app();
6998 let key = "sk-monitor-secret";
6999 let (status, _) = post_json_with_headers(
7000 &app,
7001 "/v1/chat/completions",
7002 serde_json::json!({
7003 "model": "x",
7004 "messages": [{"role": "user", "content": "hi"}],
7005 "max_tokens": 2
7006 }),
7007 &[
7008 ("authorization", &format!("Bearer {key}")),
7009 ("x-frink-client", "frink-studio"),
7010 ],
7011 )
7012 .await;
7013 assert_eq!(status, StatusCode::OK);
7014
7015 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7016 let row = stats["recent"].as_array().unwrap()[0].clone();
7017 let fingerprint = row["via_api_key"]
7018 .as_str()
7019 .expect("the row names the key that served it")
7020 .to_string();
7021 assert_eq!(fingerprint, attribution::key_fingerprint(key));
7022 assert!(!fingerprint.contains(key));
7023 assert!(
7024 !serde_json::to_string(&stats).unwrap().contains(key),
7025 "the stats payload must not carry the key in any form"
7026 );
7027 assert_eq!(row["client"], "frink-studio");
7028 }
7029
7030 #[tokio::test]
7033 async fn different_keys_are_different_callers_and_no_key_is_null() {
7034 let app = test_app();
7035 let body = serde_json::json!({
7036 "model": "x",
7037 "messages": [{"role": "user", "content": "hi"}],
7038 "max_tokens": 1
7039 });
7040 for headers in [
7041 vec![("authorization", "Bearer key-one")],
7042 vec![("authorization", "Bearer key-two")],
7043 vec![],
7044 ] {
7045 let (status, _) =
7046 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
7047 assert_eq!(status, StatusCode::OK);
7048 }
7049
7050 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7051 let recent = stats["recent"].as_array().unwrap();
7052 assert_eq!(recent.len(), 3);
7053 let one = recent[0]["via_api_key"].as_str().unwrap();
7054 let two = recent[1]["via_api_key"].as_str().unwrap();
7055 assert_ne!(one, two, "two keys must not collapse into one caller");
7056 assert!(
7057 recent[2]["via_api_key"].is_null(),
7058 "an unauthenticated call is null, not a fingerprint of nothing"
7059 );
7060 assert!(recent[2]["client"].is_null());
7061 }
7062
7063 #[tokio::test]
7068 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
7069 let state = Arc::new(test_state(
7070 named_test_model("really-loaded", 256),
7071 ResponseCache::new(4, Duration::from_secs(60)),
7072 ));
7073 let app = test_app_with_state(Arc::clone(&state));
7074
7075 let (status, _) = post_json_uri(
7076 &app,
7077 "/v1/chat/completions",
7078 serde_json::json!({
7079 "model": "gpt-4-turbo-that-is-not-here",
7080 "messages": [{"role": "user", "content": "hi"}],
7081 "max_tokens": 2
7082 }),
7083 )
7084 .await;
7085 assert_eq!(status, StatusCode::OK);
7086
7087 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7088 assert_eq!(stats["recent"][0]["model"], "really-loaded");
7089
7090 state.swap_active(None);
7093 let (status, _) = post_json_uri(
7094 &app,
7095 "/v1/chat/completions",
7096 serde_json::json!({
7097 "model": "gpt-4-turbo-that-is-not-here",
7098 "messages": [{"role": "user", "content": "hi"}]
7099 }),
7100 )
7101 .await;
7102 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7103 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7104 let recent = stats["recent"].as_array().unwrap();
7105 assert!(recent[recent.len() - 1]["model"].is_null());
7106 }
7107
7108 #[tokio::test]
7112 async fn a_streamed_row_names_the_model_it_decoded_against() {
7113 let state = Arc::new(test_state(
7114 named_test_model("model-before", 256),
7115 ResponseCache::new(4, Duration::from_secs(60)),
7116 ));
7117 let app = test_app_with_state(Arc::clone(&state));
7118 let _ = post_sse_raw(&app, resumable_request()).await;
7119 active_model(&state, "model-after");
7121
7122 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7123 assert_eq!(stats["recent"][0]["model"], "model-before");
7124 }
7125
7126 #[tokio::test]
7129 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
7130 let app = test_app();
7131 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7132 assert_eq!(status, StatusCode::OK);
7133 assert!(
7134 stats["queue_depth"].is_null(),
7135 "without continuous batching nothing queues, so there is nothing to measure"
7136 );
7137 assert!(stats["queue_rejected_total"].is_null());
7138 assert_eq!(
7139 stats["generating_now"], 0,
7140 "work in progress is measured and really is zero here"
7141 );
7142 }
7143
7144 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
7149 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
7150 }
7151
7152 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
7156 use http_body_util::BodyExt;
7157 use tower::ServiceExt;
7158
7159 let response = app
7160 .clone()
7161 .oneshot(
7162 axum::http::Request::builder()
7163 .method("POST")
7164 .uri(uri)
7165 .header("content-type", "application/json")
7166 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7167 .unwrap(),
7168 )
7169 .await
7170 .unwrap();
7171 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7172 String::from_utf8(bytes.to_vec()).unwrap()
7173 }
7174
7175 async fn get_json_with_headers(
7176 app: &Router,
7177 uri: &str,
7178 headers: &[(&str, &str)],
7179 ) -> (StatusCode, serde_json::Value) {
7180 use http_body_util::BodyExt;
7181 use tower::ServiceExt;
7182
7183 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
7184 for (name, value) in headers {
7185 builder = builder.header(*name, *value);
7186 }
7187 let response = app
7188 .clone()
7189 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
7190 .await
7191 .unwrap();
7192 let status = response.status();
7193 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7194 (
7195 status,
7196 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
7197 )
7198 }
7199
7200 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
7201 body.lines()
7202 .filter_map(|line| line.strip_prefix(field))
7203 .map(str::trim)
7204 .collect()
7205 }
7206
7207 fn resumable_request() -> serde_json::Value {
7208 serde_json::json!({
7209 "model": "m",
7210 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7211 "max_tokens": 4,
7212 "temperature": 0,
7213 "stream": true,
7214 "stream_resumable": true,
7215 })
7216 }
7217
7218 #[tokio::test]
7223 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
7224 let app = test_app();
7225 let body = post_sse_raw(&app, resumable_request()).await;
7226
7227 let request_id = body
7228 .lines()
7229 .find_map(|l| l.strip_prefix("data: "))
7230 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
7231 .and_then(|v| v["request_id"].as_str().map(str::to_string))
7232 .expect("the first chunk names the request");
7233
7234 let ids = sse_field(&body, "id:");
7235 let datas = sse_field(&body, "data:");
7236 assert_eq!(
7237 ids.len(),
7238 datas.len(),
7239 "every event carries an id, or a reconnect cannot name where it stopped"
7240 );
7241 for (i, id) in ids.iter().enumerate() {
7242 assert_eq!(*id, format!("{request_id}:{i}"));
7243 }
7244 let retries = sse_field(&body, "retry:");
7245 assert_eq!(
7246 retries.len(),
7247 1,
7248 "the reconnect delay is stated once, not on every event"
7249 );
7250 assert_eq!(retries[0], "1500");
7251 assert!(
7252 body.contains("data: [DONE]"),
7253 "the end of stream is still stated"
7254 );
7255 }
7256
7257 #[tokio::test]
7261 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
7262 let app = test_app();
7263 let mut request = resumable_request();
7264 request["stream_resumable"] = serde_json::json!(false);
7265 let body = post_sse_raw(&app, request).await;
7266 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
7267 assert!(
7268 sse_field(&body, "id:").is_empty(),
7269 "an id promises a replay this stream cannot serve"
7270 );
7271 assert!(sse_field(&body, "retry:").is_empty());
7272 }
7273
7274 #[tokio::test]
7278 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
7279 let app = test_app();
7280 let body = post_sse_raw(&app, resumable_request()).await;
7281 let request_id = sse_field(&body, "id:")[0]
7282 .rsplit_once(':')
7283 .unwrap()
7284 .0
7285 .to_string();
7286 let streamed: Vec<String> = sse_field(&body, "data:")
7287 .iter()
7288 .map(|d| d.to_string())
7289 .collect();
7290
7291 let (status, polled) = get_json(
7292 &app,
7293 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
7294 )
7295 .await;
7296 assert_eq!(status, StatusCode::OK);
7297 let events: Vec<String> = polled["events"]
7298 .as_array()
7299 .unwrap()
7300 .iter()
7301 .map(|e| e["data"].as_str().unwrap().to_string())
7302 .collect();
7303 assert_eq!(
7304 events, streamed,
7305 "the fallback must deliver the same answer, not a re-run of it"
7306 );
7307 assert_eq!(polled["request_id"], request_id);
7308 assert_eq!(
7309 polled["done"], false,
7310 "events were still being handed out, so the client must ask again"
7311 );
7312
7313 let next = polled["next_index"].as_u64().unwrap();
7316 let (_, drained) = get_json(
7317 &app,
7318 &format!(
7319 "{}?from={next}",
7320 frink_api::routes::v1_stream_poll(&request_id)
7321 ),
7322 )
7323 .await;
7324 assert_eq!(drained["done"], true);
7325 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
7326 }
7327
7328 #[tokio::test]
7332 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
7333 let app = test_app();
7334 let body = post_sse_raw(&app, resumable_request()).await;
7335 let ids = sse_field(&body, "id:");
7336 let datas: Vec<String> = sse_field(&body, "data:")
7337 .iter()
7338 .map(|d| d.to_string())
7339 .collect();
7340 assert!(
7341 ids.len() >= 3,
7342 "need a few events to resume into the middle"
7343 );
7344 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
7345
7346 let (status, resumed) = get_json_with_headers(
7347 &app,
7348 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
7349 &[],
7350 )
7351 .await;
7352 assert_eq!(status, StatusCode::OK);
7353 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
7354
7355 let (_, tail) = get_json(
7357 &app,
7358 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
7359 )
7360 .await;
7361 let tail_events: Vec<String> = tail["events"]
7362 .as_array()
7363 .unwrap()
7364 .iter()
7365 .map(|e| e["data"].as_str().unwrap().to_string())
7366 .collect();
7367 assert_eq!(tail_events, datas[2..].to_vec());
7368 }
7369
7370 #[tokio::test]
7373 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
7374 use http_body_util::BodyExt;
7375 use tower::ServiceExt;
7376
7377 let app = test_app();
7378 let body = post_sse_raw(&app, resumable_request()).await;
7379 let ids = sse_field(&body, "id:");
7380 let datas: Vec<String> = sse_field(&body, "data:")
7381 .iter()
7382 .map(|d| d.to_string())
7383 .collect();
7384 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
7385
7386 let response = app
7387 .clone()
7388 .oneshot(
7389 axum::http::Request::builder()
7390 .method("GET")
7391 .uri(frink_api::routes::v1_stream(&request_id))
7392 .header("last-event-id", format!("{request_id}:0"))
7393 .body(axum::body::Body::empty())
7394 .unwrap(),
7395 )
7396 .await
7397 .unwrap();
7398 assert_eq!(response.status(), StatusCode::OK);
7399 assert_eq!(
7400 response
7401 .headers()
7402 .get("x-accel-buffering")
7403 .and_then(|v| v.to_str().ok()),
7404 Some("no"),
7405 "the reconnect needs the same anti-buffering header as the stream"
7406 );
7407 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7408 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
7409 assert_eq!(
7410 sse_field(&resumed, "data:")
7411 .iter()
7412 .map(|d| d.to_string())
7413 .collect::<Vec<_>>(),
7414 datas[1..].to_vec()
7415 );
7416 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
7417 }
7418
7419 #[tokio::test]
7423 async fn a_last_event_id_from_another_stream_is_refused() {
7424 let app = test_app();
7425 let body = post_sse_raw(&app, resumable_request()).await;
7426 let request_id = sse_field(&body, "id:")[0]
7427 .rsplit_once(':')
7428 .unwrap()
7429 .0
7430 .to_string();
7431
7432 let (status, err) = get_json_with_headers(
7433 &app,
7434 &frink_api::routes::v1_stream(&request_id),
7435 &[("last-event-id", "chatcmpl-someone-else:3")],
7436 )
7437 .await;
7438 assert_eq!(status, StatusCode::BAD_REQUEST);
7439 assert_eq!(err["error"]["code"], "bad_last_event_id");
7440 }
7441
7442 #[tokio::test]
7446 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
7447 let app = test_app();
7448 let mut request = resumable_request();
7449 request["stream_resumable"] = serde_json::json!(false);
7450 let body = post_sse_raw(&app, request).await;
7451 let request_id = body
7452 .lines()
7453 .find_map(|l| l.strip_prefix("data: "))
7454 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
7455 .and_then(|v| v["request_id"].as_str().map(str::to_string))
7456 .unwrap();
7457
7458 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
7459 assert_eq!(status, StatusCode::NOT_FOUND);
7460 assert_eq!(err["error"]["code"], "stream_not_found");
7461 assert!(err["error"]["message"]
7462 .as_str()
7463 .unwrap()
7464 .contains("stream_resumable"));
7465 }
7466
7467 #[test]
7471 fn the_axum_stream_patterns_match_the_published_templates() {
7472 assert_eq!(
7473 axum_path(frink_api::routes::V1_STREAM),
7474 "/v1/stream/:request_id"
7475 );
7476 assert_eq!(
7477 axum_path(frink_api::routes::V1_STREAM_POLL),
7478 "/v1/stream/:request_id/poll"
7479 );
7480 assert_eq!(
7481 frink_api::routes::v1_stream("abc"),
7482 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
7483 );
7484 }
7485
7486 #[test]
7496 fn no_published_template_reaches_the_router_with_its_braces() {
7497 for template in [
7498 frink_api::routes::V1_STREAM,
7499 frink_api::routes::V1_STREAM_POLL,
7500 frink_api::routes::V1_RESPONSE,
7501 frink_api::routes::V1_RESPONSE_CANCEL,
7502 frink_api::routes::ADMIN_TASK_CANCEL,
7503 ] {
7504 assert!(
7505 template.contains('{'),
7506 "{template} is in the template list but has no placeholder"
7507 );
7508 let mounted = axum_path(template);
7509 assert!(
7510 !mounted.contains('{') && !mounted.contains('}'),
7511 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
7512 );
7513 assert!(
7514 mounted.contains(':'),
7515 "{template} lost its placeholder entirely and would match one path only"
7516 );
7517 }
7518 }
7519
7520 #[tokio::test]
7528 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
7529 let app = test_app();
7530 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
7531 assert_eq!(status, StatusCode::NOT_FOUND);
7532 assert!(
7533 !body.is_null(),
7534 "empty body means axum never matched the route, so the id was read as a literal segment"
7535 );
7536 }
7537
7538 #[tokio::test]
7541 async fn the_task_list_starts_empty_rather_than_absent() {
7542 let app = test_app();
7543 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
7544 assert_eq!(status, StatusCode::OK);
7545 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
7546 }
7547
7548 #[tokio::test]
7557 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
7558 assert!(
7559 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
7560 "this test asserts the unconfigured behaviour"
7561 );
7562 let app = test_app();
7563 let (status, body) = post_json_uri(
7564 &app,
7565 &format!("{}?action=save", frink_api::routes::slots_id(0)),
7566 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
7567 )
7568 .await;
7569 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
7570 assert!(
7571 body["error"]["message"]
7572 .as_str()
7573 .unwrap()
7574 .contains("--slot-save-path"),
7575 "{body}"
7576 );
7577 }
7578
7579 pub(crate) async fn post_json_uri(
7580 app: &Router,
7581 uri: &str,
7582 body: serde_json::Value,
7583 ) -> (StatusCode, serde_json::Value) {
7584 use http_body_util::BodyExt;
7585 use tower::ServiceExt;
7586
7587 let response = app
7588 .clone()
7589 .oneshot(
7590 axum::http::Request::builder()
7591 .method("POST")
7592 .uri(uri)
7593 .header("content-type", "application/json")
7594 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7595 .unwrap(),
7596 )
7597 .await
7598 .unwrap();
7599 let status = response.status();
7600 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7601 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
7602 (status, json)
7603 }
7604
7605 pub(crate) async fn get_json_uri(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7608 use http_body_util::BodyExt;
7609 use tower::ServiceExt;
7610
7611 let response = app
7612 .clone()
7613 .oneshot(
7614 axum::http::Request::builder()
7615 .method("GET")
7616 .uri(uri)
7617 .body(axum::body::Body::empty())
7618 .unwrap(),
7619 )
7620 .await
7621 .unwrap();
7622 let status = response.status();
7623 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7624 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
7625 (status, json)
7626 }
7627
7628 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
7629 post_json_uri(app, "/v1/chat/completions", body).await.1
7630 }
7631
7632 #[tokio::test]
7641 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
7642 let app = test_app();
7643 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
7644 assert_eq!(status, StatusCode::OK);
7645
7646 let memory = &body["memory"];
7647 if memory.is_null() {
7648 return;
7651 }
7652 assert!(
7653 memory["bytes"].as_u64().is_some_and(|b| b > 0),
7654 "a read that produced a zero is a broken read, not an idle \
7655 engine: {memory}"
7656 );
7657 assert!(
7658 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
7659 "the quantity must travel with the number: {memory}"
7660 );
7661 }
7662
7663 #[tokio::test]
7669 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
7670 let app = test_app();
7671 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
7672 assert_eq!(status, StatusCode::OK);
7673 for pool in ["kv_pages", "window_slots", "state_slots"] {
7674 assert!(
7675 body["pools"][pool].is_null(),
7676 "{pool} must be null rather than a zero row: {}",
7677 body["pools"]
7678 );
7679 }
7680 }
7681
7682 #[tokio::test]
7690 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
7691 use http_body_util::BodyExt;
7692 use tower::ServiceExt;
7693
7694 let app = test_app();
7695 let send = |body: serde_json::Value| {
7696 let app = app.clone();
7697 async move {
7698 app.oneshot(
7699 axum::http::Request::builder()
7700 .method("POST")
7701 .uri(frink_api::routes::V1_MESSAGES)
7702 .header("content-type", "application/json")
7703 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7704 .unwrap(),
7705 )
7706 .await
7707 .unwrap()
7708 }
7709 };
7710
7711 let ok = send(serde_json::json!({
7712 "model": "test",
7713 "max_tokens": 1,
7714 "messages": [{"role": "user", "content": "hi"}],
7715 }))
7716 .await;
7717 assert_eq!(ok.status(), StatusCode::OK);
7718 let id = ok
7719 .headers()
7720 .get("request-id")
7721 .expect("a served message names its id")
7722 .to_str()
7723 .unwrap()
7724 .to_string();
7725 assert!(!id.is_empty());
7726
7727 let bad = send(serde_json::json!({"model": "test"})).await;
7730 assert!(bad.status().is_client_error());
7731 let other = bad.headers().get("request-id").expect("errors too");
7732 assert_ne!(other.to_str().unwrap(), id);
7733 let _ = bad.into_body().collect().await.unwrap();
7734 }
7735
7736 #[tokio::test]
7744 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
7745 let state = Arc::new(test_state(
7746 test_model_full_byte_vocab(),
7747 ResponseCache::new(1000, Duration::from_secs(3600)),
7748 ));
7749 let app = test_app_with_state(Arc::clone(&state));
7750 let body = serde_json::json!({
7751 "model": "test",
7752 "messages": [{"role": "user", "content": "hi"}],
7753 "max_tokens": 1,
7754 });
7755
7756 state
7757 .maintenance
7758 .lock()
7759 .unwrap()
7760 .begin_rebuild()
7761 .expect("a fresh server is serving, so the rebuild starts");
7762 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
7763 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
7764 assert_eq!(refused["error"]["type"], "cache_rebuilding");
7765
7766 state.maintenance.lock().unwrap().finish_rebuild(true);
7767 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
7768 assert_eq!(
7769 status,
7770 StatusCode::OK,
7771 "the gate reopens; a rebuild is not a latch"
7772 );
7773 }
7774
7775 #[tokio::test]
7780 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
7781 let app = test_app();
7782 let (status, body) = post_json_uri(
7783 &app,
7784 frink_api::routes::V1_CANCEL,
7785 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
7786 )
7787 .await;
7788 assert_eq!(status, StatusCode::NOT_FOUND);
7789 assert_eq!(body["cancelled"], serde_json::json!(false));
7790 assert_eq!(body["request_id"], "chatcmpl-never-issued");
7791 assert!(
7792 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
7793 "the verdict must carry a human reason: {body}"
7794 );
7795 }
7796
7797 #[tokio::test]
7801 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
7802 let state = Arc::new(test_state(
7803 test_model_full_byte_vocab(),
7804 ResponseCache::new(1000, Duration::from_secs(3600)),
7805 ));
7806 let app = test_app_with_state(Arc::clone(&state));
7807 let (token, _guard) = state.cancels.register("chatcmpl-live");
7808
7809 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
7810 assert_eq!(status, StatusCode::OK);
7811 assert_eq!(before["generating_now"], serde_json::json!(1));
7812
7813 let (status, body) = post_json_uri(
7814 &app,
7815 frink_api::routes::V1_CANCEL,
7816 serde_json::json!({ "request_id": "chatcmpl-live" }),
7817 )
7818 .await;
7819 assert_eq!(status, StatusCode::OK);
7820 assert_eq!(body["cancelled"], serde_json::json!(true));
7821 assert!(
7822 token.is_cancelled(),
7823 "the endpoint answered ok without setting the flag the decode loop reads"
7824 );
7825 }
7826
7827 #[tokio::test]
7828 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
7829 let app = test_app();
7830 let (status, tok) =
7831 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
7832 assert_eq!(status, StatusCode::OK);
7833 let tokens = tok["tokens"].as_array().unwrap();
7834 assert_eq!(tok["count"], tokens.len());
7835 assert!(!tokens.is_empty());
7836
7837 let (status, detok) = post_json_uri(
7838 &app,
7839 "/v1/detokenize",
7840 serde_json::json!({ "tokens": tokens }),
7841 )
7842 .await;
7843 assert_eq!(status, StatusCode::OK);
7844 assert_eq!(detok["text"], "Hi");
7845
7846 let (status, emb) = post_json_uri(
7847 &app,
7848 "/v1/embeddings",
7849 serde_json::json!({
7850 "input": "Hi",
7851 "embedding_type": "mean"
7852 }),
7853 )
7854 .await;
7855 assert_eq!(status, StatusCode::OK);
7856 let vec = emb["data"][0]["embedding"].as_array().unwrap();
7857 assert!(!vec.is_empty());
7858 assert!(vec.iter().all(|v| v.as_f64().is_some()));
7859 }
7860
7861 #[tokio::test]
7867 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
7868 let app = test_app();
7869 let (status, body) = post_json_uri(
7870 &app,
7871 "/v1/embeddings",
7872 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
7873 )
7874 .await;
7875 assert_eq!(status, StatusCode::BAD_REQUEST);
7876 let msg = body["error"]["message"].as_str().unwrap();
7877 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
7878 }
7879
7880 #[tokio::test]
7885 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7886 async fn a_real_embedding_model_serves_v1_embeddings() {
7887 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7888 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7889 if !path.exists() {
7890 eprintln!("SKIP: {} not present", path.display());
7891 return;
7892 }
7893 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
7894 let mut state = test_state(
7895 test_model_full_byte_vocab(),
7896 ResponseCache::new(1000, Duration::from_secs(3600)),
7897 );
7898 state.embedding = Some(Arc::new(encoder));
7899 let app = test_app_with_state(Arc::new(state));
7900
7901 let (status, body) = post_json_uri(
7902 &app,
7903 "/v1/embeddings",
7904 serde_json::json!({ "input": ["Hello world", "a second input"] }),
7905 )
7906 .await;
7907 assert_eq!(status, StatusCode::OK, "{body}");
7908 assert_eq!(body["model"], "bge-small-en-v1.5");
7909 let data = body["data"].as_array().unwrap();
7910 assert_eq!(data.len(), 2);
7911 for (i, row) in data.iter().enumerate() {
7912 assert_eq!(row["index"], i);
7913 let v: Vec<f64> = row["embedding"]
7914 .as_array()
7915 .unwrap()
7916 .iter()
7917 .map(|x| x.as_f64().unwrap())
7918 .collect();
7919 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
7920 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
7921 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
7922 }
7923 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
7926
7927 let (status, mean) = post_json_uri(
7931 &app,
7932 "/v1/embeddings",
7933 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
7934 )
7935 .await;
7936 assert_eq!(status, StatusCode::OK);
7937 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
7938 }
7939
7940 #[tokio::test]
7952 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7953 async fn an_encoder_can_be_the_loaded_model() {
7954 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7955 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7956 if !path.exists() {
7957 eprintln!("SKIP: {} not present", path.display());
7958 return;
7959 }
7960
7961 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
7965 assert!(
7966 matches!(loaded, model::LoadedModel::Encoder(_)),
7967 "an encoder-only GGUF reached a decoder loader"
7968 );
7969 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
7970 assert!(
7971 matches!(loaded, Loaded::Encoder(_)),
7972 "the encoder did not stay an encoder through activation"
7973 );
7974 assert!(
7975 batcher.is_none() && ceiling.is_none(),
7976 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7977 );
7978
7979 let state = test_state(
7980 test_model_full_byte_vocab(),
7981 ResponseCache::new(1000, Duration::from_secs(3600)),
7982 );
7983 state.swap_active(Some(Arc::new(ActiveModel {
7984 id: None,
7985 loaded,
7986 batcher,
7987 ceiling,
7988 checkpoint_path: None,
7989 })));
7990 let app = test_app_with_state(Arc::new(state));
7991
7992 let (status, body) = post_json_uri(
7994 &app,
7995 "/v1/embeddings",
7996 serde_json::json!({ "input": "Hello world" }),
7997 )
7998 .await;
7999 assert_eq!(status, StatusCode::OK, "{body}");
8000 assert_eq!(body["model"], "bge-small-en-v1.5");
8001 let v = body["data"][0]["embedding"].as_array().unwrap();
8002 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
8003
8004 let (status, body) = post_json_uri(
8006 &app,
8007 "/v1/chat/completions",
8008 serde_json::json!({
8009 "model": "bge-small-en-v1.5",
8010 "messages": [{"role": "user", "content": "hi"}],
8011 }),
8012 )
8013 .await;
8014 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
8015 let msg = body["error"]["message"].as_str().unwrap();
8016 for fact in [
8017 "bge-small-en-v1.5",
8018 "bert",
8019 "embedding model",
8020 "/v1/embeddings",
8021 ] {
8022 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
8023 }
8024
8025 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
8027 assert_eq!(status, StatusCode::OK);
8028 let entry = &models["data"][0];
8029 assert_eq!(entry["id"], "bge-small-en-v1.5");
8030 assert_eq!(entry["frink_model_kind"], "embedding");
8031 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
8032 assert_eq!(entry["frink_n_embd"], 384);
8033 assert_eq!(entry["frink_pooling"], "CLS");
8034 assert_eq!(
8035 entry["frink_endpoints"],
8036 serde_json::json!(["/v1/embeddings"])
8037 );
8038 assert!(entry.get("supported_reasoning_efforts").is_none());
8041
8042 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
8044 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
8045 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
8046 assert_eq!(health["model"]["synthetic_weights"], false);
8047 let weights = health["capabilities"]
8048 .as_array()
8049 .unwrap()
8050 .iter()
8051 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
8052 .expect("a real-weights capability row");
8053 let detail = weights["detail"].as_str().unwrap_or_default();
8054 assert!(detail.contains("ENCODER"), "{detail}");
8055 let (status, body) = post_json_uri(
8063 &app,
8064 frink_api::routes::V1_TOKENIZE,
8065 serde_json::json!({ "content": "hello world" }),
8066 )
8067 .await;
8068 assert_eq!(
8069 status,
8070 StatusCode::OK,
8071 "an encoder has a real tokenizer: {body}"
8072 );
8073 let tokens = body["tokens"].as_array().expect("tokens array").clone();
8074 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
8075
8076 let (status, body) = post_json_uri(
8077 &app,
8078 frink_api::routes::V1_DETOKENIZE,
8079 serde_json::json!({ "tokens": tokens }),
8080 )
8081 .await;
8082 assert_eq!(status, StatusCode::OK, "{body}");
8083 let round_tripped = body["content"].as_str().expect("content").to_string();
8084 assert!(
8085 round_tripped.contains("hello") && round_tripped.contains("world"),
8086 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
8087 );
8088
8089 let (status, _) = post_json_uri(
8092 &app,
8093 "/v1/completions",
8094 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
8095 )
8096 .await;
8097 assert_eq!(
8098 status,
8099 StatusCode::NOT_IMPLEMENTED,
8100 "tokenizing an encoder must not have opened a path to generating with one"
8101 );
8102 }
8103
8104 #[tokio::test]
8109 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
8110 use http_body_util::BodyExt;
8111 use tower::ServiceExt;
8112
8113 let fixture = concat!(
8114 "../frink-models/tests/fixtures/",
8115 "frink_real_moe_test.gguf"
8116 );
8117 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
8118 let decoder = Decoder::from_gguf_with_expert_cache(
8119 &fixture,
8120 frink_models::config::test_moe_fixture(),
8121 Some(1024 * 1024),
8122 )
8123 .expect("MoE fixture must load store-backed");
8124
8125 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
8129 decoder.forward_token(1, 0, &mut caches);
8130
8131 let model = Model::Gguf(GgufModel {
8132 decoder: Arc::new(decoder),
8133 tokenizer: Arc::new(ServerTokenizer::Byte),
8134 stop_tokens: StopTokens::default(),
8135 bos_id: None,
8136 is_synthetic: false,
8137 chat_template: chat_template::PromptTemplate::plain(),
8138 });
8139 let state = Arc::new(test_state(
8140 model,
8141 ResponseCache::new(16, Duration::from_secs(60)),
8142 ));
8143 let app = Router::new()
8144 .route("/metrics", axum::routing::get(metrics))
8145 .route("/v1/chat/completions", post(chat_completions))
8146 .with_state(state);
8147
8148 let fetch_metrics = |app: Router| async move {
8149 let resp = app
8150 .oneshot(
8151 axum::http::Request::builder()
8152 .method("GET")
8153 .uri("/metrics")
8154 .body(axum::body::Body::empty())
8155 .unwrap(),
8156 )
8157 .await
8158 .unwrap();
8159 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
8160 String::from_utf8(bytes.to_vec()).unwrap()
8161 };
8162
8163 let after = fetch_metrics(app.clone()).await;
8164 assert!(
8165 after.contains("frink_expert_cache_misses_total"),
8166 "streaming model must expose expert-cache metrics: {after}"
8167 );
8168 let misses: u64 = after
8169 .lines()
8170 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
8171 .and_then(|l| l.split_whitespace().nth(1))
8172 .and_then(|v| v.parse().ok())
8173 .expect("misses metric line must parse");
8174 assert!(
8175 misses > 0,
8176 "decode must have read experts through the store: {after}"
8177 );
8178 }
8179
8180 fn weather_tool() -> serde_json::Value {
8181 serde_json::json!({
8182 "type": "function",
8183 "function": {
8184 "name": "get_weather",
8185 "description": "Get the current weather for a location.",
8186 "parameters": {
8187 "type": "object",
8188 "properties": {"location": {"type": "string"}},
8189 "required": ["location"]
8190 }
8191 }
8192 })
8193 }
8194
8195 fn weather_tool_def() -> ToolDef {
8196 ToolDef {
8197 kind: "function".to_string(),
8198 function: ToolFunctionDef {
8199 name: "get_weather".to_string(),
8200 description: Some("Get the current weather for a location.".to_string()),
8201 parameters: Some(serde_json::json!({
8202 "type": "object",
8203 "properties": {"location": {"type": "string"}},
8204 "required": ["location"]
8205 })),
8206 },
8207 }
8208 }
8209
8210 #[test]
8211 fn tool_preamble_mentions_every_tool_name_and_description() {
8212 let preamble = tool_preamble(&[weather_tool_def()]);
8213 assert!(preamble.contains("get_weather"));
8214 assert!(preamble.contains("Get the current weather for a location."));
8215 assert!(preamble.contains("<tool_call>"));
8216 assert!(preamble.contains("</tool_call>"));
8217 }
8218
8219 #[test]
8220 fn a_real_marker_becomes_a_structured_tool_call() {
8221 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
8222 let (message, finish) = build_response_message(
8223 text.to_string(),
8224 &[weather_tool_def()],
8225 output::OutputPosture::for_model("test-model"),
8226 "stop",
8227 );
8228 assert_eq!(finish, "tool_calls");
8229 let calls = message.tool_calls.expect("must carry a tool call");
8230 assert_eq!(calls[0].function.name, "get_weather");
8231 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
8232 assert_eq!(parsed["location"], "Paris");
8233 }
8234
8235 #[test]
8236 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
8237 let (message, finish) = build_response_message(
8238 "just an answer".to_string(),
8239 &[weather_tool_def()],
8240 output::OutputPosture::for_model("test-model"),
8241 "stop",
8242 );
8243 assert_eq!(finish, "stop");
8244 assert!(message.tool_calls.is_none());
8245 assert_eq!(message.content.as_deref(), Some("just an answer"));
8246 }
8247
8248 #[test]
8251 fn a_malformed_payload_is_not_a_tool_call() {
8252 let (message, finish) = build_response_message(
8253 "<tool_call>not valid json at all</tool_call>".to_string(),
8254 &[weather_tool_def()],
8255 output::OutputPosture::for_model("test-model"),
8256 "stop",
8257 );
8258 assert_eq!(finish, "stop");
8259 assert!(message.tool_calls.is_none());
8260 }
8261
8262 #[test]
8265 fn a_tool_that_was_never_offered_is_not_returned() {
8266 let (message, finish) = build_response_message(
8267 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
8268 &[weather_tool_def()],
8269 output::OutputPosture::for_model("test-model"),
8270 "stop",
8271 );
8272 assert_eq!(finish, "stop");
8273 assert!(message.tool_calls.is_none());
8274 }
8275
8276 #[test]
8278 fn marker_text_with_no_tools_offered_stays_content() {
8279 let (message, finish) = build_response_message(
8280 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
8281 &[],
8282 output::OutputPosture::for_model("test-model"),
8283 "stop",
8284 );
8285 assert_eq!(finish, "stop");
8286 assert!(message.tool_calls.is_none());
8287 assert!(message.content.is_some());
8288 }
8289
8290 #[test]
8294 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
8295 let opened = std::cell::Cell::new(0usize);
8296 let mut parser = crate::policy::parser::ToolCallParser::new(
8297 crate::policy::parser::ToolCallFormat::Qwen3Coder,
8298 vec![
8299 crate::policy::parser::tool_call::ToolSchema::with_parameters(
8300 "write_file",
8301 serde_json::json!({"type": "object", "properties": {
8302 "path": {"type": "string"},
8303 "contents": {"type": "string"}
8304 }}),
8305 ),
8306 ],
8307 );
8308 let wire = "<tool_call><function=write_file>\
8309 <parameter=path>\n/tmp/x\n</parameter>\
8310 <parameter=contents>\nhello world\n</parameter>\
8311 </function></tool_call>";
8312
8313 let mut deltas = Vec::new();
8314 let mut text = String::new();
8315 for piece in wire.as_bytes().chunks(7) {
8316 let chunk = String::from_utf8_lossy(piece).into_owned();
8317 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
8318 text.push_str(&more_text);
8319 deltas.extend(more);
8320 }
8321 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
8322 text.push_str(&more_text);
8323 deltas.extend(more);
8324
8325 assert_eq!(opened.get(), 1, "one call opened");
8326 assert!(text.is_empty(), "the markers are not content: {text:?}");
8327
8328 let first = &deltas[0];
8329 assert_eq!(first.index, 0);
8330 assert_eq!(first.id.as_deref(), Some("call_0"));
8331 assert_eq!(first.kind, Some("function"));
8332 assert_eq!(first.function.name.as_deref(), Some("write_file"));
8333
8334 let joined: String = deltas
8337 .iter()
8338 .filter_map(|d| d.function.arguments.clone())
8339 .collect();
8340 let parsed: serde_json::Value =
8341 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
8342 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
8343 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
8344 assert!(
8345 deltas.len() >= 3,
8346 "the arguments arrived in pieces, not whole: {}",
8347 deltas.len()
8348 );
8349 assert!(
8350 deltas[1..].iter().all(|d| d.function.name.is_none()),
8351 "only the opening delta carries identity"
8352 );
8353 }
8354
8355 #[test]
8357 fn text_around_a_streamed_call_is_still_content() {
8358 let opened = std::cell::Cell::new(0usize);
8359 let mut parser = crate::policy::parser::ToolCallParser::new(
8360 crate::policy::parser::ToolCallFormat::Qwen25,
8361 vec![crate::policy::parser::tool_call::ToolSchema::new(
8362 "get_weather",
8363 )],
8364 );
8365 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
8366 \"arguments\": {}}</tool_call> done";
8367 let mut text = String::new();
8368 for piece in wire.as_bytes().chunks(5) {
8369 let chunk = String::from_utf8_lossy(piece).into_owned();
8370 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
8371 text.push_str(&more);
8372 }
8373 let (more, _) = tool_call_deltas(parser.finish(), &opened);
8374 text.push_str(&more);
8375
8376 assert_eq!(opened.get(), 1);
8377 assert!(text.starts_with("let me check. "), "{text:?}");
8378 assert!(text.ends_with(" done"), "{text:?}");
8379 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
8380 }
8381
8382 #[test]
8385 fn a_reasoning_block_is_split_out_of_the_answer() {
8386 let (message, finish) = build_response_message(
8387 "<think>weighing it up</think>The answer is 4.".to_string(),
8388 &[],
8389 output::OutputPosture::for_model("Qwen3-8B"),
8390 "stop",
8391 );
8392 assert_eq!(finish, "stop");
8393 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
8394 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8395 }
8396
8397 #[test]
8400 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
8401 let (message, _) = build_response_message(
8402 "Use the <think> tag like this.".to_string(),
8403 &[],
8404 output::OutputPosture::for_model("llama-3.1-8b"),
8405 "stop",
8406 );
8407 assert_eq!(
8408 message.content.as_deref(),
8409 Some("Use the <think> tag like this.")
8410 );
8411 assert!(message.reasoning_content.is_none());
8412 }
8413
8414 #[tokio::test]
8422 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
8423 let app = test_app();
8424 let body = serde_json::json!({
8425 "model": "m",
8426 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8427 "max_tokens": 4,
8428 "temperature": 0,
8429 });
8430 let resp = post_json(&app, body).await;
8431 let message = &resp["choices"][0]["message"];
8432 assert!(message["content"].is_string());
8433 assert!(message.get("tool_calls").is_none());
8434 assert_eq!(resp["choices"][0]["finish_reason"], "length");
8435 assert_eq!(resp["usage"]["completion_tokens"], 4);
8436 assert_eq!(
8437 resp["usage"]["total_tokens"],
8438 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
8439 );
8440 }
8441
8442 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
8443 use http_body_util::BodyExt;
8444 use tower::ServiceExt;
8445
8446 let response = app
8447 .clone()
8448 .oneshot(
8449 axum::http::Request::builder()
8450 .method("GET")
8451 .uri(uri)
8452 .body(axum::body::Body::empty())
8453 .unwrap(),
8454 )
8455 .await
8456 .unwrap();
8457 let status = response.status();
8458 let bytes = response.into_body().collect().await.unwrap().to_bytes();
8459 (status, serde_json::from_slice(&bytes).unwrap())
8460 }
8461
8462 #[tokio::test]
8463 async fn health_answers_a_capability_handshake_not_a_boolean() {
8464 let app = test_app();
8465 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
8466 assert_eq!(status, StatusCode::OK);
8467
8468 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
8469 assert_eq!(health.state, frink_api::HealthState::Ready);
8470 assert!(health.pid > 0);
8471 assert!(health.server_time_unix_ms > 0);
8472 assert_eq!(health.last_request_age_seconds, None);
8475
8476 for id in [
8479 frink_api::health::capability::CPU,
8480 frink_api::health::capability::METAL,
8481 frink_api::health::capability::CUDA,
8482 frink_api::health::capability::REAL_WEIGHTS,
8483 frink_api::health::capability::CONTINUOUS_BATCHING,
8484 ] {
8485 let cap = health
8486 .capability(id)
8487 .unwrap_or_else(|| panic!("{id} missing"));
8488 assert!(!cap.reason.is_empty(), "{cap:?}");
8489 assert!(!cap.detail.is_empty(), "{cap:?}");
8490 }
8491 let weights = health
8495 .capability(frink_api::health::capability::REAL_WEIGHTS)
8496 .unwrap();
8497 assert!(!weights.available);
8498 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
8499 assert!(health.model.as_ref().unwrap().synthetic_weights);
8500 }
8501
8502 #[tokio::test]
8503 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
8504 let app = test_app();
8505 let _ = post_json(
8506 &app,
8507 serde_json::json!({
8508 "model": "m",
8509 "messages": [{"role": "user", "content": "\u{1}"}],
8510 "max_tokens": 1,
8511 "temperature": 0,
8512 }),
8513 )
8514 .await;
8515 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
8516 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
8517 let age = health
8518 .last_request_age_seconds
8519 .expect("a served request is evidence of liveness");
8520 assert!((0.0..5.0).contains(&age), "implausible age {age}");
8521 }
8522
8523 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
8525 use http_body_util::BodyExt;
8526 use tower::ServiceExt;
8527
8528 let response = app
8529 .clone()
8530 .oneshot(
8531 axum::http::Request::builder()
8532 .method("POST")
8533 .uri("/v1/chat/completions")
8534 .header("content-type", "application/json")
8535 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
8536 .unwrap(),
8537 )
8538 .await
8539 .unwrap();
8540 let bytes = response.into_body().collect().await.unwrap().to_bytes();
8541 String::from_utf8(bytes.to_vec())
8542 .unwrap()
8543 .lines()
8544 .filter_map(|line| line.strip_prefix("data: "))
8545 .filter(|payload| *payload != "[DONE]")
8546 .map(|payload| serde_json::from_str(payload).unwrap())
8547 .collect()
8548 }
8549
8550 #[tokio::test]
8551 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
8552 let app = test_app();
8553 let chunks = post_sse_chunks(
8554 &app,
8555 serde_json::json!({
8556 "model": "m",
8557 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8558 "max_tokens": 4,
8559 "temperature": 0,
8560 "stream": true,
8561 }),
8562 )
8563 .await;
8564
8565 assert!(!chunks.is_empty());
8566 let request_id = chunks[0]["request_id"]
8567 .as_str()
8568 .expect("the first chunk names the request")
8569 .to_string();
8570 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
8571 for (i, chunk) in chunks.iter().enumerate().skip(1) {
8574 assert!(
8575 chunk.get("request_id").is_none(),
8576 "chunk {i} repeats request_id"
8577 );
8578 }
8579 for chunk in &chunks {
8582 assert_eq!(chunk["id"], serde_json::json!(request_id));
8583 }
8584
8585 let other = post_sse_chunks(
8586 &app,
8587 serde_json::json!({
8588 "model": "m",
8589 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8590 "max_tokens": 4,
8591 "temperature": 0,
8592 "stream": true,
8593 }),
8594 )
8595 .await;
8596 assert_ne!(
8597 other[0]["request_id"].as_str().unwrap(),
8598 request_id,
8599 "two concurrent chats must not share an id"
8600 );
8601 }
8602
8603 #[tokio::test]
8604 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
8605 let app = test_app();
8606 let resp = post_json(
8607 &app,
8608 serde_json::json!({
8609 "model": "m",
8610 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8611 "max_tokens": 2,
8612 "temperature": 0,
8613 }),
8614 )
8615 .await;
8616 assert_eq!(resp["id"], resp["request_id"]);
8617 assert!(resp["request_id"]
8618 .as_str()
8619 .unwrap()
8620 .starts_with("chatcmpl-"));
8621 }
8622
8623 #[tokio::test]
8626 async fn usage_carries_separate_prefill_and_decode_timings() {
8627 let app = test_app();
8628 let resp = post_json(
8629 &app,
8630 serde_json::json!({
8631 "model": "m",
8632 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8633 "max_tokens": 4,
8634 "temperature": 0,
8635 }),
8636 )
8637 .await;
8638 let usage = &resp["usage"];
8639 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
8640 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
8641 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
8642 assert!(usage["predicted_per_second"].is_number(), "{usage}");
8643 assert!(usage.get("cached_tokens").is_none(), "{usage}");
8645 }
8646
8647 #[tokio::test]
8655 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
8656 let app = test_app();
8657 let body = serde_json::json!({
8658 "model": "m",
8659 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8660 "max_tokens": 4,
8661 "temperature": 0,
8662 "tools": [weather_tool()],
8663 });
8664 let resp = post_json(&app, body).await;
8665 let message = &resp["choices"][0]["message"];
8666 assert!(
8667 message["content"].is_string(),
8668 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
8669 );
8670 assert!(message.get("tool_calls").is_none());
8671 assert_eq!(resp["choices"][0]["finish_reason"], "length");
8675 }
8676
8677 #[tokio::test]
8681 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
8682 let app = test_app();
8683 let body = serde_json::json!({
8684 "model": "m",
8685 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8686 "max_tokens": 3,
8687 "temperature": 0,
8688 });
8689 let first = post_json(&app, body.clone()).await;
8690 assert_eq!(first["frink_cache"], "miss");
8691 let second = post_json(&app, body).await;
8692 assert_eq!(second["frink_cache"], "hit");
8693 assert_eq!(
8694 first["choices"][0]["message"]["content"],
8695 second["choices"][0]["message"]["content"]
8696 );
8697 assert_eq!(
8698 first["choices"][0]["finish_reason"],
8699 second["choices"][0]["finish_reason"]
8700 );
8701 assert_eq!(first["usage"], second["usage"]);
8702 assert_eq!(second["usage"]["completion_tokens"], 3);
8703 }
8704
8705 #[tokio::test]
8716 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
8717 let app = test_app();
8718 let plain = serde_json::json!({
8719 "model": "m",
8720 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8721 "max_tokens": 3,
8722 "temperature": 0,
8723 });
8724
8725 let first = post_json(&app, plain.clone()).await;
8726 assert_eq!(first["frink_cache"], "miss");
8727 let unconstrained = first["choices"][0]["message"]["content"]
8728 .as_str()
8729 .expect("content")
8730 .to_string();
8731
8732 let mut constrained = plain.clone();
8733 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
8734 let second = post_json(&app, constrained).await;
8735 assert_eq!(
8736 second["frink_cache"], "miss",
8737 "a grammar is part of the key, so this body has never been answered"
8738 );
8739 let constrained_answer = second["choices"][0]["message"]["content"]
8743 .as_str()
8744 .expect("content")
8745 .to_string();
8746 assert!(
8747 constrained_answer.contains("-> \"yes\"]"),
8748 "the grammar must have been compiled AND applied, not skipped \
8749 by a cache hit: {constrained_answer}"
8750 );
8751 assert_ne!(
8752 constrained_answer, unconstrained,
8753 "the constrained request was served the unconstrained answer"
8754 );
8755
8756 let third = post_json(&app, plain).await;
8760 assert_eq!(third["frink_cache"], "hit");
8761 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
8762 }
8763
8764 #[tokio::test]
8786 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
8787 let state = Arc::new(test_state(
8788 test_model_full_byte_vocab(),
8789 ResponseCache::new(1000, Duration::from_secs(3600)),
8790 ));
8791 let app = test_app_with_state(state.clone());
8792 let plain = serde_json::json!({
8793 "model": "m",
8794 "messages": [
8795 {"role": "system", "content": "Answer in JSON when it helps."},
8796 {"role": "user", "content": "\u{1}\u{2}"},
8797 ],
8798 "max_tokens": 3,
8799 "temperature": 0,
8800 });
8801
8802 let first = post_json(&app, plain.clone()).await;
8803 assert_eq!(first["frink_cache"], "miss");
8804 assert_eq!(state.cache_stats().entries, 1);
8805
8806 let mut as_json = plain.clone();
8807 as_json["response_format"] = serde_json::json!({"type": "json_object"});
8808 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
8809 assert_eq!(
8810 status,
8811 StatusCode::BAD_REQUEST,
8812 "the demo banner is not a JSON object, whoever generated it"
8813 );
8814 assert_eq!(
8815 state.cache_stats().hits,
8816 0,
8817 "a json_object request must not be answered from an entry the \
8818 JSON mask never produced"
8819 );
8820 assert_eq!(
8821 state.cache_stats().entries,
8822 2,
8823 "json_object must key its own entry, not reuse the unconstrained \
8824 one it happens to render the same prompt as"
8825 );
8826 }
8827
8828 #[tokio::test]
8840 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
8841 let app = test_app_with_state(Arc::new(test_state(
8842 test_model_full_byte_vocab_with_eos(Some(0x77)),
8843 ResponseCache::new(1000, Duration::from_secs(3600)),
8844 )));
8845 let body = serde_json::json!({
8846 "model": "m",
8847 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
8848 "max_tokens": 6,
8849 "temperature": 0,
8850 });
8851
8852 let stopped = post_json(&app, body.clone()).await;
8853 assert_eq!(stopped["frink_cache"], "miss");
8854 assert_eq!(
8855 stopped["choices"][0]["finish_reason"], "stop",
8856 "the fixture is only meaningful if the model's EOS really fires here"
8857 );
8858 assert_eq!(stopped["usage"]["completion_tokens"], 1);
8859
8860 let mut ignoring = body.clone();
8861 ignoring["ignore_eos"] = serde_json::json!(true);
8862 let ran_on = post_json(&app, ignoring).await;
8863 assert_eq!(
8864 ran_on["frink_cache"], "miss",
8865 "ignore_eos is part of the key, so this body has never been answered"
8866 );
8867 assert_eq!(
8868 ran_on["usage"]["completion_tokens"], 6,
8869 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
8870 );
8871 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
8872 assert_ne!(
8873 ran_on["choices"][0]["message"]["content"],
8874 stopped["choices"][0]["message"]["content"]
8875 );
8876 }
8877
8878 #[tokio::test]
8884 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
8885 let session_app = test_app();
8886 let manual_app = test_app();
8887
8888 let turn1 = post_json(
8890 &session_app,
8891 serde_json::json!({
8892 "model": "m",
8893 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8894 "session_id": "s1",
8895 "max_tokens": 5,
8896 "temperature": 0,
8897 }),
8898 )
8899 .await;
8900 let reply1 = turn1["choices"][0]["message"]["content"]
8901 .as_str()
8902 .unwrap()
8903 .to_string();
8904
8905 let manual_turn1 = post_json(
8910 &manual_app,
8911 serde_json::json!({
8912 "model": "m",
8913 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8914 "max_tokens": 5,
8915 "temperature": 0,
8916 }),
8917 )
8918 .await;
8919 assert_eq!(
8920 manual_turn1["choices"][0]["message"]["content"]
8921 .as_str()
8922 .unwrap(),
8923 reply1
8924 );
8925
8926 let turn2 = post_json(
8928 &session_app,
8929 serde_json::json!({
8930 "model": "m",
8931 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
8932 "session_id": "s1",
8933 "max_tokens": 5,
8934 "temperature": 0,
8935 }),
8936 )
8937 .await;
8938 let reply2 = turn2["choices"][0]["message"]["content"]
8939 .as_str()
8940 .unwrap()
8941 .to_string();
8942
8943 let manual_turn2 = post_json(
8947 &manual_app,
8948 serde_json::json!({
8949 "model": "m",
8950 "messages": [
8951 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
8952 {"role": "assistant", "content": reply1},
8953 {"role": "user", "content": "\u{4}\u{5}"},
8954 ],
8955 "max_tokens": 5,
8956 "temperature": 0,
8957 }),
8958 )
8959 .await;
8960 assert_eq!(
8961 manual_turn2["choices"][0]["message"]["content"]
8962 .as_str()
8963 .unwrap(),
8964 reply2,
8965 "resuming a session must produce identical output to manually resending the full history"
8966 );
8967 }
8968
8969 #[test]
8972 fn lock_cache_recovers_from_a_poisoned_mutex() {
8973 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
8974
8975 let poison_cache = Arc::clone(&cache);
8976 let _ = std::thread::spawn(move || {
8977 let _guard = poison_cache.lock().unwrap();
8978 panic!("simulated panic while holding the lock");
8979 })
8980 .join();
8981
8982 let recovered = lock_cache(&cache);
8984 assert_eq!(recovered.stats().entries, 0);
8985 }
8986
8987 #[test]
8988 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8989 let mut req_body = serde_json::json!({
8990 "model": "m",
8991 "messages": [{"role": "user", "content": "hi"}],
8992 });
8993 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8994 assert!(
8995 req.is_cacheable(),
8996 "default (temperature 0) must be cacheable"
8997 );
8998
8999 req_body["temperature"] = serde_json::json!(0.8);
9000 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
9001 assert!(
9002 !req.is_cacheable(),
9003 "unseeded sampling must never be cacheable"
9004 );
9005
9006 req_body["seed"] = serde_json::json!(42);
9007 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
9008 assert!(
9009 req.is_cacheable(),
9010 "sampling with an explicit seed is deterministic and must be cacheable"
9011 );
9012 }
9013
9014 const GRADED: &str = "{% if reasoning_effort %}\
9018 {% if reasoning_effort not in ['low','medium','high'] %}\
9019 {{ raise_exception('unsupported effort') }}\
9020 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
9021 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
9022
9023 fn graded_template() -> chat_template::PromptTemplate {
9024 chat_template::PromptTemplate::from_gguf_metadata(
9025 Some(GRADED),
9026 Some("qwen3"),
9027 false,
9028 true,
9029 None,
9030 None,
9031 )
9032 }
9033
9034 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
9035 serde_json::from_value(value).expect("request")
9036 }
9037
9038 #[test]
9045 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
9046 let req = chat_request(serde_json::json!({
9047 "model": "m",
9048 "messages": [{"role": "user", "content": "hi"}],
9049 "grammar": "root ::= \"a\"+",
9050 }));
9051 req.validate_supported_fields()
9052 .expect("a valid grammar is a valid request");
9053 let params = req
9054 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9055 .expect("a valid grammar compiles at params time too");
9056 assert!(
9057 params.grammar.is_some(),
9058 "the grammar was dropped between the wire and the sampler"
9059 );
9060 assert!(
9061 params.needs_vocab_logits(),
9062 "a grammar request that may fold lm_head into a GPU argmax is \
9063 a grammar request served unconstrained"
9064 );
9065
9066 let plain = chat_request(serde_json::json!({
9067 "model": "m",
9068 "messages": [{"role": "user", "content": "hi"}],
9069 }));
9070 assert!(plain
9071 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9072 .unwrap()
9073 .grammar
9074 .is_none());
9075 }
9076
9077 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
9078 chat_request(serde_json::json!({
9079 "model": "m",
9080 "messages": [{"role": "user", "content": "weather in Rome?"}],
9081 "tools": [weather_tool()],
9082 "tool_choice": tool_choice,
9083 }))
9084 }
9085
9086 #[test]
9090 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
9091 for choice in [
9092 serde_json::json!("required"),
9093 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
9094 ] {
9095 let req = tool_request(choice.clone());
9096 req.validate_supported_fields()
9097 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
9098 let params = req
9099 .generation_params_for_template(
9100 &graded_template(),
9101 "Qwen3-8B",
9102 crate::sampling_knobs::SamplerModel::absent(),
9103 )
9104 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
9105 let grammar = params
9106 .grammar
9107 .as_ref()
9108 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
9109 assert!(
9110 grammar.is_awaiting_trigger(),
9111 "the model must be free to think before it calls"
9112 );
9113 assert!(
9114 !grammar.allows_eog(),
9115 "{choice} must not be able to end the turn without a call"
9116 );
9117 assert!(
9123 params.needs_vocab_logits(),
9124 "{choice} would let a backend return a token id instead of logits"
9125 );
9126 assert!(
9127 !generate::greedy_gpu_fold_allowed(¶ms),
9128 "{choice} at temperature 0 must still refuse the greedy GPU fold"
9129 );
9130 }
9131 }
9132
9133 #[test]
9135 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
9136 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
9137 let req = tool_request(choice.clone());
9138 req.validate_supported_fields().expect("still supported");
9139 let params = match req.generation_params_for_template(
9140 &graded_template(),
9141 "Qwen3-8B",
9142 crate::sampling_knobs::SamplerModel::absent(),
9143 ) {
9144 Ok(p) => p,
9145 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
9146 };
9147 assert!(
9148 params.grammar.is_none(),
9149 "{choice} does not force a call and must not be constrained"
9150 );
9151 }
9152 }
9153
9154 #[test]
9157 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
9158 let req = chat_request(serde_json::json!({
9160 "model": "m",
9161 "messages": [{"role": "user", "content": "hi"}],
9162 "tool_choice": "required",
9163 }));
9164 let (status, _) = req
9165 .validate_supported_fields()
9166 .expect_err("nothing to call");
9167 assert_eq!(status, StatusCode::BAD_REQUEST);
9168
9169 let req =
9171 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
9172 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
9173 assert_eq!(status, StatusCode::BAD_REQUEST);
9174 assert_eq!(body["error"]["param"], "tool_choice");
9175
9176 let req = tool_request(serde_json::json!({"type": "function"}));
9178 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
9179 assert_eq!(status, StatusCode::BAD_REQUEST);
9180
9181 let req = chat_request(serde_json::json!({
9183 "model": "m",
9184 "messages": [{"role": "user", "content": "hi"}],
9185 "tools": [weather_tool()],
9186 "tool_choice": "required",
9187 "grammar": "root ::= \"a\"+",
9188 }));
9189 let (status, _) = req
9190 .validate_supported_fields()
9191 .expect_err("a grammar and a forced call are two constraints");
9192 assert_eq!(status, StatusCode::BAD_REQUEST);
9193
9194 let req = tool_request(serde_json::json!("required"));
9200 let (status, Json(body)) = match req.generation_params_for_template(
9201 &graded_template(),
9202 "muse-glimmer-8b",
9203 crate::sampling_knobs::SamplerModel::absent(),
9204 ) {
9205 Err(e) => e,
9206 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
9207 };
9208 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
9209 assert!(
9210 body["error"]["message"]
9211 .as_str()
9212 .unwrap()
9213 .contains("muse_glimmer"),
9214 "{body}"
9215 );
9216
9217 let req = tool_request(serde_json::json!("required"));
9223 let params = req
9224 .generation_params_for_template(
9225 &graded_template(),
9226 "gemma-4-E2B-it",
9227 crate::sampling_knobs::SamplerModel::absent(),
9228 )
9229 .expect("a gemma4 forced tool_choice is served");
9230 assert!(
9231 params.grammar.is_some(),
9232 "a forced tool_choice must arrive as the generation's grammar"
9233 );
9234 }
9235
9236 #[test]
9239 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
9240 let req = chat_request(serde_json::json!({
9241 "model": "m",
9242 "messages": [{"role": "user", "content": "hi"}],
9243 "grammar": "root ::= \"a",
9244 }));
9245 let (status, Json(body)) = req
9246 .validate_supported_fields()
9247 .expect_err("this does not parse");
9248 assert_eq!(status, StatusCode::BAD_REQUEST);
9249 assert_eq!(body["error"]["param"], "grammar");
9250 assert!(
9251 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
9252 .is_err(),
9253 "and again at params time"
9254 );
9255 }
9256
9257 #[test]
9264 fn response_format_json_schema_becomes_the_requests_grammar() {
9265 let req = chat_request(serde_json::json!({
9266 "model": "m",
9267 "messages": [{"role": "user", "content": "hi"}],
9268 "response_format": {
9269 "type": "json_schema",
9270 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
9271 },
9272 }));
9273 req.validate_supported_fields()
9274 .expect("a boolean schema converts");
9275 let params = req
9276 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9277 .expect("and compiles");
9278 let grammar = params.grammar.expect("the schema is the grammar");
9279 let mut g = (*grammar).clone();
9280 g.accept_token(0, b"true").expect("a boolean is accepted");
9281 assert!(g.allows_eog(), "and completes the parse");
9282 assert!(
9283 !params.json_object,
9284 "a schema is not the json_object character-class mask"
9285 );
9286 }
9287
9288 #[test]
9292 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
9293 let req = chat_request(serde_json::json!({
9294 "model": "m",
9295 "messages": [{"role": "user", "content": "hi"}],
9296 "response_format": {
9297 "type": "json_schema",
9298 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
9299 },
9300 }));
9301 let (status, Json(body)) = req
9302 .validate_supported_fields()
9303 .expect_err("minimum has no grammar in this port");
9304 assert_eq!(status, StatusCode::BAD_REQUEST);
9305 assert!(
9306 body["error"]["message"]
9307 .as_str()
9308 .expect("a message")
9309 .contains("minimum"),
9310 "the refusal must name the keyword: {body}"
9311 );
9312 assert!(
9313 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
9314 .is_err(),
9315 "and again at params time"
9316 );
9317 }
9318
9319 #[test]
9325 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
9326 let req = chat_request(serde_json::json!({
9327 "model": "m",
9328 "messages": [{"role": "user", "content": "hi"}],
9329 "tool_choice": "required",
9330 "tools": [{
9331 "type": "function",
9332 "function": {"name": "f", "parameters": {"type": "object"}},
9333 }],
9334 "response_format": {
9335 "type": "json_schema",
9336 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
9337 },
9338 }));
9339 let (status, Json(body)) = req
9340 .validate_supported_fields()
9341 .expect_err("two constraints, one generation");
9342 assert_eq!(status, StatusCode::BAD_REQUEST);
9343 assert_eq!(body["error"]["param"], "tool_choice");
9344 }
9345
9346 #[test]
9349 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
9350 let req = chat_request(serde_json::json!({
9351 "model": "m",
9352 "messages": [{"role": "user", "content": "hi"}],
9353 }));
9354 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
9355 }
9356
9357 #[test]
9362 fn min_p_reaches_the_sampler_from_the_chat_wire() {
9363 let asked = chat_request(serde_json::json!({
9364 "model": "m",
9365 "messages": [{"role": "user", "content": "hi"}],
9366 "min_p": 0.07,
9367 }));
9368 assert_eq!(
9369 asked
9370 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
9371 .expect("knobs")
9372 .min_p,
9373 0.07
9374 );
9375
9376 let silent = chat_request(serde_json::json!({
9377 "model": "m",
9378 "messages": [{"role": "user", "content": "hi"}],
9379 }));
9380 assert_eq!(
9381 silent
9382 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
9383 .expect("knobs")
9384 .min_p,
9385 0.0,
9386 "an unset min_p must be off, not llama.cpp's CLI default"
9387 );
9388 }
9389
9390 #[test]
9399 fn no_sampler_knob_is_missing_from_the_cache_key() {
9400 let base = serde_json::json!({
9401 "model": "m",
9402 "messages": [{"role": "user", "content": "hi"}],
9403 "seed": 1,
9404 });
9405 let key_for = |body: serde_json::Value| {
9406 let req = chat_request(body);
9407 let params = req
9408 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9409 .expect("params");
9410 req.cache_key("prompt", ¶ms)
9411 };
9412 let baseline = key_for(base.clone());
9413 for (knob, value) in [
9414 ("temperature", serde_json::json!(0.5)),
9415 ("top_p", serde_json::json!(0.9)),
9416 ("min_p", serde_json::json!(0.05)),
9417 ("top_k", serde_json::json!(40)),
9418 ("repetition_penalty", serde_json::json!(1.1)),
9419 ("presence_penalty", serde_json::json!(0.3)),
9420 ("frequency_penalty", serde_json::json!(0.3)),
9421 (
9422 "samplers",
9423 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
9424 ),
9425 ] {
9426 let mut body = base.clone();
9427 body[knob] = value;
9428 assert_ne!(
9429 key_for(body),
9430 baseline,
9431 "`{knob}` is not in the cache key: two requests differing \
9432 only in it would share one cached answer"
9433 );
9434 }
9435 }
9436
9437 #[test]
9447 fn no_constraint_is_missing_from_the_cache_key() {
9448 let base = serde_json::json!({
9449 "model": "m",
9450 "messages": [{"role": "user", "content": "pick one"}],
9451 });
9452 let key_for = |body: serde_json::Value| {
9453 let req = chat_request(body);
9454 let params = req
9455 .generation_params(crate::sampling_knobs::SamplerModel::absent())
9456 .expect("params");
9457 req.cache_key("prompt", ¶ms)
9458 };
9459 let baseline = key_for(base.clone());
9460 for (field, value) in [
9461 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
9462 (
9463 "response_format",
9464 serde_json::json!({"type": "json_object"}),
9465 ),
9466 (
9467 "response_format",
9468 serde_json::json!({"type": "json_schema", "json_schema": {
9469 "name": "answer",
9470 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
9471 }}),
9472 ),
9473 ("ignore_eos", serde_json::json!(true)),
9474 ("stop", serde_json::json!(["\n"])),
9475 ("max_tokens", serde_json::json!(7)),
9476 ] {
9477 let mut body = base.clone();
9478 body[field] = value.clone();
9479 assert_ne!(
9480 key_for(body),
9481 baseline,
9482 "`{field}: {value}` is not in the cache key: two requests \
9483 differing only in it would share one cached answer"
9484 );
9485 }
9486 }
9487
9488 #[test]
9492 fn an_explicit_zero_output_budget_is_a_client_error() {
9493 let req = chat_request(serde_json::json!({
9494 "model": "m",
9495 "messages": [{"role": "user", "content": "hi"}],
9496 "max_tokens": 0,
9497 }));
9498 let (status, body) = req.validate_supported_fields().expect_err("rejected");
9499 assert_eq!(status, StatusCode::BAD_REQUEST);
9500 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
9501 }
9502
9503 #[test]
9506 fn a_request_can_turn_thinking_off() {
9507 let template = graded_template();
9508 for body in [
9509 serde_json::json!({
9510 "model": "m",
9511 "messages": [{"role": "user", "content": "hi"}],
9512 "reasoning_effort": "none",
9513 }),
9514 serde_json::json!({
9515 "model": "m",
9516 "messages": [{"role": "user", "content": "hi"}],
9517 "thinking": {"type": "disabled"},
9518 }),
9519 ] {
9520 let kwargs = chat_request(body).resolve_template_kwargs(&template);
9521 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
9522 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
9523 assert!(!kwargs.contains_key("reasoning_effort"));
9526 }
9527 }
9528
9529 #[test]
9532 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
9533 let template = graded_template();
9534 let kwargs = chat_request(serde_json::json!({
9535 "model": "m",
9536 "messages": [{"role": "user", "content": "hi"}],
9537 "reasoning_effort": "high",
9538 "thinking": {"type": "disabled"},
9539 }))
9540 .resolve_template_kwargs(&template);
9541 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
9542 assert!(!kwargs.contains_key("reasoning_effort"));
9543 }
9544
9545 #[test]
9548 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
9549 let req = chat_request(serde_json::json!({
9550 "model": "m",
9551 "messages": [{"role": "user", "content": "hi"}],
9552 "thinking": {"type": "disable"},
9553 }));
9554 let (status, _) = req.validate_supported_fields().expect_err("rejected");
9555 assert_eq!(status, StatusCode::BAD_REQUEST);
9556 }
9557
9558 #[test]
9561 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
9562 let template = graded_template();
9563 let kwargs = chat_request(serde_json::json!({
9564 "model": "m",
9565 "messages": [{"role": "user", "content": "hi"}],
9566 "reasoning_effort": "none",
9567 "chat_template_kwargs": {"enable_thinking": true},
9568 }))
9569 .resolve_template_kwargs(&template);
9570 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
9571 }
9572
9573 #[test]
9577 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
9578 let template = graded_template();
9579 let req = chat_request(serde_json::json!({
9580 "model": "m",
9581 "messages": [{"role": "user", "content": "hi"}],
9582 "reasoning_effort": "minimal",
9583 }));
9584 let kwargs = req.resolve_template_kwargs(&template);
9585 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
9586 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
9587 assert!(prompt.starts_with("E:low|"), "{prompt}");
9588 }
9589
9590 #[test]
9594 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
9595 let template = graded_template();
9596 let req = chat_request(serde_json::json!({
9597 "model": "m",
9598 "messages": [{"role": "user", "content": "hi"}],
9599 "chat_template_kwargs": {"reasoning_effort": "none"},
9600 }));
9601 let kwargs = req.resolve_template_kwargs(&template);
9602 assert!(!kwargs.contains_key("reasoning_effort"));
9603 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
9604 assert_eq!(prompt, "hi");
9605 }
9606
9607 #[test]
9610 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
9611 let template = graded_template();
9612 let req = chat_request(serde_json::json!({
9613 "model": "m",
9614 "messages": [{"role": "user", "content": "hi"}],
9615 "reasoning_effort": "low",
9616 "chat_template_kwargs": {"reasoning_effort": "high"},
9617 }));
9618 assert_eq!(
9619 req.resolve_template_kwargs(&template)["reasoning_effort"],
9620 serde_json::json!("high")
9621 );
9622 }
9623
9624 #[test]
9628 fn offering_tools_turns_thinking_on_by_itself() {
9629 let template = graded_template();
9630 let quiet = chat_request(serde_json::json!({
9631 "model": "m",
9632 "messages": [{"role": "user", "content": "hi"}],
9633 }));
9634 assert!(!quiet
9635 .resolve_template_kwargs(&template)
9636 .contains_key("enable_thinking"));
9637
9638 let with_tools = chat_request(serde_json::json!({
9639 "model": "m",
9640 "messages": [{"role": "user", "content": "hi"}],
9641 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
9642 }));
9643 let kwargs = with_tools.resolve_template_kwargs(&template);
9644 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
9645 let prompt =
9646 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
9647 assert!(prompt.starts_with("THINK|"), "{prompt}");
9648 }
9649
9650 #[test]
9655 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
9656 let opener = chat_template::PromptTemplate::from_gguf_metadata(
9657 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
9658 Some("qwen3"),
9659 false,
9660 true,
9661 None,
9662 None,
9663 );
9664 let req = chat_request(serde_json::json!({
9665 "model": "m",
9666 "messages": [{"role": "user", "content": "hi"}],
9667 "chat_template_kwargs": {"enable_thinking": true},
9668 }));
9669 let kwargs = req.resolve_template_kwargs(&opener);
9670 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
9671 assert!(prompt.ends_with("<think>"), "{prompt}");
9672
9673 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
9676 let (message, _) = build_response_message(
9677 "weighing it up</think>Paris.".to_string(),
9678 &[],
9679 posture,
9680 "stop",
9681 );
9682 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
9683 assert_eq!(message.content.as_deref(), Some("Paris."));
9684
9685 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
9688 let (message, _) = build_response_message(
9689 "weighing it up</think>Paris.".to_string(),
9690 &[],
9691 closed,
9692 "stop",
9693 );
9694 assert_eq!(message.reasoning_content, None);
9695 }
9696
9697 #[test]
9698 fn stop_param_accepts_both_single_string_and_array() {
9699 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
9700 "model": "m",
9701 "messages": [{"role": "user", "content": "hi"}],
9702 "stop": "END",
9703 }))
9704 .unwrap();
9705 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
9706
9707 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
9708 "model": "m",
9709 "messages": [{"role": "user", "content": "hi"}],
9710 "stop": ["A", "B"],
9711 }))
9712 .unwrap();
9713 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
9714 }
9715
9716 #[test]
9717 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
9718 let model = test_model();
9719 let result = run_generation(
9720 &model,
9721 "hello",
9722 &greedy_params(4),
9723 None,
9724 None,
9725 None,
9726 None,
9727 None,
9728 None,
9729 );
9730 assert!(matches!(
9731 result,
9732 Err(generate::DecodeError::TokenOutOfVocab { .. })
9733 ));
9734 }
9735
9736 #[test]
9740 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
9741 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9743 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
9744
9745 let holder_pool = Arc::clone(&pool);
9746 let holder = std::thread::spawn(move || {
9747 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
9748 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
9750 drop(held);
9751 });
9752 std::thread::sleep(Duration::from_millis(15));
9753
9754 let config = generate::KvPoolConfig {
9755 pool,
9756 queue_wait: Duration::ZERO,
9757 };
9758 let result = run_generation(
9759 &model,
9760 &prompt,
9761 &greedy_params(4),
9762 Some(&config),
9763 None,
9764 None,
9765 None,
9766 None,
9767 None,
9768 );
9769 assert!(matches!(
9770 result,
9771 Err(generate::DecodeError::KvPoolExhausted)
9772 ));
9773
9774 let (status, _body) = decode_error_response(result.unwrap_err());
9775 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
9776 holder.join().unwrap();
9777 }
9778
9779 #[test]
9787 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
9788 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9790 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
9792 let config = generate::KvPoolConfig {
9793 pool,
9794 queue_wait: Duration::ZERO,
9795 };
9796
9797 let result = run_generation(
9798 &model,
9799 &prompt,
9800 &greedy_params(4),
9801 Some(&config),
9802 None,
9803 None,
9804 None,
9805 None,
9806 None,
9807 );
9808 let err = result.expect_err("one block cannot hold two layers' caches");
9809 assert!(
9810 matches!(
9811 &err,
9812 generate::DecodeError::KvBudgetExceeded { binding, .. }
9813 if *binding == frink_models::Ceiling::DeviceMemory.code()
9814 ),
9815 "expected an immovable device-memory refusal, got {err:?}"
9816 );
9817 let (status, _body) = decode_error_response(err);
9818 assert_eq!(status, StatusCode::BAD_REQUEST);
9819 }
9820
9821 #[test]
9827 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
9828 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
9829 queued: 512,
9830 cap: 512,
9831 });
9832 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
9833 assert_eq!(body["error"]["retry_after_seconds"], 1);
9834 let message = body["error"]["message"].as_str().expect("message");
9835 assert!(message.contains("512"), "{message}");
9836 }
9837
9838 #[test]
9839 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
9840 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
9841 token: 99,
9842 vocab_size: 32,
9843 });
9844 assert!(
9845 body["error"]["retry_after_seconds"].is_null(),
9846 "retrying a prompt this model cannot tokenize never helps"
9847 );
9848 }
9849
9850 #[test]
9851 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
9852 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
9853 token: 99,
9854 vocab_size: 32,
9855 });
9856 assert_eq!(status, StatusCode::BAD_REQUEST);
9857 }
9858
9859 #[test]
9860 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
9861 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9863 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
9864 let config = generate::KvPoolConfig {
9865 pool: pool.clone(),
9866 queue_wait: Duration::ZERO,
9867 };
9868
9869 let produced = run_generation(
9870 &model,
9871 &prompt,
9872 &greedy_params(4),
9873 Some(&config),
9874 None,
9875 None,
9876 None,
9877 None,
9878 None,
9879 )
9880 .unwrap();
9881 assert_eq!(produced.choices[0].finish, FinishReason::Length);
9882 assert_eq!(
9883 pool.lock().unwrap().free_blocks(),
9884 2,
9885 "a completed request must return its blocks to the pool"
9886 );
9887 }
9888
9889 #[tokio::test]
9894 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
9895 let model = Arc::new(test_model());
9896 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9897
9898 let mut handles = Vec::new();
9899 for _ in 0..8 {
9900 let model = Arc::clone(&model);
9901 let prompt = prompt.clone();
9902 handles.push(tokio::task::spawn_blocking(move || {
9903 run_generation(
9904 &model,
9905 &prompt,
9906 &greedy_params(6),
9907 None,
9908 None,
9909 None,
9910 None,
9911 None,
9912 None,
9913 )
9914 .unwrap()
9915 }));
9916 }
9917
9918 let mut results = Vec::new();
9919 for h in handles {
9920 results.push(h.await.unwrap());
9921 }
9922 for r in &results[1..] {
9927 assert_eq!(r.choices, results[0].choices, "choices must match");
9931 assert_eq!(
9932 r.usage.prompt_tokens, results[0].usage.prompt_tokens,
9933 "prompt token count must match"
9934 );
9935 assert_eq!(
9936 r.usage.completion_tokens, results[0].usage.completion_tokens,
9937 "completion token count must match"
9938 );
9939 }
9940 }
9941
9942 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
9949 let mut header_entries = Vec::new();
9950 let mut data = Vec::new();
9951 for (name, shape, values) in tensors {
9952 let start = data.len();
9953 for v in values {
9954 data.extend_from_slice(&v.to_le_bytes());
9955 }
9956 let end = data.len();
9957 let shape_str = shape
9958 .iter()
9959 .map(|d| d.to_string())
9960 .collect::<Vec<_>>()
9961 .join(",");
9962 header_entries.push(format!(
9963 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
9964 ));
9965 }
9966 let header = format!("{{{}}}", header_entries.join(","));
9967 let header_bytes = header.as_bytes();
9968 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
9969 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
9970 out.extend_from_slice(header_bytes);
9971 out.extend_from_slice(&data);
9972 out
9973 }
9974
9975 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9986 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9987 use frink_models::kimi_loader::KimiRealHparams;
9988 use frink_moe::{GatingFunction, MoeLayerConfig};
9989
9990 let hidden_dim = 8;
9991 let kda_num_heads = 2;
9992 let kda_head_dim = 3;
9993 let kda_proj = kda_num_heads * kda_head_dim;
9994 let conv_kernel = 4;
9995 let dense_intermediate = 5;
9996 let vocab_size = 256;
10000 let mla_num_heads = 1;
10001 let mla_q_lora_rank = 2;
10002 let mla_kv_lora_rank = 2;
10003 let mla_qk_nope_head_dim = 2;
10004 let mla_qk_rope_head_dim = 2;
10005 let mla_v_head_dim = 2;
10006
10007 let model_cfg = frink_models::ModelConfig {
10008 rope_layers: frink_models::rope_layers::RopeLayers::All,
10009 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
10010 name: "synthetic-kimi-server-test",
10011 n_layers: 1,
10012 n_mtp_blocks: 0,
10013 hidden_dim,
10014 n_heads: 1,
10015 n_kv_heads: 1,
10016 head_dim: 4,
10017 v_head_dim: None,
10018 vocab_size,
10019 rope_theta: 10000.0,
10020 rms_norm_eps: 1e-5,
10021 post_norm_eps: 1e-5,
10022 sliding_window: None,
10023 moe: MoeLayerConfig {
10024 expert_weights_scale: 1.0,
10025 routed_weight_before_ffn: false,
10026 n_experts: 1,
10027 n_experts_active: 1,
10028 n_shared_experts: 0,
10029 hidden_dim,
10030 expert_ffn_dim: 4,
10031 gating: GatingFunction::Sigmoid,
10032 norm_topk_prob: true,
10033 expert_group_count: None,
10034 expert_group_used_count: None,
10035 },
10036 n_dense_leading_layers: 1,
10041 moe_interleave_step: None,
10042 norm_function: frink_models::norm::NormFunction::Rms,
10043 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
10044 kda_layers: vec![1],
10045 full_attn_layers: vec![],
10046 mla: MlaConfig {
10047 num_heads: mla_num_heads,
10048 q_lora_rank: mla_q_lora_rank,
10049 kv_lora_rank: mla_kv_lora_rank,
10050 qk_nope_head_dim: mla_qk_nope_head_dim,
10051 qk_rope_head_dim: mla_qk_rope_head_dim,
10052 v_head_dim: mla_v_head_dim,
10053 use_output_gate: true,
10054 rope: None,
10055 },
10056 kda: KdaConfig {
10057 num_heads: kda_num_heads,
10058 head_dim: kda_head_dim,
10059 short_conv_kernel_size: conv_kernel,
10060 gate_lower_bound: -5.0,
10061 use_full_rank_gate: true,
10062 },
10063 }),
10064 rope_freqs: None,
10065 rope_attn_factor: 1.0,
10066 rope_dim: None,
10067 rope_dim_swa: None,
10068 rope_freqs_long: None,
10069 rope_freqs_short: None,
10070 rope_orig_ctx: None,
10071 rope_layout: frink_models::config::RopeLayout::Neox,
10072 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
10073 swa_layers: frink_models::swa_layers::SwaLayers::All,
10074 attn_logit_softcap: None,
10075 final_logit_softcap: None,
10076 embedding_scale: None,
10077 residual_scale: None,
10078 normed_residual_scale: None,
10079 clamp_kqv: None,
10080 attn_temperature: None,
10081 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
10082 block_sub_norms: false,
10083 parallel_residual: false,
10084 learned_positions: false,
10085 attn_value_scale: None,
10086 alibi_max_bias: None,
10087 layer_loops: None,
10088 skip_stream: false,
10089 parallel_ssm: false,
10090 swa_chunked: false,
10091 weightless_qk_norm: false,
10092 logit_multiplier: None,
10093 attention_scale: None,
10094 rope_theta_swa: None,
10095 ffn_activation: frink_models::config::FfnActivation::Swiglu,
10096 best_effort_fields: &["synthetic test config, not a real preset"],
10097 };
10098 let hp = KimiRealHparams {
10099 hidden_dim,
10100 kda_num_heads,
10101 kda_head_dim,
10102 mla_num_heads,
10103 mla_q_lora_rank,
10104 mla_kv_lora_rank,
10105 mla_qk_nope_head_dim,
10106 mla_qk_rope_head_dim,
10107 mla_v_head_dim,
10108 dense_intermediate_dim: dense_intermediate,
10109 moe_hidden_dim: hidden_dim,
10110 moe_intermediate_dim: 4,
10111 n_experts: 1,
10112 num_shared_experts: 0,
10113 };
10114
10115 let prefix = "language_model.model.layers.0";
10119 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
10120 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
10121 name: String,
10122 shape: Vec<usize>,
10123 n: usize| {
10124 tensors.push((name, shape, vec![0.01f32; n]));
10125 };
10126 push(
10127 &mut tensors,
10128 format!("{prefix}.input_layernorm.weight"),
10129 vec![hidden_dim],
10130 hidden_dim,
10131 );
10132 push(
10133 &mut tensors,
10134 format!("{prefix}.post_attention_layernorm.weight"),
10135 vec![hidden_dim],
10136 hidden_dim,
10137 );
10138 push(
10139 &mut tensors,
10140 format!("{prefix}.self_attention_res_norm.weight"),
10141 vec![hidden_dim],
10142 hidden_dim,
10143 );
10144 push(
10145 &mut tensors,
10146 format!("{prefix}.self_attention_res_proj.weight"),
10147 vec![1, hidden_dim],
10148 hidden_dim,
10149 );
10150 push(
10151 &mut tensors,
10152 format!("{prefix}.mlp_res_norm.weight"),
10153 vec![hidden_dim],
10154 hidden_dim,
10155 );
10156 push(
10157 &mut tensors,
10158 format!("{prefix}.mlp_res_proj.weight"),
10159 vec![1, hidden_dim],
10160 hidden_dim,
10161 );
10162 push(
10163 &mut tensors,
10164 format!("{prefix}.self_attn.q_proj.weight"),
10165 vec![kda_proj, hidden_dim],
10166 kda_proj * hidden_dim,
10167 );
10168 push(
10169 &mut tensors,
10170 format!("{prefix}.self_attn.k_proj.weight"),
10171 vec![kda_proj, hidden_dim],
10172 kda_proj * hidden_dim,
10173 );
10174 push(
10175 &mut tensors,
10176 format!("{prefix}.self_attn.v_proj.weight"),
10177 vec![kda_proj, hidden_dim],
10178 kda_proj * hidden_dim,
10179 );
10180 push(
10181 &mut tensors,
10182 format!("{prefix}.self_attn.q_conv1d.weight"),
10183 vec![kda_proj, 1, conv_kernel],
10184 kda_proj * conv_kernel,
10185 );
10186 push(
10187 &mut tensors,
10188 format!("{prefix}.self_attn.k_conv1d.weight"),
10189 vec![kda_proj, 1, conv_kernel],
10190 kda_proj * conv_kernel,
10191 );
10192 push(
10193 &mut tensors,
10194 format!("{prefix}.self_attn.v_conv1d.weight"),
10195 vec![kda_proj, 1, conv_kernel],
10196 kda_proj * conv_kernel,
10197 );
10198 push(
10199 &mut tensors,
10200 format!("{prefix}.self_attn.A_log"),
10201 vec![kda_num_heads],
10202 kda_num_heads,
10203 );
10204 push(
10205 &mut tensors,
10206 format!("{prefix}.self_attn.f_a_proj.weight"),
10207 vec![kda_head_dim, hidden_dim],
10208 kda_head_dim * hidden_dim,
10209 );
10210 push(
10211 &mut tensors,
10212 format!("{prefix}.self_attn.f_b_proj.weight"),
10213 vec![kda_proj, kda_head_dim],
10214 kda_proj * kda_head_dim,
10215 );
10216 push(
10217 &mut tensors,
10218 format!("{prefix}.self_attn.dt_bias"),
10219 vec![kda_proj],
10220 kda_proj,
10221 );
10222 push(
10223 &mut tensors,
10224 format!("{prefix}.self_attn.b_proj.weight"),
10225 vec![kda_num_heads, hidden_dim],
10226 kda_num_heads * hidden_dim,
10227 );
10228 push(
10229 &mut tensors,
10230 format!("{prefix}.self_attn.g_proj.weight"),
10231 vec![kda_proj, hidden_dim],
10232 kda_proj * hidden_dim,
10233 );
10234 push(
10235 &mut tensors,
10236 format!("{prefix}.self_attn.o_norm.weight"),
10237 vec![kda_head_dim],
10238 kda_head_dim,
10239 );
10240 push(
10241 &mut tensors,
10242 format!("{prefix}.self_attn.o_proj.weight"),
10243 vec![hidden_dim, kda_proj],
10244 hidden_dim * kda_proj,
10245 );
10246 push(
10247 &mut tensors,
10248 format!("{prefix}.mlp.gate_proj.weight"),
10249 vec![dense_intermediate, hidden_dim],
10250 dense_intermediate * hidden_dim,
10251 );
10252 push(
10253 &mut tensors,
10254 format!("{prefix}.mlp.up_proj.weight"),
10255 vec![dense_intermediate, hidden_dim],
10256 dense_intermediate * hidden_dim,
10257 );
10258 push(
10259 &mut tensors,
10260 format!("{prefix}.mlp.down_proj.weight"),
10261 vec![hidden_dim, dense_intermediate],
10262 hidden_dim * dense_intermediate,
10263 );
10264 push(
10265 &mut tensors,
10266 "language_model.model.embed_tokens.weight".to_string(),
10267 vec![vocab_size, hidden_dim],
10268 vocab_size * hidden_dim,
10269 );
10270 push(
10271 &mut tensors,
10272 "language_model.lm_head.weight".to_string(),
10273 vec![vocab_size, hidden_dim],
10274 vocab_size * hidden_dim,
10275 );
10276 push(
10277 &mut tensors,
10278 "language_model.model.norm.weight".to_string(),
10279 vec![hidden_dim],
10280 hidden_dim,
10281 );
10282 push(
10283 &mut tensors,
10284 "language_model.model.output_attn_res_norm.weight".to_string(),
10285 vec![hidden_dim],
10286 hidden_dim,
10287 );
10288 push(
10289 &mut tensors,
10290 "language_model.model.output_attn_res_proj.weight".to_string(),
10291 vec![1, hidden_dim],
10292 hidden_dim,
10293 );
10294
10295 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
10308 let dir = std::env::temp_dir().join(format!(
10309 "frink_server_kimi_e2e_test_{}_{}",
10310 std::process::id(),
10311 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
10312 ));
10313 std::fs::create_dir_all(&dir).unwrap();
10314 let shard_bytes = write_safetensors_shard(&tensors);
10315 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
10316 let map_entries: Vec<String> = tensors
10317 .iter()
10318 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
10319 .collect();
10320 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
10321 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
10322
10323 use base64::Engine;
10327 let vocab_lines: Vec<String> = (0..vocab_size as u32)
10328 .map(|b| {
10329 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
10330 format!("{b64} {b}")
10331 })
10332 .collect();
10333 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
10334
10335 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
10336 .expect("must load the synthetic Kimi checkpoint end to end");
10337 std::fs::remove_dir_all(&dir).ok();
10338 loaded
10339 }
10340
10341 #[test]
10349 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
10350 let loaded = build_synthetic_kimi_loaded();
10351 let state = build_app_state(
10352 StartupModels {
10353 loaded: model::LoadedModel::Kimi(loaded),
10354 embedding: None,
10355 },
10356 None,
10357 None,
10358 None,
10359 false,
10360 None,
10361 Arc::new(health::Detection::ready(health::probe_backends())),
10362 );
10363 let active = state.active().expect("a freshly built state has a model");
10364 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
10365 assert!(!active.is_synthetic());
10366
10367 let produced = run_generation(
10368 active.generative().unwrap(),
10369 "hi",
10370 &greedy_params(5),
10371 None,
10372 None,
10373 None,
10374 None,
10375 None,
10376 None,
10377 )
10378 .expect("a real Kimi checkpoint must generate without error");
10379 assert!(matches!(
10380 produced.choices[0].finish,
10381 FinishReason::Length | FinishReason::Stop
10382 ));
10383 }
10384
10385 #[test]
10398 fn a_grammar_constrains_the_engine_decode_path() {
10399 let loaded = build_synthetic_kimi_loaded();
10400 let state = build_app_state(
10401 StartupModels {
10402 loaded: model::LoadedModel::Kimi(loaded),
10403 embedding: None,
10404 },
10405 None,
10406 None,
10407 None,
10408 false,
10409 None,
10410 Arc::new(health::Detection::ready(health::probe_backends())),
10411 );
10412 let active = state.active().expect("a freshly built state has a model");
10413
10414 let run = |grammar: Option<&str>| {
10415 let mut params = greedy_params(6);
10416 params.grammar = grammar.map(|src| {
10417 Arc::new(
10418 frink_models::grammar::Grammar::from_str_with_root(src, "root")
10419 .expect("test grammar parses"),
10420 )
10421 });
10422 run_generation(
10423 active.generative().unwrap(),
10424 "hi",
10425 ¶ms,
10426 None,
10427 None,
10428 None,
10429 None,
10430 None,
10431 None,
10432 )
10433 };
10434
10435 let produced = run(None).expect("the unconstrained run must serve");
10436 let unconstrained = produced.choices[0].text.clone();
10437 assert!(
10438 unconstrained.chars().any(|c| c != 'a'),
10439 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
10440 constrained run below would prove nothing"
10441 );
10442
10443 let produced =
10444 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
10445 let one = produced.choices.into_iter().next().unwrap();
10446 let (finish, constrained) = (one.finish, one.text);
10447 assert!(
10448 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
10449 "the engine decode path served text its grammar forbids ({constrained:?}): \
10450 the constraint was dropped between `generate_engine` and the sampler"
10451 );
10452 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
10453 }
10454
10455 #[test]
10465 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
10466 let loaded = build_synthetic_kimi_loaded();
10467 let state = build_app_state(
10468 StartupModels {
10469 loaded: model::LoadedModel::Kimi(loaded),
10470 embedding: None,
10471 },
10472 None,
10473 None,
10474 None,
10475 false,
10476 None,
10477 Arc::new(health::Detection::ready(health::probe_backends())),
10478 );
10479
10480 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
10481 let kv_pool_config = generate::KvPoolConfig {
10482 pool: pool.clone(),
10483 queue_wait: Duration::ZERO,
10484 };
10485 let pc = Mutex::new(PrefixCache::new(4));
10486
10487 run_generation(
10488 state
10489 .active()
10490 .expect("a freshly built state has a model")
10491 .generative()
10492 .unwrap(),
10493 "hi",
10494 &greedy_params(5),
10495 Some(&kv_pool_config),
10496 None,
10497 Some(&pc),
10498 None,
10499 None,
10500 None,
10501 )
10502 .expect("a real Kimi checkpoint must generate without error");
10503
10504 assert_eq!(
10505 pool.lock().unwrap().free_blocks(),
10506 4,
10507 "the KV pool must be completely untouched by a Kimi request"
10508 );
10509 let stats = pc.lock().unwrap().stats();
10510 assert_eq!(
10511 stats.hits + stats.misses,
10512 0,
10513 "the prefix cache must never be consulted for a Kimi request"
10514 );
10515 }
10516}