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