1mod admin;
36mod anthropic;
37mod attribution;
38mod budget;
39mod cache_admin;
40mod cancel;
41mod chat_params;
42mod chat_template;
43mod cli;
44mod completion;
45mod continuation;
46mod conversations;
47mod decode_task;
48mod embeddings;
49mod generate;
50mod grammar_request;
51mod health;
52mod journal;
53mod json_mode;
54mod limits;
55mod loaded;
56mod lora;
57mod mcp;
58mod model;
59mod openai_extra;
60mod output;
61mod policy;
62mod prefill_batch;
63mod reasoning_budget;
64mod reasoning_tokens;
65mod request_tail;
66mod rerank;
67mod response_cache;
68pub(crate) mod responses;
69mod resume;
70mod sample_step;
71mod sampling_knobs;
72mod sampling_loop;
73mod security;
74mod serving;
75mod session;
76mod slots;
77mod sse;
78mod stats;
79mod stop;
80mod stream_events;
81mod tasks;
82mod tool_grammar;
83mod unimplemented_fields;
84mod unsupported_sampling;
85mod utf8_stream;
86
87use std::cell::RefCell;
88use std::convert::Infallible;
89use std::net::SocketAddr;
90use std::path::PathBuf;
91use std::rc::Rc;
92use std::sync::{Arc, Mutex, MutexGuard};
93use std::time::Duration;
94
95use axum::{
96 extract::State,
97 http::StatusCode,
98 response::sse::{Event, Sse},
99 response::{IntoResponse, Response},
100 routing::{get, post},
101 Json, Router,
102};
103use serde::{Deserialize, Serialize};
104
105use cli::apply_cli_overrides;
106pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
107
108use frink_core::cache::KvBlockPool;
109use frink_models::kimi_tokenizer::KimiTokenizer;
110use frink_models::sampling::SamplingParams;
111use frink_models::tokenizer::{SpecialTokens, StopTokens};
112use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
113use generate::{FinishReason, GenerationParams};
114pub(crate) use loaded::{ActiveModel, Loaded};
115use model::ServerTokenizer;
116use rerank::encoder_endpoints;
117use response_cache::ResponseCache;
118use sampling_knobs::SamplingKnobs;
119
120#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
129 Gguf(GgufModel),
130 Kimi(KimiModel),
131 Mla(MlaModel),
132 Gemma4(Gemma4Model),
133 Glm52(Glm52Model),
134}
135
136pub(crate) struct GgufModel {
137 decoder: Arc<Decoder>,
138 tokenizer: Arc<ServerTokenizer>,
139 stop_tokens: StopTokens,
140 bos_id: Option<usize>,
141 is_synthetic: bool,
142 chat_template: chat_template::PromptTemplate,
143}
144
145pub(crate) struct KimiModel {
146 engine: KimiEngine,
147 tokenizer: KimiTokenizer,
148 stop_tokens: StopTokens,
149 chat_template: chat_template::PromptTemplate,
150}
151
152pub(crate) struct MlaModel {
153 engine: MlaEngine,
154 tokenizer: ServerTokenizer,
155 stop_tokens: StopTokens,
156 bos_id: Option<usize>,
157 name: String,
158 chat_template: chat_template::PromptTemplate,
159}
160
161pub(crate) struct Gemma4Model {
162 engine: Gemma4Engine,
163 tokenizer: ServerTokenizer,
164 stop_tokens: StopTokens,
165 bos_id: Option<usize>,
166 name: String,
167 chat_template: chat_template::PromptTemplate,
168}
169
170pub(crate) struct Glm52Model {
171 engine: frink_models::Glm52Engine,
172 tokenizer: ServerTokenizer,
173 stop_tokens: StopTokens,
174 bos_id: Option<usize>,
175 name: String,
176 chat_template: chat_template::PromptTemplate,
177}
178
179impl Model {
180 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
181 match self {
182 Model::Gguf(m) => m.chat_template.clone(),
183 Model::Kimi(m) => m.chat_template.clone(),
184 Model::Mla(m) => m.chat_template.clone(),
185 Model::Gemma4(m) => m.chat_template.clone(),
186 Model::Glm52(m) => m.chat_template.clone(),
187 }
188 }
189
190 fn is_synthetic(&self) -> bool {
195 match self {
196 Model::Gguf(m) => m.is_synthetic,
197 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
198 }
199 }
200
201 fn tokenizer_kind(&self) -> &'static str {
202 match self {
203 Model::Gguf(m) => m.tokenizer.kind(),
204 Model::Kimi(_) => "kimi-tiktoken-bpe",
205 Model::Mla(m) => m.tokenizer.kind(),
206 Model::Gemma4(m) => m.tokenizer.kind(),
207 Model::Glm52(m) => m.tokenizer.kind(),
208 }
209 }
210
211 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
215 match self {
216 Model::Gguf(m) => m.decoder.expert_store_stats(),
217 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
218 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
219 }
220 }
221
222 pub(crate) fn name(&self) -> &str {
223 match self {
224 Model::Gguf(m) => m.decoder.config.name,
225 Model::Kimi(_) => "kimi-k3",
226 Model::Mla(m) => m.name.as_str(),
227 Model::Gemma4(m) => m.name.as_str(),
228 Model::Glm52(m) => m.name.as_str(),
229 }
230 }
231
232 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
259 match self {
260 Model::Gguf(m) => m.tokenizer.encode(text, specials),
261 Model::Kimi(m) => m
262 .tokenizer
263 .encode(text, specials)
264 .into_iter()
265 .map(|id| id as usize)
266 .collect(),
267 Model::Mla(m) => m.tokenizer.encode(text, specials),
268 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
269 Model::Glm52(m) => m.tokenizer.encode(text, specials),
270 }
271 }
272
273 pub(crate) fn bos_id(&self) -> Option<usize> {
282 match self {
283 Model::Gguf(m) => m.bos_id,
284 Model::Kimi(_) => None,
285 Model::Mla(m) => m.bos_id,
286 Model::Gemma4(m) => m.bos_id,
287 Model::Glm52(m) => m.bos_id,
288 }
289 }
290
291 pub(crate) fn decode(&self, ids: &[usize]) -> String {
292 match self {
293 Model::Gguf(m) => m.tokenizer.decode(ids),
294 Model::Kimi(m) => {
295 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
296 m.tokenizer.decode(&ids32)
297 }
298 Model::Mla(m) => m.tokenizer.decode(ids),
299 Model::Gemma4(m) => m.tokenizer.decode(ids),
300 Model::Glm52(m) => m.tokenizer.decode(ids),
301 }
302 }
303
304 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
307 match self {
308 Model::Gguf(m) => {
309 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
310 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
311 }
312 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
313 }
314 }
315
316 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
322 match self {
323 Model::Gguf(m) => Some(&m.decoder),
324 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
325 }
326 }
327
328 pub(crate) fn vocab_size(&self) -> Option<usize> {
329 match self {
330 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
331 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
332 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
333 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
334 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
335 }
336 }
337
338 fn has_real_vocabulary(&self) -> bool {
345 match self {
346 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
347 Model::Kimi(_) => true,
348 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
349 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
350 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
351 }
352 }
353}
354
355impl frink_models::dry::DryVocab for Model {
361 fn n_tokens(&self) -> usize {
362 self.vocab_size().unwrap_or(0)
363 }
364
365 fn detokenize(&self, token: usize) -> String {
366 self.decode(&[token])
367 }
368
369 fn tokenize(&self, text: &str) -> Vec<usize> {
370 self.encode(text, SpecialTokens::AsText)
371 }
372}
373
374pub(crate) struct AppState {
375 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
389 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
408 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
412 pub(crate) tasks: Arc<tasks::TaskRegistry>,
414 pub(crate) cancels: Arc<cancel::CancelRegistry>,
418 pub(crate) stats: stats::Stats,
421 pub(crate) streams: resume::StreamRegistry,
424 pub(crate) model_dir: Option<PathBuf>,
426 response_cache: Mutex<ResponseCache>,
430 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
442 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
455 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
465 sessions: session::SessionStore,
471 requests_total: std::sync::atomic::AtomicU64,
472 request_errors_total: std::sync::atomic::AtomicU64,
473 started_at: std::time::Instant,
474 last_request_ms: std::sync::atomic::AtomicU64,
479 detection: Arc<health::Detection>,
481 mcp: Option<mcp::LoadedMcpConfig>,
483 pub(crate) continuous_batching_enabled: bool,
487 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
491 loading_model: Mutex<Option<String>>,
495 last_load_error: Mutex<Option<(String, String)>>,
499 pub(crate) serving: Mutex<crate::stats::ServingStats>,
504 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
509 pub(crate) footprint:
515 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
516 pub(crate) started_unix: u64,
526}
527
528pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
534
535impl AppState {
536 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
540 self.active
541 .read()
542 .unwrap_or_else(|p| p.into_inner())
543 .clone()
544 }
545
546 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
552 self.active().ok_or_else(|| {
553 (
554 StatusCode::SERVICE_UNAVAILABLE,
555 Json(serde_json::json!({"error": {
556 "message": "no model is loaded; POST /admin/models/load with an id from \
557 GET /admin/models",
558 "type": "model_not_loaded"
559 }})),
560 )
561 })
562 }
563
564 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
572 Ok(Arc::clone(self.require_active()?.generative()?))
573 }
574
575 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
583 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
584 std::mem::replace(&mut *guard, next)
585 }
586
587 fn mark_request_finished(&self) {
591 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
592 self.last_request_ms
593 .store(ms, std::sync::atomic::Ordering::Relaxed);
594 }
595
596 pub(crate) fn uptime(&self) -> Duration {
597 self.started_at.elapsed()
598 }
599
600 pub(crate) fn requests_total(&self) -> u64 {
601 self.requests_total
602 .load(std::sync::atomic::Ordering::Relaxed)
603 }
604
605 pub(crate) fn errors_total(&self) -> u64 {
606 self.request_errors_total
607 .load(std::sync::atomic::Ordering::Relaxed)
608 }
609
610 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
611 lock_cache(&self.response_cache).stats()
612 }
613
614 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
617 let last = self
618 .last_request_ms
619 .load(std::sync::atomic::Ordering::Relaxed);
620 (last > 0)
621 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
622 .map(|age| age.max(0.0))
623 }
624
625 pub(crate) fn loading_model_id(&self) -> Option<String> {
626 self.loading_model
627 .lock()
628 .unwrap_or_else(|p| p.into_inner())
629 .clone()
630 }
631
632 pub(crate) fn set_loading_model(&self, id: Option<String>) {
633 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
634 }
635
636 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
637 self.last_load_error
638 .lock()
639 .unwrap_or_else(|p| p.into_inner())
640 .clone()
641 }
642
643 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
644 *self
645 .last_load_error
646 .lock()
647 .unwrap_or_else(|p| p.into_inner()) = error;
648 }
649
650 pub(crate) fn active_model_name(&self) -> Option<String> {
659 self.active().map(|a| a.name().to_string())
660 }
661
662 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
672 self.embedding
673 .clone()
674 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
675 }
676
677 pub(crate) fn embedding_model_name(&self) -> Option<String> {
681 match self.embedding_model() {
682 Some(e) => Some(e.name().to_string()),
683 None => self.active_model_name(),
684 }
685 }
686
687 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
688 self.stats.record(stats::entry(record));
689 }
690}
691
692fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
697 cache
698 .lock()
699 .unwrap_or_else(|poisoned| poisoned.into_inner())
700}
701
702#[derive(Debug, Clone, Deserialize)]
703#[serde(untagged)]
704pub(crate) enum MessageContent {
705 Text(String),
706 Parts(Vec<ContentPart>),
707}
708
709#[derive(Debug, Clone, Deserialize)]
710struct ContentPart {
711 #[serde(rename = "type")]
712 kind: String,
713 #[serde(default)]
714 text: Option<String>,
715 #[serde(default)]
716 image_url: Option<serde_json::Value>,
717}
718
719impl MessageContent {
720 fn as_text(&self) -> String {
721 match self {
722 Self::Text(s) => s.clone(),
723 Self::Parts(parts) => parts
724 .iter()
725 .filter_map(|p| p.text.as_deref())
726 .collect::<Vec<_>>()
727 .join(""),
728 }
729 }
730
731 fn has_image(&self) -> bool {
732 match self {
733 Self::Text(_) => false,
734 Self::Parts(parts) => parts
735 .iter()
736 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
737 }
738 }
739}
740
741#[derive(Debug, Clone, Deserialize)]
742pub(crate) struct ChatMessage {
743 pub(crate) role: String,
744 #[serde(default)]
748 pub(crate) content: Option<MessageContent>,
749 #[serde(default)]
753 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
754 #[serde(default)]
759 #[allow(dead_code)]
760 pub(crate) tool_call_id: Option<String>,
761 #[serde(default, alias = "reasoning")]
777 pub(crate) reasoning_content: Option<String>,
778}
779
780impl ChatMessage {
781 fn rendered_content(&self) -> String {
788 let mut out = self
789 .content
790 .as_ref()
791 .map(MessageContent::as_text)
792 .unwrap_or_default();
793 if let Some(calls) = &self.tool_calls {
794 for call in calls {
795 out.push_str(&format!(
796 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
797 call.function.name, call.function.arguments
798 ));
799 }
800 }
801 out
802 }
803}
804
805#[derive(Debug, Clone, Deserialize)]
806pub(crate) struct ToolCallIn {
807 #[serde(default)]
808 #[allow(dead_code)]
809 id: String,
810 #[serde(rename = "type", default)]
811 #[allow(dead_code)]
812 kind: String,
813 function: ToolCallFunctionIn,
814}
815
816#[derive(Debug, Clone, Deserialize)]
817struct ToolCallFunctionIn {
818 name: String,
819 arguments: String,
824}
825
826#[derive(Debug, Clone, Deserialize)]
829struct ToolDef {
830 #[serde(rename = "type", default)]
831 #[allow(dead_code)]
832 kind: String,
833 function: ToolFunctionDef,
834}
835
836#[derive(Debug, Clone, Deserialize)]
837struct ToolFunctionDef {
838 name: String,
839 #[serde(default)]
840 description: Option<String>,
841 #[serde(default)]
842 parameters: Option<serde_json::Value>,
843}
844
845#[derive(Debug, Clone, Deserialize)]
856#[serde(untagged)]
857enum ToolChoice {
858 Mode(String),
859 Specific(serde_json::Value),
860}
861
862#[derive(Deserialize)]
865#[serde(untagged)]
866enum StopParam {
867 One(String),
868 Many(Vec<String>),
869}
870
871#[derive(Deserialize)]
872struct ChatCompletionRequest {
873 model: String,
874 messages: Vec<ChatMessage>,
875 #[serde(default = "default_max_tokens")]
876 max_tokens: usize,
877 #[serde(default)]
878 temperature: Option<f32>,
879 #[serde(default)]
880 top_p: Option<f32>,
881 #[serde(default)]
886 min_p: Option<f32>,
887 #[serde(default)]
888 top_k: Option<usize>,
889 #[serde(default)]
890 repetition_penalty: Option<f32>,
891 #[serde(flatten)]
895 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
896 #[serde(flatten)]
900 unimplemented: crate::unimplemented_fields::UnimplementedFields,
901 #[serde(default)]
902 seed: Option<u64>,
903 #[serde(default)]
904 stop: Option<StopParam>,
905 #[serde(default)]
906 stream: Option<bool>,
907 #[serde(default)]
920 stream_resumable: Option<bool>,
921 #[serde(default)]
932 ignore_eos: Option<bool>,
933 #[serde(default)]
934 tools: Vec<ToolDef>,
935 #[serde(default)]
936 tool_choice: Option<ToolChoice>,
937 #[serde(default)]
945 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
946 #[serde(default)]
954 reasoning_effort: Option<String>,
955 #[serde(default)]
959 thinking: Option<ThinkingSwitch>,
960 #[serde(default)]
965 session_id: Option<String>,
966 #[serde(default, deserialize_with = "continuation::deserialize")]
974 continue_final_message: continuation::ContinueFinalMessage,
975 #[serde(default, alias = "thinking_budget_tokens")]
983 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
984 #[serde(default)]
986 logprobs: Option<bool>,
987 #[serde(default)]
988 top_logprobs: Option<u32>,
989 #[serde(default)]
990 presence_penalty: Option<f32>,
991 #[serde(default)]
992 frequency_penalty: Option<f32>,
993 #[serde(default)]
994 response_format: Option<serde_json::Value>,
995 #[serde(default)]
1002 logit_bias: Option<serde_json::Value>,
1003 #[serde(default)]
1007 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1008 #[serde(default)]
1018 samplers: Option<serde_json::Value>,
1019 #[serde(default)]
1029 grammar: Option<String>,
1030}
1031
1032const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1044
1045#[derive(Debug, Clone, Deserialize)]
1047pub(crate) struct ThinkingSwitch {
1048 #[serde(rename = "type")]
1049 pub(crate) kind: String,
1050}
1051
1052const THINKING_KWARG_KEYS: [&str; 4] = [
1056 "enable_thinking",
1057 "thinking",
1058 "thinking_mode",
1059 "reasoning_effort",
1060];
1061
1062const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1066
1067fn default_max_tokens() -> usize {
1068 DEFAULT_CHAT_MAX_TOKENS
1069}
1070
1071impl ChatCompletionRequest {
1072 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1081 let mut knobs = SamplingKnobs {
1082 temperature: self.temperature,
1083 top_p: self.top_p,
1084 min_p: self.min_p,
1085 top_k: self.top_k,
1086 repetition_penalty: self.repetition_penalty,
1087 presence_penalty: self.presence_penalty,
1088 frequency_penalty: self.frequency_penalty,
1089 penalty_last_n: None,
1093 sampler_order: unsupported_sampling::parse_sampler_order(
1094 self.samplers.as_ref(),
1095 "/v1/chat/completions",
1096 )?,
1097 ..SamplingKnobs::default()
1098 };
1099 self.extra_samplers.apply(&mut knobs);
1100 Ok(knobs)
1101 }
1102
1103 fn sampling_params(
1104 &self,
1105 model: crate::sampling_knobs::SamplerModel<'_>,
1106 ) -> Result<SamplingParams, ApiError> {
1107 self.sampling_knobs()?.resolve(model).map_err(|e| {
1108 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1109 })
1110 }
1111
1112 fn stop_sequences(&self) -> Vec<String> {
1113 self.stop
1114 .as_ref()
1115 .map(|s| match s {
1116 StopParam::One(v) => vec![v.clone()],
1117 StopParam::Many(v) => v.clone(),
1118 })
1119 .unwrap_or_default()
1120 }
1121
1122 fn several_choices(&self) -> bool {
1131 self.unimplemented.n.is_some_and(|n| n > 1)
1132 }
1133
1134 fn tools_active(&self) -> bool {
1135 !self.tools.is_empty()
1136 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1137 }
1138
1139 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1151 match &self.tool_choice {
1152 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1153 Some(ToolChoice::Specific(value)) => {
1154 let name = value
1158 .get("function")
1159 .and_then(|f| f.get("name"))
1160 .or_else(|| value.get("name"))
1161 .and_then(|n| n.as_str());
1162 match name {
1163 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1164 None => Err(invalid_request(
1165 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1166 function.name",
1167 "tool_choice",
1168 )),
1169 }
1170 }
1171 _ => Ok(None),
1172 }
1173 }
1174
1175 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1177 self.tools
1178 .iter()
1179 .map(|t| tool_grammar::ToolSpec {
1180 name: &t.function.name,
1181 parameters: t.function.parameters.as_ref(),
1182 })
1183 .collect()
1184 }
1185
1186 fn thinking_direction(&self) -> Option<bool> {
1229 if let Some(switch) = &self.thinking {
1230 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1231 "disabled" => Some(false),
1232 "enabled" => Some(true),
1233 _ => None,
1236 };
1237 }
1238 let effort = self.reasoning_effort.as_ref()?;
1239 DISABLE_EFFORTS
1240 .contains(&effort.trim().to_ascii_lowercase().as_str())
1241 .then_some(false)
1242 }
1243
1244 fn resolve_template_kwargs(
1245 &self,
1246 template: &chat_template::PromptTemplate,
1247 ) -> serde_json::Map<String, serde_json::Value> {
1248 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1249 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1253
1254 if !caller_steered {
1255 match self.thinking_direction() {
1256 Some(false) => {
1257 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1258 kwargs.insert(k, v);
1259 }
1260 return kwargs;
1263 }
1264 Some(true) => {
1265 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1266 kwargs.insert(k, v);
1267 }
1268 }
1269 None => {}
1270 }
1271 if let Some(effort) = &self.reasoning_effort {
1272 kwargs
1273 .entry("reasoning_effort".to_string())
1274 .or_insert_with(|| serde_json::json!(effort));
1275 }
1276 }
1277
1278 let offered: Vec<serde_json::Value> = if self.tools_active() {
1279 self.tools.iter().map(chat_template::tool_json).collect()
1280 } else {
1281 Vec::new()
1282 };
1283 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1284 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1285 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1286 kwargs.entry(k).or_insert(v);
1287 }
1288 }
1289 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1290 crate::policy::effort::EffortMapping::Mapped(to) => {
1291 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1292 }
1293 crate::policy::effort::EffortMapping::Dropped => {
1294 tracing::debug!(
1295 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1296 enough, so its own default applies"
1297 );
1298 }
1299 crate::policy::effort::EffortMapping::Unchanged => {}
1300 }
1301 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1302 kwargs
1303 }
1304
1305 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1308 if self.max_tokens == 0 {
1315 return Err(invalid_request(
1316 "max_tokens must be at least 1",
1317 "max_tokens",
1318 ));
1319 }
1320 if let Some(switch) = &self.thinking {
1325 let kind = switch.kind.trim().to_ascii_lowercase();
1326 if kind != "enabled" && kind != "disabled" {
1327 return Err(invalid_request(
1328 "thinking.type must be \"enabled\" or \"disabled\"",
1329 "thinking.type",
1330 ));
1331 }
1332 }
1333 for msg in &self.messages {
1334 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1335 return Err(unsupported_feature(
1336 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1337 ));
1338 }
1339 }
1340 if self.logprobs == Some(true) || self.top_logprobs.is_some() {
1341 return Err(unsupported_feature(
1342 "logprobs / top_logprobs are not implemented yet (see docs/API.md)",
1343 ));
1344 }
1345 self.unimplemented.refuse("/v1/chat/completions")?;
1349 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1350 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1354 let stated_grammar =
1369 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1370 if let Some(forced) = self.forced_tool_choice()? {
1378 if self.tools.is_empty() {
1379 return Err(invalid_request(
1380 "tool_choice forces a tool call, but no tools were offered",
1381 "tool_choice",
1382 ));
1383 }
1384 if let tool_grammar::Forced::Named(name) = forced {
1385 if !self.tools.iter().any(|t| t.function.name == name) {
1386 return Err(invalid_request(
1387 &format!(
1388 "tool_choice names {name:?}, which is not one of the tools offered"
1389 ),
1390 "tool_choice",
1391 ));
1392 }
1393 }
1394 if stated_grammar.is_some() {
1404 return Err(invalid_request(
1405 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1406 are two different constraints on the same generation; send one",
1407 "tool_choice",
1408 ));
1409 }
1410 if self.json_object_mode() {
1411 return Err(invalid_request(
1412 "a forced tool_choice cannot be combined with response_format json_object: \
1413 the tool-call markers are not JSON",
1414 "tool_choice",
1415 ));
1416 }
1417 }
1418 Ok(())
1419 }
1420
1421 fn effective_stop_sequences(&self) -> Vec<String> {
1428 let mut stop = self.stop_sequences();
1429 if self.tools_active() {
1430 stop.push("</tool_call>".to_string());
1431 }
1432 stop
1433 }
1434
1435 fn json_object_mode(&self) -> bool {
1436 self.response_format
1437 .as_ref()
1438 .and_then(|v| v.get("type"))
1439 .and_then(|v| v.as_str())
1440 == Some("json_object")
1441 }
1442}
1443
1444#[derive(Serialize)]
1445struct ChatCompletionChoice {
1446 index: usize,
1447 message: ChatCompletionResponseMessage,
1448 finish_reason: &'static str,
1449}
1450
1451#[derive(Serialize)]
1452struct ChatCompletionResponseMessage {
1453 role: &'static str,
1454 #[serde(skip_serializing_if = "Option::is_none")]
1455 content: Option<String>,
1456 #[serde(skip_serializing_if = "Option::is_none")]
1460 reasoning_content: Option<String>,
1461 #[serde(skip_serializing_if = "Option::is_none")]
1462 tool_calls: Option<Vec<ToolCallOut>>,
1463}
1464
1465#[derive(Serialize, Clone)]
1466struct ToolCallOut {
1467 id: String,
1468 #[serde(rename = "type")]
1469 kind: &'static str,
1470 function: ToolCallFunctionOut,
1471}
1472
1473#[derive(Serialize, Clone)]
1481struct ToolCallDelta {
1482 index: usize,
1483 #[serde(skip_serializing_if = "Option::is_none")]
1484 id: Option<String>,
1485 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1486 kind: Option<&'static str>,
1487 function: ToolCallFunctionDelta,
1488}
1489
1490#[derive(Serialize, Clone, Default)]
1491struct ToolCallFunctionDelta {
1492 #[serde(skip_serializing_if = "Option::is_none")]
1493 name: Option<String>,
1494 #[serde(skip_serializing_if = "Option::is_none")]
1497 arguments: Option<String>,
1498}
1499
1500impl ToolCallDelta {
1501 fn whole(index: usize, name: String, arguments: String) -> Self {
1503 ToolCallDelta {
1504 index,
1505 id: Some(format!("call_{index}")),
1506 kind: Some("function"),
1507 function: ToolCallFunctionDelta {
1508 name: Some(name),
1509 arguments: Some(arguments),
1510 },
1511 }
1512 }
1513
1514 fn opening(index: usize, name: String) -> Self {
1516 ToolCallDelta {
1517 index,
1518 id: Some(format!("call_{index}")),
1519 kind: Some("function"),
1520 function: ToolCallFunctionDelta {
1521 name: Some(name),
1522 arguments: Some(String::new()),
1523 },
1524 }
1525 }
1526
1527 fn arguments(index: usize, fragment: String) -> Self {
1529 ToolCallDelta {
1530 index,
1531 id: None,
1532 kind: None,
1533 function: ToolCallFunctionDelta {
1534 name: None,
1535 arguments: Some(fragment),
1536 },
1537 }
1538 }
1539}
1540
1541#[derive(Serialize, Clone)]
1542struct ToolCallFunctionOut {
1543 name: String,
1544 arguments: String,
1548}
1549
1550#[derive(Serialize)]
1551struct ChatCompletionResponse {
1552 id: String,
1553 request_id: String,
1559 object: &'static str,
1560 model: String,
1561 choices: Vec<ChatCompletionChoice>,
1562 usage: generate::Usage,
1567 frink_cache: &'static str,
1578}
1579
1580#[derive(Serialize)]
1581struct ChatCompletionChunkDelta {
1582 #[serde(skip_serializing_if = "Option::is_none")]
1583 role: Option<&'static str>,
1584 #[serde(skip_serializing_if = "Option::is_none")]
1585 content: Option<String>,
1586 #[serde(skip_serializing_if = "Option::is_none")]
1588 reasoning_content: Option<String>,
1589 #[serde(skip_serializing_if = "Option::is_none")]
1590 tool_calls: Option<Vec<ToolCallDelta>>,
1591}
1592
1593#[derive(Serialize)]
1594struct ChatCompletionChunkChoice {
1595 index: usize,
1596 delta: ChatCompletionChunkDelta,
1597 finish_reason: Option<&'static str>,
1598}
1599
1600#[derive(Serialize)]
1601struct ChatCompletionChunk {
1602 id: String,
1603 #[serde(skip_serializing_if = "Option::is_none")]
1610 request_id: Option<String>,
1611 object: &'static str,
1612 model: String,
1613 choices: Vec<ChatCompletionChunkChoice>,
1614 #[serde(skip_serializing_if = "Option::is_none")]
1617 usage: Option<generate::Usage>,
1618}
1619
1620async fn health(State(state): State<Arc<AppState>>) -> Response {
1626 let snapshot = state.detection.snapshot();
1627 let mut capabilities = snapshot.capabilities;
1628 let active = state.active();
1629
1630 capabilities.push(match active.as_deref() {
1633 None => frink_api::Capability::unavailable(
1638 frink_api::health::capability::REAL_WEIGHTS,
1639 frink_api::health::reason::MODEL_NOT_LOADED,
1640 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1641 ),
1642 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1643 frink_api::health::capability::REAL_WEIGHTS,
1644 frink_api::health::reason::MODEL_NOT_LOADED,
1645 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1646 checkpoint. Output from this model is noise.",
1647 ),
1648 Some(active) if active.encoder().is_some() => {
1658 let endpoints = active
1659 .encoder()
1660 .map(|e| encoder_endpoints(e))
1661 .unwrap_or_default();
1662 let served_by = match endpoints.is_empty() {
1663 true => "no endpoint in this build serves it".to_string(),
1664 false => format!("served by {}", endpoints.join(" and ")),
1665 };
1666 frink_api::Capability::available(
1667 frink_api::health::capability::REAL_WEIGHTS,
1668 format!(
1669 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1670 {served_by}; generation endpoints refuse it.",
1671 active.name(),
1672 ),
1673 )
1674 }
1675 Some(active) => frink_api::Capability::available(
1676 frink_api::health::capability::REAL_WEIGHTS,
1677 format!("Serving the real checkpoint '{}'.", active.name()),
1678 ),
1679 });
1680 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1681 frink_api::Capability::available(
1682 frink_api::health::capability::CONTINUOUS_BATCHING,
1683 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1684 "On by default on Metal. Concurrent requests share one batched decode worker."
1685 } else {
1686 "Concurrent requests share one batched decode step."
1687 },
1688 )
1689 } else if state.metal_private_decode_gate.is_some() {
1690 frink_api::Capability::unavailable(
1691 frink_api::health::capability::CONTINUOUS_BATCHING,
1692 frink_api::health::reason::DISABLED,
1693 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1694 )
1695 } else {
1696 frink_api::Capability::unavailable(
1697 frink_api::health::capability::CONTINUOUS_BATCHING,
1698 frink_api::health::reason::DISABLED,
1699 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1700 )
1701 });
1702
1703 let last_request_ms = state
1704 .last_request_ms
1705 .load(std::sync::atomic::Ordering::Relaxed);
1706 let uptime = state.started_at.elapsed();
1707 let health_state = if active.is_none() {
1715 frink_api::HealthState::Unavailable
1716 } else {
1717 snapshot.state
1718 };
1719 let body = frink_api::HealthResponse {
1720 state: health_state,
1721 reason: match health_state {
1722 frink_api::HealthState::Ready => None,
1723 frink_api::HealthState::Unavailable => {
1724 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1725 }
1726 frink_api::HealthState::Detecting => {
1727 Some(frink_api::health::reason::DETECTING.to_string())
1728 }
1729 },
1730 detail: match health_state {
1731 frink_api::HealthState::Ready => None,
1732 frink_api::HealthState::Unavailable => Some(
1733 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1734 .to_string(),
1735 ),
1736 frink_api::HealthState::Detecting => {
1737 Some("Probing available compute backends.".to_string())
1738 }
1739 },
1740 model: active
1741 .as_deref()
1742 .map(|active| frink_api::health::ModelSummary {
1743 id: active.name().to_string(),
1744 tokenizer: active.tokenizer_kind().to_string(),
1745 synthetic_weights: active.is_synthetic(),
1746 }),
1747 capabilities,
1748 version: env!("CARGO_PKG_VERSION").to_string(),
1749 pid: std::process::id(),
1750 uptime_seconds: uptime.as_secs_f64(),
1751 server_time_unix_ms: std::time::SystemTime::now()
1752 .duration_since(std::time::UNIX_EPOCH)
1753 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1754 .unwrap_or(0),
1755 last_request_age_seconds: (last_request_ms > 0)
1756 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1757 .map(|age| age.max(0.0)),
1758 };
1759
1760 let status =
1761 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1762 (status, Json(body)).into_response()
1763}
1764
1765async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1766 let Some(active) = state.active() else {
1770 return Json(serde_json::json!({ "object": "list", "data": [] }));
1771 };
1772 let mut model_entry = serde_json::json!({
1773 "id": active.name(),
1774 "object": "model",
1775 "frink_synthetic_weights": active.is_synthetic(),
1776 "frink_tokenizer": active.tokenizer_kind(),
1777 });
1778 if let Some(encoder) = active.encoder() {
1784 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1785 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1786 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1787 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1788 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1789 }
1790 if let Some(model) = active.generative_opt() {
1798 let parser_configured = active.reasoning_format().is_some();
1799 let gears = model.chat_template().think_gears(parser_configured);
1800 if !gears.is_empty() {
1801 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1802 if let Some(default) = &gears.default {
1803 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1804 }
1805 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1808 }
1809 }
1810 if let Some(mcp) = &state.mcp {
1811 model_entry["frink_mcp"] = mcp.models_metadata();
1812 }
1813 Json(serde_json::json!({
1814 "object": "list",
1815 "data": [model_entry]
1816 }))
1817}
1818
1819async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1833 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1834 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1835 let active = state.active();
1836 Json(serde_json::json!({
1837 "model": active.as_ref().map(|a| a.name()),
1838 "state": state
1839 .maintenance
1840 .lock()
1841 .unwrap_or_else(|p| p.into_inner())
1842 .state()
1843 .as_str(),
1844 "uptime_s": state.uptime().as_secs(),
1845 "throughput": {
1846 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1847 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1848 },
1849 "requests": {
1850 "active": state.cancels.live_count(),
1851 "completed": state.stats.recorded_total(),
1852 "p95_ms": state.stats.p95_duration_ms(),
1853 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1854 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1855 "completion_tokens_total": state.stats.tokens_generated_total(),
1856 },
1857 "pools": cache_admin::pool_gauges(&state),
1863 "memory": cache_admin::footprint_json(&state),
1866 }))
1867}
1868
1869#[derive(Deserialize)]
1870struct RequestsQuery {
1871 #[serde(default)]
1872 since: u64,
1873 #[serde(default = "default_requests_limit")]
1874 limit: usize,
1875}
1876
1877fn default_requests_limit() -> usize {
1878 stats::MAX_PAGE
1879}
1880
1881async fn recent_requests(
1889 State(state): State<Arc<AppState>>,
1890 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1891) -> Json<serde_json::Value> {
1892 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1893 Json(serde_json::json!({
1894 "requests": rows,
1895 "next_cursor": cursor,
1896 "missed": missed,
1897 "total": state.stats.recorded_total(),
1898 }))
1899}
1900
1901#[derive(Serialize)]
1902struct CombinedCacheStats {
1903 response_cache: response_cache::CacheStats,
1904 prefix_cache: Option<frink_models::PrefixCacheStats>,
1906}
1907
1908async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1909 Json(CombinedCacheStats {
1910 response_cache: lock_cache(&state.response_cache).stats(),
1911 prefix_cache: state
1912 .prefix_cache
1913 .as_ref()
1914 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1915 })
1916}
1917
1918async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1923 use std::sync::atomic::Ordering;
1924
1925 let cache_stats = lock_cache(&state.response_cache).stats();
1926 let active = state.active();
1927 let requests_total = state.requests_total.load(Ordering::Relaxed);
1928 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1929 let uptime = state.started_at.elapsed().as_secs_f64();
1930
1931 let body = format!(
1932 "# HELP frink_requests_total Total chat completion requests received.\n\
1933 # TYPE frink_requests_total counter\n\
1934 frink_requests_total {requests_total}\n\
1935 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1936 # TYPE frink_request_errors_total counter\n\
1937 frink_request_errors_total {errors_total}\n\
1938 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1939 # TYPE frink_cache_hits_total counter\n\
1940 frink_cache_hits_total {}\n\
1941 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1942 # TYPE frink_cache_misses_total counter\n\
1943 frink_cache_misses_total {}\n\
1944 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1945 # TYPE frink_cache_entries gauge\n\
1946 frink_cache_entries {}\n\
1947 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1948 # TYPE frink_synthetic_weights gauge\n\
1949 frink_synthetic_weights {}\n\
1950 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1951 # TYPE frink_uptime_seconds gauge\n\
1952 frink_uptime_seconds {uptime}\n",
1953 cache_stats.hits,
1954 cache_stats.misses,
1955 cache_stats.entries,
1956 active
1960 .as_ref()
1961 .map(|a| a.is_synthetic() as u8)
1962 .unwrap_or(0),
1963 );
1964
1965 let body = match active
1969 .as_ref()
1970 .and_then(|a| a.expert_store_stats())
1971 {
1972 Some(es) => format!(
1973 "{body}\
1974 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
1975 # TYPE frink_expert_cache_hits_total counter\n\
1976 frink_expert_cache_hits_total {}\n\
1977 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
1978 # TYPE frink_expert_cache_misses_total counter\n\
1979 frink_expert_cache_misses_total {}\n\
1980 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
1981 # TYPE frink_expert_cache_evictions_total counter\n\
1982 frink_expert_cache_evictions_total {}\n\
1983 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
1984 # TYPE frink_expert_cache_pass_throughs_total counter\n\
1985 frink_expert_cache_pass_throughs_total {}\n\
1986 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
1987 # TYPE frink_expert_cache_bytes_read_total counter\n\
1988 frink_expert_cache_bytes_read_total {}\n\
1989 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
1990 # TYPE frink_expert_cache_resident_bytes gauge\n\
1991 frink_expert_cache_resident_bytes {}\n",
1992 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
1993 ),
1994 None => body,
1995 };
1996
1997 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2002 Some(batcher) => {
2003 let sched = batcher.stats();
2004 format!(
2005 "{body}\
2006 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2007 # TYPE frink_prefill_chunks_total counter\n\
2008 frink_prefill_chunks_total {}\n\
2009 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2010 # TYPE frink_prefill_tokens_total counter\n\
2011 frink_prefill_tokens_total {}\n\
2012 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2013 # TYPE frink_decode_steps_total counter\n\
2014 frink_decode_steps_total {}\n\
2015 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2016 # TYPE frink_scheduler_queue_depth gauge\n\
2017 frink_scheduler_queue_depth {}\n\
2018 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2019 # TYPE frink_scheduler_queue_rejected_total counter\n\
2020 frink_scheduler_queue_rejected_total {}\n\
2021 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2022 # TYPE frink_kv_blocks_total gauge\n\
2023 frink_kv_blocks_total {}\n\
2024 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2025 # TYPE frink_kv_blocks_free gauge\n\
2026 frink_kv_blocks_free {}\n\
2027 # HELP frink_kv_block_size Token positions per KV block.\n\
2028 # TYPE frink_kv_block_size gauge\n\
2029 frink_kv_block_size {}\n\
2030 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2031 # TYPE frink_kv_rejected_too_large_total counter\n\
2032 frink_kv_rejected_too_large_total {}\n\
2033 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2034 # TYPE frink_kv_rejected_context_length_total counter\n\
2035 frink_kv_rejected_context_length_total {}\n\
2036 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2037 # TYPE frink_scheduler_aborted_total counter\n\
2038 frink_scheduler_aborted_total {}\n\
2039 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2040 # TYPE frink_scheduler_max_seqs gauge\n\
2041 frink_scheduler_max_seqs {}\n\
2042 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2043 # TYPE frink_scheduler_prefill_chunk gauge\n\
2044 frink_scheduler_prefill_chunk {}\n",
2045 sched.prefill_chunks,
2046 sched.prefill_tokens,
2047 sched.decode_steps,
2048 sched.queue_depth,
2049 sched.queue_rejected,
2050 sched.kv_blocks_total,
2051 sched.kv_blocks_free,
2052 sched.kv_block_size,
2053 sched.kv_rejected_too_large,
2054 sched.kv_rejected_context_length,
2055 sched.aborted,
2056 sched.max_seqs,
2057 sched.prefill_chunk,
2058 )
2059 }
2060 None => body,
2061 };
2062
2063 (
2064 [(
2065 axum::http::header::CONTENT_TYPE,
2066 "text/plain; version=0.0.4",
2067 )],
2068 body,
2069 )
2070 .into_response()
2071}
2072
2073pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2074
2075pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2081 (
2082 StatusCode::BAD_REQUEST,
2083 Json(serde_json::json!({"error": {
2084 "message": message,
2085 "type": "invalid_request_error",
2086 "param": param,
2087 "code": null,
2088 }})),
2089 )
2090}
2091
2092pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2093 (
2094 StatusCode::NOT_IMPLEMENTED,
2095 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2096 )
2097}
2098
2099pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2100 let status = match e {
2101 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2102 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2106 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2111 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2117 StatusCode::SERVICE_UNAVAILABLE
2118 }
2119 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2123 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2127 };
2128 tracing::warn!("decode error: {e}");
2129 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2130 if let generate::DecodeError::KvBudgetExceeded {
2136 binding,
2137 estimated_bytes,
2138 limit_bytes,
2139 positions,
2140 positions_limit,
2141 ..
2142 } = &e
2143 {
2144 body["error"]["type"] = serde_json::json!("invalid_request_error");
2145 body["error"]["code"] = serde_json::json!(binding);
2146 body["error"]["binding"] = serde_json::json!(binding);
2147 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2148 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2149 body["error"]["positions"] = serde_json::json!(positions);
2150 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2151 }
2152 if let Some(secs) = e.retry_after_secs() {
2156 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2157 }
2158 (status, Json(body))
2159}
2160
2161pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2162 tracing::error!("generation task panicked: {e}");
2163 (
2164 StatusCode::INTERNAL_SERVER_ERROR,
2165 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2166 )
2167}
2168
2169#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2179 model: &Model,
2180 prompt: &str,
2181 params: &GenerationParams,
2182 kv_pool: Option<&generate::KvPoolConfig>,
2183 paged_kv: Option<&generate::PagedKvConfig>,
2184 prefix_cache: Option<&Mutex<PrefixCache>>,
2185 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2186 ceiling: Option<&budget::ContextCeiling>,
2187 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2188 mut emit: impl FnMut(&str),
2189 ) -> Result<(Vec<(FinishReason, String)>, generate::Usage), generate::DecodeError> {
2195 let synthetic = model.is_synthetic();
2196 let _lora_lease = lora::lease(model, params.lora.as_deref());
2201 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2202 let params = &{
2208 let mut resolved = params.clone();
2209 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2210 model.encode(text, SpecialTokens::Parse)
2211 });
2212 resolved.reasoning_budget = resolved
2218 .reasoning_budget
2219 .armed(resolved.reasoning, prompt, |text| {
2220 model.encode(text, SpecialTokens::Parse)
2221 })
2222 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2223 resolved
2224 };
2225 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2226 let _metal_private_guard =
2227 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2228 let (finishes, usage) = match model {
2229 Model::Gguf(m) => {
2230 if let Some(batcher) = continuous_batcher {
2231 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2232 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2233 let (finish, _generated_ids, text, usage) = if synthetic {
2234 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2235 } else {
2236 batcher.generate_streaming(
2237 tokens,
2238 params.clone(),
2239 m.stop_tokens.clone(),
2240 Some(|chunk: &str| {
2241 if !chunk.is_empty() {
2242 chunks[0].push(chunk.to_string());
2243 emit(chunk);
2244 }
2245 }),
2246 )?
2247 };
2248 if !text.is_empty() && chunks[0].is_empty() {
2249 chunks[0].push(text);
2250 }
2251 (vec![finish], usage)
2255 } else {
2256 generate::generate(
2257 &m.decoder,
2258 m.tokenizer.as_ref(),
2259 &m.stop_tokens,
2260 m.bos_id,
2261 prompt,
2262 params,
2263 kv_pool,
2264 paged_kv,
2265 prefix_cache,
2266 ceiling,
2267 |choice, chunk| {
2268 chunks[choice].push(chunk.to_string());
2269 if !synthetic && choice == 0 {
2273 emit(chunk);
2274 }
2275 },
2276 )?
2277 }
2278 }
2279 Model::Kimi(m) => generate::generate_engine(
2280 &m.engine,
2281 &m.tokenizer,
2282 &m.stop_tokens,
2283 None,
2284 prompt,
2285 params,
2286 |chunk| {
2287 chunks[0].push(chunk.to_string());
2288 if !synthetic {
2289 emit(chunk);
2290 }
2291 },
2292 )?,
2293 Model::Mla(m) => generate::generate_engine(
2294 &m.engine,
2295 &m.tokenizer,
2296 &m.stop_tokens,
2297 m.bos_id,
2298 prompt,
2299 params,
2300 |chunk| {
2301 chunks[0].push(chunk.to_string());
2302 if !synthetic {
2303 emit(chunk);
2304 }
2305 },
2306 )?,
2307 Model::Gemma4(m) => generate::generate_engine(
2308 &m.engine,
2309 &m.tokenizer,
2310 &m.stop_tokens,
2311 m.bos_id,
2312 prompt,
2313 params,
2314 |chunk| {
2315 chunks[0].push(chunk.to_string());
2316 if !synthetic {
2317 emit(chunk);
2318 }
2319 },
2320 )?,
2321 Model::Glm52(m) => generate::generate_engine(
2322 &m.engine,
2323 &m.tokenizer,
2324 &m.stop_tokens,
2325 m.bos_id,
2326 prompt,
2327 params,
2328 |chunk| {
2329 chunks[0].push(chunk.to_string());
2330 if !synthetic {
2331 emit(chunk);
2332 }
2333 },
2334 )?,
2335 };
2336
2337 let mut full = chunks[0].concat();
2338 if synthetic {
2339 full = format!(
2340 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2341 to serve a real model. Decoded ids -> {full:?}]"
2342 );
2343 emit(&full);
2344 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2345 emit(&full);
2346 }
2347
2348 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2354 let mut out: Vec<(FinishReason, String)> = finishes
2355 .into_iter()
2356 .zip(chunks.into_iter().map(|c| c.concat()))
2357 .collect();
2358 if let Some(first) = out.first_mut() {
2359 first.1 = full;
2360 }
2361 Ok((out, usage))
2362}
2363
2364#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2369 model: &Model,
2370 prompt: &str,
2371 params: &GenerationParams,
2372 kv_pool: Option<&generate::KvPoolConfig>,
2373 paged_kv: Option<&generate::PagedKvConfig>,
2374 prefix_cache: Option<&Mutex<PrefixCache>>,
2375 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2376 ceiling: Option<&budget::ContextCeiling>,
2377 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2378 ) -> Result<(Vec<(FinishReason, String)>, generate::Usage), generate::DecodeError> {
2381 run_generation_emit(
2382 model,
2383 prompt,
2384 params,
2385 kv_pool,
2386 paged_kv,
2387 prefix_cache,
2388 continuous_batcher,
2389 ceiling,
2390 metal_private_decode_gate,
2391 |_| {},
2392 )
2393}
2394
2395pub(crate) fn prompt_from_messages(
2406 messages: &[ChatMessage],
2407 template: &chat_template::PromptTemplate,
2408 tools: &[ToolDef],
2409 extra: serde_json::Map<String, serde_json::Value>,
2410) -> Result<String, ApiError> {
2411 let rendered = if tools.is_empty() || template.handles_tools() {
2412 template.render(messages, tools, extra)
2413 } else {
2414 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2415 with_preamble.push(ChatMessage {
2416 role: "system".to_string(),
2417 content: Some(MessageContent::Text(tool_preamble(tools))),
2418 tool_calls: None,
2419 tool_call_id: None,
2420 reasoning_content: None,
2421 });
2422 with_preamble.extend_from_slice(messages);
2423 template.render(&with_preamble, &[], extra)
2424 };
2425 rendered.map_err(template_error_response)
2426}
2427
2428fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2433 (
2434 StatusCode::BAD_REQUEST,
2435 Json(serde_json::json!({
2436 "error": {
2437 "message": format!("chat template failed to render: {err}"),
2438 "type": "invalid_request_error",
2439 "param": "messages",
2440 "code": null,
2441 }
2442 })),
2443 )
2444}
2445
2446fn tool_preamble(tools: &[ToolDef]) -> String {
2459 let mut out = String::from(
2460 "You can call tools to help answer the user. To call a tool, respond with \
2461 EXACTLY one line in this format and nothing else:\n\
2462 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2463 object matching that tool's parameters>}}</tool_call>\n\n\
2464 Available tools:\n",
2465 );
2466 for t in tools {
2467 out.push_str(&format!(
2468 "- {}: {}\n parameters (JSON schema): {}\n",
2469 t.function.name,
2470 t.function.description.as_deref().unwrap_or(""),
2471 t.function
2472 .parameters
2473 .as_ref()
2474 .map(|v| v.to_string())
2475 .unwrap_or_else(|| "{}".to_string()),
2476 ));
2477 }
2478 out
2479}
2480
2481fn tool_call_deltas(
2491 events: Vec<crate::policy::parser::ToolCallEvent>,
2492 opened: &std::cell::Cell<usize>,
2493) -> (String, Vec<ToolCallDelta>) {
2494 let mut text = String::new();
2495 let mut deltas = Vec::new();
2496 for event in events {
2497 match event {
2498 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2499 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2500 opened.set(opened.get().max(index + 1));
2501 deltas.push(ToolCallDelta::opening(index, name));
2502 }
2503 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2504 if !fragment.is_empty() {
2505 deltas.push(ToolCallDelta::arguments(index, fragment));
2506 }
2507 }
2508 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2509 }
2510 }
2511 (text, deltas)
2512}
2513
2514fn build_response_message(
2525 text: String,
2526 tools: &[ToolDef],
2527 posture: output::OutputPosture,
2528 base_finish: &'static str,
2529) -> (ChatCompletionResponseMessage, &'static str) {
2530 let parsed = output::parse_output(&text, tools, posture);
2531 let calls: Vec<ToolCallOut> = parsed
2532 .calls
2533 .into_iter()
2534 .enumerate()
2535 .map(|(index, call)| ToolCallOut {
2536 id: format!("call_{index}"),
2537 kind: "function",
2538 function: ToolCallFunctionOut {
2539 name: call.name,
2540 arguments: call.arguments,
2541 },
2542 })
2543 .collect();
2544 if !calls.is_empty() {
2545 return (
2546 ChatCompletionResponseMessage {
2547 role: "assistant",
2548 content: None,
2549 reasoning_content: parsed.reasoning,
2550 tool_calls: Some(calls),
2551 },
2552 "tool_calls",
2553 );
2554 }
2555 (
2556 ChatCompletionResponseMessage {
2557 role: "assistant",
2558 content: Some(parsed.content),
2559 reasoning_content: parsed.reasoning,
2560 tool_calls: None,
2561 },
2562 base_finish,
2563 )
2564}
2565
2566fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2571 let mut history = match &req.session_id {
2572 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2573 None => req.messages.clone(),
2574 };
2575 if req.json_object_mode() {
2576 inject_json_object_system_hint(&mut history);
2577 }
2578 history
2579}
2580
2581fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2582 const HINT: &str =
2583 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2584 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2585 match &mut sys.content {
2586 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2587 s.push_str("\n\n");
2588 s.push_str(HINT);
2589 }
2590 None => {
2591 sys.content = Some(MessageContent::Text(HINT.to_string()));
2592 }
2593 _ => {}
2594 }
2595 } else {
2596 messages.insert(
2597 0,
2598 ChatMessage {
2599 role: "system".to_string(),
2600 content: Some(MessageContent::Text(HINT.to_string())),
2601 tool_calls: None,
2602 tool_call_id: None,
2603 reasoning_content: None,
2604 },
2605 );
2606 }
2607}
2608
2609async fn chat_completions(
2610 State(state): State<Arc<AppState>>,
2611 headers: axum::http::HeaderMap,
2612 Json(req): Json<ChatCompletionRequest>,
2613) -> Response {
2614 let attribution = attribution::Attribution::from_headers(&headers);
2615 state
2616 .requests_total
2617 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2618 let started = std::time::Instant::now();
2619
2620 let request_id = frink_api::next_request_id();
2624 let stream = req.stream.unwrap_or(false);
2625
2626 let refusal = cache_admin::check_admission(&state)
2632 .err()
2633 .or_else(|| req.validate_supported_fields().err());
2634 if let Some(err) = refusal {
2635 state
2636 .request_errors_total
2637 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2638 let response = err.into_response();
2639 state.record_request(stats::Record {
2640 request_id: &request_id,
2641 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2642 model: state.active_model_name(),
2643 status: response.status().as_u16(),
2644 stream,
2645 duration_ms: started.elapsed().as_millis() as u64,
2646 usage: None,
2647 attribution: &attribution,
2648 });
2649 return response;
2650 }
2651
2652 let response = if stream {
2653 chat_completions_stream(
2654 Arc::clone(&state),
2655 req,
2656 request_id.clone(),
2657 started,
2658 attribution.clone(),
2659 )
2660 .await
2661 .into_response()
2662 } else {
2663 chat_completions_full(
2664 Arc::clone(&state),
2665 req,
2666 request_id.clone(),
2667 started,
2668 attribution.clone(),
2669 )
2670 .await
2671 .into_response()
2672 };
2673
2674 if response.status().is_client_error() || response.status().is_server_error() {
2675 state
2676 .request_errors_total
2677 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2678 state.record_request(stats::Record {
2682 request_id: &request_id,
2683 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2684 model: state.active_model_name(),
2687 status: response.status().as_u16(),
2688 stream,
2689 duration_ms: started.elapsed().as_millis() as u64,
2690 usage: None,
2691 attribution: &attribution,
2692 });
2693 }
2694 state.mark_request_finished();
2695
2696 response
2697}
2698
2699async fn chat_completions_full(
2700 state: Arc<AppState>,
2701 req: ChatCompletionRequest,
2702 request_id: String,
2703 started: std::time::Instant,
2704 attribution: attribution::Attribution,
2705) -> Result<Json<ChatCompletionResponse>, ApiError> {
2706 let tools_active = req.tools_active();
2707 let active = state.require_active()?;
2711 let history = resolve_history(&state, &req);
2712 let template = active.generative()?.chat_template();
2713 let kwargs = req.resolve_template_kwargs(&template);
2714 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2715 let mut params =
2723 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2724 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2725 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2726
2727 let (completion, cache_status) = if let Some(cached) = key
2728 .as_ref()
2729 .and_then(|key| lock_cache(&state.response_cache).get(key))
2730 {
2731 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2732 (cached, "hit")
2733 } else {
2734 let (choices, usage) = decode_task::buffered(
2735 decode_task::DecodeHandles::take(&state, &active)?,
2736 prompt.clone(),
2737 params,
2738 )
2739 .await?;
2740
2741 let completion = response_cache::CachedCompletion { choices, usage };
2742 let cache_status = match key {
2748 Some(key) => match completion.clone().cacheable() {
2752 Some(cacheable) => {
2753 tracing::debug!("cache miss for key {}", key.digest());
2754 lock_cache(&state.response_cache).put(key, cacheable);
2755 "miss"
2756 }
2757 None => "skip",
2758 },
2759 None => "skip",
2760 };
2761 (completion, cache_status)
2762 };
2763 let content = completion.first_text().to_string();
2766
2767 if req.json_object_mode() {
2768 json_mode::validate_json_object_output(&content)?;
2769 }
2770
2771 if let Some(id) = &req.session_id {
2775 state.sessions.store_reply(
2776 id,
2777 ChatMessage {
2778 role: "assistant".to_string(),
2779 content: Some(MessageContent::Text(content.clone())),
2780 tool_calls: None,
2781 tool_call_id: None,
2782 reasoning_content: None,
2783 },
2784 );
2785 }
2786
2787 let posture = output::OutputPosture::resolve_full(
2792 active.reasoning_format(),
2793 active.tool_call_format(),
2794 &prompt,
2795 );
2796 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2797 let rendered: Vec<ChatCompletionChoice> = completion
2798 .choices
2799 .into_iter()
2800 .enumerate()
2801 .map(|(index, (finish, text))| {
2802 let (message, finish_reason) =
2803 build_response_message(text, tools, posture, finish.as_str());
2804 ChatCompletionChoice {
2805 index,
2806 message,
2807 finish_reason,
2808 }
2809 })
2810 .collect();
2811
2812 state.record_request(stats::Record {
2813 request_id: &request_id,
2814 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2815 model: Some(active.name().to_string()),
2818 status: 200,
2819 stream: false,
2820 duration_ms: started.elapsed().as_millis() as u64,
2821 usage: Some(&completion.usage),
2822 attribution: &attribution,
2823 });
2824
2825 Ok(Json(ChatCompletionResponse {
2826 id: request_id.clone(),
2827 request_id,
2828 object: "chat.completion",
2829 model: req.model,
2830 choices: rendered,
2831 usage: completion.usage,
2832 frink_cache: cache_status,
2833 }))
2834}
2835
2836async fn chat_completions_stream(
2837 state: Arc<AppState>,
2838 req: ChatCompletionRequest,
2839 request_id: String,
2840 started: std::time::Instant,
2841 attribution: attribution::Attribution,
2842) -> Result<Response, ApiError> {
2843 if req.several_choices() {
2853 return Err(unsupported_feature(
2854 "`n` > 1 with `stream` is not implemented: the choices would arrive one after \
2855 another rather than interleaved by `choices[].index`. Send the request without \
2856 `stream`, which serves `n` on this route.",
2857 ));
2858 }
2859 let tools_active = req.tools_active();
2860 let active = state.require_active()?;
2864 let history = resolve_history(&state, &req);
2865 let template = active.generative()?.chat_template();
2866 let kwargs = req.resolve_template_kwargs(&template);
2867 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2868 let model_name = req.model.clone();
2869 let session_id = req.session_id.clone();
2870 let sessions = state.sessions.clone();
2871
2872 let model = Arc::clone(active.generative()?);
2873 let kv_pool = state.kv_pool.clone();
2874 let paged_kv = state.paged_kv.clone();
2875 let prefix_cache = state.prefix_cache.clone();
2876 let batcher = active.batcher.clone();
2877 let ceiling = active.ceiling.clone();
2878 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2879 let mut params =
2880 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2881 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2882 let stats_state = Arc::clone(&state);
2883 let served_model = active.name().to_string();
2886 let posture = output::OutputPosture::resolve_full(
2890 active.reasoning_format(),
2891 active.tool_call_format(),
2892 &prompt,
2893 );
2894 let offered_tools: Vec<ToolDef> = if tools_active {
2897 req.tools.clone()
2898 } else {
2899 Vec::new()
2900 };
2901
2902 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2907 params.cancel = Some(cancel_token.clone());
2908
2909 let overlap = true;
2918
2919 let slot = req
2923 .stream_resumable
2924 .unwrap_or(false)
2925 .then(|| state.streams.register(&request_id));
2926 let emitter = resume::Emitter::new(slot);
2927
2928 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
2929 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
2933 id: request_id.clone(),
2934 request_id: None,
2935 object: "chat.completion.chunk",
2936 model: model_name.clone(),
2937 choices: vec![ChatCompletionChunkChoice {
2938 index: 0,
2939 delta: ChatCompletionChunkDelta {
2940 role: None,
2941 content: None,
2942 reasoning_content: None,
2943 tool_calls: None,
2944 },
2945 finish_reason: None,
2946 }],
2947 usage: None,
2948 });
2949
2950 tokio::task::spawn_blocking(move || {
2951 let _cancel_guard = cancel_guard;
2954 let tx_chunks = tx.clone();
2955 let orphan_timeout = sse::orphan_timeout_from_env();
2960 let mut first = true;
2961 let head_request_id = request_id.clone();
2962 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
2972 Rc::new(RefCell::new(posture.reasoning_parser()));
2973 let emit_reasoning = Rc::clone(&stream_reasoning);
2974 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
2978 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
2979 );
2980 let emit_tools = Rc::clone(&stream_tools);
2981 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
2985 let emit_streamed_calls = Rc::clone(&streamed_calls);
2986 let result = run_generation_emit(
2987 &model,
2988 &prompt,
2989 ¶ms,
2990 kv_pool.as_ref(),
2991 paged_kv.as_ref(),
2992 prefix_cache.as_deref(),
2993 batcher.as_ref(),
2994 ceiling.as_deref(),
2995 metal_private_decode_gate.as_deref(),
2996 |chunk| {
2997 if !overlap || chunk.is_empty() {
2998 return;
2999 }
3000 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3001 Some(parser) => {
3002 let delta = parser.push(chunk);
3003 (delta.reasoning, delta.content)
3004 }
3005 None => (String::new(), chunk.to_string()),
3006 };
3007 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3011 Some(parser) => {
3012 let (text, calls) =
3013 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3014 (text, calls)
3015 }
3016 None => (content, Vec::new()),
3017 };
3018 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3021 return;
3022 }
3023 let role = if first { Some("assistant") } else { None };
3024 let request_id = first.then(|| head_request_id.clone());
3025 first = false;
3026 let payload = ChatCompletionChunk {
3027 id: head_request_id.clone(),
3028 request_id,
3029 object: "chat.completion.chunk",
3030 model: model_name.clone(),
3031 choices: vec![ChatCompletionChunkChoice {
3032 index: 0,
3033 delta: ChatCompletionChunkDelta {
3034 role,
3035 content: (!content.is_empty()).then_some(content),
3036 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3037 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3038 },
3039 finish_reason: None,
3040 }],
3041 usage: None,
3042 };
3043 if let Err(why) =
3052 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3053 {
3054 if why == sse::SendFailure::Orphaned {
3055 tracing::warn!(
3056 "SSE stream {head_request_id} accepted nothing for the orphan \
3057 deadline; treating it as abandoned"
3058 );
3059 }
3060 if !emitter.is_resumable() {
3072 cancel_token.cancel();
3073 }
3074 }
3075 },
3076 );
3077
3078 let mut pending_request_id = first.then(|| request_id.clone());
3084
3085 match result {
3086 Ok((choices, usage)) => {
3089 let (finish, full_text) = choices
3090 .into_iter()
3091 .next()
3092 .expect("a generation produces at least one choice");
3093 if let Some(id) = &session_id {
3094 sessions.store_reply(
3095 id,
3096 ChatMessage {
3097 role: "assistant".to_string(),
3098 content: Some(MessageContent::Text(full_text.clone())),
3099 tool_calls: None,
3100 tool_call_id: None,
3101 reasoning_content: None,
3102 },
3103 );
3104 }
3105 let mut streamed_finish: Option<&'static str> = None;
3111 if overlap {
3112 let tail = stream_reasoning
3113 .borrow_mut()
3114 .as_mut()
3115 .map(|parser| parser.flush())
3116 .unwrap_or_default();
3117 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3118 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3119 let mut events = parser.push(&content);
3120 events.extend(parser.finish());
3121 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3122 content = text;
3123 tool_calls = calls;
3124 }
3125 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3126 let payload = ChatCompletionChunk {
3127 id: request_id.clone(),
3128 request_id: pending_request_id.take(),
3129 object: "chat.completion.chunk",
3130 model: model_name.clone(),
3131 choices: vec![ChatCompletionChunkChoice {
3132 index: 0,
3133 delta: ChatCompletionChunkDelta {
3134 role: None,
3135 content: (!content.is_empty()).then_some(content),
3136 reasoning_content: (!tail.reasoning.is_empty())
3137 .then_some(tail.reasoning),
3138 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3139 },
3140 finish_reason: None,
3141 }],
3142 usage: None,
3143 };
3144 let _ =
3145 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3146 }
3147 if streamed_calls.get() > 0 {
3148 streamed_finish = Some("tool_calls");
3149 }
3150 } else {
3151 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3154 let tool_calls: Vec<ToolCallDelta> = parsed
3155 .calls
3156 .iter()
3157 .enumerate()
3158 .map(|(index, call)| {
3159 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3160 })
3161 .collect();
3162 if !tool_calls.is_empty() {
3163 streamed_finish = Some("tool_calls");
3164 }
3165 if !tool_calls.is_empty()
3166 || !parsed.content.is_empty()
3167 || parsed.reasoning.is_some()
3168 {
3169 let payload = ChatCompletionChunk {
3170 id: request_id.clone(),
3171 request_id: pending_request_id.take(),
3172 object: "chat.completion.chunk",
3173 model: model_name.clone(),
3174 choices: vec![ChatCompletionChunkChoice {
3175 index: 0,
3176 delta: ChatCompletionChunkDelta {
3177 role: Some("assistant"),
3178 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3179 .then(|| parsed.content.clone()),
3180 reasoning_content: parsed.reasoning.clone(),
3181 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3182 },
3183 finish_reason: None,
3184 }],
3185 usage: None,
3186 };
3187 let _ =
3188 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3189 }
3190 }
3191 let final_finish_reason = match streamed_finish {
3195 Some(reason) if finish.as_str() != "length" => reason,
3196 _ => finish.as_str(),
3197 };
3198 let final_payload = ChatCompletionChunk {
3199 id: request_id.clone(),
3200 request_id: pending_request_id.take(),
3201 object: "chat.completion.chunk",
3202 model: model_name,
3203 choices: vec![ChatCompletionChunkChoice {
3204 index: 0,
3205 delta: ChatCompletionChunkDelta {
3206 role: None,
3207 content: None,
3208 reasoning_content: None,
3209 tool_calls: None,
3210 },
3211 finish_reason: Some(final_finish_reason),
3212 }],
3213 usage: Some(usage.clone()),
3214 };
3215 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3216 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3217 stats_state.record_request(stats::Record {
3222 request_id: &request_id,
3223 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3224 model: Some(served_model.clone()),
3225 status: 200,
3226 stream: true,
3227 duration_ms: started.elapsed().as_millis() as u64,
3228 usage: Some(&usage),
3229 attribution: &attribution,
3230 });
3231 }
3232 Err(e) => {
3233 tracing::warn!("decode error on streamed request {request_id}: {e}");
3234 stats_state.record_request(stats::Record {
3240 request_id: &request_id,
3241 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3242 model: Some(served_model.clone()),
3243 status: 500,
3244 stream: true,
3245 duration_ms: started.elapsed().as_millis() as u64,
3246 usage: None,
3247 attribution: &attribution,
3248 });
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,
3254 choices: vec![ChatCompletionChunkChoice {
3255 index: 0,
3256 delta: ChatCompletionChunkDelta {
3257 role: Some("assistant"),
3258 content: Some(format!("[error: {e}]")),
3259 reasoning_content: None,
3260 tool_calls: None,
3261 },
3262 finish_reason: Some("stop"),
3263 }],
3264 usage: None,
3265 };
3266 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3267 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3268 }
3269 }
3270 drop(emitter);
3274 });
3275
3276 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3277 Ok((
3300 [(
3301 axum::http::HeaderName::from_static("x-accel-buffering"),
3302 axum::http::HeaderValue::from_static("no"),
3303 )],
3304 Sse::new(stream),
3305 )
3306 .into_response())
3307}
3308
3309fn protected_routes() -> Router<Arc<AppState>> {
3340 use frink_api::routes;
3341
3342 Router::new()
3343 .route(routes::V1_MODELS, get(list_models))
3344 .route(routes::V1_RESPONSES, post(responses::responses))
3348 .route(
3349 &axum_path(routes::V1_RESPONSE),
3350 get(responses::responses_get),
3351 )
3352 .route(
3353 &axum_path(routes::V1_RESPONSE_CANCEL),
3354 post(responses::responses_cancel),
3355 )
3356 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3357 .route(routes::V1_STATS, get(serving_stats))
3358 .route(routes::V1_REQUESTS, get(recent_requests))
3359 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3360 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3361 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3362 .route(
3363 routes::LORA_ADAPTERS,
3364 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3365 )
3366 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3367 .route(routes::V1_CANCEL, post(cancel_generation))
3371 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3376 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3377 .route(routes::V1_MESSAGES, post(anthropic::messages))
3378 .route(
3379 routes::V1_MESSAGES_COUNT_TOKENS,
3380 post(anthropic::count_tokens),
3381 )
3382 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3383 .route(routes::COMPLETION, post(completion::completion))
3388 .route(routes::COMPLETIONS, post(completion::completion))
3389 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3390 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3391 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3397 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3398 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3399 .route(routes::V1_RERANK, post(rerank::rerank))
3403 .route(routes::RERANK, post(rerank::rerank))
3404 .route(routes::CACHE_STATS, get(cache_stats))
3405 .route(routes::METRICS, get(metrics))
3406 .route(routes::ADMIN_MODELS, get(admin::models))
3411 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3412 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3413 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3414 .route(routes::ADMIN_TASKS, get(admin::tasks))
3415 .route(&admin::cancel_route(), post(admin::cancel_task))
3416 .route(routes::ADMIN_STATS, get(admin::stats))
3417 .merge(conversations::router())
3421}
3422
3423fn axum_path(template: &str) -> String {
3424 let mut out = String::with_capacity(template.len());
3425 let mut rest = template;
3426 while let Some(open) = rest.find('{') {
3427 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3428 break;
3429 };
3430 out.push_str(&rest[..open]);
3431 out.push(':');
3432 out.push_str(&rest[open + 1..close]);
3433 rest = &rest[close + 1..];
3434 }
3435 out.push_str(rest);
3436 out
3437}
3438
3439async fn cancel_generation(
3448 State(state): State<Arc<AppState>>,
3449 Json(req): Json<frink_api::CancelGenerationRequest>,
3450) -> Response {
3451 let cancelled = state.cancels.cancel(&req.request_id);
3452 let status = if cancelled {
3453 StatusCode::OK
3454 } else {
3455 StatusCode::NOT_FOUND
3456 };
3457 let detail = if cancelled {
3458 "the generation was asked to stop; it ends at its next token".to_string()
3459 } else {
3460 "no generation with that request_id is running -- it has already \
3461 finished, was never issued, or was served by a path that does \
3462 not register for cancellation"
3463 .to_string()
3464 };
3465 (
3466 status,
3467 Json(frink_api::CancelGenerationResponse {
3468 request_id: req.request_id,
3469 cancelled,
3470 detail,
3471 }),
3472 )
3473 .into_response()
3474}
3475
3476type Activated = (
3480 Loaded,
3481 Option<serving::batch::ContinuousBatcher>,
3482 Option<Arc<budget::ContextCeiling>>,
3483);
3484
3485fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3503 let mut batcher = serving::batch::BatcherConfig::from_env();
3504 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3505 return batcher;
3508 }
3509 let Some(path) = path else {
3510 return batcher;
3511 };
3512 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3518 let Some((priced, gguf_ctx, source)) = priced else {
3519 return batcher;
3520 };
3521 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3522 tracing::warn!(
3526 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3527 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3528 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3529 admit on a number you choose.",
3530 priced.weights_bytes,
3531 priced.device_budget_bytes,
3532 );
3533 return batcher;
3534 };
3535 tracing::info!("{source}");
3536 tracing::info!("{}", derived.fit);
3537 let adopted = budget::apply_derived(&mut batcher, &derived);
3538 if adopted.max_context {
3539 tracing::info!(
3540 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3541 override with FRINK_CB_MAX_CONTEXT",
3542 derived.max_context
3543 );
3544 }
3545 if adopted.kv_blocks {
3546 tracing::info!(
3547 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3548 derived.kv_blocks,
3549 batcher.kv_block_size
3550 );
3551 }
3552 if let Some(narrowed) = adopted.max_context_narrowed {
3553 tracing::info!(
3554 "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",
3555 batcher.kv_blocks.unwrap_or_default(),
3556 batcher.kv_block_size
3557 );
3558 }
3559 batcher
3560}
3561
3562pub(crate) fn activate_loaded_model(
3571 loaded: model::LoadedModel,
3572 enable_continuous_batching: bool,
3573 path: Option<&str>,
3574 paged_kv: Option<&generate::PagedKvConfig>,
3575) -> Activated {
3576 match loaded {
3577 model::LoadedModel::Gguf(g) => {
3578 let decoder = Arc::new(g.decoder);
3579 let tokenizer = Arc::new(g.tokenizer);
3580 let config = price_batcher_config(path);
3581 let shape =
3587 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3588 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3589 let batcher = if enable_continuous_batching {
3590 tracing::info!(
3591 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3592 (stop sequences use the same pending-buffer trim as the private generate loop)"
3593 );
3594 let tok = Arc::clone(&tokenizer);
3595 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3596 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3597 Arc::clone(&decoder),
3598 decode,
3599 config,
3600 Arc::clone(&ceiling),
3601 paged_kv.cloned(),
3602 ))
3603 } else {
3604 None
3605 };
3606 (
3607 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3608 decoder,
3609 tokenizer,
3610 stop_tokens: g.stop_tokens,
3611 bos_id: g.bos_id,
3612 is_synthetic: g.is_synthetic,
3613 chat_template: g.chat_template,
3614 }))),
3615 batcher,
3616 Some(ceiling),
3617 )
3618 }
3619 model::LoadedModel::Kimi(k) => (
3620 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3621 engine: k.engine,
3622 tokenizer: k.tokenizer,
3623 stop_tokens: k.stop_tokens,
3624 chat_template: k.chat_template,
3625 }))),
3626 None,
3627 None,
3628 ),
3629 model::LoadedModel::Mla(m) => (
3630 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3631 engine: m.engine,
3632 tokenizer: m.tokenizer,
3633 stop_tokens: m.stop_tokens,
3634 bos_id: m.bos_id,
3635 name: m.name,
3636 chat_template: m.chat_template,
3637 }))),
3638 None,
3639 None,
3640 ),
3641 model::LoadedModel::Gemma4(m) => (
3642 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3643 engine: m.engine,
3644 tokenizer: m.tokenizer,
3645 stop_tokens: m.stop_tokens,
3646 bos_id: m.bos_id,
3647 name: m.name,
3648 chat_template: m.chat_template,
3649 }))),
3650 None,
3651 None,
3652 ),
3653 model::LoadedModel::Glm52(g) => (
3654 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3655 engine: g.engine,
3656 tokenizer: g.tokenizer,
3657 stop_tokens: g.stop_tokens,
3658 bos_id: g.bos_id,
3659 name: g.name,
3660 chat_template: g.chat_template,
3661 }))),
3662 None,
3663 None,
3664 ),
3665 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3670 }
3671}
3672
3673struct StartupModels {
3680 loaded: model::LoadedModel,
3681 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3682}
3683
3684fn continuous_batching_env() -> Option<bool> {
3685 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3686 .ok()
3687 .map(|v| v.trim().to_ascii_lowercase())
3688 .as_deref()
3689 {
3690 None => None,
3691 Some("1" | "true" | "yes" | "on") => Some(true),
3692 Some("0" | "false" | "no" | "off") => Some(false),
3693 _ => None,
3694 }
3695}
3696
3697fn metal_private_decode_active() -> bool {
3698 #[cfg(feature = "metal")]
3699 {
3700 BUILT_WITH_METAL
3701 && frink_metal::attn::metal_attn_enabled()
3702 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3703 }
3704 #[cfg(not(feature = "metal"))]
3705 {
3706 false
3707 }
3708}
3709
3710fn continuous_batching_compatible(
3711 loaded: &model::LoadedModel,
3712 kv_pool: &Option<generate::KvPoolConfig>,
3713 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3714 paged_kv: &Option<generate::PagedKvConfig>,
3715) -> bool {
3716 matches!(loaded, model::LoadedModel::Gguf(_))
3717 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3718}
3719
3720fn resolve_continuous_batching_enabled(
3721 loaded: &model::LoadedModel,
3722 kv_pool: &Option<generate::KvPoolConfig>,
3723 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3724 paged_kv: &Option<generate::PagedKvConfig>,
3725) -> bool {
3726 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3727 return false;
3728 }
3729 match continuous_batching_env() {
3730 Some(true) => true,
3731 Some(false) => false,
3732 None => metal_private_decode_active(),
3733 }
3734}
3735
3736fn acquire_metal_private_decode_gate(
3737 gate: Option<&std::sync::Mutex<()>>,
3738 used_batcher: bool,
3739) -> Option<std::sync::MutexGuard<'_, ()>> {
3740 if used_batcher {
3741 None
3742 } else {
3743 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3744 }
3745}
3746
3747fn build_app_state(
3748 models: StartupModels,
3749 kv_pool: Option<generate::KvPoolConfig>,
3750 paged_kv: Option<generate::PagedKvConfig>,
3751 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3752 enable_continuous_batching: bool,
3753 mcp: Option<mcp::LoadedMcpConfig>,
3754 detection: Arc<health::Detection>,
3755) -> AppState {
3756 let StartupModels { loaded, embedding } = models;
3757 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3758 let (loaded, batcher, ceiling) = activate_loaded_model(
3759 loaded,
3760 enable_continuous_batching,
3761 configured_path.as_deref(),
3762 paged_kv.as_ref(),
3763 );
3764 let id = startup_model_id();
3770 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3771 {
3772 None
3773 } else {
3774 tracing::info!(
3775 "Metal private-loop decode will serialize concurrent requests until \
3776 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3777 );
3778 Some(Arc::new(std::sync::Mutex::new(())))
3779 };
3780 AppState {
3781 embedding,
3782 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3783 id,
3784 loaded,
3785 batcher,
3786 ceiling,
3787 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3788 }))),
3789 paged_kv,
3790 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3791 tasks: Arc::new(tasks::TaskRegistry::new()),
3792 cancels: Arc::new(cancel::CancelRegistry::new()),
3793 stats: stats::Stats::new(),
3794 streams: resume::StreamRegistry::new(),
3795 model_dir: admin::model_dirs().into_iter().next(),
3796 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3797 kv_pool,
3798 prefix_cache,
3799 sessions: session::SessionStore::new(),
3800 requests_total: std::sync::atomic::AtomicU64::new(0),
3801 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3802 started_at: std::time::Instant::now(),
3803 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3804 detection,
3805 mcp,
3806 continuous_batching_enabled: enable_continuous_batching,
3807 metal_private_decode_gate,
3808 loading_model: Mutex::new(None),
3809 last_load_error: Mutex::new(None),
3810 serving: Mutex::new(crate::stats::ServingStats::default()),
3811 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3812 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3813 started_unix: unix_now(),
3814 }
3815}
3816
3817fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3825 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3826 return Ok(None);
3827 };
3828 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3829 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3830 tracing::info!(
3831 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3832 model.name(),
3833 model.architecture(),
3834 model.n_embd(),
3835 model.pooling_type().name(),
3836 model.n_ctx_train(),
3837 );
3838 Ok(Some(Arc::new(model)))
3839}
3840
3841fn unix_now() -> u64 {
3846 std::time::SystemTime::now()
3847 .duration_since(std::time::UNIX_EPOCH)
3848 .map(|d| d.as_secs())
3849 .unwrap_or(0)
3850}
3851
3852fn startup_model_id() -> Option<String> {
3857 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3858 let configured = std::fs::canonicalize(&configured).ok()?;
3859 admin::discover(&admin::model_dirs())
3860 .into_iter()
3861 .find(|d| {
3862 std::fs::canonicalize(&d.path)
3863 .map(|p| p == configured)
3864 .unwrap_or(false)
3865 })
3866 .map(|d| d.id)
3867}
3868
3869fn init_cpu_pool() {
3878 match frink_core::threads::init_cpu_pool() {
3879 Some(n) => eprintln!(
3880 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
3881 frink_core::threads::perf_core_count()
3882 ),
3883 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
3884 }
3885}
3886
3887fn announce_ready(addr: SocketAddr, scheme: &str) {
3900 use std::io::Write;
3901 let ready =
3902 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3903 let mut stdout = std::io::stdout().lock();
3904 let _ = writeln!(stdout, "{}", ready.to_line());
3905 let _ = stdout.flush();
3906}
3907
3908async fn shutdown_signal(exit_on_stdin_close: bool) {
3920 if !exit_on_stdin_close {
3921 std::future::pending::<()>().await;
3922 return;
3923 }
3924 let _ = tokio::task::spawn_blocking(|| {
3925 use std::io::Read;
3926 let mut sink = [0u8; 256];
3927 let mut stdin = std::io::stdin().lock();
3928 loop {
3929 match stdin.read(&mut sink) {
3930 Ok(0) => break,
3932 Ok(_) => continue,
3934 Err(e) => {
3935 tracing::warn!("stdin read failed ({e}); treating it as closed");
3936 break;
3937 }
3938 }
3939 }
3940 })
3941 .await;
3942 tracing::info!("stdin closed; shutting down");
3943}
3944
3945fn tokio_worker_threads() -> usize {
3950 std::env::var("FRINK_TOKIO_WORKERS")
3951 .ok()
3952 .and_then(|v| v.trim().parse::<usize>().ok())
3953 .filter(|n| *n > 0)
3954 .unwrap_or(2)
3955}
3956
3957fn install_ring_crypto_provider() {
3974 let _ = rustls::crypto::ring::default_provider().install_default();
3975}
3976
3977pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
3983 if args.list_devices {
3984 frink_models::devices::print_available_devices();
3985 return Ok(());
3986 }
3987 apply_cli_overrides(&args)?;
3988
3989 let _instance = {
3993 use frink_core::instance::{register, InstancePolicy};
3994 let policy = if args.allow_multiple_instances {
3995 InstancePolicy::Multi
3996 } else {
3997 InstancePolicy::from_env_or(InstancePolicy::Single)
3998 };
3999 let model = std::env::var("FRINK_MODEL_PATH").ok();
4000 register(
4001 "server",
4002 model.as_deref(),
4003 frink_core::instance::current_backend(),
4004 policy,
4005 )
4006 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4007 };
4008
4009 let journal = journal::Journal::from_env();
4010 eprintln!(
4011 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4012 journal.path()
4013 );
4014 journal.append(&journal::Record::session_start(
4015 env!("CARGO_PKG_VERSION"),
4016 std::process::id(),
4017 ));
4018 journal::install_panic_hook(journal.clone());
4019
4020 let mcp_config_path = args.mcp_config.clone();
4021 let exit_on_stdin_close = args.exit_on_stdin_close
4022 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4023 .map(|v| v == "1")
4024 .unwrap_or(false);
4025
4026 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4030 init_cpu_pool();
4031
4032 let runtime = tokio::runtime::Builder::new_multi_thread()
4033 .worker_threads(tokio_worker_threads())
4034 .enable_all()
4035 .build()?;
4036 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4037
4038 let reason = match &result {
4039 Ok(()) => "normal".to_string(),
4040 Err(e) => e.to_string(),
4041 };
4042 journal.append(&journal::Record::session_exit(reason));
4043
4044 runtime.shutdown_background();
4049
4050 result
4051}
4052
4053async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4054 let _ = tracing_subscriber::fmt::try_init();
4060
4061 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4071 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4072 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4073 .map(|v| v == "1")
4074 .unwrap_or(false);
4075 if let Err(msg) =
4076 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4077 {
4078 anyhow::bail!(msg);
4079 }
4080
4081 let embedding_model = load_embedding_model()?;
4088
4089 let mut loaded = model::load()?;
4090 match &loaded {
4091 model::LoadedModel::Gguf(g) => tracing::info!(
4092 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4093 g.decoder.config.name,
4094 g.is_synthetic,
4095 g.tokenizer.kind()
4096 ),
4097 model::LoadedModel::Kimi(k) => tracing::info!(
4098 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4099 k.tokenizer.vocab_size()
4100 ),
4101 model::LoadedModel::Mla(m) => tracing::info!(
4102 "loaded MLA GGUF '{}' (tokenizer={})",
4103 m.name,
4104 m.tokenizer.kind()
4105 ),
4106 model::LoadedModel::Gemma4(m) => tracing::info!(
4107 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4108 m.name,
4109 m.tokenizer.kind()
4110 ),
4111 model::LoadedModel::Glm52(g) => tracing::info!(
4112 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4113 g.name,
4114 g.tokenizer.kind()
4115 ),
4116 model::LoadedModel::Encoder(_) => {}
4119 }
4120 let metal_default_moe_budget = {
4128 #[cfg(feature = "metal")]
4129 {
4130 frink_core::metal_dense_enabled()
4131 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4132 }
4133 #[cfg(not(feature = "metal"))]
4134 {
4135 false
4136 }
4137 };
4138 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4139 let budget: u64 = budget_str
4140 .parse()
4141 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4142 match &mut loaded {
4143 model::LoadedModel::Gguf(g) => {
4144 tracing::info!(
4145 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4146 (CUDA and/or Metal matvecs when built with the matching feature)"
4147 );
4148 g.decoder.gpu_vram_budget_bytes = Some(budget);
4149 }
4150 model::LoadedModel::Kimi(_) => {
4151 tracing::warn!(
4152 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4153 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4154 );
4155 }
4156 model::LoadedModel::Mla(_) => {
4157 tracing::warn!(
4158 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4159 FFN path only today; ignoring expert VRAM budget"
4160 );
4161 }
4162 model::LoadedModel::Gemma4(_) => {
4163 tracing::warn!(
4164 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4165 ignoring expert VRAM budget"
4166 );
4167 }
4168 model::LoadedModel::Glm52(_) => {
4169 tracing::warn!(
4170 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4171 GPU expert placement not wired yet; ignoring"
4172 );
4173 }
4174 model::LoadedModel::Encoder(_) => {
4175 tracing::warn!(
4176 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4177 it has no routed experts to place; ignoring"
4178 );
4179 }
4180 }
4181 } else if metal_default_moe_budget {
4182 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4186 if let model::LoadedModel::Gguf(g) = &mut loaded {
4187 tracing::info!(
4188 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4189 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4190 );
4191 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4192 }
4193 }
4194 #[cfg(feature = "cuda")]
4195 {
4196 if frink_core::cuda_dense_enabled() {
4197 tracing::info!(
4198 "CUDA dense matvec enabled for WeightMatrix::apply \
4199 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4200 );
4201 } else {
4202 tracing::info!(
4203 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4204 );
4205 }
4206 }
4207 #[cfg(feature = "metal")]
4208 {
4209 if frink_core::metal_dense_enabled() {
4210 tracing::info!(
4211 "Metal dense matvec enabled for WeightMatrix::apply \
4212 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4213 );
4214 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4215 Some("1") | Some("true") | Some("on") | Some("attn") => {
4216 tracing::info!(
4217 "Metal fused attention requested (FRINK_METAL_ATTN): \
4218 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4219 );
4220 }
4221 _ => {}
4222 }
4223 tracing::info!(
4224 "Metal greedy GPU argmax: temperature<=0 folds \
4225 final_norm+lm_head+argmax into the dense stack"
4226 );
4227 } else {
4228 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4229 }
4230 }
4231 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4242 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4243 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4244 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4245 panic!(
4246 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4247 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4248 );
4249 }
4250 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4251 (Ok(blocks), Ok(block_size), Err(_)) => {
4252 let total_blocks: usize = blocks
4253 .parse()
4254 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4255 let block_size: usize = block_size
4256 .parse()
4257 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4258 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4264 .ok()
4265 .map(|v| {
4266 v.parse()
4267 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4268 })
4269 .unwrap_or(0);
4270 tracing::info!(
4271 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4272 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4273 );
4274 Some(generate::KvPoolConfig {
4275 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4276 queue_wait: Duration::from_millis(queue_wait_ms),
4277 })
4278 }
4279 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4280 let block_size: usize = block_size
4281 .parse()
4282 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4283 let budget: u64 = byte_budget
4284 .parse()
4285 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4286 let cfg = match &loaded {
4287 model::LoadedModel::Gguf(g) => &g.decoder.config,
4288 model::LoadedModel::Kimi(_)
4289 | model::LoadedModel::Mla(_)
4290 | model::LoadedModel::Gemma4(_)
4291 | model::LoadedModel::Glm52(_)
4292 | model::LoadedModel::Encoder(_) => {
4293 panic!(
4294 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4295 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4296 );
4297 }
4298 };
4299 let bytes_per_block = block_size
4300 * cfg.kv_heads_all_layers()
4301 * (cfg.head_dim + cfg.v_head_dim())
4302 * std::mem::size_of::<f32>();
4303 assert!(
4304 bytes_per_block > 0,
4305 "derived KV block byte size must be positive (check model config and block size)"
4306 );
4307 let total_blocks = (budget as usize / bytes_per_block).max(1);
4308 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4309 .ok()
4310 .map(|v| {
4311 v.parse()
4312 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4313 })
4314 .unwrap_or(0);
4315 tracing::info!(
4316 "KV cache block pool enabled from byte budget: {budget} bytes / \
4317 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4318 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4319 cfg.n_layers
4320 );
4321 Some(generate::KvPoolConfig {
4322 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4323 queue_wait: Duration::from_millis(queue_wait_ms),
4324 })
4325 }
4326 (Err(_), Err(_), Err(_)) => None,
4327 (Err(_), Ok(_), Err(_)) => panic!(
4328 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4329 (or unset all three to disable KV cache pooling)"
4330 ),
4331 (Ok(_), Ok(_), Ok(_)) => {
4332 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4333 }
4334 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4335 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4336 set together (or neither, to disable KV cache pooling)"
4337 ),
4338 };
4339 let paged_kv = match (
4345 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4346 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4347 ) {
4348 (Ok(blocks), Ok(block_size)) => {
4349 assert!(
4350 kv_pool.is_none(),
4351 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4352 mutually exclusive: both bound the same KV memory, by different means. \
4353 Set one."
4354 );
4355 let blocks_per_layer: usize = blocks
4369 .parse()
4370 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4371 let block_size: usize = block_size
4372 .parse()
4373 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4374 let gguf = match &loaded {
4375 model::LoadedModel::Gguf(g) => g,
4376 _ => panic!(
4377 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4378 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4379 ),
4380 };
4381 let cfg = &gguf.decoder.config;
4382 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4383 .ok()
4384 .map(|v| {
4385 v.parse()
4386 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4387 })
4388 .unwrap_or(0);
4389 tracing::info!(
4390 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4391 layer across {} layers, shared by all concurrent requests, \
4392 {queue_wait_ms}ms admission queue wait",
4393 cfg.n_layers
4394 );
4395 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4400 block_size,
4401 ))));
4402 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4411 crate::policy::parser::ToolCallFormat::infer(
4412 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4413 )
4414 .opener(),
4415 |text| {
4416 gguf.tokenizer
4417 .encode(text, SpecialTokens::Parse)
4418 .into_iter()
4419 .map(|t| t as u32)
4420 .collect()
4421 },
4422 );
4423 if let Some(id) = anchor_token {
4424 tracing::info!(
4425 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4426 window stops short of where its next turn rejoins"
4427 );
4428 }
4429 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4430 .ok()
4431 .map(|v| {
4432 v.parse()
4433 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4434 })
4435 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4436 if let Some(window) = cfg.uniform_sliding_window() {
4437 tracing::info!(
4438 "Paged KV window slide enabled: every layer slides by {window} every \
4439 {slide_interval} decode steps, so a request holds its prompt and a \
4440 window rather than its whole context"
4441 );
4442 } else if cfg.kv_block_window().is_some() {
4443 tracing::info!(
4444 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4445 and a page group holds one block in every layer"
4446 );
4447 }
4448 Some(generate::PagedKvConfig {
4449 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4452 queue_wait: Duration::from_millis(queue_wait_ms),
4453 radix,
4454 anchor_token,
4455 slide_interval,
4456 })
4457 }
4458 (Err(_), Err(_)) => None,
4459 _ => panic!(
4460 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4461 (or neither, to disable paged KV)"
4462 ),
4463 };
4464 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4471 let max_entries: usize = v
4472 .parse()
4473 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4474 if kv_pool.is_some() {
4475 tracing::warn!(
4476 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4477 caching will never be consulted while a KV pool is configured"
4478 );
4479 }
4480 assert!(
4487 paged_kv.is_none(),
4488 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4489 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4490 produce, so the cache could never hit. Set one."
4491 );
4492 tracing::info!(
4493 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4494 all requests"
4495 );
4496 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4497 });
4498 if matches!(
4499 loaded,
4500 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4501 ) && (kv_pool.is_some() || prefix_cache.is_some())
4502 {
4503 tracing::warn!(
4504 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4505 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4506 frink_models::engine's module docs"
4507 );
4508 }
4509 let enable_cb =
4510 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4511 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4512 tracing::info!(
4513 "continuous batching enabled by default on Metal for safe parallel serving \
4514 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4515 );
4516 }
4517 if continuous_batching_env() == Some(true)
4518 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4519 && (kv_pool.is_some() || prefix_cache.is_some())
4520 {
4521 tracing::warn!(
4522 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4523 (those modes keep the private generate path)"
4524 );
4525 }
4526 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4527 if let Ok(chunk) = n.parse::<usize>() {
4528 if chunk > 0 {
4529 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4530 }
4531 }
4532 }
4533 if matches!(
4534 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4535 Some("1")
4536 ) {
4537 tracing::warn!(
4538 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4539 (minimal spill; full layer offload still planned)"
4540 );
4541 }
4542
4543 let mcp = match mcp_config_path {
4544 Some(path) => {
4545 let loaded = mcp::load_mcp_config(&path)?;
4546 tracing::info!(
4547 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4548 loaded.path,
4549 loaded.servers.len()
4550 );
4551 Some(loaded)
4552 }
4553 None => None,
4554 };
4555
4556 let detection = health::Detection::spawn();
4560
4561 let state = Arc::new(build_app_state(
4562 StartupModels {
4563 loaded,
4564 embedding: embedding_model,
4565 },
4566 kv_pool,
4567 paged_kv,
4568 prefix_cache,
4569 enable_cb,
4570 mcp,
4571 detection,
4572 ));
4573
4574 use frink_api::routes;
4578
4579 let public = Router::new().route(routes::HEALTH, get(health));
4584
4585 let mut protected = protected_routes();
4586
4587 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4592 tracing::info!("API key auth enabled");
4593 let auth = limits::AuthConfig {
4594 api_key: Arc::new(key),
4595 };
4596 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4597 auth,
4598 limits::require_api_key,
4599 ));
4600 }
4601 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4602 let rpm: u32 = rpm
4603 .parse()
4604 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4605 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4606 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4607 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4608 limiter,
4609 limits::rate_limit,
4610 ));
4611 }
4612 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4620 let origins = security::parse_cors_origins(&spec)
4621 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4622 tracing::info!(
4623 "CORS enabled: {} allow-listed origin(s) ({})",
4624 origins.len(),
4625 spec
4626 );
4627 let cors = tower_http::cors::CorsLayer::new()
4628 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4629 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4630 .allow_headers([
4631 axum::http::header::CONTENT_TYPE,
4632 axum::http::header::AUTHORIZATION,
4633 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4639 axum::http::HeaderName::from_static("last-event-id"),
4643 ]);
4644 protected = protected.route_layer(cors);
4645 }
4646
4647 let app = public
4651 .merge(protected)
4652 .layer(axum::middleware::from_fn(limits::retry_after))
4653 .with_state(state);
4654
4655 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4661 install_ring_crypto_provider();
4662 match tls_paths {
4667 Some(paths) => {
4668 let config =
4669 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4670 .await
4671 .map_err(|e| {
4672 anyhow::anyhow!(
4673 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4674 paths.cert,
4675 paths.key
4676 )
4677 })?;
4678 let socket_addr: std::net::SocketAddr = addr
4679 .parse()
4680 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4681 let listener = std::net::TcpListener::bind(socket_addr)?;
4682 listener.set_nonblocking(true)?;
4690 let bound = listener.local_addr()?;
4691 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4692 announce_ready(bound, "https");
4693
4694 let handle = axum_server::Handle::new();
4695 let shutdown_handle = handle.clone();
4696 tokio::spawn(async move {
4697 shutdown_signal(exit_on_stdin_close).await;
4698 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4699 });
4700 axum_server::from_tcp_rustls(listener, config)?
4701 .handle(handle)
4702 .serve(app.into_make_service())
4703 .await?;
4704 }
4705 None => {
4706 let listener = tokio::net::TcpListener::bind(&addr).await?;
4707 let bound = listener.local_addr()?;
4708 tracing::info!("frink-server listening on {bound}");
4709 announce_ready(bound, "http");
4710 axum::serve(listener, app)
4711 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4712 .await?;
4713 }
4714 }
4715 Ok(())
4716}
4717
4718#[cfg(test)]
4719pub(crate) mod tests {
4720 use super::*;
4721 use frink_models::config::test_dense_fixture;
4722
4723 #[test]
4724 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4725 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4726 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4727 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4728 assert_eq!(parsed.port, 51999);
4729 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4730 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4732 }
4733
4734 fn test_model() -> Model {
4735 let cfg = test_dense_fixture();
4739 Model::Gguf(GgufModel {
4740 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4741 tokenizer: Arc::new(ServerTokenizer::Byte),
4742 stop_tokens: StopTokens::default(),
4743 bos_id: None,
4744 is_synthetic: true,
4745 chat_template: chat_template::PromptTemplate::plain(),
4746 })
4747 }
4748
4749 fn greedy_params(max_tokens: usize) -> GenerationParams {
4750 GenerationParams {
4751 n: 1,
4752 reasoning: None,
4753 max_tokens,
4754 sampling: SamplingParams::default(),
4755 seed: 1,
4756 stop: Vec::new(),
4757 stop_token_ids: Vec::new(),
4758 json_object: false,
4759 grammar: None,
4760 cancel: None,
4761 ignore_eos: false,
4762 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4763 lora: None,
4764 }
4765 }
4766
4767 fn test_model_full_byte_vocab() -> Model {
4771 test_model_full_byte_vocab_with_eos(None)
4772 }
4773
4774 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4781 let mut cfg = test_dense_fixture();
4782 cfg.vocab_size = 256;
4783 Model::Gguf(GgufModel {
4784 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4785 tokenizer: Arc::new(ServerTokenizer::Byte),
4786 stop_tokens: StopTokens::from_eos(eos),
4787 bos_id: None,
4788 is_synthetic: true,
4789 chat_template: chat_template::PromptTemplate::plain(),
4790 })
4791 }
4792
4793 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4797 AppState {
4798 embedding: None,
4799 paged_kv: None,
4800 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4801 id: None,
4802 loaded: Loaded::Generative(Arc::new(model)),
4803 batcher: None,
4804 ceiling: None,
4805 checkpoint_path: None,
4806 }))),
4807 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4808 tasks: Arc::new(tasks::TaskRegistry::new()),
4809 cancels: Arc::new(cancel::CancelRegistry::new()),
4810 stats: stats::Stats::new(),
4811 streams: resume::StreamRegistry::new(),
4812 model_dir: None,
4813 response_cache: Mutex::new(response_cache),
4814 kv_pool: None,
4815 prefix_cache: None,
4816 sessions: session::SessionStore::new(),
4817 requests_total: std::sync::atomic::AtomicU64::new(0),
4818 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4819 started_at: std::time::Instant::now(),
4820 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4821 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4822 mcp: None,
4823 continuous_batching_enabled: false,
4824 metal_private_decode_gate: None,
4825 loading_model: Mutex::new(None),
4826 last_load_error: Mutex::new(None),
4827 serving: Mutex::new(crate::stats::ServingStats::default()),
4828 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4829 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4830 started_unix: unix_now(),
4831 }
4832 }
4833
4834 pub(crate) fn test_app() -> Router {
4843 test_app_with_state(Arc::new(test_state(
4844 test_model_full_byte_vocab(),
4845 ResponseCache::new(1000, Duration::from_secs(3600)),
4846 )))
4847 }
4848
4849 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4852 Router::new()
4862 .route(frink_api::routes::HEALTH, get(health))
4863 .merge(protected_routes())
4864 .with_state(state)
4865 }
4866
4867 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4868 let mut cfg = test_dense_fixture();
4869 cfg.name = name;
4870 cfg.vocab_size = vocab_size;
4871 Model::Gguf(GgufModel {
4872 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4873 tokenizer: Arc::new(ServerTokenizer::Byte),
4874 stop_tokens: StopTokens::default(),
4875 bos_id: None,
4876 is_synthetic: true,
4877 chat_template: chat_template::PromptTemplate::plain(),
4878 })
4879 }
4880
4881 fn model_with_template(name: &'static str, source: &str) -> Model {
4885 let mut cfg = test_dense_fixture();
4886 cfg.name = name;
4887 cfg.vocab_size = 256;
4888 Model::Gguf(GgufModel {
4889 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4890 tokenizer: Arc::new(ServerTokenizer::Byte),
4891 stop_tokens: StopTokens::default(),
4892 bos_id: None,
4893 is_synthetic: true,
4894 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4895 Some(source),
4896 Some("qwen3"),
4897 false,
4898 true,
4899 None,
4900 None,
4901 ),
4902 })
4903 }
4904
4905 #[tokio::test]
4914 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4915 let strict = "{% if messages | length > 1 %}\
4918 {{ raise_exception('this template takes one turn') }}\
4919 {% endif %}{{ messages[0].content }}";
4920 let state = Arc::new(test_state(
4921 model_with_template("strict", strict),
4922 ResponseCache::new(4, Duration::from_secs(60)),
4923 ));
4924 let app = test_app_with_state(state);
4925
4926 let (status, body) = post_json_uri(
4927 &app,
4928 "/v1/chat/completions",
4929 serde_json::json!({
4930 "model": "strict",
4931 "stream": true,
4932 "messages": [
4933 {"role": "user", "content": "one"},
4934 {"role": "user", "content": "two"},
4935 ],
4936 }),
4937 )
4938 .await;
4939 assert_eq!(status, StatusCode::BAD_REQUEST);
4940 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
4941 assert!(
4942 body["error"]["message"]
4943 .as_str()
4944 .unwrap()
4945 .contains("one turn"),
4946 "the template's own message must reach the caller: {body}"
4947 );
4948
4949 let (status, _) = post_json_uri(
4951 &app,
4952 "/v1/chat/completions",
4953 serde_json::json!({
4954 "model": "strict",
4955 "stream": true,
4956 "max_tokens": 1,
4957 "messages": [{"role": "user", "content": "one"}],
4958 }),
4959 )
4960 .await;
4961 assert_eq!(status, StatusCode::OK);
4962 }
4963
4964 #[tokio::test]
4966 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
4967 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
4968 {% if reasoning_effort %}\
4969 {% if reasoning_effort not in ['low','medium','high'] %}\
4970 {{ raise_exception('bad effort') }}\
4971 {% endif %}[{{ reasoning_effort }}]\
4972 {% endif %}{{ messages[0].content }}";
4973 let state = Arc::new(test_state(
4974 model_with_template("thinker", reasoning),
4975 ResponseCache::new(4, Duration::from_secs(60)),
4976 ));
4977 let app = test_app_with_state(state);
4978 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4979 assert_eq!(status, StatusCode::OK);
4980 let entry = &models["data"][0];
4981 assert_eq!(
4982 entry["supported_reasoning_efforts"],
4983 serde_json::json!(["off", "low", "medium", "high"])
4984 );
4985 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
4986 }
4987
4988 #[tokio::test]
4993 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
4994 let app = test_app();
4995 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
4996 let entry = &models["data"][0];
4997 assert!(entry.get("supported_reasoning_efforts").is_none());
4998 assert!(entry.get("default_reasoning_effort").is_none());
4999 }
5000
5001 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5002 Arc::new(ActiveModel {
5003 id: Some(name.to_string()),
5004 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5005 batcher: None,
5006 ceiling: None,
5007 checkpoint_path: None,
5008 })
5009 .tap_into(state)
5010 }
5011
5012 trait TapInto {
5014 fn tap_into(self, state: &AppState) -> Self;
5015 }
5016 impl TapInto for Arc<ActiveModel> {
5017 fn tap_into(self, state: &AppState) -> Self {
5018 state.swap_active(Some(Arc::clone(&self)));
5019 self
5020 }
5021 }
5022
5023 #[test]
5029 fn an_in_flight_request_keeps_the_model_it_started_on() {
5030 let state = test_state(
5031 named_test_model("model-a", 256),
5032 ResponseCache::new(4, Duration::from_secs(60)),
5033 );
5034
5035 let in_flight = state.active().expect("a model is loaded");
5038 assert_eq!(in_flight.name(), "model-a");
5039
5040 active_model(&state, "model-b");
5041
5042 assert_eq!(state.active().unwrap().name(), "model-b");
5044 assert_eq!(in_flight.name(), "model-a");
5046 let (choices, _usage) = run_generation(
5047 in_flight.generative().unwrap(),
5048 "hi",
5049 &greedy_params(3),
5050 None,
5051 None,
5052 None,
5053 None,
5054 None,
5055 None,
5056 )
5057 .expect("the old model must still decode after being swapped out");
5058 assert!(matches!(
5059 choices[0].0,
5060 FinishReason::Length | FinishReason::Stop
5061 ));
5062 }
5063
5064 #[test]
5069 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5070 let state = test_state(
5071 named_test_model("model-a", 256),
5072 ResponseCache::new(4, Duration::from_secs(60)),
5073 );
5074 let in_flight = state.active().expect("a model is loaded");
5075 let weights = Arc::clone(in_flight.generative().unwrap());
5076 assert!(Arc::strong_count(&weights) >= 2);
5077
5078 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5079 id: Some("model-b".to_string()),
5080 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5081 batcher: None,
5082 ceiling: None,
5083 checkpoint_path: None,
5084 })));
5085 drop(previous);
5086 assert!(Arc::strong_count(&weights) >= 2);
5088 drop(in_flight);
5089 assert_eq!(Arc::strong_count(&weights), 1);
5090 }
5091
5092 #[tokio::test]
5096 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5097 let state = Arc::new(test_state(
5098 named_test_model("model-a", 256),
5099 ResponseCache::new(4, Duration::from_secs(60)),
5100 ));
5101 let app = test_app_with_state(Arc::clone(&state));
5102
5103 let (status, body) = post_json_uri(
5104 &app,
5105 frink_api::routes::ADMIN_MODELS_UNLOAD,
5106 serde_json::json!({}),
5107 )
5108 .await;
5109 assert_eq!(status, StatusCode::OK);
5110 assert_eq!(body["ok"], true);
5111 assert!(body["active"].is_null());
5112 assert!(state.active().is_none());
5113
5114 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5115 assert_eq!(status, StatusCode::OK);
5116 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5117 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5118
5119 let (status, body) = post_json_uri(
5120 &app,
5121 "/v1/chat/completions",
5122 serde_json::json!({
5123 "model": "x",
5124 "messages": [{"role": "user", "content": "hi"}]
5125 }),
5126 )
5127 .await;
5128 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5129 assert_eq!(body["error"]["type"], "model_not_loaded");
5130 }
5131
5132 #[tokio::test]
5136 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5137 let state = Arc::new(test_state(
5138 named_test_model("model-a", 256),
5139 ResponseCache::new(4, Duration::from_secs(60)),
5140 ));
5141 let app = test_app_with_state(Arc::clone(&state));
5142 state.swap_active(None);
5143
5144 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5145 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5148 assert_eq!(body["state"], "unavailable");
5149 assert_eq!(body["reason"], "model_not_loaded");
5150 assert!(body["model"].is_null());
5151 let real_weights = body["capabilities"]
5152 .as_array()
5153 .unwrap()
5154 .iter()
5155 .find(|c| c["id"] == "real_weights")
5156 .cloned()
5157 .expect("real_weights is always reported");
5158 assert_eq!(real_weights["available"], false);
5159 assert_eq!(real_weights["reason"], "model_not_loaded");
5160 }
5161
5162 #[tokio::test]
5166 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5167 let app = test_app();
5168
5169 let (status, completion) = post_json_uri(
5170 &app,
5171 "/v1/chat/completions",
5172 serde_json::json!({
5173 "model": "x",
5174 "messages": [{"role": "user", "content": "hi"}],
5175 "max_tokens": 4
5176 }),
5177 )
5178 .await;
5179 assert_eq!(status, StatusCode::OK);
5180 let request_id = completion["request_id"].as_str().unwrap().to_string();
5181
5182 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5183 assert_eq!(status, StatusCode::OK);
5184 let recent = stats["recent"].as_array().unwrap();
5185 assert_eq!(recent.len(), 1);
5186 let row = &recent[0];
5187 assert_eq!(row["request_id"], request_id);
5188 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5189 assert_eq!(row["status"], 200);
5190 assert_eq!(row["stream"], false);
5191 assert!(row["duration_ms"].is_number());
5194 assert!(row["decode_ms"].is_number());
5195 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5196 assert_eq!(
5197 stats["tokens_prompt_total"].as_u64().unwrap(),
5198 row["prompt_tokens"].as_u64().unwrap()
5199 );
5200 }
5201
5202 #[tokio::test]
5206 async fn a_rejected_request_is_recorded_too() {
5207 let state = Arc::new(test_state(
5208 named_test_model("model-a", 256),
5209 ResponseCache::new(4, Duration::from_secs(60)),
5210 ));
5211 let app = test_app_with_state(Arc::clone(&state));
5212 state.swap_active(None);
5213
5214 let (status, _) = post_json_uri(
5215 &app,
5216 "/v1/chat/completions",
5217 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5218 )
5219 .await;
5220 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5221
5222 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5223 let recent = stats["recent"].as_array().unwrap();
5224 assert_eq!(recent.len(), 1);
5225 assert_eq!(recent[0]["status"], 503);
5226 assert_eq!(recent[0]["completion_tokens"], 0);
5227 assert!(recent[0]["decode_ms"].is_null());
5228 assert_eq!(stats["errors_total"], 1);
5229 }
5230
5231 async fn post_json_with_headers(
5235 app: &Router,
5236 uri: &str,
5237 body: serde_json::Value,
5238 headers: &[(&str, &str)],
5239 ) -> (StatusCode, serde_json::Value) {
5240 use http_body_util::BodyExt;
5241 use tower::ServiceExt;
5242
5243 let mut builder = axum::http::Request::builder()
5244 .method("POST")
5245 .uri(uri)
5246 .header("content-type", "application/json");
5247 for (name, value) in headers {
5248 builder = builder.header(*name, *value);
5249 }
5250 let response = app
5251 .clone()
5252 .oneshot(
5253 builder
5254 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5255 .unwrap(),
5256 )
5257 .await
5258 .unwrap();
5259 let status = response.status();
5260 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5261 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5262 (status, json)
5263 }
5264
5265 #[tokio::test]
5269 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5270 let app = test_app();
5271
5272 let (status, _) = post_json_uri(
5273 &app,
5274 frink_api::routes::V1_TOKENIZE,
5275 serde_json::json!({"prompt": "hello"}),
5276 )
5277 .await;
5278 assert_eq!(status, StatusCode::OK);
5279 let (status, _) = post_json_uri(
5280 &app,
5281 frink_api::routes::V1_DETOKENIZE,
5282 serde_json::json!({"tokens": [104, 105]}),
5283 )
5284 .await;
5285 assert_eq!(status, StatusCode::OK);
5286 let (status, _) = post_json_uri(
5287 &app,
5288 frink_api::routes::V1_EMBEDDINGS,
5289 serde_json::json!({"input": "hello"}),
5290 )
5291 .await;
5292 assert_eq!(status, StatusCode::OK);
5293
5294 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5295 let routes: Vec<&str> = stats["recent"]
5296 .as_array()
5297 .unwrap()
5298 .iter()
5299 .map(|row| row["route"].as_str().unwrap())
5300 .collect();
5301 for expected in [
5302 frink_api::routes::V1_TOKENIZE,
5303 frink_api::routes::V1_DETOKENIZE,
5304 frink_api::routes::V1_EMBEDDINGS,
5305 ] {
5306 assert!(
5307 routes.contains(&expected),
5308 "{expected} is missing: {routes:?}"
5309 );
5310 }
5311
5312 let row = |route: &str| {
5313 stats["recent"]
5314 .as_array()
5315 .unwrap()
5316 .iter()
5317 .find(|r| r["route"] == route)
5318 .cloned()
5319 .unwrap()
5320 };
5321 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5325 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5326 assert!(embed["decode_ms"].is_null());
5327 assert_eq!(embed["completion_tokens"], 0);
5328 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5332 assert_eq!(
5333 stats["tokens_prompt_total"].as_u64().unwrap(),
5334 embed["prompt_tokens"].as_u64().unwrap(),
5335 "only the forward pass counted"
5336 );
5337 }
5338
5339 fn streaming_test_app() -> Router {
5344 let mut cfg = test_dense_fixture();
5345 cfg.vocab_size = 256;
5346 let model = Model::Gguf(GgufModel {
5347 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5348 tokenizer: Arc::new(ServerTokenizer::Byte),
5349 stop_tokens: StopTokens::default(),
5350 bos_id: None,
5351 is_synthetic: false,
5352 chat_template: chat_template::PromptTemplate::plain(),
5353 });
5354 test_app_with_state(Arc::new(test_state(
5355 model,
5356 ResponseCache::new(1000, Duration::from_secs(3600)),
5357 )))
5358 }
5359
5360 #[tokio::test]
5368 async fn chat_serves_several_choices_and_refuses_the_streaming_pair() {
5369 let app = test_app();
5370 let body = |n: u32, stream: bool| {
5371 serde_json::json!({
5372 "model": "x",
5373 "messages": [{"role": "user", "content": "hi"}],
5374 "max_tokens": 4,
5375 "temperature": 1.0,
5376 "n": n,
5377 "stream": stream
5378 })
5379 };
5380
5381 let (status, one) =
5382 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5383 assert_eq!(status, StatusCode::OK, "{one}");
5384
5385 let (status, three) =
5386 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5387 assert_eq!(status, StatusCode::OK, "{three}");
5388 let choices = three["choices"].as_array().expect("an array");
5389 assert_eq!(choices.len(), 3, "{three}");
5390 for (i, c) in choices.iter().enumerate() {
5391 assert_eq!(c["index"], i);
5392 assert!(c["message"]["role"].is_string(), "{c}");
5393 assert!(c["finish_reason"].is_string(), "{c}");
5394 }
5395 assert_eq!(
5397 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5398 "n = 3 billed the prompt more than once"
5399 );
5400
5401 let (status, refused) =
5404 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, true)).await;
5405 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{refused}");
5406 let message = refused["error"]["message"].as_str().unwrap_or_default();
5407 assert!(
5408 message.contains('n') && message.contains("stream"),
5409 "{refused}"
5410 );
5411 }
5412
5413 #[tokio::test]
5423 async fn every_route_refuses_the_same_unimplemented_fields() {
5424 let app = test_app();
5425 let fields = [
5426 ("n", serde_json::json!(3)),
5427 ("best_of", serde_json::json!(2)),
5428 ("prompt_logprobs", serde_json::json!(1)),
5429 ("echo", serde_json::json!(true)),
5430 ("use_beam_search", serde_json::json!(true)),
5431 ("truncate_prompt_tokens", serde_json::json!(8)),
5432 ("prompt_embeds", serde_json::json!("AA==")),
5433 ("allowed_token_ids", serde_json::json!([1, 2])),
5434 ("bad_words", serde_json::json!(["x"])),
5435 ("skip_special_tokens", serde_json::json!(false)),
5436 ("return_tokens_as_token_ids", serde_json::json!(true)),
5437 ];
5438 for (field, value) in fields {
5439 for (uri, base) in [
5440 (
5441 frink_api::routes::V1_CHAT_COMPLETIONS,
5442 serde_json::json!({
5443 "model": "x",
5444 "messages": [{"role": "user", "content": "hi"}],
5445 "max_tokens": 2
5446 }),
5447 ),
5448 (
5449 frink_api::routes::V1_COMPLETIONS,
5450 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
5451 ),
5452 (
5453 frink_api::routes::COMPLETION,
5454 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5455 ),
5456 ] {
5457 let mut body = base;
5458 body[field] = value.clone();
5459 if field == "n"
5464 && (uri == frink_api::routes::V1_COMPLETIONS
5465 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
5466 {
5467 let (status, answer) = post_json_uri(&app, uri, body).await;
5468 assert_eq!(
5469 status,
5470 StatusCode::OK,
5471 "{uri} refused a served `n`: {answer}"
5472 );
5473 assert_eq!(
5474 answer["choices"].as_array().map(Vec::len),
5475 Some(3),
5476 "{answer}"
5477 );
5478 continue;
5479 }
5480 let (status, answer) = post_json_uri(&app, uri, body).await;
5481 assert_eq!(
5482 status,
5483 StatusCode::NOT_IMPLEMENTED,
5484 "{uri} served `{field}` instead of refusing it: {answer}"
5485 );
5486 assert!(
5487 answer["error"]["message"]
5488 .as_str()
5489 .is_some_and(|m| m.contains(field)),
5490 "{uri} refused `{field}` without naming it: {answer}"
5491 );
5492 }
5493 }
5494 }
5495
5496 #[tokio::test]
5497 async fn the_native_completion_wire_is_not_the_openai_one() {
5498 let app = test_app();
5499
5500 let (status, native) = post_json_uri(
5501 &app,
5502 frink_api::routes::COMPLETION,
5503 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5504 )
5505 .await;
5506 assert_eq!(status, StatusCode::OK, "{native}");
5507 assert!(native["content"].is_string(), "{native}");
5508 assert_eq!(native["stop"], true);
5509 assert_eq!(native["stop_type"], "limit");
5510 assert_eq!(native["stopping_word"], "");
5511 assert_eq!(native["truncated"], false);
5512 assert_eq!(native["id_slot"], -1);
5513 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5514 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5515 assert!(
5516 native.get("choices").is_none(),
5517 "the native shape has no `choices`: {native}"
5518 );
5519
5520 let (status, openai) = post_json_uri(
5521 &app,
5522 frink_api::routes::V1_COMPLETIONS,
5523 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5524 )
5525 .await;
5526 assert_eq!(status, StatusCode::OK);
5527 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5528 assert!(
5529 openai.get("content").is_none(),
5530 "the OpenAI shape has no top-level `content`: {openai}"
5531 );
5532 }
5533
5534 #[tokio::test]
5538 async fn both_native_spellings_reach_the_same_handler() {
5539 let app = test_app();
5540 for route in [
5541 frink_api::routes::COMPLETION,
5542 frink_api::routes::COMPLETIONS,
5543 ] {
5544 let (status, body) = post_json_uri(
5545 &app,
5546 route,
5547 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5548 )
5549 .await;
5550 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5551 assert_eq!(body["stop"], true, "{route}");
5552 assert!(body["content"].is_string(), "{route}");
5553 }
5554
5555 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5558 let routes: Vec<&str> = stats["recent"]
5559 .as_array()
5560 .unwrap()
5561 .iter()
5562 .map(|row| row["route"].as_str().unwrap())
5563 .collect();
5564 assert!(
5565 routes.contains(&frink_api::routes::COMPLETION),
5566 "{routes:?}"
5567 );
5568 assert!(
5569 routes.contains(&frink_api::routes::COMPLETIONS),
5570 "{routes:?}"
5571 );
5572 }
5573
5574 #[tokio::test]
5580 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5581 let app = streaming_test_app();
5582 let raw = post_sse_raw_uri(
5583 &app,
5584 frink_api::routes::COMPLETION,
5585 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5586 )
5587 .await;
5588
5589 assert!(
5590 !raw.contains("[DONE]"),
5591 "llama.cpp's native stream has no sentinel: {raw}"
5592 );
5593 let frames: Vec<serde_json::Value> = raw
5594 .lines()
5595 .filter_map(|line| line.strip_prefix("data: "))
5596 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5597 .collect();
5598 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5599
5600 let (last, partials) = frames.split_last().unwrap();
5601 assert_eq!(last["stop"], true, "the last frame closes the stream");
5602 assert!(last["timings"].is_object(), "{last}");
5603 assert!(last["stop_type"].is_string(), "{last}");
5604 for partial in partials {
5605 assert_eq!(partial["stop"], false, "{partial}");
5606 assert!(partial["content"].is_string(), "{partial}");
5607 assert!(partial.get("timings").is_none(), "{partial}");
5611 assert!(partial.get("generation_settings").is_none(), "{partial}");
5612 }
5613 let streamed: String = partials
5616 .iter()
5617 .filter_map(|p| p["content"].as_str())
5618 .collect();
5619 assert_eq!(last["content"].as_str().unwrap(), streamed);
5620 }
5621
5622 #[tokio::test]
5627 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5628 let app = test_app();
5629 for body in [
5630 serde_json::json!({"prompt": "hi"}),
5631 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5632 ] {
5633 let (status, refusal) =
5634 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5635 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5636 assert!(
5637 refusal["error"]["message"]
5638 .as_str()
5639 .unwrap()
5640 .contains("n_predict"),
5641 "{refusal}"
5642 );
5643 }
5644 let (status, _) = post_json_uri(
5647 &app,
5648 frink_api::routes::COMPLETION,
5649 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5650 )
5651 .await;
5652 assert_eq!(status, StatusCode::OK);
5653 }
5654
5655 #[tokio::test]
5664 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5665 let app = streaming_test_app();
5666 let ask = |stop: serde_json::Value| {
5667 let app = app.clone();
5668 async move {
5669 post_json_uri(
5670 &app,
5671 frink_api::routes::COMPLETION,
5672 serde_json::json!({
5673 "prompt": "hi",
5674 "n_predict": 64,
5675 "ignore_eos": true,
5676 "stop": stop,
5677 }),
5678 )
5679 .await
5680 .1
5681 }
5682 };
5683
5684 let baseline = ask(serde_json::json!([])).await;
5685 assert_eq!(baseline["stop_type"], "limit");
5686 assert_eq!(baseline["stopping_word"], "");
5687 let text = baseline["content"].as_str().unwrap().to_string();
5688 let sentinel: String = text.chars().skip(1).take(2).collect();
5694 assert_eq!(
5695 sentinel.chars().count(),
5696 2,
5697 "the fixture must produce enough output to cut: {text:?}"
5698 );
5699 let cut = text.find(&sentinel).expect("it came out of this text");
5700
5701 let stopped = ask(serde_json::json!([sentinel])).await;
5702 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5703 assert_eq!(stopped["stopping_word"], sentinel);
5704 assert_eq!(
5705 stopped["content"].as_str().unwrap(),
5706 &text[..cut],
5707 "the answer must be cut at the sentinel, not run past it"
5708 );
5709 }
5710
5711 #[tokio::test]
5717 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5718 let app = test_app();
5719
5720 let (v1_status, v1) = post_json_uri(
5721 &app,
5722 frink_api::routes::V1_TOKENIZE,
5723 serde_json::json!({"prompt": "hello"}),
5724 )
5725 .await;
5726 let (alias_status, alias) = post_json_uri(
5727 &app,
5728 frink_api::routes::TOKENIZE,
5729 serde_json::json!({"content": "hello"}),
5730 )
5731 .await;
5732 assert_eq!(v1_status, StatusCode::OK);
5733 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5734 assert_eq!(v1["tokens"], alias["tokens"]);
5735 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5736
5737 let (status, both_ways) = post_json_uri(
5740 &app,
5741 frink_api::routes::TOKENIZE,
5742 serde_json::json!({"prompt": "hello"}),
5743 )
5744 .await;
5745 assert_eq!(status, StatusCode::OK);
5746 assert_eq!(both_ways["tokens"], v1["tokens"]);
5747 }
5748
5749 #[tokio::test]
5754 async fn detokenize_answers_under_both_dialects_keys() {
5755 let app = test_app();
5756 for route in [
5757 frink_api::routes::DETOKENIZE,
5758 frink_api::routes::V1_DETOKENIZE,
5759 ] {
5760 let (status, body) =
5761 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5762 assert_eq!(status, StatusCode::OK, "{route}");
5763 assert_eq!(body["text"], "hi", "{route}");
5764 assert_eq!(body["content"], body["text"], "{route}");
5765 }
5766 }
5767
5768 #[tokio::test]
5772 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5773 let app = test_app();
5774 let (status, _) = post_json_uri(
5775 &app,
5776 frink_api::routes::TOKENIZE,
5777 serde_json::json!({"content": "hello"}),
5778 )
5779 .await;
5780 assert_eq!(status, StatusCode::OK);
5781
5782 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5783 let routes: Vec<&str> = stats["recent"]
5784 .as_array()
5785 .unwrap()
5786 .iter()
5787 .map(|row| row["route"].as_str().unwrap())
5788 .collect();
5789 assert!(
5790 routes.contains(&frink_api::routes::TOKENIZE),
5791 "the alias must be its own row: {routes:?}"
5792 );
5793 assert!(
5794 !routes.contains(&frink_api::routes::V1_TOKENIZE),
5795 "nothing called /v1/tokenize: {routes:?}"
5796 );
5797 }
5798
5799 #[tokio::test]
5804 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5805 let mut cfg = test_dense_fixture();
5806 cfg.vocab_size = 256;
5807 let model = Model::Gguf(GgufModel {
5808 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5809 tokenizer: Arc::new(ServerTokenizer::Byte),
5810 stop_tokens: StopTokens::default(),
5811 bos_id: Some(7),
5812 is_synthetic: true,
5813 chat_template: chat_template::PromptTemplate::plain(),
5814 });
5815 let app = test_app_with_state(Arc::new(test_state(
5816 model,
5817 ResponseCache::new(1000, Duration::from_secs(3600)),
5818 )));
5819
5820 let (_, plain) = post_json_uri(
5821 &app,
5822 frink_api::routes::TOKENIZE,
5823 serde_json::json!({"content": "hi"}),
5824 )
5825 .await;
5826 let (_, special) = post_json_uri(
5827 &app,
5828 frink_api::routes::TOKENIZE,
5829 serde_json::json!({"content": "hi", "add_special": true}),
5830 )
5831 .await;
5832
5833 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
5834 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
5835 assert_eq!(special["count"], 3);
5836 }
5837
5838 #[tokio::test]
5841 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
5842 let app = test_app();
5843 let (status, _) = post_json_uri(
5844 &app,
5845 frink_api::routes::V1_EMBEDDINGS,
5846 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
5847 )
5848 .await;
5849 assert_eq!(status, StatusCode::BAD_REQUEST);
5850
5851 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5852 let recent = stats["recent"].as_array().unwrap();
5853 assert_eq!(recent.len(), 1);
5854 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
5855 assert_eq!(recent[0]["status"], 400);
5856 assert_eq!(
5857 recent[0]["prompt_tokens"], 0,
5858 "a rejected call embedded nothing"
5859 );
5860 }
5861
5862 #[tokio::test]
5865 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
5866 let app = test_app();
5867 let key = "sk-monitor-secret";
5868 let (status, _) = post_json_with_headers(
5869 &app,
5870 "/v1/chat/completions",
5871 serde_json::json!({
5872 "model": "x",
5873 "messages": [{"role": "user", "content": "hi"}],
5874 "max_tokens": 2
5875 }),
5876 &[
5877 ("authorization", &format!("Bearer {key}")),
5878 ("x-frink-client", "frink-studio"),
5879 ],
5880 )
5881 .await;
5882 assert_eq!(status, StatusCode::OK);
5883
5884 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5885 let row = stats["recent"].as_array().unwrap()[0].clone();
5886 let fingerprint = row["via_api_key"]
5887 .as_str()
5888 .expect("the row names the key that served it")
5889 .to_string();
5890 assert_eq!(fingerprint, attribution::key_fingerprint(key));
5891 assert!(!fingerprint.contains(key));
5892 assert!(
5893 !serde_json::to_string(&stats).unwrap().contains(key),
5894 "the stats payload must not carry the key in any form"
5895 );
5896 assert_eq!(row["client"], "frink-studio");
5897 }
5898
5899 #[tokio::test]
5902 async fn different_keys_are_different_callers_and_no_key_is_null() {
5903 let app = test_app();
5904 let body = serde_json::json!({
5905 "model": "x",
5906 "messages": [{"role": "user", "content": "hi"}],
5907 "max_tokens": 1
5908 });
5909 for headers in [
5910 vec![("authorization", "Bearer key-one")],
5911 vec![("authorization", "Bearer key-two")],
5912 vec![],
5913 ] {
5914 let (status, _) =
5915 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
5916 assert_eq!(status, StatusCode::OK);
5917 }
5918
5919 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5920 let recent = stats["recent"].as_array().unwrap();
5921 assert_eq!(recent.len(), 3);
5922 let one = recent[0]["via_api_key"].as_str().unwrap();
5923 let two = recent[1]["via_api_key"].as_str().unwrap();
5924 assert_ne!(one, two, "two keys must not collapse into one caller");
5925 assert!(
5926 recent[2]["via_api_key"].is_null(),
5927 "an unauthenticated call is null, not a fingerprint of nothing"
5928 );
5929 assert!(recent[2]["client"].is_null());
5930 }
5931
5932 #[tokio::test]
5937 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
5938 let state = Arc::new(test_state(
5939 named_test_model("really-loaded", 256),
5940 ResponseCache::new(4, Duration::from_secs(60)),
5941 ));
5942 let app = test_app_with_state(Arc::clone(&state));
5943
5944 let (status, _) = post_json_uri(
5945 &app,
5946 "/v1/chat/completions",
5947 serde_json::json!({
5948 "model": "gpt-4-turbo-that-is-not-here",
5949 "messages": [{"role": "user", "content": "hi"}],
5950 "max_tokens": 2
5951 }),
5952 )
5953 .await;
5954 assert_eq!(status, StatusCode::OK);
5955
5956 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5957 assert_eq!(stats["recent"][0]["model"], "really-loaded");
5958
5959 state.swap_active(None);
5962 let (status, _) = post_json_uri(
5963 &app,
5964 "/v1/chat/completions",
5965 serde_json::json!({
5966 "model": "gpt-4-turbo-that-is-not-here",
5967 "messages": [{"role": "user", "content": "hi"}]
5968 }),
5969 )
5970 .await;
5971 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5972 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5973 let recent = stats["recent"].as_array().unwrap();
5974 assert!(recent[recent.len() - 1]["model"].is_null());
5975 }
5976
5977 #[tokio::test]
5981 async fn a_streamed_row_names_the_model_it_decoded_against() {
5982 let state = Arc::new(test_state(
5983 named_test_model("model-before", 256),
5984 ResponseCache::new(4, Duration::from_secs(60)),
5985 ));
5986 let app = test_app_with_state(Arc::clone(&state));
5987 let _ = post_sse_raw(&app, resumable_request()).await;
5988 active_model(&state, "model-after");
5990
5991 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5992 assert_eq!(stats["recent"][0]["model"], "model-before");
5993 }
5994
5995 #[tokio::test]
5998 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
5999 let app = test_app();
6000 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6001 assert_eq!(status, StatusCode::OK);
6002 assert!(
6003 stats["queue_depth"].is_null(),
6004 "without continuous batching nothing queues, so there is nothing to measure"
6005 );
6006 assert!(stats["queue_rejected_total"].is_null());
6007 assert_eq!(
6008 stats["generating_now"], 0,
6009 "work in progress is measured and really is zero here"
6010 );
6011 }
6012
6013 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6018 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
6019 }
6020
6021 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6025 use http_body_util::BodyExt;
6026 use tower::ServiceExt;
6027
6028 let response = app
6029 .clone()
6030 .oneshot(
6031 axum::http::Request::builder()
6032 .method("POST")
6033 .uri(uri)
6034 .header("content-type", "application/json")
6035 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6036 .unwrap(),
6037 )
6038 .await
6039 .unwrap();
6040 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6041 String::from_utf8(bytes.to_vec()).unwrap()
6042 }
6043
6044 async fn get_json_with_headers(
6045 app: &Router,
6046 uri: &str,
6047 headers: &[(&str, &str)],
6048 ) -> (StatusCode, serde_json::Value) {
6049 use http_body_util::BodyExt;
6050 use tower::ServiceExt;
6051
6052 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6053 for (name, value) in headers {
6054 builder = builder.header(*name, *value);
6055 }
6056 let response = app
6057 .clone()
6058 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6059 .await
6060 .unwrap();
6061 let status = response.status();
6062 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6063 (
6064 status,
6065 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6066 )
6067 }
6068
6069 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6070 body.lines()
6071 .filter_map(|line| line.strip_prefix(field))
6072 .map(str::trim)
6073 .collect()
6074 }
6075
6076 fn resumable_request() -> serde_json::Value {
6077 serde_json::json!({
6078 "model": "m",
6079 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6080 "max_tokens": 4,
6081 "temperature": 0,
6082 "stream": true,
6083 "stream_resumable": true,
6084 })
6085 }
6086
6087 #[tokio::test]
6092 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6093 let app = test_app();
6094 let body = post_sse_raw(&app, resumable_request()).await;
6095
6096 let request_id = body
6097 .lines()
6098 .find_map(|l| l.strip_prefix("data: "))
6099 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6100 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6101 .expect("the first chunk names the request");
6102
6103 let ids = sse_field(&body, "id:");
6104 let datas = sse_field(&body, "data:");
6105 assert_eq!(
6106 ids.len(),
6107 datas.len(),
6108 "every event carries an id, or a reconnect cannot name where it stopped"
6109 );
6110 for (i, id) in ids.iter().enumerate() {
6111 assert_eq!(*id, format!("{request_id}:{i}"));
6112 }
6113 let retries = sse_field(&body, "retry:");
6114 assert_eq!(
6115 retries.len(),
6116 1,
6117 "the reconnect delay is stated once, not on every event"
6118 );
6119 assert_eq!(retries[0], "1500");
6120 assert!(
6121 body.contains("data: [DONE]"),
6122 "the end of stream is still stated"
6123 );
6124 }
6125
6126 #[tokio::test]
6130 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6131 let app = test_app();
6132 let mut request = resumable_request();
6133 request["stream_resumable"] = serde_json::json!(false);
6134 let body = post_sse_raw(&app, request).await;
6135 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6136 assert!(
6137 sse_field(&body, "id:").is_empty(),
6138 "an id promises a replay this stream cannot serve"
6139 );
6140 assert!(sse_field(&body, "retry:").is_empty());
6141 }
6142
6143 #[tokio::test]
6147 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6148 let app = test_app();
6149 let body = post_sse_raw(&app, resumable_request()).await;
6150 let request_id = sse_field(&body, "id:")[0]
6151 .rsplit_once(':')
6152 .unwrap()
6153 .0
6154 .to_string();
6155 let streamed: Vec<String> = sse_field(&body, "data:")
6156 .iter()
6157 .map(|d| d.to_string())
6158 .collect();
6159
6160 let (status, polled) = get_json(
6161 &app,
6162 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6163 )
6164 .await;
6165 assert_eq!(status, StatusCode::OK);
6166 let events: Vec<String> = polled["events"]
6167 .as_array()
6168 .unwrap()
6169 .iter()
6170 .map(|e| e["data"].as_str().unwrap().to_string())
6171 .collect();
6172 assert_eq!(
6173 events, streamed,
6174 "the fallback must deliver the same answer, not a re-run of it"
6175 );
6176 assert_eq!(polled["request_id"], request_id);
6177 assert_eq!(
6178 polled["done"], false,
6179 "events were still being handed out, so the client must ask again"
6180 );
6181
6182 let next = polled["next_index"].as_u64().unwrap();
6185 let (_, drained) = get_json(
6186 &app,
6187 &format!(
6188 "{}?from={next}",
6189 frink_api::routes::v1_stream_poll(&request_id)
6190 ),
6191 )
6192 .await;
6193 assert_eq!(drained["done"], true);
6194 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6195 }
6196
6197 #[tokio::test]
6201 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6202 let app = test_app();
6203 let body = post_sse_raw(&app, resumable_request()).await;
6204 let ids = sse_field(&body, "id:");
6205 let datas: Vec<String> = sse_field(&body, "data:")
6206 .iter()
6207 .map(|d| d.to_string())
6208 .collect();
6209 assert!(
6210 ids.len() >= 3,
6211 "need a few events to resume into the middle"
6212 );
6213 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6214
6215 let (status, resumed) = get_json_with_headers(
6216 &app,
6217 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6218 &[],
6219 )
6220 .await;
6221 assert_eq!(status, StatusCode::OK);
6222 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6223
6224 let (_, tail) = get_json(
6226 &app,
6227 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6228 )
6229 .await;
6230 let tail_events: Vec<String> = tail["events"]
6231 .as_array()
6232 .unwrap()
6233 .iter()
6234 .map(|e| e["data"].as_str().unwrap().to_string())
6235 .collect();
6236 assert_eq!(tail_events, datas[2..].to_vec());
6237 }
6238
6239 #[tokio::test]
6242 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6243 use http_body_util::BodyExt;
6244 use tower::ServiceExt;
6245
6246 let app = test_app();
6247 let body = post_sse_raw(&app, resumable_request()).await;
6248 let ids = sse_field(&body, "id:");
6249 let datas: Vec<String> = sse_field(&body, "data:")
6250 .iter()
6251 .map(|d| d.to_string())
6252 .collect();
6253 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6254
6255 let response = app
6256 .clone()
6257 .oneshot(
6258 axum::http::Request::builder()
6259 .method("GET")
6260 .uri(frink_api::routes::v1_stream(&request_id))
6261 .header("last-event-id", format!("{request_id}:0"))
6262 .body(axum::body::Body::empty())
6263 .unwrap(),
6264 )
6265 .await
6266 .unwrap();
6267 assert_eq!(response.status(), StatusCode::OK);
6268 assert_eq!(
6269 response
6270 .headers()
6271 .get("x-accel-buffering")
6272 .and_then(|v| v.to_str().ok()),
6273 Some("no"),
6274 "the reconnect needs the same anti-buffering header as the stream"
6275 );
6276 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6277 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6278 assert_eq!(
6279 sse_field(&resumed, "data:")
6280 .iter()
6281 .map(|d| d.to_string())
6282 .collect::<Vec<_>>(),
6283 datas[1..].to_vec()
6284 );
6285 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6286 }
6287
6288 #[tokio::test]
6292 async fn a_last_event_id_from_another_stream_is_refused() {
6293 let app = test_app();
6294 let body = post_sse_raw(&app, resumable_request()).await;
6295 let request_id = sse_field(&body, "id:")[0]
6296 .rsplit_once(':')
6297 .unwrap()
6298 .0
6299 .to_string();
6300
6301 let (status, err) = get_json_with_headers(
6302 &app,
6303 &frink_api::routes::v1_stream(&request_id),
6304 &[("last-event-id", "chatcmpl-someone-else:3")],
6305 )
6306 .await;
6307 assert_eq!(status, StatusCode::BAD_REQUEST);
6308 assert_eq!(err["error"]["code"], "bad_last_event_id");
6309 }
6310
6311 #[tokio::test]
6315 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6316 let app = test_app();
6317 let mut request = resumable_request();
6318 request["stream_resumable"] = serde_json::json!(false);
6319 let body = post_sse_raw(&app, request).await;
6320 let request_id = body
6321 .lines()
6322 .find_map(|l| l.strip_prefix("data: "))
6323 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6324 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6325 .unwrap();
6326
6327 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6328 assert_eq!(status, StatusCode::NOT_FOUND);
6329 assert_eq!(err["error"]["code"], "stream_not_found");
6330 assert!(err["error"]["message"]
6331 .as_str()
6332 .unwrap()
6333 .contains("stream_resumable"));
6334 }
6335
6336 #[test]
6340 fn the_axum_stream_patterns_match_the_published_templates() {
6341 assert_eq!(
6342 axum_path(frink_api::routes::V1_STREAM),
6343 "/v1/stream/:request_id"
6344 );
6345 assert_eq!(
6346 axum_path(frink_api::routes::V1_STREAM_POLL),
6347 "/v1/stream/:request_id/poll"
6348 );
6349 assert_eq!(
6350 frink_api::routes::v1_stream("abc"),
6351 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6352 );
6353 }
6354
6355 #[test]
6365 fn no_published_template_reaches_the_router_with_its_braces() {
6366 for template in [
6367 frink_api::routes::V1_STREAM,
6368 frink_api::routes::V1_STREAM_POLL,
6369 frink_api::routes::V1_RESPONSE,
6370 frink_api::routes::V1_RESPONSE_CANCEL,
6371 frink_api::routes::ADMIN_TASK_CANCEL,
6372 ] {
6373 assert!(
6374 template.contains('{'),
6375 "{template} is in the template list but has no placeholder"
6376 );
6377 let mounted = axum_path(template);
6378 assert!(
6379 !mounted.contains('{') && !mounted.contains('}'),
6380 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6381 );
6382 assert!(
6383 mounted.contains(':'),
6384 "{template} lost its placeholder entirely and would match one path only"
6385 );
6386 }
6387 }
6388
6389 #[tokio::test]
6397 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6398 let app = test_app();
6399 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6400 assert_eq!(status, StatusCode::NOT_FOUND);
6401 assert!(
6402 !body.is_null(),
6403 "empty body means axum never matched the route, so the id was read as a literal segment"
6404 );
6405 }
6406
6407 #[tokio::test]
6410 async fn the_task_list_starts_empty_rather_than_absent() {
6411 let app = test_app();
6412 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6413 assert_eq!(status, StatusCode::OK);
6414 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6415 }
6416
6417 #[tokio::test]
6426 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6427 assert!(
6428 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6429 "this test asserts the unconfigured behaviour"
6430 );
6431 let app = test_app();
6432 let (status, body) = post_json_uri(
6433 &app,
6434 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6435 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6436 )
6437 .await;
6438 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6439 assert!(
6440 body["error"]["message"]
6441 .as_str()
6442 .unwrap()
6443 .contains("--slot-save-path"),
6444 "{body}"
6445 );
6446 }
6447
6448 pub(crate) async fn post_json_uri(
6449 app: &Router,
6450 uri: &str,
6451 body: serde_json::Value,
6452 ) -> (StatusCode, serde_json::Value) {
6453 use http_body_util::BodyExt;
6454 use tower::ServiceExt;
6455
6456 let response = app
6457 .clone()
6458 .oneshot(
6459 axum::http::Request::builder()
6460 .method("POST")
6461 .uri(uri)
6462 .header("content-type", "application/json")
6463 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6464 .unwrap(),
6465 )
6466 .await
6467 .unwrap();
6468 let status = response.status();
6469 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6470 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6471 (status, json)
6472 }
6473
6474 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6475 post_json_uri(app, "/v1/chat/completions", body).await.1
6476 }
6477
6478 #[tokio::test]
6487 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6488 let app = test_app();
6489 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6490 assert_eq!(status, StatusCode::OK);
6491
6492 let memory = &body["memory"];
6493 if memory.is_null() {
6494 return;
6497 }
6498 assert!(
6499 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6500 "a read that produced a zero is a broken read, not an idle \
6501 engine: {memory}"
6502 );
6503 assert!(
6504 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6505 "the quantity must travel with the number: {memory}"
6506 );
6507 }
6508
6509 #[tokio::test]
6515 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6516 let app = test_app();
6517 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6518 assert_eq!(status, StatusCode::OK);
6519 for pool in ["kv_pages", "window_slots", "state_slots"] {
6520 assert!(
6521 body["pools"][pool].is_null(),
6522 "{pool} must be null rather than a zero row: {}",
6523 body["pools"]
6524 );
6525 }
6526 }
6527
6528 #[tokio::test]
6536 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6537 use http_body_util::BodyExt;
6538 use tower::ServiceExt;
6539
6540 let app = test_app();
6541 let send = |body: serde_json::Value| {
6542 let app = app.clone();
6543 async move {
6544 app.oneshot(
6545 axum::http::Request::builder()
6546 .method("POST")
6547 .uri(frink_api::routes::V1_MESSAGES)
6548 .header("content-type", "application/json")
6549 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6550 .unwrap(),
6551 )
6552 .await
6553 .unwrap()
6554 }
6555 };
6556
6557 let ok = send(serde_json::json!({
6558 "model": "test",
6559 "max_tokens": 1,
6560 "messages": [{"role": "user", "content": "hi"}],
6561 }))
6562 .await;
6563 assert_eq!(ok.status(), StatusCode::OK);
6564 let id = ok
6565 .headers()
6566 .get("request-id")
6567 .expect("a served message names its id")
6568 .to_str()
6569 .unwrap()
6570 .to_string();
6571 assert!(!id.is_empty());
6572
6573 let bad = send(serde_json::json!({"model": "test"})).await;
6576 assert!(bad.status().is_client_error());
6577 let other = bad.headers().get("request-id").expect("errors too");
6578 assert_ne!(other.to_str().unwrap(), id);
6579 let _ = bad.into_body().collect().await.unwrap();
6580 }
6581
6582 #[tokio::test]
6590 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6591 let state = Arc::new(test_state(
6592 test_model_full_byte_vocab(),
6593 ResponseCache::new(1000, Duration::from_secs(3600)),
6594 ));
6595 let app = test_app_with_state(Arc::clone(&state));
6596 let body = serde_json::json!({
6597 "model": "test",
6598 "messages": [{"role": "user", "content": "hi"}],
6599 "max_tokens": 1,
6600 });
6601
6602 state
6603 .maintenance
6604 .lock()
6605 .unwrap()
6606 .begin_rebuild()
6607 .expect("a fresh server is serving, so the rebuild starts");
6608 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6609 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6610 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6611
6612 state.maintenance.lock().unwrap().finish_rebuild(true);
6613 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6614 assert_eq!(
6615 status,
6616 StatusCode::OK,
6617 "the gate reopens; a rebuild is not a latch"
6618 );
6619 }
6620
6621 #[tokio::test]
6626 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6627 let app = test_app();
6628 let (status, body) = post_json_uri(
6629 &app,
6630 frink_api::routes::V1_CANCEL,
6631 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6632 )
6633 .await;
6634 assert_eq!(status, StatusCode::NOT_FOUND);
6635 assert_eq!(body["cancelled"], serde_json::json!(false));
6636 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6637 assert!(
6638 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6639 "the verdict must carry a human reason: {body}"
6640 );
6641 }
6642
6643 #[tokio::test]
6647 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6648 let state = Arc::new(test_state(
6649 test_model_full_byte_vocab(),
6650 ResponseCache::new(1000, Duration::from_secs(3600)),
6651 ));
6652 let app = test_app_with_state(Arc::clone(&state));
6653 let (token, _guard) = state.cancels.register("chatcmpl-live");
6654
6655 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6656 assert_eq!(status, StatusCode::OK);
6657 assert_eq!(before["generating_now"], serde_json::json!(1));
6658
6659 let (status, body) = post_json_uri(
6660 &app,
6661 frink_api::routes::V1_CANCEL,
6662 serde_json::json!({ "request_id": "chatcmpl-live" }),
6663 )
6664 .await;
6665 assert_eq!(status, StatusCode::OK);
6666 assert_eq!(body["cancelled"], serde_json::json!(true));
6667 assert!(
6668 token.is_cancelled(),
6669 "the endpoint answered ok without setting the flag the decode loop reads"
6670 );
6671 }
6672
6673 #[tokio::test]
6674 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6675 let app = test_app();
6676 let (status, tok) =
6677 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6678 assert_eq!(status, StatusCode::OK);
6679 let tokens = tok["tokens"].as_array().unwrap();
6680 assert_eq!(tok["count"], tokens.len());
6681 assert!(!tokens.is_empty());
6682
6683 let (status, detok) = post_json_uri(
6684 &app,
6685 "/v1/detokenize",
6686 serde_json::json!({ "tokens": tokens }),
6687 )
6688 .await;
6689 assert_eq!(status, StatusCode::OK);
6690 assert_eq!(detok["text"], "Hi");
6691
6692 let (status, emb) = post_json_uri(
6693 &app,
6694 "/v1/embeddings",
6695 serde_json::json!({
6696 "input": "Hi",
6697 "embedding_type": "mean"
6698 }),
6699 )
6700 .await;
6701 assert_eq!(status, StatusCode::OK);
6702 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6703 assert!(!vec.is_empty());
6704 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6705 }
6706
6707 #[tokio::test]
6713 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6714 let app = test_app();
6715 let (status, body) = post_json_uri(
6716 &app,
6717 "/v1/embeddings",
6718 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6719 )
6720 .await;
6721 assert_eq!(status, StatusCode::BAD_REQUEST);
6722 let msg = body["error"]["message"].as_str().unwrap();
6723 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6724 }
6725
6726 #[tokio::test]
6731 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6732 async fn a_real_embedding_model_serves_v1_embeddings() {
6733 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6734 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6735 if !path.exists() {
6736 eprintln!("SKIP: {} not present", path.display());
6737 return;
6738 }
6739 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6740 let mut state = test_state(
6741 test_model_full_byte_vocab(),
6742 ResponseCache::new(1000, Duration::from_secs(3600)),
6743 );
6744 state.embedding = Some(Arc::new(encoder));
6745 let app = test_app_with_state(Arc::new(state));
6746
6747 let (status, body) = post_json_uri(
6748 &app,
6749 "/v1/embeddings",
6750 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6751 )
6752 .await;
6753 assert_eq!(status, StatusCode::OK, "{body}");
6754 assert_eq!(body["model"], "bge-small-en-v1.5");
6755 let data = body["data"].as_array().unwrap();
6756 assert_eq!(data.len(), 2);
6757 for (i, row) in data.iter().enumerate() {
6758 assert_eq!(row["index"], i);
6759 let v: Vec<f64> = row["embedding"]
6760 .as_array()
6761 .unwrap()
6762 .iter()
6763 .map(|x| x.as_f64().unwrap())
6764 .collect();
6765 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6766 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6767 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6768 }
6769 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6772
6773 let (status, mean) = post_json_uri(
6777 &app,
6778 "/v1/embeddings",
6779 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6780 )
6781 .await;
6782 assert_eq!(status, StatusCode::OK);
6783 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6784 }
6785
6786 #[tokio::test]
6798 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6799 async fn an_encoder_can_be_the_loaded_model() {
6800 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6801 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6802 if !path.exists() {
6803 eprintln!("SKIP: {} not present", path.display());
6804 return;
6805 }
6806
6807 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6811 assert!(
6812 matches!(loaded, model::LoadedModel::Encoder(_)),
6813 "an encoder-only GGUF reached a decoder loader"
6814 );
6815 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6816 assert!(
6817 matches!(loaded, Loaded::Encoder(_)),
6818 "the encoder did not stay an encoder through activation"
6819 );
6820 assert!(
6821 batcher.is_none() && ceiling.is_none(),
6822 "an encoder was given a decode batcher or a KV ceiling it has no use for"
6823 );
6824
6825 let state = test_state(
6826 test_model_full_byte_vocab(),
6827 ResponseCache::new(1000, Duration::from_secs(3600)),
6828 );
6829 state.swap_active(Some(Arc::new(ActiveModel {
6830 id: None,
6831 loaded,
6832 batcher,
6833 ceiling,
6834 checkpoint_path: None,
6835 })));
6836 let app = test_app_with_state(Arc::new(state));
6837
6838 let (status, body) = post_json_uri(
6840 &app,
6841 "/v1/embeddings",
6842 serde_json::json!({ "input": "Hello world" }),
6843 )
6844 .await;
6845 assert_eq!(status, StatusCode::OK, "{body}");
6846 assert_eq!(body["model"], "bge-small-en-v1.5");
6847 let v = body["data"][0]["embedding"].as_array().unwrap();
6848 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
6849
6850 let (status, body) = post_json_uri(
6852 &app,
6853 "/v1/chat/completions",
6854 serde_json::json!({
6855 "model": "bge-small-en-v1.5",
6856 "messages": [{"role": "user", "content": "hi"}],
6857 }),
6858 )
6859 .await;
6860 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
6861 let msg = body["error"]["message"].as_str().unwrap();
6862 for fact in [
6863 "bge-small-en-v1.5",
6864 "bert",
6865 "embedding model",
6866 "/v1/embeddings",
6867 ] {
6868 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
6869 }
6870
6871 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
6873 assert_eq!(status, StatusCode::OK);
6874 let entry = &models["data"][0];
6875 assert_eq!(entry["id"], "bge-small-en-v1.5");
6876 assert_eq!(entry["frink_model_kind"], "embedding");
6877 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
6878 assert_eq!(entry["frink_n_embd"], 384);
6879 assert_eq!(entry["frink_pooling"], "CLS");
6880 assert_eq!(
6881 entry["frink_endpoints"],
6882 serde_json::json!(["/v1/embeddings"])
6883 );
6884 assert!(entry.get("supported_reasoning_efforts").is_none());
6887
6888 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
6890 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
6891 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
6892 assert_eq!(health["model"]["synthetic_weights"], false);
6893 let weights = health["capabilities"]
6894 .as_array()
6895 .unwrap()
6896 .iter()
6897 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
6898 .expect("a real-weights capability row");
6899 let detail = weights["detail"].as_str().unwrap_or_default();
6900 assert!(detail.contains("ENCODER"), "{detail}");
6901 let (status, body) = post_json_uri(
6909 &app,
6910 frink_api::routes::V1_TOKENIZE,
6911 serde_json::json!({ "content": "hello world" }),
6912 )
6913 .await;
6914 assert_eq!(
6915 status,
6916 StatusCode::OK,
6917 "an encoder has a real tokenizer: {body}"
6918 );
6919 let tokens = body["tokens"].as_array().expect("tokens array").clone();
6920 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
6921
6922 let (status, body) = post_json_uri(
6923 &app,
6924 frink_api::routes::V1_DETOKENIZE,
6925 serde_json::json!({ "tokens": tokens }),
6926 )
6927 .await;
6928 assert_eq!(status, StatusCode::OK, "{body}");
6929 let round_tripped = body["content"].as_str().expect("content").to_string();
6930 assert!(
6931 round_tripped.contains("hello") && round_tripped.contains("world"),
6932 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
6933 );
6934
6935 let (status, _) = post_json_uri(
6938 &app,
6939 "/v1/completions",
6940 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
6941 )
6942 .await;
6943 assert_eq!(
6944 status,
6945 StatusCode::NOT_IMPLEMENTED,
6946 "tokenizing an encoder must not have opened a path to generating with one"
6947 );
6948 }
6949
6950 #[tokio::test]
6955 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
6956 use http_body_util::BodyExt;
6957 use tower::ServiceExt;
6958
6959 let fixture = concat!(
6960 "../frink-models/tests/fixtures/",
6961 "frink_real_moe_test.gguf"
6962 );
6963 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
6964 let decoder = Decoder::from_gguf_with_expert_cache(
6965 &fixture,
6966 frink_models::config::test_moe_fixture(),
6967 Some(1024 * 1024),
6968 )
6969 .expect("MoE fixture must load store-backed");
6970
6971 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
6975 decoder.forward_token(1, 0, &mut caches);
6976
6977 let model = Model::Gguf(GgufModel {
6978 decoder: Arc::new(decoder),
6979 tokenizer: Arc::new(ServerTokenizer::Byte),
6980 stop_tokens: StopTokens::default(),
6981 bos_id: None,
6982 is_synthetic: false,
6983 chat_template: chat_template::PromptTemplate::plain(),
6984 });
6985 let state = Arc::new(test_state(
6986 model,
6987 ResponseCache::new(16, Duration::from_secs(60)),
6988 ));
6989 let app = Router::new()
6990 .route("/metrics", axum::routing::get(metrics))
6991 .route("/v1/chat/completions", post(chat_completions))
6992 .with_state(state);
6993
6994 let fetch_metrics = |app: Router| async move {
6995 let resp = app
6996 .oneshot(
6997 axum::http::Request::builder()
6998 .method("GET")
6999 .uri("/metrics")
7000 .body(axum::body::Body::empty())
7001 .unwrap(),
7002 )
7003 .await
7004 .unwrap();
7005 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7006 String::from_utf8(bytes.to_vec()).unwrap()
7007 };
7008
7009 let after = fetch_metrics(app.clone()).await;
7010 assert!(
7011 after.contains("frink_expert_cache_misses_total"),
7012 "streaming model must expose expert-cache metrics: {after}"
7013 );
7014 let misses: u64 = after
7015 .lines()
7016 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
7017 .and_then(|l| l.split_whitespace().nth(1))
7018 .and_then(|v| v.parse().ok())
7019 .expect("misses metric line must parse");
7020 assert!(
7021 misses > 0,
7022 "decode must have read experts through the store: {after}"
7023 );
7024 }
7025
7026 fn weather_tool() -> serde_json::Value {
7027 serde_json::json!({
7028 "type": "function",
7029 "function": {
7030 "name": "get_weather",
7031 "description": "Get the current weather for a location.",
7032 "parameters": {
7033 "type": "object",
7034 "properties": {"location": {"type": "string"}},
7035 "required": ["location"]
7036 }
7037 }
7038 })
7039 }
7040
7041 fn weather_tool_def() -> ToolDef {
7042 ToolDef {
7043 kind: "function".to_string(),
7044 function: ToolFunctionDef {
7045 name: "get_weather".to_string(),
7046 description: Some("Get the current weather for a location.".to_string()),
7047 parameters: Some(serde_json::json!({
7048 "type": "object",
7049 "properties": {"location": {"type": "string"}},
7050 "required": ["location"]
7051 })),
7052 },
7053 }
7054 }
7055
7056 #[test]
7057 fn tool_preamble_mentions_every_tool_name_and_description() {
7058 let preamble = tool_preamble(&[weather_tool_def()]);
7059 assert!(preamble.contains("get_weather"));
7060 assert!(preamble.contains("Get the current weather for a location."));
7061 assert!(preamble.contains("<tool_call>"));
7062 assert!(preamble.contains("</tool_call>"));
7063 }
7064
7065 #[test]
7066 fn a_real_marker_becomes_a_structured_tool_call() {
7067 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7068 let (message, finish) = build_response_message(
7069 text.to_string(),
7070 &[weather_tool_def()],
7071 output::OutputPosture::for_model("test-model"),
7072 "stop",
7073 );
7074 assert_eq!(finish, "tool_calls");
7075 let calls = message.tool_calls.expect("must carry a tool call");
7076 assert_eq!(calls[0].function.name, "get_weather");
7077 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7078 assert_eq!(parsed["location"], "Paris");
7079 }
7080
7081 #[test]
7082 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7083 let (message, finish) = build_response_message(
7084 "just an answer".to_string(),
7085 &[weather_tool_def()],
7086 output::OutputPosture::for_model("test-model"),
7087 "stop",
7088 );
7089 assert_eq!(finish, "stop");
7090 assert!(message.tool_calls.is_none());
7091 assert_eq!(message.content.as_deref(), Some("just an answer"));
7092 }
7093
7094 #[test]
7097 fn a_malformed_payload_is_not_a_tool_call() {
7098 let (message, finish) = build_response_message(
7099 "<tool_call>not valid json at all</tool_call>".to_string(),
7100 &[weather_tool_def()],
7101 output::OutputPosture::for_model("test-model"),
7102 "stop",
7103 );
7104 assert_eq!(finish, "stop");
7105 assert!(message.tool_calls.is_none());
7106 }
7107
7108 #[test]
7111 fn a_tool_that_was_never_offered_is_not_returned() {
7112 let (message, finish) = build_response_message(
7113 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7114 &[weather_tool_def()],
7115 output::OutputPosture::for_model("test-model"),
7116 "stop",
7117 );
7118 assert_eq!(finish, "stop");
7119 assert!(message.tool_calls.is_none());
7120 }
7121
7122 #[test]
7124 fn marker_text_with_no_tools_offered_stays_content() {
7125 let (message, finish) = build_response_message(
7126 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7127 &[],
7128 output::OutputPosture::for_model("test-model"),
7129 "stop",
7130 );
7131 assert_eq!(finish, "stop");
7132 assert!(message.tool_calls.is_none());
7133 assert!(message.content.is_some());
7134 }
7135
7136 #[test]
7140 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7141 let opened = std::cell::Cell::new(0usize);
7142 let mut parser = crate::policy::parser::ToolCallParser::new(
7143 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7144 vec![
7145 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7146 "write_file",
7147 serde_json::json!({"type": "object", "properties": {
7148 "path": {"type": "string"},
7149 "contents": {"type": "string"}
7150 }}),
7151 ),
7152 ],
7153 );
7154 let wire = "<tool_call><function=write_file>\
7155 <parameter=path>\n/tmp/x\n</parameter>\
7156 <parameter=contents>\nhello world\n</parameter>\
7157 </function></tool_call>";
7158
7159 let mut deltas = Vec::new();
7160 let mut text = String::new();
7161 for piece in wire.as_bytes().chunks(7) {
7162 let chunk = String::from_utf8_lossy(piece).into_owned();
7163 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7164 text.push_str(&more_text);
7165 deltas.extend(more);
7166 }
7167 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7168 text.push_str(&more_text);
7169 deltas.extend(more);
7170
7171 assert_eq!(opened.get(), 1, "one call opened");
7172 assert!(text.is_empty(), "the markers are not content: {text:?}");
7173
7174 let first = &deltas[0];
7175 assert_eq!(first.index, 0);
7176 assert_eq!(first.id.as_deref(), Some("call_0"));
7177 assert_eq!(first.kind, Some("function"));
7178 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7179
7180 let joined: String = deltas
7183 .iter()
7184 .filter_map(|d| d.function.arguments.clone())
7185 .collect();
7186 let parsed: serde_json::Value =
7187 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7188 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7189 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7190 assert!(
7191 deltas.len() >= 3,
7192 "the arguments arrived in pieces, not whole: {}",
7193 deltas.len()
7194 );
7195 assert!(
7196 deltas[1..].iter().all(|d| d.function.name.is_none()),
7197 "only the opening delta carries identity"
7198 );
7199 }
7200
7201 #[test]
7203 fn text_around_a_streamed_call_is_still_content() {
7204 let opened = std::cell::Cell::new(0usize);
7205 let mut parser = crate::policy::parser::ToolCallParser::new(
7206 crate::policy::parser::ToolCallFormat::Qwen25,
7207 vec![crate::policy::parser::tool_call::ToolSchema::new(
7208 "get_weather",
7209 )],
7210 );
7211 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7212 \"arguments\": {}}</tool_call> done";
7213 let mut text = String::new();
7214 for piece in wire.as_bytes().chunks(5) {
7215 let chunk = String::from_utf8_lossy(piece).into_owned();
7216 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7217 text.push_str(&more);
7218 }
7219 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7220 text.push_str(&more);
7221
7222 assert_eq!(opened.get(), 1);
7223 assert!(text.starts_with("let me check. "), "{text:?}");
7224 assert!(text.ends_with(" done"), "{text:?}");
7225 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7226 }
7227
7228 #[test]
7231 fn a_reasoning_block_is_split_out_of_the_answer() {
7232 let (message, finish) = build_response_message(
7233 "<think>weighing it up</think>The answer is 4.".to_string(),
7234 &[],
7235 output::OutputPosture::for_model("Qwen3-8B"),
7236 "stop",
7237 );
7238 assert_eq!(finish, "stop");
7239 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7240 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7241 }
7242
7243 #[test]
7246 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7247 let (message, _) = build_response_message(
7248 "Use the <think> tag like this.".to_string(),
7249 &[],
7250 output::OutputPosture::for_model("llama-3.1-8b"),
7251 "stop",
7252 );
7253 assert_eq!(
7254 message.content.as_deref(),
7255 Some("Use the <think> tag like this.")
7256 );
7257 assert!(message.reasoning_content.is_none());
7258 }
7259
7260 #[tokio::test]
7268 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7269 let app = test_app();
7270 let body = serde_json::json!({
7271 "model": "m",
7272 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7273 "max_tokens": 4,
7274 "temperature": 0,
7275 });
7276 let resp = post_json(&app, body).await;
7277 let message = &resp["choices"][0]["message"];
7278 assert!(message["content"].is_string());
7279 assert!(message.get("tool_calls").is_none());
7280 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7281 assert_eq!(resp["usage"]["completion_tokens"], 4);
7282 assert_eq!(
7283 resp["usage"]["total_tokens"],
7284 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7285 );
7286 }
7287
7288 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7289 use http_body_util::BodyExt;
7290 use tower::ServiceExt;
7291
7292 let response = app
7293 .clone()
7294 .oneshot(
7295 axum::http::Request::builder()
7296 .method("GET")
7297 .uri(uri)
7298 .body(axum::body::Body::empty())
7299 .unwrap(),
7300 )
7301 .await
7302 .unwrap();
7303 let status = response.status();
7304 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7305 (status, serde_json::from_slice(&bytes).unwrap())
7306 }
7307
7308 #[tokio::test]
7309 async fn health_answers_a_capability_handshake_not_a_boolean() {
7310 let app = test_app();
7311 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7312 assert_eq!(status, StatusCode::OK);
7313
7314 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7315 assert_eq!(health.state, frink_api::HealthState::Ready);
7316 assert!(health.pid > 0);
7317 assert!(health.server_time_unix_ms > 0);
7318 assert_eq!(health.last_request_age_seconds, None);
7321
7322 for id in [
7325 frink_api::health::capability::CPU,
7326 frink_api::health::capability::METAL,
7327 frink_api::health::capability::CUDA,
7328 frink_api::health::capability::REAL_WEIGHTS,
7329 frink_api::health::capability::CONTINUOUS_BATCHING,
7330 ] {
7331 let cap = health
7332 .capability(id)
7333 .unwrap_or_else(|| panic!("{id} missing"));
7334 assert!(!cap.reason.is_empty(), "{cap:?}");
7335 assert!(!cap.detail.is_empty(), "{cap:?}");
7336 }
7337 let weights = health
7341 .capability(frink_api::health::capability::REAL_WEIGHTS)
7342 .unwrap();
7343 assert!(!weights.available);
7344 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7345 assert!(health.model.as_ref().unwrap().synthetic_weights);
7346 }
7347
7348 #[tokio::test]
7349 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7350 let app = test_app();
7351 let _ = post_json(
7352 &app,
7353 serde_json::json!({
7354 "model": "m",
7355 "messages": [{"role": "user", "content": "\u{1}"}],
7356 "max_tokens": 1,
7357 "temperature": 0,
7358 }),
7359 )
7360 .await;
7361 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7362 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7363 let age = health
7364 .last_request_age_seconds
7365 .expect("a served request is evidence of liveness");
7366 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7367 }
7368
7369 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7371 use http_body_util::BodyExt;
7372 use tower::ServiceExt;
7373
7374 let response = app
7375 .clone()
7376 .oneshot(
7377 axum::http::Request::builder()
7378 .method("POST")
7379 .uri("/v1/chat/completions")
7380 .header("content-type", "application/json")
7381 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7382 .unwrap(),
7383 )
7384 .await
7385 .unwrap();
7386 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7387 String::from_utf8(bytes.to_vec())
7388 .unwrap()
7389 .lines()
7390 .filter_map(|line| line.strip_prefix("data: "))
7391 .filter(|payload| *payload != "[DONE]")
7392 .map(|payload| serde_json::from_str(payload).unwrap())
7393 .collect()
7394 }
7395
7396 #[tokio::test]
7397 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7398 let app = test_app();
7399 let chunks = post_sse_chunks(
7400 &app,
7401 serde_json::json!({
7402 "model": "m",
7403 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7404 "max_tokens": 4,
7405 "temperature": 0,
7406 "stream": true,
7407 }),
7408 )
7409 .await;
7410
7411 assert!(!chunks.is_empty());
7412 let request_id = chunks[0]["request_id"]
7413 .as_str()
7414 .expect("the first chunk names the request")
7415 .to_string();
7416 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7417 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7420 assert!(
7421 chunk.get("request_id").is_none(),
7422 "chunk {i} repeats request_id"
7423 );
7424 }
7425 for chunk in &chunks {
7428 assert_eq!(chunk["id"], serde_json::json!(request_id));
7429 }
7430
7431 let other = post_sse_chunks(
7432 &app,
7433 serde_json::json!({
7434 "model": "m",
7435 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7436 "max_tokens": 4,
7437 "temperature": 0,
7438 "stream": true,
7439 }),
7440 )
7441 .await;
7442 assert_ne!(
7443 other[0]["request_id"].as_str().unwrap(),
7444 request_id,
7445 "two concurrent chats must not share an id"
7446 );
7447 }
7448
7449 #[tokio::test]
7450 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7451 let app = test_app();
7452 let resp = post_json(
7453 &app,
7454 serde_json::json!({
7455 "model": "m",
7456 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7457 "max_tokens": 2,
7458 "temperature": 0,
7459 }),
7460 )
7461 .await;
7462 assert_eq!(resp["id"], resp["request_id"]);
7463 assert!(resp["request_id"]
7464 .as_str()
7465 .unwrap()
7466 .starts_with("chatcmpl-"));
7467 }
7468
7469 #[tokio::test]
7472 async fn usage_carries_separate_prefill_and_decode_timings() {
7473 let app = test_app();
7474 let resp = post_json(
7475 &app,
7476 serde_json::json!({
7477 "model": "m",
7478 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7479 "max_tokens": 4,
7480 "temperature": 0,
7481 }),
7482 )
7483 .await;
7484 let usage = &resp["usage"];
7485 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7486 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7487 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7488 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7489 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7491 }
7492
7493 #[tokio::test]
7501 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7502 let app = test_app();
7503 let body = serde_json::json!({
7504 "model": "m",
7505 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7506 "max_tokens": 4,
7507 "temperature": 0,
7508 "tools": [weather_tool()],
7509 });
7510 let resp = post_json(&app, body).await;
7511 let message = &resp["choices"][0]["message"];
7512 assert!(
7513 message["content"].is_string(),
7514 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7515 );
7516 assert!(message.get("tool_calls").is_none());
7517 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7521 }
7522
7523 #[tokio::test]
7527 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7528 let app = test_app();
7529 let body = serde_json::json!({
7530 "model": "m",
7531 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7532 "max_tokens": 3,
7533 "temperature": 0,
7534 });
7535 let first = post_json(&app, body.clone()).await;
7536 assert_eq!(first["frink_cache"], "miss");
7537 let second = post_json(&app, body).await;
7538 assert_eq!(second["frink_cache"], "hit");
7539 assert_eq!(
7540 first["choices"][0]["message"]["content"],
7541 second["choices"][0]["message"]["content"]
7542 );
7543 assert_eq!(
7544 first["choices"][0]["finish_reason"],
7545 second["choices"][0]["finish_reason"]
7546 );
7547 assert_eq!(first["usage"], second["usage"]);
7548 assert_eq!(second["usage"]["completion_tokens"], 3);
7549 }
7550
7551 #[tokio::test]
7562 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7563 let app = test_app();
7564 let plain = serde_json::json!({
7565 "model": "m",
7566 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7567 "max_tokens": 3,
7568 "temperature": 0,
7569 });
7570
7571 let first = post_json(&app, plain.clone()).await;
7572 assert_eq!(first["frink_cache"], "miss");
7573 let unconstrained = first["choices"][0]["message"]["content"]
7574 .as_str()
7575 .expect("content")
7576 .to_string();
7577
7578 let mut constrained = plain.clone();
7579 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7580 let second = post_json(&app, constrained).await;
7581 assert_eq!(
7582 second["frink_cache"], "miss",
7583 "a grammar is part of the key, so this body has never been answered"
7584 );
7585 let constrained_answer = second["choices"][0]["message"]["content"]
7589 .as_str()
7590 .expect("content")
7591 .to_string();
7592 assert!(
7593 constrained_answer.contains("-> \"yes\"]"),
7594 "the grammar must have been compiled AND applied, not skipped \
7595 by a cache hit: {constrained_answer}"
7596 );
7597 assert_ne!(
7598 constrained_answer, unconstrained,
7599 "the constrained request was served the unconstrained answer"
7600 );
7601
7602 let third = post_json(&app, plain).await;
7606 assert_eq!(third["frink_cache"], "hit");
7607 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7608 }
7609
7610 #[tokio::test]
7632 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7633 let state = Arc::new(test_state(
7634 test_model_full_byte_vocab(),
7635 ResponseCache::new(1000, Duration::from_secs(3600)),
7636 ));
7637 let app = test_app_with_state(state.clone());
7638 let plain = serde_json::json!({
7639 "model": "m",
7640 "messages": [
7641 {"role": "system", "content": "Answer in JSON when it helps."},
7642 {"role": "user", "content": "\u{1}\u{2}"},
7643 ],
7644 "max_tokens": 3,
7645 "temperature": 0,
7646 });
7647
7648 let first = post_json(&app, plain.clone()).await;
7649 assert_eq!(first["frink_cache"], "miss");
7650 assert_eq!(state.cache_stats().entries, 1);
7651
7652 let mut as_json = plain.clone();
7653 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7654 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7655 assert_eq!(
7656 status,
7657 StatusCode::BAD_REQUEST,
7658 "the demo banner is not a JSON object, whoever generated it"
7659 );
7660 assert_eq!(
7661 state.cache_stats().hits,
7662 0,
7663 "a json_object request must not be answered from an entry the \
7664 JSON mask never produced"
7665 );
7666 assert_eq!(
7667 state.cache_stats().entries,
7668 2,
7669 "json_object must key its own entry, not reuse the unconstrained \
7670 one it happens to render the same prompt as"
7671 );
7672 }
7673
7674 #[tokio::test]
7686 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7687 let app = test_app_with_state(Arc::new(test_state(
7688 test_model_full_byte_vocab_with_eos(Some(0x77)),
7689 ResponseCache::new(1000, Duration::from_secs(3600)),
7690 )));
7691 let body = serde_json::json!({
7692 "model": "m",
7693 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7694 "max_tokens": 6,
7695 "temperature": 0,
7696 });
7697
7698 let stopped = post_json(&app, body.clone()).await;
7699 assert_eq!(stopped["frink_cache"], "miss");
7700 assert_eq!(
7701 stopped["choices"][0]["finish_reason"], "stop",
7702 "the fixture is only meaningful if the model's EOS really fires here"
7703 );
7704 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7705
7706 let mut ignoring = body.clone();
7707 ignoring["ignore_eos"] = serde_json::json!(true);
7708 let ran_on = post_json(&app, ignoring).await;
7709 assert_eq!(
7710 ran_on["frink_cache"], "miss",
7711 "ignore_eos is part of the key, so this body has never been answered"
7712 );
7713 assert_eq!(
7714 ran_on["usage"]["completion_tokens"], 6,
7715 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7716 );
7717 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7718 assert_ne!(
7719 ran_on["choices"][0]["message"]["content"],
7720 stopped["choices"][0]["message"]["content"]
7721 );
7722 }
7723
7724 #[tokio::test]
7730 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7731 let session_app = test_app();
7732 let manual_app = test_app();
7733
7734 let turn1 = post_json(
7736 &session_app,
7737 serde_json::json!({
7738 "model": "m",
7739 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7740 "session_id": "s1",
7741 "max_tokens": 5,
7742 "temperature": 0,
7743 }),
7744 )
7745 .await;
7746 let reply1 = turn1["choices"][0]["message"]["content"]
7747 .as_str()
7748 .unwrap()
7749 .to_string();
7750
7751 let manual_turn1 = post_json(
7756 &manual_app,
7757 serde_json::json!({
7758 "model": "m",
7759 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7760 "max_tokens": 5,
7761 "temperature": 0,
7762 }),
7763 )
7764 .await;
7765 assert_eq!(
7766 manual_turn1["choices"][0]["message"]["content"]
7767 .as_str()
7768 .unwrap(),
7769 reply1
7770 );
7771
7772 let turn2 = post_json(
7774 &session_app,
7775 serde_json::json!({
7776 "model": "m",
7777 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7778 "session_id": "s1",
7779 "max_tokens": 5,
7780 "temperature": 0,
7781 }),
7782 )
7783 .await;
7784 let reply2 = turn2["choices"][0]["message"]["content"]
7785 .as_str()
7786 .unwrap()
7787 .to_string();
7788
7789 let manual_turn2 = post_json(
7793 &manual_app,
7794 serde_json::json!({
7795 "model": "m",
7796 "messages": [
7797 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7798 {"role": "assistant", "content": reply1},
7799 {"role": "user", "content": "\u{4}\u{5}"},
7800 ],
7801 "max_tokens": 5,
7802 "temperature": 0,
7803 }),
7804 )
7805 .await;
7806 assert_eq!(
7807 manual_turn2["choices"][0]["message"]["content"]
7808 .as_str()
7809 .unwrap(),
7810 reply2,
7811 "resuming a session must produce identical output to manually resending the full history"
7812 );
7813 }
7814
7815 #[test]
7818 fn lock_cache_recovers_from_a_poisoned_mutex() {
7819 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7820
7821 let poison_cache = Arc::clone(&cache);
7822 let _ = std::thread::spawn(move || {
7823 let _guard = poison_cache.lock().unwrap();
7824 panic!("simulated panic while holding the lock");
7825 })
7826 .join();
7827
7828 let recovered = lock_cache(&cache);
7830 assert_eq!(recovered.stats().entries, 0);
7831 }
7832
7833 #[test]
7834 fn is_cacheable_true_for_greedy_or_seeded_requests() {
7835 let mut req_body = serde_json::json!({
7836 "model": "m",
7837 "messages": [{"role": "user", "content": "hi"}],
7838 });
7839 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7840 assert!(
7841 req.is_cacheable(),
7842 "default (temperature 0) must be cacheable"
7843 );
7844
7845 req_body["temperature"] = serde_json::json!(0.8);
7846 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
7847 assert!(
7848 !req.is_cacheable(),
7849 "unseeded sampling must never be cacheable"
7850 );
7851
7852 req_body["seed"] = serde_json::json!(42);
7853 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
7854 assert!(
7855 req.is_cacheable(),
7856 "sampling with an explicit seed is deterministic and must be cacheable"
7857 );
7858 }
7859
7860 const GRADED: &str = "{% if reasoning_effort %}\
7864 {% if reasoning_effort not in ['low','medium','high'] %}\
7865 {{ raise_exception('unsupported effort') }}\
7866 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
7867 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
7868
7869 fn graded_template() -> chat_template::PromptTemplate {
7870 chat_template::PromptTemplate::from_gguf_metadata(
7871 Some(GRADED),
7872 Some("qwen3"),
7873 false,
7874 true,
7875 None,
7876 None,
7877 )
7878 }
7879
7880 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
7881 serde_json::from_value(value).expect("request")
7882 }
7883
7884 #[test]
7891 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
7892 let req = chat_request(serde_json::json!({
7893 "model": "m",
7894 "messages": [{"role": "user", "content": "hi"}],
7895 "grammar": "root ::= \"a\"+",
7896 }));
7897 req.validate_supported_fields()
7898 .expect("a valid grammar is a valid request");
7899 let params = req
7900 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7901 .expect("a valid grammar compiles at params time too");
7902 assert!(
7903 params.grammar.is_some(),
7904 "the grammar was dropped between the wire and the sampler"
7905 );
7906 assert!(
7907 params.needs_vocab_logits(),
7908 "a grammar request that may fold lm_head into a GPU argmax is \
7909 a grammar request served unconstrained"
7910 );
7911
7912 let plain = chat_request(serde_json::json!({
7913 "model": "m",
7914 "messages": [{"role": "user", "content": "hi"}],
7915 }));
7916 assert!(plain
7917 .generation_params(crate::sampling_knobs::SamplerModel::absent())
7918 .unwrap()
7919 .grammar
7920 .is_none());
7921 }
7922
7923 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
7924 chat_request(serde_json::json!({
7925 "model": "m",
7926 "messages": [{"role": "user", "content": "weather in Rome?"}],
7927 "tools": [weather_tool()],
7928 "tool_choice": tool_choice,
7929 }))
7930 }
7931
7932 #[test]
7936 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
7937 for choice in [
7938 serde_json::json!("required"),
7939 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
7940 ] {
7941 let req = tool_request(choice.clone());
7942 req.validate_supported_fields()
7943 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
7944 let params = req
7945 .generation_params_for_template(
7946 &graded_template(),
7947 "Qwen3-8B",
7948 crate::sampling_knobs::SamplerModel::absent(),
7949 )
7950 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
7951 let grammar = params
7952 .grammar
7953 .as_ref()
7954 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
7955 assert!(
7956 grammar.is_awaiting_trigger(),
7957 "the model must be free to think before it calls"
7958 );
7959 assert!(
7960 !grammar.allows_eog(),
7961 "{choice} must not be able to end the turn without a call"
7962 );
7963 assert!(
7969 params.needs_vocab_logits(),
7970 "{choice} would let a backend return a token id instead of logits"
7971 );
7972 assert!(
7973 !generate::greedy_gpu_fold_allowed(¶ms),
7974 "{choice} at temperature 0 must still refuse the greedy GPU fold"
7975 );
7976 }
7977 }
7978
7979 #[test]
7981 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
7982 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
7983 let req = tool_request(choice.clone());
7984 req.validate_supported_fields().expect("still supported");
7985 let params = match req.generation_params_for_template(
7986 &graded_template(),
7987 "Qwen3-8B",
7988 crate::sampling_knobs::SamplerModel::absent(),
7989 ) {
7990 Ok(p) => p,
7991 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
7992 };
7993 assert!(
7994 params.grammar.is_none(),
7995 "{choice} does not force a call and must not be constrained"
7996 );
7997 }
7998 }
7999
8000 #[test]
8003 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8004 let req = chat_request(serde_json::json!({
8006 "model": "m",
8007 "messages": [{"role": "user", "content": "hi"}],
8008 "tool_choice": "required",
8009 }));
8010 let (status, _) = req
8011 .validate_supported_fields()
8012 .expect_err("nothing to call");
8013 assert_eq!(status, StatusCode::BAD_REQUEST);
8014
8015 let req =
8017 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8018 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8019 assert_eq!(status, StatusCode::BAD_REQUEST);
8020 assert_eq!(body["error"]["param"], "tool_choice");
8021
8022 let req = tool_request(serde_json::json!({"type": "function"}));
8024 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8025 assert_eq!(status, StatusCode::BAD_REQUEST);
8026
8027 let req = chat_request(serde_json::json!({
8029 "model": "m",
8030 "messages": [{"role": "user", "content": "hi"}],
8031 "tools": [weather_tool()],
8032 "tool_choice": "required",
8033 "grammar": "root ::= \"a\"+",
8034 }));
8035 let (status, _) = req
8036 .validate_supported_fields()
8037 .expect_err("a grammar and a forced call are two constraints");
8038 assert_eq!(status, StatusCode::BAD_REQUEST);
8039
8040 let req = tool_request(serde_json::json!("required"));
8046 let (status, Json(body)) = match req.generation_params_for_template(
8047 &graded_template(),
8048 "muse-glimmer-8b",
8049 crate::sampling_knobs::SamplerModel::absent(),
8050 ) {
8051 Err(e) => e,
8052 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
8053 };
8054 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8055 assert!(
8056 body["error"]["message"]
8057 .as_str()
8058 .unwrap()
8059 .contains("muse_glimmer"),
8060 "{body}"
8061 );
8062
8063 let req = tool_request(serde_json::json!("required"));
8069 let params = req
8070 .generation_params_for_template(
8071 &graded_template(),
8072 "gemma-4-E2B-it",
8073 crate::sampling_knobs::SamplerModel::absent(),
8074 )
8075 .expect("a gemma4 forced tool_choice is served");
8076 assert!(
8077 params.grammar.is_some(),
8078 "a forced tool_choice must arrive as the generation's grammar"
8079 );
8080 }
8081
8082 #[test]
8085 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8086 let req = chat_request(serde_json::json!({
8087 "model": "m",
8088 "messages": [{"role": "user", "content": "hi"}],
8089 "grammar": "root ::= \"a",
8090 }));
8091 let (status, Json(body)) = req
8092 .validate_supported_fields()
8093 .expect_err("this does not parse");
8094 assert_eq!(status, StatusCode::BAD_REQUEST);
8095 assert_eq!(body["error"]["param"], "grammar");
8096 assert!(
8097 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8098 .is_err(),
8099 "and again at params time"
8100 );
8101 }
8102
8103 #[test]
8110 fn response_format_json_schema_becomes_the_requests_grammar() {
8111 let req = chat_request(serde_json::json!({
8112 "model": "m",
8113 "messages": [{"role": "user", "content": "hi"}],
8114 "response_format": {
8115 "type": "json_schema",
8116 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8117 },
8118 }));
8119 req.validate_supported_fields()
8120 .expect("a boolean schema converts");
8121 let params = req
8122 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8123 .expect("and compiles");
8124 let grammar = params.grammar.expect("the schema is the grammar");
8125 let mut g = (*grammar).clone();
8126 g.accept_token(0, b"true").expect("a boolean is accepted");
8127 assert!(g.allows_eog(), "and completes the parse");
8128 assert!(
8129 !params.json_object,
8130 "a schema is not the json_object character-class mask"
8131 );
8132 }
8133
8134 #[test]
8138 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8139 let req = chat_request(serde_json::json!({
8140 "model": "m",
8141 "messages": [{"role": "user", "content": "hi"}],
8142 "response_format": {
8143 "type": "json_schema",
8144 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8145 },
8146 }));
8147 let (status, Json(body)) = req
8148 .validate_supported_fields()
8149 .expect_err("minimum has no grammar in this port");
8150 assert_eq!(status, StatusCode::BAD_REQUEST);
8151 assert!(
8152 body["error"]["message"]
8153 .as_str()
8154 .expect("a message")
8155 .contains("minimum"),
8156 "the refusal must name the keyword: {body}"
8157 );
8158 assert!(
8159 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8160 .is_err(),
8161 "and again at params time"
8162 );
8163 }
8164
8165 #[test]
8171 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8172 let req = chat_request(serde_json::json!({
8173 "model": "m",
8174 "messages": [{"role": "user", "content": "hi"}],
8175 "tool_choice": "required",
8176 "tools": [{
8177 "type": "function",
8178 "function": {"name": "f", "parameters": {"type": "object"}},
8179 }],
8180 "response_format": {
8181 "type": "json_schema",
8182 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8183 },
8184 }));
8185 let (status, Json(body)) = req
8186 .validate_supported_fields()
8187 .expect_err("two constraints, one generation");
8188 assert_eq!(status, StatusCode::BAD_REQUEST);
8189 assert_eq!(body["error"]["param"], "tool_choice");
8190 }
8191
8192 #[test]
8195 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8196 let req = chat_request(serde_json::json!({
8197 "model": "m",
8198 "messages": [{"role": "user", "content": "hi"}],
8199 }));
8200 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8201 }
8202
8203 #[test]
8208 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8209 let asked = chat_request(serde_json::json!({
8210 "model": "m",
8211 "messages": [{"role": "user", "content": "hi"}],
8212 "min_p": 0.07,
8213 }));
8214 assert_eq!(
8215 asked
8216 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8217 .expect("knobs")
8218 .min_p,
8219 0.07
8220 );
8221
8222 let silent = chat_request(serde_json::json!({
8223 "model": "m",
8224 "messages": [{"role": "user", "content": "hi"}],
8225 }));
8226 assert_eq!(
8227 silent
8228 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8229 .expect("knobs")
8230 .min_p,
8231 0.0,
8232 "an unset min_p must be off, not llama.cpp's CLI default"
8233 );
8234 }
8235
8236 #[test]
8245 fn no_sampler_knob_is_missing_from_the_cache_key() {
8246 let base = serde_json::json!({
8247 "model": "m",
8248 "messages": [{"role": "user", "content": "hi"}],
8249 "seed": 1,
8250 });
8251 let key_for = |body: serde_json::Value| {
8252 let req = chat_request(body);
8253 let params = req
8254 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8255 .expect("params");
8256 req.cache_key("prompt", ¶ms)
8257 };
8258 let baseline = key_for(base.clone());
8259 for (knob, value) in [
8260 ("temperature", serde_json::json!(0.5)),
8261 ("top_p", serde_json::json!(0.9)),
8262 ("min_p", serde_json::json!(0.05)),
8263 ("top_k", serde_json::json!(40)),
8264 ("repetition_penalty", serde_json::json!(1.1)),
8265 ("presence_penalty", serde_json::json!(0.3)),
8266 ("frequency_penalty", serde_json::json!(0.3)),
8267 (
8268 "samplers",
8269 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8270 ),
8271 ] {
8272 let mut body = base.clone();
8273 body[knob] = value;
8274 assert_ne!(
8275 key_for(body),
8276 baseline,
8277 "`{knob}` is not in the cache key: two requests differing \
8278 only in it would share one cached answer"
8279 );
8280 }
8281 }
8282
8283 #[test]
8293 fn no_constraint_is_missing_from_the_cache_key() {
8294 let base = serde_json::json!({
8295 "model": "m",
8296 "messages": [{"role": "user", "content": "pick one"}],
8297 });
8298 let key_for = |body: serde_json::Value| {
8299 let req = chat_request(body);
8300 let params = req
8301 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8302 .expect("params");
8303 req.cache_key("prompt", ¶ms)
8304 };
8305 let baseline = key_for(base.clone());
8306 for (field, value) in [
8307 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8308 (
8309 "response_format",
8310 serde_json::json!({"type": "json_object"}),
8311 ),
8312 (
8313 "response_format",
8314 serde_json::json!({"type": "json_schema", "json_schema": {
8315 "name": "answer",
8316 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8317 }}),
8318 ),
8319 ("ignore_eos", serde_json::json!(true)),
8320 ("stop", serde_json::json!(["\n"])),
8321 ("max_tokens", serde_json::json!(7)),
8322 ] {
8323 let mut body = base.clone();
8324 body[field] = value.clone();
8325 assert_ne!(
8326 key_for(body),
8327 baseline,
8328 "`{field}: {value}` is not in the cache key: two requests \
8329 differing only in it would share one cached answer"
8330 );
8331 }
8332 }
8333
8334 #[test]
8338 fn an_explicit_zero_output_budget_is_a_client_error() {
8339 let req = chat_request(serde_json::json!({
8340 "model": "m",
8341 "messages": [{"role": "user", "content": "hi"}],
8342 "max_tokens": 0,
8343 }));
8344 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8345 assert_eq!(status, StatusCode::BAD_REQUEST);
8346 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8347 }
8348
8349 #[test]
8352 fn a_request_can_turn_thinking_off() {
8353 let template = graded_template();
8354 for body in [
8355 serde_json::json!({
8356 "model": "m",
8357 "messages": [{"role": "user", "content": "hi"}],
8358 "reasoning_effort": "none",
8359 }),
8360 serde_json::json!({
8361 "model": "m",
8362 "messages": [{"role": "user", "content": "hi"}],
8363 "thinking": {"type": "disabled"},
8364 }),
8365 ] {
8366 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8367 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8368 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8369 assert!(!kwargs.contains_key("reasoning_effort"));
8372 }
8373 }
8374
8375 #[test]
8378 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8379 let template = graded_template();
8380 let kwargs = chat_request(serde_json::json!({
8381 "model": "m",
8382 "messages": [{"role": "user", "content": "hi"}],
8383 "reasoning_effort": "high",
8384 "thinking": {"type": "disabled"},
8385 }))
8386 .resolve_template_kwargs(&template);
8387 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8388 assert!(!kwargs.contains_key("reasoning_effort"));
8389 }
8390
8391 #[test]
8394 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8395 let req = chat_request(serde_json::json!({
8396 "model": "m",
8397 "messages": [{"role": "user", "content": "hi"}],
8398 "thinking": {"type": "disable"},
8399 }));
8400 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8401 assert_eq!(status, StatusCode::BAD_REQUEST);
8402 }
8403
8404 #[test]
8407 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8408 let template = graded_template();
8409 let kwargs = chat_request(serde_json::json!({
8410 "model": "m",
8411 "messages": [{"role": "user", "content": "hi"}],
8412 "reasoning_effort": "none",
8413 "chat_template_kwargs": {"enable_thinking": true},
8414 }))
8415 .resolve_template_kwargs(&template);
8416 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8417 }
8418
8419 #[test]
8423 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8424 let template = graded_template();
8425 let req = chat_request(serde_json::json!({
8426 "model": "m",
8427 "messages": [{"role": "user", "content": "hi"}],
8428 "reasoning_effort": "minimal",
8429 }));
8430 let kwargs = req.resolve_template_kwargs(&template);
8431 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8432 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8433 assert!(prompt.starts_with("E:low|"), "{prompt}");
8434 }
8435
8436 #[test]
8440 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8441 let template = graded_template();
8442 let req = chat_request(serde_json::json!({
8443 "model": "m",
8444 "messages": [{"role": "user", "content": "hi"}],
8445 "chat_template_kwargs": {"reasoning_effort": "none"},
8446 }));
8447 let kwargs = req.resolve_template_kwargs(&template);
8448 assert!(!kwargs.contains_key("reasoning_effort"));
8449 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8450 assert_eq!(prompt, "hi");
8451 }
8452
8453 #[test]
8456 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8457 let template = graded_template();
8458 let req = chat_request(serde_json::json!({
8459 "model": "m",
8460 "messages": [{"role": "user", "content": "hi"}],
8461 "reasoning_effort": "low",
8462 "chat_template_kwargs": {"reasoning_effort": "high"},
8463 }));
8464 assert_eq!(
8465 req.resolve_template_kwargs(&template)["reasoning_effort"],
8466 serde_json::json!("high")
8467 );
8468 }
8469
8470 #[test]
8474 fn offering_tools_turns_thinking_on_by_itself() {
8475 let template = graded_template();
8476 let quiet = chat_request(serde_json::json!({
8477 "model": "m",
8478 "messages": [{"role": "user", "content": "hi"}],
8479 }));
8480 assert!(!quiet
8481 .resolve_template_kwargs(&template)
8482 .contains_key("enable_thinking"));
8483
8484 let with_tools = chat_request(serde_json::json!({
8485 "model": "m",
8486 "messages": [{"role": "user", "content": "hi"}],
8487 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8488 }));
8489 let kwargs = with_tools.resolve_template_kwargs(&template);
8490 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8491 let prompt =
8492 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8493 assert!(prompt.starts_with("THINK|"), "{prompt}");
8494 }
8495
8496 #[test]
8501 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8502 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8503 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8504 Some("qwen3"),
8505 false,
8506 true,
8507 None,
8508 None,
8509 );
8510 let req = chat_request(serde_json::json!({
8511 "model": "m",
8512 "messages": [{"role": "user", "content": "hi"}],
8513 "chat_template_kwargs": {"enable_thinking": true},
8514 }));
8515 let kwargs = req.resolve_template_kwargs(&opener);
8516 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8517 assert!(prompt.ends_with("<think>"), "{prompt}");
8518
8519 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8522 let (message, _) = build_response_message(
8523 "weighing it up</think>Paris.".to_string(),
8524 &[],
8525 posture,
8526 "stop",
8527 );
8528 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8529 assert_eq!(message.content.as_deref(), Some("Paris."));
8530
8531 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8534 let (message, _) = build_response_message(
8535 "weighing it up</think>Paris.".to_string(),
8536 &[],
8537 closed,
8538 "stop",
8539 );
8540 assert_eq!(message.reasoning_content, None);
8541 }
8542
8543 #[test]
8544 fn stop_param_accepts_both_single_string_and_array() {
8545 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8546 "model": "m",
8547 "messages": [{"role": "user", "content": "hi"}],
8548 "stop": "END",
8549 }))
8550 .unwrap();
8551 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8552
8553 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8554 "model": "m",
8555 "messages": [{"role": "user", "content": "hi"}],
8556 "stop": ["A", "B"],
8557 }))
8558 .unwrap();
8559 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8560 }
8561
8562 #[test]
8563 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8564 let model = test_model();
8565 let result = run_generation(
8566 &model,
8567 "hello",
8568 &greedy_params(4),
8569 None,
8570 None,
8571 None,
8572 None,
8573 None,
8574 None,
8575 );
8576 assert!(matches!(
8577 result,
8578 Err(generate::DecodeError::TokenOutOfVocab { .. })
8579 ));
8580 }
8581
8582 #[test]
8586 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8587 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8589 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8590
8591 let holder_pool = Arc::clone(&pool);
8592 let holder = std::thread::spawn(move || {
8593 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8594 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8596 drop(held);
8597 });
8598 std::thread::sleep(Duration::from_millis(15));
8599
8600 let config = generate::KvPoolConfig {
8601 pool,
8602 queue_wait: Duration::ZERO,
8603 };
8604 let result = run_generation(
8605 &model,
8606 &prompt,
8607 &greedy_params(4),
8608 Some(&config),
8609 None,
8610 None,
8611 None,
8612 None,
8613 None,
8614 );
8615 assert!(matches!(
8616 result,
8617 Err(generate::DecodeError::KvPoolExhausted)
8618 ));
8619
8620 let (status, _body) = decode_error_response(result.unwrap_err());
8621 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8622 holder.join().unwrap();
8623 }
8624
8625 #[test]
8633 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8634 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8636 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8638 let config = generate::KvPoolConfig {
8639 pool,
8640 queue_wait: Duration::ZERO,
8641 };
8642
8643 let result = run_generation(
8644 &model,
8645 &prompt,
8646 &greedy_params(4),
8647 Some(&config),
8648 None,
8649 None,
8650 None,
8651 None,
8652 None,
8653 );
8654 let err = result.expect_err("one block cannot hold two layers' caches");
8655 assert!(
8656 matches!(
8657 &err,
8658 generate::DecodeError::KvBudgetExceeded { binding, .. }
8659 if *binding == frink_models::Ceiling::DeviceMemory.code()
8660 ),
8661 "expected an immovable device-memory refusal, got {err:?}"
8662 );
8663 let (status, _body) = decode_error_response(err);
8664 assert_eq!(status, StatusCode::BAD_REQUEST);
8665 }
8666
8667 #[test]
8673 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8674 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8675 queued: 512,
8676 cap: 512,
8677 });
8678 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8679 assert_eq!(body["error"]["retry_after_seconds"], 1);
8680 let message = body["error"]["message"].as_str().expect("message");
8681 assert!(message.contains("512"), "{message}");
8682 }
8683
8684 #[test]
8685 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8686 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8687 token: 99,
8688 vocab_size: 32,
8689 });
8690 assert!(
8691 body["error"]["retry_after_seconds"].is_null(),
8692 "retrying a prompt this model cannot tokenize never helps"
8693 );
8694 }
8695
8696 #[test]
8697 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8698 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8699 token: 99,
8700 vocab_size: 32,
8701 });
8702 assert_eq!(status, StatusCode::BAD_REQUEST);
8703 }
8704
8705 #[test]
8706 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8707 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8709 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8710 let config = generate::KvPoolConfig {
8711 pool: pool.clone(),
8712 queue_wait: Duration::ZERO,
8713 };
8714
8715 let (choices, _usage) = run_generation(
8716 &model,
8717 &prompt,
8718 &greedy_params(4),
8719 Some(&config),
8720 None,
8721 None,
8722 None,
8723 None,
8724 None,
8725 )
8726 .unwrap();
8727 assert_eq!(choices[0].0, FinishReason::Length);
8728 assert_eq!(
8729 pool.lock().unwrap().free_blocks(),
8730 2,
8731 "a completed request must return its blocks to the pool"
8732 );
8733 }
8734
8735 #[tokio::test]
8740 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8741 let model = Arc::new(test_model());
8742 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8743
8744 let mut handles = Vec::new();
8745 for _ in 0..8 {
8746 let model = Arc::clone(&model);
8747 let prompt = prompt.clone();
8748 handles.push(tokio::task::spawn_blocking(move || {
8749 run_generation(
8750 &model,
8751 &prompt,
8752 &greedy_params(6),
8753 None,
8754 None,
8755 None,
8756 None,
8757 None,
8758 None,
8759 )
8760 .unwrap()
8761 }));
8762 }
8763
8764 let mut results = Vec::new();
8765 for h in handles {
8766 results.push(h.await.unwrap());
8767 }
8768 for r in &results[1..] {
8773 assert_eq!(r.0, results[0].0, "choices must match");
8777 assert_eq!(
8778 r.1.prompt_tokens, results[0].1.prompt_tokens,
8779 "prompt token count must match"
8780 );
8781 assert_eq!(
8782 r.1.completion_tokens, results[0].1.completion_tokens,
8783 "completion token count must match"
8784 );
8785 }
8786 }
8787
8788 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8795 let mut header_entries = Vec::new();
8796 let mut data = Vec::new();
8797 for (name, shape, values) in tensors {
8798 let start = data.len();
8799 for v in values {
8800 data.extend_from_slice(&v.to_le_bytes());
8801 }
8802 let end = data.len();
8803 let shape_str = shape
8804 .iter()
8805 .map(|d| d.to_string())
8806 .collect::<Vec<_>>()
8807 .join(",");
8808 header_entries.push(format!(
8809 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8810 ));
8811 }
8812 let header = format!("{{{}}}", header_entries.join(","));
8813 let header_bytes = header.as_bytes();
8814 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8815 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8816 out.extend_from_slice(header_bytes);
8817 out.extend_from_slice(&data);
8818 out
8819 }
8820
8821 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
8832 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
8833 use frink_models::kimi_loader::KimiRealHparams;
8834 use frink_moe::{GatingFunction, MoeLayerConfig};
8835
8836 let hidden_dim = 8;
8837 let kda_num_heads = 2;
8838 let kda_head_dim = 3;
8839 let kda_proj = kda_num_heads * kda_head_dim;
8840 let conv_kernel = 4;
8841 let dense_intermediate = 5;
8842 let vocab_size = 256;
8846 let mla_num_heads = 1;
8847 let mla_q_lora_rank = 2;
8848 let mla_kv_lora_rank = 2;
8849 let mla_qk_nope_head_dim = 2;
8850 let mla_qk_rope_head_dim = 2;
8851 let mla_v_head_dim = 2;
8852
8853 let model_cfg = frink_models::ModelConfig {
8854 rope_layers: frink_models::rope_layers::RopeLayers::All,
8855 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
8856 name: "synthetic-kimi-server-test",
8857 n_layers: 1,
8858 n_mtp_blocks: 0,
8859 hidden_dim,
8860 n_heads: 1,
8861 n_kv_heads: 1,
8862 head_dim: 4,
8863 v_head_dim: None,
8864 vocab_size,
8865 rope_theta: 10000.0,
8866 rms_norm_eps: 1e-5,
8867 post_norm_eps: 1e-5,
8868 sliding_window: None,
8869 moe: MoeLayerConfig {
8870 expert_weights_scale: 1.0,
8871 routed_weight_before_ffn: false,
8872 n_experts: 1,
8873 n_experts_active: 1,
8874 n_shared_experts: 0,
8875 hidden_dim,
8876 expert_ffn_dim: 4,
8877 gating: GatingFunction::Sigmoid,
8878 norm_topk_prob: true,
8879 expert_group_count: None,
8880 expert_group_used_count: None,
8881 },
8882 n_dense_leading_layers: 1,
8887 moe_interleave_step: None,
8888 norm_function: frink_models::norm::NormFunction::Rms,
8889 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
8890 kda_layers: vec![1],
8891 full_attn_layers: vec![],
8892 mla: MlaConfig {
8893 num_heads: mla_num_heads,
8894 q_lora_rank: mla_q_lora_rank,
8895 kv_lora_rank: mla_kv_lora_rank,
8896 qk_nope_head_dim: mla_qk_nope_head_dim,
8897 qk_rope_head_dim: mla_qk_rope_head_dim,
8898 v_head_dim: mla_v_head_dim,
8899 use_output_gate: true,
8900 rope: None,
8901 },
8902 kda: KdaConfig {
8903 num_heads: kda_num_heads,
8904 head_dim: kda_head_dim,
8905 short_conv_kernel_size: conv_kernel,
8906 gate_lower_bound: -5.0,
8907 use_full_rank_gate: true,
8908 },
8909 }),
8910 rope_freqs: None,
8911 rope_attn_factor: 1.0,
8912 rope_dim: None,
8913 rope_dim_swa: None,
8914 rope_freqs_long: None,
8915 rope_freqs_short: None,
8916 rope_orig_ctx: None,
8917 rope_layout: frink_models::config::RopeLayout::Neox,
8918 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
8919 swa_layers: frink_models::swa_layers::SwaLayers::All,
8920 attn_logit_softcap: None,
8921 final_logit_softcap: None,
8922 embedding_scale: None,
8923 residual_scale: None,
8924 normed_residual_scale: None,
8925 clamp_kqv: None,
8926 attn_temperature: None,
8927 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
8928 block_sub_norms: false,
8929 parallel_residual: false,
8930 learned_positions: false,
8931 attn_value_scale: None,
8932 alibi_max_bias: None,
8933 layer_loops: None,
8934 skip_stream: false,
8935 parallel_ssm: false,
8936 swa_chunked: false,
8937 weightless_qk_norm: false,
8938 logit_multiplier: None,
8939 attention_scale: None,
8940 rope_theta_swa: None,
8941 ffn_activation: frink_models::config::FfnActivation::Swiglu,
8942 best_effort_fields: &["synthetic test config, not a real preset"],
8943 };
8944 let hp = KimiRealHparams {
8945 hidden_dim,
8946 kda_num_heads,
8947 kda_head_dim,
8948 mla_num_heads,
8949 mla_q_lora_rank,
8950 mla_kv_lora_rank,
8951 mla_qk_nope_head_dim,
8952 mla_qk_rope_head_dim,
8953 mla_v_head_dim,
8954 dense_intermediate_dim: dense_intermediate,
8955 moe_hidden_dim: hidden_dim,
8956 moe_intermediate_dim: 4,
8957 n_experts: 1,
8958 num_shared_experts: 0,
8959 };
8960
8961 let prefix = "language_model.model.layers.0";
8965 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
8966 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
8967 name: String,
8968 shape: Vec<usize>,
8969 n: usize| {
8970 tensors.push((name, shape, vec![0.01f32; n]));
8971 };
8972 push(
8973 &mut tensors,
8974 format!("{prefix}.input_layernorm.weight"),
8975 vec![hidden_dim],
8976 hidden_dim,
8977 );
8978 push(
8979 &mut tensors,
8980 format!("{prefix}.post_attention_layernorm.weight"),
8981 vec![hidden_dim],
8982 hidden_dim,
8983 );
8984 push(
8985 &mut tensors,
8986 format!("{prefix}.self_attention_res_norm.weight"),
8987 vec![hidden_dim],
8988 hidden_dim,
8989 );
8990 push(
8991 &mut tensors,
8992 format!("{prefix}.self_attention_res_proj.weight"),
8993 vec![1, hidden_dim],
8994 hidden_dim,
8995 );
8996 push(
8997 &mut tensors,
8998 format!("{prefix}.mlp_res_norm.weight"),
8999 vec![hidden_dim],
9000 hidden_dim,
9001 );
9002 push(
9003 &mut tensors,
9004 format!("{prefix}.mlp_res_proj.weight"),
9005 vec![1, hidden_dim],
9006 hidden_dim,
9007 );
9008 push(
9009 &mut tensors,
9010 format!("{prefix}.self_attn.q_proj.weight"),
9011 vec![kda_proj, hidden_dim],
9012 kda_proj * hidden_dim,
9013 );
9014 push(
9015 &mut tensors,
9016 format!("{prefix}.self_attn.k_proj.weight"),
9017 vec![kda_proj, hidden_dim],
9018 kda_proj * hidden_dim,
9019 );
9020 push(
9021 &mut tensors,
9022 format!("{prefix}.self_attn.v_proj.weight"),
9023 vec![kda_proj, hidden_dim],
9024 kda_proj * hidden_dim,
9025 );
9026 push(
9027 &mut tensors,
9028 format!("{prefix}.self_attn.q_conv1d.weight"),
9029 vec![kda_proj, 1, conv_kernel],
9030 kda_proj * conv_kernel,
9031 );
9032 push(
9033 &mut tensors,
9034 format!("{prefix}.self_attn.k_conv1d.weight"),
9035 vec![kda_proj, 1, conv_kernel],
9036 kda_proj * conv_kernel,
9037 );
9038 push(
9039 &mut tensors,
9040 format!("{prefix}.self_attn.v_conv1d.weight"),
9041 vec![kda_proj, 1, conv_kernel],
9042 kda_proj * conv_kernel,
9043 );
9044 push(
9045 &mut tensors,
9046 format!("{prefix}.self_attn.A_log"),
9047 vec![kda_num_heads],
9048 kda_num_heads,
9049 );
9050 push(
9051 &mut tensors,
9052 format!("{prefix}.self_attn.f_a_proj.weight"),
9053 vec![kda_head_dim, hidden_dim],
9054 kda_head_dim * hidden_dim,
9055 );
9056 push(
9057 &mut tensors,
9058 format!("{prefix}.self_attn.f_b_proj.weight"),
9059 vec![kda_proj, kda_head_dim],
9060 kda_proj * kda_head_dim,
9061 );
9062 push(
9063 &mut tensors,
9064 format!("{prefix}.self_attn.dt_bias"),
9065 vec![kda_proj],
9066 kda_proj,
9067 );
9068 push(
9069 &mut tensors,
9070 format!("{prefix}.self_attn.b_proj.weight"),
9071 vec![kda_num_heads, hidden_dim],
9072 kda_num_heads * hidden_dim,
9073 );
9074 push(
9075 &mut tensors,
9076 format!("{prefix}.self_attn.g_proj.weight"),
9077 vec![kda_proj, hidden_dim],
9078 kda_proj * hidden_dim,
9079 );
9080 push(
9081 &mut tensors,
9082 format!("{prefix}.self_attn.o_norm.weight"),
9083 vec![kda_head_dim],
9084 kda_head_dim,
9085 );
9086 push(
9087 &mut tensors,
9088 format!("{prefix}.self_attn.o_proj.weight"),
9089 vec![hidden_dim, kda_proj],
9090 hidden_dim * kda_proj,
9091 );
9092 push(
9093 &mut tensors,
9094 format!("{prefix}.mlp.gate_proj.weight"),
9095 vec![dense_intermediate, hidden_dim],
9096 dense_intermediate * hidden_dim,
9097 );
9098 push(
9099 &mut tensors,
9100 format!("{prefix}.mlp.up_proj.weight"),
9101 vec![dense_intermediate, hidden_dim],
9102 dense_intermediate * hidden_dim,
9103 );
9104 push(
9105 &mut tensors,
9106 format!("{prefix}.mlp.down_proj.weight"),
9107 vec![hidden_dim, dense_intermediate],
9108 hidden_dim * dense_intermediate,
9109 );
9110 push(
9111 &mut tensors,
9112 "language_model.model.embed_tokens.weight".to_string(),
9113 vec![vocab_size, hidden_dim],
9114 vocab_size * hidden_dim,
9115 );
9116 push(
9117 &mut tensors,
9118 "language_model.lm_head.weight".to_string(),
9119 vec![vocab_size, hidden_dim],
9120 vocab_size * hidden_dim,
9121 );
9122 push(
9123 &mut tensors,
9124 "language_model.model.norm.weight".to_string(),
9125 vec![hidden_dim],
9126 hidden_dim,
9127 );
9128 push(
9129 &mut tensors,
9130 "language_model.model.output_attn_res_norm.weight".to_string(),
9131 vec![hidden_dim],
9132 hidden_dim,
9133 );
9134 push(
9135 &mut tensors,
9136 "language_model.model.output_attn_res_proj.weight".to_string(),
9137 vec![1, hidden_dim],
9138 hidden_dim,
9139 );
9140
9141 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9154 let dir = std::env::temp_dir().join(format!(
9155 "frink_server_kimi_e2e_test_{}_{}",
9156 std::process::id(),
9157 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9158 ));
9159 std::fs::create_dir_all(&dir).unwrap();
9160 let shard_bytes = write_safetensors_shard(&tensors);
9161 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9162 let map_entries: Vec<String> = tensors
9163 .iter()
9164 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9165 .collect();
9166 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9167 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9168
9169 use base64::Engine;
9173 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9174 .map(|b| {
9175 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9176 format!("{b64} {b}")
9177 })
9178 .collect();
9179 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9180
9181 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9182 .expect("must load the synthetic Kimi checkpoint end to end");
9183 std::fs::remove_dir_all(&dir).ok();
9184 loaded
9185 }
9186
9187 #[test]
9195 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9196 let loaded = build_synthetic_kimi_loaded();
9197 let state = build_app_state(
9198 StartupModels {
9199 loaded: model::LoadedModel::Kimi(loaded),
9200 embedding: None,
9201 },
9202 None,
9203 None,
9204 None,
9205 false,
9206 None,
9207 Arc::new(health::Detection::ready(health::probe_backends())),
9208 );
9209 let active = state.active().expect("a freshly built state has a model");
9210 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9211 assert!(!active.is_synthetic());
9212
9213 let (choices, _usage) = run_generation(
9214 active.generative().unwrap(),
9215 "hi",
9216 &greedy_params(5),
9217 None,
9218 None,
9219 None,
9220 None,
9221 None,
9222 None,
9223 )
9224 .expect("a real Kimi checkpoint must generate without error");
9225 assert!(matches!(
9226 choices[0].0,
9227 FinishReason::Length | FinishReason::Stop
9228 ));
9229 }
9230
9231 #[test]
9244 fn a_grammar_constrains_the_engine_decode_path() {
9245 let loaded = build_synthetic_kimi_loaded();
9246 let state = build_app_state(
9247 StartupModels {
9248 loaded: model::LoadedModel::Kimi(loaded),
9249 embedding: None,
9250 },
9251 None,
9252 None,
9253 None,
9254 false,
9255 None,
9256 Arc::new(health::Detection::ready(health::probe_backends())),
9257 );
9258 let active = state.active().expect("a freshly built state has a model");
9259
9260 let run = |grammar: Option<&str>| {
9261 let mut params = greedy_params(6);
9262 params.grammar = grammar.map(|src| {
9263 Arc::new(
9264 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9265 .expect("test grammar parses"),
9266 )
9267 });
9268 run_generation(
9269 active.generative().unwrap(),
9270 "hi",
9271 ¶ms,
9272 None,
9273 None,
9274 None,
9275 None,
9276 None,
9277 None,
9278 )
9279 };
9280
9281 let (choices, _) = run(None).expect("the unconstrained run must serve");
9282 let unconstrained = choices[0].1.clone();
9283 assert!(
9284 unconstrained.chars().any(|c| c != 'a'),
9285 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9286 constrained run below would prove nothing"
9287 );
9288
9289 let (choices, _) =
9290 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9291 let (finish, constrained) = choices.into_iter().next().unwrap();
9292 assert!(
9293 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9294 "the engine decode path served text its grammar forbids ({constrained:?}): \
9295 the constraint was dropped between `generate_engine` and the sampler"
9296 );
9297 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9298 }
9299
9300 #[test]
9310 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9311 let loaded = build_synthetic_kimi_loaded();
9312 let state = build_app_state(
9313 StartupModels {
9314 loaded: model::LoadedModel::Kimi(loaded),
9315 embedding: None,
9316 },
9317 None,
9318 None,
9319 None,
9320 false,
9321 None,
9322 Arc::new(health::Detection::ready(health::probe_backends())),
9323 );
9324
9325 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9326 let kv_pool_config = generate::KvPoolConfig {
9327 pool: pool.clone(),
9328 queue_wait: Duration::ZERO,
9329 };
9330 let pc = Mutex::new(PrefixCache::new(4));
9331
9332 run_generation(
9333 state
9334 .active()
9335 .expect("a freshly built state has a model")
9336 .generative()
9337 .unwrap(),
9338 "hi",
9339 &greedy_params(5),
9340 Some(&kv_pool_config),
9341 None,
9342 Some(&pc),
9343 None,
9344 None,
9345 None,
9346 )
9347 .expect("a real Kimi checkpoint must generate without error");
9348
9349 assert_eq!(
9350 pool.lock().unwrap().free_blocks(),
9351 4,
9352 "the KV pool must be completely untouched by a Kimi request"
9353 );
9354 let stats = pc.lock().unwrap().stats();
9355 assert_eq!(
9356 stats.hits + stats.misses,
9357 0,
9358 "the prefix cache must never be consulted for a Kimi request"
9359 );
9360 }
9361}