1mod admin;
36mod anthropic;
37mod attribution;
38mod best_of;
39mod budget;
40mod cache_admin;
41mod cache_salt;
42mod cancel;
43mod chat_params;
44mod chat_template;
45mod cli;
46mod completion;
47mod continuation;
48mod conversations;
49mod decode_task;
50mod embeddings;
51mod generate;
52mod grammar_request;
53mod health;
54mod journal;
55mod json_mode;
56mod limits;
57mod loaded;
58mod logprobs;
59mod lora;
60mod mcp;
61mod model;
62mod openai_extra;
63mod output;
64mod policy;
65mod prefill_batch;
66mod reasoning_budget;
67mod reasoning_tokens;
68mod request_tail;
69mod rerank;
70mod response_cache;
71pub(crate) mod responses;
72mod resume;
73mod sample_step;
74mod sampling_knobs;
75mod sampling_loop;
76mod security;
77mod serving;
78mod session;
79mod slots;
80mod sse;
81mod stats;
82mod stop;
83mod stream_events;
84mod tasks;
85mod tool_grammar;
86mod unimplemented_fields;
87mod unsupported_sampling;
88mod utf8_stream;
89
90use std::cell::RefCell;
91use std::convert::Infallible;
92use std::net::SocketAddr;
93use std::path::PathBuf;
94use std::rc::Rc;
95use std::sync::{Arc, Mutex, MutexGuard};
96use std::time::Duration;
97
98use axum::{
99 extract::State,
100 http::StatusCode,
101 response::sse::{Event, Sse},
102 response::{IntoResponse, Response},
103 routing::{get, post},
104 Json, Router,
105};
106use serde::{Deserialize, Serialize};
107
108use cli::apply_cli_overrides;
109pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
110
111use frink_core::cache::KvBlockPool;
112use frink_models::kimi_tokenizer::KimiTokenizer;
113use frink_models::sampling::SamplingParams;
114use frink_models::tokenizer::{SpecialTokens, StopTokens};
115use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
116#[cfg(test)]
117use generate::FinishReason;
118use generate::GenerationParams;
119pub(crate) use loaded::{ActiveModel, Loaded};
120use model::ServerTokenizer;
121use rerank::encoder_endpoints;
122use response_cache::ResponseCache;
123use sampling_knobs::SamplingKnobs;
124
125#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
134 Gguf(GgufModel),
135 Kimi(KimiModel),
136 Mla(MlaModel),
137 Gemma4(Gemma4Model),
138 Glm52(Glm52Model),
139}
140
141pub(crate) struct GgufModel {
142 decoder: Arc<Decoder>,
143 tokenizer: Arc<ServerTokenizer>,
144 stop_tokens: StopTokens,
145 bos_id: Option<usize>,
146 is_synthetic: bool,
147 chat_template: chat_template::PromptTemplate,
148}
149
150pub(crate) struct KimiModel {
151 engine: KimiEngine,
152 tokenizer: KimiTokenizer,
153 stop_tokens: StopTokens,
154 chat_template: chat_template::PromptTemplate,
155}
156
157pub(crate) struct MlaModel {
158 engine: MlaEngine,
159 tokenizer: ServerTokenizer,
160 stop_tokens: StopTokens,
161 bos_id: Option<usize>,
162 name: String,
163 chat_template: chat_template::PromptTemplate,
164}
165
166pub(crate) struct Gemma4Model {
167 engine: Gemma4Engine,
168 tokenizer: ServerTokenizer,
169 stop_tokens: StopTokens,
170 bos_id: Option<usize>,
171 name: String,
172 chat_template: chat_template::PromptTemplate,
173}
174
175pub(crate) struct Glm52Model {
176 engine: frink_models::Glm52Engine,
177 tokenizer: ServerTokenizer,
178 stop_tokens: StopTokens,
179 bos_id: Option<usize>,
180 name: String,
181 chat_template: chat_template::PromptTemplate,
182}
183
184impl Model {
185 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
186 match self {
187 Model::Gguf(m) => m.chat_template.clone(),
188 Model::Kimi(m) => m.chat_template.clone(),
189 Model::Mla(m) => m.chat_template.clone(),
190 Model::Gemma4(m) => m.chat_template.clone(),
191 Model::Glm52(m) => m.chat_template.clone(),
192 }
193 }
194
195 fn is_synthetic(&self) -> bool {
200 match self {
201 Model::Gguf(m) => m.is_synthetic,
202 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
203 }
204 }
205
206 fn tokenizer_kind(&self) -> &'static str {
207 match self {
208 Model::Gguf(m) => m.tokenizer.kind(),
209 Model::Kimi(_) => "kimi-tiktoken-bpe",
210 Model::Mla(m) => m.tokenizer.kind(),
211 Model::Gemma4(m) => m.tokenizer.kind(),
212 Model::Glm52(m) => m.tokenizer.kind(),
213 }
214 }
215
216 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
220 match self {
221 Model::Gguf(m) => m.decoder.expert_store_stats(),
222 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
223 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
224 }
225 }
226
227 pub(crate) fn name(&self) -> &str {
228 match self {
229 Model::Gguf(m) => m.decoder.config.name,
230 Model::Kimi(_) => "kimi-k3",
231 Model::Mla(m) => m.name.as_str(),
232 Model::Gemma4(m) => m.name.as_str(),
233 Model::Glm52(m) => m.name.as_str(),
234 }
235 }
236
237 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
264 match self {
265 Model::Gguf(m) => m.tokenizer.encode(text, specials),
266 Model::Kimi(m) => m
267 .tokenizer
268 .encode(text, specials)
269 .into_iter()
270 .map(|id| id as usize)
271 .collect(),
272 Model::Mla(m) => m.tokenizer.encode(text, specials),
273 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
274 Model::Glm52(m) => m.tokenizer.encode(text, specials),
275 }
276 }
277
278 pub(crate) fn bos_id(&self) -> Option<usize> {
287 match self {
288 Model::Gguf(m) => m.bos_id,
289 Model::Kimi(_) => None,
290 Model::Mla(m) => m.bos_id,
291 Model::Gemma4(m) => m.bos_id,
292 Model::Glm52(m) => m.bos_id,
293 }
294 }
295
296 pub(crate) fn decode(&self, ids: &[usize]) -> String {
297 match self {
298 Model::Gguf(m) => m.tokenizer.decode(ids),
299 Model::Kimi(m) => {
300 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
301 m.tokenizer.decode(&ids32)
302 }
303 Model::Mla(m) => m.tokenizer.decode(ids),
304 Model::Gemma4(m) => m.tokenizer.decode(ids),
305 Model::Glm52(m) => m.tokenizer.decode(ids),
306 }
307 }
308
309 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
312 match self {
313 Model::Gguf(m) => {
314 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
315 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
316 }
317 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
318 }
319 }
320
321 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
327 match self {
328 Model::Gguf(m) => Some(&m.decoder),
329 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
330 }
331 }
332
333 pub(crate) fn vocab_size(&self) -> Option<usize> {
334 match self {
335 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
336 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
337 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
338 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
339 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
340 }
341 }
342
343 fn has_real_vocabulary(&self) -> bool {
350 match self {
351 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
352 Model::Kimi(_) => true,
353 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
354 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
355 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
356 }
357 }
358}
359
360impl frink_models::dry::DryVocab for Model {
366 fn n_tokens(&self) -> usize {
367 self.vocab_size().unwrap_or(0)
368 }
369
370 fn detokenize(&self, token: usize) -> String {
371 self.decode(&[token])
372 }
373
374 fn tokenize(&self, text: &str) -> Vec<usize> {
375 self.encode(text, SpecialTokens::AsText)
376 }
377}
378
379pub(crate) struct AppState {
380 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
394 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
413 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
417 pub(crate) tasks: Arc<tasks::TaskRegistry>,
419 pub(crate) cancels: Arc<cancel::CancelRegistry>,
423 pub(crate) stats: stats::Stats,
426 pub(crate) streams: resume::StreamRegistry,
429 pub(crate) model_dir: Option<PathBuf>,
431 response_cache: Mutex<ResponseCache>,
435 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
447 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
460 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
470 sessions: session::SessionStore,
476 requests_total: std::sync::atomic::AtomicU64,
477 request_errors_total: std::sync::atomic::AtomicU64,
478 started_at: std::time::Instant,
479 last_request_ms: std::sync::atomic::AtomicU64,
484 detection: Arc<health::Detection>,
486 mcp: Option<mcp::LoadedMcpConfig>,
488 pub(crate) continuous_batching_enabled: bool,
492 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
496 loading_model: Mutex<Option<String>>,
500 last_load_error: Mutex<Option<(String, String)>>,
504 pub(crate) serving: Mutex<crate::stats::ServingStats>,
509 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
514 pub(crate) footprint:
520 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
521 pub(crate) started_unix: u64,
531}
532
533pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
539
540impl AppState {
541 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
545 self.active
546 .read()
547 .unwrap_or_else(|p| p.into_inner())
548 .clone()
549 }
550
551 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
557 self.active().ok_or_else(|| {
558 (
559 StatusCode::SERVICE_UNAVAILABLE,
560 Json(serde_json::json!({"error": {
561 "message": "no model is loaded; POST /admin/models/load with an id from \
562 GET /admin/models",
563 "type": "model_not_loaded"
564 }})),
565 )
566 })
567 }
568
569 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
577 Ok(Arc::clone(self.require_active()?.generative()?))
578 }
579
580 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
588 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
589 std::mem::replace(&mut *guard, next)
590 }
591
592 fn mark_request_finished(&self) {
596 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
597 self.last_request_ms
598 .store(ms, std::sync::atomic::Ordering::Relaxed);
599 }
600
601 pub(crate) fn uptime(&self) -> Duration {
602 self.started_at.elapsed()
603 }
604
605 pub(crate) fn requests_total(&self) -> u64 {
606 self.requests_total
607 .load(std::sync::atomic::Ordering::Relaxed)
608 }
609
610 pub(crate) fn errors_total(&self) -> u64 {
611 self.request_errors_total
612 .load(std::sync::atomic::Ordering::Relaxed)
613 }
614
615 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
616 lock_cache(&self.response_cache).stats()
617 }
618
619 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
622 let last = self
623 .last_request_ms
624 .load(std::sync::atomic::Ordering::Relaxed);
625 (last > 0)
626 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
627 .map(|age| age.max(0.0))
628 }
629
630 pub(crate) fn loading_model_id(&self) -> Option<String> {
631 self.loading_model
632 .lock()
633 .unwrap_or_else(|p| p.into_inner())
634 .clone()
635 }
636
637 pub(crate) fn set_loading_model(&self, id: Option<String>) {
638 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
639 }
640
641 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
642 self.last_load_error
643 .lock()
644 .unwrap_or_else(|p| p.into_inner())
645 .clone()
646 }
647
648 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
649 *self
650 .last_load_error
651 .lock()
652 .unwrap_or_else(|p| p.into_inner()) = error;
653 }
654
655 pub(crate) fn active_model_name(&self) -> Option<String> {
664 self.active().map(|a| a.name().to_string())
665 }
666
667 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
677 self.embedding
678 .clone()
679 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
680 }
681
682 pub(crate) fn embedding_model_name(&self) -> Option<String> {
686 match self.embedding_model() {
687 Some(e) => Some(e.name().to_string()),
688 None => self.active_model_name(),
689 }
690 }
691
692 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
693 self.stats.record(stats::entry(record));
694 }
695}
696
697fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
702 cache
703 .lock()
704 .unwrap_or_else(|poisoned| poisoned.into_inner())
705}
706
707#[derive(Debug, Clone, Deserialize)]
708#[serde(untagged)]
709pub(crate) enum MessageContent {
710 Text(String),
711 Parts(Vec<ContentPart>),
712}
713
714#[derive(Debug, Clone, Deserialize)]
715struct ContentPart {
716 #[serde(rename = "type")]
717 kind: String,
718 #[serde(default)]
719 text: Option<String>,
720 #[serde(default)]
721 image_url: Option<serde_json::Value>,
722}
723
724impl MessageContent {
725 fn as_text(&self) -> String {
726 match self {
727 Self::Text(s) => s.clone(),
728 Self::Parts(parts) => parts
729 .iter()
730 .filter_map(|p| p.text.as_deref())
731 .collect::<Vec<_>>()
732 .join(""),
733 }
734 }
735
736 fn has_image(&self) -> bool {
737 match self {
738 Self::Text(_) => false,
739 Self::Parts(parts) => parts
740 .iter()
741 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
742 }
743 }
744}
745
746#[derive(Debug, Clone, Deserialize)]
747pub(crate) struct ChatMessage {
748 pub(crate) role: String,
749 #[serde(default)]
753 pub(crate) content: Option<MessageContent>,
754 #[serde(default)]
758 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
759 #[serde(default)]
764 #[allow(dead_code)]
765 pub(crate) tool_call_id: Option<String>,
766 #[serde(default, alias = "reasoning")]
782 pub(crate) reasoning_content: Option<String>,
783}
784
785impl ChatMessage {
786 fn rendered_content(&self) -> String {
793 let mut out = self
794 .content
795 .as_ref()
796 .map(MessageContent::as_text)
797 .unwrap_or_default();
798 if let Some(calls) = &self.tool_calls {
799 for call in calls {
800 out.push_str(&format!(
801 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
802 call.function.name, call.function.arguments
803 ));
804 }
805 }
806 out
807 }
808}
809
810#[derive(Debug, Clone, Deserialize)]
811pub(crate) struct ToolCallIn {
812 #[serde(default)]
813 #[allow(dead_code)]
814 id: String,
815 #[serde(rename = "type", default)]
816 #[allow(dead_code)]
817 kind: String,
818 function: ToolCallFunctionIn,
819}
820
821#[derive(Debug, Clone, Deserialize)]
822struct ToolCallFunctionIn {
823 name: String,
824 arguments: String,
829}
830
831#[derive(Debug, Clone, Deserialize)]
834struct ToolDef {
835 #[serde(rename = "type", default)]
836 #[allow(dead_code)]
837 kind: String,
838 function: ToolFunctionDef,
839}
840
841#[derive(Debug, Clone, Deserialize)]
842struct ToolFunctionDef {
843 name: String,
844 #[serde(default)]
845 description: Option<String>,
846 #[serde(default)]
847 parameters: Option<serde_json::Value>,
848}
849
850#[derive(Debug, Clone, Deserialize)]
861#[serde(untagged)]
862enum ToolChoice {
863 Mode(String),
864 Specific(serde_json::Value),
865}
866
867#[derive(Deserialize)]
870#[serde(untagged)]
871enum StopParam {
872 One(String),
873 Many(Vec<String>),
874}
875
876#[derive(Deserialize)]
877struct ChatCompletionRequest {
878 model: String,
879 messages: Vec<ChatMessage>,
880 #[serde(default = "default_max_tokens")]
881 max_tokens: usize,
882 #[serde(default)]
883 temperature: Option<f32>,
884 #[serde(default)]
885 top_p: Option<f32>,
886 #[serde(default)]
891 min_p: Option<f32>,
892 #[serde(default)]
893 top_k: Option<usize>,
894 #[serde(default)]
895 repetition_penalty: Option<f32>,
896 #[serde(flatten)]
900 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
901 #[serde(flatten)]
905 unimplemented: crate::unimplemented_fields::UnimplementedFields,
906 #[serde(default)]
907 seed: Option<u64>,
908 #[serde(default)]
909 stop: Option<StopParam>,
910 #[serde(default)]
911 stream: Option<bool>,
912 #[serde(default)]
925 stream_resumable: Option<bool>,
926 #[serde(default)]
937 ignore_eos: Option<bool>,
938 #[serde(default)]
939 tools: Vec<ToolDef>,
940 #[serde(default)]
941 tool_choice: Option<ToolChoice>,
942 #[serde(default)]
950 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
951 #[serde(default)]
959 reasoning_effort: Option<String>,
960 #[serde(default)]
964 thinking: Option<ThinkingSwitch>,
965 #[serde(default)]
970 session_id: Option<String>,
971 #[serde(default, deserialize_with = "continuation::deserialize")]
979 continue_final_message: continuation::ContinueFinalMessage,
980 #[serde(default, alias = "thinking_budget_tokens")]
988 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
989 #[serde(default)]
991 logprobs: Option<bool>,
992 #[serde(default)]
993 top_logprobs: Option<u32>,
994 #[serde(default)]
995 presence_penalty: Option<f32>,
996 #[serde(default)]
997 frequency_penalty: Option<f32>,
998 #[serde(default)]
999 response_format: Option<serde_json::Value>,
1000 #[serde(default)]
1007 logit_bias: Option<serde_json::Value>,
1008 #[serde(default)]
1012 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1013 #[serde(default)]
1023 samplers: Option<serde_json::Value>,
1024 #[serde(default)]
1034 grammar: Option<String>,
1035}
1036
1037const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1049
1050#[derive(Debug, Clone, Deserialize)]
1052pub(crate) struct ThinkingSwitch {
1053 #[serde(rename = "type")]
1054 pub(crate) kind: String,
1055}
1056
1057const THINKING_KWARG_KEYS: [&str; 4] = [
1061 "enable_thinking",
1062 "thinking",
1063 "thinking_mode",
1064 "reasoning_effort",
1065];
1066
1067const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1071
1072fn default_max_tokens() -> usize {
1073 DEFAULT_CHAT_MAX_TOKENS
1074}
1075
1076impl ChatCompletionRequest {
1077 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1086 let mut knobs = SamplingKnobs {
1087 temperature: self.temperature,
1088 top_p: self.top_p,
1089 min_p: self.min_p,
1090 top_k: self.top_k,
1091 repetition_penalty: self.repetition_penalty,
1092 presence_penalty: self.presence_penalty,
1093 frequency_penalty: self.frequency_penalty,
1094 penalty_last_n: None,
1098 sampler_order: unsupported_sampling::parse_sampler_order(
1099 self.samplers.as_ref(),
1100 "/v1/chat/completions",
1101 )?,
1102 ..SamplingKnobs::default()
1103 };
1104 self.extra_samplers.apply(&mut knobs);
1105 Ok(knobs)
1106 }
1107
1108 fn sampling_params(
1109 &self,
1110 model: crate::sampling_knobs::SamplerModel<'_>,
1111 ) -> Result<SamplingParams, ApiError> {
1112 self.sampling_knobs()?.resolve(model).map_err(|e| {
1113 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1114 })
1115 }
1116
1117 fn stop_sequences(&self) -> Vec<String> {
1118 self.stop
1119 .as_ref()
1120 .map(|s| match s {
1121 StopParam::One(v) => vec![v.clone()],
1122 StopParam::Many(v) => v.clone(),
1123 })
1124 .unwrap_or_default()
1125 }
1126
1127 fn n_logprobs(&self) -> Result<Option<usize>, ApiError> {
1141 const MAX: u32 = 20;
1142 match (self.logprobs, self.top_logprobs) {
1143 (Some(true), Some(n)) if n > MAX => Err(invalid_request(
1144 &format!(
1145 "`top_logprobs` is {n}; this server reports at most {MAX} alternatives per \
1146 position, as upstream does"
1147 ),
1148 "top_logprobs",
1149 )),
1150 (Some(true), Some(n)) => Ok(Some(n as usize)),
1151 (Some(true), None) => Ok(Some(0)),
1155 (_, Some(_)) => Err(invalid_request(
1156 "`top_logprobs` requires `logprobs: true`",
1157 "top_logprobs",
1158 )),
1159 _ => Ok(None),
1160 }
1161 }
1162
1163 fn several_choices(&self) -> bool {
1168 self.unimplemented.n.is_some_and(|n| n > 1)
1169 }
1170
1171 fn tools_active(&self) -> bool {
1172 !self.tools.is_empty()
1173 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1174 }
1175
1176 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1188 match &self.tool_choice {
1189 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1190 Some(ToolChoice::Specific(value)) => {
1191 let name = value
1195 .get("function")
1196 .and_then(|f| f.get("name"))
1197 .or_else(|| value.get("name"))
1198 .and_then(|n| n.as_str());
1199 match name {
1200 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1201 None => Err(invalid_request(
1202 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1203 function.name",
1204 "tool_choice",
1205 )),
1206 }
1207 }
1208 _ => Ok(None),
1209 }
1210 }
1211
1212 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1214 self.tools
1215 .iter()
1216 .map(|t| tool_grammar::ToolSpec {
1217 name: &t.function.name,
1218 parameters: t.function.parameters.as_ref(),
1219 })
1220 .collect()
1221 }
1222
1223 fn thinking_direction(&self) -> Option<bool> {
1266 if let Some(switch) = &self.thinking {
1267 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1268 "disabled" => Some(false),
1269 "enabled" => Some(true),
1270 _ => None,
1273 };
1274 }
1275 let effort = self.reasoning_effort.as_ref()?;
1276 DISABLE_EFFORTS
1277 .contains(&effort.trim().to_ascii_lowercase().as_str())
1278 .then_some(false)
1279 }
1280
1281 fn resolve_template_kwargs(
1282 &self,
1283 template: &chat_template::PromptTemplate,
1284 ) -> serde_json::Map<String, serde_json::Value> {
1285 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1286 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1290
1291 if !caller_steered {
1292 match self.thinking_direction() {
1293 Some(false) => {
1294 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1295 kwargs.insert(k, v);
1296 }
1297 return kwargs;
1300 }
1301 Some(true) => {
1302 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1303 kwargs.insert(k, v);
1304 }
1305 }
1306 None => {}
1307 }
1308 if let Some(effort) = &self.reasoning_effort {
1309 kwargs
1310 .entry("reasoning_effort".to_string())
1311 .or_insert_with(|| serde_json::json!(effort));
1312 }
1313 }
1314
1315 let offered: Vec<serde_json::Value> = if self.tools_active() {
1316 self.tools.iter().map(chat_template::tool_json).collect()
1317 } else {
1318 Vec::new()
1319 };
1320 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1321 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1322 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1323 kwargs.entry(k).or_insert(v);
1324 }
1325 }
1326 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1327 crate::policy::effort::EffortMapping::Mapped(to) => {
1328 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1329 }
1330 crate::policy::effort::EffortMapping::Dropped => {
1331 tracing::debug!(
1332 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1333 enough, so its own default applies"
1334 );
1335 }
1336 crate::policy::effort::EffortMapping::Unchanged => {}
1337 }
1338 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1339 kwargs
1340 }
1341
1342 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1345 if self.max_tokens == 0 {
1352 return Err(invalid_request(
1353 "max_tokens must be at least 1",
1354 "max_tokens",
1355 ));
1356 }
1357 if let Some(switch) = &self.thinking {
1362 let kind = switch.kind.trim().to_ascii_lowercase();
1363 if kind != "enabled" && kind != "disabled" {
1364 return Err(invalid_request(
1365 "thinking.type must be \"enabled\" or \"disabled\"",
1366 "thinking.type",
1367 ));
1368 }
1369 }
1370 for msg in &self.messages {
1371 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1372 return Err(unsupported_feature(
1373 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1374 ));
1375 }
1376 }
1377 self.n_logprobs()?;
1381 self.unimplemented.refuse("/v1/chat/completions")?;
1385 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1386 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1390 let stated_grammar =
1405 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1406 if let Some(forced) = self.forced_tool_choice()? {
1414 if self.tools.is_empty() {
1415 return Err(invalid_request(
1416 "tool_choice forces a tool call, but no tools were offered",
1417 "tool_choice",
1418 ));
1419 }
1420 if let tool_grammar::Forced::Named(name) = forced {
1421 if !self.tools.iter().any(|t| t.function.name == name) {
1422 return Err(invalid_request(
1423 &format!(
1424 "tool_choice names {name:?}, which is not one of the tools offered"
1425 ),
1426 "tool_choice",
1427 ));
1428 }
1429 }
1430 if stated_grammar.is_some() {
1440 return Err(invalid_request(
1441 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1442 are two different constraints on the same generation; send one",
1443 "tool_choice",
1444 ));
1445 }
1446 if self.json_object_mode() {
1447 return Err(invalid_request(
1448 "a forced tool_choice cannot be combined with response_format json_object: \
1449 the tool-call markers are not JSON",
1450 "tool_choice",
1451 ));
1452 }
1453 }
1454 Ok(())
1455 }
1456
1457 fn effective_stop_sequences(&self) -> Vec<String> {
1464 let mut stop = self.stop_sequences();
1465 if self.tools_active() {
1466 stop.push("</tool_call>".to_string());
1467 }
1468 stop
1469 }
1470
1471 fn json_object_mode(&self) -> bool {
1472 self.response_format
1473 .as_ref()
1474 .and_then(|v| v.get("type"))
1475 .and_then(|v| v.as_str())
1476 == Some("json_object")
1477 }
1478}
1479
1480#[derive(Serialize)]
1481struct ChatCompletionChoice {
1482 index: usize,
1483 message: ChatCompletionResponseMessage,
1484 finish_reason: &'static str,
1485 #[serde(skip_serializing_if = "Option::is_none")]
1490 logprobs: Option<serde_json::Value>,
1491}
1492
1493#[derive(Serialize)]
1494struct ChatCompletionResponseMessage {
1495 role: &'static str,
1496 #[serde(skip_serializing_if = "Option::is_none")]
1497 content: Option<String>,
1498 #[serde(skip_serializing_if = "Option::is_none")]
1502 reasoning_content: Option<String>,
1503 #[serde(skip_serializing_if = "Option::is_none")]
1504 tool_calls: Option<Vec<ToolCallOut>>,
1505}
1506
1507#[derive(Serialize, Clone)]
1508struct ToolCallOut {
1509 id: String,
1510 #[serde(rename = "type")]
1511 kind: &'static str,
1512 function: ToolCallFunctionOut,
1513}
1514
1515#[derive(Serialize, Clone)]
1523struct ToolCallDelta {
1524 index: usize,
1525 #[serde(skip_serializing_if = "Option::is_none")]
1526 id: Option<String>,
1527 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1528 kind: Option<&'static str>,
1529 function: ToolCallFunctionDelta,
1530}
1531
1532#[derive(Serialize, Clone, Default)]
1533struct ToolCallFunctionDelta {
1534 #[serde(skip_serializing_if = "Option::is_none")]
1535 name: Option<String>,
1536 #[serde(skip_serializing_if = "Option::is_none")]
1539 arguments: Option<String>,
1540}
1541
1542impl ToolCallDelta {
1543 fn whole(index: usize, name: String, arguments: String) -> Self {
1545 ToolCallDelta {
1546 index,
1547 id: Some(format!("call_{index}")),
1548 kind: Some("function"),
1549 function: ToolCallFunctionDelta {
1550 name: Some(name),
1551 arguments: Some(arguments),
1552 },
1553 }
1554 }
1555
1556 fn opening(index: usize, name: String) -> Self {
1558 ToolCallDelta {
1559 index,
1560 id: Some(format!("call_{index}")),
1561 kind: Some("function"),
1562 function: ToolCallFunctionDelta {
1563 name: Some(name),
1564 arguments: Some(String::new()),
1565 },
1566 }
1567 }
1568
1569 fn arguments(index: usize, fragment: String) -> Self {
1571 ToolCallDelta {
1572 index,
1573 id: None,
1574 kind: None,
1575 function: ToolCallFunctionDelta {
1576 name: None,
1577 arguments: Some(fragment),
1578 },
1579 }
1580 }
1581}
1582
1583#[derive(Serialize, Clone)]
1584struct ToolCallFunctionOut {
1585 name: String,
1586 arguments: String,
1590}
1591
1592#[derive(Serialize)]
1593struct ChatCompletionResponse {
1594 id: String,
1595 request_id: String,
1601 object: &'static str,
1602 model: String,
1603 choices: Vec<ChatCompletionChoice>,
1604 usage: generate::Usage,
1609 frink_cache: &'static str,
1620}
1621
1622#[derive(Serialize)]
1623struct ChatCompletionChunkDelta {
1624 #[serde(skip_serializing_if = "Option::is_none")]
1625 role: Option<&'static str>,
1626 #[serde(skip_serializing_if = "Option::is_none")]
1627 content: Option<String>,
1628 #[serde(skip_serializing_if = "Option::is_none")]
1630 reasoning_content: Option<String>,
1631 #[serde(skip_serializing_if = "Option::is_none")]
1632 tool_calls: Option<Vec<ToolCallDelta>>,
1633}
1634
1635#[derive(Serialize)]
1636struct ChatCompletionChunkChoice {
1637 index: usize,
1638 delta: ChatCompletionChunkDelta,
1639 finish_reason: Option<&'static str>,
1640}
1641
1642#[derive(Serialize)]
1643struct ChatCompletionChunk {
1644 id: String,
1645 #[serde(skip_serializing_if = "Option::is_none")]
1652 request_id: Option<String>,
1653 object: &'static str,
1654 model: String,
1655 choices: Vec<ChatCompletionChunkChoice>,
1656 #[serde(skip_serializing_if = "Option::is_none")]
1659 usage: Option<generate::Usage>,
1660}
1661
1662async fn health(State(state): State<Arc<AppState>>) -> Response {
1668 let snapshot = state.detection.snapshot();
1669 let mut capabilities = snapshot.capabilities;
1670 let active = state.active();
1671
1672 capabilities.push(match active.as_deref() {
1675 None => frink_api::Capability::unavailable(
1680 frink_api::health::capability::REAL_WEIGHTS,
1681 frink_api::health::reason::MODEL_NOT_LOADED,
1682 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1683 ),
1684 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1685 frink_api::health::capability::REAL_WEIGHTS,
1686 frink_api::health::reason::MODEL_NOT_LOADED,
1687 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1688 checkpoint. Output from this model is noise.",
1689 ),
1690 Some(active) if active.encoder().is_some() => {
1700 let endpoints = active
1701 .encoder()
1702 .map(|e| encoder_endpoints(e))
1703 .unwrap_or_default();
1704 let served_by = match endpoints.is_empty() {
1705 true => "no endpoint in this build serves it".to_string(),
1706 false => format!("served by {}", endpoints.join(" and ")),
1707 };
1708 frink_api::Capability::available(
1709 frink_api::health::capability::REAL_WEIGHTS,
1710 format!(
1711 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1712 {served_by}; generation endpoints refuse it.",
1713 active.name(),
1714 ),
1715 )
1716 }
1717 Some(active) => frink_api::Capability::available(
1718 frink_api::health::capability::REAL_WEIGHTS,
1719 format!("Serving the real checkpoint '{}'.", active.name()),
1720 ),
1721 });
1722 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1723 frink_api::Capability::available(
1724 frink_api::health::capability::CONTINUOUS_BATCHING,
1725 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1726 "On by default on Metal. Concurrent requests share one batched decode worker."
1727 } else {
1728 "Concurrent requests share one batched decode step."
1729 },
1730 )
1731 } else if state.metal_private_decode_gate.is_some() {
1732 frink_api::Capability::unavailable(
1733 frink_api::health::capability::CONTINUOUS_BATCHING,
1734 frink_api::health::reason::DISABLED,
1735 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1736 )
1737 } else {
1738 frink_api::Capability::unavailable(
1739 frink_api::health::capability::CONTINUOUS_BATCHING,
1740 frink_api::health::reason::DISABLED,
1741 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1742 )
1743 });
1744
1745 let last_request_ms = state
1746 .last_request_ms
1747 .load(std::sync::atomic::Ordering::Relaxed);
1748 let uptime = state.started_at.elapsed();
1749 let health_state = if active.is_none() {
1757 frink_api::HealthState::Unavailable
1758 } else {
1759 snapshot.state
1760 };
1761 let body = frink_api::HealthResponse {
1762 state: health_state,
1763 reason: match health_state {
1764 frink_api::HealthState::Ready => None,
1765 frink_api::HealthState::Unavailable => {
1766 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1767 }
1768 frink_api::HealthState::Detecting => {
1769 Some(frink_api::health::reason::DETECTING.to_string())
1770 }
1771 },
1772 detail: match health_state {
1773 frink_api::HealthState::Ready => None,
1774 frink_api::HealthState::Unavailable => Some(
1775 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1776 .to_string(),
1777 ),
1778 frink_api::HealthState::Detecting => {
1779 Some("Probing available compute backends.".to_string())
1780 }
1781 },
1782 model: active
1783 .as_deref()
1784 .map(|active| frink_api::health::ModelSummary {
1785 id: active.name().to_string(),
1786 tokenizer: active.tokenizer_kind().to_string(),
1787 synthetic_weights: active.is_synthetic(),
1788 }),
1789 capabilities,
1790 version: env!("CARGO_PKG_VERSION").to_string(),
1791 pid: std::process::id(),
1792 uptime_seconds: uptime.as_secs_f64(),
1793 server_time_unix_ms: std::time::SystemTime::now()
1794 .duration_since(std::time::UNIX_EPOCH)
1795 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1796 .unwrap_or(0),
1797 last_request_age_seconds: (last_request_ms > 0)
1798 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1799 .map(|age| age.max(0.0)),
1800 };
1801
1802 let status =
1803 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1804 (status, Json(body)).into_response()
1805}
1806
1807async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1808 let Some(active) = state.active() else {
1812 return Json(serde_json::json!({ "object": "list", "data": [] }));
1813 };
1814 let mut model_entry = serde_json::json!({
1815 "id": active.name(),
1816 "object": "model",
1817 "frink_synthetic_weights": active.is_synthetic(),
1818 "frink_tokenizer": active.tokenizer_kind(),
1819 });
1820 if let Some(encoder) = active.encoder() {
1826 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1827 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1828 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1829 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1830 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1831 }
1832 if let Some(model) = active.generative_opt() {
1840 let parser_configured = active.reasoning_format().is_some();
1841 let gears = model.chat_template().think_gears(parser_configured);
1842 if !gears.is_empty() {
1843 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1844 if let Some(default) = &gears.default {
1845 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1846 }
1847 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1850 }
1851 }
1852 if let Some(mcp) = &state.mcp {
1853 model_entry["frink_mcp"] = mcp.models_metadata();
1854 }
1855 Json(serde_json::json!({
1856 "object": "list",
1857 "data": [model_entry]
1858 }))
1859}
1860
1861async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1875 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1876 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1877 let active = state.active();
1878 Json(serde_json::json!({
1879 "model": active.as_ref().map(|a| a.name()),
1880 "state": state
1881 .maintenance
1882 .lock()
1883 .unwrap_or_else(|p| p.into_inner())
1884 .state()
1885 .as_str(),
1886 "uptime_s": state.uptime().as_secs(),
1887 "throughput": {
1888 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1889 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1890 },
1891 "requests": {
1892 "active": state.cancels.live_count(),
1893 "completed": state.stats.recorded_total(),
1894 "p95_ms": state.stats.p95_duration_ms(),
1895 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1896 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1897 "completion_tokens_total": state.stats.tokens_generated_total(),
1898 },
1899 "pools": cache_admin::pool_gauges(&state),
1905 "memory": cache_admin::footprint_json(&state),
1908 }))
1909}
1910
1911#[derive(Deserialize)]
1912struct RequestsQuery {
1913 #[serde(default)]
1914 since: u64,
1915 #[serde(default = "default_requests_limit")]
1916 limit: usize,
1917}
1918
1919fn default_requests_limit() -> usize {
1920 stats::MAX_PAGE
1921}
1922
1923async fn recent_requests(
1931 State(state): State<Arc<AppState>>,
1932 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1933) -> Json<serde_json::Value> {
1934 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1935 Json(serde_json::json!({
1936 "requests": rows,
1937 "next_cursor": cursor,
1938 "missed": missed,
1939 "total": state.stats.recorded_total(),
1940 }))
1941}
1942
1943#[derive(Serialize)]
1944struct CombinedCacheStats {
1945 response_cache: response_cache::CacheStats,
1946 prefix_cache: Option<frink_models::PrefixCacheStats>,
1948}
1949
1950async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1951 Json(CombinedCacheStats {
1952 response_cache: lock_cache(&state.response_cache).stats(),
1953 prefix_cache: state
1954 .prefix_cache
1955 .as_ref()
1956 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1957 })
1958}
1959
1960async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1965 use std::sync::atomic::Ordering;
1966
1967 let cache_stats = lock_cache(&state.response_cache).stats();
1968 let active = state.active();
1969 let requests_total = state.requests_total.load(Ordering::Relaxed);
1970 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1971 let uptime = state.started_at.elapsed().as_secs_f64();
1972
1973 let body = format!(
1974 "# HELP frink_requests_total Total chat completion requests received.\n\
1975 # TYPE frink_requests_total counter\n\
1976 frink_requests_total {requests_total}\n\
1977 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1978 # TYPE frink_request_errors_total counter\n\
1979 frink_request_errors_total {errors_total}\n\
1980 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1981 # TYPE frink_cache_hits_total counter\n\
1982 frink_cache_hits_total {}\n\
1983 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1984 # TYPE frink_cache_misses_total counter\n\
1985 frink_cache_misses_total {}\n\
1986 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1987 # TYPE frink_cache_entries gauge\n\
1988 frink_cache_entries {}\n\
1989 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1990 # TYPE frink_synthetic_weights gauge\n\
1991 frink_synthetic_weights {}\n\
1992 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1993 # TYPE frink_uptime_seconds gauge\n\
1994 frink_uptime_seconds {uptime}\n",
1995 cache_stats.hits,
1996 cache_stats.misses,
1997 cache_stats.entries,
1998 active
2002 .as_ref()
2003 .map(|a| a.is_synthetic() as u8)
2004 .unwrap_or(0),
2005 );
2006
2007 let body = match active
2011 .as_ref()
2012 .and_then(|a| a.expert_store_stats())
2013 {
2014 Some(es) => format!(
2015 "{body}\
2016 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
2017 # TYPE frink_expert_cache_hits_total counter\n\
2018 frink_expert_cache_hits_total {}\n\
2019 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2020 # TYPE frink_expert_cache_misses_total counter\n\
2021 frink_expert_cache_misses_total {}\n\
2022 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
2023 # TYPE frink_expert_cache_evictions_total counter\n\
2024 frink_expert_cache_evictions_total {}\n\
2025 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2026 # TYPE frink_expert_cache_pass_throughs_total counter\n\
2027 frink_expert_cache_pass_throughs_total {}\n\
2028 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2029 # TYPE frink_expert_cache_bytes_read_total counter\n\
2030 frink_expert_cache_bytes_read_total {}\n\
2031 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2032 # TYPE frink_expert_cache_resident_bytes gauge\n\
2033 frink_expert_cache_resident_bytes {}\n",
2034 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2035 ),
2036 None => body,
2037 };
2038
2039 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2044 Some(batcher) => {
2045 let sched = batcher.stats();
2046 format!(
2047 "{body}\
2048 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2049 # TYPE frink_prefill_chunks_total counter\n\
2050 frink_prefill_chunks_total {}\n\
2051 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2052 # TYPE frink_prefill_tokens_total counter\n\
2053 frink_prefill_tokens_total {}\n\
2054 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2055 # TYPE frink_decode_steps_total counter\n\
2056 frink_decode_steps_total {}\n\
2057 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2058 # TYPE frink_scheduler_queue_depth gauge\n\
2059 frink_scheduler_queue_depth {}\n\
2060 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2061 # TYPE frink_scheduler_queue_rejected_total counter\n\
2062 frink_scheduler_queue_rejected_total {}\n\
2063 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2064 # TYPE frink_kv_blocks_total gauge\n\
2065 frink_kv_blocks_total {}\n\
2066 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2067 # TYPE frink_kv_blocks_free gauge\n\
2068 frink_kv_blocks_free {}\n\
2069 # HELP frink_kv_block_size Token positions per KV block.\n\
2070 # TYPE frink_kv_block_size gauge\n\
2071 frink_kv_block_size {}\n\
2072 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2073 # TYPE frink_kv_rejected_too_large_total counter\n\
2074 frink_kv_rejected_too_large_total {}\n\
2075 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2076 # TYPE frink_kv_rejected_context_length_total counter\n\
2077 frink_kv_rejected_context_length_total {}\n\
2078 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2079 # TYPE frink_scheduler_aborted_total counter\n\
2080 frink_scheduler_aborted_total {}\n\
2081 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2082 # TYPE frink_scheduler_max_seqs gauge\n\
2083 frink_scheduler_max_seqs {}\n\
2084 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2085 # TYPE frink_scheduler_prefill_chunk gauge\n\
2086 frink_scheduler_prefill_chunk {}\n",
2087 sched.prefill_chunks,
2088 sched.prefill_tokens,
2089 sched.decode_steps,
2090 sched.queue_depth,
2091 sched.queue_rejected,
2092 sched.kv_blocks_total,
2093 sched.kv_blocks_free,
2094 sched.kv_block_size,
2095 sched.kv_rejected_too_large,
2096 sched.kv_rejected_context_length,
2097 sched.aborted,
2098 sched.max_seqs,
2099 sched.prefill_chunk,
2100 )
2101 }
2102 None => body,
2103 };
2104
2105 (
2106 [(
2107 axum::http::header::CONTENT_TYPE,
2108 "text/plain; version=0.0.4",
2109 )],
2110 body,
2111 )
2112 .into_response()
2113}
2114
2115pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2116
2117pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2123 (
2124 StatusCode::BAD_REQUEST,
2125 Json(serde_json::json!({"error": {
2126 "message": message,
2127 "type": "invalid_request_error",
2128 "param": param,
2129 "code": null,
2130 }})),
2131 )
2132}
2133
2134pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2135 (
2136 StatusCode::NOT_IMPLEMENTED,
2137 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2138 )
2139}
2140
2141pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2142 let status = match e {
2143 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2144 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2148 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2153 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2159 StatusCode::SERVICE_UNAVAILABLE
2160 }
2161 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2165 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2169 };
2170 tracing::warn!("decode error: {e}");
2171 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2172 if let generate::DecodeError::KvBudgetExceeded {
2178 binding,
2179 estimated_bytes,
2180 limit_bytes,
2181 positions,
2182 positions_limit,
2183 ..
2184 } = &e
2185 {
2186 body["error"]["type"] = serde_json::json!("invalid_request_error");
2187 body["error"]["code"] = serde_json::json!(binding);
2188 body["error"]["binding"] = serde_json::json!(binding);
2189 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2190 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2191 body["error"]["positions"] = serde_json::json!(positions);
2192 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2193 }
2194 if let Some(secs) = e.retry_after_secs() {
2198 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2199 }
2200 (status, Json(body))
2201}
2202
2203pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2204 tracing::error!("generation task panicked: {e}");
2205 (
2206 StatusCode::INTERNAL_SERVER_ERROR,
2207 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2208 )
2209}
2210
2211#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2221 model: &Model,
2222 prompt: &str,
2223 params: &GenerationParams,
2224 kv_pool: Option<&generate::KvPoolConfig>,
2225 paged_kv: Option<&generate::PagedKvConfig>,
2226 prefix_cache: Option<&Mutex<PrefixCache>>,
2227 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2228 ceiling: Option<&budget::ContextCeiling>,
2229 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2230 mut emit: impl FnMut(&str),
2231 ) -> Result<generate::Generated, generate::DecodeError> {
2237 let synthetic = model.is_synthetic();
2238 let _lora_lease = lora::lease(model, params.lora.as_deref());
2243 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2244 let params = &{
2250 let mut resolved = params.clone();
2251 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2252 model.encode(text, SpecialTokens::Parse)
2253 });
2254 resolved.reasoning_budget = resolved
2260 .reasoning_budget
2261 .armed(resolved.reasoning, prompt, |text| {
2262 model.encode(text, SpecialTokens::Parse)
2263 })
2264 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2265 resolved
2266 };
2267 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2268 let _metal_private_guard =
2269 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2270 let (finishes, prompt_rows, prompt_ids, usage) = match model {
2271 Model::Gguf(m) => {
2272 if let Some(batcher) = continuous_batcher {
2273 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2274 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2275 let (finish, _generated_ids, text, usage) = if synthetic {
2276 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2277 } else {
2278 batcher.generate_streaming(
2279 tokens,
2280 params.clone(),
2281 m.stop_tokens.clone(),
2282 Some(|chunk: &str| {
2283 if !chunk.is_empty() {
2284 chunks[0].push(chunk.to_string());
2285 emit(chunk);
2286 }
2287 }),
2288 )?
2289 };
2290 if !text.is_empty() && chunks[0].is_empty() {
2291 chunks[0].push(text);
2292 }
2293 (vec![(finish, Vec::new())], Vec::new(), Vec::new(), usage)
2303 } else {
2304 generate::generate(
2305 &m.decoder,
2306 m.tokenizer.as_ref(),
2307 &m.stop_tokens,
2308 m.bos_id,
2309 prompt,
2310 params,
2311 kv_pool,
2312 paged_kv,
2313 prefix_cache,
2314 ceiling,
2315 |choice, chunk| {
2316 chunks[choice].push(chunk.to_string());
2317 if !synthetic && choice == 0 {
2321 emit(chunk);
2322 }
2323 },
2324 )?
2325 }
2326 }
2327 Model::Kimi(m) => generate::generate_engine(
2328 &m.engine,
2329 &m.tokenizer,
2330 &m.stop_tokens,
2331 None,
2332 prompt,
2333 params,
2334 |chunk| {
2335 chunks[0].push(chunk.to_string());
2336 if !synthetic {
2337 emit(chunk);
2338 }
2339 },
2340 )?,
2341 Model::Mla(m) => generate::generate_engine(
2342 &m.engine,
2343 &m.tokenizer,
2344 &m.stop_tokens,
2345 m.bos_id,
2346 prompt,
2347 params,
2348 |chunk| {
2349 chunks[0].push(chunk.to_string());
2350 if !synthetic {
2351 emit(chunk);
2352 }
2353 },
2354 )?,
2355 Model::Gemma4(m) => generate::generate_engine(
2356 &m.engine,
2357 &m.tokenizer,
2358 &m.stop_tokens,
2359 m.bos_id,
2360 prompt,
2361 params,
2362 |chunk| {
2363 chunks[0].push(chunk.to_string());
2364 if !synthetic {
2365 emit(chunk);
2366 }
2367 },
2368 )?,
2369 Model::Glm52(m) => generate::generate_engine(
2370 &m.engine,
2371 &m.tokenizer,
2372 &m.stop_tokens,
2373 m.bos_id,
2374 prompt,
2375 params,
2376 |chunk| {
2377 chunks[0].push(chunk.to_string());
2378 if !synthetic {
2379 emit(chunk);
2380 }
2381 },
2382 )?,
2383 };
2384
2385 let mut full = chunks[0].concat();
2386 if synthetic {
2387 full = format!(
2388 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2389 to serve a real model. Decoded ids -> {full:?}]"
2390 );
2391 emit(&full);
2392 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2393 emit(&full);
2394 }
2395
2396 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2402 let mut out: Vec<generate::GeneratedChoice> = finishes
2403 .into_iter()
2404 .zip(chunks.into_iter().map(|c| c.concat()))
2405 .map(|((finish, logprobs), text)| generate::GeneratedChoice {
2406 finish,
2407 text,
2408 logprobs,
2409 })
2410 .collect();
2411 if let Some(first) = out.first_mut() {
2412 if synthetic {
2420 first.logprobs.clear();
2421 }
2422 first.text = full;
2423 }
2424 Ok(generate::Generated {
2425 choices: out,
2426 prompt_rows,
2427 prompt_ids,
2428 usage,
2429 })
2430}
2431
2432#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2437 model: &Model,
2438 prompt: &str,
2439 params: &GenerationParams,
2440 kv_pool: Option<&generate::KvPoolConfig>,
2441 paged_kv: Option<&generate::PagedKvConfig>,
2442 prefix_cache: Option<&Mutex<PrefixCache>>,
2443 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2444 ceiling: Option<&budget::ContextCeiling>,
2445 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2446 ) -> Result<generate::Generated, generate::DecodeError> {
2449 run_generation_emit(
2450 model,
2451 prompt,
2452 params,
2453 kv_pool,
2454 paged_kv,
2455 prefix_cache,
2456 continuous_batcher,
2457 ceiling,
2458 metal_private_decode_gate,
2459 |_| {},
2460 )
2461}
2462
2463pub(crate) fn prompt_from_messages(
2474 messages: &[ChatMessage],
2475 template: &chat_template::PromptTemplate,
2476 tools: &[ToolDef],
2477 extra: serde_json::Map<String, serde_json::Value>,
2478) -> Result<String, ApiError> {
2479 let rendered = if tools.is_empty() || template.handles_tools() {
2480 template.render(messages, tools, extra)
2481 } else {
2482 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2483 with_preamble.push(ChatMessage {
2484 role: "system".to_string(),
2485 content: Some(MessageContent::Text(tool_preamble(tools))),
2486 tool_calls: None,
2487 tool_call_id: None,
2488 reasoning_content: None,
2489 });
2490 with_preamble.extend_from_slice(messages);
2491 template.render(&with_preamble, &[], extra)
2492 };
2493 rendered.map_err(template_error_response)
2494}
2495
2496fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2501 (
2502 StatusCode::BAD_REQUEST,
2503 Json(serde_json::json!({
2504 "error": {
2505 "message": format!("chat template failed to render: {err}"),
2506 "type": "invalid_request_error",
2507 "param": "messages",
2508 "code": null,
2509 }
2510 })),
2511 )
2512}
2513
2514fn tool_preamble(tools: &[ToolDef]) -> String {
2527 let mut out = String::from(
2528 "You can call tools to help answer the user. To call a tool, respond with \
2529 EXACTLY one line in this format and nothing else:\n\
2530 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2531 object matching that tool's parameters>}}</tool_call>\n\n\
2532 Available tools:\n",
2533 );
2534 for t in tools {
2535 out.push_str(&format!(
2536 "- {}: {}\n parameters (JSON schema): {}\n",
2537 t.function.name,
2538 t.function.description.as_deref().unwrap_or(""),
2539 t.function
2540 .parameters
2541 .as_ref()
2542 .map(|v| v.to_string())
2543 .unwrap_or_else(|| "{}".to_string()),
2544 ));
2545 }
2546 out
2547}
2548
2549fn tool_call_deltas(
2559 events: Vec<crate::policy::parser::ToolCallEvent>,
2560 opened: &std::cell::Cell<usize>,
2561) -> (String, Vec<ToolCallDelta>) {
2562 let mut text = String::new();
2563 let mut deltas = Vec::new();
2564 for event in events {
2565 match event {
2566 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2567 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2568 opened.set(opened.get().max(index + 1));
2569 deltas.push(ToolCallDelta::opening(index, name));
2570 }
2571 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2572 if !fragment.is_empty() {
2573 deltas.push(ToolCallDelta::arguments(index, fragment));
2574 }
2575 }
2576 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2577 }
2578 }
2579 (text, deltas)
2580}
2581
2582fn build_response_message(
2593 text: String,
2594 tools: &[ToolDef],
2595 posture: output::OutputPosture,
2596 base_finish: &'static str,
2597) -> (ChatCompletionResponseMessage, &'static str) {
2598 let parsed = output::parse_output(&text, tools, posture);
2599 let calls: Vec<ToolCallOut> = parsed
2600 .calls
2601 .into_iter()
2602 .enumerate()
2603 .map(|(index, call)| ToolCallOut {
2604 id: format!("call_{index}"),
2605 kind: "function",
2606 function: ToolCallFunctionOut {
2607 name: call.name,
2608 arguments: call.arguments,
2609 },
2610 })
2611 .collect();
2612 if !calls.is_empty() {
2613 return (
2614 ChatCompletionResponseMessage {
2615 role: "assistant",
2616 content: None,
2617 reasoning_content: parsed.reasoning,
2618 tool_calls: Some(calls),
2619 },
2620 "tool_calls",
2621 );
2622 }
2623 (
2624 ChatCompletionResponseMessage {
2625 role: "assistant",
2626 content: Some(parsed.content),
2627 reasoning_content: parsed.reasoning,
2628 tool_calls: None,
2629 },
2630 base_finish,
2631 )
2632}
2633
2634fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2639 let mut history = match &req.session_id {
2640 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2641 None => req.messages.clone(),
2642 };
2643 if req.json_object_mode() {
2644 inject_json_object_system_hint(&mut history);
2645 }
2646 history
2647}
2648
2649fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2650 const HINT: &str =
2651 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2652 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2653 match &mut sys.content {
2654 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2655 s.push_str("\n\n");
2656 s.push_str(HINT);
2657 }
2658 None => {
2659 sys.content = Some(MessageContent::Text(HINT.to_string()));
2660 }
2661 _ => {}
2662 }
2663 } else {
2664 messages.insert(
2665 0,
2666 ChatMessage {
2667 role: "system".to_string(),
2668 content: Some(MessageContent::Text(HINT.to_string())),
2669 tool_calls: None,
2670 tool_call_id: None,
2671 reasoning_content: None,
2672 },
2673 );
2674 }
2675}
2676
2677async fn chat_completions(
2678 State(state): State<Arc<AppState>>,
2679 headers: axum::http::HeaderMap,
2680 Json(req): Json<ChatCompletionRequest>,
2681) -> Response {
2682 let attribution = attribution::Attribution::from_headers(&headers);
2683 state
2684 .requests_total
2685 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2686 let started = std::time::Instant::now();
2687
2688 let request_id = frink_api::next_request_id();
2692 let stream = req.stream.unwrap_or(false);
2693
2694 let refusal = cache_admin::check_admission(&state)
2700 .err()
2701 .or_else(|| req.validate_supported_fields().err());
2702 if let Some(err) = refusal {
2703 state
2704 .request_errors_total
2705 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2706 let response = err.into_response();
2707 state.record_request(stats::Record {
2708 request_id: &request_id,
2709 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2710 model: state.active_model_name(),
2711 status: response.status().as_u16(),
2712 stream,
2713 duration_ms: started.elapsed().as_millis() as u64,
2714 usage: None,
2715 attribution: &attribution,
2716 });
2717 return response;
2718 }
2719
2720 let response = if stream {
2721 chat_completions_stream(
2722 Arc::clone(&state),
2723 req,
2724 request_id.clone(),
2725 started,
2726 attribution.clone(),
2727 )
2728 .await
2729 .into_response()
2730 } else {
2731 chat_completions_full(
2732 Arc::clone(&state),
2733 req,
2734 request_id.clone(),
2735 started,
2736 attribution.clone(),
2737 )
2738 .await
2739 .into_response()
2740 };
2741
2742 if response.status().is_client_error() || response.status().is_server_error() {
2743 state
2744 .request_errors_total
2745 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2746 state.record_request(stats::Record {
2750 request_id: &request_id,
2751 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2752 model: state.active_model_name(),
2755 status: response.status().as_u16(),
2756 stream,
2757 duration_ms: started.elapsed().as_millis() as u64,
2758 usage: None,
2759 attribution: &attribution,
2760 });
2761 }
2762 state.mark_request_finished();
2763
2764 response
2765}
2766
2767async fn chat_completions_full(
2768 state: Arc<AppState>,
2769 req: ChatCompletionRequest,
2770 request_id: String,
2771 started: std::time::Instant,
2772 attribution: attribution::Attribution,
2773) -> Result<Json<ChatCompletionResponse>, ApiError> {
2774 let tools_active = req.tools_active();
2775 let active = state.require_active()?;
2779 let history = resolve_history(&state, &req);
2780 let template = active.generative()?.chat_template();
2781 let kwargs = req.resolve_template_kwargs(&template);
2782 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2783 let mut params =
2791 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2792 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2793 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2794
2795 let mut generated_logprobs: Vec<crate::sampling_loop::PerTokenProbs> = Vec::new();
2799 let n_logprobs = req.n_logprobs()?;
2802 let decode_piece = |id: usize| active.decode_any(&[id]);
2804 let (completion, cache_status) = if let Some(cached) = key
2805 .as_ref()
2806 .and_then(|key| lock_cache(&state.response_cache).get(key))
2807 {
2808 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2809 (cached, "hit")
2810 } else {
2811 let produced = decode_task::buffered(
2812 decode_task::DecodeHandles::take(&state, &active)?,
2813 prompt.clone(),
2814 params,
2815 )
2816 .await?;
2817 let usage = produced.usage;
2818 let choices = produced.choices;
2819
2820 generated_logprobs = choices.iter().map(|c| c.logprobs.clone()).collect();
2825 let completion = response_cache::CachedCompletion {
2826 choices: choices.into_iter().map(|c| (c.finish, c.text)).collect(),
2827 usage,
2828 };
2829 let cache_status = match key {
2835 Some(key) => match completion.clone().cacheable() {
2839 Some(cacheable) => {
2840 tracing::debug!("cache miss for key {}", key.digest());
2841 lock_cache(&state.response_cache).put(key, cacheable);
2842 "miss"
2843 }
2844 None => "skip",
2845 },
2846 None => "skip",
2847 };
2848 (completion, cache_status)
2849 };
2850 let content = completion.first_text().to_string();
2853
2854 if req.json_object_mode() {
2855 json_mode::validate_json_object_output(&content)?;
2856 }
2857
2858 if let Some(id) = &req.session_id {
2862 state.sessions.store_reply(
2863 id,
2864 ChatMessage {
2865 role: "assistant".to_string(),
2866 content: Some(MessageContent::Text(content.clone())),
2867 tool_calls: None,
2868 tool_call_id: None,
2869 reasoning_content: None,
2870 },
2871 );
2872 }
2873
2874 let posture = output::OutputPosture::resolve_full(
2879 active.reasoning_format(),
2880 active.tool_call_format(),
2881 &prompt,
2882 );
2883 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2884 let wanted = req.unimplemented.n.unwrap_or(1).max(1) as usize;
2888 let ranked: Vec<(generate::FinishReason, String)> = if completion.choices.len() > wanted {
2889 let scored: Vec<crate::generate::GeneratedChoice> = completion
2890 .choices
2891 .into_iter()
2892 .zip(
2893 generated_logprobs
2894 .iter()
2895 .cloned()
2896 .chain(std::iter::repeat(Vec::new())),
2897 )
2898 .map(
2899 |((finish, text), logprobs)| crate::generate::GeneratedChoice {
2900 finish,
2901 text,
2902 logprobs,
2903 },
2904 )
2905 .collect();
2906 let best = crate::best_of::take_best(scored, wanted);
2907 generated_logprobs = best.iter().map(|c| c.logprobs.clone()).collect();
2908 best.into_iter().map(|c| (c.finish, c.text)).collect()
2909 } else {
2910 completion.choices
2911 };
2912 let rendered: Vec<ChatCompletionChoice> = ranked
2913 .into_iter()
2914 .enumerate()
2915 .map(|(index, (finish, text))| {
2916 let (message, finish_reason) =
2917 build_response_message(text, tools, posture, finish.as_str());
2918 ChatCompletionChoice {
2919 index,
2920 message,
2921 finish_reason,
2922 logprobs: n_logprobs.map(|k| {
2923 crate::logprobs::render_chat(
2924 generated_logprobs.get(index).unwrap_or(&Vec::new()),
2925 Some(k),
2926 &decode_piece,
2927 )
2928 }),
2929 }
2930 })
2931 .collect();
2932
2933 state.record_request(stats::Record {
2934 request_id: &request_id,
2935 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2936 model: Some(active.name().to_string()),
2939 status: 200,
2940 stream: false,
2941 duration_ms: started.elapsed().as_millis() as u64,
2942 usage: Some(&completion.usage),
2943 attribution: &attribution,
2944 });
2945
2946 Ok(Json(ChatCompletionResponse {
2947 id: request_id.clone(),
2948 request_id,
2949 object: "chat.completion",
2950 model: req.model,
2951 choices: rendered,
2952 usage: completion.usage,
2953 frink_cache: cache_status,
2954 }))
2955}
2956
2957async fn chat_completions_stream(
2958 state: Arc<AppState>,
2959 req: ChatCompletionRequest,
2960 request_id: String,
2961 started: std::time::Instant,
2962 attribution: attribution::Attribution,
2963) -> Result<Response, ApiError> {
2964 if req.several_choices() {
2974 return Err(unsupported_feature(
2975 "`n` > 1 with `stream` is not implemented: the choices would arrive one after \
2976 another rather than interleaved by `choices[].index`. Send the request without \
2977 `stream`, which serves `n` on this route.",
2978 ));
2979 }
2980 let tools_active = req.tools_active();
2981 let active = state.require_active()?;
2985 let history = resolve_history(&state, &req);
2986 let template = active.generative()?.chat_template();
2987 let kwargs = req.resolve_template_kwargs(&template);
2988 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2989 let model_name = req.model.clone();
2990 let session_id = req.session_id.clone();
2991 let sessions = state.sessions.clone();
2992
2993 let model = Arc::clone(active.generative()?);
2994 let kv_pool = state.kv_pool.clone();
2995 let paged_kv = state.paged_kv.clone();
2996 let prefix_cache = state.prefix_cache.clone();
2997 let batcher = active.batcher.clone();
2998 let ceiling = active.ceiling.clone();
2999 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
3000 let mut params =
3001 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
3002 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
3003 let stats_state = Arc::clone(&state);
3004 let served_model = active.name().to_string();
3007 let posture = output::OutputPosture::resolve_full(
3011 active.reasoning_format(),
3012 active.tool_call_format(),
3013 &prompt,
3014 );
3015 let offered_tools: Vec<ToolDef> = if tools_active {
3018 req.tools.clone()
3019 } else {
3020 Vec::new()
3021 };
3022
3023 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
3028 params.cancel = Some(cancel_token.clone());
3029
3030 let overlap = true;
3039
3040 let slot = req
3044 .stream_resumable
3045 .unwrap_or(false)
3046 .then(|| state.streams.register(&request_id));
3047 let emitter = resume::Emitter::new(slot);
3048
3049 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3050 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3054 id: request_id.clone(),
3055 request_id: None,
3056 object: "chat.completion.chunk",
3057 model: model_name.clone(),
3058 choices: vec![ChatCompletionChunkChoice {
3059 index: 0,
3060 delta: ChatCompletionChunkDelta {
3061 role: None,
3062 content: None,
3063 reasoning_content: None,
3064 tool_calls: None,
3065 },
3066 finish_reason: None,
3067 }],
3068 usage: None,
3069 });
3070
3071 tokio::task::spawn_blocking(move || {
3072 let _cancel_guard = cancel_guard;
3075 let tx_chunks = tx.clone();
3076 let orphan_timeout = sse::orphan_timeout_from_env();
3081 let mut first = true;
3082 let head_request_id = request_id.clone();
3083 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3093 Rc::new(RefCell::new(posture.reasoning_parser()));
3094 let emit_reasoning = Rc::clone(&stream_reasoning);
3095 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3099 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3100 );
3101 let emit_tools = Rc::clone(&stream_tools);
3102 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3106 let emit_streamed_calls = Rc::clone(&streamed_calls);
3107 let result = run_generation_emit(
3108 &model,
3109 &prompt,
3110 ¶ms,
3111 kv_pool.as_ref(),
3112 paged_kv.as_ref(),
3113 prefix_cache.as_deref(),
3114 batcher.as_ref(),
3115 ceiling.as_deref(),
3116 metal_private_decode_gate.as_deref(),
3117 |chunk| {
3118 if !overlap || chunk.is_empty() {
3119 return;
3120 }
3121 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3122 Some(parser) => {
3123 let delta = parser.push(chunk);
3124 (delta.reasoning, delta.content)
3125 }
3126 None => (String::new(), chunk.to_string()),
3127 };
3128 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3132 Some(parser) => {
3133 let (text, calls) =
3134 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3135 (text, calls)
3136 }
3137 None => (content, Vec::new()),
3138 };
3139 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3142 return;
3143 }
3144 let role = if first { Some("assistant") } else { None };
3145 let request_id = first.then(|| head_request_id.clone());
3146 first = false;
3147 let payload = ChatCompletionChunk {
3148 id: head_request_id.clone(),
3149 request_id,
3150 object: "chat.completion.chunk",
3151 model: model_name.clone(),
3152 choices: vec![ChatCompletionChunkChoice {
3153 index: 0,
3154 delta: ChatCompletionChunkDelta {
3155 role,
3156 content: (!content.is_empty()).then_some(content),
3157 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3158 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3159 },
3160 finish_reason: None,
3161 }],
3162 usage: None,
3163 };
3164 if let Err(why) =
3173 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3174 {
3175 if why == sse::SendFailure::Orphaned {
3176 tracing::warn!(
3177 "SSE stream {head_request_id} accepted nothing for the orphan \
3178 deadline; treating it as abandoned"
3179 );
3180 }
3181 if !emitter.is_resumable() {
3193 cancel_token.cancel();
3194 }
3195 }
3196 },
3197 );
3198
3199 let mut pending_request_id = first.then(|| request_id.clone());
3205
3206 match result {
3207 Ok(generated) => {
3210 let usage = generated.usage;
3211 let one = generated
3212 .choices
3213 .into_iter()
3214 .next()
3215 .expect("a generation produces at least one choice");
3216 let (finish, full_text) = (one.finish, one.text);
3217 if let Some(id) = &session_id {
3218 sessions.store_reply(
3219 id,
3220 ChatMessage {
3221 role: "assistant".to_string(),
3222 content: Some(MessageContent::Text(full_text.clone())),
3223 tool_calls: None,
3224 tool_call_id: None,
3225 reasoning_content: None,
3226 },
3227 );
3228 }
3229 let mut streamed_finish: Option<&'static str> = None;
3235 if overlap {
3236 let tail = stream_reasoning
3237 .borrow_mut()
3238 .as_mut()
3239 .map(|parser| parser.flush())
3240 .unwrap_or_default();
3241 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3242 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3243 let mut events = parser.push(&content);
3244 events.extend(parser.finish());
3245 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3246 content = text;
3247 tool_calls = calls;
3248 }
3249 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3250 let payload = ChatCompletionChunk {
3251 id: request_id.clone(),
3252 request_id: pending_request_id.take(),
3253 object: "chat.completion.chunk",
3254 model: model_name.clone(),
3255 choices: vec![ChatCompletionChunkChoice {
3256 index: 0,
3257 delta: ChatCompletionChunkDelta {
3258 role: None,
3259 content: (!content.is_empty()).then_some(content),
3260 reasoning_content: (!tail.reasoning.is_empty())
3261 .then_some(tail.reasoning),
3262 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3263 },
3264 finish_reason: None,
3265 }],
3266 usage: None,
3267 };
3268 let _ =
3269 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3270 }
3271 if streamed_calls.get() > 0 {
3272 streamed_finish = Some("tool_calls");
3273 }
3274 } else {
3275 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3278 let tool_calls: Vec<ToolCallDelta> = parsed
3279 .calls
3280 .iter()
3281 .enumerate()
3282 .map(|(index, call)| {
3283 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3284 })
3285 .collect();
3286 if !tool_calls.is_empty() {
3287 streamed_finish = Some("tool_calls");
3288 }
3289 if !tool_calls.is_empty()
3290 || !parsed.content.is_empty()
3291 || parsed.reasoning.is_some()
3292 {
3293 let payload = ChatCompletionChunk {
3294 id: request_id.clone(),
3295 request_id: pending_request_id.take(),
3296 object: "chat.completion.chunk",
3297 model: model_name.clone(),
3298 choices: vec![ChatCompletionChunkChoice {
3299 index: 0,
3300 delta: ChatCompletionChunkDelta {
3301 role: Some("assistant"),
3302 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3303 .then(|| parsed.content.clone()),
3304 reasoning_content: parsed.reasoning.clone(),
3305 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3306 },
3307 finish_reason: None,
3308 }],
3309 usage: None,
3310 };
3311 let _ =
3312 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3313 }
3314 }
3315 let final_finish_reason = match streamed_finish {
3319 Some(reason) if finish.as_str() != "length" => reason,
3320 _ => finish.as_str(),
3321 };
3322 let final_payload = ChatCompletionChunk {
3323 id: request_id.clone(),
3324 request_id: pending_request_id.take(),
3325 object: "chat.completion.chunk",
3326 model: model_name,
3327 choices: vec![ChatCompletionChunkChoice {
3328 index: 0,
3329 delta: ChatCompletionChunkDelta {
3330 role: None,
3331 content: None,
3332 reasoning_content: None,
3333 tool_calls: None,
3334 },
3335 finish_reason: Some(final_finish_reason),
3336 }],
3337 usage: Some(usage.clone()),
3338 };
3339 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3340 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3341 stats_state.record_request(stats::Record {
3346 request_id: &request_id,
3347 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3348 model: Some(served_model.clone()),
3349 status: 200,
3350 stream: true,
3351 duration_ms: started.elapsed().as_millis() as u64,
3352 usage: Some(&usage),
3353 attribution: &attribution,
3354 });
3355 }
3356 Err(e) => {
3357 tracing::warn!("decode error on streamed request {request_id}: {e}");
3358 stats_state.record_request(stats::Record {
3364 request_id: &request_id,
3365 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3366 model: Some(served_model.clone()),
3367 status: 500,
3368 stream: true,
3369 duration_ms: started.elapsed().as_millis() as u64,
3370 usage: None,
3371 attribution: &attribution,
3372 });
3373 let payload = ChatCompletionChunk {
3374 id: request_id.clone(),
3375 request_id: pending_request_id.take(),
3376 object: "chat.completion.chunk",
3377 model: model_name,
3378 choices: vec![ChatCompletionChunkChoice {
3379 index: 0,
3380 delta: ChatCompletionChunkDelta {
3381 role: Some("assistant"),
3382 content: Some(format!("[error: {e}]")),
3383 reasoning_content: None,
3384 tool_calls: None,
3385 },
3386 finish_reason: Some("stop"),
3387 }],
3388 usage: None,
3389 };
3390 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3391 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3392 }
3393 }
3394 drop(emitter);
3398 });
3399
3400 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3401 Ok((
3424 [(
3425 axum::http::HeaderName::from_static("x-accel-buffering"),
3426 axum::http::HeaderValue::from_static("no"),
3427 )],
3428 Sse::new(stream),
3429 )
3430 .into_response())
3431}
3432
3433fn protected_routes() -> Router<Arc<AppState>> {
3464 use frink_api::routes;
3465
3466 Router::new()
3467 .route(routes::V1_MODELS, get(list_models))
3468 .route(routes::V1_RESPONSES, post(responses::responses))
3472 .route(
3473 &axum_path(routes::V1_RESPONSE),
3474 get(responses::responses_get),
3475 )
3476 .route(
3477 &axum_path(routes::V1_RESPONSE_CANCEL),
3478 post(responses::responses_cancel),
3479 )
3480 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3481 .route(routes::V1_STATS, get(serving_stats))
3482 .route(routes::V1_REQUESTS, get(recent_requests))
3483 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3484 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3485 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3486 .route(
3487 routes::LORA_ADAPTERS,
3488 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3489 )
3490 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3491 .route(routes::V1_CANCEL, post(cancel_generation))
3495 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3500 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3501 .route(routes::V1_MESSAGES, post(anthropic::messages))
3502 .route(
3503 routes::V1_MESSAGES_COUNT_TOKENS,
3504 post(anthropic::count_tokens),
3505 )
3506 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3507 .route(routes::COMPLETION, post(completion::completion))
3512 .route(routes::COMPLETIONS, post(completion::completion))
3513 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3514 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3515 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3521 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3522 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3523 .route(routes::V1_RERANK, post(rerank::rerank))
3527 .route(routes::RERANK, post(rerank::rerank))
3528 .route(routes::CACHE_STATS, get(cache_stats))
3529 .route(routes::METRICS, get(metrics))
3530 .route(routes::ADMIN_MODELS, get(admin::models))
3535 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3536 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3537 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3538 .route(routes::ADMIN_TASKS, get(admin::tasks))
3539 .route(&admin::cancel_route(), post(admin::cancel_task))
3540 .route(routes::ADMIN_STATS, get(admin::stats))
3541 .merge(conversations::router())
3545}
3546
3547fn axum_path(template: &str) -> String {
3548 let mut out = String::with_capacity(template.len());
3549 let mut rest = template;
3550 while let Some(open) = rest.find('{') {
3551 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3552 break;
3553 };
3554 out.push_str(&rest[..open]);
3555 out.push(':');
3556 out.push_str(&rest[open + 1..close]);
3557 rest = &rest[close + 1..];
3558 }
3559 out.push_str(rest);
3560 out
3561}
3562
3563async fn cancel_generation(
3572 State(state): State<Arc<AppState>>,
3573 Json(req): Json<frink_api::CancelGenerationRequest>,
3574) -> Response {
3575 let cancelled = state.cancels.cancel(&req.request_id);
3576 let status = if cancelled {
3577 StatusCode::OK
3578 } else {
3579 StatusCode::NOT_FOUND
3580 };
3581 let detail = if cancelled {
3582 "the generation was asked to stop; it ends at its next token".to_string()
3583 } else {
3584 "no generation with that request_id is running -- it has already \
3585 finished, was never issued, or was served by a path that does \
3586 not register for cancellation"
3587 .to_string()
3588 };
3589 (
3590 status,
3591 Json(frink_api::CancelGenerationResponse {
3592 request_id: req.request_id,
3593 cancelled,
3594 detail,
3595 }),
3596 )
3597 .into_response()
3598}
3599
3600type Activated = (
3604 Loaded,
3605 Option<serving::batch::ContinuousBatcher>,
3606 Option<Arc<budget::ContextCeiling>>,
3607);
3608
3609fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3627 let mut batcher = serving::batch::BatcherConfig::from_env();
3628 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3629 return batcher;
3632 }
3633 let Some(path) = path else {
3634 return batcher;
3635 };
3636 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3642 let Some((priced, gguf_ctx, source)) = priced else {
3643 return batcher;
3644 };
3645 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3646 tracing::warn!(
3650 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3651 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3652 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3653 admit on a number you choose.",
3654 priced.weights_bytes,
3655 priced.device_budget_bytes,
3656 );
3657 return batcher;
3658 };
3659 tracing::info!("{source}");
3660 tracing::info!("{}", derived.fit);
3661 let adopted = budget::apply_derived(&mut batcher, &derived);
3662 if adopted.max_context {
3663 tracing::info!(
3664 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3665 override with FRINK_CB_MAX_CONTEXT",
3666 derived.max_context
3667 );
3668 }
3669 if adopted.kv_blocks {
3670 tracing::info!(
3671 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3672 derived.kv_blocks,
3673 batcher.kv_block_size
3674 );
3675 }
3676 if let Some(narrowed) = adopted.max_context_narrowed {
3677 tracing::info!(
3678 "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",
3679 batcher.kv_blocks.unwrap_or_default(),
3680 batcher.kv_block_size
3681 );
3682 }
3683 batcher
3684}
3685
3686pub(crate) fn activate_loaded_model(
3695 loaded: model::LoadedModel,
3696 enable_continuous_batching: bool,
3697 path: Option<&str>,
3698 paged_kv: Option<&generate::PagedKvConfig>,
3699) -> Activated {
3700 match loaded {
3701 model::LoadedModel::Gguf(g) => {
3702 let decoder = Arc::new(g.decoder);
3703 let tokenizer = Arc::new(g.tokenizer);
3704 let config = price_batcher_config(path);
3705 let shape =
3711 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3712 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3713 let batcher = if enable_continuous_batching {
3714 tracing::info!(
3715 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3716 (stop sequences use the same pending-buffer trim as the private generate loop)"
3717 );
3718 let tok = Arc::clone(&tokenizer);
3719 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3720 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3721 Arc::clone(&decoder),
3722 decode,
3723 config,
3724 Arc::clone(&ceiling),
3725 paged_kv.cloned(),
3726 ))
3727 } else {
3728 None
3729 };
3730 (
3731 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3732 decoder,
3733 tokenizer,
3734 stop_tokens: g.stop_tokens,
3735 bos_id: g.bos_id,
3736 is_synthetic: g.is_synthetic,
3737 chat_template: g.chat_template,
3738 }))),
3739 batcher,
3740 Some(ceiling),
3741 )
3742 }
3743 model::LoadedModel::Kimi(k) => (
3744 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3745 engine: k.engine,
3746 tokenizer: k.tokenizer,
3747 stop_tokens: k.stop_tokens,
3748 chat_template: k.chat_template,
3749 }))),
3750 None,
3751 None,
3752 ),
3753 model::LoadedModel::Mla(m) => (
3754 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3755 engine: m.engine,
3756 tokenizer: m.tokenizer,
3757 stop_tokens: m.stop_tokens,
3758 bos_id: m.bos_id,
3759 name: m.name,
3760 chat_template: m.chat_template,
3761 }))),
3762 None,
3763 None,
3764 ),
3765 model::LoadedModel::Gemma4(m) => (
3766 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3767 engine: m.engine,
3768 tokenizer: m.tokenizer,
3769 stop_tokens: m.stop_tokens,
3770 bos_id: m.bos_id,
3771 name: m.name,
3772 chat_template: m.chat_template,
3773 }))),
3774 None,
3775 None,
3776 ),
3777 model::LoadedModel::Glm52(g) => (
3778 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3779 engine: g.engine,
3780 tokenizer: g.tokenizer,
3781 stop_tokens: g.stop_tokens,
3782 bos_id: g.bos_id,
3783 name: g.name,
3784 chat_template: g.chat_template,
3785 }))),
3786 None,
3787 None,
3788 ),
3789 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3794 }
3795}
3796
3797struct StartupModels {
3804 loaded: model::LoadedModel,
3805 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3806}
3807
3808fn continuous_batching_env() -> Option<bool> {
3809 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3810 .ok()
3811 .map(|v| v.trim().to_ascii_lowercase())
3812 .as_deref()
3813 {
3814 None => None,
3815 Some("1" | "true" | "yes" | "on") => Some(true),
3816 Some("0" | "false" | "no" | "off") => Some(false),
3817 _ => None,
3818 }
3819}
3820
3821fn metal_private_decode_active() -> bool {
3822 #[cfg(feature = "metal")]
3823 {
3824 BUILT_WITH_METAL
3825 && frink_metal::attn::metal_attn_enabled()
3826 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3827 }
3828 #[cfg(not(feature = "metal"))]
3829 {
3830 false
3831 }
3832}
3833
3834fn continuous_batching_compatible(
3835 loaded: &model::LoadedModel,
3836 kv_pool: &Option<generate::KvPoolConfig>,
3837 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3838 paged_kv: &Option<generate::PagedKvConfig>,
3839) -> bool {
3840 matches!(loaded, model::LoadedModel::Gguf(_))
3841 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3842}
3843
3844fn resolve_continuous_batching_enabled(
3845 loaded: &model::LoadedModel,
3846 kv_pool: &Option<generate::KvPoolConfig>,
3847 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3848 paged_kv: &Option<generate::PagedKvConfig>,
3849) -> bool {
3850 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3851 return false;
3852 }
3853 match continuous_batching_env() {
3854 Some(true) => true,
3855 Some(false) => false,
3856 None => metal_private_decode_active(),
3857 }
3858}
3859
3860fn acquire_metal_private_decode_gate(
3861 gate: Option<&std::sync::Mutex<()>>,
3862 used_batcher: bool,
3863) -> Option<std::sync::MutexGuard<'_, ()>> {
3864 if used_batcher {
3865 None
3866 } else {
3867 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3868 }
3869}
3870
3871fn build_app_state(
3872 models: StartupModels,
3873 kv_pool: Option<generate::KvPoolConfig>,
3874 paged_kv: Option<generate::PagedKvConfig>,
3875 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3876 enable_continuous_batching: bool,
3877 mcp: Option<mcp::LoadedMcpConfig>,
3878 detection: Arc<health::Detection>,
3879) -> AppState {
3880 let StartupModels { loaded, embedding } = models;
3881 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3882 let (loaded, batcher, ceiling) = activate_loaded_model(
3883 loaded,
3884 enable_continuous_batching,
3885 configured_path.as_deref(),
3886 paged_kv.as_ref(),
3887 );
3888 let id = startup_model_id();
3894 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3895 {
3896 None
3897 } else {
3898 tracing::info!(
3899 "Metal private-loop decode will serialize concurrent requests until \
3900 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3901 );
3902 Some(Arc::new(std::sync::Mutex::new(())))
3903 };
3904 AppState {
3905 embedding,
3906 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3907 id,
3908 loaded,
3909 batcher,
3910 ceiling,
3911 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3912 }))),
3913 paged_kv,
3914 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3915 tasks: Arc::new(tasks::TaskRegistry::new()),
3916 cancels: Arc::new(cancel::CancelRegistry::new()),
3917 stats: stats::Stats::new(),
3918 streams: resume::StreamRegistry::new(),
3919 model_dir: admin::model_dirs().into_iter().next(),
3920 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3921 kv_pool,
3922 prefix_cache,
3923 sessions: session::SessionStore::new(),
3924 requests_total: std::sync::atomic::AtomicU64::new(0),
3925 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3926 started_at: std::time::Instant::now(),
3927 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3928 detection,
3929 mcp,
3930 continuous_batching_enabled: enable_continuous_batching,
3931 metal_private_decode_gate,
3932 loading_model: Mutex::new(None),
3933 last_load_error: Mutex::new(None),
3934 serving: Mutex::new(crate::stats::ServingStats::default()),
3935 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3936 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3937 started_unix: unix_now(),
3938 }
3939}
3940
3941fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3949 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3950 return Ok(None);
3951 };
3952 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3953 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3954 tracing::info!(
3955 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3956 model.name(),
3957 model.architecture(),
3958 model.n_embd(),
3959 model.pooling_type().name(),
3960 model.n_ctx_train(),
3961 );
3962 Ok(Some(Arc::new(model)))
3963}
3964
3965fn unix_now() -> u64 {
3970 std::time::SystemTime::now()
3971 .duration_since(std::time::UNIX_EPOCH)
3972 .map(|d| d.as_secs())
3973 .unwrap_or(0)
3974}
3975
3976fn startup_model_id() -> Option<String> {
3981 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3982 let configured = std::fs::canonicalize(&configured).ok()?;
3983 admin::discover(&admin::model_dirs())
3984 .into_iter()
3985 .find(|d| {
3986 std::fs::canonicalize(&d.path)
3987 .map(|p| p == configured)
3988 .unwrap_or(false)
3989 })
3990 .map(|d| d.id)
3991}
3992
3993fn init_cpu_pool() {
4002 match frink_core::threads::init_cpu_pool() {
4003 Some(n) => eprintln!(
4004 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
4005 frink_core::threads::perf_core_count()
4006 ),
4007 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
4008 }
4009}
4010
4011fn announce_ready(addr: SocketAddr, scheme: &str) {
4024 use std::io::Write;
4025 let ready =
4026 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
4027 let mut stdout = std::io::stdout().lock();
4028 let _ = writeln!(stdout, "{}", ready.to_line());
4029 let _ = stdout.flush();
4030}
4031
4032async fn shutdown_signal(exit_on_stdin_close: bool) {
4044 if !exit_on_stdin_close {
4045 std::future::pending::<()>().await;
4046 return;
4047 }
4048 let _ = tokio::task::spawn_blocking(|| {
4049 use std::io::Read;
4050 let mut sink = [0u8; 256];
4051 let mut stdin = std::io::stdin().lock();
4052 loop {
4053 match stdin.read(&mut sink) {
4054 Ok(0) => break,
4056 Ok(_) => continue,
4058 Err(e) => {
4059 tracing::warn!("stdin read failed ({e}); treating it as closed");
4060 break;
4061 }
4062 }
4063 }
4064 })
4065 .await;
4066 tracing::info!("stdin closed; shutting down");
4067}
4068
4069fn tokio_worker_threads() -> usize {
4074 std::env::var("FRINK_TOKIO_WORKERS")
4075 .ok()
4076 .and_then(|v| v.trim().parse::<usize>().ok())
4077 .filter(|n| *n > 0)
4078 .unwrap_or(2)
4079}
4080
4081fn install_ring_crypto_provider() {
4098 let _ = rustls::crypto::ring::default_provider().install_default();
4099}
4100
4101pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4107 if args.list_devices {
4108 frink_models::devices::print_available_devices();
4109 return Ok(());
4110 }
4111 apply_cli_overrides(&args)?;
4112
4113 let _instance = {
4117 use frink_core::instance::{register, InstancePolicy};
4118 let policy = if args.allow_multiple_instances {
4119 InstancePolicy::Multi
4120 } else {
4121 InstancePolicy::from_env_or(InstancePolicy::Single)
4122 };
4123 let model = std::env::var("FRINK_MODEL_PATH").ok();
4124 register(
4125 "server",
4126 model.as_deref(),
4127 frink_core::instance::current_backend(),
4128 policy,
4129 )
4130 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4131 };
4132
4133 let journal = journal::Journal::from_env();
4134 eprintln!(
4135 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4136 journal.path()
4137 );
4138 journal.append(&journal::Record::session_start(
4139 env!("CARGO_PKG_VERSION"),
4140 std::process::id(),
4141 ));
4142 journal::install_panic_hook(journal.clone());
4143
4144 let mcp_config_path = args.mcp_config.clone();
4145 let exit_on_stdin_close = args.exit_on_stdin_close
4146 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4147 .map(|v| v == "1")
4148 .unwrap_or(false);
4149
4150 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4154 init_cpu_pool();
4155
4156 let runtime = tokio::runtime::Builder::new_multi_thread()
4157 .worker_threads(tokio_worker_threads())
4158 .enable_all()
4159 .build()?;
4160 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4161
4162 let reason = match &result {
4163 Ok(()) => "normal".to_string(),
4164 Err(e) => e.to_string(),
4165 };
4166 journal.append(&journal::Record::session_exit(reason));
4167
4168 runtime.shutdown_background();
4173
4174 result
4175}
4176
4177async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4178 let _ = tracing_subscriber::fmt::try_init();
4184
4185 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4195 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4196 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4197 .map(|v| v == "1")
4198 .unwrap_or(false);
4199 if let Err(msg) =
4200 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4201 {
4202 anyhow::bail!(msg);
4203 }
4204
4205 let embedding_model = load_embedding_model()?;
4212
4213 let mut loaded = model::load()?;
4214 match &loaded {
4215 model::LoadedModel::Gguf(g) => tracing::info!(
4216 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4217 g.decoder.config.name,
4218 g.is_synthetic,
4219 g.tokenizer.kind()
4220 ),
4221 model::LoadedModel::Kimi(k) => tracing::info!(
4222 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4223 k.tokenizer.vocab_size()
4224 ),
4225 model::LoadedModel::Mla(m) => tracing::info!(
4226 "loaded MLA GGUF '{}' (tokenizer={})",
4227 m.name,
4228 m.tokenizer.kind()
4229 ),
4230 model::LoadedModel::Gemma4(m) => tracing::info!(
4231 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4232 m.name,
4233 m.tokenizer.kind()
4234 ),
4235 model::LoadedModel::Glm52(g) => tracing::info!(
4236 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4237 g.name,
4238 g.tokenizer.kind()
4239 ),
4240 model::LoadedModel::Encoder(_) => {}
4243 }
4244 let metal_default_moe_budget = {
4252 #[cfg(feature = "metal")]
4253 {
4254 frink_core::metal_dense_enabled()
4255 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4256 }
4257 #[cfg(not(feature = "metal"))]
4258 {
4259 false
4260 }
4261 };
4262 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4263 let budget: u64 = budget_str
4264 .parse()
4265 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4266 match &mut loaded {
4267 model::LoadedModel::Gguf(g) => {
4268 tracing::info!(
4269 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4270 (CUDA and/or Metal matvecs when built with the matching feature)"
4271 );
4272 g.decoder.gpu_vram_budget_bytes = Some(budget);
4273 }
4274 model::LoadedModel::Kimi(_) => {
4275 tracing::warn!(
4276 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4277 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4278 );
4279 }
4280 model::LoadedModel::Mla(_) => {
4281 tracing::warn!(
4282 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4283 FFN path only today; ignoring expert VRAM budget"
4284 );
4285 }
4286 model::LoadedModel::Gemma4(_) => {
4287 tracing::warn!(
4288 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4289 ignoring expert VRAM budget"
4290 );
4291 }
4292 model::LoadedModel::Glm52(_) => {
4293 tracing::warn!(
4294 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4295 GPU expert placement not wired yet; ignoring"
4296 );
4297 }
4298 model::LoadedModel::Encoder(_) => {
4299 tracing::warn!(
4300 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4301 it has no routed experts to place; ignoring"
4302 );
4303 }
4304 }
4305 } else if metal_default_moe_budget {
4306 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4310 if let model::LoadedModel::Gguf(g) = &mut loaded {
4311 tracing::info!(
4312 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4313 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4314 );
4315 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4316 }
4317 }
4318 #[cfg(feature = "cuda")]
4319 {
4320 if frink_core::cuda_dense_enabled() {
4321 tracing::info!(
4322 "CUDA dense matvec enabled for WeightMatrix::apply \
4323 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4324 );
4325 } else {
4326 tracing::info!(
4327 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4328 );
4329 }
4330 }
4331 #[cfg(feature = "metal")]
4332 {
4333 if frink_core::metal_dense_enabled() {
4334 tracing::info!(
4335 "Metal dense matvec enabled for WeightMatrix::apply \
4336 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4337 );
4338 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4339 Some("1") | Some("true") | Some("on") | Some("attn") => {
4340 tracing::info!(
4341 "Metal fused attention requested (FRINK_METAL_ATTN): \
4342 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4343 );
4344 }
4345 _ => {}
4346 }
4347 tracing::info!(
4348 "Metal greedy GPU argmax: temperature<=0 folds \
4349 final_norm+lm_head+argmax into the dense stack"
4350 );
4351 } else {
4352 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4353 }
4354 }
4355 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4366 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4367 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4368 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4369 panic!(
4370 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4371 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4372 );
4373 }
4374 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4375 (Ok(blocks), Ok(block_size), Err(_)) => {
4376 let total_blocks: usize = blocks
4377 .parse()
4378 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4379 let block_size: usize = block_size
4380 .parse()
4381 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4382 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4388 .ok()
4389 .map(|v| {
4390 v.parse()
4391 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4392 })
4393 .unwrap_or(0);
4394 tracing::info!(
4395 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4396 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4397 );
4398 Some(generate::KvPoolConfig {
4399 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4400 queue_wait: Duration::from_millis(queue_wait_ms),
4401 })
4402 }
4403 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4404 let block_size: usize = block_size
4405 .parse()
4406 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4407 let budget: u64 = byte_budget
4408 .parse()
4409 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4410 let cfg = match &loaded {
4411 model::LoadedModel::Gguf(g) => &g.decoder.config,
4412 model::LoadedModel::Kimi(_)
4413 | model::LoadedModel::Mla(_)
4414 | model::LoadedModel::Gemma4(_)
4415 | model::LoadedModel::Glm52(_)
4416 | model::LoadedModel::Encoder(_) => {
4417 panic!(
4418 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4419 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4420 );
4421 }
4422 };
4423 let bytes_per_block = block_size
4424 * cfg.kv_heads_all_layers()
4425 * (cfg.head_dim + cfg.v_head_dim())
4426 * std::mem::size_of::<f32>();
4427 assert!(
4428 bytes_per_block > 0,
4429 "derived KV block byte size must be positive (check model config and block size)"
4430 );
4431 let total_blocks = (budget as usize / bytes_per_block).max(1);
4432 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 from byte budget: {budget} bytes / \
4441 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4442 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4443 cfg.n_layers
4444 );
4445 Some(generate::KvPoolConfig {
4446 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4447 queue_wait: Duration::from_millis(queue_wait_ms),
4448 })
4449 }
4450 (Err(_), Err(_), Err(_)) => None,
4451 (Err(_), Ok(_), Err(_)) => panic!(
4452 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4453 (or unset all three to disable KV cache pooling)"
4454 ),
4455 (Ok(_), Ok(_), Ok(_)) => {
4456 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4457 }
4458 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4459 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4460 set together (or neither, to disable KV cache pooling)"
4461 ),
4462 };
4463 let paged_kv = match (
4469 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4470 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4471 ) {
4472 (Ok(blocks), Ok(block_size)) => {
4473 assert!(
4474 kv_pool.is_none(),
4475 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4476 mutually exclusive: both bound the same KV memory, by different means. \
4477 Set one."
4478 );
4479 let blocks_per_layer: usize = blocks
4493 .parse()
4494 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4495 let block_size: usize = block_size
4496 .parse()
4497 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4498 let gguf = match &loaded {
4499 model::LoadedModel::Gguf(g) => g,
4500 _ => panic!(
4501 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4502 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4503 ),
4504 };
4505 let cfg = &gguf.decoder.config;
4506 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4507 .ok()
4508 .map(|v| {
4509 v.parse()
4510 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4511 })
4512 .unwrap_or(0);
4513 tracing::info!(
4514 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4515 layer across {} layers, shared by all concurrent requests, \
4516 {queue_wait_ms}ms admission queue wait",
4517 cfg.n_layers
4518 );
4519 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4524 block_size,
4525 ))));
4526 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4535 crate::policy::parser::ToolCallFormat::infer(
4536 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4537 )
4538 .opener(),
4539 |text| {
4540 gguf.tokenizer
4541 .encode(text, SpecialTokens::Parse)
4542 .into_iter()
4543 .map(|t| t as u32)
4544 .collect()
4545 },
4546 );
4547 if let Some(id) = anchor_token {
4548 tracing::info!(
4549 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4550 window stops short of where its next turn rejoins"
4551 );
4552 }
4553 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4554 .ok()
4555 .map(|v| {
4556 v.parse()
4557 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4558 })
4559 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4560 if let Some(window) = cfg.uniform_sliding_window() {
4561 tracing::info!(
4562 "Paged KV window slide enabled: every layer slides by {window} every \
4563 {slide_interval} decode steps, so a request holds its prompt and a \
4564 window rather than its whole context"
4565 );
4566 } else if cfg.kv_block_window().is_some() {
4567 tracing::info!(
4568 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4569 and a page group holds one block in every layer"
4570 );
4571 }
4572 Some(generate::PagedKvConfig {
4573 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4576 queue_wait: Duration::from_millis(queue_wait_ms),
4577 radix,
4578 anchor_token,
4579 slide_interval,
4580 })
4581 }
4582 (Err(_), Err(_)) => None,
4583 _ => panic!(
4584 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4585 (or neither, to disable paged KV)"
4586 ),
4587 };
4588 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4595 let max_entries: usize = v
4596 .parse()
4597 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4598 if kv_pool.is_some() {
4599 tracing::warn!(
4600 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4601 caching will never be consulted while a KV pool is configured"
4602 );
4603 }
4604 assert!(
4611 paged_kv.is_none(),
4612 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4613 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4614 produce, so the cache could never hit. Set one."
4615 );
4616 tracing::info!(
4617 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4618 all requests"
4619 );
4620 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4621 });
4622 if matches!(
4623 loaded,
4624 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4625 ) && (kv_pool.is_some() || prefix_cache.is_some())
4626 {
4627 tracing::warn!(
4628 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4629 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4630 frink_models::engine's module docs"
4631 );
4632 }
4633 let enable_cb =
4634 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4635 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4636 tracing::info!(
4637 "continuous batching enabled by default on Metal for safe parallel serving \
4638 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4639 );
4640 }
4641 if continuous_batching_env() == Some(true)
4642 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4643 && (kv_pool.is_some() || prefix_cache.is_some())
4644 {
4645 tracing::warn!(
4646 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4647 (those modes keep the private generate path)"
4648 );
4649 }
4650 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4651 if let Ok(chunk) = n.parse::<usize>() {
4652 if chunk > 0 {
4653 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4654 }
4655 }
4656 }
4657 if matches!(
4658 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4659 Some("1")
4660 ) {
4661 tracing::warn!(
4662 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4663 (minimal spill; full layer offload still planned)"
4664 );
4665 }
4666
4667 let mcp = match mcp_config_path {
4668 Some(path) => {
4669 let loaded = mcp::load_mcp_config(&path)?;
4670 tracing::info!(
4671 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4672 loaded.path,
4673 loaded.servers.len()
4674 );
4675 Some(loaded)
4676 }
4677 None => None,
4678 };
4679
4680 let detection = health::Detection::spawn();
4684
4685 let state = Arc::new(build_app_state(
4686 StartupModels {
4687 loaded,
4688 embedding: embedding_model,
4689 },
4690 kv_pool,
4691 paged_kv,
4692 prefix_cache,
4693 enable_cb,
4694 mcp,
4695 detection,
4696 ));
4697
4698 use frink_api::routes;
4702
4703 let public = Router::new().route(routes::HEALTH, get(health));
4708
4709 let mut protected = protected_routes();
4710
4711 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4716 tracing::info!("API key auth enabled");
4717 let auth = limits::AuthConfig {
4718 api_key: Arc::new(key),
4719 };
4720 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4721 auth,
4722 limits::require_api_key,
4723 ));
4724 }
4725 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4726 let rpm: u32 = rpm
4727 .parse()
4728 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4729 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4730 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4731 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4732 limiter,
4733 limits::rate_limit,
4734 ));
4735 }
4736 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4744 let origins = security::parse_cors_origins(&spec)
4745 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4746 tracing::info!(
4747 "CORS enabled: {} allow-listed origin(s) ({})",
4748 origins.len(),
4749 spec
4750 );
4751 let cors = tower_http::cors::CorsLayer::new()
4752 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4753 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4754 .allow_headers([
4755 axum::http::header::CONTENT_TYPE,
4756 axum::http::header::AUTHORIZATION,
4757 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4763 axum::http::HeaderName::from_static("last-event-id"),
4767 ]);
4768 protected = protected.route_layer(cors);
4769 }
4770
4771 let app = public
4775 .merge(protected)
4776 .layer(axum::middleware::from_fn(limits::retry_after))
4777 .with_state(state);
4778
4779 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4785 install_ring_crypto_provider();
4786 match tls_paths {
4791 Some(paths) => {
4792 let config =
4793 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4794 .await
4795 .map_err(|e| {
4796 anyhow::anyhow!(
4797 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4798 paths.cert,
4799 paths.key
4800 )
4801 })?;
4802 let socket_addr: std::net::SocketAddr = addr
4803 .parse()
4804 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4805 let listener = std::net::TcpListener::bind(socket_addr)?;
4806 listener.set_nonblocking(true)?;
4814 let bound = listener.local_addr()?;
4815 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4816 announce_ready(bound, "https");
4817
4818 let handle = axum_server::Handle::new();
4819 let shutdown_handle = handle.clone();
4820 tokio::spawn(async move {
4821 shutdown_signal(exit_on_stdin_close).await;
4822 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4823 });
4824 axum_server::from_tcp_rustls(listener, config)?
4825 .handle(handle)
4826 .serve(app.into_make_service())
4827 .await?;
4828 }
4829 None => {
4830 let listener = tokio::net::TcpListener::bind(&addr).await?;
4831 let bound = listener.local_addr()?;
4832 tracing::info!("frink-server listening on {bound}");
4833 announce_ready(bound, "http");
4834 axum::serve(listener, app)
4835 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4836 .await?;
4837 }
4838 }
4839 Ok(())
4840}
4841
4842#[cfg(test)]
4843pub(crate) mod tests {
4844 use super::*;
4845 use frink_models::config::test_dense_fixture;
4846
4847 #[test]
4848 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4849 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4850 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4851 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4852 assert_eq!(parsed.port, 51999);
4853 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4854 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4856 }
4857
4858 fn test_model() -> Model {
4859 let cfg = test_dense_fixture();
4863 Model::Gguf(GgufModel {
4864 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4865 tokenizer: Arc::new(ServerTokenizer::Byte),
4866 stop_tokens: StopTokens::default(),
4867 bos_id: None,
4868 is_synthetic: true,
4869 chat_template: chat_template::PromptTemplate::plain(),
4870 })
4871 }
4872
4873 fn greedy_params(max_tokens: usize) -> GenerationParams {
4874 GenerationParams {
4875 cache_salt: None,
4876 prompt_logprobs: None,
4877 wants_logprobs: false,
4878 n: 1,
4879 reasoning: None,
4880 max_tokens,
4881 sampling: SamplingParams::default(),
4882 seed: 1,
4883 stop: Vec::new(),
4884 stop_token_ids: Vec::new(),
4885 json_object: false,
4886 grammar: None,
4887 cancel: None,
4888 ignore_eos: false,
4889 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4890 lora: None,
4891 }
4892 }
4893
4894 fn test_model_full_byte_vocab() -> Model {
4898 test_model_full_byte_vocab_with_eos(None)
4899 }
4900
4901 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4908 let mut cfg = test_dense_fixture();
4909 cfg.vocab_size = 256;
4910 Model::Gguf(GgufModel {
4911 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4912 tokenizer: Arc::new(ServerTokenizer::Byte),
4913 stop_tokens: StopTokens::from_eos(eos),
4914 bos_id: None,
4915 is_synthetic: true,
4916 chat_template: chat_template::PromptTemplate::plain(),
4917 })
4918 }
4919
4920 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4924 AppState {
4925 embedding: None,
4926 paged_kv: None,
4927 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4928 id: None,
4929 loaded: Loaded::Generative(Arc::new(model)),
4930 batcher: None,
4931 ceiling: None,
4932 checkpoint_path: None,
4933 }))),
4934 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4935 tasks: Arc::new(tasks::TaskRegistry::new()),
4936 cancels: Arc::new(cancel::CancelRegistry::new()),
4937 stats: stats::Stats::new(),
4938 streams: resume::StreamRegistry::new(),
4939 model_dir: None,
4940 response_cache: Mutex::new(response_cache),
4941 kv_pool: None,
4942 prefix_cache: None,
4943 sessions: session::SessionStore::new(),
4944 requests_total: std::sync::atomic::AtomicU64::new(0),
4945 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4946 started_at: std::time::Instant::now(),
4947 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4948 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4949 mcp: None,
4950 continuous_batching_enabled: false,
4951 metal_private_decode_gate: None,
4952 loading_model: Mutex::new(None),
4953 last_load_error: Mutex::new(None),
4954 serving: Mutex::new(crate::stats::ServingStats::default()),
4955 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4956 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4957 started_unix: unix_now(),
4958 }
4959 }
4960
4961 pub(crate) fn test_app() -> Router {
4970 test_app_with_state(Arc::new(test_state(
4971 test_model_full_byte_vocab(),
4972 ResponseCache::new(1000, Duration::from_secs(3600)),
4973 )))
4974 }
4975
4976 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4979 Router::new()
4989 .route(frink_api::routes::HEALTH, get(health))
4990 .merge(protected_routes())
4991 .with_state(state)
4992 }
4993
4994 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4995 let mut cfg = test_dense_fixture();
4996 cfg.name = name;
4997 cfg.vocab_size = vocab_size;
4998 Model::Gguf(GgufModel {
4999 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5000 tokenizer: Arc::new(ServerTokenizer::Byte),
5001 stop_tokens: StopTokens::default(),
5002 bos_id: None,
5003 is_synthetic: true,
5004 chat_template: chat_template::PromptTemplate::plain(),
5005 })
5006 }
5007
5008 fn model_with_template(name: &'static str, source: &str) -> Model {
5012 let mut cfg = test_dense_fixture();
5013 cfg.name = name;
5014 cfg.vocab_size = 256;
5015 Model::Gguf(GgufModel {
5016 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5017 tokenizer: Arc::new(ServerTokenizer::Byte),
5018 stop_tokens: StopTokens::default(),
5019 bos_id: None,
5020 is_synthetic: true,
5021 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
5022 Some(source),
5023 Some("qwen3"),
5024 false,
5025 true,
5026 None,
5027 None,
5028 ),
5029 })
5030 }
5031
5032 #[tokio::test]
5041 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
5042 let strict = "{% if messages | length > 1 %}\
5045 {{ raise_exception('this template takes one turn') }}\
5046 {% endif %}{{ messages[0].content }}";
5047 let state = Arc::new(test_state(
5048 model_with_template("strict", strict),
5049 ResponseCache::new(4, Duration::from_secs(60)),
5050 ));
5051 let app = test_app_with_state(state);
5052
5053 let (status, body) = post_json_uri(
5054 &app,
5055 "/v1/chat/completions",
5056 serde_json::json!({
5057 "model": "strict",
5058 "stream": true,
5059 "messages": [
5060 {"role": "user", "content": "one"},
5061 {"role": "user", "content": "two"},
5062 ],
5063 }),
5064 )
5065 .await;
5066 assert_eq!(status, StatusCode::BAD_REQUEST);
5067 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5068 assert!(
5069 body["error"]["message"]
5070 .as_str()
5071 .unwrap()
5072 .contains("one turn"),
5073 "the template's own message must reach the caller: {body}"
5074 );
5075
5076 let (status, _) = post_json_uri(
5078 &app,
5079 "/v1/chat/completions",
5080 serde_json::json!({
5081 "model": "strict",
5082 "stream": true,
5083 "max_tokens": 1,
5084 "messages": [{"role": "user", "content": "one"}],
5085 }),
5086 )
5087 .await;
5088 assert_eq!(status, StatusCode::OK);
5089 }
5090
5091 #[tokio::test]
5093 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5094 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5095 {% if reasoning_effort %}\
5096 {% if reasoning_effort not in ['low','medium','high'] %}\
5097 {{ raise_exception('bad effort') }}\
5098 {% endif %}[{{ reasoning_effort }}]\
5099 {% endif %}{{ messages[0].content }}";
5100 let state = Arc::new(test_state(
5101 model_with_template("thinker", reasoning),
5102 ResponseCache::new(4, Duration::from_secs(60)),
5103 ));
5104 let app = test_app_with_state(state);
5105 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5106 assert_eq!(status, StatusCode::OK);
5107 let entry = &models["data"][0];
5108 assert_eq!(
5109 entry["supported_reasoning_efforts"],
5110 serde_json::json!(["off", "low", "medium", "high"])
5111 );
5112 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5113 }
5114
5115 #[tokio::test]
5120 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5121 let app = test_app();
5122 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5123 let entry = &models["data"][0];
5124 assert!(entry.get("supported_reasoning_efforts").is_none());
5125 assert!(entry.get("default_reasoning_effort").is_none());
5126 }
5127
5128 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5129 Arc::new(ActiveModel {
5130 id: Some(name.to_string()),
5131 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5132 batcher: None,
5133 ceiling: None,
5134 checkpoint_path: None,
5135 })
5136 .tap_into(state)
5137 }
5138
5139 trait TapInto {
5141 fn tap_into(self, state: &AppState) -> Self;
5142 }
5143 impl TapInto for Arc<ActiveModel> {
5144 fn tap_into(self, state: &AppState) -> Self {
5145 state.swap_active(Some(Arc::clone(&self)));
5146 self
5147 }
5148 }
5149
5150 #[test]
5156 fn an_in_flight_request_keeps_the_model_it_started_on() {
5157 let state = test_state(
5158 named_test_model("model-a", 256),
5159 ResponseCache::new(4, Duration::from_secs(60)),
5160 );
5161
5162 let in_flight = state.active().expect("a model is loaded");
5165 assert_eq!(in_flight.name(), "model-a");
5166
5167 active_model(&state, "model-b");
5168
5169 assert_eq!(state.active().unwrap().name(), "model-b");
5171 assert_eq!(in_flight.name(), "model-a");
5173 let produced = run_generation(
5174 in_flight.generative().unwrap(),
5175 "hi",
5176 &greedy_params(3),
5177 None,
5178 None,
5179 None,
5180 None,
5181 None,
5182 None,
5183 )
5184 .expect("the old model must still decode after being swapped out");
5185 assert!(matches!(
5186 produced.choices[0].finish,
5187 FinishReason::Length | FinishReason::Stop
5188 ));
5189 }
5190
5191 #[test]
5196 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5197 let state = test_state(
5198 named_test_model("model-a", 256),
5199 ResponseCache::new(4, Duration::from_secs(60)),
5200 );
5201 let in_flight = state.active().expect("a model is loaded");
5202 let weights = Arc::clone(in_flight.generative().unwrap());
5203 assert!(Arc::strong_count(&weights) >= 2);
5204
5205 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5206 id: Some("model-b".to_string()),
5207 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5208 batcher: None,
5209 ceiling: None,
5210 checkpoint_path: None,
5211 })));
5212 drop(previous);
5213 assert!(Arc::strong_count(&weights) >= 2);
5215 drop(in_flight);
5216 assert_eq!(Arc::strong_count(&weights), 1);
5217 }
5218
5219 #[tokio::test]
5223 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5224 let state = Arc::new(test_state(
5225 named_test_model("model-a", 256),
5226 ResponseCache::new(4, Duration::from_secs(60)),
5227 ));
5228 let app = test_app_with_state(Arc::clone(&state));
5229
5230 let (status, body) = post_json_uri(
5231 &app,
5232 frink_api::routes::ADMIN_MODELS_UNLOAD,
5233 serde_json::json!({}),
5234 )
5235 .await;
5236 assert_eq!(status, StatusCode::OK);
5237 assert_eq!(body["ok"], true);
5238 assert!(body["active"].is_null());
5239 assert!(state.active().is_none());
5240
5241 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5242 assert_eq!(status, StatusCode::OK);
5243 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5244 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5245
5246 let (status, body) = post_json_uri(
5247 &app,
5248 "/v1/chat/completions",
5249 serde_json::json!({
5250 "model": "x",
5251 "messages": [{"role": "user", "content": "hi"}]
5252 }),
5253 )
5254 .await;
5255 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5256 assert_eq!(body["error"]["type"], "model_not_loaded");
5257 }
5258
5259 #[tokio::test]
5263 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5264 let state = Arc::new(test_state(
5265 named_test_model("model-a", 256),
5266 ResponseCache::new(4, Duration::from_secs(60)),
5267 ));
5268 let app = test_app_with_state(Arc::clone(&state));
5269 state.swap_active(None);
5270
5271 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5272 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5275 assert_eq!(body["state"], "unavailable");
5276 assert_eq!(body["reason"], "model_not_loaded");
5277 assert!(body["model"].is_null());
5278 let real_weights = body["capabilities"]
5279 .as_array()
5280 .unwrap()
5281 .iter()
5282 .find(|c| c["id"] == "real_weights")
5283 .cloned()
5284 .expect("real_weights is always reported");
5285 assert_eq!(real_weights["available"], false);
5286 assert_eq!(real_weights["reason"], "model_not_loaded");
5287 }
5288
5289 #[tokio::test]
5293 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5294 let app = test_app();
5295
5296 let (status, completion) = post_json_uri(
5297 &app,
5298 "/v1/chat/completions",
5299 serde_json::json!({
5300 "model": "x",
5301 "messages": [{"role": "user", "content": "hi"}],
5302 "max_tokens": 4
5303 }),
5304 )
5305 .await;
5306 assert_eq!(status, StatusCode::OK);
5307 let request_id = completion["request_id"].as_str().unwrap().to_string();
5308
5309 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5310 assert_eq!(status, StatusCode::OK);
5311 let recent = stats["recent"].as_array().unwrap();
5312 assert_eq!(recent.len(), 1);
5313 let row = &recent[0];
5314 assert_eq!(row["request_id"], request_id);
5315 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5316 assert_eq!(row["status"], 200);
5317 assert_eq!(row["stream"], false);
5318 assert!(row["duration_ms"].is_number());
5321 assert!(row["decode_ms"].is_number());
5322 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5323 assert_eq!(
5324 stats["tokens_prompt_total"].as_u64().unwrap(),
5325 row["prompt_tokens"].as_u64().unwrap()
5326 );
5327 }
5328
5329 #[tokio::test]
5333 async fn a_rejected_request_is_recorded_too() {
5334 let state = Arc::new(test_state(
5335 named_test_model("model-a", 256),
5336 ResponseCache::new(4, Duration::from_secs(60)),
5337 ));
5338 let app = test_app_with_state(Arc::clone(&state));
5339 state.swap_active(None);
5340
5341 let (status, _) = post_json_uri(
5342 &app,
5343 "/v1/chat/completions",
5344 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5345 )
5346 .await;
5347 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5348
5349 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5350 let recent = stats["recent"].as_array().unwrap();
5351 assert_eq!(recent.len(), 1);
5352 assert_eq!(recent[0]["status"], 503);
5353 assert_eq!(recent[0]["completion_tokens"], 0);
5354 assert!(recent[0]["decode_ms"].is_null());
5355 assert_eq!(stats["errors_total"], 1);
5356 }
5357
5358 async fn post_json_with_headers(
5362 app: &Router,
5363 uri: &str,
5364 body: serde_json::Value,
5365 headers: &[(&str, &str)],
5366 ) -> (StatusCode, serde_json::Value) {
5367 use http_body_util::BodyExt;
5368 use tower::ServiceExt;
5369
5370 let mut builder = axum::http::Request::builder()
5371 .method("POST")
5372 .uri(uri)
5373 .header("content-type", "application/json");
5374 for (name, value) in headers {
5375 builder = builder.header(*name, *value);
5376 }
5377 let response = app
5378 .clone()
5379 .oneshot(
5380 builder
5381 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5382 .unwrap(),
5383 )
5384 .await
5385 .unwrap();
5386 let status = response.status();
5387 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5388 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5389 (status, json)
5390 }
5391
5392 #[tokio::test]
5396 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5397 let app = test_app();
5398
5399 let (status, _) = post_json_uri(
5400 &app,
5401 frink_api::routes::V1_TOKENIZE,
5402 serde_json::json!({"prompt": "hello"}),
5403 )
5404 .await;
5405 assert_eq!(status, StatusCode::OK);
5406 let (status, _) = post_json_uri(
5407 &app,
5408 frink_api::routes::V1_DETOKENIZE,
5409 serde_json::json!({"tokens": [104, 105]}),
5410 )
5411 .await;
5412 assert_eq!(status, StatusCode::OK);
5413 let (status, _) = post_json_uri(
5414 &app,
5415 frink_api::routes::V1_EMBEDDINGS,
5416 serde_json::json!({"input": "hello"}),
5417 )
5418 .await;
5419 assert_eq!(status, StatusCode::OK);
5420
5421 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5422 let routes: Vec<&str> = stats["recent"]
5423 .as_array()
5424 .unwrap()
5425 .iter()
5426 .map(|row| row["route"].as_str().unwrap())
5427 .collect();
5428 for expected in [
5429 frink_api::routes::V1_TOKENIZE,
5430 frink_api::routes::V1_DETOKENIZE,
5431 frink_api::routes::V1_EMBEDDINGS,
5432 ] {
5433 assert!(
5434 routes.contains(&expected),
5435 "{expected} is missing: {routes:?}"
5436 );
5437 }
5438
5439 let row = |route: &str| {
5440 stats["recent"]
5441 .as_array()
5442 .unwrap()
5443 .iter()
5444 .find(|r| r["route"] == route)
5445 .cloned()
5446 .unwrap()
5447 };
5448 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5452 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5453 assert!(embed["decode_ms"].is_null());
5454 assert_eq!(embed["completion_tokens"], 0);
5455 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5459 assert_eq!(
5460 stats["tokens_prompt_total"].as_u64().unwrap(),
5461 embed["prompt_tokens"].as_u64().unwrap(),
5462 "only the forward pass counted"
5463 );
5464 }
5465
5466 fn streaming_test_app() -> Router {
5471 let mut cfg = test_dense_fixture();
5472 cfg.vocab_size = 256;
5473 let model = Model::Gguf(GgufModel {
5474 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5475 tokenizer: Arc::new(ServerTokenizer::Byte),
5476 stop_tokens: StopTokens::default(),
5477 bos_id: None,
5478 is_synthetic: false,
5479 chat_template: chat_template::PromptTemplate::plain(),
5480 });
5481 test_app_with_state(Arc::new(test_state(
5482 model,
5483 ResponseCache::new(1000, Duration::from_secs(3600)),
5484 )))
5485 }
5486
5487 #[tokio::test]
5496 async fn chat_logprobs_are_rendered_and_are_never_served_from_cache() {
5497 let app = test_app();
5498 let body = |logprobs: Option<(bool, Option<u32>)>| {
5499 let mut b = serde_json::json!({
5500 "model": "x",
5501 "messages": [{"role": "user", "content": "hi"}],
5502 "max_tokens": 4
5503 });
5504 if let Some((on, top)) = logprobs {
5505 b["logprobs"] = serde_json::json!(on);
5506 if let Some(n) = top {
5507 b["top_logprobs"] = serde_json::json!(n);
5508 }
5509 }
5510 b
5511 };
5512
5513 let (status, plain) =
5515 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(None)).await;
5516 assert_eq!(status, StatusCode::OK, "{plain}");
5517 assert!(plain["choices"][0]["logprobs"].is_null(), "{plain}");
5518
5519 for attempt in 0..2 {
5523 let (status, with) = post_json_uri(
5524 &app,
5525 frink_api::routes::V1_CHAT_COMPLETIONS,
5526 body(Some((true, Some(2)))),
5527 )
5528 .await;
5529 assert_eq!(status, StatusCode::OK, "{with}");
5530 assert_ne!(
5531 with["frink_cache"], "hit",
5532 "attempt {attempt} replayed a cached answer for a logprobs request: {with}"
5533 );
5534 let lp = &with["choices"][0]["logprobs"];
5535 assert!(lp.is_object(), "attempt {attempt}: {with}");
5536 let content = lp["content"].as_array().expect("content");
5537 assert!(lp["tokens"].is_null(), "completions shape leaked: {lp}");
5539 for entry in content {
5540 assert!(entry["token"].is_string(), "{entry}");
5541 assert!(entry["bytes"].is_array(), "{entry}");
5542 let v = entry["logprob"].as_f64().expect("a real number");
5543 assert!(v <= 0.0 && v.is_finite(), "{entry}");
5544 let top = entry["top_logprobs"].as_array().expect("top_logprobs");
5545 assert!(top.len() <= 2, "asked for 2, got {}", top.len());
5546 }
5547 }
5548 }
5549
5550 #[tokio::test]
5556 async fn the_chat_logprobs_pair_is_validated() {
5557 let app = test_app();
5558 for (extra, why) in [
5559 (serde_json::json!({"top_logprobs": 3}), "without logprobs"),
5560 (
5561 serde_json::json!({"logprobs": true, "top_logprobs": 21}),
5562 "above the cap",
5563 ),
5564 ] {
5565 let mut body = serde_json::json!({
5566 "model": "x",
5567 "messages": [{"role": "user", "content": "hi"}],
5568 "max_tokens": 2
5569 });
5570 for (k, v) in extra.as_object().unwrap() {
5571 body[k] = v.clone();
5572 }
5573 let (status, answer) =
5574 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await;
5575 assert_eq!(status, StatusCode::BAD_REQUEST, "{why}: {answer}");
5576 assert!(
5577 answer["error"]["message"]
5578 .as_str()
5579 .is_some_and(|m| m.contains("top_logprobs")),
5580 "{why}: {answer}"
5581 );
5582 }
5583 }
5584
5585 #[tokio::test]
5592 async fn a_salt_keeps_one_callers_cached_answer_from_another() {
5593 let app = test_app();
5594 let body = |salt: Option<&str>| {
5595 let mut b = serde_json::json!({
5596 "model": "x",
5597 "messages": [{"role": "user", "content": "the same prompt"}],
5598 "max_tokens": 4,
5599 "seed": 1
5600 });
5601 if let Some(s) = salt {
5602 b["cache_salt"] = serde_json::json!(s);
5603 }
5604 b
5605 };
5606 let post = |b: serde_json::Value| {
5607 let app = app.clone();
5608 async move { post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, b).await }
5609 };
5610
5611 let (status, _) = post(body(Some("tenant-a"))).await;
5613 assert_eq!(status, StatusCode::OK);
5614 let (_, again) = post(body(Some("tenant-a"))).await;
5615 assert_eq!(
5616 again["frink_cache"], "hit",
5617 "the owner did not get its own entry back: {again}"
5618 );
5619
5620 let (_, other) = post(body(Some("tenant-b"))).await;
5622 assert_ne!(
5623 other["frink_cache"], "hit",
5624 "a different caller was served tenant-a's answer: {other}"
5625 );
5626
5627 let (_, shared) = post(body(None)).await;
5629 assert_ne!(
5630 shared["frink_cache"], "hit",
5631 "an unsalted request was served a salted answer: {shared}"
5632 );
5633 }
5634
5635 #[tokio::test]
5640 async fn chat_serves_several_choices_and_refuses_the_streaming_pair() {
5641 let app = test_app();
5642 let body = |n: u32, stream: bool| {
5643 serde_json::json!({
5644 "model": "x",
5645 "messages": [{"role": "user", "content": "hi"}],
5646 "max_tokens": 4,
5647 "temperature": 1.0,
5648 "n": n,
5649 "stream": stream
5650 })
5651 };
5652
5653 let (status, one) =
5654 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5655 assert_eq!(status, StatusCode::OK, "{one}");
5656
5657 let (status, three) =
5658 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5659 assert_eq!(status, StatusCode::OK, "{three}");
5660 let choices = three["choices"].as_array().expect("an array");
5661 assert_eq!(choices.len(), 3, "{three}");
5662 for (i, c) in choices.iter().enumerate() {
5663 assert_eq!(c["index"], i);
5664 assert!(c["message"]["role"].is_string(), "{c}");
5665 assert!(c["finish_reason"].is_string(), "{c}");
5666 }
5667 assert_eq!(
5669 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5670 "n = 3 billed the prompt more than once"
5671 );
5672
5673 let (status, refused) =
5676 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, true)).await;
5677 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{refused}");
5678 let message = refused["error"]["message"].as_str().unwrap_or_default();
5679 assert!(
5680 message.contains('n') && message.contains("stream"),
5681 "{refused}"
5682 );
5683 }
5684
5685 #[tokio::test]
5695 async fn every_route_refuses_the_same_unimplemented_fields() {
5696 let app = test_app();
5697 let fields = [
5698 ("n", serde_json::json!(3)),
5699 ("best_of", serde_json::json!(2)),
5700 ("prompt_logprobs", serde_json::json!(1)),
5701 ("echo", serde_json::json!(true)),
5702 ("use_beam_search", serde_json::json!(true)),
5703 ("truncate_prompt_tokens", serde_json::json!(8)),
5704 ("prompt_embeds", serde_json::json!("AA==")),
5705 ("allowed_token_ids", serde_json::json!([1, 2])),
5706 ("bad_words", serde_json::json!(["x"])),
5707 ("skip_special_tokens", serde_json::json!(false)),
5708 ("return_tokens_as_token_ids", serde_json::json!(true)),
5709 ];
5710 for (field, value) in fields {
5711 for (uri, base) in [
5712 (
5713 frink_api::routes::V1_CHAT_COMPLETIONS,
5714 serde_json::json!({
5715 "model": "x",
5716 "messages": [{"role": "user", "content": "hi"}],
5717 "max_tokens": 2
5718 }),
5719 ),
5720 (
5721 frink_api::routes::V1_COMPLETIONS,
5722 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
5723 ),
5724 (
5725 frink_api::routes::COMPLETION,
5726 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5727 ),
5728 ] {
5729 let mut body = base;
5730 body[field] = value.clone();
5731 if field == "prompt_logprobs" && uri == frink_api::routes::V1_COMPLETIONS {
5738 let (status, answer) = post_json_uri(&app, uri, body).await;
5739 assert_eq!(status, StatusCode::OK, "{uri} refused it: {answer}");
5740 assert!(
5741 answer["prompt_logprobs"].is_array(),
5742 "served without the field: {answer}"
5743 );
5744 continue;
5745 }
5746 if (field == "n" || field == "best_of")
5747 && (uri == frink_api::routes::V1_COMPLETIONS
5748 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
5749 {
5750 let (status, answer) = post_json_uri(&app, uri, body).await;
5751 assert_eq!(
5752 status,
5753 StatusCode::OK,
5754 "{uri} refused a served `{field}`: {answer}"
5755 );
5756 let want = if field == "n" { 3 } else { 1 };
5760 assert_eq!(
5761 answer["choices"].as_array().map(Vec::len),
5762 Some(want),
5763 "{field}: {answer}"
5764 );
5765 continue;
5766 }
5767 let (status, answer) = post_json_uri(&app, uri, body).await;
5768 assert_eq!(
5769 status,
5770 StatusCode::NOT_IMPLEMENTED,
5771 "{uri} served `{field}` instead of refusing it: {answer}"
5772 );
5773 assert!(
5774 answer["error"]["message"]
5775 .as_str()
5776 .is_some_and(|m| m.contains(field)),
5777 "{uri} refused `{field}` without naming it: {answer}"
5778 );
5779 }
5780 }
5781 }
5782
5783 #[tokio::test]
5784 async fn the_native_completion_wire_is_not_the_openai_one() {
5785 let app = test_app();
5786
5787 let (status, native) = post_json_uri(
5788 &app,
5789 frink_api::routes::COMPLETION,
5790 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5791 )
5792 .await;
5793 assert_eq!(status, StatusCode::OK, "{native}");
5794 assert!(native["content"].is_string(), "{native}");
5795 assert_eq!(native["stop"], true);
5796 assert_eq!(native["stop_type"], "limit");
5797 assert_eq!(native["stopping_word"], "");
5798 assert_eq!(native["truncated"], false);
5799 assert_eq!(native["id_slot"], -1);
5800 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5801 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5802 assert!(
5803 native.get("choices").is_none(),
5804 "the native shape has no `choices`: {native}"
5805 );
5806
5807 let (status, openai) = post_json_uri(
5808 &app,
5809 frink_api::routes::V1_COMPLETIONS,
5810 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5811 )
5812 .await;
5813 assert_eq!(status, StatusCode::OK);
5814 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5815 assert!(
5816 openai.get("content").is_none(),
5817 "the OpenAI shape has no top-level `content`: {openai}"
5818 );
5819 }
5820
5821 #[tokio::test]
5825 async fn both_native_spellings_reach_the_same_handler() {
5826 let app = test_app();
5827 for route in [
5828 frink_api::routes::COMPLETION,
5829 frink_api::routes::COMPLETIONS,
5830 ] {
5831 let (status, body) = post_json_uri(
5832 &app,
5833 route,
5834 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5835 )
5836 .await;
5837 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5838 assert_eq!(body["stop"], true, "{route}");
5839 assert!(body["content"].is_string(), "{route}");
5840 }
5841
5842 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5845 let routes: Vec<&str> = stats["recent"]
5846 .as_array()
5847 .unwrap()
5848 .iter()
5849 .map(|row| row["route"].as_str().unwrap())
5850 .collect();
5851 assert!(
5852 routes.contains(&frink_api::routes::COMPLETION),
5853 "{routes:?}"
5854 );
5855 assert!(
5856 routes.contains(&frink_api::routes::COMPLETIONS),
5857 "{routes:?}"
5858 );
5859 }
5860
5861 #[tokio::test]
5867 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5868 let app = streaming_test_app();
5869 let raw = post_sse_raw_uri(
5870 &app,
5871 frink_api::routes::COMPLETION,
5872 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5873 )
5874 .await;
5875
5876 assert!(
5877 !raw.contains("[DONE]"),
5878 "llama.cpp's native stream has no sentinel: {raw}"
5879 );
5880 let frames: Vec<serde_json::Value> = raw
5881 .lines()
5882 .filter_map(|line| line.strip_prefix("data: "))
5883 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5884 .collect();
5885 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5886
5887 let (last, partials) = frames.split_last().unwrap();
5888 assert_eq!(last["stop"], true, "the last frame closes the stream");
5889 assert!(last["timings"].is_object(), "{last}");
5890 assert!(last["stop_type"].is_string(), "{last}");
5891 for partial in partials {
5892 assert_eq!(partial["stop"], false, "{partial}");
5893 assert!(partial["content"].is_string(), "{partial}");
5894 assert!(partial.get("timings").is_none(), "{partial}");
5898 assert!(partial.get("generation_settings").is_none(), "{partial}");
5899 }
5900 let streamed: String = partials
5903 .iter()
5904 .filter_map(|p| p["content"].as_str())
5905 .collect();
5906 assert_eq!(last["content"].as_str().unwrap(), streamed);
5907 }
5908
5909 #[tokio::test]
5914 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5915 let app = test_app();
5916 for body in [
5917 serde_json::json!({"prompt": "hi"}),
5918 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5919 ] {
5920 let (status, refusal) =
5921 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5922 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5923 assert!(
5924 refusal["error"]["message"]
5925 .as_str()
5926 .unwrap()
5927 .contains("n_predict"),
5928 "{refusal}"
5929 );
5930 }
5931 let (status, _) = post_json_uri(
5934 &app,
5935 frink_api::routes::COMPLETION,
5936 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5937 )
5938 .await;
5939 assert_eq!(status, StatusCode::OK);
5940 }
5941
5942 #[tokio::test]
5951 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5952 let app = streaming_test_app();
5953 let ask = |stop: serde_json::Value| {
5954 let app = app.clone();
5955 async move {
5956 post_json_uri(
5957 &app,
5958 frink_api::routes::COMPLETION,
5959 serde_json::json!({
5960 "prompt": "hi",
5961 "n_predict": 64,
5962 "ignore_eos": true,
5963 "stop": stop,
5964 }),
5965 )
5966 .await
5967 .1
5968 }
5969 };
5970
5971 let baseline = ask(serde_json::json!([])).await;
5972 assert_eq!(baseline["stop_type"], "limit");
5973 assert_eq!(baseline["stopping_word"], "");
5974 let text = baseline["content"].as_str().unwrap().to_string();
5975 let sentinel: String = text.chars().skip(1).take(2).collect();
5981 assert_eq!(
5982 sentinel.chars().count(),
5983 2,
5984 "the fixture must produce enough output to cut: {text:?}"
5985 );
5986 let cut = text.find(&sentinel).expect("it came out of this text");
5987
5988 let stopped = ask(serde_json::json!([sentinel])).await;
5989 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5990 assert_eq!(stopped["stopping_word"], sentinel);
5991 assert_eq!(
5992 stopped["content"].as_str().unwrap(),
5993 &text[..cut],
5994 "the answer must be cut at the sentinel, not run past it"
5995 );
5996 }
5997
5998 #[tokio::test]
6004 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
6005 let app = test_app();
6006
6007 let (v1_status, v1) = post_json_uri(
6008 &app,
6009 frink_api::routes::V1_TOKENIZE,
6010 serde_json::json!({"prompt": "hello"}),
6011 )
6012 .await;
6013 let (alias_status, alias) = post_json_uri(
6014 &app,
6015 frink_api::routes::TOKENIZE,
6016 serde_json::json!({"content": "hello"}),
6017 )
6018 .await;
6019 assert_eq!(v1_status, StatusCode::OK);
6020 assert_eq!(alias_status, StatusCode::OK, "{alias}");
6021 assert_eq!(v1["tokens"], alias["tokens"]);
6022 assert!(!alias["tokens"].as_array().unwrap().is_empty());
6023
6024 let (status, both_ways) = post_json_uri(
6027 &app,
6028 frink_api::routes::TOKENIZE,
6029 serde_json::json!({"prompt": "hello"}),
6030 )
6031 .await;
6032 assert_eq!(status, StatusCode::OK);
6033 assert_eq!(both_ways["tokens"], v1["tokens"]);
6034 }
6035
6036 #[tokio::test]
6041 async fn detokenize_answers_under_both_dialects_keys() {
6042 let app = test_app();
6043 for route in [
6044 frink_api::routes::DETOKENIZE,
6045 frink_api::routes::V1_DETOKENIZE,
6046 ] {
6047 let (status, body) =
6048 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
6049 assert_eq!(status, StatusCode::OK, "{route}");
6050 assert_eq!(body["text"], "hi", "{route}");
6051 assert_eq!(body["content"], body["text"], "{route}");
6052 }
6053 }
6054
6055 #[tokio::test]
6059 async fn the_alias_is_recorded_under_the_path_the_client_called() {
6060 let app = test_app();
6061 let (status, _) = post_json_uri(
6062 &app,
6063 frink_api::routes::TOKENIZE,
6064 serde_json::json!({"content": "hello"}),
6065 )
6066 .await;
6067 assert_eq!(status, StatusCode::OK);
6068
6069 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6070 let routes: Vec<&str> = stats["recent"]
6071 .as_array()
6072 .unwrap()
6073 .iter()
6074 .map(|row| row["route"].as_str().unwrap())
6075 .collect();
6076 assert!(
6077 routes.contains(&frink_api::routes::TOKENIZE),
6078 "the alias must be its own row: {routes:?}"
6079 );
6080 assert!(
6081 !routes.contains(&frink_api::routes::V1_TOKENIZE),
6082 "nothing called /v1/tokenize: {routes:?}"
6083 );
6084 }
6085
6086 #[tokio::test]
6091 async fn add_special_prepends_the_same_bos_the_decoder_would() {
6092 let mut cfg = test_dense_fixture();
6093 cfg.vocab_size = 256;
6094 let model = Model::Gguf(GgufModel {
6095 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
6096 tokenizer: Arc::new(ServerTokenizer::Byte),
6097 stop_tokens: StopTokens::default(),
6098 bos_id: Some(7),
6099 is_synthetic: true,
6100 chat_template: chat_template::PromptTemplate::plain(),
6101 });
6102 let app = test_app_with_state(Arc::new(test_state(
6103 model,
6104 ResponseCache::new(1000, Duration::from_secs(3600)),
6105 )));
6106
6107 let (_, plain) = post_json_uri(
6108 &app,
6109 frink_api::routes::TOKENIZE,
6110 serde_json::json!({"content": "hi"}),
6111 )
6112 .await;
6113 let (_, special) = post_json_uri(
6114 &app,
6115 frink_api::routes::TOKENIZE,
6116 serde_json::json!({"content": "hi", "add_special": true}),
6117 )
6118 .await;
6119
6120 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6121 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6122 assert_eq!(special["count"], 3);
6123 }
6124
6125 #[tokio::test]
6128 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6129 let app = test_app();
6130 let (status, _) = post_json_uri(
6131 &app,
6132 frink_api::routes::V1_EMBEDDINGS,
6133 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6134 )
6135 .await;
6136 assert_eq!(status, StatusCode::BAD_REQUEST);
6137
6138 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6139 let recent = stats["recent"].as_array().unwrap();
6140 assert_eq!(recent.len(), 1);
6141 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
6142 assert_eq!(recent[0]["status"], 400);
6143 assert_eq!(
6144 recent[0]["prompt_tokens"], 0,
6145 "a rejected call embedded nothing"
6146 );
6147 }
6148
6149 #[tokio::test]
6152 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6153 let app = test_app();
6154 let key = "sk-monitor-secret";
6155 let (status, _) = post_json_with_headers(
6156 &app,
6157 "/v1/chat/completions",
6158 serde_json::json!({
6159 "model": "x",
6160 "messages": [{"role": "user", "content": "hi"}],
6161 "max_tokens": 2
6162 }),
6163 &[
6164 ("authorization", &format!("Bearer {key}")),
6165 ("x-frink-client", "frink-studio"),
6166 ],
6167 )
6168 .await;
6169 assert_eq!(status, StatusCode::OK);
6170
6171 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6172 let row = stats["recent"].as_array().unwrap()[0].clone();
6173 let fingerprint = row["via_api_key"]
6174 .as_str()
6175 .expect("the row names the key that served it")
6176 .to_string();
6177 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6178 assert!(!fingerprint.contains(key));
6179 assert!(
6180 !serde_json::to_string(&stats).unwrap().contains(key),
6181 "the stats payload must not carry the key in any form"
6182 );
6183 assert_eq!(row["client"], "frink-studio");
6184 }
6185
6186 #[tokio::test]
6189 async fn different_keys_are_different_callers_and_no_key_is_null() {
6190 let app = test_app();
6191 let body = serde_json::json!({
6192 "model": "x",
6193 "messages": [{"role": "user", "content": "hi"}],
6194 "max_tokens": 1
6195 });
6196 for headers in [
6197 vec![("authorization", "Bearer key-one")],
6198 vec![("authorization", "Bearer key-two")],
6199 vec![],
6200 ] {
6201 let (status, _) =
6202 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6203 assert_eq!(status, StatusCode::OK);
6204 }
6205
6206 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6207 let recent = stats["recent"].as_array().unwrap();
6208 assert_eq!(recent.len(), 3);
6209 let one = recent[0]["via_api_key"].as_str().unwrap();
6210 let two = recent[1]["via_api_key"].as_str().unwrap();
6211 assert_ne!(one, two, "two keys must not collapse into one caller");
6212 assert!(
6213 recent[2]["via_api_key"].is_null(),
6214 "an unauthenticated call is null, not a fingerprint of nothing"
6215 );
6216 assert!(recent[2]["client"].is_null());
6217 }
6218
6219 #[tokio::test]
6224 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6225 let state = Arc::new(test_state(
6226 named_test_model("really-loaded", 256),
6227 ResponseCache::new(4, Duration::from_secs(60)),
6228 ));
6229 let app = test_app_with_state(Arc::clone(&state));
6230
6231 let (status, _) = post_json_uri(
6232 &app,
6233 "/v1/chat/completions",
6234 serde_json::json!({
6235 "model": "gpt-4-turbo-that-is-not-here",
6236 "messages": [{"role": "user", "content": "hi"}],
6237 "max_tokens": 2
6238 }),
6239 )
6240 .await;
6241 assert_eq!(status, StatusCode::OK);
6242
6243 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6244 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6245
6246 state.swap_active(None);
6249 let (status, _) = post_json_uri(
6250 &app,
6251 "/v1/chat/completions",
6252 serde_json::json!({
6253 "model": "gpt-4-turbo-that-is-not-here",
6254 "messages": [{"role": "user", "content": "hi"}]
6255 }),
6256 )
6257 .await;
6258 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6259 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6260 let recent = stats["recent"].as_array().unwrap();
6261 assert!(recent[recent.len() - 1]["model"].is_null());
6262 }
6263
6264 #[tokio::test]
6268 async fn a_streamed_row_names_the_model_it_decoded_against() {
6269 let state = Arc::new(test_state(
6270 named_test_model("model-before", 256),
6271 ResponseCache::new(4, Duration::from_secs(60)),
6272 ));
6273 let app = test_app_with_state(Arc::clone(&state));
6274 let _ = post_sse_raw(&app, resumable_request()).await;
6275 active_model(&state, "model-after");
6277
6278 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6279 assert_eq!(stats["recent"][0]["model"], "model-before");
6280 }
6281
6282 #[tokio::test]
6285 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6286 let app = test_app();
6287 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6288 assert_eq!(status, StatusCode::OK);
6289 assert!(
6290 stats["queue_depth"].is_null(),
6291 "without continuous batching nothing queues, so there is nothing to measure"
6292 );
6293 assert!(stats["queue_rejected_total"].is_null());
6294 assert_eq!(
6295 stats["generating_now"], 0,
6296 "work in progress is measured and really is zero here"
6297 );
6298 }
6299
6300 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6305 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
6306 }
6307
6308 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6312 use http_body_util::BodyExt;
6313 use tower::ServiceExt;
6314
6315 let response = app
6316 .clone()
6317 .oneshot(
6318 axum::http::Request::builder()
6319 .method("POST")
6320 .uri(uri)
6321 .header("content-type", "application/json")
6322 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6323 .unwrap(),
6324 )
6325 .await
6326 .unwrap();
6327 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6328 String::from_utf8(bytes.to_vec()).unwrap()
6329 }
6330
6331 async fn get_json_with_headers(
6332 app: &Router,
6333 uri: &str,
6334 headers: &[(&str, &str)],
6335 ) -> (StatusCode, serde_json::Value) {
6336 use http_body_util::BodyExt;
6337 use tower::ServiceExt;
6338
6339 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6340 for (name, value) in headers {
6341 builder = builder.header(*name, *value);
6342 }
6343 let response = app
6344 .clone()
6345 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6346 .await
6347 .unwrap();
6348 let status = response.status();
6349 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6350 (
6351 status,
6352 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6353 )
6354 }
6355
6356 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6357 body.lines()
6358 .filter_map(|line| line.strip_prefix(field))
6359 .map(str::trim)
6360 .collect()
6361 }
6362
6363 fn resumable_request() -> serde_json::Value {
6364 serde_json::json!({
6365 "model": "m",
6366 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6367 "max_tokens": 4,
6368 "temperature": 0,
6369 "stream": true,
6370 "stream_resumable": true,
6371 })
6372 }
6373
6374 #[tokio::test]
6379 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6380 let app = test_app();
6381 let body = post_sse_raw(&app, resumable_request()).await;
6382
6383 let request_id = body
6384 .lines()
6385 .find_map(|l| l.strip_prefix("data: "))
6386 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6387 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6388 .expect("the first chunk names the request");
6389
6390 let ids = sse_field(&body, "id:");
6391 let datas = sse_field(&body, "data:");
6392 assert_eq!(
6393 ids.len(),
6394 datas.len(),
6395 "every event carries an id, or a reconnect cannot name where it stopped"
6396 );
6397 for (i, id) in ids.iter().enumerate() {
6398 assert_eq!(*id, format!("{request_id}:{i}"));
6399 }
6400 let retries = sse_field(&body, "retry:");
6401 assert_eq!(
6402 retries.len(),
6403 1,
6404 "the reconnect delay is stated once, not on every event"
6405 );
6406 assert_eq!(retries[0], "1500");
6407 assert!(
6408 body.contains("data: [DONE]"),
6409 "the end of stream is still stated"
6410 );
6411 }
6412
6413 #[tokio::test]
6417 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6418 let app = test_app();
6419 let mut request = resumable_request();
6420 request["stream_resumable"] = serde_json::json!(false);
6421 let body = post_sse_raw(&app, request).await;
6422 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6423 assert!(
6424 sse_field(&body, "id:").is_empty(),
6425 "an id promises a replay this stream cannot serve"
6426 );
6427 assert!(sse_field(&body, "retry:").is_empty());
6428 }
6429
6430 #[tokio::test]
6434 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6435 let app = test_app();
6436 let body = post_sse_raw(&app, resumable_request()).await;
6437 let request_id = sse_field(&body, "id:")[0]
6438 .rsplit_once(':')
6439 .unwrap()
6440 .0
6441 .to_string();
6442 let streamed: Vec<String> = sse_field(&body, "data:")
6443 .iter()
6444 .map(|d| d.to_string())
6445 .collect();
6446
6447 let (status, polled) = get_json(
6448 &app,
6449 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6450 )
6451 .await;
6452 assert_eq!(status, StatusCode::OK);
6453 let events: Vec<String> = polled["events"]
6454 .as_array()
6455 .unwrap()
6456 .iter()
6457 .map(|e| e["data"].as_str().unwrap().to_string())
6458 .collect();
6459 assert_eq!(
6460 events, streamed,
6461 "the fallback must deliver the same answer, not a re-run of it"
6462 );
6463 assert_eq!(polled["request_id"], request_id);
6464 assert_eq!(
6465 polled["done"], false,
6466 "events were still being handed out, so the client must ask again"
6467 );
6468
6469 let next = polled["next_index"].as_u64().unwrap();
6472 let (_, drained) = get_json(
6473 &app,
6474 &format!(
6475 "{}?from={next}",
6476 frink_api::routes::v1_stream_poll(&request_id)
6477 ),
6478 )
6479 .await;
6480 assert_eq!(drained["done"], true);
6481 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6482 }
6483
6484 #[tokio::test]
6488 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6489 let app = test_app();
6490 let body = post_sse_raw(&app, resumable_request()).await;
6491 let ids = sse_field(&body, "id:");
6492 let datas: Vec<String> = sse_field(&body, "data:")
6493 .iter()
6494 .map(|d| d.to_string())
6495 .collect();
6496 assert!(
6497 ids.len() >= 3,
6498 "need a few events to resume into the middle"
6499 );
6500 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6501
6502 let (status, resumed) = get_json_with_headers(
6503 &app,
6504 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6505 &[],
6506 )
6507 .await;
6508 assert_eq!(status, StatusCode::OK);
6509 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6510
6511 let (_, tail) = get_json(
6513 &app,
6514 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6515 )
6516 .await;
6517 let tail_events: Vec<String> = tail["events"]
6518 .as_array()
6519 .unwrap()
6520 .iter()
6521 .map(|e| e["data"].as_str().unwrap().to_string())
6522 .collect();
6523 assert_eq!(tail_events, datas[2..].to_vec());
6524 }
6525
6526 #[tokio::test]
6529 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6530 use http_body_util::BodyExt;
6531 use tower::ServiceExt;
6532
6533 let app = test_app();
6534 let body = post_sse_raw(&app, resumable_request()).await;
6535 let ids = sse_field(&body, "id:");
6536 let datas: Vec<String> = sse_field(&body, "data:")
6537 .iter()
6538 .map(|d| d.to_string())
6539 .collect();
6540 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6541
6542 let response = app
6543 .clone()
6544 .oneshot(
6545 axum::http::Request::builder()
6546 .method("GET")
6547 .uri(frink_api::routes::v1_stream(&request_id))
6548 .header("last-event-id", format!("{request_id}:0"))
6549 .body(axum::body::Body::empty())
6550 .unwrap(),
6551 )
6552 .await
6553 .unwrap();
6554 assert_eq!(response.status(), StatusCode::OK);
6555 assert_eq!(
6556 response
6557 .headers()
6558 .get("x-accel-buffering")
6559 .and_then(|v| v.to_str().ok()),
6560 Some("no"),
6561 "the reconnect needs the same anti-buffering header as the stream"
6562 );
6563 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6564 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6565 assert_eq!(
6566 sse_field(&resumed, "data:")
6567 .iter()
6568 .map(|d| d.to_string())
6569 .collect::<Vec<_>>(),
6570 datas[1..].to_vec()
6571 );
6572 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6573 }
6574
6575 #[tokio::test]
6579 async fn a_last_event_id_from_another_stream_is_refused() {
6580 let app = test_app();
6581 let body = post_sse_raw(&app, resumable_request()).await;
6582 let request_id = sse_field(&body, "id:")[0]
6583 .rsplit_once(':')
6584 .unwrap()
6585 .0
6586 .to_string();
6587
6588 let (status, err) = get_json_with_headers(
6589 &app,
6590 &frink_api::routes::v1_stream(&request_id),
6591 &[("last-event-id", "chatcmpl-someone-else:3")],
6592 )
6593 .await;
6594 assert_eq!(status, StatusCode::BAD_REQUEST);
6595 assert_eq!(err["error"]["code"], "bad_last_event_id");
6596 }
6597
6598 #[tokio::test]
6602 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6603 let app = test_app();
6604 let mut request = resumable_request();
6605 request["stream_resumable"] = serde_json::json!(false);
6606 let body = post_sse_raw(&app, request).await;
6607 let request_id = body
6608 .lines()
6609 .find_map(|l| l.strip_prefix("data: "))
6610 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6611 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6612 .unwrap();
6613
6614 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6615 assert_eq!(status, StatusCode::NOT_FOUND);
6616 assert_eq!(err["error"]["code"], "stream_not_found");
6617 assert!(err["error"]["message"]
6618 .as_str()
6619 .unwrap()
6620 .contains("stream_resumable"));
6621 }
6622
6623 #[test]
6627 fn the_axum_stream_patterns_match_the_published_templates() {
6628 assert_eq!(
6629 axum_path(frink_api::routes::V1_STREAM),
6630 "/v1/stream/:request_id"
6631 );
6632 assert_eq!(
6633 axum_path(frink_api::routes::V1_STREAM_POLL),
6634 "/v1/stream/:request_id/poll"
6635 );
6636 assert_eq!(
6637 frink_api::routes::v1_stream("abc"),
6638 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6639 );
6640 }
6641
6642 #[test]
6652 fn no_published_template_reaches_the_router_with_its_braces() {
6653 for template in [
6654 frink_api::routes::V1_STREAM,
6655 frink_api::routes::V1_STREAM_POLL,
6656 frink_api::routes::V1_RESPONSE,
6657 frink_api::routes::V1_RESPONSE_CANCEL,
6658 frink_api::routes::ADMIN_TASK_CANCEL,
6659 ] {
6660 assert!(
6661 template.contains('{'),
6662 "{template} is in the template list but has no placeholder"
6663 );
6664 let mounted = axum_path(template);
6665 assert!(
6666 !mounted.contains('{') && !mounted.contains('}'),
6667 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6668 );
6669 assert!(
6670 mounted.contains(':'),
6671 "{template} lost its placeholder entirely and would match one path only"
6672 );
6673 }
6674 }
6675
6676 #[tokio::test]
6684 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6685 let app = test_app();
6686 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6687 assert_eq!(status, StatusCode::NOT_FOUND);
6688 assert!(
6689 !body.is_null(),
6690 "empty body means axum never matched the route, so the id was read as a literal segment"
6691 );
6692 }
6693
6694 #[tokio::test]
6697 async fn the_task_list_starts_empty_rather_than_absent() {
6698 let app = test_app();
6699 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6700 assert_eq!(status, StatusCode::OK);
6701 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6702 }
6703
6704 #[tokio::test]
6713 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6714 assert!(
6715 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6716 "this test asserts the unconfigured behaviour"
6717 );
6718 let app = test_app();
6719 let (status, body) = post_json_uri(
6720 &app,
6721 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6722 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6723 )
6724 .await;
6725 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6726 assert!(
6727 body["error"]["message"]
6728 .as_str()
6729 .unwrap()
6730 .contains("--slot-save-path"),
6731 "{body}"
6732 );
6733 }
6734
6735 pub(crate) async fn post_json_uri(
6736 app: &Router,
6737 uri: &str,
6738 body: serde_json::Value,
6739 ) -> (StatusCode, serde_json::Value) {
6740 use http_body_util::BodyExt;
6741 use tower::ServiceExt;
6742
6743 let response = app
6744 .clone()
6745 .oneshot(
6746 axum::http::Request::builder()
6747 .method("POST")
6748 .uri(uri)
6749 .header("content-type", "application/json")
6750 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6751 .unwrap(),
6752 )
6753 .await
6754 .unwrap();
6755 let status = response.status();
6756 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6757 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6758 (status, json)
6759 }
6760
6761 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6762 post_json_uri(app, "/v1/chat/completions", body).await.1
6763 }
6764
6765 #[tokio::test]
6774 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6775 let app = test_app();
6776 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6777 assert_eq!(status, StatusCode::OK);
6778
6779 let memory = &body["memory"];
6780 if memory.is_null() {
6781 return;
6784 }
6785 assert!(
6786 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6787 "a read that produced a zero is a broken read, not an idle \
6788 engine: {memory}"
6789 );
6790 assert!(
6791 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6792 "the quantity must travel with the number: {memory}"
6793 );
6794 }
6795
6796 #[tokio::test]
6802 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6803 let app = test_app();
6804 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6805 assert_eq!(status, StatusCode::OK);
6806 for pool in ["kv_pages", "window_slots", "state_slots"] {
6807 assert!(
6808 body["pools"][pool].is_null(),
6809 "{pool} must be null rather than a zero row: {}",
6810 body["pools"]
6811 );
6812 }
6813 }
6814
6815 #[tokio::test]
6823 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6824 use http_body_util::BodyExt;
6825 use tower::ServiceExt;
6826
6827 let app = test_app();
6828 let send = |body: serde_json::Value| {
6829 let app = app.clone();
6830 async move {
6831 app.oneshot(
6832 axum::http::Request::builder()
6833 .method("POST")
6834 .uri(frink_api::routes::V1_MESSAGES)
6835 .header("content-type", "application/json")
6836 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6837 .unwrap(),
6838 )
6839 .await
6840 .unwrap()
6841 }
6842 };
6843
6844 let ok = send(serde_json::json!({
6845 "model": "test",
6846 "max_tokens": 1,
6847 "messages": [{"role": "user", "content": "hi"}],
6848 }))
6849 .await;
6850 assert_eq!(ok.status(), StatusCode::OK);
6851 let id = ok
6852 .headers()
6853 .get("request-id")
6854 .expect("a served message names its id")
6855 .to_str()
6856 .unwrap()
6857 .to_string();
6858 assert!(!id.is_empty());
6859
6860 let bad = send(serde_json::json!({"model": "test"})).await;
6863 assert!(bad.status().is_client_error());
6864 let other = bad.headers().get("request-id").expect("errors too");
6865 assert_ne!(other.to_str().unwrap(), id);
6866 let _ = bad.into_body().collect().await.unwrap();
6867 }
6868
6869 #[tokio::test]
6877 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6878 let state = Arc::new(test_state(
6879 test_model_full_byte_vocab(),
6880 ResponseCache::new(1000, Duration::from_secs(3600)),
6881 ));
6882 let app = test_app_with_state(Arc::clone(&state));
6883 let body = serde_json::json!({
6884 "model": "test",
6885 "messages": [{"role": "user", "content": "hi"}],
6886 "max_tokens": 1,
6887 });
6888
6889 state
6890 .maintenance
6891 .lock()
6892 .unwrap()
6893 .begin_rebuild()
6894 .expect("a fresh server is serving, so the rebuild starts");
6895 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6896 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6897 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6898
6899 state.maintenance.lock().unwrap().finish_rebuild(true);
6900 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6901 assert_eq!(
6902 status,
6903 StatusCode::OK,
6904 "the gate reopens; a rebuild is not a latch"
6905 );
6906 }
6907
6908 #[tokio::test]
6913 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6914 let app = test_app();
6915 let (status, body) = post_json_uri(
6916 &app,
6917 frink_api::routes::V1_CANCEL,
6918 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6919 )
6920 .await;
6921 assert_eq!(status, StatusCode::NOT_FOUND);
6922 assert_eq!(body["cancelled"], serde_json::json!(false));
6923 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6924 assert!(
6925 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6926 "the verdict must carry a human reason: {body}"
6927 );
6928 }
6929
6930 #[tokio::test]
6934 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6935 let state = Arc::new(test_state(
6936 test_model_full_byte_vocab(),
6937 ResponseCache::new(1000, Duration::from_secs(3600)),
6938 ));
6939 let app = test_app_with_state(Arc::clone(&state));
6940 let (token, _guard) = state.cancels.register("chatcmpl-live");
6941
6942 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6943 assert_eq!(status, StatusCode::OK);
6944 assert_eq!(before["generating_now"], serde_json::json!(1));
6945
6946 let (status, body) = post_json_uri(
6947 &app,
6948 frink_api::routes::V1_CANCEL,
6949 serde_json::json!({ "request_id": "chatcmpl-live" }),
6950 )
6951 .await;
6952 assert_eq!(status, StatusCode::OK);
6953 assert_eq!(body["cancelled"], serde_json::json!(true));
6954 assert!(
6955 token.is_cancelled(),
6956 "the endpoint answered ok without setting the flag the decode loop reads"
6957 );
6958 }
6959
6960 #[tokio::test]
6961 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6962 let app = test_app();
6963 let (status, tok) =
6964 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6965 assert_eq!(status, StatusCode::OK);
6966 let tokens = tok["tokens"].as_array().unwrap();
6967 assert_eq!(tok["count"], tokens.len());
6968 assert!(!tokens.is_empty());
6969
6970 let (status, detok) = post_json_uri(
6971 &app,
6972 "/v1/detokenize",
6973 serde_json::json!({ "tokens": tokens }),
6974 )
6975 .await;
6976 assert_eq!(status, StatusCode::OK);
6977 assert_eq!(detok["text"], "Hi");
6978
6979 let (status, emb) = post_json_uri(
6980 &app,
6981 "/v1/embeddings",
6982 serde_json::json!({
6983 "input": "Hi",
6984 "embedding_type": "mean"
6985 }),
6986 )
6987 .await;
6988 assert_eq!(status, StatusCode::OK);
6989 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6990 assert!(!vec.is_empty());
6991 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6992 }
6993
6994 #[tokio::test]
7000 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
7001 let app = test_app();
7002 let (status, body) = post_json_uri(
7003 &app,
7004 "/v1/embeddings",
7005 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
7006 )
7007 .await;
7008 assert_eq!(status, StatusCode::BAD_REQUEST);
7009 let msg = body["error"]["message"].as_str().unwrap();
7010 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
7011 }
7012
7013 #[tokio::test]
7018 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7019 async fn a_real_embedding_model_serves_v1_embeddings() {
7020 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7021 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7022 if !path.exists() {
7023 eprintln!("SKIP: {} not present", path.display());
7024 return;
7025 }
7026 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
7027 let mut state = test_state(
7028 test_model_full_byte_vocab(),
7029 ResponseCache::new(1000, Duration::from_secs(3600)),
7030 );
7031 state.embedding = Some(Arc::new(encoder));
7032 let app = test_app_with_state(Arc::new(state));
7033
7034 let (status, body) = post_json_uri(
7035 &app,
7036 "/v1/embeddings",
7037 serde_json::json!({ "input": ["Hello world", "a second input"] }),
7038 )
7039 .await;
7040 assert_eq!(status, StatusCode::OK, "{body}");
7041 assert_eq!(body["model"], "bge-small-en-v1.5");
7042 let data = body["data"].as_array().unwrap();
7043 assert_eq!(data.len(), 2);
7044 for (i, row) in data.iter().enumerate() {
7045 assert_eq!(row["index"], i);
7046 let v: Vec<f64> = row["embedding"]
7047 .as_array()
7048 .unwrap()
7049 .iter()
7050 .map(|x| x.as_f64().unwrap())
7051 .collect();
7052 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
7053 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
7054 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
7055 }
7056 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
7059
7060 let (status, mean) = post_json_uri(
7064 &app,
7065 "/v1/embeddings",
7066 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
7067 )
7068 .await;
7069 assert_eq!(status, StatusCode::OK);
7070 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
7071 }
7072
7073 #[tokio::test]
7085 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
7086 async fn an_encoder_can_be_the_loaded_model() {
7087 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
7088 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
7089 if !path.exists() {
7090 eprintln!("SKIP: {} not present", path.display());
7091 return;
7092 }
7093
7094 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
7098 assert!(
7099 matches!(loaded, model::LoadedModel::Encoder(_)),
7100 "an encoder-only GGUF reached a decoder loader"
7101 );
7102 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
7103 assert!(
7104 matches!(loaded, Loaded::Encoder(_)),
7105 "the encoder did not stay an encoder through activation"
7106 );
7107 assert!(
7108 batcher.is_none() && ceiling.is_none(),
7109 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7110 );
7111
7112 let state = test_state(
7113 test_model_full_byte_vocab(),
7114 ResponseCache::new(1000, Duration::from_secs(3600)),
7115 );
7116 state.swap_active(Some(Arc::new(ActiveModel {
7117 id: None,
7118 loaded,
7119 batcher,
7120 ceiling,
7121 checkpoint_path: None,
7122 })));
7123 let app = test_app_with_state(Arc::new(state));
7124
7125 let (status, body) = post_json_uri(
7127 &app,
7128 "/v1/embeddings",
7129 serde_json::json!({ "input": "Hello world" }),
7130 )
7131 .await;
7132 assert_eq!(status, StatusCode::OK, "{body}");
7133 assert_eq!(body["model"], "bge-small-en-v1.5");
7134 let v = body["data"][0]["embedding"].as_array().unwrap();
7135 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
7136
7137 let (status, body) = post_json_uri(
7139 &app,
7140 "/v1/chat/completions",
7141 serde_json::json!({
7142 "model": "bge-small-en-v1.5",
7143 "messages": [{"role": "user", "content": "hi"}],
7144 }),
7145 )
7146 .await;
7147 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
7148 let msg = body["error"]["message"].as_str().unwrap();
7149 for fact in [
7150 "bge-small-en-v1.5",
7151 "bert",
7152 "embedding model",
7153 "/v1/embeddings",
7154 ] {
7155 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
7156 }
7157
7158 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
7160 assert_eq!(status, StatusCode::OK);
7161 let entry = &models["data"][0];
7162 assert_eq!(entry["id"], "bge-small-en-v1.5");
7163 assert_eq!(entry["frink_model_kind"], "embedding");
7164 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
7165 assert_eq!(entry["frink_n_embd"], 384);
7166 assert_eq!(entry["frink_pooling"], "CLS");
7167 assert_eq!(
7168 entry["frink_endpoints"],
7169 serde_json::json!(["/v1/embeddings"])
7170 );
7171 assert!(entry.get("supported_reasoning_efforts").is_none());
7174
7175 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
7177 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
7178 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
7179 assert_eq!(health["model"]["synthetic_weights"], false);
7180 let weights = health["capabilities"]
7181 .as_array()
7182 .unwrap()
7183 .iter()
7184 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
7185 .expect("a real-weights capability row");
7186 let detail = weights["detail"].as_str().unwrap_or_default();
7187 assert!(detail.contains("ENCODER"), "{detail}");
7188 let (status, body) = post_json_uri(
7196 &app,
7197 frink_api::routes::V1_TOKENIZE,
7198 serde_json::json!({ "content": "hello world" }),
7199 )
7200 .await;
7201 assert_eq!(
7202 status,
7203 StatusCode::OK,
7204 "an encoder has a real tokenizer: {body}"
7205 );
7206 let tokens = body["tokens"].as_array().expect("tokens array").clone();
7207 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
7208
7209 let (status, body) = post_json_uri(
7210 &app,
7211 frink_api::routes::V1_DETOKENIZE,
7212 serde_json::json!({ "tokens": tokens }),
7213 )
7214 .await;
7215 assert_eq!(status, StatusCode::OK, "{body}");
7216 let round_tripped = body["content"].as_str().expect("content").to_string();
7217 assert!(
7218 round_tripped.contains("hello") && round_tripped.contains("world"),
7219 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
7220 );
7221
7222 let (status, _) = post_json_uri(
7225 &app,
7226 "/v1/completions",
7227 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
7228 )
7229 .await;
7230 assert_eq!(
7231 status,
7232 StatusCode::NOT_IMPLEMENTED,
7233 "tokenizing an encoder must not have opened a path to generating with one"
7234 );
7235 }
7236
7237 #[tokio::test]
7242 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7243 use http_body_util::BodyExt;
7244 use tower::ServiceExt;
7245
7246 let fixture = concat!(
7247 "../frink-models/tests/fixtures/",
7248 "frink_real_moe_test.gguf"
7249 );
7250 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7251 let decoder = Decoder::from_gguf_with_expert_cache(
7252 &fixture,
7253 frink_models::config::test_moe_fixture(),
7254 Some(1024 * 1024),
7255 )
7256 .expect("MoE fixture must load store-backed");
7257
7258 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
7262 decoder.forward_token(1, 0, &mut caches);
7263
7264 let model = Model::Gguf(GgufModel {
7265 decoder: Arc::new(decoder),
7266 tokenizer: Arc::new(ServerTokenizer::Byte),
7267 stop_tokens: StopTokens::default(),
7268 bos_id: None,
7269 is_synthetic: false,
7270 chat_template: chat_template::PromptTemplate::plain(),
7271 });
7272 let state = Arc::new(test_state(
7273 model,
7274 ResponseCache::new(16, Duration::from_secs(60)),
7275 ));
7276 let app = Router::new()
7277 .route("/metrics", axum::routing::get(metrics))
7278 .route("/v1/chat/completions", post(chat_completions))
7279 .with_state(state);
7280
7281 let fetch_metrics = |app: Router| async move {
7282 let resp = app
7283 .oneshot(
7284 axum::http::Request::builder()
7285 .method("GET")
7286 .uri("/metrics")
7287 .body(axum::body::Body::empty())
7288 .unwrap(),
7289 )
7290 .await
7291 .unwrap();
7292 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7293 String::from_utf8(bytes.to_vec()).unwrap()
7294 };
7295
7296 let after = fetch_metrics(app.clone()).await;
7297 assert!(
7298 after.contains("frink_expert_cache_misses_total"),
7299 "streaming model must expose expert-cache metrics: {after}"
7300 );
7301 let misses: u64 = after
7302 .lines()
7303 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
7304 .and_then(|l| l.split_whitespace().nth(1))
7305 .and_then(|v| v.parse().ok())
7306 .expect("misses metric line must parse");
7307 assert!(
7308 misses > 0,
7309 "decode must have read experts through the store: {after}"
7310 );
7311 }
7312
7313 fn weather_tool() -> serde_json::Value {
7314 serde_json::json!({
7315 "type": "function",
7316 "function": {
7317 "name": "get_weather",
7318 "description": "Get the current weather for a location.",
7319 "parameters": {
7320 "type": "object",
7321 "properties": {"location": {"type": "string"}},
7322 "required": ["location"]
7323 }
7324 }
7325 })
7326 }
7327
7328 fn weather_tool_def() -> ToolDef {
7329 ToolDef {
7330 kind: "function".to_string(),
7331 function: ToolFunctionDef {
7332 name: "get_weather".to_string(),
7333 description: Some("Get the current weather for a location.".to_string()),
7334 parameters: Some(serde_json::json!({
7335 "type": "object",
7336 "properties": {"location": {"type": "string"}},
7337 "required": ["location"]
7338 })),
7339 },
7340 }
7341 }
7342
7343 #[test]
7344 fn tool_preamble_mentions_every_tool_name_and_description() {
7345 let preamble = tool_preamble(&[weather_tool_def()]);
7346 assert!(preamble.contains("get_weather"));
7347 assert!(preamble.contains("Get the current weather for a location."));
7348 assert!(preamble.contains("<tool_call>"));
7349 assert!(preamble.contains("</tool_call>"));
7350 }
7351
7352 #[test]
7353 fn a_real_marker_becomes_a_structured_tool_call() {
7354 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7355 let (message, finish) = build_response_message(
7356 text.to_string(),
7357 &[weather_tool_def()],
7358 output::OutputPosture::for_model("test-model"),
7359 "stop",
7360 );
7361 assert_eq!(finish, "tool_calls");
7362 let calls = message.tool_calls.expect("must carry a tool call");
7363 assert_eq!(calls[0].function.name, "get_weather");
7364 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7365 assert_eq!(parsed["location"], "Paris");
7366 }
7367
7368 #[test]
7369 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7370 let (message, finish) = build_response_message(
7371 "just an answer".to_string(),
7372 &[weather_tool_def()],
7373 output::OutputPosture::for_model("test-model"),
7374 "stop",
7375 );
7376 assert_eq!(finish, "stop");
7377 assert!(message.tool_calls.is_none());
7378 assert_eq!(message.content.as_deref(), Some("just an answer"));
7379 }
7380
7381 #[test]
7384 fn a_malformed_payload_is_not_a_tool_call() {
7385 let (message, finish) = build_response_message(
7386 "<tool_call>not valid json at all</tool_call>".to_string(),
7387 &[weather_tool_def()],
7388 output::OutputPosture::for_model("test-model"),
7389 "stop",
7390 );
7391 assert_eq!(finish, "stop");
7392 assert!(message.tool_calls.is_none());
7393 }
7394
7395 #[test]
7398 fn a_tool_that_was_never_offered_is_not_returned() {
7399 let (message, finish) = build_response_message(
7400 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7401 &[weather_tool_def()],
7402 output::OutputPosture::for_model("test-model"),
7403 "stop",
7404 );
7405 assert_eq!(finish, "stop");
7406 assert!(message.tool_calls.is_none());
7407 }
7408
7409 #[test]
7411 fn marker_text_with_no_tools_offered_stays_content() {
7412 let (message, finish) = build_response_message(
7413 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7414 &[],
7415 output::OutputPosture::for_model("test-model"),
7416 "stop",
7417 );
7418 assert_eq!(finish, "stop");
7419 assert!(message.tool_calls.is_none());
7420 assert!(message.content.is_some());
7421 }
7422
7423 #[test]
7427 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7428 let opened = std::cell::Cell::new(0usize);
7429 let mut parser = crate::policy::parser::ToolCallParser::new(
7430 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7431 vec![
7432 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7433 "write_file",
7434 serde_json::json!({"type": "object", "properties": {
7435 "path": {"type": "string"},
7436 "contents": {"type": "string"}
7437 }}),
7438 ),
7439 ],
7440 );
7441 let wire = "<tool_call><function=write_file>\
7442 <parameter=path>\n/tmp/x\n</parameter>\
7443 <parameter=contents>\nhello world\n</parameter>\
7444 </function></tool_call>";
7445
7446 let mut deltas = Vec::new();
7447 let mut text = String::new();
7448 for piece in wire.as_bytes().chunks(7) {
7449 let chunk = String::from_utf8_lossy(piece).into_owned();
7450 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7451 text.push_str(&more_text);
7452 deltas.extend(more);
7453 }
7454 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7455 text.push_str(&more_text);
7456 deltas.extend(more);
7457
7458 assert_eq!(opened.get(), 1, "one call opened");
7459 assert!(text.is_empty(), "the markers are not content: {text:?}");
7460
7461 let first = &deltas[0];
7462 assert_eq!(first.index, 0);
7463 assert_eq!(first.id.as_deref(), Some("call_0"));
7464 assert_eq!(first.kind, Some("function"));
7465 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7466
7467 let joined: String = deltas
7470 .iter()
7471 .filter_map(|d| d.function.arguments.clone())
7472 .collect();
7473 let parsed: serde_json::Value =
7474 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7475 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7476 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7477 assert!(
7478 deltas.len() >= 3,
7479 "the arguments arrived in pieces, not whole: {}",
7480 deltas.len()
7481 );
7482 assert!(
7483 deltas[1..].iter().all(|d| d.function.name.is_none()),
7484 "only the opening delta carries identity"
7485 );
7486 }
7487
7488 #[test]
7490 fn text_around_a_streamed_call_is_still_content() {
7491 let opened = std::cell::Cell::new(0usize);
7492 let mut parser = crate::policy::parser::ToolCallParser::new(
7493 crate::policy::parser::ToolCallFormat::Qwen25,
7494 vec![crate::policy::parser::tool_call::ToolSchema::new(
7495 "get_weather",
7496 )],
7497 );
7498 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7499 \"arguments\": {}}</tool_call> done";
7500 let mut text = String::new();
7501 for piece in wire.as_bytes().chunks(5) {
7502 let chunk = String::from_utf8_lossy(piece).into_owned();
7503 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7504 text.push_str(&more);
7505 }
7506 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7507 text.push_str(&more);
7508
7509 assert_eq!(opened.get(), 1);
7510 assert!(text.starts_with("let me check. "), "{text:?}");
7511 assert!(text.ends_with(" done"), "{text:?}");
7512 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7513 }
7514
7515 #[test]
7518 fn a_reasoning_block_is_split_out_of_the_answer() {
7519 let (message, finish) = build_response_message(
7520 "<think>weighing it up</think>The answer is 4.".to_string(),
7521 &[],
7522 output::OutputPosture::for_model("Qwen3-8B"),
7523 "stop",
7524 );
7525 assert_eq!(finish, "stop");
7526 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7527 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7528 }
7529
7530 #[test]
7533 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7534 let (message, _) = build_response_message(
7535 "Use the <think> tag like this.".to_string(),
7536 &[],
7537 output::OutputPosture::for_model("llama-3.1-8b"),
7538 "stop",
7539 );
7540 assert_eq!(
7541 message.content.as_deref(),
7542 Some("Use the <think> tag like this.")
7543 );
7544 assert!(message.reasoning_content.is_none());
7545 }
7546
7547 #[tokio::test]
7555 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7556 let app = test_app();
7557 let body = serde_json::json!({
7558 "model": "m",
7559 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7560 "max_tokens": 4,
7561 "temperature": 0,
7562 });
7563 let resp = post_json(&app, body).await;
7564 let message = &resp["choices"][0]["message"];
7565 assert!(message["content"].is_string());
7566 assert!(message.get("tool_calls").is_none());
7567 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7568 assert_eq!(resp["usage"]["completion_tokens"], 4);
7569 assert_eq!(
7570 resp["usage"]["total_tokens"],
7571 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7572 );
7573 }
7574
7575 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7576 use http_body_util::BodyExt;
7577 use tower::ServiceExt;
7578
7579 let response = app
7580 .clone()
7581 .oneshot(
7582 axum::http::Request::builder()
7583 .method("GET")
7584 .uri(uri)
7585 .body(axum::body::Body::empty())
7586 .unwrap(),
7587 )
7588 .await
7589 .unwrap();
7590 let status = response.status();
7591 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7592 (status, serde_json::from_slice(&bytes).unwrap())
7593 }
7594
7595 #[tokio::test]
7596 async fn health_answers_a_capability_handshake_not_a_boolean() {
7597 let app = test_app();
7598 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7599 assert_eq!(status, StatusCode::OK);
7600
7601 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7602 assert_eq!(health.state, frink_api::HealthState::Ready);
7603 assert!(health.pid > 0);
7604 assert!(health.server_time_unix_ms > 0);
7605 assert_eq!(health.last_request_age_seconds, None);
7608
7609 for id in [
7612 frink_api::health::capability::CPU,
7613 frink_api::health::capability::METAL,
7614 frink_api::health::capability::CUDA,
7615 frink_api::health::capability::REAL_WEIGHTS,
7616 frink_api::health::capability::CONTINUOUS_BATCHING,
7617 ] {
7618 let cap = health
7619 .capability(id)
7620 .unwrap_or_else(|| panic!("{id} missing"));
7621 assert!(!cap.reason.is_empty(), "{cap:?}");
7622 assert!(!cap.detail.is_empty(), "{cap:?}");
7623 }
7624 let weights = health
7628 .capability(frink_api::health::capability::REAL_WEIGHTS)
7629 .unwrap();
7630 assert!(!weights.available);
7631 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7632 assert!(health.model.as_ref().unwrap().synthetic_weights);
7633 }
7634
7635 #[tokio::test]
7636 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7637 let app = test_app();
7638 let _ = post_json(
7639 &app,
7640 serde_json::json!({
7641 "model": "m",
7642 "messages": [{"role": "user", "content": "\u{1}"}],
7643 "max_tokens": 1,
7644 "temperature": 0,
7645 }),
7646 )
7647 .await;
7648 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7649 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7650 let age = health
7651 .last_request_age_seconds
7652 .expect("a served request is evidence of liveness");
7653 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7654 }
7655
7656 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7658 use http_body_util::BodyExt;
7659 use tower::ServiceExt;
7660
7661 let response = app
7662 .clone()
7663 .oneshot(
7664 axum::http::Request::builder()
7665 .method("POST")
7666 .uri("/v1/chat/completions")
7667 .header("content-type", "application/json")
7668 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7669 .unwrap(),
7670 )
7671 .await
7672 .unwrap();
7673 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7674 String::from_utf8(bytes.to_vec())
7675 .unwrap()
7676 .lines()
7677 .filter_map(|line| line.strip_prefix("data: "))
7678 .filter(|payload| *payload != "[DONE]")
7679 .map(|payload| serde_json::from_str(payload).unwrap())
7680 .collect()
7681 }
7682
7683 #[tokio::test]
7684 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7685 let app = test_app();
7686 let chunks = post_sse_chunks(
7687 &app,
7688 serde_json::json!({
7689 "model": "m",
7690 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7691 "max_tokens": 4,
7692 "temperature": 0,
7693 "stream": true,
7694 }),
7695 )
7696 .await;
7697
7698 assert!(!chunks.is_empty());
7699 let request_id = chunks[0]["request_id"]
7700 .as_str()
7701 .expect("the first chunk names the request")
7702 .to_string();
7703 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7704 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7707 assert!(
7708 chunk.get("request_id").is_none(),
7709 "chunk {i} repeats request_id"
7710 );
7711 }
7712 for chunk in &chunks {
7715 assert_eq!(chunk["id"], serde_json::json!(request_id));
7716 }
7717
7718 let other = post_sse_chunks(
7719 &app,
7720 serde_json::json!({
7721 "model": "m",
7722 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7723 "max_tokens": 4,
7724 "temperature": 0,
7725 "stream": true,
7726 }),
7727 )
7728 .await;
7729 assert_ne!(
7730 other[0]["request_id"].as_str().unwrap(),
7731 request_id,
7732 "two concurrent chats must not share an id"
7733 );
7734 }
7735
7736 #[tokio::test]
7737 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7738 let app = test_app();
7739 let resp = post_json(
7740 &app,
7741 serde_json::json!({
7742 "model": "m",
7743 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7744 "max_tokens": 2,
7745 "temperature": 0,
7746 }),
7747 )
7748 .await;
7749 assert_eq!(resp["id"], resp["request_id"]);
7750 assert!(resp["request_id"]
7751 .as_str()
7752 .unwrap()
7753 .starts_with("chatcmpl-"));
7754 }
7755
7756 #[tokio::test]
7759 async fn usage_carries_separate_prefill_and_decode_timings() {
7760 let app = test_app();
7761 let resp = post_json(
7762 &app,
7763 serde_json::json!({
7764 "model": "m",
7765 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7766 "max_tokens": 4,
7767 "temperature": 0,
7768 }),
7769 )
7770 .await;
7771 let usage = &resp["usage"];
7772 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7773 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7774 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7775 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7776 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7778 }
7779
7780 #[tokio::test]
7788 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7789 let app = test_app();
7790 let body = serde_json::json!({
7791 "model": "m",
7792 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7793 "max_tokens": 4,
7794 "temperature": 0,
7795 "tools": [weather_tool()],
7796 });
7797 let resp = post_json(&app, body).await;
7798 let message = &resp["choices"][0]["message"];
7799 assert!(
7800 message["content"].is_string(),
7801 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7802 );
7803 assert!(message.get("tool_calls").is_none());
7804 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7808 }
7809
7810 #[tokio::test]
7814 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7815 let app = test_app();
7816 let body = serde_json::json!({
7817 "model": "m",
7818 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7819 "max_tokens": 3,
7820 "temperature": 0,
7821 });
7822 let first = post_json(&app, body.clone()).await;
7823 assert_eq!(first["frink_cache"], "miss");
7824 let second = post_json(&app, body).await;
7825 assert_eq!(second["frink_cache"], "hit");
7826 assert_eq!(
7827 first["choices"][0]["message"]["content"],
7828 second["choices"][0]["message"]["content"]
7829 );
7830 assert_eq!(
7831 first["choices"][0]["finish_reason"],
7832 second["choices"][0]["finish_reason"]
7833 );
7834 assert_eq!(first["usage"], second["usage"]);
7835 assert_eq!(second["usage"]["completion_tokens"], 3);
7836 }
7837
7838 #[tokio::test]
7849 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7850 let app = test_app();
7851 let plain = serde_json::json!({
7852 "model": "m",
7853 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7854 "max_tokens": 3,
7855 "temperature": 0,
7856 });
7857
7858 let first = post_json(&app, plain.clone()).await;
7859 assert_eq!(first["frink_cache"], "miss");
7860 let unconstrained = first["choices"][0]["message"]["content"]
7861 .as_str()
7862 .expect("content")
7863 .to_string();
7864
7865 let mut constrained = plain.clone();
7866 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7867 let second = post_json(&app, constrained).await;
7868 assert_eq!(
7869 second["frink_cache"], "miss",
7870 "a grammar is part of the key, so this body has never been answered"
7871 );
7872 let constrained_answer = second["choices"][0]["message"]["content"]
7876 .as_str()
7877 .expect("content")
7878 .to_string();
7879 assert!(
7880 constrained_answer.contains("-> \"yes\"]"),
7881 "the grammar must have been compiled AND applied, not skipped \
7882 by a cache hit: {constrained_answer}"
7883 );
7884 assert_ne!(
7885 constrained_answer, unconstrained,
7886 "the constrained request was served the unconstrained answer"
7887 );
7888
7889 let third = post_json(&app, plain).await;
7893 assert_eq!(third["frink_cache"], "hit");
7894 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7895 }
7896
7897 #[tokio::test]
7919 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7920 let state = Arc::new(test_state(
7921 test_model_full_byte_vocab(),
7922 ResponseCache::new(1000, Duration::from_secs(3600)),
7923 ));
7924 let app = test_app_with_state(state.clone());
7925 let plain = serde_json::json!({
7926 "model": "m",
7927 "messages": [
7928 {"role": "system", "content": "Answer in JSON when it helps."},
7929 {"role": "user", "content": "\u{1}\u{2}"},
7930 ],
7931 "max_tokens": 3,
7932 "temperature": 0,
7933 });
7934
7935 let first = post_json(&app, plain.clone()).await;
7936 assert_eq!(first["frink_cache"], "miss");
7937 assert_eq!(state.cache_stats().entries, 1);
7938
7939 let mut as_json = plain.clone();
7940 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7941 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7942 assert_eq!(
7943 status,
7944 StatusCode::BAD_REQUEST,
7945 "the demo banner is not a JSON object, whoever generated it"
7946 );
7947 assert_eq!(
7948 state.cache_stats().hits,
7949 0,
7950 "a json_object request must not be answered from an entry the \
7951 JSON mask never produced"
7952 );
7953 assert_eq!(
7954 state.cache_stats().entries,
7955 2,
7956 "json_object must key its own entry, not reuse the unconstrained \
7957 one it happens to render the same prompt as"
7958 );
7959 }
7960
7961 #[tokio::test]
7973 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7974 let app = test_app_with_state(Arc::new(test_state(
7975 test_model_full_byte_vocab_with_eos(Some(0x77)),
7976 ResponseCache::new(1000, Duration::from_secs(3600)),
7977 )));
7978 let body = serde_json::json!({
7979 "model": "m",
7980 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7981 "max_tokens": 6,
7982 "temperature": 0,
7983 });
7984
7985 let stopped = post_json(&app, body.clone()).await;
7986 assert_eq!(stopped["frink_cache"], "miss");
7987 assert_eq!(
7988 stopped["choices"][0]["finish_reason"], "stop",
7989 "the fixture is only meaningful if the model's EOS really fires here"
7990 );
7991 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7992
7993 let mut ignoring = body.clone();
7994 ignoring["ignore_eos"] = serde_json::json!(true);
7995 let ran_on = post_json(&app, ignoring).await;
7996 assert_eq!(
7997 ran_on["frink_cache"], "miss",
7998 "ignore_eos is part of the key, so this body has never been answered"
7999 );
8000 assert_eq!(
8001 ran_on["usage"]["completion_tokens"], 6,
8002 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
8003 );
8004 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
8005 assert_ne!(
8006 ran_on["choices"][0]["message"]["content"],
8007 stopped["choices"][0]["message"]["content"]
8008 );
8009 }
8010
8011 #[tokio::test]
8017 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
8018 let session_app = test_app();
8019 let manual_app = test_app();
8020
8021 let turn1 = post_json(
8023 &session_app,
8024 serde_json::json!({
8025 "model": "m",
8026 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8027 "session_id": "s1",
8028 "max_tokens": 5,
8029 "temperature": 0,
8030 }),
8031 )
8032 .await;
8033 let reply1 = turn1["choices"][0]["message"]["content"]
8034 .as_str()
8035 .unwrap()
8036 .to_string();
8037
8038 let manual_turn1 = post_json(
8043 &manual_app,
8044 serde_json::json!({
8045 "model": "m",
8046 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
8047 "max_tokens": 5,
8048 "temperature": 0,
8049 }),
8050 )
8051 .await;
8052 assert_eq!(
8053 manual_turn1["choices"][0]["message"]["content"]
8054 .as_str()
8055 .unwrap(),
8056 reply1
8057 );
8058
8059 let turn2 = post_json(
8061 &session_app,
8062 serde_json::json!({
8063 "model": "m",
8064 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
8065 "session_id": "s1",
8066 "max_tokens": 5,
8067 "temperature": 0,
8068 }),
8069 )
8070 .await;
8071 let reply2 = turn2["choices"][0]["message"]["content"]
8072 .as_str()
8073 .unwrap()
8074 .to_string();
8075
8076 let manual_turn2 = post_json(
8080 &manual_app,
8081 serde_json::json!({
8082 "model": "m",
8083 "messages": [
8084 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
8085 {"role": "assistant", "content": reply1},
8086 {"role": "user", "content": "\u{4}\u{5}"},
8087 ],
8088 "max_tokens": 5,
8089 "temperature": 0,
8090 }),
8091 )
8092 .await;
8093 assert_eq!(
8094 manual_turn2["choices"][0]["message"]["content"]
8095 .as_str()
8096 .unwrap(),
8097 reply2,
8098 "resuming a session must produce identical output to manually resending the full history"
8099 );
8100 }
8101
8102 #[test]
8105 fn lock_cache_recovers_from_a_poisoned_mutex() {
8106 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
8107
8108 let poison_cache = Arc::clone(&cache);
8109 let _ = std::thread::spawn(move || {
8110 let _guard = poison_cache.lock().unwrap();
8111 panic!("simulated panic while holding the lock");
8112 })
8113 .join();
8114
8115 let recovered = lock_cache(&cache);
8117 assert_eq!(recovered.stats().entries, 0);
8118 }
8119
8120 #[test]
8121 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8122 let mut req_body = serde_json::json!({
8123 "model": "m",
8124 "messages": [{"role": "user", "content": "hi"}],
8125 });
8126 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8127 assert!(
8128 req.is_cacheable(),
8129 "default (temperature 0) must be cacheable"
8130 );
8131
8132 req_body["temperature"] = serde_json::json!(0.8);
8133 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8134 assert!(
8135 !req.is_cacheable(),
8136 "unseeded sampling must never be cacheable"
8137 );
8138
8139 req_body["seed"] = serde_json::json!(42);
8140 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
8141 assert!(
8142 req.is_cacheable(),
8143 "sampling with an explicit seed is deterministic and must be cacheable"
8144 );
8145 }
8146
8147 const GRADED: &str = "{% if reasoning_effort %}\
8151 {% if reasoning_effort not in ['low','medium','high'] %}\
8152 {{ raise_exception('unsupported effort') }}\
8153 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
8154 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
8155
8156 fn graded_template() -> chat_template::PromptTemplate {
8157 chat_template::PromptTemplate::from_gguf_metadata(
8158 Some(GRADED),
8159 Some("qwen3"),
8160 false,
8161 true,
8162 None,
8163 None,
8164 )
8165 }
8166
8167 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
8168 serde_json::from_value(value).expect("request")
8169 }
8170
8171 #[test]
8178 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
8179 let req = chat_request(serde_json::json!({
8180 "model": "m",
8181 "messages": [{"role": "user", "content": "hi"}],
8182 "grammar": "root ::= \"a\"+",
8183 }));
8184 req.validate_supported_fields()
8185 .expect("a valid grammar is a valid request");
8186 let params = req
8187 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8188 .expect("a valid grammar compiles at params time too");
8189 assert!(
8190 params.grammar.is_some(),
8191 "the grammar was dropped between the wire and the sampler"
8192 );
8193 assert!(
8194 params.needs_vocab_logits(),
8195 "a grammar request that may fold lm_head into a GPU argmax is \
8196 a grammar request served unconstrained"
8197 );
8198
8199 let plain = chat_request(serde_json::json!({
8200 "model": "m",
8201 "messages": [{"role": "user", "content": "hi"}],
8202 }));
8203 assert!(plain
8204 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8205 .unwrap()
8206 .grammar
8207 .is_none());
8208 }
8209
8210 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
8211 chat_request(serde_json::json!({
8212 "model": "m",
8213 "messages": [{"role": "user", "content": "weather in Rome?"}],
8214 "tools": [weather_tool()],
8215 "tool_choice": tool_choice,
8216 }))
8217 }
8218
8219 #[test]
8223 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
8224 for choice in [
8225 serde_json::json!("required"),
8226 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
8227 ] {
8228 let req = tool_request(choice.clone());
8229 req.validate_supported_fields()
8230 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
8231 let params = req
8232 .generation_params_for_template(
8233 &graded_template(),
8234 "Qwen3-8B",
8235 crate::sampling_knobs::SamplerModel::absent(),
8236 )
8237 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
8238 let grammar = params
8239 .grammar
8240 .as_ref()
8241 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
8242 assert!(
8243 grammar.is_awaiting_trigger(),
8244 "the model must be free to think before it calls"
8245 );
8246 assert!(
8247 !grammar.allows_eog(),
8248 "{choice} must not be able to end the turn without a call"
8249 );
8250 assert!(
8256 params.needs_vocab_logits(),
8257 "{choice} would let a backend return a token id instead of logits"
8258 );
8259 assert!(
8260 !generate::greedy_gpu_fold_allowed(¶ms),
8261 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8262 );
8263 }
8264 }
8265
8266 #[test]
8268 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8269 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8270 let req = tool_request(choice.clone());
8271 req.validate_supported_fields().expect("still supported");
8272 let params = match req.generation_params_for_template(
8273 &graded_template(),
8274 "Qwen3-8B",
8275 crate::sampling_knobs::SamplerModel::absent(),
8276 ) {
8277 Ok(p) => p,
8278 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
8279 };
8280 assert!(
8281 params.grammar.is_none(),
8282 "{choice} does not force a call and must not be constrained"
8283 );
8284 }
8285 }
8286
8287 #[test]
8290 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8291 let req = chat_request(serde_json::json!({
8293 "model": "m",
8294 "messages": [{"role": "user", "content": "hi"}],
8295 "tool_choice": "required",
8296 }));
8297 let (status, _) = req
8298 .validate_supported_fields()
8299 .expect_err("nothing to call");
8300 assert_eq!(status, StatusCode::BAD_REQUEST);
8301
8302 let req =
8304 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8305 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8306 assert_eq!(status, StatusCode::BAD_REQUEST);
8307 assert_eq!(body["error"]["param"], "tool_choice");
8308
8309 let req = tool_request(serde_json::json!({"type": "function"}));
8311 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8312 assert_eq!(status, StatusCode::BAD_REQUEST);
8313
8314 let req = chat_request(serde_json::json!({
8316 "model": "m",
8317 "messages": [{"role": "user", "content": "hi"}],
8318 "tools": [weather_tool()],
8319 "tool_choice": "required",
8320 "grammar": "root ::= \"a\"+",
8321 }));
8322 let (status, _) = req
8323 .validate_supported_fields()
8324 .expect_err("a grammar and a forced call are two constraints");
8325 assert_eq!(status, StatusCode::BAD_REQUEST);
8326
8327 let req = tool_request(serde_json::json!("required"));
8333 let (status, Json(body)) = match req.generation_params_for_template(
8334 &graded_template(),
8335 "muse-glimmer-8b",
8336 crate::sampling_knobs::SamplerModel::absent(),
8337 ) {
8338 Err(e) => e,
8339 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
8340 };
8341 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8342 assert!(
8343 body["error"]["message"]
8344 .as_str()
8345 .unwrap()
8346 .contains("muse_glimmer"),
8347 "{body}"
8348 );
8349
8350 let req = tool_request(serde_json::json!("required"));
8356 let params = req
8357 .generation_params_for_template(
8358 &graded_template(),
8359 "gemma-4-E2B-it",
8360 crate::sampling_knobs::SamplerModel::absent(),
8361 )
8362 .expect("a gemma4 forced tool_choice is served");
8363 assert!(
8364 params.grammar.is_some(),
8365 "a forced tool_choice must arrive as the generation's grammar"
8366 );
8367 }
8368
8369 #[test]
8372 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8373 let req = chat_request(serde_json::json!({
8374 "model": "m",
8375 "messages": [{"role": "user", "content": "hi"}],
8376 "grammar": "root ::= \"a",
8377 }));
8378 let (status, Json(body)) = req
8379 .validate_supported_fields()
8380 .expect_err("this does not parse");
8381 assert_eq!(status, StatusCode::BAD_REQUEST);
8382 assert_eq!(body["error"]["param"], "grammar");
8383 assert!(
8384 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8385 .is_err(),
8386 "and again at params time"
8387 );
8388 }
8389
8390 #[test]
8397 fn response_format_json_schema_becomes_the_requests_grammar() {
8398 let req = chat_request(serde_json::json!({
8399 "model": "m",
8400 "messages": [{"role": "user", "content": "hi"}],
8401 "response_format": {
8402 "type": "json_schema",
8403 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8404 },
8405 }));
8406 req.validate_supported_fields()
8407 .expect("a boolean schema converts");
8408 let params = req
8409 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8410 .expect("and compiles");
8411 let grammar = params.grammar.expect("the schema is the grammar");
8412 let mut g = (*grammar).clone();
8413 g.accept_token(0, b"true").expect("a boolean is accepted");
8414 assert!(g.allows_eog(), "and completes the parse");
8415 assert!(
8416 !params.json_object,
8417 "a schema is not the json_object character-class mask"
8418 );
8419 }
8420
8421 #[test]
8425 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8426 let req = chat_request(serde_json::json!({
8427 "model": "m",
8428 "messages": [{"role": "user", "content": "hi"}],
8429 "response_format": {
8430 "type": "json_schema",
8431 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8432 },
8433 }));
8434 let (status, Json(body)) = req
8435 .validate_supported_fields()
8436 .expect_err("minimum has no grammar in this port");
8437 assert_eq!(status, StatusCode::BAD_REQUEST);
8438 assert!(
8439 body["error"]["message"]
8440 .as_str()
8441 .expect("a message")
8442 .contains("minimum"),
8443 "the refusal must name the keyword: {body}"
8444 );
8445 assert!(
8446 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8447 .is_err(),
8448 "and again at params time"
8449 );
8450 }
8451
8452 #[test]
8458 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8459 let req = chat_request(serde_json::json!({
8460 "model": "m",
8461 "messages": [{"role": "user", "content": "hi"}],
8462 "tool_choice": "required",
8463 "tools": [{
8464 "type": "function",
8465 "function": {"name": "f", "parameters": {"type": "object"}},
8466 }],
8467 "response_format": {
8468 "type": "json_schema",
8469 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8470 },
8471 }));
8472 let (status, Json(body)) = req
8473 .validate_supported_fields()
8474 .expect_err("two constraints, one generation");
8475 assert_eq!(status, StatusCode::BAD_REQUEST);
8476 assert_eq!(body["error"]["param"], "tool_choice");
8477 }
8478
8479 #[test]
8482 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8483 let req = chat_request(serde_json::json!({
8484 "model": "m",
8485 "messages": [{"role": "user", "content": "hi"}],
8486 }));
8487 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8488 }
8489
8490 #[test]
8495 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8496 let asked = chat_request(serde_json::json!({
8497 "model": "m",
8498 "messages": [{"role": "user", "content": "hi"}],
8499 "min_p": 0.07,
8500 }));
8501 assert_eq!(
8502 asked
8503 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8504 .expect("knobs")
8505 .min_p,
8506 0.07
8507 );
8508
8509 let silent = chat_request(serde_json::json!({
8510 "model": "m",
8511 "messages": [{"role": "user", "content": "hi"}],
8512 }));
8513 assert_eq!(
8514 silent
8515 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8516 .expect("knobs")
8517 .min_p,
8518 0.0,
8519 "an unset min_p must be off, not llama.cpp's CLI default"
8520 );
8521 }
8522
8523 #[test]
8532 fn no_sampler_knob_is_missing_from_the_cache_key() {
8533 let base = serde_json::json!({
8534 "model": "m",
8535 "messages": [{"role": "user", "content": "hi"}],
8536 "seed": 1,
8537 });
8538 let key_for = |body: serde_json::Value| {
8539 let req = chat_request(body);
8540 let params = req
8541 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8542 .expect("params");
8543 req.cache_key("prompt", ¶ms)
8544 };
8545 let baseline = key_for(base.clone());
8546 for (knob, value) in [
8547 ("temperature", serde_json::json!(0.5)),
8548 ("top_p", serde_json::json!(0.9)),
8549 ("min_p", serde_json::json!(0.05)),
8550 ("top_k", serde_json::json!(40)),
8551 ("repetition_penalty", serde_json::json!(1.1)),
8552 ("presence_penalty", serde_json::json!(0.3)),
8553 ("frequency_penalty", serde_json::json!(0.3)),
8554 (
8555 "samplers",
8556 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8557 ),
8558 ] {
8559 let mut body = base.clone();
8560 body[knob] = value;
8561 assert_ne!(
8562 key_for(body),
8563 baseline,
8564 "`{knob}` is not in the cache key: two requests differing \
8565 only in it would share one cached answer"
8566 );
8567 }
8568 }
8569
8570 #[test]
8580 fn no_constraint_is_missing_from_the_cache_key() {
8581 let base = serde_json::json!({
8582 "model": "m",
8583 "messages": [{"role": "user", "content": "pick one"}],
8584 });
8585 let key_for = |body: serde_json::Value| {
8586 let req = chat_request(body);
8587 let params = req
8588 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8589 .expect("params");
8590 req.cache_key("prompt", ¶ms)
8591 };
8592 let baseline = key_for(base.clone());
8593 for (field, value) in [
8594 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8595 (
8596 "response_format",
8597 serde_json::json!({"type": "json_object"}),
8598 ),
8599 (
8600 "response_format",
8601 serde_json::json!({"type": "json_schema", "json_schema": {
8602 "name": "answer",
8603 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8604 }}),
8605 ),
8606 ("ignore_eos", serde_json::json!(true)),
8607 ("stop", serde_json::json!(["\n"])),
8608 ("max_tokens", serde_json::json!(7)),
8609 ] {
8610 let mut body = base.clone();
8611 body[field] = value.clone();
8612 assert_ne!(
8613 key_for(body),
8614 baseline,
8615 "`{field}: {value}` is not in the cache key: two requests \
8616 differing only in it would share one cached answer"
8617 );
8618 }
8619 }
8620
8621 #[test]
8625 fn an_explicit_zero_output_budget_is_a_client_error() {
8626 let req = chat_request(serde_json::json!({
8627 "model": "m",
8628 "messages": [{"role": "user", "content": "hi"}],
8629 "max_tokens": 0,
8630 }));
8631 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8632 assert_eq!(status, StatusCode::BAD_REQUEST);
8633 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8634 }
8635
8636 #[test]
8639 fn a_request_can_turn_thinking_off() {
8640 let template = graded_template();
8641 for body in [
8642 serde_json::json!({
8643 "model": "m",
8644 "messages": [{"role": "user", "content": "hi"}],
8645 "reasoning_effort": "none",
8646 }),
8647 serde_json::json!({
8648 "model": "m",
8649 "messages": [{"role": "user", "content": "hi"}],
8650 "thinking": {"type": "disabled"},
8651 }),
8652 ] {
8653 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8654 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8655 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8656 assert!(!kwargs.contains_key("reasoning_effort"));
8659 }
8660 }
8661
8662 #[test]
8665 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8666 let template = graded_template();
8667 let kwargs = chat_request(serde_json::json!({
8668 "model": "m",
8669 "messages": [{"role": "user", "content": "hi"}],
8670 "reasoning_effort": "high",
8671 "thinking": {"type": "disabled"},
8672 }))
8673 .resolve_template_kwargs(&template);
8674 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8675 assert!(!kwargs.contains_key("reasoning_effort"));
8676 }
8677
8678 #[test]
8681 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8682 let req = chat_request(serde_json::json!({
8683 "model": "m",
8684 "messages": [{"role": "user", "content": "hi"}],
8685 "thinking": {"type": "disable"},
8686 }));
8687 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8688 assert_eq!(status, StatusCode::BAD_REQUEST);
8689 }
8690
8691 #[test]
8694 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8695 let template = graded_template();
8696 let kwargs = chat_request(serde_json::json!({
8697 "model": "m",
8698 "messages": [{"role": "user", "content": "hi"}],
8699 "reasoning_effort": "none",
8700 "chat_template_kwargs": {"enable_thinking": true},
8701 }))
8702 .resolve_template_kwargs(&template);
8703 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8704 }
8705
8706 #[test]
8710 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8711 let template = graded_template();
8712 let req = chat_request(serde_json::json!({
8713 "model": "m",
8714 "messages": [{"role": "user", "content": "hi"}],
8715 "reasoning_effort": "minimal",
8716 }));
8717 let kwargs = req.resolve_template_kwargs(&template);
8718 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8719 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8720 assert!(prompt.starts_with("E:low|"), "{prompt}");
8721 }
8722
8723 #[test]
8727 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8728 let template = graded_template();
8729 let req = chat_request(serde_json::json!({
8730 "model": "m",
8731 "messages": [{"role": "user", "content": "hi"}],
8732 "chat_template_kwargs": {"reasoning_effort": "none"},
8733 }));
8734 let kwargs = req.resolve_template_kwargs(&template);
8735 assert!(!kwargs.contains_key("reasoning_effort"));
8736 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8737 assert_eq!(prompt, "hi");
8738 }
8739
8740 #[test]
8743 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8744 let template = graded_template();
8745 let req = chat_request(serde_json::json!({
8746 "model": "m",
8747 "messages": [{"role": "user", "content": "hi"}],
8748 "reasoning_effort": "low",
8749 "chat_template_kwargs": {"reasoning_effort": "high"},
8750 }));
8751 assert_eq!(
8752 req.resolve_template_kwargs(&template)["reasoning_effort"],
8753 serde_json::json!("high")
8754 );
8755 }
8756
8757 #[test]
8761 fn offering_tools_turns_thinking_on_by_itself() {
8762 let template = graded_template();
8763 let quiet = chat_request(serde_json::json!({
8764 "model": "m",
8765 "messages": [{"role": "user", "content": "hi"}],
8766 }));
8767 assert!(!quiet
8768 .resolve_template_kwargs(&template)
8769 .contains_key("enable_thinking"));
8770
8771 let with_tools = chat_request(serde_json::json!({
8772 "model": "m",
8773 "messages": [{"role": "user", "content": "hi"}],
8774 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8775 }));
8776 let kwargs = with_tools.resolve_template_kwargs(&template);
8777 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8778 let prompt =
8779 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8780 assert!(prompt.starts_with("THINK|"), "{prompt}");
8781 }
8782
8783 #[test]
8788 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8789 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8790 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8791 Some("qwen3"),
8792 false,
8793 true,
8794 None,
8795 None,
8796 );
8797 let req = chat_request(serde_json::json!({
8798 "model": "m",
8799 "messages": [{"role": "user", "content": "hi"}],
8800 "chat_template_kwargs": {"enable_thinking": true},
8801 }));
8802 let kwargs = req.resolve_template_kwargs(&opener);
8803 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8804 assert!(prompt.ends_with("<think>"), "{prompt}");
8805
8806 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8809 let (message, _) = build_response_message(
8810 "weighing it up</think>Paris.".to_string(),
8811 &[],
8812 posture,
8813 "stop",
8814 );
8815 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8816 assert_eq!(message.content.as_deref(), Some("Paris."));
8817
8818 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8821 let (message, _) = build_response_message(
8822 "weighing it up</think>Paris.".to_string(),
8823 &[],
8824 closed,
8825 "stop",
8826 );
8827 assert_eq!(message.reasoning_content, None);
8828 }
8829
8830 #[test]
8831 fn stop_param_accepts_both_single_string_and_array() {
8832 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8833 "model": "m",
8834 "messages": [{"role": "user", "content": "hi"}],
8835 "stop": "END",
8836 }))
8837 .unwrap();
8838 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8839
8840 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8841 "model": "m",
8842 "messages": [{"role": "user", "content": "hi"}],
8843 "stop": ["A", "B"],
8844 }))
8845 .unwrap();
8846 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8847 }
8848
8849 #[test]
8850 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8851 let model = test_model();
8852 let result = run_generation(
8853 &model,
8854 "hello",
8855 &greedy_params(4),
8856 None,
8857 None,
8858 None,
8859 None,
8860 None,
8861 None,
8862 );
8863 assert!(matches!(
8864 result,
8865 Err(generate::DecodeError::TokenOutOfVocab { .. })
8866 ));
8867 }
8868
8869 #[test]
8873 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8874 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8876 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8877
8878 let holder_pool = Arc::clone(&pool);
8879 let holder = std::thread::spawn(move || {
8880 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8881 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8883 drop(held);
8884 });
8885 std::thread::sleep(Duration::from_millis(15));
8886
8887 let config = generate::KvPoolConfig {
8888 pool,
8889 queue_wait: Duration::ZERO,
8890 };
8891 let result = run_generation(
8892 &model,
8893 &prompt,
8894 &greedy_params(4),
8895 Some(&config),
8896 None,
8897 None,
8898 None,
8899 None,
8900 None,
8901 );
8902 assert!(matches!(
8903 result,
8904 Err(generate::DecodeError::KvPoolExhausted)
8905 ));
8906
8907 let (status, _body) = decode_error_response(result.unwrap_err());
8908 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8909 holder.join().unwrap();
8910 }
8911
8912 #[test]
8920 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8921 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8923 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8925 let config = generate::KvPoolConfig {
8926 pool,
8927 queue_wait: Duration::ZERO,
8928 };
8929
8930 let result = run_generation(
8931 &model,
8932 &prompt,
8933 &greedy_params(4),
8934 Some(&config),
8935 None,
8936 None,
8937 None,
8938 None,
8939 None,
8940 );
8941 let err = result.expect_err("one block cannot hold two layers' caches");
8942 assert!(
8943 matches!(
8944 &err,
8945 generate::DecodeError::KvBudgetExceeded { binding, .. }
8946 if *binding == frink_models::Ceiling::DeviceMemory.code()
8947 ),
8948 "expected an immovable device-memory refusal, got {err:?}"
8949 );
8950 let (status, _body) = decode_error_response(err);
8951 assert_eq!(status, StatusCode::BAD_REQUEST);
8952 }
8953
8954 #[test]
8960 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8961 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8962 queued: 512,
8963 cap: 512,
8964 });
8965 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8966 assert_eq!(body["error"]["retry_after_seconds"], 1);
8967 let message = body["error"]["message"].as_str().expect("message");
8968 assert!(message.contains("512"), "{message}");
8969 }
8970
8971 #[test]
8972 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8973 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8974 token: 99,
8975 vocab_size: 32,
8976 });
8977 assert!(
8978 body["error"]["retry_after_seconds"].is_null(),
8979 "retrying a prompt this model cannot tokenize never helps"
8980 );
8981 }
8982
8983 #[test]
8984 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8985 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8986 token: 99,
8987 vocab_size: 32,
8988 });
8989 assert_eq!(status, StatusCode::BAD_REQUEST);
8990 }
8991
8992 #[test]
8993 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8994 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8996 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8997 let config = generate::KvPoolConfig {
8998 pool: pool.clone(),
8999 queue_wait: Duration::ZERO,
9000 };
9001
9002 let produced = run_generation(
9003 &model,
9004 &prompt,
9005 &greedy_params(4),
9006 Some(&config),
9007 None,
9008 None,
9009 None,
9010 None,
9011 None,
9012 )
9013 .unwrap();
9014 assert_eq!(produced.choices[0].finish, FinishReason::Length);
9015 assert_eq!(
9016 pool.lock().unwrap().free_blocks(),
9017 2,
9018 "a completed request must return its blocks to the pool"
9019 );
9020 }
9021
9022 #[tokio::test]
9027 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
9028 let model = Arc::new(test_model());
9029 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
9030
9031 let mut handles = Vec::new();
9032 for _ in 0..8 {
9033 let model = Arc::clone(&model);
9034 let prompt = prompt.clone();
9035 handles.push(tokio::task::spawn_blocking(move || {
9036 run_generation(
9037 &model,
9038 &prompt,
9039 &greedy_params(6),
9040 None,
9041 None,
9042 None,
9043 None,
9044 None,
9045 None,
9046 )
9047 .unwrap()
9048 }));
9049 }
9050
9051 let mut results = Vec::new();
9052 for h in handles {
9053 results.push(h.await.unwrap());
9054 }
9055 for r in &results[1..] {
9060 assert_eq!(r.choices, results[0].choices, "choices must match");
9064 assert_eq!(
9065 r.usage.prompt_tokens, results[0].usage.prompt_tokens,
9066 "prompt token count must match"
9067 );
9068 assert_eq!(
9069 r.usage.completion_tokens, results[0].usage.completion_tokens,
9070 "completion token count must match"
9071 );
9072 }
9073 }
9074
9075 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
9082 let mut header_entries = Vec::new();
9083 let mut data = Vec::new();
9084 for (name, shape, values) in tensors {
9085 let start = data.len();
9086 for v in values {
9087 data.extend_from_slice(&v.to_le_bytes());
9088 }
9089 let end = data.len();
9090 let shape_str = shape
9091 .iter()
9092 .map(|d| d.to_string())
9093 .collect::<Vec<_>>()
9094 .join(",");
9095 header_entries.push(format!(
9096 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
9097 ));
9098 }
9099 let header = format!("{{{}}}", header_entries.join(","));
9100 let header_bytes = header.as_bytes();
9101 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
9102 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
9103 out.extend_from_slice(header_bytes);
9104 out.extend_from_slice(&data);
9105 out
9106 }
9107
9108 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9119 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9120 use frink_models::kimi_loader::KimiRealHparams;
9121 use frink_moe::{GatingFunction, MoeLayerConfig};
9122
9123 let hidden_dim = 8;
9124 let kda_num_heads = 2;
9125 let kda_head_dim = 3;
9126 let kda_proj = kda_num_heads * kda_head_dim;
9127 let conv_kernel = 4;
9128 let dense_intermediate = 5;
9129 let vocab_size = 256;
9133 let mla_num_heads = 1;
9134 let mla_q_lora_rank = 2;
9135 let mla_kv_lora_rank = 2;
9136 let mla_qk_nope_head_dim = 2;
9137 let mla_qk_rope_head_dim = 2;
9138 let mla_v_head_dim = 2;
9139
9140 let model_cfg = frink_models::ModelConfig {
9141 rope_layers: frink_models::rope_layers::RopeLayers::All,
9142 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
9143 name: "synthetic-kimi-server-test",
9144 n_layers: 1,
9145 n_mtp_blocks: 0,
9146 hidden_dim,
9147 n_heads: 1,
9148 n_kv_heads: 1,
9149 head_dim: 4,
9150 v_head_dim: None,
9151 vocab_size,
9152 rope_theta: 10000.0,
9153 rms_norm_eps: 1e-5,
9154 post_norm_eps: 1e-5,
9155 sliding_window: None,
9156 moe: MoeLayerConfig {
9157 expert_weights_scale: 1.0,
9158 routed_weight_before_ffn: false,
9159 n_experts: 1,
9160 n_experts_active: 1,
9161 n_shared_experts: 0,
9162 hidden_dim,
9163 expert_ffn_dim: 4,
9164 gating: GatingFunction::Sigmoid,
9165 norm_topk_prob: true,
9166 expert_group_count: None,
9167 expert_group_used_count: None,
9168 },
9169 n_dense_leading_layers: 1,
9174 moe_interleave_step: None,
9175 norm_function: frink_models::norm::NormFunction::Rms,
9176 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
9177 kda_layers: vec![1],
9178 full_attn_layers: vec![],
9179 mla: MlaConfig {
9180 num_heads: mla_num_heads,
9181 q_lora_rank: mla_q_lora_rank,
9182 kv_lora_rank: mla_kv_lora_rank,
9183 qk_nope_head_dim: mla_qk_nope_head_dim,
9184 qk_rope_head_dim: mla_qk_rope_head_dim,
9185 v_head_dim: mla_v_head_dim,
9186 use_output_gate: true,
9187 rope: None,
9188 },
9189 kda: KdaConfig {
9190 num_heads: kda_num_heads,
9191 head_dim: kda_head_dim,
9192 short_conv_kernel_size: conv_kernel,
9193 gate_lower_bound: -5.0,
9194 use_full_rank_gate: true,
9195 },
9196 }),
9197 rope_freqs: None,
9198 rope_attn_factor: 1.0,
9199 rope_dim: None,
9200 rope_dim_swa: None,
9201 rope_freqs_long: None,
9202 rope_freqs_short: None,
9203 rope_orig_ctx: None,
9204 rope_layout: frink_models::config::RopeLayout::Neox,
9205 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
9206 swa_layers: frink_models::swa_layers::SwaLayers::All,
9207 attn_logit_softcap: None,
9208 final_logit_softcap: None,
9209 embedding_scale: None,
9210 residual_scale: None,
9211 normed_residual_scale: None,
9212 clamp_kqv: None,
9213 attn_temperature: None,
9214 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
9215 block_sub_norms: false,
9216 parallel_residual: false,
9217 learned_positions: false,
9218 attn_value_scale: None,
9219 alibi_max_bias: None,
9220 layer_loops: None,
9221 skip_stream: false,
9222 parallel_ssm: false,
9223 swa_chunked: false,
9224 weightless_qk_norm: false,
9225 logit_multiplier: None,
9226 attention_scale: None,
9227 rope_theta_swa: None,
9228 ffn_activation: frink_models::config::FfnActivation::Swiglu,
9229 best_effort_fields: &["synthetic test config, not a real preset"],
9230 };
9231 let hp = KimiRealHparams {
9232 hidden_dim,
9233 kda_num_heads,
9234 kda_head_dim,
9235 mla_num_heads,
9236 mla_q_lora_rank,
9237 mla_kv_lora_rank,
9238 mla_qk_nope_head_dim,
9239 mla_qk_rope_head_dim,
9240 mla_v_head_dim,
9241 dense_intermediate_dim: dense_intermediate,
9242 moe_hidden_dim: hidden_dim,
9243 moe_intermediate_dim: 4,
9244 n_experts: 1,
9245 num_shared_experts: 0,
9246 };
9247
9248 let prefix = "language_model.model.layers.0";
9252 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
9253 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
9254 name: String,
9255 shape: Vec<usize>,
9256 n: usize| {
9257 tensors.push((name, shape, vec![0.01f32; n]));
9258 };
9259 push(
9260 &mut tensors,
9261 format!("{prefix}.input_layernorm.weight"),
9262 vec![hidden_dim],
9263 hidden_dim,
9264 );
9265 push(
9266 &mut tensors,
9267 format!("{prefix}.post_attention_layernorm.weight"),
9268 vec![hidden_dim],
9269 hidden_dim,
9270 );
9271 push(
9272 &mut tensors,
9273 format!("{prefix}.self_attention_res_norm.weight"),
9274 vec![hidden_dim],
9275 hidden_dim,
9276 );
9277 push(
9278 &mut tensors,
9279 format!("{prefix}.self_attention_res_proj.weight"),
9280 vec![1, hidden_dim],
9281 hidden_dim,
9282 );
9283 push(
9284 &mut tensors,
9285 format!("{prefix}.mlp_res_norm.weight"),
9286 vec![hidden_dim],
9287 hidden_dim,
9288 );
9289 push(
9290 &mut tensors,
9291 format!("{prefix}.mlp_res_proj.weight"),
9292 vec![1, hidden_dim],
9293 hidden_dim,
9294 );
9295 push(
9296 &mut tensors,
9297 format!("{prefix}.self_attn.q_proj.weight"),
9298 vec![kda_proj, hidden_dim],
9299 kda_proj * hidden_dim,
9300 );
9301 push(
9302 &mut tensors,
9303 format!("{prefix}.self_attn.k_proj.weight"),
9304 vec![kda_proj, hidden_dim],
9305 kda_proj * hidden_dim,
9306 );
9307 push(
9308 &mut tensors,
9309 format!("{prefix}.self_attn.v_proj.weight"),
9310 vec![kda_proj, hidden_dim],
9311 kda_proj * hidden_dim,
9312 );
9313 push(
9314 &mut tensors,
9315 format!("{prefix}.self_attn.q_conv1d.weight"),
9316 vec![kda_proj, 1, conv_kernel],
9317 kda_proj * conv_kernel,
9318 );
9319 push(
9320 &mut tensors,
9321 format!("{prefix}.self_attn.k_conv1d.weight"),
9322 vec![kda_proj, 1, conv_kernel],
9323 kda_proj * conv_kernel,
9324 );
9325 push(
9326 &mut tensors,
9327 format!("{prefix}.self_attn.v_conv1d.weight"),
9328 vec![kda_proj, 1, conv_kernel],
9329 kda_proj * conv_kernel,
9330 );
9331 push(
9332 &mut tensors,
9333 format!("{prefix}.self_attn.A_log"),
9334 vec![kda_num_heads],
9335 kda_num_heads,
9336 );
9337 push(
9338 &mut tensors,
9339 format!("{prefix}.self_attn.f_a_proj.weight"),
9340 vec![kda_head_dim, hidden_dim],
9341 kda_head_dim * hidden_dim,
9342 );
9343 push(
9344 &mut tensors,
9345 format!("{prefix}.self_attn.f_b_proj.weight"),
9346 vec![kda_proj, kda_head_dim],
9347 kda_proj * kda_head_dim,
9348 );
9349 push(
9350 &mut tensors,
9351 format!("{prefix}.self_attn.dt_bias"),
9352 vec![kda_proj],
9353 kda_proj,
9354 );
9355 push(
9356 &mut tensors,
9357 format!("{prefix}.self_attn.b_proj.weight"),
9358 vec![kda_num_heads, hidden_dim],
9359 kda_num_heads * hidden_dim,
9360 );
9361 push(
9362 &mut tensors,
9363 format!("{prefix}.self_attn.g_proj.weight"),
9364 vec![kda_proj, hidden_dim],
9365 kda_proj * hidden_dim,
9366 );
9367 push(
9368 &mut tensors,
9369 format!("{prefix}.self_attn.o_norm.weight"),
9370 vec![kda_head_dim],
9371 kda_head_dim,
9372 );
9373 push(
9374 &mut tensors,
9375 format!("{prefix}.self_attn.o_proj.weight"),
9376 vec![hidden_dim, kda_proj],
9377 hidden_dim * kda_proj,
9378 );
9379 push(
9380 &mut tensors,
9381 format!("{prefix}.mlp.gate_proj.weight"),
9382 vec![dense_intermediate, hidden_dim],
9383 dense_intermediate * hidden_dim,
9384 );
9385 push(
9386 &mut tensors,
9387 format!("{prefix}.mlp.up_proj.weight"),
9388 vec![dense_intermediate, hidden_dim],
9389 dense_intermediate * hidden_dim,
9390 );
9391 push(
9392 &mut tensors,
9393 format!("{prefix}.mlp.down_proj.weight"),
9394 vec![hidden_dim, dense_intermediate],
9395 hidden_dim * dense_intermediate,
9396 );
9397 push(
9398 &mut tensors,
9399 "language_model.model.embed_tokens.weight".to_string(),
9400 vec![vocab_size, hidden_dim],
9401 vocab_size * hidden_dim,
9402 );
9403 push(
9404 &mut tensors,
9405 "language_model.lm_head.weight".to_string(),
9406 vec![vocab_size, hidden_dim],
9407 vocab_size * hidden_dim,
9408 );
9409 push(
9410 &mut tensors,
9411 "language_model.model.norm.weight".to_string(),
9412 vec![hidden_dim],
9413 hidden_dim,
9414 );
9415 push(
9416 &mut tensors,
9417 "language_model.model.output_attn_res_norm.weight".to_string(),
9418 vec![hidden_dim],
9419 hidden_dim,
9420 );
9421 push(
9422 &mut tensors,
9423 "language_model.model.output_attn_res_proj.weight".to_string(),
9424 vec![1, hidden_dim],
9425 hidden_dim,
9426 );
9427
9428 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9441 let dir = std::env::temp_dir().join(format!(
9442 "frink_server_kimi_e2e_test_{}_{}",
9443 std::process::id(),
9444 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9445 ));
9446 std::fs::create_dir_all(&dir).unwrap();
9447 let shard_bytes = write_safetensors_shard(&tensors);
9448 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9449 let map_entries: Vec<String> = tensors
9450 .iter()
9451 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9452 .collect();
9453 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9454 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9455
9456 use base64::Engine;
9460 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9461 .map(|b| {
9462 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9463 format!("{b64} {b}")
9464 })
9465 .collect();
9466 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9467
9468 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9469 .expect("must load the synthetic Kimi checkpoint end to end");
9470 std::fs::remove_dir_all(&dir).ok();
9471 loaded
9472 }
9473
9474 #[test]
9482 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9483 let loaded = build_synthetic_kimi_loaded();
9484 let state = build_app_state(
9485 StartupModels {
9486 loaded: model::LoadedModel::Kimi(loaded),
9487 embedding: None,
9488 },
9489 None,
9490 None,
9491 None,
9492 false,
9493 None,
9494 Arc::new(health::Detection::ready(health::probe_backends())),
9495 );
9496 let active = state.active().expect("a freshly built state has a model");
9497 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9498 assert!(!active.is_synthetic());
9499
9500 let produced = run_generation(
9501 active.generative().unwrap(),
9502 "hi",
9503 &greedy_params(5),
9504 None,
9505 None,
9506 None,
9507 None,
9508 None,
9509 None,
9510 )
9511 .expect("a real Kimi checkpoint must generate without error");
9512 assert!(matches!(
9513 produced.choices[0].finish,
9514 FinishReason::Length | FinishReason::Stop
9515 ));
9516 }
9517
9518 #[test]
9531 fn a_grammar_constrains_the_engine_decode_path() {
9532 let loaded = build_synthetic_kimi_loaded();
9533 let state = build_app_state(
9534 StartupModels {
9535 loaded: model::LoadedModel::Kimi(loaded),
9536 embedding: None,
9537 },
9538 None,
9539 None,
9540 None,
9541 false,
9542 None,
9543 Arc::new(health::Detection::ready(health::probe_backends())),
9544 );
9545 let active = state.active().expect("a freshly built state has a model");
9546
9547 let run = |grammar: Option<&str>| {
9548 let mut params = greedy_params(6);
9549 params.grammar = grammar.map(|src| {
9550 Arc::new(
9551 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9552 .expect("test grammar parses"),
9553 )
9554 });
9555 run_generation(
9556 active.generative().unwrap(),
9557 "hi",
9558 ¶ms,
9559 None,
9560 None,
9561 None,
9562 None,
9563 None,
9564 None,
9565 )
9566 };
9567
9568 let produced = run(None).expect("the unconstrained run must serve");
9569 let unconstrained = produced.choices[0].text.clone();
9570 assert!(
9571 unconstrained.chars().any(|c| c != 'a'),
9572 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9573 constrained run below would prove nothing"
9574 );
9575
9576 let produced =
9577 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9578 let one = produced.choices.into_iter().next().unwrap();
9579 let (finish, constrained) = (one.finish, one.text);
9580 assert!(
9581 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9582 "the engine decode path served text its grammar forbids ({constrained:?}): \
9583 the constraint was dropped between `generate_engine` and the sampler"
9584 );
9585 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9586 }
9587
9588 #[test]
9598 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9599 let loaded = build_synthetic_kimi_loaded();
9600 let state = build_app_state(
9601 StartupModels {
9602 loaded: model::LoadedModel::Kimi(loaded),
9603 embedding: None,
9604 },
9605 None,
9606 None,
9607 None,
9608 false,
9609 None,
9610 Arc::new(health::Detection::ready(health::probe_backends())),
9611 );
9612
9613 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9614 let kv_pool_config = generate::KvPoolConfig {
9615 pool: pool.clone(),
9616 queue_wait: Duration::ZERO,
9617 };
9618 let pc = Mutex::new(PrefixCache::new(4));
9619
9620 run_generation(
9621 state
9622 .active()
9623 .expect("a freshly built state has a model")
9624 .generative()
9625 .unwrap(),
9626 "hi",
9627 &greedy_params(5),
9628 Some(&kv_pool_config),
9629 None,
9630 Some(&pc),
9631 None,
9632 None,
9633 None,
9634 )
9635 .expect("a real Kimi checkpoint must generate without error");
9636
9637 assert_eq!(
9638 pool.lock().unwrap().free_blocks(),
9639 4,
9640 "the KV pool must be completely untouched by a Kimi request"
9641 );
9642 let stats = pc.lock().unwrap().stats();
9643 assert_eq!(
9644 stats.hits + stats.misses,
9645 0,
9646 "the prefix cache must never be consulted for a Kimi request"
9647 );
9648 }
9649}