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 logprobs;
57mod lora;
58mod mcp;
59mod model;
60mod openai_extra;
61mod output;
62mod policy;
63mod prefill_batch;
64mod reasoning_budget;
65mod reasoning_tokens;
66mod request_tail;
67mod rerank;
68mod response_cache;
69pub(crate) mod responses;
70mod resume;
71mod sample_step;
72mod sampling_knobs;
73mod sampling_loop;
74mod security;
75mod serving;
76mod session;
77mod slots;
78mod sse;
79mod stats;
80mod stop;
81mod stream_events;
82mod tasks;
83mod tool_grammar;
84mod unimplemented_fields;
85mod unsupported_sampling;
86mod utf8_stream;
87
88use std::cell::RefCell;
89use std::convert::Infallible;
90use std::net::SocketAddr;
91use std::path::PathBuf;
92use std::rc::Rc;
93use std::sync::{Arc, Mutex, MutexGuard};
94use std::time::Duration;
95
96use axum::{
97 extract::State,
98 http::StatusCode,
99 response::sse::{Event, Sse},
100 response::{IntoResponse, Response},
101 routing::{get, post},
102 Json, Router,
103};
104use serde::{Deserialize, Serialize};
105
106use cli::apply_cli_overrides;
107pub use cli::{ServerArgs, BUILT_WITH_CUDA, BUILT_WITH_METAL};
108
109use frink_core::cache::KvBlockPool;
110use frink_models::kimi_tokenizer::KimiTokenizer;
111use frink_models::sampling::SamplingParams;
112use frink_models::tokenizer::{SpecialTokens, StopTokens};
113use frink_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
114#[cfg(test)]
115use generate::FinishReason;
116use generate::GenerationParams;
117pub(crate) use loaded::{ActiveModel, Loaded};
118use model::ServerTokenizer;
119use rerank::encoder_endpoints;
120use response_cache::ResponseCache;
121use sampling_knobs::SamplingKnobs;
122
123#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
132 Gguf(GgufModel),
133 Kimi(KimiModel),
134 Mla(MlaModel),
135 Gemma4(Gemma4Model),
136 Glm52(Glm52Model),
137}
138
139pub(crate) struct GgufModel {
140 decoder: Arc<Decoder>,
141 tokenizer: Arc<ServerTokenizer>,
142 stop_tokens: StopTokens,
143 bos_id: Option<usize>,
144 is_synthetic: bool,
145 chat_template: chat_template::PromptTemplate,
146}
147
148pub(crate) struct KimiModel {
149 engine: KimiEngine,
150 tokenizer: KimiTokenizer,
151 stop_tokens: StopTokens,
152 chat_template: chat_template::PromptTemplate,
153}
154
155pub(crate) struct MlaModel {
156 engine: MlaEngine,
157 tokenizer: ServerTokenizer,
158 stop_tokens: StopTokens,
159 bos_id: Option<usize>,
160 name: String,
161 chat_template: chat_template::PromptTemplate,
162}
163
164pub(crate) struct Gemma4Model {
165 engine: Gemma4Engine,
166 tokenizer: ServerTokenizer,
167 stop_tokens: StopTokens,
168 bos_id: Option<usize>,
169 name: String,
170 chat_template: chat_template::PromptTemplate,
171}
172
173pub(crate) struct Glm52Model {
174 engine: frink_models::Glm52Engine,
175 tokenizer: ServerTokenizer,
176 stop_tokens: StopTokens,
177 bos_id: Option<usize>,
178 name: String,
179 chat_template: chat_template::PromptTemplate,
180}
181
182impl Model {
183 pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
184 match self {
185 Model::Gguf(m) => m.chat_template.clone(),
186 Model::Kimi(m) => m.chat_template.clone(),
187 Model::Mla(m) => m.chat_template.clone(),
188 Model::Gemma4(m) => m.chat_template.clone(),
189 Model::Glm52(m) => m.chat_template.clone(),
190 }
191 }
192
193 fn is_synthetic(&self) -> bool {
198 match self {
199 Model::Gguf(m) => m.is_synthetic,
200 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
201 }
202 }
203
204 fn tokenizer_kind(&self) -> &'static str {
205 match self {
206 Model::Gguf(m) => m.tokenizer.kind(),
207 Model::Kimi(_) => "kimi-tiktoken-bpe",
208 Model::Mla(m) => m.tokenizer.kind(),
209 Model::Gemma4(m) => m.tokenizer.kind(),
210 Model::Glm52(m) => m.tokenizer.kind(),
211 }
212 }
213
214 fn expert_store_stats(&self) -> Option<frink_core::expert_store::ExpertStoreStats> {
218 match self {
219 Model::Gguf(m) => m.decoder.expert_store_stats(),
220 Model::Kimi(m) => m.engine.weights.expert_store_stats(),
221 Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
222 }
223 }
224
225 pub(crate) fn name(&self) -> &str {
226 match self {
227 Model::Gguf(m) => m.decoder.config.name,
228 Model::Kimi(_) => "kimi-k3",
229 Model::Mla(m) => m.name.as_str(),
230 Model::Gemma4(m) => m.name.as_str(),
231 Model::Glm52(m) => m.name.as_str(),
232 }
233 }
234
235 pub(crate) fn encode(&self, text: &str, specials: SpecialTokens) -> Vec<usize> {
262 match self {
263 Model::Gguf(m) => m.tokenizer.encode(text, specials),
264 Model::Kimi(m) => m
265 .tokenizer
266 .encode(text, specials)
267 .into_iter()
268 .map(|id| id as usize)
269 .collect(),
270 Model::Mla(m) => m.tokenizer.encode(text, specials),
271 Model::Gemma4(m) => m.tokenizer.encode(text, specials),
272 Model::Glm52(m) => m.tokenizer.encode(text, specials),
273 }
274 }
275
276 pub(crate) fn bos_id(&self) -> Option<usize> {
285 match self {
286 Model::Gguf(m) => m.bos_id,
287 Model::Kimi(_) => None,
288 Model::Mla(m) => m.bos_id,
289 Model::Gemma4(m) => m.bos_id,
290 Model::Glm52(m) => m.bos_id,
291 }
292 }
293
294 pub(crate) fn decode(&self, ids: &[usize]) -> String {
295 match self {
296 Model::Gguf(m) => m.tokenizer.decode(ids),
297 Model::Kimi(m) => {
298 let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
299 m.tokenizer.decode(&ids32)
300 }
301 Model::Mla(m) => m.tokenizer.decode(ids),
302 Model::Gemma4(m) => m.tokenizer.decode(ids),
303 Model::Glm52(m) => m.tokenizer.decode(ids),
304 }
305 }
306
307 pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
310 match self {
311 Model::Gguf(m) => {
312 let mut caches: Vec<_> = m.decoder.config.new_kv_caches();
313 Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
314 }
315 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
316 }
317 }
318
319 pub(crate) fn gguf_decoder(&self) -> Option<&Arc<Decoder>> {
325 match self {
326 Model::Gguf(m) => Some(&m.decoder),
327 Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
328 }
329 }
330
331 pub(crate) fn vocab_size(&self) -> Option<usize> {
332 match self {
333 Model::Gguf(m) => Some(m.decoder.config.vocab_size),
334 Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
335 Model::Mla(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
336 Model::Gemma4(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
337 Model::Glm52(m) => Some(frink_models::Engine::vocab_size(&m.engine)),
338 }
339 }
340
341 fn has_real_vocabulary(&self) -> bool {
348 match self {
349 Model::Gguf(m) => !matches!(*m.tokenizer, model::ServerTokenizer::Byte),
350 Model::Kimi(_) => true,
351 Model::Mla(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
352 Model::Gemma4(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
353 Model::Glm52(m) => !matches!(m.tokenizer, model::ServerTokenizer::Byte),
354 }
355 }
356}
357
358impl frink_models::dry::DryVocab for Model {
364 fn n_tokens(&self) -> usize {
365 self.vocab_size().unwrap_or(0)
366 }
367
368 fn detokenize(&self, token: usize) -> String {
369 self.decode(&[token])
370 }
371
372 fn tokenize(&self, text: &str) -> Vec<usize> {
373 self.encode(text, SpecialTokens::AsText)
374 }
375}
376
377pub(crate) struct AppState {
378 pub(crate) embedding: Option<Arc<frink_models::EmbeddingModel>>,
392 active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
411 pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
415 pub(crate) tasks: Arc<tasks::TaskRegistry>,
417 pub(crate) cancels: Arc<cancel::CancelRegistry>,
421 pub(crate) stats: stats::Stats,
424 pub(crate) streams: resume::StreamRegistry,
427 pub(crate) model_dir: Option<PathBuf>,
429 response_cache: Mutex<ResponseCache>,
433 pub(crate) kv_pool: Option<generate::KvPoolConfig>,
445 pub(crate) paged_kv: Option<generate::PagedKvConfig>,
458 pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
468 sessions: session::SessionStore,
474 requests_total: std::sync::atomic::AtomicU64,
475 request_errors_total: std::sync::atomic::AtomicU64,
476 started_at: std::time::Instant,
477 last_request_ms: std::sync::atomic::AtomicU64,
482 detection: Arc<health::Detection>,
484 mcp: Option<mcp::LoadedMcpConfig>,
486 pub(crate) continuous_batching_enabled: bool,
490 pub(crate) metal_private_decode_gate: Option<Arc<std::sync::Mutex<()>>>,
494 loading_model: Mutex<Option<String>>,
498 last_load_error: Mutex<Option<(String, String)>>,
502 pub(crate) serving: Mutex<crate::stats::ServingStats>,
507 pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
512 pub(crate) footprint:
518 Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
519 pub(crate) started_unix: u64,
529}
530
531pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
537
538impl AppState {
539 pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
543 self.active
544 .read()
545 .unwrap_or_else(|p| p.into_inner())
546 .clone()
547 }
548
549 pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
555 self.active().ok_or_else(|| {
556 (
557 StatusCode::SERVICE_UNAVAILABLE,
558 Json(serde_json::json!({"error": {
559 "message": "no model is loaded; POST /admin/models/load with an id from \
560 GET /admin/models",
561 "type": "model_not_loaded"
562 }})),
563 )
564 })
565 }
566
567 pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
575 Ok(Arc::clone(self.require_active()?.generative()?))
576 }
577
578 pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
586 let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
587 std::mem::replace(&mut *guard, next)
588 }
589
590 fn mark_request_finished(&self) {
594 let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
595 self.last_request_ms
596 .store(ms, std::sync::atomic::Ordering::Relaxed);
597 }
598
599 pub(crate) fn uptime(&self) -> Duration {
600 self.started_at.elapsed()
601 }
602
603 pub(crate) fn requests_total(&self) -> u64 {
604 self.requests_total
605 .load(std::sync::atomic::Ordering::Relaxed)
606 }
607
608 pub(crate) fn errors_total(&self) -> u64 {
609 self.request_errors_total
610 .load(std::sync::atomic::Ordering::Relaxed)
611 }
612
613 pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
614 lock_cache(&self.response_cache).stats()
615 }
616
617 pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
620 let last = self
621 .last_request_ms
622 .load(std::sync::atomic::Ordering::Relaxed);
623 (last > 0)
624 .then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
625 .map(|age| age.max(0.0))
626 }
627
628 pub(crate) fn loading_model_id(&self) -> Option<String> {
629 self.loading_model
630 .lock()
631 .unwrap_or_else(|p| p.into_inner())
632 .clone()
633 }
634
635 pub(crate) fn set_loading_model(&self, id: Option<String>) {
636 *self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
637 }
638
639 pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
640 self.last_load_error
641 .lock()
642 .unwrap_or_else(|p| p.into_inner())
643 .clone()
644 }
645
646 pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
647 *self
648 .last_load_error
649 .lock()
650 .unwrap_or_else(|p| p.into_inner()) = error;
651 }
652
653 pub(crate) fn active_model_name(&self) -> Option<String> {
662 self.active().map(|a| a.name().to_string())
663 }
664
665 pub(crate) fn embedding_model(&self) -> Option<Arc<frink_models::EmbeddingModel>> {
675 self.embedding
676 .clone()
677 .or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
678 }
679
680 pub(crate) fn embedding_model_name(&self) -> Option<String> {
684 match self.embedding_model() {
685 Some(e) => Some(e.name().to_string()),
686 None => self.active_model_name(),
687 }
688 }
689
690 pub(crate) fn record_request(&self, record: stats::Record<'_>) {
691 self.stats.record(stats::entry(record));
692 }
693}
694
695fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
700 cache
701 .lock()
702 .unwrap_or_else(|poisoned| poisoned.into_inner())
703}
704
705#[derive(Debug, Clone, Deserialize)]
706#[serde(untagged)]
707pub(crate) enum MessageContent {
708 Text(String),
709 Parts(Vec<ContentPart>),
710}
711
712#[derive(Debug, Clone, Deserialize)]
713struct ContentPart {
714 #[serde(rename = "type")]
715 kind: String,
716 #[serde(default)]
717 text: Option<String>,
718 #[serde(default)]
719 image_url: Option<serde_json::Value>,
720}
721
722impl MessageContent {
723 fn as_text(&self) -> String {
724 match self {
725 Self::Text(s) => s.clone(),
726 Self::Parts(parts) => parts
727 .iter()
728 .filter_map(|p| p.text.as_deref())
729 .collect::<Vec<_>>()
730 .join(""),
731 }
732 }
733
734 fn has_image(&self) -> bool {
735 match self {
736 Self::Text(_) => false,
737 Self::Parts(parts) => parts
738 .iter()
739 .any(|p| p.kind == "image_url" || p.image_url.is_some()),
740 }
741 }
742}
743
744#[derive(Debug, Clone, Deserialize)]
745pub(crate) struct ChatMessage {
746 pub(crate) role: String,
747 #[serde(default)]
751 pub(crate) content: Option<MessageContent>,
752 #[serde(default)]
756 pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
757 #[serde(default)]
762 #[allow(dead_code)]
763 pub(crate) tool_call_id: Option<String>,
764 #[serde(default, alias = "reasoning")]
780 pub(crate) reasoning_content: Option<String>,
781}
782
783impl ChatMessage {
784 fn rendered_content(&self) -> String {
791 let mut out = self
792 .content
793 .as_ref()
794 .map(MessageContent::as_text)
795 .unwrap_or_default();
796 if let Some(calls) = &self.tool_calls {
797 for call in calls {
798 out.push_str(&format!(
799 "<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
800 call.function.name, call.function.arguments
801 ));
802 }
803 }
804 out
805 }
806}
807
808#[derive(Debug, Clone, Deserialize)]
809pub(crate) struct ToolCallIn {
810 #[serde(default)]
811 #[allow(dead_code)]
812 id: String,
813 #[serde(rename = "type", default)]
814 #[allow(dead_code)]
815 kind: String,
816 function: ToolCallFunctionIn,
817}
818
819#[derive(Debug, Clone, Deserialize)]
820struct ToolCallFunctionIn {
821 name: String,
822 arguments: String,
827}
828
829#[derive(Debug, Clone, Deserialize)]
832struct ToolDef {
833 #[serde(rename = "type", default)]
834 #[allow(dead_code)]
835 kind: String,
836 function: ToolFunctionDef,
837}
838
839#[derive(Debug, Clone, Deserialize)]
840struct ToolFunctionDef {
841 name: String,
842 #[serde(default)]
843 description: Option<String>,
844 #[serde(default)]
845 parameters: Option<serde_json::Value>,
846}
847
848#[derive(Debug, Clone, Deserialize)]
859#[serde(untagged)]
860enum ToolChoice {
861 Mode(String),
862 Specific(serde_json::Value),
863}
864
865#[derive(Deserialize)]
868#[serde(untagged)]
869enum StopParam {
870 One(String),
871 Many(Vec<String>),
872}
873
874#[derive(Deserialize)]
875struct ChatCompletionRequest {
876 model: String,
877 messages: Vec<ChatMessage>,
878 #[serde(default = "default_max_tokens")]
879 max_tokens: usize,
880 #[serde(default)]
881 temperature: Option<f32>,
882 #[serde(default)]
883 top_p: Option<f32>,
884 #[serde(default)]
889 min_p: Option<f32>,
890 #[serde(default)]
891 top_k: Option<usize>,
892 #[serde(default)]
893 repetition_penalty: Option<f32>,
894 #[serde(flatten)]
898 extra_samplers: crate::sampling_knobs::ExtraSamplerFields,
899 #[serde(flatten)]
903 unimplemented: crate::unimplemented_fields::UnimplementedFields,
904 #[serde(default)]
905 seed: Option<u64>,
906 #[serde(default)]
907 stop: Option<StopParam>,
908 #[serde(default)]
909 stream: Option<bool>,
910 #[serde(default)]
923 stream_resumable: Option<bool>,
924 #[serde(default)]
935 ignore_eos: Option<bool>,
936 #[serde(default)]
937 tools: Vec<ToolDef>,
938 #[serde(default)]
939 tool_choice: Option<ToolChoice>,
940 #[serde(default)]
948 chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
949 #[serde(default)]
957 reasoning_effort: Option<String>,
958 #[serde(default)]
962 thinking: Option<ThinkingSwitch>,
963 #[serde(default)]
968 session_id: Option<String>,
969 #[serde(default, deserialize_with = "continuation::deserialize")]
977 continue_final_message: continuation::ContinueFinalMessage,
978 #[serde(default, alias = "thinking_budget_tokens")]
986 reasoning_budget_tokens: Option<reasoning_budget::BudgetTokens>,
987 #[serde(default)]
989 logprobs: Option<bool>,
990 #[serde(default)]
991 top_logprobs: Option<u32>,
992 #[serde(default)]
993 presence_penalty: Option<f32>,
994 #[serde(default)]
995 frequency_penalty: Option<f32>,
996 #[serde(default)]
997 response_format: Option<serde_json::Value>,
998 #[serde(default)]
1005 logit_bias: Option<serde_json::Value>,
1006 #[serde(default)]
1010 lora: Option<Vec<frink_api::LoraScaleRequest>>,
1011 #[serde(default)]
1021 samplers: Option<serde_json::Value>,
1022 #[serde(default)]
1032 grammar: Option<String>,
1033}
1034
1035const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
1047
1048#[derive(Debug, Clone, Deserialize)]
1050pub(crate) struct ThinkingSwitch {
1051 #[serde(rename = "type")]
1052 pub(crate) kind: String,
1053}
1054
1055const THINKING_KWARG_KEYS: [&str; 4] = [
1059 "enable_thinking",
1060 "thinking",
1061 "thinking_mode",
1062 "reasoning_effort",
1063];
1064
1065const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
1069
1070fn default_max_tokens() -> usize {
1071 DEFAULT_CHAT_MAX_TOKENS
1072}
1073
1074impl ChatCompletionRequest {
1075 fn sampling_knobs(&self) -> Result<SamplingKnobs, ApiError> {
1084 let mut knobs = SamplingKnobs {
1085 temperature: self.temperature,
1086 top_p: self.top_p,
1087 min_p: self.min_p,
1088 top_k: self.top_k,
1089 repetition_penalty: self.repetition_penalty,
1090 presence_penalty: self.presence_penalty,
1091 frequency_penalty: self.frequency_penalty,
1092 penalty_last_n: None,
1096 sampler_order: unsupported_sampling::parse_sampler_order(
1097 self.samplers.as_ref(),
1098 "/v1/chat/completions",
1099 )?,
1100 ..SamplingKnobs::default()
1101 };
1102 self.extra_samplers.apply(&mut knobs);
1103 Ok(knobs)
1104 }
1105
1106 fn sampling_params(
1107 &self,
1108 model: crate::sampling_knobs::SamplerModel<'_>,
1109 ) -> Result<SamplingParams, ApiError> {
1110 self.sampling_knobs()?.resolve(model).map_err(|e| {
1111 unsupported_feature(&format!("`dry_multiplier` on /v1/chat/completions: {e}"))
1112 })
1113 }
1114
1115 fn stop_sequences(&self) -> Vec<String> {
1116 self.stop
1117 .as_ref()
1118 .map(|s| match s {
1119 StopParam::One(v) => vec![v.clone()],
1120 StopParam::Many(v) => v.clone(),
1121 })
1122 .unwrap_or_default()
1123 }
1124
1125 fn n_logprobs(&self) -> Result<Option<usize>, ApiError> {
1139 const MAX: u32 = 20;
1140 match (self.logprobs, self.top_logprobs) {
1141 (Some(true), Some(n)) if n > MAX => Err(invalid_request(
1142 &format!(
1143 "`top_logprobs` is {n}; this server reports at most {MAX} alternatives per \
1144 position, as upstream does"
1145 ),
1146 "top_logprobs",
1147 )),
1148 (Some(true), Some(n)) => Ok(Some(n as usize)),
1149 (Some(true), None) => Ok(Some(0)),
1153 (_, Some(_)) => Err(invalid_request(
1154 "`top_logprobs` requires `logprobs: true`",
1155 "top_logprobs",
1156 )),
1157 _ => Ok(None),
1158 }
1159 }
1160
1161 fn several_choices(&self) -> bool {
1166 self.unimplemented.n.is_some_and(|n| n > 1)
1167 }
1168
1169 fn tools_active(&self) -> bool {
1170 !self.tools.is_empty()
1171 && !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
1172 }
1173
1174 fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
1186 match &self.tool_choice {
1187 Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
1188 Some(ToolChoice::Specific(value)) => {
1189 let name = value
1193 .get("function")
1194 .and_then(|f| f.get("name"))
1195 .or_else(|| value.get("name"))
1196 .and_then(|n| n.as_str());
1197 match name {
1198 Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
1199 None => Err(invalid_request(
1200 "tool_choice must be \"auto\", \"none\", \"required\", or an object with \
1201 function.name",
1202 "tool_choice",
1203 )),
1204 }
1205 }
1206 _ => Ok(None),
1207 }
1208 }
1209
1210 fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
1212 self.tools
1213 .iter()
1214 .map(|t| tool_grammar::ToolSpec {
1215 name: &t.function.name,
1216 parameters: t.function.parameters.as_ref(),
1217 })
1218 .collect()
1219 }
1220
1221 fn thinking_direction(&self) -> Option<bool> {
1264 if let Some(switch) = &self.thinking {
1265 return match switch.kind.trim().to_ascii_lowercase().as_str() {
1266 "disabled" => Some(false),
1267 "enabled" => Some(true),
1268 _ => None,
1271 };
1272 }
1273 let effort = self.reasoning_effort.as_ref()?;
1274 DISABLE_EFFORTS
1275 .contains(&effort.trim().to_ascii_lowercase().as_str())
1276 .then_some(false)
1277 }
1278
1279 fn resolve_template_kwargs(
1280 &self,
1281 template: &chat_template::PromptTemplate,
1282 ) -> serde_json::Map<String, serde_json::Value> {
1283 let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
1284 let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
1288
1289 if !caller_steered {
1290 match self.thinking_direction() {
1291 Some(false) => {
1292 for (k, v) in crate::policy::effort::thinking_off_kwargs() {
1293 kwargs.insert(k, v);
1294 }
1295 return kwargs;
1298 }
1299 Some(true) => {
1300 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1301 kwargs.insert(k, v);
1302 }
1303 }
1304 None => {}
1305 }
1306 if let Some(effort) = &self.reasoning_effort {
1307 kwargs
1308 .entry("reasoning_effort".to_string())
1309 .or_insert_with(|| serde_json::json!(effort));
1310 }
1311 }
1312
1313 let offered: Vec<serde_json::Value> = if self.tools_active() {
1314 self.tools.iter().map(chat_template::tool_json).collect()
1315 } else {
1316 Vec::new()
1317 };
1318 let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
1319 if thinking == crate::policy::effort::ThinkingMode::Thinking {
1320 for (k, v) in crate::policy::effort::thinking_on_kwargs() {
1321 kwargs.entry(k).or_insert(v);
1322 }
1323 }
1324 match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
1325 crate::policy::effort::EffortMapping::Mapped(to) => {
1326 tracing::debug!("reasoning_effort quantized to {}", to.as_str());
1327 }
1328 crate::policy::effort::EffortMapping::Dropped => {
1329 tracing::debug!(
1330 "reasoning_effort dropped: this checkpoint's template grades no gear close \
1331 enough, so its own default applies"
1332 );
1333 }
1334 crate::policy::effort::EffortMapping::Unchanged => {}
1335 }
1336 crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
1337 kwargs
1338 }
1339
1340 fn validate_supported_fields(&self) -> Result<(), ApiError> {
1343 if self.max_tokens == 0 {
1350 return Err(invalid_request(
1351 "max_tokens must be at least 1",
1352 "max_tokens",
1353 ));
1354 }
1355 if let Some(switch) = &self.thinking {
1360 let kind = switch.kind.trim().to_ascii_lowercase();
1361 if kind != "enabled" && kind != "disabled" {
1362 return Err(invalid_request(
1363 "thinking.type must be \"enabled\" or \"disabled\"",
1364 "thinking.type",
1365 ));
1366 }
1367 }
1368 for msg in &self.messages {
1369 if msg.content.as_ref().is_some_and(MessageContent::has_image) {
1370 return Err(unsupported_feature(
1371 "image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
1372 ));
1373 }
1374 }
1375 self.n_logprobs()?;
1379 self.unimplemented.refuse("/v1/chat/completions")?;
1383 unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
1384 unsupported_sampling::parse_sampler_order(self.samplers.as_ref(), "/v1/chat/completions")?;
1388 let stated_grammar =
1403 grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
1404 if let Some(forced) = self.forced_tool_choice()? {
1412 if self.tools.is_empty() {
1413 return Err(invalid_request(
1414 "tool_choice forces a tool call, but no tools were offered",
1415 "tool_choice",
1416 ));
1417 }
1418 if let tool_grammar::Forced::Named(name) = forced {
1419 if !self.tools.iter().any(|t| t.function.name == name) {
1420 return Err(invalid_request(
1421 &format!(
1422 "tool_choice names {name:?}, which is not one of the tools offered"
1423 ),
1424 "tool_choice",
1425 ));
1426 }
1427 }
1428 if stated_grammar.is_some() {
1438 return Err(invalid_request(
1439 "a forced tool_choice and a \"grammar\" or response_format \"json_schema\" \
1440 are two different constraints on the same generation; send one",
1441 "tool_choice",
1442 ));
1443 }
1444 if self.json_object_mode() {
1445 return Err(invalid_request(
1446 "a forced tool_choice cannot be combined with response_format json_object: \
1447 the tool-call markers are not JSON",
1448 "tool_choice",
1449 ));
1450 }
1451 }
1452 Ok(())
1453 }
1454
1455 fn effective_stop_sequences(&self) -> Vec<String> {
1462 let mut stop = self.stop_sequences();
1463 if self.tools_active() {
1464 stop.push("</tool_call>".to_string());
1465 }
1466 stop
1467 }
1468
1469 fn json_object_mode(&self) -> bool {
1470 self.response_format
1471 .as_ref()
1472 .and_then(|v| v.get("type"))
1473 .and_then(|v| v.as_str())
1474 == Some("json_object")
1475 }
1476}
1477
1478#[derive(Serialize)]
1479struct ChatCompletionChoice {
1480 index: usize,
1481 message: ChatCompletionResponseMessage,
1482 finish_reason: &'static str,
1483 #[serde(skip_serializing_if = "Option::is_none")]
1488 logprobs: Option<serde_json::Value>,
1489}
1490
1491#[derive(Serialize)]
1492struct ChatCompletionResponseMessage {
1493 role: &'static str,
1494 #[serde(skip_serializing_if = "Option::is_none")]
1495 content: Option<String>,
1496 #[serde(skip_serializing_if = "Option::is_none")]
1500 reasoning_content: Option<String>,
1501 #[serde(skip_serializing_if = "Option::is_none")]
1502 tool_calls: Option<Vec<ToolCallOut>>,
1503}
1504
1505#[derive(Serialize, Clone)]
1506struct ToolCallOut {
1507 id: String,
1508 #[serde(rename = "type")]
1509 kind: &'static str,
1510 function: ToolCallFunctionOut,
1511}
1512
1513#[derive(Serialize, Clone)]
1521struct ToolCallDelta {
1522 index: usize,
1523 #[serde(skip_serializing_if = "Option::is_none")]
1524 id: Option<String>,
1525 #[serde(rename = "type", skip_serializing_if = "Option::is_none")]
1526 kind: Option<&'static str>,
1527 function: ToolCallFunctionDelta,
1528}
1529
1530#[derive(Serialize, Clone, Default)]
1531struct ToolCallFunctionDelta {
1532 #[serde(skip_serializing_if = "Option::is_none")]
1533 name: Option<String>,
1534 #[serde(skip_serializing_if = "Option::is_none")]
1537 arguments: Option<String>,
1538}
1539
1540impl ToolCallDelta {
1541 fn whole(index: usize, name: String, arguments: String) -> Self {
1543 ToolCallDelta {
1544 index,
1545 id: Some(format!("call_{index}")),
1546 kind: Some("function"),
1547 function: ToolCallFunctionDelta {
1548 name: Some(name),
1549 arguments: Some(arguments),
1550 },
1551 }
1552 }
1553
1554 fn opening(index: usize, name: String) -> Self {
1556 ToolCallDelta {
1557 index,
1558 id: Some(format!("call_{index}")),
1559 kind: Some("function"),
1560 function: ToolCallFunctionDelta {
1561 name: Some(name),
1562 arguments: Some(String::new()),
1563 },
1564 }
1565 }
1566
1567 fn arguments(index: usize, fragment: String) -> Self {
1569 ToolCallDelta {
1570 index,
1571 id: None,
1572 kind: None,
1573 function: ToolCallFunctionDelta {
1574 name: None,
1575 arguments: Some(fragment),
1576 },
1577 }
1578 }
1579}
1580
1581#[derive(Serialize, Clone)]
1582struct ToolCallFunctionOut {
1583 name: String,
1584 arguments: String,
1588}
1589
1590#[derive(Serialize)]
1591struct ChatCompletionResponse {
1592 id: String,
1593 request_id: String,
1599 object: &'static str,
1600 model: String,
1601 choices: Vec<ChatCompletionChoice>,
1602 usage: generate::Usage,
1607 frink_cache: &'static str,
1618}
1619
1620#[derive(Serialize)]
1621struct ChatCompletionChunkDelta {
1622 #[serde(skip_serializing_if = "Option::is_none")]
1623 role: Option<&'static str>,
1624 #[serde(skip_serializing_if = "Option::is_none")]
1625 content: Option<String>,
1626 #[serde(skip_serializing_if = "Option::is_none")]
1628 reasoning_content: Option<String>,
1629 #[serde(skip_serializing_if = "Option::is_none")]
1630 tool_calls: Option<Vec<ToolCallDelta>>,
1631}
1632
1633#[derive(Serialize)]
1634struct ChatCompletionChunkChoice {
1635 index: usize,
1636 delta: ChatCompletionChunkDelta,
1637 finish_reason: Option<&'static str>,
1638}
1639
1640#[derive(Serialize)]
1641struct ChatCompletionChunk {
1642 id: String,
1643 #[serde(skip_serializing_if = "Option::is_none")]
1650 request_id: Option<String>,
1651 object: &'static str,
1652 model: String,
1653 choices: Vec<ChatCompletionChunkChoice>,
1654 #[serde(skip_serializing_if = "Option::is_none")]
1657 usage: Option<generate::Usage>,
1658}
1659
1660async fn health(State(state): State<Arc<AppState>>) -> Response {
1666 let snapshot = state.detection.snapshot();
1667 let mut capabilities = snapshot.capabilities;
1668 let active = state.active();
1669
1670 capabilities.push(match active.as_deref() {
1673 None => frink_api::Capability::unavailable(
1678 frink_api::health::capability::REAL_WEIGHTS,
1679 frink_api::health::reason::MODEL_NOT_LOADED,
1680 "No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
1681 ),
1682 Some(active) if active.is_synthetic() => frink_api::Capability::unavailable(
1683 frink_api::health::capability::REAL_WEIGHTS,
1684 frink_api::health::reason::MODEL_NOT_LOADED,
1685 "Serving synthetic random weights: set FRINK_MODEL_PATH (or -m) to a real \
1686 checkpoint. Output from this model is noise.",
1687 ),
1688 Some(active) if active.encoder().is_some() => {
1698 let endpoints = active
1699 .encoder()
1700 .map(|e| encoder_endpoints(e))
1701 .unwrap_or_default();
1702 let served_by = match endpoints.is_empty() {
1703 true => "no endpoint in this build serves it".to_string(),
1704 false => format!("served by {}", endpoints.join(" and ")),
1705 };
1706 frink_api::Capability::available(
1707 frink_api::health::capability::REAL_WEIGHTS,
1708 format!(
1709 "Serving the real embedding checkpoint '{}'. This is an ENCODER, \
1710 {served_by}; generation endpoints refuse it.",
1711 active.name(),
1712 ),
1713 )
1714 }
1715 Some(active) => frink_api::Capability::available(
1716 frink_api::health::capability::REAL_WEIGHTS,
1717 format!("Serving the real checkpoint '{}'.", active.name()),
1718 ),
1719 });
1720 capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
1721 frink_api::Capability::available(
1722 frink_api::health::capability::CONTINUOUS_BATCHING,
1723 if state.continuous_batching_enabled && continuous_batching_env().is_none() {
1724 "On by default on Metal. Concurrent requests share one batched decode worker."
1725 } else {
1726 "Concurrent requests share one batched decode step."
1727 },
1728 )
1729 } else if state.metal_private_decode_gate.is_some() {
1730 frink_api::Capability::unavailable(
1731 frink_api::health::capability::CONTINUOUS_BATCHING,
1732 frink_api::health::reason::DISABLED,
1733 "Off; private Metal decodes serialize (one at a time). Set FRINK_CONTINUOUS_BATCHING=1 or --cont-batching for parallel serving.",
1734 )
1735 } else {
1736 frink_api::Capability::unavailable(
1737 frink_api::health::capability::CONTINUOUS_BATCHING,
1738 frink_api::health::reason::DISABLED,
1739 "Off; set FRINK_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
1740 )
1741 });
1742
1743 let last_request_ms = state
1744 .last_request_ms
1745 .load(std::sync::atomic::Ordering::Relaxed);
1746 let uptime = state.started_at.elapsed();
1747 let health_state = if active.is_none() {
1755 frink_api::HealthState::Unavailable
1756 } else {
1757 snapshot.state
1758 };
1759 let body = frink_api::HealthResponse {
1760 state: health_state,
1761 reason: match health_state {
1762 frink_api::HealthState::Ready => None,
1763 frink_api::HealthState::Unavailable => {
1764 Some(frink_api::health::reason::MODEL_NOT_LOADED.to_string())
1765 }
1766 frink_api::HealthState::Detecting => {
1767 Some(frink_api::health::reason::DETECTING.to_string())
1768 }
1769 },
1770 detail: match health_state {
1771 frink_api::HealthState::Ready => None,
1772 frink_api::HealthState::Unavailable => Some(
1773 "No model is loaded. POST /admin/models/load with an id from GET /admin/models."
1774 .to_string(),
1775 ),
1776 frink_api::HealthState::Detecting => {
1777 Some("Probing available compute backends.".to_string())
1778 }
1779 },
1780 model: active
1781 .as_deref()
1782 .map(|active| frink_api::health::ModelSummary {
1783 id: active.name().to_string(),
1784 tokenizer: active.tokenizer_kind().to_string(),
1785 synthetic_weights: active.is_synthetic(),
1786 }),
1787 capabilities,
1788 version: env!("CARGO_PKG_VERSION").to_string(),
1789 pid: std::process::id(),
1790 uptime_seconds: uptime.as_secs_f64(),
1791 server_time_unix_ms: std::time::SystemTime::now()
1792 .duration_since(std::time::UNIX_EPOCH)
1793 .map(|d| d.as_millis().min(u64::MAX as u128) as u64)
1794 .unwrap_or(0),
1795 last_request_age_seconds: (last_request_ms > 0)
1796 .then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
1797 .map(|age| age.max(0.0)),
1798 };
1799
1800 let status =
1801 StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1802 (status, Json(body)).into_response()
1803}
1804
1805async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1806 let Some(active) = state.active() else {
1810 return Json(serde_json::json!({ "object": "list", "data": [] }));
1811 };
1812 let mut model_entry = serde_json::json!({
1813 "id": active.name(),
1814 "object": "model",
1815 "frink_synthetic_weights": active.is_synthetic(),
1816 "frink_tokenizer": active.tokenizer_kind(),
1817 });
1818 if let Some(encoder) = active.encoder() {
1824 model_entry["frink_model_kind"] = serde_json::json!("embedding");
1825 model_entry["frink_endpoints"] = serde_json::json!(encoder_endpoints(encoder));
1826 model_entry["frink_n_embd"] = serde_json::json!(encoder.n_embd());
1827 model_entry["frink_pooling"] = serde_json::json!(encoder.pooling_type().name());
1828 model_entry["frink_context_length"] = serde_json::json!(encoder.n_ctx_train());
1829 }
1830 if let Some(model) = active.generative_opt() {
1838 let parser_configured = active.reasoning_format().is_some();
1839 let gears = model.chat_template().think_gears(parser_configured);
1840 if !gears.is_empty() {
1841 model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
1842 if let Some(default) = &gears.default {
1843 model_entry["default_reasoning_effort"] = serde_json::json!(default);
1844 }
1845 model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
1848 }
1849 }
1850 if let Some(mcp) = &state.mcp {
1851 model_entry["frink_mcp"] = mcp.models_metadata();
1852 }
1853 Json(serde_json::json!({
1854 "object": "list",
1855 "data": [model_entry]
1856 }))
1857}
1858
1859async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
1873 let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
1874 let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
1875 let active = state.active();
1876 Json(serde_json::json!({
1877 "model": active.as_ref().map(|a| a.name()),
1878 "state": state
1879 .maintenance
1880 .lock()
1881 .unwrap_or_else(|p| p.into_inner())
1882 .state()
1883 .as_str(),
1884 "uptime_s": state.uptime().as_secs(),
1885 "throughput": {
1886 "decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1887 "prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
1888 },
1889 "requests": {
1890 "active": state.cancels.live_count(),
1891 "completed": state.stats.recorded_total(),
1892 "p95_ms": state.stats.p95_duration_ms(),
1893 "ttft_mean_ms": state.stats.ttft_mean_ms(),
1894 "prompt_tokens_total": state.stats.tokens_prompt_total(),
1895 "completion_tokens_total": state.stats.tokens_generated_total(),
1896 },
1897 "pools": cache_admin::pool_gauges(&state),
1903 "memory": cache_admin::footprint_json(&state),
1906 }))
1907}
1908
1909#[derive(Deserialize)]
1910struct RequestsQuery {
1911 #[serde(default)]
1912 since: u64,
1913 #[serde(default = "default_requests_limit")]
1914 limit: usize,
1915}
1916
1917fn default_requests_limit() -> usize {
1918 stats::MAX_PAGE
1919}
1920
1921async fn recent_requests(
1929 State(state): State<Arc<AppState>>,
1930 axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
1931) -> Json<serde_json::Value> {
1932 let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
1933 Json(serde_json::json!({
1934 "requests": rows,
1935 "next_cursor": cursor,
1936 "missed": missed,
1937 "total": state.stats.recorded_total(),
1938 }))
1939}
1940
1941#[derive(Serialize)]
1942struct CombinedCacheStats {
1943 response_cache: response_cache::CacheStats,
1944 prefix_cache: Option<frink_models::PrefixCacheStats>,
1946}
1947
1948async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
1949 Json(CombinedCacheStats {
1950 response_cache: lock_cache(&state.response_cache).stats(),
1951 prefix_cache: state
1952 .prefix_cache
1953 .as_ref()
1954 .map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
1955 })
1956}
1957
1958async fn metrics(State(state): State<Arc<AppState>>) -> Response {
1963 use std::sync::atomic::Ordering;
1964
1965 let cache_stats = lock_cache(&state.response_cache).stats();
1966 let active = state.active();
1967 let requests_total = state.requests_total.load(Ordering::Relaxed);
1968 let errors_total = state.request_errors_total.load(Ordering::Relaxed);
1969 let uptime = state.started_at.elapsed().as_secs_f64();
1970
1971 let body = format!(
1972 "# HELP frink_requests_total Total chat completion requests received.\n\
1973 # TYPE frink_requests_total counter\n\
1974 frink_requests_total {requests_total}\n\
1975 # HELP frink_request_errors_total Total chat completion requests that returned an error.\n\
1976 # TYPE frink_request_errors_total counter\n\
1977 frink_request_errors_total {errors_total}\n\
1978 # HELP frink_cache_hits_total Whole-response cache hits.\n\
1979 # TYPE frink_cache_hits_total counter\n\
1980 frink_cache_hits_total {}\n\
1981 # HELP frink_cache_misses_total Whole-response cache misses.\n\
1982 # TYPE frink_cache_misses_total counter\n\
1983 frink_cache_misses_total {}\n\
1984 # HELP frink_cache_entries Current whole-response cache entry count.\n\
1985 # TYPE frink_cache_entries gauge\n\
1986 frink_cache_entries {}\n\
1987 # HELP frink_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
1988 # TYPE frink_synthetic_weights gauge\n\
1989 frink_synthetic_weights {}\n\
1990 # HELP frink_uptime_seconds Seconds since this server process started.\n\
1991 # TYPE frink_uptime_seconds gauge\n\
1992 frink_uptime_seconds {uptime}\n",
1993 cache_stats.hits,
1994 cache_stats.misses,
1995 cache_stats.entries,
1996 active
2000 .as_ref()
2001 .map(|a| a.is_synthetic() as u8)
2002 .unwrap_or(0),
2003 );
2004
2005 let body = match active
2009 .as_ref()
2010 .and_then(|a| a.expert_store_stats())
2011 {
2012 Some(es) => format!(
2013 "{body}\
2014 # HELP frink_expert_cache_hits_total Expert-store cache hits.\n\
2015 # TYPE frink_expert_cache_hits_total counter\n\
2016 frink_expert_cache_hits_total {}\n\
2017 # HELP frink_expert_cache_misses_total Expert-store cache misses (source reads).\n\
2018 # TYPE frink_expert_cache_misses_total counter\n\
2019 frink_expert_cache_misses_total {}\n\
2020 # HELP frink_expert_cache_evictions_total Expert-store LRU evictions.\n\
2021 # TYPE frink_expert_cache_evictions_total counter\n\
2022 frink_expert_cache_evictions_total {}\n\
2023 # HELP frink_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
2024 # TYPE frink_expert_cache_pass_throughs_total counter\n\
2025 frink_expert_cache_pass_throughs_total {}\n\
2026 # HELP frink_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
2027 # TYPE frink_expert_cache_bytes_read_total counter\n\
2028 frink_expert_cache_bytes_read_total {}\n\
2029 # HELP frink_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
2030 # TYPE frink_expert_cache_resident_bytes gauge\n\
2031 frink_expert_cache_resident_bytes {}\n",
2032 es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
2033 ),
2034 None => body,
2035 };
2036
2037 let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
2042 Some(batcher) => {
2043 let sched = batcher.stats();
2044 format!(
2045 "{body}\
2046 # HELP frink_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
2047 # TYPE frink_prefill_chunks_total counter\n\
2048 frink_prefill_chunks_total {}\n\
2049 # HELP frink_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
2050 # TYPE frink_prefill_tokens_total counter\n\
2051 frink_prefill_tokens_total {}\n\
2052 # HELP frink_decode_steps_total Batched decode steps the batch scheduler has run.\n\
2053 # TYPE frink_decode_steps_total counter\n\
2054 frink_decode_steps_total {}\n\
2055 # HELP frink_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
2056 # TYPE frink_scheduler_queue_depth gauge\n\
2057 frink_scheduler_queue_depth {}\n\
2058 # HELP frink_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
2059 # TYPE frink_scheduler_queue_rejected_total counter\n\
2060 frink_scheduler_queue_rejected_total {}\n\
2061 # HELP frink_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
2062 # TYPE frink_kv_blocks_total gauge\n\
2063 frink_kv_blocks_total {}\n\
2064 # HELP frink_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
2065 # TYPE frink_kv_blocks_free gauge\n\
2066 frink_kv_blocks_free {}\n\
2067 # HELP frink_kv_block_size Token positions per KV block.\n\
2068 # TYPE frink_kv_block_size gauge\n\
2069 frink_kv_block_size {}\n\
2070 # HELP frink_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
2071 # TYPE frink_kv_rejected_too_large_total counter\n\
2072 frink_kv_rejected_too_large_total {}\n\
2073 # HELP frink_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
2074 # TYPE frink_kv_rejected_context_length_total counter\n\
2075 frink_kv_rejected_context_length_total {}\n\
2076 # HELP frink_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
2077 # TYPE frink_scheduler_aborted_total counter\n\
2078 frink_scheduler_aborted_total {}\n\
2079 # HELP frink_scheduler_max_seqs Cap on in-flight sequences (-np / FRINK_CB_MAX_SEQS); 0 when unlimited.\n\
2080 # TYPE frink_scheduler_max_seqs gauge\n\
2081 frink_scheduler_max_seqs {}\n\
2082 # HELP frink_scheduler_prefill_chunk Prompt tokens per prefill chunk (-b / -ub / FRINK_CB_PREFILL_CHUNK).\n\
2083 # TYPE frink_scheduler_prefill_chunk gauge\n\
2084 frink_scheduler_prefill_chunk {}\n",
2085 sched.prefill_chunks,
2086 sched.prefill_tokens,
2087 sched.decode_steps,
2088 sched.queue_depth,
2089 sched.queue_rejected,
2090 sched.kv_blocks_total,
2091 sched.kv_blocks_free,
2092 sched.kv_block_size,
2093 sched.kv_rejected_too_large,
2094 sched.kv_rejected_context_length,
2095 sched.aborted,
2096 sched.max_seqs,
2097 sched.prefill_chunk,
2098 )
2099 }
2100 None => body,
2101 };
2102
2103 (
2104 [(
2105 axum::http::header::CONTENT_TYPE,
2106 "text/plain; version=0.0.4",
2107 )],
2108 body,
2109 )
2110 .into_response()
2111}
2112
2113pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
2114
2115pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
2121 (
2122 StatusCode::BAD_REQUEST,
2123 Json(serde_json::json!({"error": {
2124 "message": message,
2125 "type": "invalid_request_error",
2126 "param": param,
2127 "code": null,
2128 }})),
2129 )
2130}
2131
2132pub(crate) fn unsupported_feature(message: &str) -> ApiError {
2133 (
2134 StatusCode::NOT_IMPLEMENTED,
2135 Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
2136 )
2137}
2138
2139pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
2140 let status = match e {
2141 generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
2142 generate::DecodeError::Unsupported(_) => StatusCode::NOT_IMPLEMENTED,
2146 generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
2151 generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
2157 StatusCode::SERVICE_UNAVAILABLE
2158 }
2159 generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
2163 generate::DecodeError::ReasoningBudget { .. } => StatusCode::INTERNAL_SERVER_ERROR,
2167 };
2168 tracing::warn!("decode error: {e}");
2169 let mut body = serde_json::json!({"error": {"message": e.to_string()}});
2170 if let generate::DecodeError::KvBudgetExceeded {
2176 binding,
2177 estimated_bytes,
2178 limit_bytes,
2179 positions,
2180 positions_limit,
2181 ..
2182 } = &e
2183 {
2184 body["error"]["type"] = serde_json::json!("invalid_request_error");
2185 body["error"]["code"] = serde_json::json!(binding);
2186 body["error"]["binding"] = serde_json::json!(binding);
2187 body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
2188 body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
2189 body["error"]["positions"] = serde_json::json!(positions);
2190 body["error"]["positions_limit"] = serde_json::json!(positions_limit);
2191 }
2192 if let Some(secs) = e.retry_after_secs() {
2196 body["error"]["retry_after_seconds"] = serde_json::json!(secs);
2197 }
2198 (status, Json(body))
2199}
2200
2201pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
2202 tracing::error!("generation task panicked: {e}");
2203 (
2204 StatusCode::INTERNAL_SERVER_ERROR,
2205 Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
2206 )
2207}
2208
2209#[allow(clippy::too_many_arguments)] fn run_generation_emit(
2219 model: &Model,
2220 prompt: &str,
2221 params: &GenerationParams,
2222 kv_pool: Option<&generate::KvPoolConfig>,
2223 paged_kv: Option<&generate::PagedKvConfig>,
2224 prefix_cache: Option<&Mutex<PrefixCache>>,
2225 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2226 ceiling: Option<&budget::ContextCeiling>,
2227 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2228 mut emit: impl FnMut(&str),
2229 ) -> Result<(Vec<generate::GeneratedChoice>, generate::Usage), generate::DecodeError> {
2235 let synthetic = model.is_synthetic();
2236 let _lora_lease = lora::lease(model, params.lora.as_deref());
2241 let mut chunks: Vec<Vec<String>> = vec![Vec::new(); params.n.max(1)];
2242 let params = &{
2248 let mut resolved = params.clone();
2249 resolved.stop_token_ids = crate::stop::resolve_stop_tokens(&resolved.stop, |text| {
2250 model.encode(text, SpecialTokens::Parse)
2251 });
2252 resolved.reasoning_budget = resolved
2258 .reasoning_budget
2259 .armed(resolved.reasoning, prompt, |text| {
2260 model.encode(text, SpecialTokens::Parse)
2261 })
2262 .map_err(|detail| generate::DecodeError::ReasoningBudget { detail })?;
2263 resolved
2264 };
2265 let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
2266 let _metal_private_guard =
2267 acquire_metal_private_decode_gate(metal_private_decode_gate, used_batcher);
2268 let (finishes, usage) = match model {
2269 Model::Gguf(m) => {
2270 if let Some(batcher) = continuous_batcher {
2271 let mut tokens = m.tokenizer.encode(prompt, SpecialTokens::Parse);
2272 frink_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
2273 let (finish, _generated_ids, text, usage) = if synthetic {
2274 batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?
2275 } else {
2276 batcher.generate_streaming(
2277 tokens,
2278 params.clone(),
2279 m.stop_tokens.clone(),
2280 Some(|chunk: &str| {
2281 if !chunk.is_empty() {
2282 chunks[0].push(chunk.to_string());
2283 emit(chunk);
2284 }
2285 }),
2286 )?
2287 };
2288 if !text.is_empty() && chunks[0].is_empty() {
2289 chunks[0].push(text);
2290 }
2291 (vec![(finish, Vec::new())], usage)
2298 } else {
2299 generate::generate(
2300 &m.decoder,
2301 m.tokenizer.as_ref(),
2302 &m.stop_tokens,
2303 m.bos_id,
2304 prompt,
2305 params,
2306 kv_pool,
2307 paged_kv,
2308 prefix_cache,
2309 ceiling,
2310 |choice, chunk| {
2311 chunks[choice].push(chunk.to_string());
2312 if !synthetic && choice == 0 {
2316 emit(chunk);
2317 }
2318 },
2319 )?
2320 }
2321 }
2322 Model::Kimi(m) => generate::generate_engine(
2323 &m.engine,
2324 &m.tokenizer,
2325 &m.stop_tokens,
2326 None,
2327 prompt,
2328 params,
2329 |chunk| {
2330 chunks[0].push(chunk.to_string());
2331 if !synthetic {
2332 emit(chunk);
2333 }
2334 },
2335 )?,
2336 Model::Mla(m) => generate::generate_engine(
2337 &m.engine,
2338 &m.tokenizer,
2339 &m.stop_tokens,
2340 m.bos_id,
2341 prompt,
2342 params,
2343 |chunk| {
2344 chunks[0].push(chunk.to_string());
2345 if !synthetic {
2346 emit(chunk);
2347 }
2348 },
2349 )?,
2350 Model::Gemma4(m) => generate::generate_engine(
2351 &m.engine,
2352 &m.tokenizer,
2353 &m.stop_tokens,
2354 m.bos_id,
2355 prompt,
2356 params,
2357 |chunk| {
2358 chunks[0].push(chunk.to_string());
2359 if !synthetic {
2360 emit(chunk);
2361 }
2362 },
2363 )?,
2364 Model::Glm52(m) => generate::generate_engine(
2365 &m.engine,
2366 &m.tokenizer,
2367 &m.stop_tokens,
2368 m.bos_id,
2369 prompt,
2370 params,
2371 |chunk| {
2372 chunks[0].push(chunk.to_string());
2373 if !synthetic {
2374 emit(chunk);
2375 }
2376 },
2377 )?,
2378 };
2379
2380 let mut full = chunks[0].concat();
2381 if synthetic {
2382 full = format!(
2383 "[frink synthetic-weight demo: no real checkpoint loaded -- set FRINK_MODEL_PATH \
2384 to serve a real model. Decoded ids -> {full:?}]"
2385 );
2386 emit(&full);
2387 } else if used_batcher && !full.is_empty() && chunks[0].is_empty() {
2388 emit(&full);
2389 }
2390
2391 debug_assert_eq!(finishes.len(), chunks.len(), "one finish reason per choice");
2397 let mut out: Vec<generate::GeneratedChoice> = finishes
2398 .into_iter()
2399 .zip(chunks.into_iter().map(|c| c.concat()))
2400 .map(|((finish, logprobs), text)| generate::GeneratedChoice {
2401 finish,
2402 text,
2403 logprobs,
2404 })
2405 .collect();
2406 if let Some(first) = out.first_mut() {
2407 if synthetic {
2415 first.logprobs.clear();
2416 }
2417 first.text = full;
2418 }
2419 Ok((out, usage))
2420}
2421
2422#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
2427 model: &Model,
2428 prompt: &str,
2429 params: &GenerationParams,
2430 kv_pool: Option<&generate::KvPoolConfig>,
2431 paged_kv: Option<&generate::PagedKvConfig>,
2432 prefix_cache: Option<&Mutex<PrefixCache>>,
2433 continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
2434 ceiling: Option<&budget::ContextCeiling>,
2435 metal_private_decode_gate: Option<&std::sync::Mutex<()>>,
2436 ) -> Result<(Vec<generate::GeneratedChoice>, generate::Usage), generate::DecodeError> {
2439 run_generation_emit(
2440 model,
2441 prompt,
2442 params,
2443 kv_pool,
2444 paged_kv,
2445 prefix_cache,
2446 continuous_batcher,
2447 ceiling,
2448 metal_private_decode_gate,
2449 |_| {},
2450 )
2451}
2452
2453pub(crate) fn prompt_from_messages(
2464 messages: &[ChatMessage],
2465 template: &chat_template::PromptTemplate,
2466 tools: &[ToolDef],
2467 extra: serde_json::Map<String, serde_json::Value>,
2468) -> Result<String, ApiError> {
2469 let rendered = if tools.is_empty() || template.handles_tools() {
2470 template.render(messages, tools, extra)
2471 } else {
2472 let mut with_preamble = Vec::with_capacity(messages.len() + 1);
2473 with_preamble.push(ChatMessage {
2474 role: "system".to_string(),
2475 content: Some(MessageContent::Text(tool_preamble(tools))),
2476 tool_calls: None,
2477 tool_call_id: None,
2478 reasoning_content: None,
2479 });
2480 with_preamble.extend_from_slice(messages);
2481 template.render(&with_preamble, &[], extra)
2482 };
2483 rendered.map_err(template_error_response)
2484}
2485
2486fn template_error_response(err: frink_models::chat_template::TemplateError) -> ApiError {
2491 (
2492 StatusCode::BAD_REQUEST,
2493 Json(serde_json::json!({
2494 "error": {
2495 "message": format!("chat template failed to render: {err}"),
2496 "type": "invalid_request_error",
2497 "param": "messages",
2498 "code": null,
2499 }
2500 })),
2501 )
2502}
2503
2504fn tool_preamble(tools: &[ToolDef]) -> String {
2517 let mut out = String::from(
2518 "You can call tools to help answer the user. To call a tool, respond with \
2519 EXACTLY one line in this format and nothing else:\n\
2520 <tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
2521 object matching that tool's parameters>}}</tool_call>\n\n\
2522 Available tools:\n",
2523 );
2524 for t in tools {
2525 out.push_str(&format!(
2526 "- {}: {}\n parameters (JSON schema): {}\n",
2527 t.function.name,
2528 t.function.description.as_deref().unwrap_or(""),
2529 t.function
2530 .parameters
2531 .as_ref()
2532 .map(|v| v.to_string())
2533 .unwrap_or_else(|| "{}".to_string()),
2534 ));
2535 }
2536 out
2537}
2538
2539fn tool_call_deltas(
2549 events: Vec<crate::policy::parser::ToolCallEvent>,
2550 opened: &std::cell::Cell<usize>,
2551) -> (String, Vec<ToolCallDelta>) {
2552 let mut text = String::new();
2553 let mut deltas = Vec::new();
2554 for event in events {
2555 match event {
2556 crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
2557 crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
2558 opened.set(opened.get().max(index + 1));
2559 deltas.push(ToolCallDelta::opening(index, name));
2560 }
2561 crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
2562 if !fragment.is_empty() {
2563 deltas.push(ToolCallDelta::arguments(index, fragment));
2564 }
2565 }
2566 crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
2567 }
2568 }
2569 (text, deltas)
2570}
2571
2572fn build_response_message(
2583 text: String,
2584 tools: &[ToolDef],
2585 posture: output::OutputPosture,
2586 base_finish: &'static str,
2587) -> (ChatCompletionResponseMessage, &'static str) {
2588 let parsed = output::parse_output(&text, tools, posture);
2589 let calls: Vec<ToolCallOut> = parsed
2590 .calls
2591 .into_iter()
2592 .enumerate()
2593 .map(|(index, call)| ToolCallOut {
2594 id: format!("call_{index}"),
2595 kind: "function",
2596 function: ToolCallFunctionOut {
2597 name: call.name,
2598 arguments: call.arguments,
2599 },
2600 })
2601 .collect();
2602 if !calls.is_empty() {
2603 return (
2604 ChatCompletionResponseMessage {
2605 role: "assistant",
2606 content: None,
2607 reasoning_content: parsed.reasoning,
2608 tool_calls: Some(calls),
2609 },
2610 "tool_calls",
2611 );
2612 }
2613 (
2614 ChatCompletionResponseMessage {
2615 role: "assistant",
2616 content: Some(parsed.content),
2617 reasoning_content: parsed.reasoning,
2618 tool_calls: None,
2619 },
2620 base_finish,
2621 )
2622}
2623
2624fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
2629 let mut history = match &req.session_id {
2630 Some(id) => state.sessions.extend_and_get(id, &req.messages),
2631 None => req.messages.clone(),
2632 };
2633 if req.json_object_mode() {
2634 inject_json_object_system_hint(&mut history);
2635 }
2636 history
2637}
2638
2639fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
2640 const HINT: &str =
2641 "You must respond with valid JSON only (a single JSON object, no markdown fences).";
2642 if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
2643 match &mut sys.content {
2644 Some(MessageContent::Text(s)) if !s.contains("JSON") => {
2645 s.push_str("\n\n");
2646 s.push_str(HINT);
2647 }
2648 None => {
2649 sys.content = Some(MessageContent::Text(HINT.to_string()));
2650 }
2651 _ => {}
2652 }
2653 } else {
2654 messages.insert(
2655 0,
2656 ChatMessage {
2657 role: "system".to_string(),
2658 content: Some(MessageContent::Text(HINT.to_string())),
2659 tool_calls: None,
2660 tool_call_id: None,
2661 reasoning_content: None,
2662 },
2663 );
2664 }
2665}
2666
2667async fn chat_completions(
2668 State(state): State<Arc<AppState>>,
2669 headers: axum::http::HeaderMap,
2670 Json(req): Json<ChatCompletionRequest>,
2671) -> Response {
2672 let attribution = attribution::Attribution::from_headers(&headers);
2673 state
2674 .requests_total
2675 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2676 let started = std::time::Instant::now();
2677
2678 let request_id = frink_api::next_request_id();
2682 let stream = req.stream.unwrap_or(false);
2683
2684 let refusal = cache_admin::check_admission(&state)
2690 .err()
2691 .or_else(|| req.validate_supported_fields().err());
2692 if let Some(err) = refusal {
2693 state
2694 .request_errors_total
2695 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2696 let response = err.into_response();
2697 state.record_request(stats::Record {
2698 request_id: &request_id,
2699 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2700 model: state.active_model_name(),
2701 status: response.status().as_u16(),
2702 stream,
2703 duration_ms: started.elapsed().as_millis() as u64,
2704 usage: None,
2705 attribution: &attribution,
2706 });
2707 return response;
2708 }
2709
2710 let response = if stream {
2711 chat_completions_stream(
2712 Arc::clone(&state),
2713 req,
2714 request_id.clone(),
2715 started,
2716 attribution.clone(),
2717 )
2718 .await
2719 .into_response()
2720 } else {
2721 chat_completions_full(
2722 Arc::clone(&state),
2723 req,
2724 request_id.clone(),
2725 started,
2726 attribution.clone(),
2727 )
2728 .await
2729 .into_response()
2730 };
2731
2732 if response.status().is_client_error() || response.status().is_server_error() {
2733 state
2734 .request_errors_total
2735 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2736 state.record_request(stats::Record {
2740 request_id: &request_id,
2741 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2742 model: state.active_model_name(),
2745 status: response.status().as_u16(),
2746 stream,
2747 duration_ms: started.elapsed().as_millis() as u64,
2748 usage: None,
2749 attribution: &attribution,
2750 });
2751 }
2752 state.mark_request_finished();
2753
2754 response
2755}
2756
2757async fn chat_completions_full(
2758 state: Arc<AppState>,
2759 req: ChatCompletionRequest,
2760 request_id: String,
2761 started: std::time::Instant,
2762 attribution: attribution::Attribution,
2763) -> Result<Json<ChatCompletionResponse>, ApiError> {
2764 let tools_active = req.tools_active();
2765 let active = state.require_active()?;
2769 let history = resolve_history(&state, &req);
2770 let template = active.generative()?.chat_template();
2771 let kwargs = req.resolve_template_kwargs(&template);
2772 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2773 let mut params =
2781 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2782 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2783 let key = req.is_cacheable().then(|| req.cache_key(&prompt, ¶ms));
2784
2785 let mut generated_logprobs: Vec<crate::sampling_loop::PerTokenProbs> = Vec::new();
2789 let n_logprobs = req.n_logprobs()?;
2792 let decode_piece = |id: usize| active.decode_any(&[id]);
2794 let (completion, cache_status) = if let Some(cached) = key
2795 .as_ref()
2796 .and_then(|key| lock_cache(&state.response_cache).get(key))
2797 {
2798 tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
2799 (cached, "hit")
2800 } else {
2801 let (choices, usage) = decode_task::buffered(
2802 decode_task::DecodeHandles::take(&state, &active)?,
2803 prompt.clone(),
2804 params,
2805 )
2806 .await?;
2807
2808 generated_logprobs = choices.iter().map(|c| c.logprobs.clone()).collect();
2813 let completion = response_cache::CachedCompletion {
2814 choices: choices.into_iter().map(|c| (c.finish, c.text)).collect(),
2815 usage,
2816 };
2817 let cache_status = match key {
2823 Some(key) => match completion.clone().cacheable() {
2827 Some(cacheable) => {
2828 tracing::debug!("cache miss for key {}", key.digest());
2829 lock_cache(&state.response_cache).put(key, cacheable);
2830 "miss"
2831 }
2832 None => "skip",
2833 },
2834 None => "skip",
2835 };
2836 (completion, cache_status)
2837 };
2838 let content = completion.first_text().to_string();
2841
2842 if req.json_object_mode() {
2843 json_mode::validate_json_object_output(&content)?;
2844 }
2845
2846 if let Some(id) = &req.session_id {
2850 state.sessions.store_reply(
2851 id,
2852 ChatMessage {
2853 role: "assistant".to_string(),
2854 content: Some(MessageContent::Text(content.clone())),
2855 tool_calls: None,
2856 tool_call_id: None,
2857 reasoning_content: None,
2858 },
2859 );
2860 }
2861
2862 let posture = output::OutputPosture::resolve_full(
2867 active.reasoning_format(),
2868 active.tool_call_format(),
2869 &prompt,
2870 );
2871 let tools: &[_] = if tools_active { &req.tools } else { &[] };
2872 let rendered: Vec<ChatCompletionChoice> = completion
2873 .choices
2874 .into_iter()
2875 .enumerate()
2876 .map(|(index, (finish, text))| {
2877 let (message, finish_reason) =
2878 build_response_message(text, tools, posture, finish.as_str());
2879 ChatCompletionChoice {
2880 index,
2881 message,
2882 finish_reason,
2883 logprobs: n_logprobs.map(|k| {
2884 crate::logprobs::render_chat(
2885 generated_logprobs.get(index).unwrap_or(&Vec::new()),
2886 Some(k),
2887 &decode_piece,
2888 )
2889 }),
2890 }
2891 })
2892 .collect();
2893
2894 state.record_request(stats::Record {
2895 request_id: &request_id,
2896 route: frink_api::routes::V1_CHAT_COMPLETIONS,
2897 model: Some(active.name().to_string()),
2900 status: 200,
2901 stream: false,
2902 duration_ms: started.elapsed().as_millis() as u64,
2903 usage: Some(&completion.usage),
2904 attribution: &attribution,
2905 });
2906
2907 Ok(Json(ChatCompletionResponse {
2908 id: request_id.clone(),
2909 request_id,
2910 object: "chat.completion",
2911 model: req.model,
2912 choices: rendered,
2913 usage: completion.usage,
2914 frink_cache: cache_status,
2915 }))
2916}
2917
2918async fn chat_completions_stream(
2919 state: Arc<AppState>,
2920 req: ChatCompletionRequest,
2921 request_id: String,
2922 started: std::time::Instant,
2923 attribution: attribution::Attribution,
2924) -> Result<Response, ApiError> {
2925 if req.several_choices() {
2935 return Err(unsupported_feature(
2936 "`n` > 1 with `stream` is not implemented: the choices would arrive one after \
2937 another rather than interleaved by `choices[].index`. Send the request without \
2938 `stream`, which serves `n` on this route.",
2939 ));
2940 }
2941 let tools_active = req.tools_active();
2942 let active = state.require_active()?;
2946 let history = resolve_history(&state, &req);
2947 let template = active.generative()?.chat_template();
2948 let kwargs = req.resolve_template_kwargs(&template);
2949 let prompt = req.render_prompt(&history, &template, &req.tools, kwargs, active.name())?;
2950 let model_name = req.model.clone();
2951 let session_id = req.session_id.clone();
2952 let sessions = state.sessions.clone();
2953
2954 let model = Arc::clone(active.generative()?);
2955 let kv_pool = state.kv_pool.clone();
2956 let paged_kv = state.paged_kv.clone();
2957 let prefix_cache = state.prefix_cache.clone();
2958 let batcher = active.batcher.clone();
2959 let ceiling = active.ceiling.clone();
2960 let metal_private_decode_gate = state.metal_private_decode_gate.clone();
2961 let mut params =
2962 req.generation_params_for_template(&template, active.name(), active.sampler_model())?;
2963 params.lora = lora::resolve_request(active.generative()?, req.lora.as_deref())?;
2964 let stats_state = Arc::clone(&state);
2965 let served_model = active.name().to_string();
2968 let posture = output::OutputPosture::resolve_full(
2972 active.reasoning_format(),
2973 active.tool_call_format(),
2974 &prompt,
2975 );
2976 let offered_tools: Vec<ToolDef> = if tools_active {
2979 req.tools.clone()
2980 } else {
2981 Vec::new()
2982 };
2983
2984 let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
2989 params.cancel = Some(cancel_token.clone());
2990
2991 let overlap = true;
3000
3001 let slot = req
3005 .stream_resumable
3006 .unwrap_or(false)
3007 .then(|| state.streams.register(&request_id));
3008 let emitter = resume::Emitter::new(slot);
3009
3010 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
3011 let keepalive = sse::keepalive_event(&ChatCompletionChunk {
3015 id: request_id.clone(),
3016 request_id: None,
3017 object: "chat.completion.chunk",
3018 model: model_name.clone(),
3019 choices: vec![ChatCompletionChunkChoice {
3020 index: 0,
3021 delta: ChatCompletionChunkDelta {
3022 role: None,
3023 content: None,
3024 reasoning_content: None,
3025 tool_calls: None,
3026 },
3027 finish_reason: None,
3028 }],
3029 usage: None,
3030 });
3031
3032 tokio::task::spawn_blocking(move || {
3033 let _cancel_guard = cancel_guard;
3036 let tx_chunks = tx.clone();
3037 let orphan_timeout = sse::orphan_timeout_from_env();
3042 let mut first = true;
3043 let head_request_id = request_id.clone();
3044 let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
3054 Rc::new(RefCell::new(posture.reasoning_parser()));
3055 let emit_reasoning = Rc::clone(&stream_reasoning);
3056 let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
3060 RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
3061 );
3062 let emit_tools = Rc::clone(&stream_tools);
3063 let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
3067 let emit_streamed_calls = Rc::clone(&streamed_calls);
3068 let result = run_generation_emit(
3069 &model,
3070 &prompt,
3071 ¶ms,
3072 kv_pool.as_ref(),
3073 paged_kv.as_ref(),
3074 prefix_cache.as_deref(),
3075 batcher.as_ref(),
3076 ceiling.as_deref(),
3077 metal_private_decode_gate.as_deref(),
3078 |chunk| {
3079 if !overlap || chunk.is_empty() {
3080 return;
3081 }
3082 let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
3083 Some(parser) => {
3084 let delta = parser.push(chunk);
3085 (delta.reasoning, delta.content)
3086 }
3087 None => (String::new(), chunk.to_string()),
3088 };
3089 let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
3093 Some(parser) => {
3094 let (text, calls) =
3095 tool_call_deltas(parser.push(&content), &emit_streamed_calls);
3096 (text, calls)
3097 }
3098 None => (content, Vec::new()),
3099 };
3100 if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
3103 return;
3104 }
3105 let role = if first { Some("assistant") } else { None };
3106 let request_id = first.then(|| head_request_id.clone());
3107 first = false;
3108 let payload = ChatCompletionChunk {
3109 id: head_request_id.clone(),
3110 request_id,
3111 object: "chat.completion.chunk",
3112 model: model_name.clone(),
3113 choices: vec![ChatCompletionChunkChoice {
3114 index: 0,
3115 delta: ChatCompletionChunkDelta {
3116 role,
3117 content: (!content.is_empty()).then_some(content),
3118 reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
3119 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3120 },
3121 finish_reason: None,
3122 }],
3123 usage: None,
3124 };
3125 if let Err(why) =
3134 sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
3135 {
3136 if why == sse::SendFailure::Orphaned {
3137 tracing::warn!(
3138 "SSE stream {head_request_id} accepted nothing for the orphan \
3139 deadline; treating it as abandoned"
3140 );
3141 }
3142 if !emitter.is_resumable() {
3154 cancel_token.cancel();
3155 }
3156 }
3157 },
3158 );
3159
3160 let mut pending_request_id = first.then(|| request_id.clone());
3166
3167 match result {
3168 Ok((choices, usage)) => {
3171 let one = choices
3172 .into_iter()
3173 .next()
3174 .expect("a generation produces at least one choice");
3175 let (finish, full_text) = (one.finish, one.text);
3176 if let Some(id) = &session_id {
3177 sessions.store_reply(
3178 id,
3179 ChatMessage {
3180 role: "assistant".to_string(),
3181 content: Some(MessageContent::Text(full_text.clone())),
3182 tool_calls: None,
3183 tool_call_id: None,
3184 reasoning_content: None,
3185 },
3186 );
3187 }
3188 let mut streamed_finish: Option<&'static str> = None;
3194 if overlap {
3195 let tail = stream_reasoning
3196 .borrow_mut()
3197 .as_mut()
3198 .map(|parser| parser.flush())
3199 .unwrap_or_default();
3200 let (mut content, mut tool_calls) = (tail.content, Vec::new());
3201 if let Some(parser) = stream_tools.borrow_mut().as_mut() {
3202 let mut events = parser.push(&content);
3203 events.extend(parser.finish());
3204 let (text, calls) = tool_call_deltas(events, &streamed_calls);
3205 content = text;
3206 tool_calls = calls;
3207 }
3208 if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
3209 let payload = ChatCompletionChunk {
3210 id: request_id.clone(),
3211 request_id: pending_request_id.take(),
3212 object: "chat.completion.chunk",
3213 model: model_name.clone(),
3214 choices: vec![ChatCompletionChunkChoice {
3215 index: 0,
3216 delta: ChatCompletionChunkDelta {
3217 role: None,
3218 content: (!content.is_empty()).then_some(content),
3219 reasoning_content: (!tail.reasoning.is_empty())
3220 .then_some(tail.reasoning),
3221 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3222 },
3223 finish_reason: None,
3224 }],
3225 usage: None,
3226 };
3227 let _ =
3228 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3229 }
3230 if streamed_calls.get() > 0 {
3231 streamed_finish = Some("tool_calls");
3232 }
3233 } else {
3234 let parsed = output::parse_output(&full_text, &offered_tools, posture);
3237 let tool_calls: Vec<ToolCallDelta> = parsed
3238 .calls
3239 .iter()
3240 .enumerate()
3241 .map(|(index, call)| {
3242 ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
3243 })
3244 .collect();
3245 if !tool_calls.is_empty() {
3246 streamed_finish = Some("tool_calls");
3247 }
3248 if !tool_calls.is_empty()
3249 || !parsed.content.is_empty()
3250 || parsed.reasoning.is_some()
3251 {
3252 let payload = ChatCompletionChunk {
3253 id: request_id.clone(),
3254 request_id: pending_request_id.take(),
3255 object: "chat.completion.chunk",
3256 model: model_name.clone(),
3257 choices: vec![ChatCompletionChunkChoice {
3258 index: 0,
3259 delta: ChatCompletionChunkDelta {
3260 role: Some("assistant"),
3261 content: (!parsed.content.is_empty() && tool_calls.is_empty())
3262 .then(|| parsed.content.clone()),
3263 reasoning_content: parsed.reasoning.clone(),
3264 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
3265 },
3266 finish_reason: None,
3267 }],
3268 usage: None,
3269 };
3270 let _ =
3271 sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3272 }
3273 }
3274 let final_finish_reason = match streamed_finish {
3278 Some(reason) if finish.as_str() != "length" => reason,
3279 _ => finish.as_str(),
3280 };
3281 let final_payload = ChatCompletionChunk {
3282 id: request_id.clone(),
3283 request_id: pending_request_id.take(),
3284 object: "chat.completion.chunk",
3285 model: model_name,
3286 choices: vec![ChatCompletionChunkChoice {
3287 index: 0,
3288 delta: ChatCompletionChunkDelta {
3289 role: None,
3290 content: None,
3291 reasoning_content: None,
3292 tool_calls: None,
3293 },
3294 finish_reason: Some(final_finish_reason),
3295 }],
3296 usage: Some(usage.clone()),
3297 };
3298 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
3299 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3300 stats_state.record_request(stats::Record {
3305 request_id: &request_id,
3306 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3307 model: Some(served_model.clone()),
3308 status: 200,
3309 stream: true,
3310 duration_ms: started.elapsed().as_millis() as u64,
3311 usage: Some(&usage),
3312 attribution: &attribution,
3313 });
3314 }
3315 Err(e) => {
3316 tracing::warn!("decode error on streamed request {request_id}: {e}");
3317 stats_state.record_request(stats::Record {
3323 request_id: &request_id,
3324 route: frink_api::routes::V1_CHAT_COMPLETIONS,
3325 model: Some(served_model.clone()),
3326 status: 500,
3327 stream: true,
3328 duration_ms: started.elapsed().as_millis() as u64,
3329 usage: None,
3330 attribution: &attribution,
3331 });
3332 let payload = ChatCompletionChunk {
3333 id: request_id.clone(),
3334 request_id: pending_request_id.take(),
3335 object: "chat.completion.chunk",
3336 model: model_name,
3337 choices: vec![ChatCompletionChunkChoice {
3338 index: 0,
3339 delta: ChatCompletionChunkDelta {
3340 role: Some("assistant"),
3341 content: Some(format!("[error: {e}]")),
3342 reasoning_content: None,
3343 tool_calls: None,
3344 },
3345 finish_reason: Some("stop"),
3346 }],
3347 usage: None,
3348 };
3349 let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
3350 let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
3351 }
3352 }
3353 drop(emitter);
3357 });
3358
3359 let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
3360 Ok((
3383 [(
3384 axum::http::HeaderName::from_static("x-accel-buffering"),
3385 axum::http::HeaderValue::from_static("no"),
3386 )],
3387 Sse::new(stream),
3388 )
3389 .into_response())
3390}
3391
3392fn protected_routes() -> Router<Arc<AppState>> {
3423 use frink_api::routes;
3424
3425 Router::new()
3426 .route(routes::V1_MODELS, get(list_models))
3427 .route(routes::V1_RESPONSES, post(responses::responses))
3431 .route(
3432 &axum_path(routes::V1_RESPONSE),
3433 get(responses::responses_get),
3434 )
3435 .route(
3436 &axum_path(routes::V1_RESPONSE_CANCEL),
3437 post(responses::responses_cancel),
3438 )
3439 .route(&axum_path(routes::SLOTS_ID), post(slots::post_slot))
3440 .route(routes::V1_STATS, get(serving_stats))
3441 .route(routes::V1_REQUESTS, get(recent_requests))
3442 .route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
3443 .route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
3444 .route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
3445 .route(
3446 routes::LORA_ADAPTERS,
3447 get(lora::get_lora_adapters).post(lora::post_lora_adapters),
3448 )
3449 .route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
3450 .route(routes::V1_CANCEL, post(cancel_generation))
3454 .route(&axum_path(routes::V1_STREAM), get(resume::resume))
3459 .route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
3460 .route(routes::V1_MESSAGES, post(anthropic::messages))
3461 .route(
3462 routes::V1_MESSAGES_COUNT_TOKENS,
3463 post(anthropic::count_tokens),
3464 )
3465 .route(routes::V1_COMPLETIONS, post(openai_extra::completions))
3466 .route(routes::COMPLETION, post(completion::completion))
3471 .route(routes::COMPLETIONS, post(completion::completion))
3472 .route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
3473 .route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
3474 .route(routes::TOKENIZE, post(openai_extra::tokenize))
3480 .route(routes::DETOKENIZE, post(openai_extra::detokenize))
3481 .route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
3482 .route(routes::V1_RERANK, post(rerank::rerank))
3486 .route(routes::RERANK, post(rerank::rerank))
3487 .route(routes::CACHE_STATS, get(cache_stats))
3488 .route(routes::METRICS, get(metrics))
3489 .route(routes::ADMIN_MODELS, get(admin::models))
3494 .route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
3495 .route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
3496 .route(routes::ADMIN_DOWNLOAD, post(admin::download))
3497 .route(routes::ADMIN_TASKS, get(admin::tasks))
3498 .route(&admin::cancel_route(), post(admin::cancel_task))
3499 .route(routes::ADMIN_STATS, get(admin::stats))
3500 .merge(conversations::router())
3504}
3505
3506fn axum_path(template: &str) -> String {
3507 let mut out = String::with_capacity(template.len());
3508 let mut rest = template;
3509 while let Some(open) = rest.find('{') {
3510 let Some(close) = rest[open..].find('}').map(|c| open + c) else {
3511 break;
3512 };
3513 out.push_str(&rest[..open]);
3514 out.push(':');
3515 out.push_str(&rest[open + 1..close]);
3516 rest = &rest[close + 1..];
3517 }
3518 out.push_str(rest);
3519 out
3520}
3521
3522async fn cancel_generation(
3531 State(state): State<Arc<AppState>>,
3532 Json(req): Json<frink_api::CancelGenerationRequest>,
3533) -> Response {
3534 let cancelled = state.cancels.cancel(&req.request_id);
3535 let status = if cancelled {
3536 StatusCode::OK
3537 } else {
3538 StatusCode::NOT_FOUND
3539 };
3540 let detail = if cancelled {
3541 "the generation was asked to stop; it ends at its next token".to_string()
3542 } else {
3543 "no generation with that request_id is running -- it has already \
3544 finished, was never issued, or was served by a path that does \
3545 not register for cancellation"
3546 .to_string()
3547 };
3548 (
3549 status,
3550 Json(frink_api::CancelGenerationResponse {
3551 request_id: req.request_id,
3552 cancelled,
3553 detail,
3554 }),
3555 )
3556 .into_response()
3557}
3558
3559type Activated = (
3563 Loaded,
3564 Option<serving::batch::ContinuousBatcher>,
3565 Option<Arc<budget::ContextCeiling>>,
3566);
3567
3568fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
3586 let mut batcher = serving::batch::BatcherConfig::from_env();
3587 if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
3588 return batcher;
3591 }
3592 let Some(path) = path else {
3593 return batcher;
3594 };
3595 let priced = budget::price_gguf(path, frink_models::KvElem::F32, 1);
3601 let Some((priced, gguf_ctx, source)) = priced else {
3602 return batcher;
3603 };
3604 let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
3605 tracing::warn!(
3609 "this checkpoint's weights leave no room for KV inside the {source}: {} weight \
3610 bytes against a {} byte budget. Serving with no derived context ceiling -- set \
3611 FRINK_DEVICE_BUDGET_BYTES if the probe is wrong, or FRINK_CB_MAX_CONTEXT to \
3612 admit on a number you choose.",
3613 priced.weights_bytes,
3614 priced.device_budget_bytes,
3615 );
3616 return batcher;
3617 };
3618 tracing::info!("{source}");
3619 tracing::info!("{}", derived.fit);
3620 let adopted = budget::apply_derived(&mut batcher, &derived);
3621 if adopted.max_context {
3622 tracing::info!(
3623 "derived per-request context ceiling: {} token positions (prompt + max_tokens); \
3624 override with FRINK_CB_MAX_CONTEXT",
3625 derived.max_context
3626 );
3627 }
3628 if adopted.kv_blocks {
3629 tracing::info!(
3630 "derived KV block budget: {} blocks x {} positions; override with FRINK_CB_KV_BLOCKS",
3631 derived.kv_blocks,
3632 batcher.kv_block_size
3633 );
3634 }
3635 if let Some(narrowed) = adopted.max_context_narrowed {
3636 tracing::info!(
3637 "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",
3638 batcher.kv_blocks.unwrap_or_default(),
3639 batcher.kv_block_size
3640 );
3641 }
3642 batcher
3643}
3644
3645pub(crate) fn activate_loaded_model(
3654 loaded: model::LoadedModel,
3655 enable_continuous_batching: bool,
3656 path: Option<&str>,
3657 paged_kv: Option<&generate::PagedKvConfig>,
3658) -> Activated {
3659 match loaded {
3660 model::LoadedModel::Gguf(g) => {
3661 let decoder = Arc::new(g.decoder);
3662 let tokenizer = Arc::new(g.tokenizer);
3663 let config = price_batcher_config(path);
3664 let shape =
3670 frink_models::KvShape::from_config(&decoder.config, frink_models::KvElem::F32);
3671 let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
3672 let batcher = if enable_continuous_batching {
3673 tracing::info!(
3674 "continuous batching enabled: decode steps share Decoder::forward_multi_seq \
3675 (stop sequences use the same pending-buffer trim as the private generate loop)"
3676 );
3677 let tok = Arc::clone(&tokenizer);
3678 let decode = Arc::new(move |ids: &[usize]| tok.decode_bytes(ids));
3679 Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
3680 Arc::clone(&decoder),
3681 decode,
3682 config,
3683 Arc::clone(&ceiling),
3684 paged_kv.cloned(),
3685 ))
3686 } else {
3687 None
3688 };
3689 (
3690 Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
3691 decoder,
3692 tokenizer,
3693 stop_tokens: g.stop_tokens,
3694 bos_id: g.bos_id,
3695 is_synthetic: g.is_synthetic,
3696 chat_template: g.chat_template,
3697 }))),
3698 batcher,
3699 Some(ceiling),
3700 )
3701 }
3702 model::LoadedModel::Kimi(k) => (
3703 Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
3704 engine: k.engine,
3705 tokenizer: k.tokenizer,
3706 stop_tokens: k.stop_tokens,
3707 chat_template: k.chat_template,
3708 }))),
3709 None,
3710 None,
3711 ),
3712 model::LoadedModel::Mla(m) => (
3713 Loaded::Generative(Arc::new(Model::Mla(MlaModel {
3714 engine: m.engine,
3715 tokenizer: m.tokenizer,
3716 stop_tokens: m.stop_tokens,
3717 bos_id: m.bos_id,
3718 name: m.name,
3719 chat_template: m.chat_template,
3720 }))),
3721 None,
3722 None,
3723 ),
3724 model::LoadedModel::Gemma4(m) => (
3725 Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
3726 engine: m.engine,
3727 tokenizer: m.tokenizer,
3728 stop_tokens: m.stop_tokens,
3729 bos_id: m.bos_id,
3730 name: m.name,
3731 chat_template: m.chat_template,
3732 }))),
3733 None,
3734 None,
3735 ),
3736 model::LoadedModel::Glm52(g) => (
3737 Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
3738 engine: g.engine,
3739 tokenizer: g.tokenizer,
3740 stop_tokens: g.stop_tokens,
3741 bos_id: g.bos_id,
3742 name: g.name,
3743 chat_template: g.chat_template,
3744 }))),
3745 None,
3746 None,
3747 ),
3748 model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
3753 }
3754}
3755
3756struct StartupModels {
3763 loaded: model::LoadedModel,
3764 embedding: Option<Arc<frink_models::EmbeddingModel>>,
3765}
3766
3767fn continuous_batching_env() -> Option<bool> {
3768 match std::env::var("FRINK_CONTINUOUS_BATCHING")
3769 .ok()
3770 .map(|v| v.trim().to_ascii_lowercase())
3771 .as_deref()
3772 {
3773 None => None,
3774 Some("1" | "true" | "yes" | "on") => Some(true),
3775 Some("0" | "false" | "no" | "off") => Some(false),
3776 _ => None,
3777 }
3778}
3779
3780fn metal_private_decode_active() -> bool {
3781 #[cfg(feature = "metal")]
3782 {
3783 BUILT_WITH_METAL
3784 && frink_metal::attn::metal_attn_enabled()
3785 && std::env::var("FRINK_METAL").ok().as_deref() != Some("0")
3786 }
3787 #[cfg(not(feature = "metal"))]
3788 {
3789 false
3790 }
3791}
3792
3793fn continuous_batching_compatible(
3794 loaded: &model::LoadedModel,
3795 kv_pool: &Option<generate::KvPoolConfig>,
3796 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3797 paged_kv: &Option<generate::PagedKvConfig>,
3798) -> bool {
3799 matches!(loaded, model::LoadedModel::Gguf(_))
3800 && (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
3801}
3802
3803fn resolve_continuous_batching_enabled(
3804 loaded: &model::LoadedModel,
3805 kv_pool: &Option<generate::KvPoolConfig>,
3806 prefix_cache: &Option<Arc<Mutex<PrefixCache>>>,
3807 paged_kv: &Option<generate::PagedKvConfig>,
3808) -> bool {
3809 if !continuous_batching_compatible(loaded, kv_pool, prefix_cache, paged_kv) {
3810 return false;
3811 }
3812 match continuous_batching_env() {
3813 Some(true) => true,
3814 Some(false) => false,
3815 None => metal_private_decode_active(),
3816 }
3817}
3818
3819fn acquire_metal_private_decode_gate(
3820 gate: Option<&std::sync::Mutex<()>>,
3821 used_batcher: bool,
3822) -> Option<std::sync::MutexGuard<'_, ()>> {
3823 if used_batcher {
3824 None
3825 } else {
3826 gate.map(|g| g.lock().unwrap_or_else(|p| p.into_inner()))
3827 }
3828}
3829
3830fn build_app_state(
3831 models: StartupModels,
3832 kv_pool: Option<generate::KvPoolConfig>,
3833 paged_kv: Option<generate::PagedKvConfig>,
3834 prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
3835 enable_continuous_batching: bool,
3836 mcp: Option<mcp::LoadedMcpConfig>,
3837 detection: Arc<health::Detection>,
3838) -> AppState {
3839 let StartupModels { loaded, embedding } = models;
3840 let configured_path = std::env::var("FRINK_MODEL_PATH").ok();
3841 let (loaded, batcher, ceiling) = activate_loaded_model(
3842 loaded,
3843 enable_continuous_batching,
3844 configured_path.as_deref(),
3845 paged_kv.as_ref(),
3846 );
3847 let id = startup_model_id();
3853 let metal_private_decode_gate = if enable_continuous_batching || !metal_private_decode_active()
3854 {
3855 None
3856 } else {
3857 tracing::info!(
3858 "Metal private-loop decode will serialize concurrent requests until \
3859 continuous batching is enabled (FRINK_CONTINUOUS_BATCHING=1 or --cont-batching)"
3860 );
3861 Some(Arc::new(std::sync::Mutex::new(())))
3862 };
3863 AppState {
3864 embedding,
3865 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
3866 id,
3867 loaded,
3868 batcher,
3869 ceiling,
3870 checkpoint_path: configured_path.as_deref().map(PathBuf::from),
3871 }))),
3872 paged_kv,
3873 load_in_progress: std::sync::atomic::AtomicBool::new(false),
3874 tasks: Arc::new(tasks::TaskRegistry::new()),
3875 cancels: Arc::new(cancel::CancelRegistry::new()),
3876 stats: stats::Stats::new(),
3877 streams: resume::StreamRegistry::new(),
3878 model_dir: admin::model_dirs().into_iter().next(),
3879 response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
3880 kv_pool,
3881 prefix_cache,
3882 sessions: session::SessionStore::new(),
3883 requests_total: std::sync::atomic::AtomicU64::new(0),
3884 request_errors_total: std::sync::atomic::AtomicU64::new(0),
3885 started_at: std::time::Instant::now(),
3886 last_request_ms: std::sync::atomic::AtomicU64::new(0),
3887 detection,
3888 mcp,
3889 continuous_batching_enabled: enable_continuous_batching,
3890 metal_private_decode_gate,
3891 loading_model: Mutex::new(None),
3892 last_load_error: Mutex::new(None),
3893 serving: Mutex::new(crate::stats::ServingStats::default()),
3894 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
3895 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
3896 started_unix: unix_now(),
3897 }
3898}
3899
3900fn load_embedding_model() -> anyhow::Result<Option<Arc<frink_models::EmbeddingModel>>> {
3908 let Ok(path) = std::env::var("FRINK_EMBEDDING_MODEL_PATH") else {
3909 return Ok(None);
3910 };
3911 let model = frink_models::EmbeddingModel::from_gguf_path(&path)
3912 .map_err(|e| anyhow::anyhow!("FRINK_EMBEDDING_MODEL_PATH={path}: {e}"))?;
3913 tracing::info!(
3914 "loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
3915 model.name(),
3916 model.architecture(),
3917 model.n_embd(),
3918 model.pooling_type().name(),
3919 model.n_ctx_train(),
3920 );
3921 Ok(Some(Arc::new(model)))
3922}
3923
3924fn unix_now() -> u64 {
3929 std::time::SystemTime::now()
3930 .duration_since(std::time::UNIX_EPOCH)
3931 .map(|d| d.as_secs())
3932 .unwrap_or(0)
3933}
3934
3935fn startup_model_id() -> Option<String> {
3940 let configured = std::env::var("FRINK_MODEL_PATH").ok()?;
3941 let configured = std::fs::canonicalize(&configured).ok()?;
3942 admin::discover(&admin::model_dirs())
3943 .into_iter()
3944 .find(|d| {
3945 std::fs::canonicalize(&d.path)
3946 .map(|p| p == configured)
3947 .unwrap_or(false)
3948 })
3949 .map(|d| d.id)
3950}
3951
3952fn init_cpu_pool() {
3961 match frink_core::threads::init_cpu_pool() {
3962 Some(n) => eprintln!(
3963 "frink-server: rayon pool {n} threads (perf cores {}; override with FRINK_CPU_THREADS)",
3964 frink_core::threads::perf_core_count()
3965 ),
3966 None => eprintln!("frink-server: global rayon pool already built; leaving it alone"),
3967 }
3968}
3969
3970fn announce_ready(addr: SocketAddr, scheme: &str) {
3983 use std::io::Write;
3984 let ready =
3985 frink_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
3986 let mut stdout = std::io::stdout().lock();
3987 let _ = writeln!(stdout, "{}", ready.to_line());
3988 let _ = stdout.flush();
3989}
3990
3991async fn shutdown_signal(exit_on_stdin_close: bool) {
4003 if !exit_on_stdin_close {
4004 std::future::pending::<()>().await;
4005 return;
4006 }
4007 let _ = tokio::task::spawn_blocking(|| {
4008 use std::io::Read;
4009 let mut sink = [0u8; 256];
4010 let mut stdin = std::io::stdin().lock();
4011 loop {
4012 match stdin.read(&mut sink) {
4013 Ok(0) => break,
4015 Ok(_) => continue,
4017 Err(e) => {
4018 tracing::warn!("stdin read failed ({e}); treating it as closed");
4019 break;
4020 }
4021 }
4022 }
4023 })
4024 .await;
4025 tracing::info!("stdin closed; shutting down");
4026}
4027
4028fn tokio_worker_threads() -> usize {
4033 std::env::var("FRINK_TOKIO_WORKERS")
4034 .ok()
4035 .and_then(|v| v.trim().parse::<usize>().ok())
4036 .filter(|n| *n > 0)
4037 .unwrap_or(2)
4038}
4039
4040fn install_ring_crypto_provider() {
4057 let _ = rustls::crypto::ring::default_provider().install_default();
4058}
4059
4060pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
4066 if args.list_devices {
4067 frink_models::devices::print_available_devices();
4068 return Ok(());
4069 }
4070 apply_cli_overrides(&args)?;
4071
4072 let _instance = {
4076 use frink_core::instance::{register, InstancePolicy};
4077 let policy = if args.allow_multiple_instances {
4078 InstancePolicy::Multi
4079 } else {
4080 InstancePolicy::from_env_or(InstancePolicy::Single)
4081 };
4082 let model = std::env::var("FRINK_MODEL_PATH").ok();
4083 register(
4084 "server",
4085 model.as_deref(),
4086 frink_core::instance::current_backend(),
4087 policy,
4088 )
4089 .map_err(|conflict| anyhow::anyhow!("{conflict}"))?
4090 };
4091
4092 let journal = journal::Journal::from_env();
4093 eprintln!(
4094 "frink-server: process lifecycle journal at {:?} (override with FRINK_JOURNAL_PATH)",
4095 journal.path()
4096 );
4097 journal.append(&journal::Record::session_start(
4098 env!("CARGO_PKG_VERSION"),
4099 std::process::id(),
4100 ));
4101 journal::install_panic_hook(journal.clone());
4102
4103 let mcp_config_path = args.mcp_config.clone();
4104 let exit_on_stdin_close = args.exit_on_stdin_close
4105 || std::env::var("FRINK_EXIT_ON_STDIN_CLOSE")
4106 .map(|v| v == "1")
4107 .unwrap_or(false);
4108
4109 unsafe { frink_core::weight_matrix::default_cpu_int_dot_on() };
4113 init_cpu_pool();
4114
4115 let runtime = tokio::runtime::Builder::new_multi_thread()
4116 .worker_threads(tokio_worker_threads())
4117 .enable_all()
4118 .build()?;
4119 let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
4120
4121 let reason = match &result {
4122 Ok(()) => "normal".to_string(),
4123 Err(e) => e.to_string(),
4124 };
4125 journal.append(&journal::Record::session_exit(reason));
4126
4127 runtime.shutdown_background();
4132
4133 result
4134}
4135
4136async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
4137 let _ = tracing_subscriber::fmt::try_init();
4143
4144 let addr = std::env::var("FRINK_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
4154 let api_key_configured = std::env::var("FRINK_API_KEY").is_ok();
4155 let allow_unauthenticated_remote = std::env::var("FRINK_ALLOW_UNAUTHENTICATED_REMOTE")
4156 .map(|v| v == "1")
4157 .unwrap_or(false);
4158 if let Err(msg) =
4159 security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
4160 {
4161 anyhow::bail!(msg);
4162 }
4163
4164 let embedding_model = load_embedding_model()?;
4171
4172 let mut loaded = model::load()?;
4173 match &loaded {
4174 model::LoadedModel::Gguf(g) => tracing::info!(
4175 "loaded GGUF model '{}' (synthetic={}, tokenizer={})",
4176 g.decoder.config.name,
4177 g.is_synthetic,
4178 g.tokenizer.kind()
4179 ),
4180 model::LoadedModel::Kimi(k) => tracing::info!(
4181 "loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
4182 k.tokenizer.vocab_size()
4183 ),
4184 model::LoadedModel::Mla(m) => tracing::info!(
4185 "loaded MLA GGUF '{}' (tokenizer={})",
4186 m.name,
4187 m.tokenizer.kind()
4188 ),
4189 model::LoadedModel::Gemma4(m) => tracing::info!(
4190 "loaded Gemma4 GGUF '{}' (tokenizer={})",
4191 m.name,
4192 m.tokenizer.kind()
4193 ),
4194 model::LoadedModel::Glm52(g) => tracing::info!(
4195 "loaded GLM-5.2 GGUF '{}' (tokenizer={})",
4196 g.name,
4197 g.tokenizer.kind()
4198 ),
4199 model::LoadedModel::Encoder(_) => {}
4202 }
4203 let metal_default_moe_budget = {
4211 #[cfg(feature = "metal")]
4212 {
4213 frink_core::metal_dense_enabled()
4214 && std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES").is_err()
4215 }
4216 #[cfg(not(feature = "metal"))]
4217 {
4218 false
4219 }
4220 };
4221 if let Ok(budget_str) = std::env::var("FRINK_GPU_VRAM_BUDGET_BYTES") {
4222 let budget: u64 = budget_str
4223 .parse()
4224 .expect("FRINK_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
4225 match &mut loaded {
4226 model::LoadedModel::Gguf(g) => {
4227 tracing::info!(
4228 "GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
4229 (CUDA and/or Metal matvecs when built with the matching feature)"
4230 );
4231 g.decoder.gpu_vram_budget_bytes = Some(budget);
4232 }
4233 model::LoadedModel::Kimi(_) => {
4234 tracing::warn!(
4235 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
4236 supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
4237 );
4238 }
4239 model::LoadedModel::Mla(_) => {
4240 tracing::warn!(
4241 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
4242 FFN path only today; ignoring expert VRAM budget"
4243 );
4244 }
4245 model::LoadedModel::Gemma4(_) => {
4246 tracing::warn!(
4247 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
4248 ignoring expert VRAM budget"
4249 );
4250 }
4251 model::LoadedModel::Glm52(_) => {
4252 tracing::warn!(
4253 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
4254 GPU expert placement not wired yet; ignoring"
4255 );
4256 }
4257 model::LoadedModel::Encoder(_) => {
4258 tracing::warn!(
4259 "FRINK_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
4260 it has no routed experts to place; ignoring"
4261 );
4262 }
4263 }
4264 } else if metal_default_moe_budget {
4265 const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
4269 if let model::LoadedModel::Gguf(g) = &mut loaded {
4270 tracing::info!(
4271 "Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
4272 set FRINK_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
4273 );
4274 g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
4275 }
4276 }
4277 #[cfg(feature = "cuda")]
4278 {
4279 if frink_core::cuda_dense_enabled() {
4280 tracing::info!(
4281 "CUDA dense matvec enabled for WeightMatrix::apply \
4282 (FRINK_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
4283 );
4284 } else {
4285 tracing::info!(
4286 "CUDA dense matvec disabled (FRINK_CUDA); dense decode uses CPU or Metal"
4287 );
4288 }
4289 }
4290 #[cfg(feature = "metal")]
4291 {
4292 if frink_core::metal_dense_enabled() {
4293 tracing::info!(
4294 "Metal dense matvec enabled for WeightMatrix::apply \
4295 (FRINK_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
4296 );
4297 match std::env::var("FRINK_METAL_ATTN").ok().as_deref() {
4298 Some("1") | Some("true") | Some("on") | Some("attn") => {
4299 tracing::info!(
4300 "Metal fused attention requested (FRINK_METAL_ATTN): \
4301 QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
4302 );
4303 }
4304 _ => {}
4305 }
4306 tracing::info!(
4307 "Metal greedy GPU argmax: temperature<=0 folds \
4308 final_norm+lm_head+argmax into the dense stack"
4309 );
4310 } else {
4311 tracing::info!("Metal dense matvec disabled (FRINK_METAL); dense decode uses CPU");
4312 }
4313 }
4314 let blocks_env = std::env::var("FRINK_KV_POOL_BLOCKS");
4325 let block_size_env = std::env::var("FRINK_KV_POOL_BLOCK_SIZE");
4326 let byte_budget_env = std::env::var("FRINK_KV_BYTE_BUDGET");
4327 if blocks_env.is_ok() && byte_budget_env.is_ok() {
4328 panic!(
4329 "FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive \
4330 (set one block-count source plus FRINK_KV_POOL_BLOCK_SIZE, or neither to disable)"
4331 );
4332 }
4333 let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
4334 (Ok(blocks), Ok(block_size), Err(_)) => {
4335 let total_blocks: usize = blocks
4336 .parse()
4337 .expect("FRINK_KV_POOL_BLOCKS must be a positive integer");
4338 let block_size: usize = block_size
4339 .parse()
4340 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4341 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4347 .ok()
4348 .map(|v| {
4349 v.parse()
4350 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4351 })
4352 .unwrap_or(0);
4353 tracing::info!(
4354 "KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
4355 each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
4356 );
4357 Some(generate::KvPoolConfig {
4358 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4359 queue_wait: Duration::from_millis(queue_wait_ms),
4360 })
4361 }
4362 (Err(_), Ok(block_size), Ok(byte_budget)) => {
4363 let block_size: usize = block_size
4364 .parse()
4365 .expect("FRINK_KV_POOL_BLOCK_SIZE must be a positive integer");
4366 let budget: u64 = byte_budget
4367 .parse()
4368 .expect("FRINK_KV_BYTE_BUDGET must be a positive integer");
4369 let cfg = match &loaded {
4370 model::LoadedModel::Gguf(g) => &g.decoder.config,
4371 model::LoadedModel::Kimi(_)
4372 | model::LoadedModel::Mla(_)
4373 | model::LoadedModel::Gemma4(_)
4374 | model::LoadedModel::Glm52(_)
4375 | model::LoadedModel::Encoder(_) => {
4376 panic!(
4377 "FRINK_KV_BYTE_BUDGET requires a GGUF decoder model \
4378 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4379 );
4380 }
4381 };
4382 let bytes_per_block = block_size
4383 * cfg.kv_heads_all_layers()
4384 * (cfg.head_dim + cfg.v_head_dim())
4385 * std::mem::size_of::<f32>();
4386 assert!(
4387 bytes_per_block > 0,
4388 "derived KV block byte size must be positive (check model config and block size)"
4389 );
4390 let total_blocks = (budget as usize / bytes_per_block).max(1);
4391 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4392 .ok()
4393 .map(|v| {
4394 v.parse()
4395 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4396 })
4397 .unwrap_or(0);
4398 tracing::info!(
4399 "KV cache block pool enabled from byte budget: {budget} bytes / \
4400 {bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
4401 {total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
4402 cfg.n_layers
4403 );
4404 Some(generate::KvPoolConfig {
4405 pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
4406 queue_wait: Duration::from_millis(queue_wait_ms),
4407 })
4408 }
4409 (Err(_), Err(_), Err(_)) => None,
4410 (Err(_), Ok(_), Err(_)) => panic!(
4411 "FRINK_KV_POOL_BLOCK_SIZE requires FRINK_KV_POOL_BLOCKS or FRINK_KV_BYTE_BUDGET \
4412 (or unset all three to disable KV cache pooling)"
4413 ),
4414 (Ok(_), Ok(_), Ok(_)) => {
4415 unreachable!("FRINK_KV_POOL_BLOCKS and FRINK_KV_BYTE_BUDGET are mutually exclusive")
4416 }
4417 (Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
4418 "FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET and FRINK_KV_POOL_BLOCK_SIZE must be \
4419 set together (or neither, to disable KV cache pooling)"
4420 ),
4421 };
4422 let paged_kv = match (
4428 std::env::var("FRINK_PAGED_KV_BLOCKS"),
4429 std::env::var("FRINK_PAGED_KV_BLOCK_SIZE"),
4430 ) {
4431 (Ok(blocks), Ok(block_size)) => {
4432 assert!(
4433 kv_pool.is_none(),
4434 "FRINK_PAGED_KV_BLOCKS and FRINK_KV_POOL_BLOCKS/FRINK_KV_BYTE_BUDGET are \
4435 mutually exclusive: both bound the same KV memory, by different means. \
4436 Set one."
4437 );
4438 let blocks_per_layer: usize = blocks
4452 .parse()
4453 .expect("FRINK_PAGED_KV_BLOCKS must be a positive integer");
4454 let block_size: usize = block_size
4455 .parse()
4456 .expect("FRINK_PAGED_KV_BLOCK_SIZE must be a positive integer");
4457 let gguf = match &loaded {
4458 model::LoadedModel::Gguf(g) => g,
4459 _ => panic!(
4460 "FRINK_PAGED_KV_BLOCKS requires a GGUF decoder model \
4461 (set FRINK_MODEL_PATH to a generic-decoder .gguf file)"
4462 ),
4463 };
4464 let cfg = &gguf.decoder.config;
4465 let queue_wait_ms: u64 = std::env::var("FRINK_KV_POOL_QUEUE_TIMEOUT_MS")
4466 .ok()
4467 .map(|v| {
4468 v.parse()
4469 .expect("FRINK_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
4470 })
4471 .unwrap_or(0);
4472 tracing::info!(
4473 "Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
4474 layer across {} layers, shared by all concurrent requests, \
4475 {queue_wait_ms}ms admission queue wait",
4476 cfg.n_layers
4477 );
4478 let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
4483 block_size,
4484 ))));
4485 let anchor_token = crate::policy::anchor::resolve_anchor_token(
4494 crate::policy::parser::ToolCallFormat::infer(
4495 &std::env::var("FRINK_MODEL_PATH").unwrap_or_default(),
4496 )
4497 .opener(),
4498 |text| {
4499 gguf.tokenizer
4500 .encode(text, SpecialTokens::Parse)
4501 .into_iter()
4502 .map(|t| t as u32)
4503 .collect()
4504 },
4505 );
4506 if let Some(id) = anchor_token {
4507 tracing::info!(
4508 "Paged KV window slide: tool-call anchor is token {id}, so a turn's \
4509 window stops short of where its next turn rejoins"
4510 );
4511 }
4512 let slide_interval: usize = std::env::var("FRINK_PAGED_KV_SLIDE_INTERVAL")
4513 .ok()
4514 .map(|v| {
4515 v.parse()
4516 .expect("FRINK_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
4517 })
4518 .unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
4519 if let Some(window) = cfg.uniform_sliding_window() {
4520 tracing::info!(
4521 "Paged KV window slide enabled: every layer slides by {window} every \
4522 {slide_interval} decode steps, so a request holds its prompt and a \
4523 window rather than its whole context"
4524 );
4525 } else if cfg.kv_block_window().is_some() {
4526 tracing::info!(
4527 "Paged KV window slide NOT enabled: this model has full-attention layers, \
4528 and a page group holds one block in every layer"
4529 );
4530 }
4531 Some(generate::PagedKvConfig {
4532 store: Arc::new(cfg.new_paged_kv(block_size, blocks_per_layer)),
4535 queue_wait: Duration::from_millis(queue_wait_ms),
4536 radix,
4537 anchor_token,
4538 slide_interval,
4539 })
4540 }
4541 (Err(_), Err(_)) => None,
4542 _ => panic!(
4543 "FRINK_PAGED_KV_BLOCKS and FRINK_PAGED_KV_BLOCK_SIZE must be set together \
4544 (or neither, to disable paged KV)"
4545 ),
4546 };
4547 let prefix_cache = std::env::var("FRINK_PREFIX_CACHE_ENTRIES").ok().map(|v| {
4554 let max_entries: usize = v
4555 .parse()
4556 .expect("FRINK_PREFIX_CACHE_ENTRIES must be a positive integer");
4557 if kv_pool.is_some() {
4558 tracing::warn!(
4559 "FRINK_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
4560 caching will never be consulted while a KV pool is configured"
4561 );
4562 }
4563 assert!(
4570 paged_kv.is_none(),
4571 "FRINK_PREFIX_CACHE_ENTRIES and FRINK_PAGED_KV_BLOCKS are mutually exclusive: \
4572 the prefix cache stores contiguous KV snapshots, which a paged request does not \
4573 produce, so the cache could never hit. Set one."
4574 );
4575 tracing::info!(
4576 "KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
4577 all requests"
4578 );
4579 Arc::new(Mutex::new(PrefixCache::new(max_entries)))
4580 });
4581 if matches!(
4582 loaded,
4583 model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
4584 ) && (kv_pool.is_some() || prefix_cache.is_some())
4585 {
4586 tracing::warn!(
4587 "KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
4588 neither is consulted for those engines (state shapes differ from Decoder KV); see \
4589 frink_models::engine's module docs"
4590 );
4591 }
4592 let enable_cb =
4593 resolve_continuous_batching_enabled(&loaded, &kv_pool, &prefix_cache, &paged_kv);
4594 if enable_cb && continuous_batching_env().is_none() && metal_private_decode_active() {
4595 tracing::info!(
4596 "continuous batching enabled by default on Metal for safe parallel serving \
4597 (set FRINK_CONTINUOUS_BATCHING=0 or --no-cont-batching to use the private path)"
4598 );
4599 }
4600 if continuous_batching_env() == Some(true)
4601 && !continuous_batching_compatible(&loaded, &kv_pool, &prefix_cache, &paged_kv)
4602 && (kv_pool.is_some() || prefix_cache.is_some())
4603 {
4604 tracing::warn!(
4605 "FRINK_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
4606 (those modes keep the private generate path)"
4607 );
4608 }
4609 if let Ok(n) = std::env::var("FRINK_CHUNKED_PREFILL") {
4610 if let Ok(chunk) = n.parse::<usize>() {
4611 if chunk > 0 {
4612 tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
4613 }
4614 }
4615 }
4616 if matches!(
4617 std::env::var("FRINK_CPU_KV_OFFLOAD").ok().as_deref(),
4618 Some("1")
4619 ) {
4620 tracing::warn!(
4621 "FRINK_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
4622 (minimal spill; full layer offload still planned)"
4623 );
4624 }
4625
4626 let mcp = match mcp_config_path {
4627 Some(path) => {
4628 let loaded = mcp::load_mcp_config(&path)?;
4629 tracing::info!(
4630 "MCP config loaded from {} ({} server(s); invocation not wired yet)",
4631 loaded.path,
4632 loaded.servers.len()
4633 );
4634 Some(loaded)
4635 }
4636 None => None,
4637 };
4638
4639 let detection = health::Detection::spawn();
4643
4644 let state = Arc::new(build_app_state(
4645 StartupModels {
4646 loaded,
4647 embedding: embedding_model,
4648 },
4649 kv_pool,
4650 paged_kv,
4651 prefix_cache,
4652 enable_cb,
4653 mcp,
4654 detection,
4655 ));
4656
4657 use frink_api::routes;
4661
4662 let public = Router::new().route(routes::HEALTH, get(health));
4667
4668 let mut protected = protected_routes();
4669
4670 if let Ok(key) = std::env::var("FRINK_API_KEY") {
4675 tracing::info!("API key auth enabled");
4676 let auth = limits::AuthConfig {
4677 api_key: Arc::new(key),
4678 };
4679 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4680 auth,
4681 limits::require_api_key,
4682 ));
4683 }
4684 if let Ok(rpm) = std::env::var("FRINK_RATE_LIMIT_PER_MINUTE") {
4685 let rpm: u32 = rpm
4686 .parse()
4687 .expect("FRINK_RATE_LIMIT_PER_MINUTE must be a positive integer");
4688 tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
4689 let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
4690 protected = protected.route_layer(axum::middleware::from_fn_with_state(
4691 limiter,
4692 limits::rate_limit,
4693 ));
4694 }
4695 if let Ok(spec) = std::env::var("FRINK_CORS_ORIGINS") {
4703 let origins = security::parse_cors_origins(&spec)
4704 .unwrap_or_else(|e| panic!("FRINK_CORS_ORIGINS: {e}"));
4705 tracing::info!(
4706 "CORS enabled: {} allow-listed origin(s) ({})",
4707 origins.len(),
4708 spec
4709 );
4710 let cors = tower_http::cors::CorsLayer::new()
4711 .allow_origin(tower_http::cors::AllowOrigin::list(origins))
4712 .allow_methods([axum::http::Method::GET, axum::http::Method::POST])
4713 .allow_headers([
4714 axum::http::header::CONTENT_TYPE,
4715 axum::http::header::AUTHORIZATION,
4716 axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
4722 axum::http::HeaderName::from_static("last-event-id"),
4726 ]);
4727 protected = protected.route_layer(cors);
4728 }
4729
4730 let app = public
4734 .merge(protected)
4735 .layer(axum::middleware::from_fn(limits::retry_after))
4736 .with_state(state);
4737
4738 let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
4744 install_ring_crypto_provider();
4745 match tls_paths {
4750 Some(paths) => {
4751 let config =
4752 axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
4753 .await
4754 .map_err(|e| {
4755 anyhow::anyhow!(
4756 "failed to load TLS cert/key ({:?}, {:?}): {e}",
4757 paths.cert,
4758 paths.key
4759 )
4760 })?;
4761 let socket_addr: std::net::SocketAddr = addr
4762 .parse()
4763 .map_err(|e| anyhow::anyhow!("invalid FRINK_ADDR {addr:?} for TLS: {e}"))?;
4764 let listener = std::net::TcpListener::bind(socket_addr)?;
4765 listener.set_nonblocking(true)?;
4773 let bound = listener.local_addr()?;
4774 tracing::info!("TLS enabled: frink-server listening on https://{bound}");
4775 announce_ready(bound, "https");
4776
4777 let handle = axum_server::Handle::new();
4778 let shutdown_handle = handle.clone();
4779 tokio::spawn(async move {
4780 shutdown_signal(exit_on_stdin_close).await;
4781 shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
4782 });
4783 axum_server::from_tcp_rustls(listener, config)?
4784 .handle(handle)
4785 .serve(app.into_make_service())
4786 .await?;
4787 }
4788 None => {
4789 let listener = tokio::net::TcpListener::bind(&addr).await?;
4790 let bound = listener.local_addr()?;
4791 tracing::info!("frink-server listening on {bound}");
4792 announce_ready(bound, "http");
4793 axum::serve(listener, app)
4794 .with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
4795 .await?;
4796 }
4797 }
4798 Ok(())
4799}
4800
4801#[cfg(test)]
4802pub(crate) mod tests {
4803 use super::*;
4804 use frink_models::config::test_dense_fixture;
4805
4806 #[test]
4807 fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
4808 let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
4809 let ready = frink_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
4810 let parsed = frink_api::ServerReady::from_line(&ready.to_line()).unwrap();
4811 assert_eq!(parsed.port, 51999);
4812 assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
4813 assert!(frink_api::ServerReady::from_line("INFO frink-server listening").is_none());
4815 }
4816
4817 fn test_model() -> Model {
4818 let cfg = test_dense_fixture();
4822 Model::Gguf(GgufModel {
4823 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
4824 tokenizer: Arc::new(ServerTokenizer::Byte),
4825 stop_tokens: StopTokens::default(),
4826 bos_id: None,
4827 is_synthetic: true,
4828 chat_template: chat_template::PromptTemplate::plain(),
4829 })
4830 }
4831
4832 fn greedy_params(max_tokens: usize) -> GenerationParams {
4833 GenerationParams {
4834 wants_logprobs: false,
4835 n: 1,
4836 reasoning: None,
4837 max_tokens,
4838 sampling: SamplingParams::default(),
4839 seed: 1,
4840 stop: Vec::new(),
4841 stop_token_ids: Vec::new(),
4842 json_object: false,
4843 grammar: None,
4844 cancel: None,
4845 ignore_eos: false,
4846 reasoning_budget: crate::reasoning_budget::ReasoningBudget::Unrestricted,
4847 lora: None,
4848 }
4849 }
4850
4851 fn test_model_full_byte_vocab() -> Model {
4855 test_model_full_byte_vocab_with_eos(None)
4856 }
4857
4858 fn test_model_full_byte_vocab_with_eos(eos: Option<usize>) -> Model {
4865 let mut cfg = test_dense_fixture();
4866 cfg.vocab_size = 256;
4867 Model::Gguf(GgufModel {
4868 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4869 tokenizer: Arc::new(ServerTokenizer::Byte),
4870 stop_tokens: StopTokens::from_eos(eos),
4871 bos_id: None,
4872 is_synthetic: true,
4873 chat_template: chat_template::PromptTemplate::plain(),
4874 })
4875 }
4876
4877 pub(crate) fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
4881 AppState {
4882 embedding: None,
4883 paged_kv: None,
4884 active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
4885 id: None,
4886 loaded: Loaded::Generative(Arc::new(model)),
4887 batcher: None,
4888 ceiling: None,
4889 checkpoint_path: None,
4890 }))),
4891 load_in_progress: std::sync::atomic::AtomicBool::new(false),
4892 tasks: Arc::new(tasks::TaskRegistry::new()),
4893 cancels: Arc::new(cancel::CancelRegistry::new()),
4894 stats: stats::Stats::new(),
4895 streams: resume::StreamRegistry::new(),
4896 model_dir: None,
4897 response_cache: Mutex::new(response_cache),
4898 kv_pool: None,
4899 prefix_cache: None,
4900 sessions: session::SessionStore::new(),
4901 requests_total: std::sync::atomic::AtomicU64::new(0),
4902 request_errors_total: std::sync::atomic::AtomicU64::new(0),
4903 started_at: std::time::Instant::now(),
4904 last_request_ms: std::sync::atomic::AtomicU64::new(0),
4905 detection: Arc::new(health::Detection::ready(health::probe_backends())),
4906 mcp: None,
4907 continuous_batching_enabled: false,
4908 metal_private_decode_gate: None,
4909 loading_model: Mutex::new(None),
4910 last_load_error: Mutex::new(None),
4911 serving: Mutex::new(crate::stats::ServingStats::default()),
4912 maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
4913 footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
4914 started_unix: unix_now(),
4915 }
4916 }
4917
4918 pub(crate) fn test_app() -> Router {
4927 test_app_with_state(Arc::new(test_state(
4928 test_model_full_byte_vocab(),
4929 ResponseCache::new(1000, Duration::from_secs(3600)),
4930 )))
4931 }
4932
4933 pub(crate) fn test_app_with_state(state: Arc<AppState>) -> Router {
4936 Router::new()
4946 .route(frink_api::routes::HEALTH, get(health))
4947 .merge(protected_routes())
4948 .with_state(state)
4949 }
4950
4951 fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
4952 let mut cfg = test_dense_fixture();
4953 cfg.name = name;
4954 cfg.vocab_size = vocab_size;
4955 Model::Gguf(GgufModel {
4956 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4957 tokenizer: Arc::new(ServerTokenizer::Byte),
4958 stop_tokens: StopTokens::default(),
4959 bos_id: None,
4960 is_synthetic: true,
4961 chat_template: chat_template::PromptTemplate::plain(),
4962 })
4963 }
4964
4965 fn model_with_template(name: &'static str, source: &str) -> Model {
4969 let mut cfg = test_dense_fixture();
4970 cfg.name = name;
4971 cfg.vocab_size = 256;
4972 Model::Gguf(GgufModel {
4973 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
4974 tokenizer: Arc::new(ServerTokenizer::Byte),
4975 stop_tokens: StopTokens::default(),
4976 bos_id: None,
4977 is_synthetic: true,
4978 chat_template: chat_template::PromptTemplate::from_gguf_metadata(
4979 Some(source),
4980 Some("qwen3"),
4981 false,
4982 true,
4983 None,
4984 None,
4985 ),
4986 })
4987 }
4988
4989 #[tokio::test]
4998 async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
4999 let strict = "{% if messages | length > 1 %}\
5002 {{ raise_exception('this template takes one turn') }}\
5003 {% endif %}{{ messages[0].content }}";
5004 let state = Arc::new(test_state(
5005 model_with_template("strict", strict),
5006 ResponseCache::new(4, Duration::from_secs(60)),
5007 ));
5008 let app = test_app_with_state(state);
5009
5010 let (status, body) = post_json_uri(
5011 &app,
5012 "/v1/chat/completions",
5013 serde_json::json!({
5014 "model": "strict",
5015 "stream": true,
5016 "messages": [
5017 {"role": "user", "content": "one"},
5018 {"role": "user", "content": "two"},
5019 ],
5020 }),
5021 )
5022 .await;
5023 assert_eq!(status, StatusCode::BAD_REQUEST);
5024 assert_eq!(body["error"]["param"], serde_json::json!("messages"));
5025 assert!(
5026 body["error"]["message"]
5027 .as_str()
5028 .unwrap()
5029 .contains("one turn"),
5030 "the template's own message must reach the caller: {body}"
5031 );
5032
5033 let (status, _) = post_json_uri(
5035 &app,
5036 "/v1/chat/completions",
5037 serde_json::json!({
5038 "model": "strict",
5039 "stream": true,
5040 "max_tokens": 1,
5041 "messages": [{"role": "user", "content": "one"}],
5042 }),
5043 )
5044 .await;
5045 assert_eq!(status, StatusCode::OK);
5046 }
5047
5048 #[tokio::test]
5050 async fn models_advertises_the_gears_this_checkpoint_actually_has() {
5051 let reasoning = "{% if enable_thinking %}<think>{% endif %}\
5052 {% if reasoning_effort %}\
5053 {% if reasoning_effort not in ['low','medium','high'] %}\
5054 {{ raise_exception('bad effort') }}\
5055 {% endif %}[{{ reasoning_effort }}]\
5056 {% endif %}{{ messages[0].content }}";
5057 let state = Arc::new(test_state(
5058 model_with_template("thinker", reasoning),
5059 ResponseCache::new(4, Duration::from_secs(60)),
5060 ));
5061 let app = test_app_with_state(state);
5062 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5063 assert_eq!(status, StatusCode::OK);
5064 let entry = &models["data"][0];
5065 assert_eq!(
5066 entry["supported_reasoning_efforts"],
5067 serde_json::json!(["off", "low", "medium", "high"])
5068 );
5069 assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
5070 }
5071
5072 #[tokio::test]
5077 async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
5078 let app = test_app();
5079 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5080 let entry = &models["data"][0];
5081 assert!(entry.get("supported_reasoning_efforts").is_none());
5082 assert!(entry.get("default_reasoning_effort").is_none());
5083 }
5084
5085 fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
5086 Arc::new(ActiveModel {
5087 id: Some(name.to_string()),
5088 loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
5089 batcher: None,
5090 ceiling: None,
5091 checkpoint_path: None,
5092 })
5093 .tap_into(state)
5094 }
5095
5096 trait TapInto {
5098 fn tap_into(self, state: &AppState) -> Self;
5099 }
5100 impl TapInto for Arc<ActiveModel> {
5101 fn tap_into(self, state: &AppState) -> Self {
5102 state.swap_active(Some(Arc::clone(&self)));
5103 self
5104 }
5105 }
5106
5107 #[test]
5113 fn an_in_flight_request_keeps_the_model_it_started_on() {
5114 let state = test_state(
5115 named_test_model("model-a", 256),
5116 ResponseCache::new(4, Duration::from_secs(60)),
5117 );
5118
5119 let in_flight = state.active().expect("a model is loaded");
5122 assert_eq!(in_flight.name(), "model-a");
5123
5124 active_model(&state, "model-b");
5125
5126 assert_eq!(state.active().unwrap().name(), "model-b");
5128 assert_eq!(in_flight.name(), "model-a");
5130 let (choices, _usage) = run_generation(
5131 in_flight.generative().unwrap(),
5132 "hi",
5133 &greedy_params(3),
5134 None,
5135 None,
5136 None,
5137 None,
5138 None,
5139 None,
5140 )
5141 .expect("the old model must still decode after being swapped out");
5142 assert!(matches!(
5143 choices[0].finish,
5144 FinishReason::Length | FinishReason::Stop
5145 ));
5146 }
5147
5148 #[test]
5153 fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
5154 let state = test_state(
5155 named_test_model("model-a", 256),
5156 ResponseCache::new(4, Duration::from_secs(60)),
5157 );
5158 let in_flight = state.active().expect("a model is loaded");
5159 let weights = Arc::clone(in_flight.generative().unwrap());
5160 assert!(Arc::strong_count(&weights) >= 2);
5161
5162 let previous = state.swap_active(Some(Arc::new(ActiveModel {
5163 id: Some("model-b".to_string()),
5164 loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
5165 batcher: None,
5166 ceiling: None,
5167 checkpoint_path: None,
5168 })));
5169 drop(previous);
5170 assert!(Arc::strong_count(&weights) >= 2);
5172 drop(in_flight);
5173 assert_eq!(Arc::strong_count(&weights), 1);
5174 }
5175
5176 #[tokio::test]
5180 async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
5181 let state = Arc::new(test_state(
5182 named_test_model("model-a", 256),
5183 ResponseCache::new(4, Duration::from_secs(60)),
5184 ));
5185 let app = test_app_with_state(Arc::clone(&state));
5186
5187 let (status, body) = post_json_uri(
5188 &app,
5189 frink_api::routes::ADMIN_MODELS_UNLOAD,
5190 serde_json::json!({}),
5191 )
5192 .await;
5193 assert_eq!(status, StatusCode::OK);
5194 assert_eq!(body["ok"], true);
5195 assert!(body["active"].is_null());
5196 assert!(state.active().is_none());
5197
5198 let (status, _) = get_json(&app, frink_api::routes::V1_MODELS).await;
5199 assert_eq!(status, StatusCode::OK);
5200 let (_, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
5201 assert_eq!(models["data"].as_array().unwrap().len(), 0);
5202
5203 let (status, body) = post_json_uri(
5204 &app,
5205 "/v1/chat/completions",
5206 serde_json::json!({
5207 "model": "x",
5208 "messages": [{"role": "user", "content": "hi"}]
5209 }),
5210 )
5211 .await;
5212 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5213 assert_eq!(body["error"]["type"], "model_not_loaded");
5214 }
5215
5216 #[tokio::test]
5220 async fn health_reports_the_unloaded_state_rather_than_going_silent() {
5221 let state = Arc::new(test_state(
5222 named_test_model("model-a", 256),
5223 ResponseCache::new(4, Duration::from_secs(60)),
5224 ));
5225 let app = test_app_with_state(Arc::clone(&state));
5226 state.swap_active(None);
5227
5228 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
5229 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5232 assert_eq!(body["state"], "unavailable");
5233 assert_eq!(body["reason"], "model_not_loaded");
5234 assert!(body["model"].is_null());
5235 let real_weights = body["capabilities"]
5236 .as_array()
5237 .unwrap()
5238 .iter()
5239 .find(|c| c["id"] == "real_weights")
5240 .cloned()
5241 .expect("real_weights is always reported");
5242 assert_eq!(real_weights["available"], false);
5243 assert_eq!(real_weights["reason"], "model_not_loaded");
5244 }
5245
5246 #[tokio::test]
5250 async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
5251 let app = test_app();
5252
5253 let (status, completion) = post_json_uri(
5254 &app,
5255 "/v1/chat/completions",
5256 serde_json::json!({
5257 "model": "x",
5258 "messages": [{"role": "user", "content": "hi"}],
5259 "max_tokens": 4
5260 }),
5261 )
5262 .await;
5263 assert_eq!(status, StatusCode::OK);
5264 let request_id = completion["request_id"].as_str().unwrap().to_string();
5265
5266 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5267 assert_eq!(status, StatusCode::OK);
5268 let recent = stats["recent"].as_array().unwrap();
5269 assert_eq!(recent.len(), 1);
5270 let row = &recent[0];
5271 assert_eq!(row["request_id"], request_id);
5272 assert_eq!(row["route"], frink_api::routes::V1_CHAT_COMPLETIONS);
5273 assert_eq!(row["status"], 200);
5274 assert_eq!(row["stream"], false);
5275 assert!(row["duration_ms"].is_number());
5278 assert!(row["decode_ms"].is_number());
5279 assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
5280 assert_eq!(
5281 stats["tokens_prompt_total"].as_u64().unwrap(),
5282 row["prompt_tokens"].as_u64().unwrap()
5283 );
5284 }
5285
5286 #[tokio::test]
5290 async fn a_rejected_request_is_recorded_too() {
5291 let state = Arc::new(test_state(
5292 named_test_model("model-a", 256),
5293 ResponseCache::new(4, Duration::from_secs(60)),
5294 ));
5295 let app = test_app_with_state(Arc::clone(&state));
5296 state.swap_active(None);
5297
5298 let (status, _) = post_json_uri(
5299 &app,
5300 "/v1/chat/completions",
5301 serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
5302 )
5303 .await;
5304 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
5305
5306 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5307 let recent = stats["recent"].as_array().unwrap();
5308 assert_eq!(recent.len(), 1);
5309 assert_eq!(recent[0]["status"], 503);
5310 assert_eq!(recent[0]["completion_tokens"], 0);
5311 assert!(recent[0]["decode_ms"].is_null());
5312 assert_eq!(stats["errors_total"], 1);
5313 }
5314
5315 async fn post_json_with_headers(
5319 app: &Router,
5320 uri: &str,
5321 body: serde_json::Value,
5322 headers: &[(&str, &str)],
5323 ) -> (StatusCode, serde_json::Value) {
5324 use http_body_util::BodyExt;
5325 use tower::ServiceExt;
5326
5327 let mut builder = axum::http::Request::builder()
5328 .method("POST")
5329 .uri(uri)
5330 .header("content-type", "application/json");
5331 for (name, value) in headers {
5332 builder = builder.header(*name, *value);
5333 }
5334 let response = app
5335 .clone()
5336 .oneshot(
5337 builder
5338 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
5339 .unwrap(),
5340 )
5341 .await
5342 .unwrap();
5343 let status = response.status();
5344 let bytes = response.into_body().collect().await.unwrap().to_bytes();
5345 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
5346 (status, json)
5347 }
5348
5349 #[tokio::test]
5353 async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
5354 let app = test_app();
5355
5356 let (status, _) = post_json_uri(
5357 &app,
5358 frink_api::routes::V1_TOKENIZE,
5359 serde_json::json!({"prompt": "hello"}),
5360 )
5361 .await;
5362 assert_eq!(status, StatusCode::OK);
5363 let (status, _) = post_json_uri(
5364 &app,
5365 frink_api::routes::V1_DETOKENIZE,
5366 serde_json::json!({"tokens": [104, 105]}),
5367 )
5368 .await;
5369 assert_eq!(status, StatusCode::OK);
5370 let (status, _) = post_json_uri(
5371 &app,
5372 frink_api::routes::V1_EMBEDDINGS,
5373 serde_json::json!({"input": "hello"}),
5374 )
5375 .await;
5376 assert_eq!(status, StatusCode::OK);
5377
5378 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5379 let routes: Vec<&str> = stats["recent"]
5380 .as_array()
5381 .unwrap()
5382 .iter()
5383 .map(|row| row["route"].as_str().unwrap())
5384 .collect();
5385 for expected in [
5386 frink_api::routes::V1_TOKENIZE,
5387 frink_api::routes::V1_DETOKENIZE,
5388 frink_api::routes::V1_EMBEDDINGS,
5389 ] {
5390 assert!(
5391 routes.contains(&expected),
5392 "{expected} is missing: {routes:?}"
5393 );
5394 }
5395
5396 let row = |route: &str| {
5397 stats["recent"]
5398 .as_array()
5399 .unwrap()
5400 .iter()
5401 .find(|r| r["route"] == route)
5402 .cloned()
5403 .unwrap()
5404 };
5405 let embed = row(frink_api::routes::V1_EMBEDDINGS);
5409 assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
5410 assert!(embed["decode_ms"].is_null());
5411 assert_eq!(embed["completion_tokens"], 0);
5412 assert_eq!(row(frink_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
5416 assert_eq!(
5417 stats["tokens_prompt_total"].as_u64().unwrap(),
5418 embed["prompt_tokens"].as_u64().unwrap(),
5419 "only the forward pass counted"
5420 );
5421 }
5422
5423 fn streaming_test_app() -> Router {
5428 let mut cfg = test_dense_fixture();
5429 cfg.vocab_size = 256;
5430 let model = Model::Gguf(GgufModel {
5431 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5432 tokenizer: Arc::new(ServerTokenizer::Byte),
5433 stop_tokens: StopTokens::default(),
5434 bos_id: None,
5435 is_synthetic: false,
5436 chat_template: chat_template::PromptTemplate::plain(),
5437 });
5438 test_app_with_state(Arc::new(test_state(
5439 model,
5440 ResponseCache::new(1000, Duration::from_secs(3600)),
5441 )))
5442 }
5443
5444 #[tokio::test]
5453 async fn chat_logprobs_are_rendered_and_are_never_served_from_cache() {
5454 let app = test_app();
5455 let body = |logprobs: Option<(bool, Option<u32>)>| {
5456 let mut b = serde_json::json!({
5457 "model": "x",
5458 "messages": [{"role": "user", "content": "hi"}],
5459 "max_tokens": 4
5460 });
5461 if let Some((on, top)) = logprobs {
5462 b["logprobs"] = serde_json::json!(on);
5463 if let Some(n) = top {
5464 b["top_logprobs"] = serde_json::json!(n);
5465 }
5466 }
5467 b
5468 };
5469
5470 let (status, plain) =
5472 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(None)).await;
5473 assert_eq!(status, StatusCode::OK, "{plain}");
5474 assert!(plain["choices"][0]["logprobs"].is_null(), "{plain}");
5475
5476 for attempt in 0..2 {
5480 let (status, with) = post_json_uri(
5481 &app,
5482 frink_api::routes::V1_CHAT_COMPLETIONS,
5483 body(Some((true, Some(2)))),
5484 )
5485 .await;
5486 assert_eq!(status, StatusCode::OK, "{with}");
5487 assert_ne!(
5488 with["frink_cache"], "hit",
5489 "attempt {attempt} replayed a cached answer for a logprobs request: {with}"
5490 );
5491 let lp = &with["choices"][0]["logprobs"];
5492 assert!(lp.is_object(), "attempt {attempt}: {with}");
5493 let content = lp["content"].as_array().expect("content");
5494 assert!(lp["tokens"].is_null(), "completions shape leaked: {lp}");
5496 for entry in content {
5497 assert!(entry["token"].is_string(), "{entry}");
5498 assert!(entry["bytes"].is_array(), "{entry}");
5499 let v = entry["logprob"].as_f64().expect("a real number");
5500 assert!(v <= 0.0 && v.is_finite(), "{entry}");
5501 let top = entry["top_logprobs"].as_array().expect("top_logprobs");
5502 assert!(top.len() <= 2, "asked for 2, got {}", top.len());
5503 }
5504 }
5505 }
5506
5507 #[tokio::test]
5513 async fn the_chat_logprobs_pair_is_validated() {
5514 let app = test_app();
5515 for (extra, why) in [
5516 (serde_json::json!({"top_logprobs": 3}), "without logprobs"),
5517 (
5518 serde_json::json!({"logprobs": true, "top_logprobs": 21}),
5519 "above the cap",
5520 ),
5521 ] {
5522 let mut body = serde_json::json!({
5523 "model": "x",
5524 "messages": [{"role": "user", "content": "hi"}],
5525 "max_tokens": 2
5526 });
5527 for (k, v) in extra.as_object().unwrap() {
5528 body[k] = v.clone();
5529 }
5530 let (status, answer) =
5531 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await;
5532 assert_eq!(status, StatusCode::BAD_REQUEST, "{why}: {answer}");
5533 assert!(
5534 answer["error"]["message"]
5535 .as_str()
5536 .is_some_and(|m| m.contains("top_logprobs")),
5537 "{why}: {answer}"
5538 );
5539 }
5540 }
5541
5542 #[tokio::test]
5547 async fn chat_serves_several_choices_and_refuses_the_streaming_pair() {
5548 let app = test_app();
5549 let body = |n: u32, stream: bool| {
5550 serde_json::json!({
5551 "model": "x",
5552 "messages": [{"role": "user", "content": "hi"}],
5553 "max_tokens": 4,
5554 "temperature": 1.0,
5555 "n": n,
5556 "stream": stream
5557 })
5558 };
5559
5560 let (status, one) =
5561 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(1, false)).await;
5562 assert_eq!(status, StatusCode::OK, "{one}");
5563
5564 let (status, three) =
5565 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, false)).await;
5566 assert_eq!(status, StatusCode::OK, "{three}");
5567 let choices = three["choices"].as_array().expect("an array");
5568 assert_eq!(choices.len(), 3, "{three}");
5569 for (i, c) in choices.iter().enumerate() {
5570 assert_eq!(c["index"], i);
5571 assert!(c["message"]["role"].is_string(), "{c}");
5572 assert!(c["finish_reason"].is_string(), "{c}");
5573 }
5574 assert_eq!(
5576 three["usage"]["prompt_tokens"], one["usage"]["prompt_tokens"],
5577 "n = 3 billed the prompt more than once"
5578 );
5579
5580 let (status, refused) =
5583 post_json_uri(&app, frink_api::routes::V1_CHAT_COMPLETIONS, body(3, true)).await;
5584 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{refused}");
5585 let message = refused["error"]["message"].as_str().unwrap_or_default();
5586 assert!(
5587 message.contains('n') && message.contains("stream"),
5588 "{refused}"
5589 );
5590 }
5591
5592 #[tokio::test]
5602 async fn every_route_refuses_the_same_unimplemented_fields() {
5603 let app = test_app();
5604 let fields = [
5605 ("n", serde_json::json!(3)),
5606 ("best_of", serde_json::json!(2)),
5607 ("prompt_logprobs", serde_json::json!(1)),
5608 ("echo", serde_json::json!(true)),
5609 ("use_beam_search", serde_json::json!(true)),
5610 ("truncate_prompt_tokens", serde_json::json!(8)),
5611 ("prompt_embeds", serde_json::json!("AA==")),
5612 ("allowed_token_ids", serde_json::json!([1, 2])),
5613 ("bad_words", serde_json::json!(["x"])),
5614 ("skip_special_tokens", serde_json::json!(false)),
5615 ("return_tokens_as_token_ids", serde_json::json!(true)),
5616 ];
5617 for (field, value) in fields {
5618 for (uri, base) in [
5619 (
5620 frink_api::routes::V1_CHAT_COMPLETIONS,
5621 serde_json::json!({
5622 "model": "x",
5623 "messages": [{"role": "user", "content": "hi"}],
5624 "max_tokens": 2
5625 }),
5626 ),
5627 (
5628 frink_api::routes::V1_COMPLETIONS,
5629 serde_json::json!({"prompt": "hi", "max_tokens": 2}),
5630 ),
5631 (
5632 frink_api::routes::COMPLETION,
5633 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5634 ),
5635 ] {
5636 let mut body = base;
5637 body[field] = value.clone();
5638 if field == "n"
5643 && (uri == frink_api::routes::V1_COMPLETIONS
5644 || uri == frink_api::routes::V1_CHAT_COMPLETIONS)
5645 {
5646 let (status, answer) = post_json_uri(&app, uri, body).await;
5647 assert_eq!(
5648 status,
5649 StatusCode::OK,
5650 "{uri} refused a served `n`: {answer}"
5651 );
5652 assert_eq!(
5653 answer["choices"].as_array().map(Vec::len),
5654 Some(3),
5655 "{answer}"
5656 );
5657 continue;
5658 }
5659 let (status, answer) = post_json_uri(&app, uri, body).await;
5660 assert_eq!(
5661 status,
5662 StatusCode::NOT_IMPLEMENTED,
5663 "{uri} served `{field}` instead of refusing it: {answer}"
5664 );
5665 assert!(
5666 answer["error"]["message"]
5667 .as_str()
5668 .is_some_and(|m| m.contains(field)),
5669 "{uri} refused `{field}` without naming it: {answer}"
5670 );
5671 }
5672 }
5673 }
5674
5675 #[tokio::test]
5676 async fn the_native_completion_wire_is_not_the_openai_one() {
5677 let app = test_app();
5678
5679 let (status, native) = post_json_uri(
5680 &app,
5681 frink_api::routes::COMPLETION,
5682 serde_json::json!({"prompt": "hi", "n_predict": 4}),
5683 )
5684 .await;
5685 assert_eq!(status, StatusCode::OK, "{native}");
5686 assert!(native["content"].is_string(), "{native}");
5687 assert_eq!(native["stop"], true);
5688 assert_eq!(native["stop_type"], "limit");
5689 assert_eq!(native["stopping_word"], "");
5690 assert_eq!(native["truncated"], false);
5691 assert_eq!(native["id_slot"], -1);
5692 assert!(native["timings"]["prompt_n"].is_number(), "{native}");
5693 assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
5694 assert!(
5695 native.get("choices").is_none(),
5696 "the native shape has no `choices`: {native}"
5697 );
5698
5699 let (status, openai) = post_json_uri(
5700 &app,
5701 frink_api::routes::V1_COMPLETIONS,
5702 serde_json::json!({"prompt": "hi", "max_tokens": 4}),
5703 )
5704 .await;
5705 assert_eq!(status, StatusCode::OK);
5706 assert!(openai["choices"][0]["text"].is_string(), "{openai}");
5707 assert!(
5708 openai.get("content").is_none(),
5709 "the OpenAI shape has no top-level `content`: {openai}"
5710 );
5711 }
5712
5713 #[tokio::test]
5717 async fn both_native_spellings_reach_the_same_handler() {
5718 let app = test_app();
5719 for route in [
5720 frink_api::routes::COMPLETION,
5721 frink_api::routes::COMPLETIONS,
5722 ] {
5723 let (status, body) = post_json_uri(
5724 &app,
5725 route,
5726 serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
5727 )
5728 .await;
5729 assert_eq!(status, StatusCode::OK, "{route}: {body}");
5730 assert_eq!(body["stop"], true, "{route}");
5731 assert!(body["content"].is_string(), "{route}");
5732 }
5733
5734 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5737 let routes: Vec<&str> = stats["recent"]
5738 .as_array()
5739 .unwrap()
5740 .iter()
5741 .map(|row| row["route"].as_str().unwrap())
5742 .collect();
5743 assert!(
5744 routes.contains(&frink_api::routes::COMPLETION),
5745 "{routes:?}"
5746 );
5747 assert!(
5748 routes.contains(&frink_api::routes::COMPLETIONS),
5749 "{routes:?}"
5750 );
5751 }
5752
5753 #[tokio::test]
5759 async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
5760 let app = streaming_test_app();
5761 let raw = post_sse_raw_uri(
5762 &app,
5763 frink_api::routes::COMPLETION,
5764 serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
5765 )
5766 .await;
5767
5768 assert!(
5769 !raw.contains("[DONE]"),
5770 "llama.cpp's native stream has no sentinel: {raw}"
5771 );
5772 let frames: Vec<serde_json::Value> = raw
5773 .lines()
5774 .filter_map(|line| line.strip_prefix("data: "))
5775 .map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
5776 .collect();
5777 assert!(frames.len() >= 2, "expected partials then a final: {raw}");
5778
5779 let (last, partials) = frames.split_last().unwrap();
5780 assert_eq!(last["stop"], true, "the last frame closes the stream");
5781 assert!(last["timings"].is_object(), "{last}");
5782 assert!(last["stop_type"].is_string(), "{last}");
5783 for partial in partials {
5784 assert_eq!(partial["stop"], false, "{partial}");
5785 assert!(partial["content"].is_string(), "{partial}");
5786 assert!(partial.get("timings").is_none(), "{partial}");
5790 assert!(partial.get("generation_settings").is_none(), "{partial}");
5791 }
5792 let streamed: String = partials
5795 .iter()
5796 .filter_map(|p| p["content"].as_str())
5797 .collect();
5798 assert_eq!(last["content"].as_str().unwrap(), streamed);
5799 }
5800
5801 #[tokio::test]
5806 async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
5807 let app = test_app();
5808 for body in [
5809 serde_json::json!({"prompt": "hi"}),
5810 serde_json::json!({"prompt": "hi", "n_predict": -1}),
5811 ] {
5812 let (status, refusal) =
5813 post_json_uri(&app, frink_api::routes::COMPLETION, body.clone()).await;
5814 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
5815 assert!(
5816 refusal["error"]["message"]
5817 .as_str()
5818 .unwrap()
5819 .contains("n_predict"),
5820 "{refusal}"
5821 );
5822 }
5823 let (status, _) = post_json_uri(
5826 &app,
5827 frink_api::routes::COMPLETION,
5828 serde_json::json!({"prompt": "hi", "n_predict": 2}),
5829 )
5830 .await;
5831 assert_eq!(status, StatusCode::OK);
5832 }
5833
5834 #[tokio::test]
5843 async fn a_stop_string_halts_the_answer_and_is_named_back() {
5844 let app = streaming_test_app();
5845 let ask = |stop: serde_json::Value| {
5846 let app = app.clone();
5847 async move {
5848 post_json_uri(
5849 &app,
5850 frink_api::routes::COMPLETION,
5851 serde_json::json!({
5852 "prompt": "hi",
5853 "n_predict": 64,
5854 "ignore_eos": true,
5855 "stop": stop,
5856 }),
5857 )
5858 .await
5859 .1
5860 }
5861 };
5862
5863 let baseline = ask(serde_json::json!([])).await;
5864 assert_eq!(baseline["stop_type"], "limit");
5865 assert_eq!(baseline["stopping_word"], "");
5866 let text = baseline["content"].as_str().unwrap().to_string();
5867 let sentinel: String = text.chars().skip(1).take(2).collect();
5873 assert_eq!(
5874 sentinel.chars().count(),
5875 2,
5876 "the fixture must produce enough output to cut: {text:?}"
5877 );
5878 let cut = text.find(&sentinel).expect("it came out of this text");
5879
5880 let stopped = ask(serde_json::json!([sentinel])).await;
5881 assert_eq!(stopped["stop_type"], "word", "{stopped}");
5882 assert_eq!(stopped["stopping_word"], sentinel);
5883 assert_eq!(
5884 stopped["content"].as_str().unwrap(),
5885 &text[..cut],
5886 "the answer must be cut at the sentinel, not run past it"
5887 );
5888 }
5889
5890 #[tokio::test]
5896 async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
5897 let app = test_app();
5898
5899 let (v1_status, v1) = post_json_uri(
5900 &app,
5901 frink_api::routes::V1_TOKENIZE,
5902 serde_json::json!({"prompt": "hello"}),
5903 )
5904 .await;
5905 let (alias_status, alias) = post_json_uri(
5906 &app,
5907 frink_api::routes::TOKENIZE,
5908 serde_json::json!({"content": "hello"}),
5909 )
5910 .await;
5911 assert_eq!(v1_status, StatusCode::OK);
5912 assert_eq!(alias_status, StatusCode::OK, "{alias}");
5913 assert_eq!(v1["tokens"], alias["tokens"]);
5914 assert!(!alias["tokens"].as_array().unwrap().is_empty());
5915
5916 let (status, both_ways) = post_json_uri(
5919 &app,
5920 frink_api::routes::TOKENIZE,
5921 serde_json::json!({"prompt": "hello"}),
5922 )
5923 .await;
5924 assert_eq!(status, StatusCode::OK);
5925 assert_eq!(both_ways["tokens"], v1["tokens"]);
5926 }
5927
5928 #[tokio::test]
5933 async fn detokenize_answers_under_both_dialects_keys() {
5934 let app = test_app();
5935 for route in [
5936 frink_api::routes::DETOKENIZE,
5937 frink_api::routes::V1_DETOKENIZE,
5938 ] {
5939 let (status, body) =
5940 post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
5941 assert_eq!(status, StatusCode::OK, "{route}");
5942 assert_eq!(body["text"], "hi", "{route}");
5943 assert_eq!(body["content"], body["text"], "{route}");
5944 }
5945 }
5946
5947 #[tokio::test]
5951 async fn the_alias_is_recorded_under_the_path_the_client_called() {
5952 let app = test_app();
5953 let (status, _) = post_json_uri(
5954 &app,
5955 frink_api::routes::TOKENIZE,
5956 serde_json::json!({"content": "hello"}),
5957 )
5958 .await;
5959 assert_eq!(status, StatusCode::OK);
5960
5961 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
5962 let routes: Vec<&str> = stats["recent"]
5963 .as_array()
5964 .unwrap()
5965 .iter()
5966 .map(|row| row["route"].as_str().unwrap())
5967 .collect();
5968 assert!(
5969 routes.contains(&frink_api::routes::TOKENIZE),
5970 "the alias must be its own row: {routes:?}"
5971 );
5972 assert!(
5973 !routes.contains(&frink_api::routes::V1_TOKENIZE),
5974 "nothing called /v1/tokenize: {routes:?}"
5975 );
5976 }
5977
5978 #[tokio::test]
5983 async fn add_special_prepends_the_same_bos_the_decoder_would() {
5984 let mut cfg = test_dense_fixture();
5985 cfg.vocab_size = 256;
5986 let model = Model::Gguf(GgufModel {
5987 decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
5988 tokenizer: Arc::new(ServerTokenizer::Byte),
5989 stop_tokens: StopTokens::default(),
5990 bos_id: Some(7),
5991 is_synthetic: true,
5992 chat_template: chat_template::PromptTemplate::plain(),
5993 });
5994 let app = test_app_with_state(Arc::new(test_state(
5995 model,
5996 ResponseCache::new(1000, Duration::from_secs(3600)),
5997 )));
5998
5999 let (_, plain) = post_json_uri(
6000 &app,
6001 frink_api::routes::TOKENIZE,
6002 serde_json::json!({"content": "hi"}),
6003 )
6004 .await;
6005 let (_, special) = post_json_uri(
6006 &app,
6007 frink_api::routes::TOKENIZE,
6008 serde_json::json!({"content": "hi", "add_special": true}),
6009 )
6010 .await;
6011
6012 assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
6013 assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
6014 assert_eq!(special["count"], 3);
6015 }
6016
6017 #[tokio::test]
6020 async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
6021 let app = test_app();
6022 let (status, _) = post_json_uri(
6023 &app,
6024 frink_api::routes::V1_EMBEDDINGS,
6025 serde_json::json!({"input": "hi", "encoding_format": "base64"}),
6026 )
6027 .await;
6028 assert_eq!(status, StatusCode::BAD_REQUEST);
6029
6030 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6031 let recent = stats["recent"].as_array().unwrap();
6032 assert_eq!(recent.len(), 1);
6033 assert_eq!(recent[0]["route"], frink_api::routes::V1_EMBEDDINGS);
6034 assert_eq!(recent[0]["status"], 400);
6035 assert_eq!(
6036 recent[0]["prompt_tokens"], 0,
6037 "a rejected call embedded nothing"
6038 );
6039 }
6040
6041 #[tokio::test]
6044 async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
6045 let app = test_app();
6046 let key = "sk-monitor-secret";
6047 let (status, _) = post_json_with_headers(
6048 &app,
6049 "/v1/chat/completions",
6050 serde_json::json!({
6051 "model": "x",
6052 "messages": [{"role": "user", "content": "hi"}],
6053 "max_tokens": 2
6054 }),
6055 &[
6056 ("authorization", &format!("Bearer {key}")),
6057 ("x-frink-client", "frink-studio"),
6058 ],
6059 )
6060 .await;
6061 assert_eq!(status, StatusCode::OK);
6062
6063 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6064 let row = stats["recent"].as_array().unwrap()[0].clone();
6065 let fingerprint = row["via_api_key"]
6066 .as_str()
6067 .expect("the row names the key that served it")
6068 .to_string();
6069 assert_eq!(fingerprint, attribution::key_fingerprint(key));
6070 assert!(!fingerprint.contains(key));
6071 assert!(
6072 !serde_json::to_string(&stats).unwrap().contains(key),
6073 "the stats payload must not carry the key in any form"
6074 );
6075 assert_eq!(row["client"], "frink-studio");
6076 }
6077
6078 #[tokio::test]
6081 async fn different_keys_are_different_callers_and_no_key_is_null() {
6082 let app = test_app();
6083 let body = serde_json::json!({
6084 "model": "x",
6085 "messages": [{"role": "user", "content": "hi"}],
6086 "max_tokens": 1
6087 });
6088 for headers in [
6089 vec![("authorization", "Bearer key-one")],
6090 vec![("authorization", "Bearer key-two")],
6091 vec![],
6092 ] {
6093 let (status, _) =
6094 post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
6095 assert_eq!(status, StatusCode::OK);
6096 }
6097
6098 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6099 let recent = stats["recent"].as_array().unwrap();
6100 assert_eq!(recent.len(), 3);
6101 let one = recent[0]["via_api_key"].as_str().unwrap();
6102 let two = recent[1]["via_api_key"].as_str().unwrap();
6103 assert_ne!(one, two, "two keys must not collapse into one caller");
6104 assert!(
6105 recent[2]["via_api_key"].is_null(),
6106 "an unauthenticated call is null, not a fingerprint of nothing"
6107 );
6108 assert!(recent[2]["client"].is_null());
6109 }
6110
6111 #[tokio::test]
6116 async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
6117 let state = Arc::new(test_state(
6118 named_test_model("really-loaded", 256),
6119 ResponseCache::new(4, Duration::from_secs(60)),
6120 ));
6121 let app = test_app_with_state(Arc::clone(&state));
6122
6123 let (status, _) = post_json_uri(
6124 &app,
6125 "/v1/chat/completions",
6126 serde_json::json!({
6127 "model": "gpt-4-turbo-that-is-not-here",
6128 "messages": [{"role": "user", "content": "hi"}],
6129 "max_tokens": 2
6130 }),
6131 )
6132 .await;
6133 assert_eq!(status, StatusCode::OK);
6134
6135 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6136 assert_eq!(stats["recent"][0]["model"], "really-loaded");
6137
6138 state.swap_active(None);
6141 let (status, _) = post_json_uri(
6142 &app,
6143 "/v1/chat/completions",
6144 serde_json::json!({
6145 "model": "gpt-4-turbo-that-is-not-here",
6146 "messages": [{"role": "user", "content": "hi"}]
6147 }),
6148 )
6149 .await;
6150 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6151 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6152 let recent = stats["recent"].as_array().unwrap();
6153 assert!(recent[recent.len() - 1]["model"].is_null());
6154 }
6155
6156 #[tokio::test]
6160 async fn a_streamed_row_names_the_model_it_decoded_against() {
6161 let state = Arc::new(test_state(
6162 named_test_model("model-before", 256),
6163 ResponseCache::new(4, Duration::from_secs(60)),
6164 ));
6165 let app = test_app_with_state(Arc::clone(&state));
6166 let _ = post_sse_raw(&app, resumable_request()).await;
6167 active_model(&state, "model-after");
6169
6170 let (_, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6171 assert_eq!(stats["recent"][0]["model"], "model-before");
6172 }
6173
6174 #[tokio::test]
6177 async fn the_queue_gauge_is_null_when_nothing_can_queue() {
6178 let app = test_app();
6179 let (status, stats) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6180 assert_eq!(status, StatusCode::OK);
6181 assert!(
6182 stats["queue_depth"].is_null(),
6183 "without continuous batching nothing queues, so there is nothing to measure"
6184 );
6185 assert!(stats["queue_rejected_total"].is_null());
6186 assert_eq!(
6187 stats["generating_now"], 0,
6188 "work in progress is measured and really is zero here"
6189 );
6190 }
6191
6192 async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
6197 post_sse_raw_uri(app, frink_api::routes::V1_CHAT_COMPLETIONS, body).await
6198 }
6199
6200 async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
6204 use http_body_util::BodyExt;
6205 use tower::ServiceExt;
6206
6207 let response = app
6208 .clone()
6209 .oneshot(
6210 axum::http::Request::builder()
6211 .method("POST")
6212 .uri(uri)
6213 .header("content-type", "application/json")
6214 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6215 .unwrap(),
6216 )
6217 .await
6218 .unwrap();
6219 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6220 String::from_utf8(bytes.to_vec()).unwrap()
6221 }
6222
6223 async fn get_json_with_headers(
6224 app: &Router,
6225 uri: &str,
6226 headers: &[(&str, &str)],
6227 ) -> (StatusCode, serde_json::Value) {
6228 use http_body_util::BodyExt;
6229 use tower::ServiceExt;
6230
6231 let mut builder = axum::http::Request::builder().method("GET").uri(uri);
6232 for (name, value) in headers {
6233 builder = builder.header(*name, *value);
6234 }
6235 let response = app
6236 .clone()
6237 .oneshot(builder.body(axum::body::Body::empty()).unwrap())
6238 .await
6239 .unwrap();
6240 let status = response.status();
6241 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6242 (
6243 status,
6244 serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
6245 )
6246 }
6247
6248 fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
6249 body.lines()
6250 .filter_map(|line| line.strip_prefix(field))
6251 .map(str::trim)
6252 .collect()
6253 }
6254
6255 fn resumable_request() -> serde_json::Value {
6256 serde_json::json!({
6257 "model": "m",
6258 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
6259 "max_tokens": 4,
6260 "temperature": 0,
6261 "stream": true,
6262 "stream_resumable": true,
6263 })
6264 }
6265
6266 #[tokio::test]
6271 async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
6272 let app = test_app();
6273 let body = post_sse_raw(&app, resumable_request()).await;
6274
6275 let request_id = body
6276 .lines()
6277 .find_map(|l| l.strip_prefix("data: "))
6278 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6279 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6280 .expect("the first chunk names the request");
6281
6282 let ids = sse_field(&body, "id:");
6283 let datas = sse_field(&body, "data:");
6284 assert_eq!(
6285 ids.len(),
6286 datas.len(),
6287 "every event carries an id, or a reconnect cannot name where it stopped"
6288 );
6289 for (i, id) in ids.iter().enumerate() {
6290 assert_eq!(*id, format!("{request_id}:{i}"));
6291 }
6292 let retries = sse_field(&body, "retry:");
6293 assert_eq!(
6294 retries.len(),
6295 1,
6296 "the reconnect delay is stated once, not on every event"
6297 );
6298 assert_eq!(retries[0], "1500");
6299 assert!(
6300 body.contains("data: [DONE]"),
6301 "the end of stream is still stated"
6302 );
6303 }
6304
6305 #[tokio::test]
6309 async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
6310 let app = test_app();
6311 let mut request = resumable_request();
6312 request["stream_resumable"] = serde_json::json!(false);
6313 let body = post_sse_raw(&app, request).await;
6314 assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
6315 assert!(
6316 sse_field(&body, "id:").is_empty(),
6317 "an id promises a replay this stream cannot serve"
6318 );
6319 assert!(sse_field(&body, "retry:").is_empty());
6320 }
6321
6322 #[tokio::test]
6326 async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
6327 let app = test_app();
6328 let body = post_sse_raw(&app, resumable_request()).await;
6329 let request_id = sse_field(&body, "id:")[0]
6330 .rsplit_once(':')
6331 .unwrap()
6332 .0
6333 .to_string();
6334 let streamed: Vec<String> = sse_field(&body, "data:")
6335 .iter()
6336 .map(|d| d.to_string())
6337 .collect();
6338
6339 let (status, polled) = get_json(
6340 &app,
6341 &format!("{}?from=0", frink_api::routes::v1_stream_poll(&request_id)),
6342 )
6343 .await;
6344 assert_eq!(status, StatusCode::OK);
6345 let events: Vec<String> = polled["events"]
6346 .as_array()
6347 .unwrap()
6348 .iter()
6349 .map(|e| e["data"].as_str().unwrap().to_string())
6350 .collect();
6351 assert_eq!(
6352 events, streamed,
6353 "the fallback must deliver the same answer, not a re-run of it"
6354 );
6355 assert_eq!(polled["request_id"], request_id);
6356 assert_eq!(
6357 polled["done"], false,
6358 "events were still being handed out, so the client must ask again"
6359 );
6360
6361 let next = polled["next_index"].as_u64().unwrap();
6364 let (_, drained) = get_json(
6365 &app,
6366 &format!(
6367 "{}?from={next}",
6368 frink_api::routes::v1_stream_poll(&request_id)
6369 ),
6370 )
6371 .await;
6372 assert_eq!(drained["done"], true);
6373 assert_eq!(drained["events"].as_array().unwrap().len(), 0);
6374 }
6375
6376 #[tokio::test]
6380 async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
6381 let app = test_app();
6382 let body = post_sse_raw(&app, resumable_request()).await;
6383 let ids = sse_field(&body, "id:");
6384 let datas: Vec<String> = sse_field(&body, "data:")
6385 .iter()
6386 .map(|d| d.to_string())
6387 .collect();
6388 assert!(
6389 ids.len() >= 3,
6390 "need a few events to resume into the middle"
6391 );
6392 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6393
6394 let (status, resumed) = get_json_with_headers(
6395 &app,
6396 &format!("{}/poll", frink_api::routes::v1_stream(&request_id)),
6397 &[],
6398 )
6399 .await;
6400 assert_eq!(status, StatusCode::OK);
6401 assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
6402
6403 let (_, tail) = get_json(
6405 &app,
6406 &format!("{}?from=2", frink_api::routes::v1_stream_poll(&request_id)),
6407 )
6408 .await;
6409 let tail_events: Vec<String> = tail["events"]
6410 .as_array()
6411 .unwrap()
6412 .iter()
6413 .map(|e| e["data"].as_str().unwrap().to_string())
6414 .collect();
6415 assert_eq!(tail_events, datas[2..].to_vec());
6416 }
6417
6418 #[tokio::test]
6421 async fn an_sse_reconnect_resumes_from_the_last_event_id() {
6422 use http_body_util::BodyExt;
6423 use tower::ServiceExt;
6424
6425 let app = test_app();
6426 let body = post_sse_raw(&app, resumable_request()).await;
6427 let ids = sse_field(&body, "id:");
6428 let datas: Vec<String> = sse_field(&body, "data:")
6429 .iter()
6430 .map(|d| d.to_string())
6431 .collect();
6432 let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
6433
6434 let response = app
6435 .clone()
6436 .oneshot(
6437 axum::http::Request::builder()
6438 .method("GET")
6439 .uri(frink_api::routes::v1_stream(&request_id))
6440 .header("last-event-id", format!("{request_id}:0"))
6441 .body(axum::body::Body::empty())
6442 .unwrap(),
6443 )
6444 .await
6445 .unwrap();
6446 assert_eq!(response.status(), StatusCode::OK);
6447 assert_eq!(
6448 response
6449 .headers()
6450 .get("x-accel-buffering")
6451 .and_then(|v| v.to_str().ok()),
6452 Some("no"),
6453 "the reconnect needs the same anti-buffering header as the stream"
6454 );
6455 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6456 let resumed = String::from_utf8(bytes.to_vec()).unwrap();
6457 assert_eq!(
6458 sse_field(&resumed, "data:")
6459 .iter()
6460 .map(|d| d.to_string())
6461 .collect::<Vec<_>>(),
6462 datas[1..].to_vec()
6463 );
6464 assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
6465 }
6466
6467 #[tokio::test]
6471 async fn a_last_event_id_from_another_stream_is_refused() {
6472 let app = test_app();
6473 let body = post_sse_raw(&app, resumable_request()).await;
6474 let request_id = sse_field(&body, "id:")[0]
6475 .rsplit_once(':')
6476 .unwrap()
6477 .0
6478 .to_string();
6479
6480 let (status, err) = get_json_with_headers(
6481 &app,
6482 &frink_api::routes::v1_stream(&request_id),
6483 &[("last-event-id", "chatcmpl-someone-else:3")],
6484 )
6485 .await;
6486 assert_eq!(status, StatusCode::BAD_REQUEST);
6487 assert_eq!(err["error"]["code"], "bad_last_event_id");
6488 }
6489
6490 #[tokio::test]
6494 async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
6495 let app = test_app();
6496 let mut request = resumable_request();
6497 request["stream_resumable"] = serde_json::json!(false);
6498 let body = post_sse_raw(&app, request).await;
6499 let request_id = body
6500 .lines()
6501 .find_map(|l| l.strip_prefix("data: "))
6502 .and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
6503 .and_then(|v| v["request_id"].as_str().map(str::to_string))
6504 .unwrap();
6505
6506 let (status, err) = get_json(&app, &frink_api::routes::v1_stream_poll(&request_id)).await;
6507 assert_eq!(status, StatusCode::NOT_FOUND);
6508 assert_eq!(err["error"]["code"], "stream_not_found");
6509 assert!(err["error"]["message"]
6510 .as_str()
6511 .unwrap()
6512 .contains("stream_resumable"));
6513 }
6514
6515 #[test]
6519 fn the_axum_stream_patterns_match_the_published_templates() {
6520 assert_eq!(
6521 axum_path(frink_api::routes::V1_STREAM),
6522 "/v1/stream/:request_id"
6523 );
6524 assert_eq!(
6525 axum_path(frink_api::routes::V1_STREAM_POLL),
6526 "/v1/stream/:request_id/poll"
6527 );
6528 assert_eq!(
6529 frink_api::routes::v1_stream("abc"),
6530 axum_path(frink_api::routes::V1_STREAM).replace(":request_id", "abc")
6531 );
6532 }
6533
6534 #[test]
6544 fn no_published_template_reaches_the_router_with_its_braces() {
6545 for template in [
6546 frink_api::routes::V1_STREAM,
6547 frink_api::routes::V1_STREAM_POLL,
6548 frink_api::routes::V1_RESPONSE,
6549 frink_api::routes::V1_RESPONSE_CANCEL,
6550 frink_api::routes::ADMIN_TASK_CANCEL,
6551 ] {
6552 assert!(
6553 template.contains('{'),
6554 "{template} is in the template list but has no placeholder"
6555 );
6556 let mounted = axum_path(template);
6557 assert!(
6558 !mounted.contains('{') && !mounted.contains('}'),
6559 "{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
6560 );
6561 assert!(
6562 mounted.contains(':'),
6563 "{template} lost its placeholder entirely and would match one path only"
6564 );
6565 }
6566 }
6567
6568 #[tokio::test]
6576 async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
6577 let app = test_app();
6578 let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
6579 assert_eq!(status, StatusCode::NOT_FOUND);
6580 assert!(
6581 !body.is_null(),
6582 "empty body means axum never matched the route, so the id was read as a literal segment"
6583 );
6584 }
6585
6586 #[tokio::test]
6589 async fn the_task_list_starts_empty_rather_than_absent() {
6590 let app = test_app();
6591 let (status, body) = get_json(&app, frink_api::routes::ADMIN_TASKS).await;
6592 assert_eq!(status, StatusCode::OK);
6593 assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
6594 }
6595
6596 #[tokio::test]
6605 async fn the_slots_route_is_registered_and_refuses_by_naming_slot_save_path() {
6606 assert!(
6607 std::env::var("FRINK_SLOT_SAVE_PATH").is_err(),
6608 "this test asserts the unconfigured behaviour"
6609 );
6610 let app = test_app();
6611 let (status, body) = post_json_uri(
6612 &app,
6613 &format!("{}?action=save", frink_api::routes::slots_id(0)),
6614 serde_json::json!({"filename": "sys.fslot", "prompt": "hi"}),
6615 )
6616 .await;
6617 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
6618 assert!(
6619 body["error"]["message"]
6620 .as_str()
6621 .unwrap()
6622 .contains("--slot-save-path"),
6623 "{body}"
6624 );
6625 }
6626
6627 pub(crate) async fn post_json_uri(
6628 app: &Router,
6629 uri: &str,
6630 body: serde_json::Value,
6631 ) -> (StatusCode, serde_json::Value) {
6632 use http_body_util::BodyExt;
6633 use tower::ServiceExt;
6634
6635 let response = app
6636 .clone()
6637 .oneshot(
6638 axum::http::Request::builder()
6639 .method("POST")
6640 .uri(uri)
6641 .header("content-type", "application/json")
6642 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6643 .unwrap(),
6644 )
6645 .await
6646 .unwrap();
6647 let status = response.status();
6648 let bytes = response.into_body().collect().await.unwrap().to_bytes();
6649 let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
6650 (status, json)
6651 }
6652
6653 async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
6654 post_json_uri(app, "/v1/chat/completions", body).await.1
6655 }
6656
6657 #[tokio::test]
6666 async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
6667 let app = test_app();
6668 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6669 assert_eq!(status, StatusCode::OK);
6670
6671 let memory = &body["memory"];
6672 if memory.is_null() {
6673 return;
6676 }
6677 assert!(
6678 memory["bytes"].as_u64().is_some_and(|b| b > 0),
6679 "a read that produced a zero is a broken read, not an idle \
6680 engine: {memory}"
6681 );
6682 assert!(
6683 ["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
6684 "the quantity must travel with the number: {memory}"
6685 );
6686 }
6687
6688 #[tokio::test]
6694 async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
6695 let app = test_app();
6696 let (status, body) = get_json(&app, frink_api::routes::V1_STATS).await;
6697 assert_eq!(status, StatusCode::OK);
6698 for pool in ["kv_pages", "window_slots", "state_slots"] {
6699 assert!(
6700 body["pools"][pool].is_null(),
6701 "{pool} must be null rather than a zero row: {}",
6702 body["pools"]
6703 );
6704 }
6705 }
6706
6707 #[tokio::test]
6715 async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
6716 use http_body_util::BodyExt;
6717 use tower::ServiceExt;
6718
6719 let app = test_app();
6720 let send = |body: serde_json::Value| {
6721 let app = app.clone();
6722 async move {
6723 app.oneshot(
6724 axum::http::Request::builder()
6725 .method("POST")
6726 .uri(frink_api::routes::V1_MESSAGES)
6727 .header("content-type", "application/json")
6728 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
6729 .unwrap(),
6730 )
6731 .await
6732 .unwrap()
6733 }
6734 };
6735
6736 let ok = send(serde_json::json!({
6737 "model": "test",
6738 "max_tokens": 1,
6739 "messages": [{"role": "user", "content": "hi"}],
6740 }))
6741 .await;
6742 assert_eq!(ok.status(), StatusCode::OK);
6743 let id = ok
6744 .headers()
6745 .get("request-id")
6746 .expect("a served message names its id")
6747 .to_str()
6748 .unwrap()
6749 .to_string();
6750 assert!(!id.is_empty());
6751
6752 let bad = send(serde_json::json!({"model": "test"})).await;
6755 assert!(bad.status().is_client_error());
6756 let other = bad.headers().get("request-id").expect("errors too");
6757 assert_ne!(other.to_str().unwrap(), id);
6758 let _ = bad.into_body().collect().await.unwrap();
6759 }
6760
6761 #[tokio::test]
6769 async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
6770 let state = Arc::new(test_state(
6771 test_model_full_byte_vocab(),
6772 ResponseCache::new(1000, Duration::from_secs(3600)),
6773 ));
6774 let app = test_app_with_state(Arc::clone(&state));
6775 let body = serde_json::json!({
6776 "model": "test",
6777 "messages": [{"role": "user", "content": "hi"}],
6778 "max_tokens": 1,
6779 });
6780
6781 state
6782 .maintenance
6783 .lock()
6784 .unwrap()
6785 .begin_rebuild()
6786 .expect("a fresh server is serving, so the rebuild starts");
6787 let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
6788 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
6789 assert_eq!(refused["error"]["type"], "cache_rebuilding");
6790
6791 state.maintenance.lock().unwrap().finish_rebuild(true);
6792 let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
6793 assert_eq!(
6794 status,
6795 StatusCode::OK,
6796 "the gate reopens; a rebuild is not a latch"
6797 );
6798 }
6799
6800 #[tokio::test]
6805 async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
6806 let app = test_app();
6807 let (status, body) = post_json_uri(
6808 &app,
6809 frink_api::routes::V1_CANCEL,
6810 serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
6811 )
6812 .await;
6813 assert_eq!(status, StatusCode::NOT_FOUND);
6814 assert_eq!(body["cancelled"], serde_json::json!(false));
6815 assert_eq!(body["request_id"], "chatcmpl-never-issued");
6816 assert!(
6817 body["detail"].as_str().is_some_and(|d| !d.is_empty()),
6818 "the verdict must carry a human reason: {body}"
6819 );
6820 }
6821
6822 #[tokio::test]
6826 async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
6827 let state = Arc::new(test_state(
6828 test_model_full_byte_vocab(),
6829 ResponseCache::new(1000, Duration::from_secs(3600)),
6830 ));
6831 let app = test_app_with_state(Arc::clone(&state));
6832 let (token, _guard) = state.cancels.register("chatcmpl-live");
6833
6834 let (status, before) = get_json(&app, frink_api::routes::ADMIN_STATS).await;
6835 assert_eq!(status, StatusCode::OK);
6836 assert_eq!(before["generating_now"], serde_json::json!(1));
6837
6838 let (status, body) = post_json_uri(
6839 &app,
6840 frink_api::routes::V1_CANCEL,
6841 serde_json::json!({ "request_id": "chatcmpl-live" }),
6842 )
6843 .await;
6844 assert_eq!(status, StatusCode::OK);
6845 assert_eq!(body["cancelled"], serde_json::json!(true));
6846 assert!(
6847 token.is_cancelled(),
6848 "the endpoint answered ok without setting the flag the decode loop reads"
6849 );
6850 }
6851
6852 #[tokio::test]
6853 async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
6854 let app = test_app();
6855 let (status, tok) =
6856 post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
6857 assert_eq!(status, StatusCode::OK);
6858 let tokens = tok["tokens"].as_array().unwrap();
6859 assert_eq!(tok["count"], tokens.len());
6860 assert!(!tokens.is_empty());
6861
6862 let (status, detok) = post_json_uri(
6863 &app,
6864 "/v1/detokenize",
6865 serde_json::json!({ "tokens": tokens }),
6866 )
6867 .await;
6868 assert_eq!(status, StatusCode::OK);
6869 assert_eq!(detok["text"], "Hi");
6870
6871 let (status, emb) = post_json_uri(
6872 &app,
6873 "/v1/embeddings",
6874 serde_json::json!({
6875 "input": "Hi",
6876 "embedding_type": "mean"
6877 }),
6878 )
6879 .await;
6880 assert_eq!(status, StatusCode::OK);
6881 let vec = emb["data"][0]["embedding"].as_array().unwrap();
6882 assert!(!vec.is_empty());
6883 assert!(vec.iter().all(|v| v.as_f64().is_some()));
6884 }
6885
6886 #[tokio::test]
6892 async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
6893 let app = test_app();
6894 let (status, body) = post_json_uri(
6895 &app,
6896 "/v1/embeddings",
6897 serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
6898 )
6899 .await;
6900 assert_eq!(status, StatusCode::BAD_REQUEST);
6901 let msg = body["error"]["message"].as_str().unwrap();
6902 assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
6903 }
6904
6905 #[tokio::test]
6910 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6911 async fn a_real_embedding_model_serves_v1_embeddings() {
6912 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6913 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6914 if !path.exists() {
6915 eprintln!("SKIP: {} not present", path.display());
6916 return;
6917 }
6918 let encoder = frink_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
6919 let mut state = test_state(
6920 test_model_full_byte_vocab(),
6921 ResponseCache::new(1000, Duration::from_secs(3600)),
6922 );
6923 state.embedding = Some(Arc::new(encoder));
6924 let app = test_app_with_state(Arc::new(state));
6925
6926 let (status, body) = post_json_uri(
6927 &app,
6928 "/v1/embeddings",
6929 serde_json::json!({ "input": ["Hello world", "a second input"] }),
6930 )
6931 .await;
6932 assert_eq!(status, StatusCode::OK, "{body}");
6933 assert_eq!(body["model"], "bge-small-en-v1.5");
6934 let data = body["data"].as_array().unwrap();
6935 assert_eq!(data.len(), 2);
6936 for (i, row) in data.iter().enumerate() {
6937 assert_eq!(row["index"], i);
6938 let v: Vec<f64> = row["embedding"]
6939 .as_array()
6940 .unwrap()
6941 .iter()
6942 .map(|x| x.as_f64().unwrap())
6943 .collect();
6944 assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
6945 let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
6946 assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
6947 }
6948 assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
6951
6952 let (status, mean) = post_json_uri(
6956 &app,
6957 "/v1/embeddings",
6958 serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
6959 )
6960 .await;
6961 assert_eq!(status, StatusCode::OK);
6962 assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
6963 }
6964
6965 #[tokio::test]
6977 #[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
6978 async fn an_encoder_can_be_the_loaded_model() {
6979 let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
6980 .join("../../models/bge-small-en-v1.5-q8_0.gguf");
6981 if !path.exists() {
6982 eprintln!("SKIP: {} not present", path.display());
6983 return;
6984 }
6985
6986 let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
6990 assert!(
6991 matches!(loaded, model::LoadedModel::Encoder(_)),
6992 "an encoder-only GGUF reached a decoder loader"
6993 );
6994 let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
6995 assert!(
6996 matches!(loaded, Loaded::Encoder(_)),
6997 "the encoder did not stay an encoder through activation"
6998 );
6999 assert!(
7000 batcher.is_none() && ceiling.is_none(),
7001 "an encoder was given a decode batcher or a KV ceiling it has no use for"
7002 );
7003
7004 let state = test_state(
7005 test_model_full_byte_vocab(),
7006 ResponseCache::new(1000, Duration::from_secs(3600)),
7007 );
7008 state.swap_active(Some(Arc::new(ActiveModel {
7009 id: None,
7010 loaded,
7011 batcher,
7012 ceiling,
7013 checkpoint_path: None,
7014 })));
7015 let app = test_app_with_state(Arc::new(state));
7016
7017 let (status, body) = post_json_uri(
7019 &app,
7020 "/v1/embeddings",
7021 serde_json::json!({ "input": "Hello world" }),
7022 )
7023 .await;
7024 assert_eq!(status, StatusCode::OK, "{body}");
7025 assert_eq!(body["model"], "bge-small-en-v1.5");
7026 let v = body["data"][0]["embedding"].as_array().unwrap();
7027 assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
7028
7029 let (status, body) = post_json_uri(
7031 &app,
7032 "/v1/chat/completions",
7033 serde_json::json!({
7034 "model": "bge-small-en-v1.5",
7035 "messages": [{"role": "user", "content": "hi"}],
7036 }),
7037 )
7038 .await;
7039 assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
7040 let msg = body["error"]["message"].as_str().unwrap();
7041 for fact in [
7042 "bge-small-en-v1.5",
7043 "bert",
7044 "embedding model",
7045 "/v1/embeddings",
7046 ] {
7047 assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
7048 }
7049
7050 let (status, models) = get_json(&app, frink_api::routes::V1_MODELS).await;
7052 assert_eq!(status, StatusCode::OK);
7053 let entry = &models["data"][0];
7054 assert_eq!(entry["id"], "bge-small-en-v1.5");
7055 assert_eq!(entry["frink_model_kind"], "embedding");
7056 assert_eq!(entry["frink_tokenizer"], "gguf-wordpiece");
7057 assert_eq!(entry["frink_n_embd"], 384);
7058 assert_eq!(entry["frink_pooling"], "CLS");
7059 assert_eq!(
7060 entry["frink_endpoints"],
7061 serde_json::json!(["/v1/embeddings"])
7062 );
7063 assert!(entry.get("supported_reasoning_efforts").is_none());
7066
7067 let (status, health) = get_json(&app, frink_api::routes::HEALTH).await;
7069 assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
7070 assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
7071 assert_eq!(health["model"]["synthetic_weights"], false);
7072 let weights = health["capabilities"]
7073 .as_array()
7074 .unwrap()
7075 .iter()
7076 .find(|c| c["id"] == frink_api::health::capability::REAL_WEIGHTS)
7077 .expect("a real-weights capability row");
7078 let detail = weights["detail"].as_str().unwrap_or_default();
7079 assert!(detail.contains("ENCODER"), "{detail}");
7080 let (status, body) = post_json_uri(
7088 &app,
7089 frink_api::routes::V1_TOKENIZE,
7090 serde_json::json!({ "content": "hello world" }),
7091 )
7092 .await;
7093 assert_eq!(
7094 status,
7095 StatusCode::OK,
7096 "an encoder has a real tokenizer: {body}"
7097 );
7098 let tokens = body["tokens"].as_array().expect("tokens array").clone();
7099 assert!(!tokens.is_empty(), "WordPiece produced nothing: {body}");
7100
7101 let (status, body) = post_json_uri(
7102 &app,
7103 frink_api::routes::V1_DETOKENIZE,
7104 serde_json::json!({ "tokens": tokens }),
7105 )
7106 .await;
7107 assert_eq!(status, StatusCode::OK, "{body}");
7108 let round_tripped = body["content"].as_str().expect("content").to_string();
7109 assert!(
7110 round_tripped.contains("hello") && round_tripped.contains("world"),
7111 "the ids did not decode back through the encoder's own vocabulary: {round_tripped}"
7112 );
7113
7114 let (status, _) = post_json_uri(
7117 &app,
7118 "/v1/completions",
7119 serde_json::json!({ "model": "m", "prompt": "hi", "max_tokens": 1 }),
7120 )
7121 .await;
7122 assert_eq!(
7123 status,
7124 StatusCode::NOT_IMPLEMENTED,
7125 "tokenizing an encoder must not have opened a path to generating with one"
7126 );
7127 }
7128
7129 #[tokio::test]
7134 async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
7135 use http_body_util::BodyExt;
7136 use tower::ServiceExt;
7137
7138 let fixture = concat!(
7139 "../frink-models/tests/fixtures/",
7140 "frink_real_moe_test.gguf"
7141 );
7142 let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
7143 let decoder = Decoder::from_gguf_with_expert_cache(
7144 &fixture,
7145 frink_models::config::test_moe_fixture(),
7146 Some(1024 * 1024),
7147 )
7148 .expect("MoE fixture must load store-backed");
7149
7150 let mut caches: Vec<frink_core::cache::KvCache> = decoder.config.new_kv_caches();
7154 decoder.forward_token(1, 0, &mut caches);
7155
7156 let model = Model::Gguf(GgufModel {
7157 decoder: Arc::new(decoder),
7158 tokenizer: Arc::new(ServerTokenizer::Byte),
7159 stop_tokens: StopTokens::default(),
7160 bos_id: None,
7161 is_synthetic: false,
7162 chat_template: chat_template::PromptTemplate::plain(),
7163 });
7164 let state = Arc::new(test_state(
7165 model,
7166 ResponseCache::new(16, Duration::from_secs(60)),
7167 ));
7168 let app = Router::new()
7169 .route("/metrics", axum::routing::get(metrics))
7170 .route("/v1/chat/completions", post(chat_completions))
7171 .with_state(state);
7172
7173 let fetch_metrics = |app: Router| async move {
7174 let resp = app
7175 .oneshot(
7176 axum::http::Request::builder()
7177 .method("GET")
7178 .uri("/metrics")
7179 .body(axum::body::Body::empty())
7180 .unwrap(),
7181 )
7182 .await
7183 .unwrap();
7184 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
7185 String::from_utf8(bytes.to_vec()).unwrap()
7186 };
7187
7188 let after = fetch_metrics(app.clone()).await;
7189 assert!(
7190 after.contains("frink_expert_cache_misses_total"),
7191 "streaming model must expose expert-cache metrics: {after}"
7192 );
7193 let misses: u64 = after
7194 .lines()
7195 .find(|l| l.starts_with("frink_expert_cache_misses_total"))
7196 .and_then(|l| l.split_whitespace().nth(1))
7197 .and_then(|v| v.parse().ok())
7198 .expect("misses metric line must parse");
7199 assert!(
7200 misses > 0,
7201 "decode must have read experts through the store: {after}"
7202 );
7203 }
7204
7205 fn weather_tool() -> serde_json::Value {
7206 serde_json::json!({
7207 "type": "function",
7208 "function": {
7209 "name": "get_weather",
7210 "description": "Get the current weather for a location.",
7211 "parameters": {
7212 "type": "object",
7213 "properties": {"location": {"type": "string"}},
7214 "required": ["location"]
7215 }
7216 }
7217 })
7218 }
7219
7220 fn weather_tool_def() -> ToolDef {
7221 ToolDef {
7222 kind: "function".to_string(),
7223 function: ToolFunctionDef {
7224 name: "get_weather".to_string(),
7225 description: Some("Get the current weather for a location.".to_string()),
7226 parameters: Some(serde_json::json!({
7227 "type": "object",
7228 "properties": {"location": {"type": "string"}},
7229 "required": ["location"]
7230 })),
7231 },
7232 }
7233 }
7234
7235 #[test]
7236 fn tool_preamble_mentions_every_tool_name_and_description() {
7237 let preamble = tool_preamble(&[weather_tool_def()]);
7238 assert!(preamble.contains("get_weather"));
7239 assert!(preamble.contains("Get the current weather for a location."));
7240 assert!(preamble.contains("<tool_call>"));
7241 assert!(preamble.contains("</tool_call>"));
7242 }
7243
7244 #[test]
7245 fn a_real_marker_becomes_a_structured_tool_call() {
7246 let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
7247 let (message, finish) = build_response_message(
7248 text.to_string(),
7249 &[weather_tool_def()],
7250 output::OutputPosture::for_model("test-model"),
7251 "stop",
7252 );
7253 assert_eq!(finish, "tool_calls");
7254 let calls = message.tool_calls.expect("must carry a tool call");
7255 assert_eq!(calls[0].function.name, "get_weather");
7256 let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
7257 assert_eq!(parsed["location"], "Paris");
7258 }
7259
7260 #[test]
7261 fn a_plain_answer_is_not_promoted_to_a_tool_call() {
7262 let (message, finish) = build_response_message(
7263 "just an answer".to_string(),
7264 &[weather_tool_def()],
7265 output::OutputPosture::for_model("test-model"),
7266 "stop",
7267 );
7268 assert_eq!(finish, "stop");
7269 assert!(message.tool_calls.is_none());
7270 assert_eq!(message.content.as_deref(), Some("just an answer"));
7271 }
7272
7273 #[test]
7276 fn a_malformed_payload_is_not_a_tool_call() {
7277 let (message, finish) = build_response_message(
7278 "<tool_call>not valid json at all</tool_call>".to_string(),
7279 &[weather_tool_def()],
7280 output::OutputPosture::for_model("test-model"),
7281 "stop",
7282 );
7283 assert_eq!(finish, "stop");
7284 assert!(message.tool_calls.is_none());
7285 }
7286
7287 #[test]
7290 fn a_tool_that_was_never_offered_is_not_returned() {
7291 let (message, finish) = build_response_message(
7292 "<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
7293 &[weather_tool_def()],
7294 output::OutputPosture::for_model("test-model"),
7295 "stop",
7296 );
7297 assert_eq!(finish, "stop");
7298 assert!(message.tool_calls.is_none());
7299 }
7300
7301 #[test]
7303 fn marker_text_with_no_tools_offered_stays_content() {
7304 let (message, finish) = build_response_message(
7305 "<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
7306 &[],
7307 output::OutputPosture::for_model("test-model"),
7308 "stop",
7309 );
7310 assert_eq!(finish, "stop");
7311 assert!(message.tool_calls.is_none());
7312 assert!(message.content.is_some());
7313 }
7314
7315 #[test]
7319 fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
7320 let opened = std::cell::Cell::new(0usize);
7321 let mut parser = crate::policy::parser::ToolCallParser::new(
7322 crate::policy::parser::ToolCallFormat::Qwen3Coder,
7323 vec![
7324 crate::policy::parser::tool_call::ToolSchema::with_parameters(
7325 "write_file",
7326 serde_json::json!({"type": "object", "properties": {
7327 "path": {"type": "string"},
7328 "contents": {"type": "string"}
7329 }}),
7330 ),
7331 ],
7332 );
7333 let wire = "<tool_call><function=write_file>\
7334 <parameter=path>\n/tmp/x\n</parameter>\
7335 <parameter=contents>\nhello world\n</parameter>\
7336 </function></tool_call>";
7337
7338 let mut deltas = Vec::new();
7339 let mut text = String::new();
7340 for piece in wire.as_bytes().chunks(7) {
7341 let chunk = String::from_utf8_lossy(piece).into_owned();
7342 let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
7343 text.push_str(&more_text);
7344 deltas.extend(more);
7345 }
7346 let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
7347 text.push_str(&more_text);
7348 deltas.extend(more);
7349
7350 assert_eq!(opened.get(), 1, "one call opened");
7351 assert!(text.is_empty(), "the markers are not content: {text:?}");
7352
7353 let first = &deltas[0];
7354 assert_eq!(first.index, 0);
7355 assert_eq!(first.id.as_deref(), Some("call_0"));
7356 assert_eq!(first.kind, Some("function"));
7357 assert_eq!(first.function.name.as_deref(), Some("write_file"));
7358
7359 let joined: String = deltas
7362 .iter()
7363 .filter_map(|d| d.function.arguments.clone())
7364 .collect();
7365 let parsed: serde_json::Value =
7366 serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
7367 assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
7368 assert_eq!(parsed["contents"], serde_json::json!("hello world"));
7369 assert!(
7370 deltas.len() >= 3,
7371 "the arguments arrived in pieces, not whole: {}",
7372 deltas.len()
7373 );
7374 assert!(
7375 deltas[1..].iter().all(|d| d.function.name.is_none()),
7376 "only the opening delta carries identity"
7377 );
7378 }
7379
7380 #[test]
7382 fn text_around_a_streamed_call_is_still_content() {
7383 let opened = std::cell::Cell::new(0usize);
7384 let mut parser = crate::policy::parser::ToolCallParser::new(
7385 crate::policy::parser::ToolCallFormat::Qwen25,
7386 vec![crate::policy::parser::tool_call::ToolSchema::new(
7387 "get_weather",
7388 )],
7389 );
7390 let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
7391 \"arguments\": {}}</tool_call> done";
7392 let mut text = String::new();
7393 for piece in wire.as_bytes().chunks(5) {
7394 let chunk = String::from_utf8_lossy(piece).into_owned();
7395 let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
7396 text.push_str(&more);
7397 }
7398 let (more, _) = tool_call_deltas(parser.finish(), &opened);
7399 text.push_str(&more);
7400
7401 assert_eq!(opened.get(), 1);
7402 assert!(text.starts_with("let me check. "), "{text:?}");
7403 assert!(text.ends_with(" done"), "{text:?}");
7404 assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
7405 }
7406
7407 #[test]
7410 fn a_reasoning_block_is_split_out_of_the_answer() {
7411 let (message, finish) = build_response_message(
7412 "<think>weighing it up</think>The answer is 4.".to_string(),
7413 &[],
7414 output::OutputPosture::for_model("Qwen3-8B"),
7415 "stop",
7416 );
7417 assert_eq!(finish, "stop");
7418 assert_eq!(message.content.as_deref(), Some("The answer is 4."));
7419 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
7420 }
7421
7422 #[test]
7425 fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
7426 let (message, _) = build_response_message(
7427 "Use the <think> tag like this.".to_string(),
7428 &[],
7429 output::OutputPosture::for_model("llama-3.1-8b"),
7430 "stop",
7431 );
7432 assert_eq!(
7433 message.content.as_deref(),
7434 Some("Use the <think> tag like this.")
7435 );
7436 assert!(message.reasoning_content.is_none());
7437 }
7438
7439 #[tokio::test]
7447 async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
7448 let app = test_app();
7449 let body = serde_json::json!({
7450 "model": "m",
7451 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7452 "max_tokens": 4,
7453 "temperature": 0,
7454 });
7455 let resp = post_json(&app, body).await;
7456 let message = &resp["choices"][0]["message"];
7457 assert!(message["content"].is_string());
7458 assert!(message.get("tool_calls").is_none());
7459 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7460 assert_eq!(resp["usage"]["completion_tokens"], 4);
7461 assert_eq!(
7462 resp["usage"]["total_tokens"],
7463 resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
7464 );
7465 }
7466
7467 pub(crate) async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
7468 use http_body_util::BodyExt;
7469 use tower::ServiceExt;
7470
7471 let response = app
7472 .clone()
7473 .oneshot(
7474 axum::http::Request::builder()
7475 .method("GET")
7476 .uri(uri)
7477 .body(axum::body::Body::empty())
7478 .unwrap(),
7479 )
7480 .await
7481 .unwrap();
7482 let status = response.status();
7483 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7484 (status, serde_json::from_slice(&bytes).unwrap())
7485 }
7486
7487 #[tokio::test]
7488 async fn health_answers_a_capability_handshake_not_a_boolean() {
7489 let app = test_app();
7490 let (status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7491 assert_eq!(status, StatusCode::OK);
7492
7493 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7494 assert_eq!(health.state, frink_api::HealthState::Ready);
7495 assert!(health.pid > 0);
7496 assert!(health.server_time_unix_ms > 0);
7497 assert_eq!(health.last_request_age_seconds, None);
7500
7501 for id in [
7504 frink_api::health::capability::CPU,
7505 frink_api::health::capability::METAL,
7506 frink_api::health::capability::CUDA,
7507 frink_api::health::capability::REAL_WEIGHTS,
7508 frink_api::health::capability::CONTINUOUS_BATCHING,
7509 ] {
7510 let cap = health
7511 .capability(id)
7512 .unwrap_or_else(|| panic!("{id} missing"));
7513 assert!(!cap.reason.is_empty(), "{cap:?}");
7514 assert!(!cap.detail.is_empty(), "{cap:?}");
7515 }
7516 let weights = health
7520 .capability(frink_api::health::capability::REAL_WEIGHTS)
7521 .unwrap();
7522 assert!(!weights.available);
7523 assert_eq!(weights.reason, frink_api::health::reason::MODEL_NOT_LOADED);
7524 assert!(health.model.as_ref().unwrap().synthetic_weights);
7525 }
7526
7527 #[tokio::test]
7528 async fn health_vouches_for_liveness_after_a_request_has_been_served() {
7529 let app = test_app();
7530 let _ = post_json(
7531 &app,
7532 serde_json::json!({
7533 "model": "m",
7534 "messages": [{"role": "user", "content": "\u{1}"}],
7535 "max_tokens": 1,
7536 "temperature": 0,
7537 }),
7538 )
7539 .await;
7540 let (_status, body) = get_json(&app, frink_api::routes::HEALTH).await;
7541 let health: frink_api::HealthResponse = serde_json::from_value(body).unwrap();
7542 let age = health
7543 .last_request_age_seconds
7544 .expect("a served request is evidence of liveness");
7545 assert!((0.0..5.0).contains(&age), "implausible age {age}");
7546 }
7547
7548 async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
7550 use http_body_util::BodyExt;
7551 use tower::ServiceExt;
7552
7553 let response = app
7554 .clone()
7555 .oneshot(
7556 axum::http::Request::builder()
7557 .method("POST")
7558 .uri("/v1/chat/completions")
7559 .header("content-type", "application/json")
7560 .body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
7561 .unwrap(),
7562 )
7563 .await
7564 .unwrap();
7565 let bytes = response.into_body().collect().await.unwrap().to_bytes();
7566 String::from_utf8(bytes.to_vec())
7567 .unwrap()
7568 .lines()
7569 .filter_map(|line| line.strip_prefix("data: "))
7570 .filter(|payload| *payload != "[DONE]")
7571 .map(|payload| serde_json::from_str(payload).unwrap())
7572 .collect()
7573 }
7574
7575 #[tokio::test]
7576 async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
7577 let app = test_app();
7578 let chunks = post_sse_chunks(
7579 &app,
7580 serde_json::json!({
7581 "model": "m",
7582 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7583 "max_tokens": 4,
7584 "temperature": 0,
7585 "stream": true,
7586 }),
7587 )
7588 .await;
7589
7590 assert!(!chunks.is_empty());
7591 let request_id = chunks[0]["request_id"]
7592 .as_str()
7593 .expect("the first chunk names the request")
7594 .to_string();
7595 assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
7596 for (i, chunk) in chunks.iter().enumerate().skip(1) {
7599 assert!(
7600 chunk.get("request_id").is_none(),
7601 "chunk {i} repeats request_id"
7602 );
7603 }
7604 for chunk in &chunks {
7607 assert_eq!(chunk["id"], serde_json::json!(request_id));
7608 }
7609
7610 let other = post_sse_chunks(
7611 &app,
7612 serde_json::json!({
7613 "model": "m",
7614 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7615 "max_tokens": 4,
7616 "temperature": 0,
7617 "stream": true,
7618 }),
7619 )
7620 .await;
7621 assert_ne!(
7622 other[0]["request_id"].as_str().unwrap(),
7623 request_id,
7624 "two concurrent chats must not share an id"
7625 );
7626 }
7627
7628 #[tokio::test]
7629 async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
7630 let app = test_app();
7631 let resp = post_json(
7632 &app,
7633 serde_json::json!({
7634 "model": "m",
7635 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7636 "max_tokens": 2,
7637 "temperature": 0,
7638 }),
7639 )
7640 .await;
7641 assert_eq!(resp["id"], resp["request_id"]);
7642 assert!(resp["request_id"]
7643 .as_str()
7644 .unwrap()
7645 .starts_with("chatcmpl-"));
7646 }
7647
7648 #[tokio::test]
7651 async fn usage_carries_separate_prefill_and_decode_timings() {
7652 let app = test_app();
7653 let resp = post_json(
7654 &app,
7655 serde_json::json!({
7656 "model": "m",
7657 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7658 "max_tokens": 4,
7659 "temperature": 0,
7660 }),
7661 )
7662 .await;
7663 let usage = &resp["usage"];
7664 assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
7665 assert!(usage["generation_duration_ms"].is_number(), "{usage}");
7666 assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
7667 assert!(usage["predicted_per_second"].is_number(), "{usage}");
7668 assert!(usage.get("cached_tokens").is_none(), "{usage}");
7670 }
7671
7672 #[tokio::test]
7680 async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
7681 let app = test_app();
7682 let body = serde_json::json!({
7683 "model": "m",
7684 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7685 "max_tokens": 4,
7686 "temperature": 0,
7687 "tools": [weather_tool()],
7688 });
7689 let resp = post_json(&app, body).await;
7690 let message = &resp["choices"][0]["message"];
7691 assert!(
7692 message["content"].is_string(),
7693 "must fall back to plain content when no real tool-call marker is present: {resp:?}"
7694 );
7695 assert!(message.get("tool_calls").is_none());
7696 assert_eq!(resp["choices"][0]["finish_reason"], "length");
7700 }
7701
7702 #[tokio::test]
7706 async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
7707 let app = test_app();
7708 let body = serde_json::json!({
7709 "model": "m",
7710 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7711 "max_tokens": 3,
7712 "temperature": 0,
7713 });
7714 let first = post_json(&app, body.clone()).await;
7715 assert_eq!(first["frink_cache"], "miss");
7716 let second = post_json(&app, body).await;
7717 assert_eq!(second["frink_cache"], "hit");
7718 assert_eq!(
7719 first["choices"][0]["message"]["content"],
7720 second["choices"][0]["message"]["content"]
7721 );
7722 assert_eq!(
7723 first["choices"][0]["finish_reason"],
7724 second["choices"][0]["finish_reason"]
7725 );
7726 assert_eq!(first["usage"], second["usage"]);
7727 assert_eq!(second["usage"]["completion_tokens"], 3);
7728 }
7729
7730 #[tokio::test]
7741 async fn a_grammar_request_is_not_answered_from_an_unconstrained_cache_entry() {
7742 let app = test_app();
7743 let plain = serde_json::json!({
7744 "model": "m",
7745 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7746 "max_tokens": 3,
7747 "temperature": 0,
7748 });
7749
7750 let first = post_json(&app, plain.clone()).await;
7751 assert_eq!(first["frink_cache"], "miss");
7752 let unconstrained = first["choices"][0]["message"]["content"]
7753 .as_str()
7754 .expect("content")
7755 .to_string();
7756
7757 let mut constrained = plain.clone();
7758 constrained["grammar"] = serde_json::json!("root ::= \"yes\"");
7759 let second = post_json(&app, constrained).await;
7760 assert_eq!(
7761 second["frink_cache"], "miss",
7762 "a grammar is part of the key, so this body has never been answered"
7763 );
7764 let constrained_answer = second["choices"][0]["message"]["content"]
7768 .as_str()
7769 .expect("content")
7770 .to_string();
7771 assert!(
7772 constrained_answer.contains("-> \"yes\"]"),
7773 "the grammar must have been compiled AND applied, not skipped \
7774 by a cache hit: {constrained_answer}"
7775 );
7776 assert_ne!(
7777 constrained_answer, unconstrained,
7778 "the constrained request was served the unconstrained answer"
7779 );
7780
7781 let third = post_json(&app, plain).await;
7785 assert_eq!(third["frink_cache"], "hit");
7786 assert_eq!(third["choices"][0]["message"]["content"], unconstrained);
7787 }
7788
7789 #[tokio::test]
7811 async fn a_json_object_request_does_not_reuse_the_unconstrained_cache_entry() {
7812 let state = Arc::new(test_state(
7813 test_model_full_byte_vocab(),
7814 ResponseCache::new(1000, Duration::from_secs(3600)),
7815 ));
7816 let app = test_app_with_state(state.clone());
7817 let plain = serde_json::json!({
7818 "model": "m",
7819 "messages": [
7820 {"role": "system", "content": "Answer in JSON when it helps."},
7821 {"role": "user", "content": "\u{1}\u{2}"},
7822 ],
7823 "max_tokens": 3,
7824 "temperature": 0,
7825 });
7826
7827 let first = post_json(&app, plain.clone()).await;
7828 assert_eq!(first["frink_cache"], "miss");
7829 assert_eq!(state.cache_stats().entries, 1);
7830
7831 let mut as_json = plain.clone();
7832 as_json["response_format"] = serde_json::json!({"type": "json_object"});
7833 let (status, _) = post_json_uri(&app, "/v1/chat/completions", as_json).await;
7834 assert_eq!(
7835 status,
7836 StatusCode::BAD_REQUEST,
7837 "the demo banner is not a JSON object, whoever generated it"
7838 );
7839 assert_eq!(
7840 state.cache_stats().hits,
7841 0,
7842 "a json_object request must not be answered from an entry the \
7843 JSON mask never produced"
7844 );
7845 assert_eq!(
7846 state.cache_stats().entries,
7847 2,
7848 "json_object must key its own entry, not reuse the unconstrained \
7849 one it happens to render the same prompt as"
7850 );
7851 }
7852
7853 #[tokio::test]
7865 async fn an_ignore_eos_request_is_not_answered_from_a_cache_entry_that_stopped_at_eos() {
7866 let app = test_app_with_state(Arc::new(test_state(
7867 test_model_full_byte_vocab_with_eos(Some(0x77)),
7868 ResponseCache::new(1000, Duration::from_secs(3600)),
7869 )));
7870 let body = serde_json::json!({
7871 "model": "m",
7872 "messages": [{"role": "user", "content": "\u{1}\u{2}"}],
7873 "max_tokens": 6,
7874 "temperature": 0,
7875 });
7876
7877 let stopped = post_json(&app, body.clone()).await;
7878 assert_eq!(stopped["frink_cache"], "miss");
7879 assert_eq!(
7880 stopped["choices"][0]["finish_reason"], "stop",
7881 "the fixture is only meaningful if the model's EOS really fires here"
7882 );
7883 assert_eq!(stopped["usage"]["completion_tokens"], 1);
7884
7885 let mut ignoring = body.clone();
7886 ignoring["ignore_eos"] = serde_json::json!(true);
7887 let ran_on = post_json(&app, ignoring).await;
7888 assert_eq!(
7889 ran_on["frink_cache"], "miss",
7890 "ignore_eos is part of the key, so this body has never been answered"
7891 );
7892 assert_eq!(
7893 ran_on["usage"]["completion_tokens"], 6,
7894 "ignore_eos must run the full budget, not replay the EOS-terminated answer"
7895 );
7896 assert_eq!(ran_on["choices"][0]["finish_reason"], "length");
7897 assert_ne!(
7898 ran_on["choices"][0]["message"]["content"],
7899 stopped["choices"][0]["message"]["content"]
7900 );
7901 }
7902
7903 #[tokio::test]
7909 async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
7910 let session_app = test_app();
7911 let manual_app = test_app();
7912
7913 let turn1 = post_json(
7915 &session_app,
7916 serde_json::json!({
7917 "model": "m",
7918 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7919 "session_id": "s1",
7920 "max_tokens": 5,
7921 "temperature": 0,
7922 }),
7923 )
7924 .await;
7925 let reply1 = turn1["choices"][0]["message"]["content"]
7926 .as_str()
7927 .unwrap()
7928 .to_string();
7929
7930 let manual_turn1 = post_json(
7935 &manual_app,
7936 serde_json::json!({
7937 "model": "m",
7938 "messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
7939 "max_tokens": 5,
7940 "temperature": 0,
7941 }),
7942 )
7943 .await;
7944 assert_eq!(
7945 manual_turn1["choices"][0]["message"]["content"]
7946 .as_str()
7947 .unwrap(),
7948 reply1
7949 );
7950
7951 let turn2 = post_json(
7953 &session_app,
7954 serde_json::json!({
7955 "model": "m",
7956 "messages": [{"role": "user", "content": "\u{4}\u{5}"}],
7957 "session_id": "s1",
7958 "max_tokens": 5,
7959 "temperature": 0,
7960 }),
7961 )
7962 .await;
7963 let reply2 = turn2["choices"][0]["message"]["content"]
7964 .as_str()
7965 .unwrap()
7966 .to_string();
7967
7968 let manual_turn2 = post_json(
7972 &manual_app,
7973 serde_json::json!({
7974 "model": "m",
7975 "messages": [
7976 {"role": "user", "content": "\u{1}\u{2}\u{3}"},
7977 {"role": "assistant", "content": reply1},
7978 {"role": "user", "content": "\u{4}\u{5}"},
7979 ],
7980 "max_tokens": 5,
7981 "temperature": 0,
7982 }),
7983 )
7984 .await;
7985 assert_eq!(
7986 manual_turn2["choices"][0]["message"]["content"]
7987 .as_str()
7988 .unwrap(),
7989 reply2,
7990 "resuming a session must produce identical output to manually resending the full history"
7991 );
7992 }
7993
7994 #[test]
7997 fn lock_cache_recovers_from_a_poisoned_mutex() {
7998 let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
7999
8000 let poison_cache = Arc::clone(&cache);
8001 let _ = std::thread::spawn(move || {
8002 let _guard = poison_cache.lock().unwrap();
8003 panic!("simulated panic while holding the lock");
8004 })
8005 .join();
8006
8007 let recovered = lock_cache(&cache);
8009 assert_eq!(recovered.stats().entries, 0);
8010 }
8011
8012 #[test]
8013 fn is_cacheable_true_for_greedy_or_seeded_requests() {
8014 let mut req_body = serde_json::json!({
8015 "model": "m",
8016 "messages": [{"role": "user", "content": "hi"}],
8017 });
8018 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8019 assert!(
8020 req.is_cacheable(),
8021 "default (temperature 0) must be cacheable"
8022 );
8023
8024 req_body["temperature"] = serde_json::json!(0.8);
8025 let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
8026 assert!(
8027 !req.is_cacheable(),
8028 "unseeded sampling must never be cacheable"
8029 );
8030
8031 req_body["seed"] = serde_json::json!(42);
8032 let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
8033 assert!(
8034 req.is_cacheable(),
8035 "sampling with an explicit seed is deterministic and must be cacheable"
8036 );
8037 }
8038
8039 const GRADED: &str = "{% if reasoning_effort %}\
8043 {% if reasoning_effort not in ['low','medium','high'] %}\
8044 {{ raise_exception('unsupported effort') }}\
8045 {% endif %}E:{{ reasoning_effort }}|{% endif %}\
8046 {% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
8047
8048 fn graded_template() -> chat_template::PromptTemplate {
8049 chat_template::PromptTemplate::from_gguf_metadata(
8050 Some(GRADED),
8051 Some("qwen3"),
8052 false,
8053 true,
8054 None,
8055 None,
8056 )
8057 }
8058
8059 fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
8060 serde_json::from_value(value).expect("request")
8061 }
8062
8063 #[test]
8070 fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
8071 let req = chat_request(serde_json::json!({
8072 "model": "m",
8073 "messages": [{"role": "user", "content": "hi"}],
8074 "grammar": "root ::= \"a\"+",
8075 }));
8076 req.validate_supported_fields()
8077 .expect("a valid grammar is a valid request");
8078 let params = req
8079 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8080 .expect("a valid grammar compiles at params time too");
8081 assert!(
8082 params.grammar.is_some(),
8083 "the grammar was dropped between the wire and the sampler"
8084 );
8085 assert!(
8086 params.needs_vocab_logits(),
8087 "a grammar request that may fold lm_head into a GPU argmax is \
8088 a grammar request served unconstrained"
8089 );
8090
8091 let plain = chat_request(serde_json::json!({
8092 "model": "m",
8093 "messages": [{"role": "user", "content": "hi"}],
8094 }));
8095 assert!(plain
8096 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8097 .unwrap()
8098 .grammar
8099 .is_none());
8100 }
8101
8102 fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
8103 chat_request(serde_json::json!({
8104 "model": "m",
8105 "messages": [{"role": "user", "content": "weather in Rome?"}],
8106 "tools": [weather_tool()],
8107 "tool_choice": tool_choice,
8108 }))
8109 }
8110
8111 #[test]
8115 fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
8116 for choice in [
8117 serde_json::json!("required"),
8118 serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
8119 ] {
8120 let req = tool_request(choice.clone());
8121 req.validate_supported_fields()
8122 .unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
8123 let params = req
8124 .generation_params_for_template(
8125 &graded_template(),
8126 "Qwen3-8B",
8127 crate::sampling_knobs::SamplerModel::absent(),
8128 )
8129 .unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
8130 let grammar = params
8131 .grammar
8132 .as_ref()
8133 .unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
8134 assert!(
8135 grammar.is_awaiting_trigger(),
8136 "the model must be free to think before it calls"
8137 );
8138 assert!(
8139 !grammar.allows_eog(),
8140 "{choice} must not be able to end the turn without a call"
8141 );
8142 assert!(
8148 params.needs_vocab_logits(),
8149 "{choice} would let a backend return a token id instead of logits"
8150 );
8151 assert!(
8152 !generate::greedy_gpu_fold_allowed(¶ms),
8153 "{choice} at temperature 0 must still refuse the greedy GPU fold"
8154 );
8155 }
8156 }
8157
8158 #[test]
8160 fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
8161 for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
8162 let req = tool_request(choice.clone());
8163 req.validate_supported_fields().expect("still supported");
8164 let params = match req.generation_params_for_template(
8165 &graded_template(),
8166 "Qwen3-8B",
8167 crate::sampling_knobs::SamplerModel::absent(),
8168 ) {
8169 Ok(p) => p,
8170 Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
8171 };
8172 assert!(
8173 params.grammar.is_none(),
8174 "{choice} does not force a call and must not be constrained"
8175 );
8176 }
8177 }
8178
8179 #[test]
8182 fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
8183 let req = chat_request(serde_json::json!({
8185 "model": "m",
8186 "messages": [{"role": "user", "content": "hi"}],
8187 "tool_choice": "required",
8188 }));
8189 let (status, _) = req
8190 .validate_supported_fields()
8191 .expect_err("nothing to call");
8192 assert_eq!(status, StatusCode::BAD_REQUEST);
8193
8194 let req =
8196 tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
8197 let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
8198 assert_eq!(status, StatusCode::BAD_REQUEST);
8199 assert_eq!(body["error"]["param"], "tool_choice");
8200
8201 let req = tool_request(serde_json::json!({"type": "function"}));
8203 let (status, _) = req.validate_supported_fields().expect_err("names nothing");
8204 assert_eq!(status, StatusCode::BAD_REQUEST);
8205
8206 let req = chat_request(serde_json::json!({
8208 "model": "m",
8209 "messages": [{"role": "user", "content": "hi"}],
8210 "tools": [weather_tool()],
8211 "tool_choice": "required",
8212 "grammar": "root ::= \"a\"+",
8213 }));
8214 let (status, _) = req
8215 .validate_supported_fields()
8216 .expect_err("a grammar and a forced call are two constraints");
8217 assert_eq!(status, StatusCode::BAD_REQUEST);
8218
8219 let req = tool_request(serde_json::json!("required"));
8225 let (status, Json(body)) = match req.generation_params_for_template(
8226 &graded_template(),
8227 "muse-glimmer-8b",
8228 crate::sampling_knobs::SamplerModel::absent(),
8229 ) {
8230 Err(e) => e,
8231 Ok(_) => panic!("a muse_glimmer call's boundary is a channel, not a marker"),
8232 };
8233 assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
8234 assert!(
8235 body["error"]["message"]
8236 .as_str()
8237 .unwrap()
8238 .contains("muse_glimmer"),
8239 "{body}"
8240 );
8241
8242 let req = tool_request(serde_json::json!("required"));
8248 let params = req
8249 .generation_params_for_template(
8250 &graded_template(),
8251 "gemma-4-E2B-it",
8252 crate::sampling_knobs::SamplerModel::absent(),
8253 )
8254 .expect("a gemma4 forced tool_choice is served");
8255 assert!(
8256 params.grammar.is_some(),
8257 "a forced tool_choice must arrive as the generation's grammar"
8258 );
8259 }
8260
8261 #[test]
8264 fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
8265 let req = chat_request(serde_json::json!({
8266 "model": "m",
8267 "messages": [{"role": "user", "content": "hi"}],
8268 "grammar": "root ::= \"a",
8269 }));
8270 let (status, Json(body)) = req
8271 .validate_supported_fields()
8272 .expect_err("this does not parse");
8273 assert_eq!(status, StatusCode::BAD_REQUEST);
8274 assert_eq!(body["error"]["param"], "grammar");
8275 assert!(
8276 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8277 .is_err(),
8278 "and again at params time"
8279 );
8280 }
8281
8282 #[test]
8289 fn response_format_json_schema_becomes_the_requests_grammar() {
8290 let req = chat_request(serde_json::json!({
8291 "model": "m",
8292 "messages": [{"role": "user", "content": "hi"}],
8293 "response_format": {
8294 "type": "json_schema",
8295 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8296 },
8297 }));
8298 req.validate_supported_fields()
8299 .expect("a boolean schema converts");
8300 let params = req
8301 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8302 .expect("and compiles");
8303 let grammar = params.grammar.expect("the schema is the grammar");
8304 let mut g = (*grammar).clone();
8305 g.accept_token(0, b"true").expect("a boolean is accepted");
8306 assert!(g.allows_eog(), "and completes the parse");
8307 assert!(
8308 !params.json_object,
8309 "a schema is not the json_object character-class mask"
8310 );
8311 }
8312
8313 #[test]
8317 fn an_unconvertible_response_format_schema_is_a_400_naming_the_keyword() {
8318 let req = chat_request(serde_json::json!({
8319 "model": "m",
8320 "messages": [{"role": "user", "content": "hi"}],
8321 "response_format": {
8322 "type": "json_schema",
8323 "json_schema": {"name": "x", "schema": {"type": "integer", "minimum": 3}},
8324 },
8325 }));
8326 let (status, Json(body)) = req
8327 .validate_supported_fields()
8328 .expect_err("minimum has no grammar in this port");
8329 assert_eq!(status, StatusCode::BAD_REQUEST);
8330 assert!(
8331 body["error"]["message"]
8332 .as_str()
8333 .expect("a message")
8334 .contains("minimum"),
8335 "the refusal must name the keyword: {body}"
8336 );
8337 assert!(
8338 req.generation_params(crate::sampling_knobs::SamplerModel::absent())
8339 .is_err(),
8340 "and again at params time"
8341 );
8342 }
8343
8344 #[test]
8350 fn a_forced_tool_choice_and_a_schema_are_two_constraints() {
8351 let req = chat_request(serde_json::json!({
8352 "model": "m",
8353 "messages": [{"role": "user", "content": "hi"}],
8354 "tool_choice": "required",
8355 "tools": [{
8356 "type": "function",
8357 "function": {"name": "f", "parameters": {"type": "object"}},
8358 }],
8359 "response_format": {
8360 "type": "json_schema",
8361 "json_schema": {"name": "x", "schema": {"type": "boolean"}},
8362 },
8363 }));
8364 let (status, Json(body)) = req
8365 .validate_supported_fields()
8366 .expect_err("two constraints, one generation");
8367 assert_eq!(status, StatusCode::BAD_REQUEST);
8368 assert_eq!(body["error"]["param"], "tool_choice");
8369 }
8370
8371 #[test]
8374 fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
8375 let req = chat_request(serde_json::json!({
8376 "model": "m",
8377 "messages": [{"role": "user", "content": "hi"}],
8378 }));
8379 assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
8380 }
8381
8382 #[test]
8387 fn min_p_reaches_the_sampler_from_the_chat_wire() {
8388 let asked = chat_request(serde_json::json!({
8389 "model": "m",
8390 "messages": [{"role": "user", "content": "hi"}],
8391 "min_p": 0.07,
8392 }));
8393 assert_eq!(
8394 asked
8395 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8396 .expect("knobs")
8397 .min_p,
8398 0.07
8399 );
8400
8401 let silent = chat_request(serde_json::json!({
8402 "model": "m",
8403 "messages": [{"role": "user", "content": "hi"}],
8404 }));
8405 assert_eq!(
8406 silent
8407 .sampling_params(crate::sampling_knobs::SamplerModel::absent())
8408 .expect("knobs")
8409 .min_p,
8410 0.0,
8411 "an unset min_p must be off, not llama.cpp's CLI default"
8412 );
8413 }
8414
8415 #[test]
8424 fn no_sampler_knob_is_missing_from_the_cache_key() {
8425 let base = serde_json::json!({
8426 "model": "m",
8427 "messages": [{"role": "user", "content": "hi"}],
8428 "seed": 1,
8429 });
8430 let key_for = |body: serde_json::Value| {
8431 let req = chat_request(body);
8432 let params = req
8433 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8434 .expect("params");
8435 req.cache_key("prompt", ¶ms)
8436 };
8437 let baseline = key_for(base.clone());
8438 for (knob, value) in [
8439 ("temperature", serde_json::json!(0.5)),
8440 ("top_p", serde_json::json!(0.9)),
8441 ("min_p", serde_json::json!(0.05)),
8442 ("top_k", serde_json::json!(40)),
8443 ("repetition_penalty", serde_json::json!(1.1)),
8444 ("presence_penalty", serde_json::json!(0.3)),
8445 ("frequency_penalty", serde_json::json!(0.3)),
8446 (
8447 "samplers",
8448 serde_json::json!(["penalties", "top_p", "top_k", "min_p", "temperature"]),
8449 ),
8450 ] {
8451 let mut body = base.clone();
8452 body[knob] = value;
8453 assert_ne!(
8454 key_for(body),
8455 baseline,
8456 "`{knob}` is not in the cache key: two requests differing \
8457 only in it would share one cached answer"
8458 );
8459 }
8460 }
8461
8462 #[test]
8472 fn no_constraint_is_missing_from_the_cache_key() {
8473 let base = serde_json::json!({
8474 "model": "m",
8475 "messages": [{"role": "user", "content": "pick one"}],
8476 });
8477 let key_for = |body: serde_json::Value| {
8478 let req = chat_request(body);
8479 let params = req
8480 .generation_params(crate::sampling_knobs::SamplerModel::absent())
8481 .expect("params");
8482 req.cache_key("prompt", ¶ms)
8483 };
8484 let baseline = key_for(base.clone());
8485 for (field, value) in [
8486 ("grammar", serde_json::json!("root ::= \"yes\" | \"no\"")),
8487 (
8488 "response_format",
8489 serde_json::json!({"type": "json_object"}),
8490 ),
8491 (
8492 "response_format",
8493 serde_json::json!({"type": "json_schema", "json_schema": {
8494 "name": "answer",
8495 "schema": {"type": "object", "properties": {"a": {"type": "string"}}}
8496 }}),
8497 ),
8498 ("ignore_eos", serde_json::json!(true)),
8499 ("stop", serde_json::json!(["\n"])),
8500 ("max_tokens", serde_json::json!(7)),
8501 ] {
8502 let mut body = base.clone();
8503 body[field] = value.clone();
8504 assert_ne!(
8505 key_for(body),
8506 baseline,
8507 "`{field}: {value}` is not in the cache key: two requests \
8508 differing only in it would share one cached answer"
8509 );
8510 }
8511 }
8512
8513 #[test]
8517 fn an_explicit_zero_output_budget_is_a_client_error() {
8518 let req = chat_request(serde_json::json!({
8519 "model": "m",
8520 "messages": [{"role": "user", "content": "hi"}],
8521 "max_tokens": 0,
8522 }));
8523 let (status, body) = req.validate_supported_fields().expect_err("rejected");
8524 assert_eq!(status, StatusCode::BAD_REQUEST);
8525 assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
8526 }
8527
8528 #[test]
8531 fn a_request_can_turn_thinking_off() {
8532 let template = graded_template();
8533 for body in [
8534 serde_json::json!({
8535 "model": "m",
8536 "messages": [{"role": "user", "content": "hi"}],
8537 "reasoning_effort": "none",
8538 }),
8539 serde_json::json!({
8540 "model": "m",
8541 "messages": [{"role": "user", "content": "hi"}],
8542 "thinking": {"type": "disabled"},
8543 }),
8544 ] {
8545 let kwargs = chat_request(body).resolve_template_kwargs(&template);
8546 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8547 assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
8548 assert!(!kwargs.contains_key("reasoning_effort"));
8551 }
8552 }
8553
8554 #[test]
8557 fn a_disabled_switch_beats_an_effort_in_the_same_request() {
8558 let template = graded_template();
8559 let kwargs = chat_request(serde_json::json!({
8560 "model": "m",
8561 "messages": [{"role": "user", "content": "hi"}],
8562 "reasoning_effort": "high",
8563 "thinking": {"type": "disabled"},
8564 }))
8565 .resolve_template_kwargs(&template);
8566 assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
8567 assert!(!kwargs.contains_key("reasoning_effort"));
8568 }
8569
8570 #[test]
8573 fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
8574 let req = chat_request(serde_json::json!({
8575 "model": "m",
8576 "messages": [{"role": "user", "content": "hi"}],
8577 "thinking": {"type": "disable"},
8578 }));
8579 let (status, _) = req.validate_supported_fields().expect_err("rejected");
8580 assert_eq!(status, StatusCode::BAD_REQUEST);
8581 }
8582
8583 #[test]
8586 fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
8587 let template = graded_template();
8588 let kwargs = chat_request(serde_json::json!({
8589 "model": "m",
8590 "messages": [{"role": "user", "content": "hi"}],
8591 "reasoning_effort": "none",
8592 "chat_template_kwargs": {"enable_thinking": true},
8593 }))
8594 .resolve_template_kwargs(&template);
8595 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8596 }
8597
8598 #[test]
8602 fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
8603 let template = graded_template();
8604 let req = chat_request(serde_json::json!({
8605 "model": "m",
8606 "messages": [{"role": "user", "content": "hi"}],
8607 "reasoning_effort": "minimal",
8608 }));
8609 let kwargs = req.resolve_template_kwargs(&template);
8610 assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
8611 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8612 assert!(prompt.starts_with("E:low|"), "{prompt}");
8613 }
8614
8615 #[test]
8619 fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
8620 let template = graded_template();
8621 let req = chat_request(serde_json::json!({
8622 "model": "m",
8623 "messages": [{"role": "user", "content": "hi"}],
8624 "chat_template_kwargs": {"reasoning_effort": "none"},
8625 }));
8626 let kwargs = req.resolve_template_kwargs(&template);
8627 assert!(!kwargs.contains_key("reasoning_effort"));
8628 let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
8629 assert_eq!(prompt, "hi");
8630 }
8631
8632 #[test]
8635 fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
8636 let template = graded_template();
8637 let req = chat_request(serde_json::json!({
8638 "model": "m",
8639 "messages": [{"role": "user", "content": "hi"}],
8640 "reasoning_effort": "low",
8641 "chat_template_kwargs": {"reasoning_effort": "high"},
8642 }));
8643 assert_eq!(
8644 req.resolve_template_kwargs(&template)["reasoning_effort"],
8645 serde_json::json!("high")
8646 );
8647 }
8648
8649 #[test]
8653 fn offering_tools_turns_thinking_on_by_itself() {
8654 let template = graded_template();
8655 let quiet = chat_request(serde_json::json!({
8656 "model": "m",
8657 "messages": [{"role": "user", "content": "hi"}],
8658 }));
8659 assert!(!quiet
8660 .resolve_template_kwargs(&template)
8661 .contains_key("enable_thinking"));
8662
8663 let with_tools = chat_request(serde_json::json!({
8664 "model": "m",
8665 "messages": [{"role": "user", "content": "hi"}],
8666 "tools": [{"type": "function", "function": {"name": "get_weather"}}],
8667 }));
8668 let kwargs = with_tools.resolve_template_kwargs(&template);
8669 assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
8670 let prompt =
8671 prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
8672 assert!(prompt.starts_with("THINK|"), "{prompt}");
8673 }
8674
8675 #[test]
8680 fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
8681 let opener = chat_template::PromptTemplate::from_gguf_metadata(
8682 Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
8683 Some("qwen3"),
8684 false,
8685 true,
8686 None,
8687 None,
8688 );
8689 let req = chat_request(serde_json::json!({
8690 "model": "m",
8691 "messages": [{"role": "user", "content": "hi"}],
8692 "chat_template_kwargs": {"enable_thinking": true},
8693 }));
8694 let kwargs = req.resolve_template_kwargs(&opener);
8695 let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
8696 assert!(prompt.ends_with("<think>"), "{prompt}");
8697
8698 let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
8701 let (message, _) = build_response_message(
8702 "weighing it up</think>Paris.".to_string(),
8703 &[],
8704 posture,
8705 "stop",
8706 );
8707 assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
8708 assert_eq!(message.content.as_deref(), Some("Paris."));
8709
8710 let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
8713 let (message, _) = build_response_message(
8714 "weighing it up</think>Paris.".to_string(),
8715 &[],
8716 closed,
8717 "stop",
8718 );
8719 assert_eq!(message.reasoning_content, None);
8720 }
8721
8722 #[test]
8723 fn stop_param_accepts_both_single_string_and_array() {
8724 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8725 "model": "m",
8726 "messages": [{"role": "user", "content": "hi"}],
8727 "stop": "END",
8728 }))
8729 .unwrap();
8730 assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
8731
8732 let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
8733 "model": "m",
8734 "messages": [{"role": "user", "content": "hi"}],
8735 "stop": ["A", "B"],
8736 }))
8737 .unwrap();
8738 assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
8739 }
8740
8741 #[test]
8742 fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
8743 let model = test_model();
8744 let result = run_generation(
8745 &model,
8746 "hello",
8747 &greedy_params(4),
8748 None,
8749 None,
8750 None,
8751 None,
8752 None,
8753 None,
8754 );
8755 assert!(matches!(
8756 result,
8757 Err(generate::DecodeError::TokenOutOfVocab { .. })
8758 ));
8759 }
8760
8761 #[test]
8765 fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
8766 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8768 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8769
8770 let holder_pool = Arc::clone(&pool);
8771 let holder = std::thread::spawn(move || {
8772 let mut held = frink_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
8773 held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
8775 drop(held);
8776 });
8777 std::thread::sleep(Duration::from_millis(15));
8778
8779 let config = generate::KvPoolConfig {
8780 pool,
8781 queue_wait: Duration::ZERO,
8782 };
8783 let result = run_generation(
8784 &model,
8785 &prompt,
8786 &greedy_params(4),
8787 Some(&config),
8788 None,
8789 None,
8790 None,
8791 None,
8792 None,
8793 );
8794 assert!(matches!(
8795 result,
8796 Err(generate::DecodeError::KvPoolExhausted)
8797 ));
8798
8799 let (status, _body) = decode_error_response(result.unwrap_err());
8800 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8801 holder.join().unwrap();
8802 }
8803
8804 #[test]
8812 fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
8813 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8815 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 1)));
8817 let config = generate::KvPoolConfig {
8818 pool,
8819 queue_wait: Duration::ZERO,
8820 };
8821
8822 let result = run_generation(
8823 &model,
8824 &prompt,
8825 &greedy_params(4),
8826 Some(&config),
8827 None,
8828 None,
8829 None,
8830 None,
8831 None,
8832 );
8833 let err = result.expect_err("one block cannot hold two layers' caches");
8834 assert!(
8835 matches!(
8836 &err,
8837 generate::DecodeError::KvBudgetExceeded { binding, .. }
8838 if *binding == frink_models::Ceiling::DeviceMemory.code()
8839 ),
8840 "expected an immovable device-memory refusal, got {err:?}"
8841 );
8842 let (status, _body) = decode_error_response(err);
8843 assert_eq!(status, StatusCode::BAD_REQUEST);
8844 }
8845
8846 #[test]
8852 fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
8853 let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
8854 queued: 512,
8855 cap: 512,
8856 });
8857 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
8858 assert_eq!(body["error"]["retry_after_seconds"], 1);
8859 let message = body["error"]["message"].as_str().expect("message");
8860 assert!(message.contains("512"), "{message}");
8861 }
8862
8863 #[test]
8864 fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
8865 let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8866 token: 99,
8867 vocab_size: 32,
8868 });
8869 assert!(
8870 body["error"]["retry_after_seconds"].is_null(),
8871 "retrying a prompt this model cannot tokenize never helps"
8872 );
8873 }
8874
8875 #[test]
8876 fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
8877 let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
8878 token: 99,
8879 vocab_size: 32,
8880 });
8881 assert_eq!(status, StatusCode::BAD_REQUEST);
8882 }
8883
8884 #[test]
8885 fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
8886 let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8888 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 2)));
8889 let config = generate::KvPoolConfig {
8890 pool: pool.clone(),
8891 queue_wait: Duration::ZERO,
8892 };
8893
8894 let (choices, _usage) = run_generation(
8895 &model,
8896 &prompt,
8897 &greedy_params(4),
8898 Some(&config),
8899 None,
8900 None,
8901 None,
8902 None,
8903 None,
8904 )
8905 .unwrap();
8906 assert_eq!(choices[0].finish, FinishReason::Length);
8907 assert_eq!(
8908 pool.lock().unwrap().free_blocks(),
8909 2,
8910 "a completed request must return its blocks to the pool"
8911 );
8912 }
8913
8914 #[tokio::test]
8919 async fn concurrent_requests_against_the_same_model_do_not_interfere() {
8920 let model = Arc::new(test_model());
8921 let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
8922
8923 let mut handles = Vec::new();
8924 for _ in 0..8 {
8925 let model = Arc::clone(&model);
8926 let prompt = prompt.clone();
8927 handles.push(tokio::task::spawn_blocking(move || {
8928 run_generation(
8929 &model,
8930 &prompt,
8931 &greedy_params(6),
8932 None,
8933 None,
8934 None,
8935 None,
8936 None,
8937 None,
8938 )
8939 .unwrap()
8940 }));
8941 }
8942
8943 let mut results = Vec::new();
8944 for h in handles {
8945 results.push(h.await.unwrap());
8946 }
8947 for r in &results[1..] {
8952 assert_eq!(r.0, results[0].0, "choices must match");
8956 assert_eq!(
8957 r.1.prompt_tokens, results[0].1.prompt_tokens,
8958 "prompt token count must match"
8959 );
8960 assert_eq!(
8961 r.1.completion_tokens, results[0].1.completion_tokens,
8962 "completion token count must match"
8963 );
8964 }
8965 }
8966
8967 fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
8974 let mut header_entries = Vec::new();
8975 let mut data = Vec::new();
8976 for (name, shape, values) in tensors {
8977 let start = data.len();
8978 for v in values {
8979 data.extend_from_slice(&v.to_le_bytes());
8980 }
8981 let end = data.len();
8982 let shape_str = shape
8983 .iter()
8984 .map(|d| d.to_string())
8985 .collect::<Vec<_>>()
8986 .join(",");
8987 header_entries.push(format!(
8988 "\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
8989 ));
8990 }
8991 let header = format!("{{{}}}", header_entries.join(","));
8992 let header_bytes = header.as_bytes();
8993 let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
8994 out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
8995 out.extend_from_slice(header_bytes);
8996 out.extend_from_slice(&data);
8997 out
8998 }
8999
9000 fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
9011 use frink_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
9012 use frink_models::kimi_loader::KimiRealHparams;
9013 use frink_moe::{GatingFunction, MoeLayerConfig};
9014
9015 let hidden_dim = 8;
9016 let kda_num_heads = 2;
9017 let kda_head_dim = 3;
9018 let kda_proj = kda_num_heads * kda_head_dim;
9019 let conv_kernel = 4;
9020 let dense_intermediate = 5;
9021 let vocab_size = 256;
9025 let mla_num_heads = 1;
9026 let mla_q_lora_rank = 2;
9027 let mla_kv_lora_rank = 2;
9028 let mla_qk_nope_head_dim = 2;
9029 let mla_qk_rope_head_dim = 2;
9030 let mla_v_head_dim = 2;
9031
9032 let model_cfg = frink_models::ModelConfig {
9033 rope_layers: frink_models::rope_layers::RopeLayers::All,
9034 layer_shapes: frink_models::layer_shapes::LayerShapes::Uniform,
9035 name: "synthetic-kimi-server-test",
9036 n_layers: 1,
9037 n_mtp_blocks: 0,
9038 hidden_dim,
9039 n_heads: 1,
9040 n_kv_heads: 1,
9041 head_dim: 4,
9042 v_head_dim: None,
9043 vocab_size,
9044 rope_theta: 10000.0,
9045 rms_norm_eps: 1e-5,
9046 post_norm_eps: 1e-5,
9047 sliding_window: None,
9048 moe: MoeLayerConfig {
9049 expert_weights_scale: 1.0,
9050 routed_weight_before_ffn: false,
9051 n_experts: 1,
9052 n_experts_active: 1,
9053 n_shared_experts: 0,
9054 hidden_dim,
9055 expert_ffn_dim: 4,
9056 gating: GatingFunction::Sigmoid,
9057 norm_topk_prob: true,
9058 expert_group_count: None,
9059 expert_group_used_count: None,
9060 },
9061 n_dense_leading_layers: 1,
9066 moe_interleave_step: None,
9067 norm_function: frink_models::norm::NormFunction::Rms,
9068 attention: AttentionKind::KimiHybrid(KimiHybridAttention {
9069 kda_layers: vec![1],
9070 full_attn_layers: vec![],
9071 mla: MlaConfig {
9072 num_heads: mla_num_heads,
9073 q_lora_rank: mla_q_lora_rank,
9074 kv_lora_rank: mla_kv_lora_rank,
9075 qk_nope_head_dim: mla_qk_nope_head_dim,
9076 qk_rope_head_dim: mla_qk_rope_head_dim,
9077 v_head_dim: mla_v_head_dim,
9078 use_output_gate: true,
9079 rope: None,
9080 },
9081 kda: KdaConfig {
9082 num_heads: kda_num_heads,
9083 head_dim: kda_head_dim,
9084 short_conv_kernel_size: conv_kernel,
9085 gate_lower_bound: -5.0,
9086 use_full_rank_gate: true,
9087 },
9088 }),
9089 rope_freqs: None,
9090 rope_attn_factor: 1.0,
9091 rope_dim: None,
9092 rope_dim_swa: None,
9093 rope_freqs_long: None,
9094 rope_freqs_short: None,
9095 rope_orig_ctx: None,
9096 rope_layout: frink_models::config::RopeLayout::Neox,
9097 qk_norm_style: frink_models::capability::QkNormStyle::WholeVector,
9098 swa_layers: frink_models::swa_layers::SwaLayers::All,
9099 attn_logit_softcap: None,
9100 final_logit_softcap: None,
9101 embedding_scale: None,
9102 residual_scale: None,
9103 normed_residual_scale: None,
9104 clamp_kqv: None,
9105 attn_temperature: None,
9106 router_input: frink_models::router_input::RouterInput::NormedFfnInput,
9107 block_sub_norms: false,
9108 parallel_residual: false,
9109 learned_positions: false,
9110 attn_value_scale: None,
9111 alibi_max_bias: None,
9112 layer_loops: None,
9113 skip_stream: false,
9114 parallel_ssm: false,
9115 swa_chunked: false,
9116 weightless_qk_norm: false,
9117 logit_multiplier: None,
9118 attention_scale: None,
9119 rope_theta_swa: None,
9120 ffn_activation: frink_models::config::FfnActivation::Swiglu,
9121 best_effort_fields: &["synthetic test config, not a real preset"],
9122 };
9123 let hp = KimiRealHparams {
9124 hidden_dim,
9125 kda_num_heads,
9126 kda_head_dim,
9127 mla_num_heads,
9128 mla_q_lora_rank,
9129 mla_kv_lora_rank,
9130 mla_qk_nope_head_dim,
9131 mla_qk_rope_head_dim,
9132 mla_v_head_dim,
9133 dense_intermediate_dim: dense_intermediate,
9134 moe_hidden_dim: hidden_dim,
9135 moe_intermediate_dim: 4,
9136 n_experts: 1,
9137 num_shared_experts: 0,
9138 };
9139
9140 let prefix = "language_model.model.layers.0";
9144 let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
9145 let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
9146 name: String,
9147 shape: Vec<usize>,
9148 n: usize| {
9149 tensors.push((name, shape, vec![0.01f32; n]));
9150 };
9151 push(
9152 &mut tensors,
9153 format!("{prefix}.input_layernorm.weight"),
9154 vec![hidden_dim],
9155 hidden_dim,
9156 );
9157 push(
9158 &mut tensors,
9159 format!("{prefix}.post_attention_layernorm.weight"),
9160 vec![hidden_dim],
9161 hidden_dim,
9162 );
9163 push(
9164 &mut tensors,
9165 format!("{prefix}.self_attention_res_norm.weight"),
9166 vec![hidden_dim],
9167 hidden_dim,
9168 );
9169 push(
9170 &mut tensors,
9171 format!("{prefix}.self_attention_res_proj.weight"),
9172 vec![1, hidden_dim],
9173 hidden_dim,
9174 );
9175 push(
9176 &mut tensors,
9177 format!("{prefix}.mlp_res_norm.weight"),
9178 vec![hidden_dim],
9179 hidden_dim,
9180 );
9181 push(
9182 &mut tensors,
9183 format!("{prefix}.mlp_res_proj.weight"),
9184 vec![1, hidden_dim],
9185 hidden_dim,
9186 );
9187 push(
9188 &mut tensors,
9189 format!("{prefix}.self_attn.q_proj.weight"),
9190 vec![kda_proj, hidden_dim],
9191 kda_proj * hidden_dim,
9192 );
9193 push(
9194 &mut tensors,
9195 format!("{prefix}.self_attn.k_proj.weight"),
9196 vec![kda_proj, hidden_dim],
9197 kda_proj * hidden_dim,
9198 );
9199 push(
9200 &mut tensors,
9201 format!("{prefix}.self_attn.v_proj.weight"),
9202 vec![kda_proj, hidden_dim],
9203 kda_proj * hidden_dim,
9204 );
9205 push(
9206 &mut tensors,
9207 format!("{prefix}.self_attn.q_conv1d.weight"),
9208 vec![kda_proj, 1, conv_kernel],
9209 kda_proj * conv_kernel,
9210 );
9211 push(
9212 &mut tensors,
9213 format!("{prefix}.self_attn.k_conv1d.weight"),
9214 vec![kda_proj, 1, conv_kernel],
9215 kda_proj * conv_kernel,
9216 );
9217 push(
9218 &mut tensors,
9219 format!("{prefix}.self_attn.v_conv1d.weight"),
9220 vec![kda_proj, 1, conv_kernel],
9221 kda_proj * conv_kernel,
9222 );
9223 push(
9224 &mut tensors,
9225 format!("{prefix}.self_attn.A_log"),
9226 vec![kda_num_heads],
9227 kda_num_heads,
9228 );
9229 push(
9230 &mut tensors,
9231 format!("{prefix}.self_attn.f_a_proj.weight"),
9232 vec![kda_head_dim, hidden_dim],
9233 kda_head_dim * hidden_dim,
9234 );
9235 push(
9236 &mut tensors,
9237 format!("{prefix}.self_attn.f_b_proj.weight"),
9238 vec![kda_proj, kda_head_dim],
9239 kda_proj * kda_head_dim,
9240 );
9241 push(
9242 &mut tensors,
9243 format!("{prefix}.self_attn.dt_bias"),
9244 vec![kda_proj],
9245 kda_proj,
9246 );
9247 push(
9248 &mut tensors,
9249 format!("{prefix}.self_attn.b_proj.weight"),
9250 vec![kda_num_heads, hidden_dim],
9251 kda_num_heads * hidden_dim,
9252 );
9253 push(
9254 &mut tensors,
9255 format!("{prefix}.self_attn.g_proj.weight"),
9256 vec![kda_proj, hidden_dim],
9257 kda_proj * hidden_dim,
9258 );
9259 push(
9260 &mut tensors,
9261 format!("{prefix}.self_attn.o_norm.weight"),
9262 vec![kda_head_dim],
9263 kda_head_dim,
9264 );
9265 push(
9266 &mut tensors,
9267 format!("{prefix}.self_attn.o_proj.weight"),
9268 vec![hidden_dim, kda_proj],
9269 hidden_dim * kda_proj,
9270 );
9271 push(
9272 &mut tensors,
9273 format!("{prefix}.mlp.gate_proj.weight"),
9274 vec![dense_intermediate, hidden_dim],
9275 dense_intermediate * hidden_dim,
9276 );
9277 push(
9278 &mut tensors,
9279 format!("{prefix}.mlp.up_proj.weight"),
9280 vec![dense_intermediate, hidden_dim],
9281 dense_intermediate * hidden_dim,
9282 );
9283 push(
9284 &mut tensors,
9285 format!("{prefix}.mlp.down_proj.weight"),
9286 vec![hidden_dim, dense_intermediate],
9287 hidden_dim * dense_intermediate,
9288 );
9289 push(
9290 &mut tensors,
9291 "language_model.model.embed_tokens.weight".to_string(),
9292 vec![vocab_size, hidden_dim],
9293 vocab_size * hidden_dim,
9294 );
9295 push(
9296 &mut tensors,
9297 "language_model.lm_head.weight".to_string(),
9298 vec![vocab_size, hidden_dim],
9299 vocab_size * hidden_dim,
9300 );
9301 push(
9302 &mut tensors,
9303 "language_model.model.norm.weight".to_string(),
9304 vec![hidden_dim],
9305 hidden_dim,
9306 );
9307 push(
9308 &mut tensors,
9309 "language_model.model.output_attn_res_norm.weight".to_string(),
9310 vec![hidden_dim],
9311 hidden_dim,
9312 );
9313 push(
9314 &mut tensors,
9315 "language_model.model.output_attn_res_proj.weight".to_string(),
9316 vec![1, hidden_dim],
9317 hidden_dim,
9318 );
9319
9320 static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
9333 let dir = std::env::temp_dir().join(format!(
9334 "frink_server_kimi_e2e_test_{}_{}",
9335 std::process::id(),
9336 FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
9337 ));
9338 std::fs::create_dir_all(&dir).unwrap();
9339 let shard_bytes = write_safetensors_shard(&tensors);
9340 std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
9341 let map_entries: Vec<String> = tensors
9342 .iter()
9343 .map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
9344 .collect();
9345 let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
9346 std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
9347
9348 use base64::Engine;
9352 let vocab_lines: Vec<String> = (0..vocab_size as u32)
9353 .map(|b| {
9354 let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
9355 format!("{b64} {b}")
9356 })
9357 .collect();
9358 std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
9359
9360 let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
9361 .expect("must load the synthetic Kimi checkpoint end to end");
9362 std::fs::remove_dir_all(&dir).ok();
9363 loaded
9364 }
9365
9366 #[test]
9374 fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
9375 let loaded = build_synthetic_kimi_loaded();
9376 let state = build_app_state(
9377 StartupModels {
9378 loaded: model::LoadedModel::Kimi(loaded),
9379 embedding: None,
9380 },
9381 None,
9382 None,
9383 None,
9384 false,
9385 None,
9386 Arc::new(health::Detection::ready(health::probe_backends())),
9387 );
9388 let active = state.active().expect("a freshly built state has a model");
9389 assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
9390 assert!(!active.is_synthetic());
9391
9392 let (choices, _usage) = run_generation(
9393 active.generative().unwrap(),
9394 "hi",
9395 &greedy_params(5),
9396 None,
9397 None,
9398 None,
9399 None,
9400 None,
9401 None,
9402 )
9403 .expect("a real Kimi checkpoint must generate without error");
9404 assert!(matches!(
9405 choices[0].finish,
9406 FinishReason::Length | FinishReason::Stop
9407 ));
9408 }
9409
9410 #[test]
9423 fn a_grammar_constrains_the_engine_decode_path() {
9424 let loaded = build_synthetic_kimi_loaded();
9425 let state = build_app_state(
9426 StartupModels {
9427 loaded: model::LoadedModel::Kimi(loaded),
9428 embedding: None,
9429 },
9430 None,
9431 None,
9432 None,
9433 false,
9434 None,
9435 Arc::new(health::Detection::ready(health::probe_backends())),
9436 );
9437 let active = state.active().expect("a freshly built state has a model");
9438
9439 let run = |grammar: Option<&str>| {
9440 let mut params = greedy_params(6);
9441 params.grammar = grammar.map(|src| {
9442 Arc::new(
9443 frink_models::grammar::Grammar::from_str_with_root(src, "root")
9444 .expect("test grammar parses"),
9445 )
9446 });
9447 run_generation(
9448 active.generative().unwrap(),
9449 "hi",
9450 ¶ms,
9451 None,
9452 None,
9453 None,
9454 None,
9455 None,
9456 None,
9457 )
9458 };
9459
9460 let (choices, _) = run(None).expect("the unconstrained run must serve");
9461 let unconstrained = choices[0].text.clone();
9462 assert!(
9463 unconstrained.chars().any(|c| c != 'a'),
9464 "the unconstrained run produced only `a` ({unconstrained:?}), so the \
9465 constrained run below would prove nothing"
9466 );
9467
9468 let (choices, _) =
9469 run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
9470 let one = choices.into_iter().next().unwrap();
9471 let (finish, constrained) = (one.finish, one.text);
9472 assert!(
9473 !constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
9474 "the engine decode path served text its grammar forbids ({constrained:?}): \
9475 the constraint was dropped between `generate_engine` and the sampler"
9476 );
9477 assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
9478 }
9479
9480 #[test]
9490 fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
9491 let loaded = build_synthetic_kimi_loaded();
9492 let state = build_app_state(
9493 StartupModels {
9494 loaded: model::LoadedModel::Kimi(loaded),
9495 embedding: None,
9496 },
9497 None,
9498 None,
9499 None,
9500 false,
9501 None,
9502 Arc::new(health::Detection::ready(health::probe_backends())),
9503 );
9504
9505 let pool = Arc::new(Mutex::new(frink_core::cache::KvBlockPool::new(64, 4)));
9506 let kv_pool_config = generate::KvPoolConfig {
9507 pool: pool.clone(),
9508 queue_wait: Duration::ZERO,
9509 };
9510 let pc = Mutex::new(PrefixCache::new(4));
9511
9512 run_generation(
9513 state
9514 .active()
9515 .expect("a freshly built state has a model")
9516 .generative()
9517 .unwrap(),
9518 "hi",
9519 &greedy_params(5),
9520 Some(&kv_pool_config),
9521 None,
9522 Some(&pc),
9523 None,
9524 None,
9525 None,
9526 )
9527 .expect("a real Kimi checkpoint must generate without error");
9528
9529 assert_eq!(
9530 pool.lock().unwrap().free_blocks(),
9531 4,
9532 "the KV pool must be completely untouched by a Kimi request"
9533 );
9534 let stats = pc.lock().unwrap().stats();
9535 assert_eq!(
9536 stats.hits + stats.misses,
9537 0,
9538 "the prefix cache must never be consulted for a Kimi request"
9539 );
9540 }
9541}