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