mod admin;
mod anthropic;
mod attribution;
mod budget;
mod cache_admin;
mod cancel;
mod chat_template;
mod completion;
mod conversations;
mod decode_task;
mod embeddings;
mod generate;
mod grammar_request;
mod health;
mod journal;
mod json_mode;
mod limits;
mod loaded;
mod mcp;
mod model;
mod openai_extra;
mod output;
mod policy;
mod response_cache;
pub(crate) mod responses;
mod resume;
mod sample_step;
mod sampling_knobs;
mod security;
mod serving;
mod session;
mod sse;
mod stats;
mod stop;
mod stream_events;
mod tasks;
mod tool_grammar;
mod unsupported_sampling;
use std::cell::RefCell;
use std::convert::Infallible;
use std::fmt;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::path::PathBuf;
use std::rc::Rc;
use std::str::FromStr;
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use axum::{
extract::State,
http::StatusCode,
response::sse::{Event, Sse},
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use clap::{Parser, ValueEnum};
use serde::{Deserialize, Serialize};
use ferrox_core::cache::KvBlockPool;
use ferrox_models::kimi_tokenizer::KimiTokenizer;
use ferrox_models::sampling::SamplingParams;
use ferrox_models::tokenizer::StopTokens;
use ferrox_models::{Decoder, Gemma4Engine, KimiEngine, MlaEngine, PrefixCache};
use generate::{FinishReason, GenerationParams};
pub(crate) use loaded::{ActiveModel, Loaded};
use model::ServerTokenizer;
use response_cache::{CacheKey, ResponseCache};
use sampling_knobs::SamplingKnobs;
#[derive(Parser, Debug, PartialEq)]
#[command(
name = "ferrox-server",
about = "OpenAI-compatible Ferrox inference server"
)]
pub struct ServerArgs {
#[arg(short = 'm', long = "model", value_name = "FILE")]
model: Option<String>,
#[arg(long, value_name = "HOST")]
host: Option<IpAddr>,
#[arg(long, value_name = "PORT")]
port: Option<u16>,
#[arg(short = 't', long = "threads", value_name = "N")]
threads: Option<usize>,
#[arg(
long = "device",
visible_alias = "dev",
value_name = "DEVICE",
ignore_case = true
)]
device: Option<OffloadDevice>,
#[arg(long = "list-devices", default_value_t = false)]
list_devices: bool,
#[arg(
long = "n-gpu-layers",
visible_aliases = ["gpu-layers", "ngl"],
value_name = "N"
)]
n_gpu_layers: Option<GpuLayers>,
#[arg(long = "mcp-config", value_name = "PATH")]
mcp_config: Option<PathBuf>,
#[arg(long = "exit-on-stdin-close", default_value_t = false)]
exit_on_stdin_close: bool,
#[arg(long = "allow-multiple-instances", default_value_t = false)]
allow_multiple_instances: bool,
}
impl ServerArgs {
pub fn parse_llama_style<I>(argv: I) -> Self
where
I: IntoIterator<Item = String>,
{
Self::parse_from(rewrite_llama_style_argv(argv.into_iter().collect()))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
enum OffloadDevice {
Auto,
None,
Cpu,
Metal,
Cuda,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum GpuLayers {
Auto,
All,
Count(u32),
}
impl GpuLayers {
fn offload_enabled(self) -> bool {
!matches!(self, Self::Count(0))
}
}
impl FromStr for GpuLayers {
type Err = String;
fn from_str(value: &str) -> Result<Self, Self::Err> {
match value {
"auto" => Ok(Self::Auto),
"all" => Ok(Self::All),
_ => value
.parse::<u32>()
.map(Self::Count)
.map_err(|_| "expected 0, a positive integer, 'auto', or 'all'".into()),
}
}
}
impl fmt::Display for GpuLayers {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Auto => f.write_str("auto"),
Self::All => f.write_str("all"),
Self::Count(value) => value.fmt(f),
}
}
}
pub const BUILT_WITH_METAL: bool = cfg!(feature = "metal");
pub const BUILT_WITH_CUDA: bool = cfg!(feature = "cuda");
fn rewrite_llama_style_argv(args: Vec<String>) -> Vec<String> {
args.into_iter()
.map(|arg| match arg.as_str() {
"-ngl" => "--n-gpu-layers".into(),
"-dev" => "--device".into(),
_ => arg,
})
.collect()
}
fn print_available_devices() {
println!("Available devices:");
println!(" CPU");
let metal = ferrox_metal::MetalProfile::detect();
if let Some(name) = metal.device_name {
println!(" Metal: {name}");
}
let cuda = ferrox_cuda::HardwareProfile::detect();
if cuda.cuda_available {
let name = cuda.cuda_device_name.as_deref().unwrap_or("unknown device");
println!(" CUDA: {name}");
if cuda.cuda_device_count > 1 {
println!(" ({} devices detected)", cuda.cuda_device_count);
}
}
}
fn cli_bind_addr(args: &ServerArgs, env_addr: Option<&str>) -> Option<String> {
if args.host.is_none() && args.port.is_none() {
return None;
}
let existing = env_addr.and_then(|value| value.parse::<SocketAddr>().ok());
let host = args
.host
.or_else(|| existing.map(|addr| addr.ip()))
.unwrap_or(IpAddr::V4(Ipv4Addr::LOCALHOST));
let port = args
.port
.or_else(|| existing.map(|addr| addr.port()))
.unwrap_or(8383);
Some(SocketAddr::new(host, port).to_string())
}
fn apply_cli_overrides(args: &ServerArgs) -> anyhow::Result<()> {
if let Some(model) = &args.model {
unsafe { std::env::set_var("FERROX_MODEL_PATH", model) };
}
if let Some(addr) = cli_bind_addr(args, std::env::var("FERROX_ADDR").ok().as_deref()) {
unsafe { std::env::set_var("FERROX_ADDR", addr) };
}
if let Some(threads) = args.threads {
if threads == 0 {
anyhow::bail!("--threads must be greater than zero");
}
unsafe {
std::env::set_var("FERROX_CPU_THREADS", threads.to_string());
std::env::set_var("RAYON_NUM_THREADS", threads.to_string());
}
}
if args.device.is_none() && args.n_gpu_layers.is_none() {
} else {
let layers = args.n_gpu_layers.unwrap_or(GpuLayers::Auto);
let device = if layers.offload_enabled() {
args.device.unwrap_or(OffloadDevice::Auto)
} else {
OffloadDevice::None
};
match device {
OffloadDevice::None | OffloadDevice::Cpu => unsafe {
std::env::set_var("FERROX_METAL", "0");
std::env::set_var("FERROX_METAL_ATTN", "0");
std::env::set_var("FERROX_CUDA", "0");
},
OffloadDevice::Auto => unsafe {
std::env::set_var("FERROX_METAL", "auto");
std::env::set_var("FERROX_CUDA", "auto");
if std::env::var_os("FERROX_METAL_ATTN").is_none() {
std::env::set_var("FERROX_METAL_ATTN", "1");
}
},
OffloadDevice::Metal => {
#[cfg(not(feature = "metal"))]
{
anyhow::bail!(
"Metal requested but this binary was built without --features metal"
);
}
#[cfg(feature = "metal")]
{
if !ferrox_metal::MetalProfile::detect().available {
anyhow::bail!("Metal requested but no Metal device is available");
}
unsafe {
std::env::set_var("FERROX_METAL", "1");
if std::env::var_os("FERROX_METAL_ATTN").is_none() {
std::env::set_var("FERROX_METAL_ATTN", "1");
}
std::env::set_var("FERROX_CUDA", "0");
}
}
}
OffloadDevice::Cuda => {
#[cfg(not(feature = "cuda"))]
{
anyhow::bail!(
"CUDA requested but this binary was built without --features cuda"
);
}
#[cfg(feature = "cuda")]
{
if !ferrox_cuda::HardwareProfile::detect().cuda_available {
anyhow::bail!("CUDA requested but no CUDA device is available");
}
unsafe {
std::env::set_var("FERROX_CUDA", "1");
std::env::set_var("FERROX_METAL", "0");
std::env::set_var("FERROX_METAL_ATTN", "0");
}
}
}
}
}
Ok(())
}
#[allow(clippy::large_enum_variant)] pub(crate) enum Model {
Gguf(GgufModel),
Kimi(KimiModel),
Mla(MlaModel),
Gemma4(Gemma4Model),
Glm52(Glm52Model),
}
pub(crate) struct GgufModel {
decoder: Arc<Decoder>,
tokenizer: Arc<ServerTokenizer>,
stop_tokens: StopTokens,
bos_id: Option<usize>,
is_synthetic: bool,
chat_template: chat_template::PromptTemplate,
}
pub(crate) struct KimiModel {
engine: KimiEngine,
tokenizer: KimiTokenizer,
stop_tokens: StopTokens,
chat_template: chat_template::PromptTemplate,
}
pub(crate) struct MlaModel {
engine: MlaEngine,
tokenizer: ServerTokenizer,
stop_tokens: StopTokens,
bos_id: Option<usize>,
name: String,
chat_template: chat_template::PromptTemplate,
}
pub(crate) struct Gemma4Model {
engine: Gemma4Engine,
tokenizer: ServerTokenizer,
stop_tokens: StopTokens,
bos_id: Option<usize>,
name: String,
chat_template: chat_template::PromptTemplate,
}
pub(crate) struct Glm52Model {
engine: ferrox_models::Glm52Engine,
tokenizer: ServerTokenizer,
stop_tokens: StopTokens,
bos_id: Option<usize>,
name: String,
chat_template: chat_template::PromptTemplate,
}
impl Model {
pub(crate) fn chat_template(&self) -> chat_template::PromptTemplate {
match self {
Model::Gguf(m) => m.chat_template.clone(),
Model::Kimi(m) => m.chat_template.clone(),
Model::Mla(m) => m.chat_template.clone(),
Model::Gemma4(m) => m.chat_template.clone(),
Model::Glm52(m) => m.chat_template.clone(),
}
}
fn is_synthetic(&self) -> bool {
match self {
Model::Gguf(m) => m.is_synthetic,
Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => false,
}
}
fn tokenizer_kind(&self) -> &'static str {
match self {
Model::Gguf(m) => m.tokenizer.kind(),
Model::Kimi(_) => "kimi-tiktoken-bpe",
Model::Mla(m) => m.tokenizer.kind(),
Model::Gemma4(m) => m.tokenizer.kind(),
Model::Glm52(m) => m.tokenizer.kind(),
}
}
fn expert_store_stats(&self) -> Option<ferrox_core::expert_store::ExpertStoreStats> {
match self {
Model::Gguf(m) => m.decoder.expert_store_stats(),
Model::Kimi(m) => m.engine.weights.expert_store_stats(),
Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
}
}
pub(crate) fn name(&self) -> &str {
match self {
Model::Gguf(m) => m.decoder.config.name,
Model::Kimi(_) => "kimi-k3",
Model::Mla(m) => m.name.as_str(),
Model::Gemma4(m) => m.name.as_str(),
Model::Glm52(m) => m.name.as_str(),
}
}
pub(crate) fn encode(&self, text: &str) -> Vec<usize> {
match self {
Model::Gguf(m) => m.tokenizer.encode(text),
Model::Kimi(m) => m
.tokenizer
.encode(text)
.into_iter()
.map(|id| id as usize)
.collect(),
Model::Mla(m) => m.tokenizer.encode(text),
Model::Gemma4(m) => m.tokenizer.encode(text),
Model::Glm52(m) => m.tokenizer.encode(text),
}
}
pub(crate) fn bos_id(&self) -> Option<usize> {
match self {
Model::Gguf(m) => m.bos_id,
Model::Kimi(_) => None,
Model::Mla(m) => m.bos_id,
Model::Gemma4(m) => m.bos_id,
Model::Glm52(m) => m.bos_id,
}
}
pub(crate) fn decode(&self, ids: &[usize]) -> String {
match self {
Model::Gguf(m) => m.tokenizer.decode(ids),
Model::Kimi(m) => {
let ids32: Vec<u32> = ids.iter().map(|&id| id as u32).collect();
m.tokenizer.decode(&ids32)
}
Model::Mla(m) => m.tokenizer.decode(ids),
Model::Gemma4(m) => m.tokenizer.decode(ids),
Model::Glm52(m) => m.tokenizer.decode(ids),
}
}
pub(crate) fn embed_tokens(&self, tokens: &[usize]) -> Option<Vec<Vec<f32>>> {
match self {
Model::Gguf(m) => {
let mut caches: Vec<_> = (0..m.decoder.layers.len())
.map(|_| {
ferrox_core::cache::KvCache::new(
m.decoder.config.n_kv_heads,
m.decoder.config.head_dim,
)
})
.collect();
Some(m.decoder.forward_hidden_batch(tokens, 0, &mut caches))
}
Model::Kimi(_) | Model::Mla(_) | Model::Gemma4(_) | Model::Glm52(_) => None,
}
}
pub(crate) fn vocab_size(&self) -> Option<usize> {
match self {
Model::Gguf(m) => Some(m.decoder.config.vocab_size),
Model::Kimi(m) => Some(m.tokenizer.vocab_size()),
Model::Mla(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
Model::Gemma4(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
Model::Glm52(m) => Some(ferrox_models::Engine::vocab_size(&m.engine)),
}
}
}
pub(crate) struct AppState {
pub(crate) embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
active: std::sync::RwLock<Option<Arc<ActiveModel>>>,
pub(crate) load_in_progress: std::sync::atomic::AtomicBool,
pub(crate) tasks: Arc<tasks::TaskRegistry>,
pub(crate) cancels: Arc<cancel::CancelRegistry>,
pub(crate) stats: stats::Stats,
pub(crate) streams: resume::StreamRegistry,
pub(crate) model_dir: Option<PathBuf>,
response_cache: Mutex<ResponseCache>,
pub(crate) kv_pool: Option<generate::KvPoolConfig>,
pub(crate) paged_kv: Option<generate::PagedKvConfig>,
pub(crate) prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
sessions: session::SessionStore,
requests_total: std::sync::atomic::AtomicU64,
request_errors_total: std::sync::atomic::AtomicU64,
started_at: std::time::Instant,
last_request_ms: std::sync::atomic::AtomicU64,
detection: Arc<health::Detection>,
mcp: Option<mcp::LoadedMcpConfig>,
pub(crate) continuous_batching_enabled: bool,
loading_model: Mutex<Option<String>>,
last_load_error: Mutex<Option<(String, String)>>,
pub(crate) serving: Mutex<crate::stats::ServingStats>,
pub(crate) maintenance: Mutex<crate::policy::maintenance::MaintenanceGate>,
pub(crate) footprint:
Mutex<crate::policy::footprint::ProbeCache<crate::policy::footprint::Footprint>>,
pub(crate) started_unix: u64,
}
pub(crate) const FOOTPRINT_TTL_MS: u64 = 2_000;
impl AppState {
pub(crate) fn active(&self) -> Option<Arc<ActiveModel>> {
self.active
.read()
.unwrap_or_else(|p| p.into_inner())
.clone()
}
pub(crate) fn require_active(&self) -> Result<Arc<ActiveModel>, ApiError> {
self.active().ok_or_else(|| {
(
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({"error": {
"message": "no model is loaded; POST /admin/models/load with an id from \
GET /admin/models",
"type": "model_not_loaded"
}})),
)
})
}
pub(crate) fn require_model(&self) -> Result<Arc<Model>, ApiError> {
Ok(Arc::clone(self.require_active()?.generative()?))
}
pub(crate) fn swap_active(&self, next: Option<Arc<ActiveModel>>) -> Option<Arc<ActiveModel>> {
let mut guard = self.active.write().unwrap_or_else(|p| p.into_inner());
std::mem::replace(&mut *guard, next)
}
fn mark_request_finished(&self) {
let ms = self.started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
self.last_request_ms
.store(ms, std::sync::atomic::Ordering::Relaxed);
}
pub(crate) fn uptime(&self) -> Duration {
self.started_at.elapsed()
}
pub(crate) fn requests_total(&self) -> u64 {
self.requests_total
.load(std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn errors_total(&self) -> u64 {
self.request_errors_total
.load(std::sync::atomic::Ordering::Relaxed)
}
pub(crate) fn cache_stats(&self) -> response_cache::CacheStats {
lock_cache(&self.response_cache).stats()
}
pub(crate) fn last_request_age_seconds(&self) -> Option<f64> {
let last = self
.last_request_ms
.load(std::sync::atomic::Ordering::Relaxed);
(last > 0)
.then(|| self.uptime().as_secs_f64() - (last as f64 / 1000.0))
.map(|age| age.max(0.0))
}
pub(crate) fn loading_model_id(&self) -> Option<String> {
self.loading_model
.lock()
.unwrap_or_else(|p| p.into_inner())
.clone()
}
pub(crate) fn set_loading_model(&self, id: Option<String>) {
*self.loading_model.lock().unwrap_or_else(|p| p.into_inner()) = id;
}
pub(crate) fn last_load_error(&self) -> Option<(String, String)> {
self.last_load_error
.lock()
.unwrap_or_else(|p| p.into_inner())
.clone()
}
pub(crate) fn set_last_load_error(&self, error: Option<(String, String)>) {
*self
.last_load_error
.lock()
.unwrap_or_else(|p| p.into_inner()) = error;
}
pub(crate) fn active_model_name(&self) -> Option<String> {
self.active().map(|a| a.name().to_string())
}
pub(crate) fn embedding_model(&self) -> Option<Arc<ferrox_models::EmbeddingModel>> {
self.embedding
.clone()
.or_else(|| self.active().and_then(|a| a.encoder().map(Arc::clone)))
}
pub(crate) fn embedding_model_name(&self) -> Option<String> {
match self.embedding_model() {
Some(e) => Some(e.name().to_string()),
None => self.active_model_name(),
}
}
pub(crate) fn record_request(&self, record: stats::Record<'_>) {
self.stats.record(stats::entry(record));
}
}
fn lock_cache(cache: &Mutex<ResponseCache>) -> MutexGuard<'_, ResponseCache> {
cache
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
pub(crate) enum MessageContent {
Text(String),
Parts(Vec<ContentPart>),
}
#[derive(Debug, Clone, Deserialize)]
struct ContentPart {
#[serde(rename = "type")]
kind: String,
#[serde(default)]
text: Option<String>,
#[serde(default)]
image_url: Option<serde_json::Value>,
}
impl MessageContent {
fn as_text(&self) -> String {
match self {
Self::Text(s) => s.clone(),
Self::Parts(parts) => parts
.iter()
.filter_map(|p| p.text.as_deref())
.collect::<Vec<_>>()
.join(""),
}
}
fn has_image(&self) -> bool {
match self {
Self::Text(_) => false,
Self::Parts(parts) => parts
.iter()
.any(|p| p.kind == "image_url" || p.image_url.is_some()),
}
}
}
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct ChatMessage {
pub(crate) role: String,
#[serde(default)]
pub(crate) content: Option<MessageContent>,
#[serde(default)]
pub(crate) tool_calls: Option<Vec<ToolCallIn>>,
#[serde(default)]
#[allow(dead_code)]
pub(crate) tool_call_id: Option<String>,
#[serde(default, alias = "reasoning")]
pub(crate) reasoning_content: Option<String>,
}
impl ChatMessage {
fn rendered_content(&self) -> String {
let mut out = self
.content
.as_ref()
.map(MessageContent::as_text)
.unwrap_or_default();
if let Some(calls) = &self.tool_calls {
for call in calls {
out.push_str(&format!(
"<tool_call>{{\"name\": \"{}\", \"arguments\": {}}}</tool_call>",
call.function.name, call.function.arguments
));
}
}
out
}
}
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct ToolCallIn {
#[serde(default)]
#[allow(dead_code)]
id: String,
#[serde(rename = "type", default)]
#[allow(dead_code)]
kind: String,
function: ToolCallFunctionIn,
}
#[derive(Debug, Clone, Deserialize)]
struct ToolCallFunctionIn {
name: String,
arguments: String,
}
#[derive(Debug, Clone, Deserialize)]
struct ToolDef {
#[serde(rename = "type", default)]
#[allow(dead_code)]
kind: String,
function: ToolFunctionDef,
}
#[derive(Debug, Clone, Deserialize)]
struct ToolFunctionDef {
name: String,
#[serde(default)]
description: Option<String>,
#[serde(default)]
parameters: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
enum ToolChoice {
Mode(String),
Specific(serde_json::Value),
}
#[derive(Deserialize)]
#[serde(untagged)]
enum StopParam {
One(String),
Many(Vec<String>),
}
#[derive(Deserialize)]
struct ChatCompletionRequest {
model: String,
messages: Vec<ChatMessage>,
#[serde(default = "default_max_tokens")]
max_tokens: usize,
#[serde(default)]
temperature: Option<f32>,
#[serde(default)]
top_p: Option<f32>,
#[serde(default)]
min_p: Option<f32>,
#[serde(default)]
top_k: Option<usize>,
#[serde(default)]
repetition_penalty: Option<f32>,
#[serde(default)]
seed: Option<u64>,
#[serde(default)]
stop: Option<StopParam>,
#[serde(default)]
stream: Option<bool>,
#[serde(default)]
stream_resumable: Option<bool>,
#[serde(default)]
ignore_eos: Option<bool>,
#[serde(default)]
tools: Vec<ToolDef>,
#[serde(default)]
tool_choice: Option<ToolChoice>,
#[serde(default)]
chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
#[serde(default)]
reasoning_effort: Option<String>,
#[serde(default)]
thinking: Option<ThinkingSwitch>,
#[serde(default)]
session_id: Option<String>,
#[serde(default)]
logprobs: Option<bool>,
#[serde(default)]
top_logprobs: Option<u32>,
#[serde(default)]
n: Option<u32>,
#[serde(default)]
presence_penalty: Option<f32>,
#[serde(default)]
frequency_penalty: Option<f32>,
#[serde(default)]
response_format: Option<serde_json::Value>,
#[serde(default)]
logit_bias: Option<serde_json::Value>,
#[serde(default)]
grammar: Option<String>,
}
const DEFAULT_CHAT_MAX_TOKENS: usize = 32_768;
#[derive(Debug, Clone, Deserialize)]
pub(crate) struct ThinkingSwitch {
#[serde(rename = "type")]
pub(crate) kind: String,
}
const THINKING_KWARG_KEYS: [&str; 4] = [
"enable_thinking",
"thinking",
"thinking_mode",
"reasoning_effort",
];
const DISABLE_EFFORTS: [&str; 2] = ["none", "off"];
fn default_max_tokens() -> usize {
DEFAULT_CHAT_MAX_TOKENS
}
impl ChatCompletionRequest {
fn sampling_knobs(&self) -> SamplingKnobs {
SamplingKnobs {
temperature: self.temperature,
top_p: self.top_p,
min_p: self.min_p,
top_k: self.top_k,
repetition_penalty: self.repetition_penalty,
presence_penalty: self.presence_penalty,
frequency_penalty: self.frequency_penalty,
penalty_last_n: None,
}
}
fn sampling_params(&self) -> SamplingParams {
self.sampling_knobs().resolve()
}
fn stop_sequences(&self) -> Vec<String> {
self.stop
.as_ref()
.map(|s| match s {
StopParam::One(v) => vec![v.clone()],
StopParam::Many(v) => v.clone(),
})
.unwrap_or_default()
}
fn tools_active(&self) -> bool {
!self.tools.is_empty()
&& !matches!(&self.tool_choice, Some(ToolChoice::Mode(m)) if m == "none")
}
fn forced_tool_choice(&self) -> Result<Option<tool_grammar::Forced<'_>>, ApiError> {
match &self.tool_choice {
Some(ToolChoice::Mode(m)) if m == "required" => Ok(Some(tool_grammar::Forced::Any)),
Some(ToolChoice::Specific(value)) => {
let name = value
.get("function")
.and_then(|f| f.get("name"))
.or_else(|| value.get("name"))
.and_then(|n| n.as_str());
match name {
Some(name) => Ok(Some(tool_grammar::Forced::Named(name))),
None => Err(invalid_request(
"tool_choice must be \"auto\", \"none\", \"required\", or an object with \
function.name",
"tool_choice",
)),
}
}
_ => Ok(None),
}
}
fn tool_specs(&self) -> Vec<tool_grammar::ToolSpec<'_>> {
self.tools
.iter()
.map(|t| tool_grammar::ToolSpec {
name: &t.function.name,
parameters: t.function.parameters.as_ref(),
})
.collect()
}
fn thinking_direction(&self) -> Option<bool> {
if let Some(switch) = &self.thinking {
return match switch.kind.trim().to_ascii_lowercase().as_str() {
"disabled" => Some(false),
"enabled" => Some(true),
_ => None,
};
}
let effort = self.reasoning_effort.as_ref()?;
DISABLE_EFFORTS
.contains(&effort.trim().to_ascii_lowercase().as_str())
.then_some(false)
}
fn resolve_template_kwargs(
&self,
template: &chat_template::PromptTemplate,
) -> serde_json::Map<String, serde_json::Value> {
let mut kwargs = self.chat_template_kwargs.clone().unwrap_or_default();
let caller_steered = THINKING_KWARG_KEYS.iter().any(|k| kwargs.contains_key(*k));
if !caller_steered {
match self.thinking_direction() {
Some(false) => {
for (k, v) in crate::policy::effort::thinking_off_kwargs() {
kwargs.insert(k, v);
}
return kwargs;
}
Some(true) => {
for (k, v) in crate::policy::effort::thinking_on_kwargs() {
kwargs.insert(k, v);
}
}
None => {}
}
if let Some(effort) = &self.reasoning_effort {
kwargs
.entry("reasoning_effort".to_string())
.or_insert_with(|| serde_json::json!(effort));
}
}
let offered: Vec<serde_json::Value> = if self.tools_active() {
self.tools.iter().map(chat_template::tool_json).collect()
} else {
Vec::new()
};
let thinking = crate::policy::effort::resolve_thinking_mode(Some(&kwargs), Some(&offered));
if thinking == crate::policy::effort::ThinkingMode::Thinking {
for (k, v) in crate::policy::effort::thinking_on_kwargs() {
kwargs.entry(k).or_insert(v);
}
}
match crate::policy::effort::sanitize_effort(&mut kwargs, template.efforts()) {
crate::policy::effort::EffortMapping::Mapped(to) => {
tracing::debug!("reasoning_effort quantized to {}", to.as_str());
}
crate::policy::effort::EffortMapping::Dropped => {
tracing::debug!(
"reasoning_effort dropped: this checkpoint's template grades no gear close \
enough, so its own default applies"
);
}
crate::policy::effort::EffortMapping::Unchanged => {}
}
crate::policy::effort::broadcast_effort_spellings(&mut kwargs);
kwargs
}
fn validate_supported_fields(&self) -> Result<(), ApiError> {
if self.max_tokens == 0 {
return Err(invalid_request(
"max_tokens must be at least 1",
"max_tokens",
));
}
if let Some(switch) = &self.thinking {
let kind = switch.kind.trim().to_ascii_lowercase();
if kind != "enabled" && kind != "disabled" {
return Err(invalid_request(
"thinking.type must be \"enabled\" or \"disabled\"",
"thinking.type",
));
}
}
for msg in &self.messages {
if msg.content.as_ref().is_some_and(MessageContent::has_image) {
return Err(unsupported_feature(
"image_url content parts are not implemented (multimodal/VL deferred, see docs/API.md)",
));
}
}
if self.logprobs == Some(true) || self.top_logprobs.is_some() {
return Err(unsupported_feature(
"logprobs / top_logprobs are not implemented yet (see docs/API.md)",
));
}
if self.n.is_some_and(|n| n > 1) {
return Err(unsupported_feature(
"n > 1 is not implemented (single completion only)",
));
}
unsupported_sampling::refuse_logit_bias(self.logit_bias.as_ref(), "/v1/chat/completions")?;
grammar_request::for_request(self.grammar.as_deref(), self.response_format.as_ref())?;
if let Some(fmt) = &self.response_format {
match fmt.get("type").and_then(|v| v.as_str()) {
Some("json_object") => {}
Some(other) => {
return Err((
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": {
"message": format!(
"response_format type {other:?} is not supported (only json_object)"
)
}
})),
));
}
None => {
return Err((
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": {
"message": "response_format must include \"type\" (only json_object is supported)"
}
})),
));
}
}
}
if let Some(forced) = self.forced_tool_choice()? {
if self.tools.is_empty() {
return Err(invalid_request(
"tool_choice forces a tool call, but no tools were offered",
"tool_choice",
));
}
if let tool_grammar::Forced::Named(name) = forced {
if !self.tools.iter().any(|t| t.function.name == name) {
return Err(invalid_request(
&format!(
"tool_choice names {name:?}, which is not one of the tools offered"
),
"tool_choice",
));
}
}
if self.grammar.is_some() {
return Err(invalid_request(
"a forced tool_choice and a \"grammar\" are two different constraints on the \
same generation; send one",
"tool_choice",
));
}
if self.json_object_mode() {
return Err(invalid_request(
"a forced tool_choice cannot be combined with response_format json_object: \
the tool-call markers are not JSON",
"tool_choice",
));
}
}
Ok(())
}
fn effective_stop_sequences(&self) -> Vec<String> {
let mut stop = self.stop_sequences();
if self.tools_active() {
stop.push("</tool_call>".to_string());
}
stop
}
fn json_object_mode(&self) -> bool {
self.response_format
.as_ref()
.and_then(|v| v.get("type"))
.and_then(|v| v.as_str())
== Some("json_object")
}
fn generation_params(&self) -> Result<GenerationParams, ApiError> {
Ok(GenerationParams {
max_tokens: self.max_tokens,
sampling: self.sampling_params(),
seed: self.resolved_seed(),
stop: self.effective_stop_sequences(),
stop_token_ids: Vec::new(),
json_object: self.json_object_mode(),
grammar: grammar_request::for_request(
self.grammar.as_deref(),
self.response_format.as_ref(),
)?,
cancel: None,
ignore_eos: self.ignore_eos.unwrap_or(false),
})
}
fn generation_params_for_template(
&self,
template: &chat_template::PromptTemplate,
served_model: &str,
) -> Result<GenerationParams, ApiError> {
let mut params = self.generation_params()?;
if let Some(stop) = template.end_of_turn() {
if !params.stop.iter().any(|s| s == stop) {
params.stop.push(stop.to_string());
}
}
if let Some(forced) = self.forced_tool_choice()? {
params.grammar = Some(tool_grammar::build(
forced,
&self.tool_specs(),
policy::parser::ToolCallFormat::infer(served_model),
)?);
}
Ok(params)
}
fn is_cacheable(&self) -> bool {
self.temperature.unwrap_or(0.0) <= 0.0 || self.seed.is_some()
}
fn cache_key(&self, prompt: &str) -> CacheKey {
CacheKey {
model: self.model.clone(),
prompt: prompt.to_string(),
max_tokens: self.max_tokens,
sampling: response_cache::sampling_key(&self.sampling_params()),
seed: self.seed,
stop: self.effective_stop_sequences(),
}
}
fn resolved_seed(&self) -> u64 {
self.seed.unwrap_or_else(|| {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos() as u64)
.unwrap_or(0xDEFA017)
})
}
}
#[derive(Serialize)]
struct ChatCompletionChoice {
index: usize,
message: ChatCompletionResponseMessage,
finish_reason: &'static str,
}
#[derive(Serialize)]
struct ChatCompletionResponseMessage {
role: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
reasoning_content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_calls: Option<Vec<ToolCallOut>>,
}
#[derive(Serialize, Clone)]
struct ToolCallOut {
id: String,
#[serde(rename = "type")]
kind: &'static str,
function: ToolCallFunctionOut,
}
#[derive(Serialize, Clone)]
struct ToolCallDelta {
index: usize,
#[serde(skip_serializing_if = "Option::is_none")]
id: Option<String>,
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
kind: Option<&'static str>,
function: ToolCallFunctionDelta,
}
#[derive(Serialize, Clone, Default)]
struct ToolCallFunctionDelta {
#[serde(skip_serializing_if = "Option::is_none")]
name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
arguments: Option<String>,
}
impl ToolCallDelta {
fn whole(index: usize, name: String, arguments: String) -> Self {
ToolCallDelta {
index,
id: Some(format!("call_{index}")),
kind: Some("function"),
function: ToolCallFunctionDelta {
name: Some(name),
arguments: Some(arguments),
},
}
}
fn opening(index: usize, name: String) -> Self {
ToolCallDelta {
index,
id: Some(format!("call_{index}")),
kind: Some("function"),
function: ToolCallFunctionDelta {
name: Some(name),
arguments: Some(String::new()),
},
}
}
fn arguments(index: usize, fragment: String) -> Self {
ToolCallDelta {
index,
id: None,
kind: None,
function: ToolCallFunctionDelta {
name: None,
arguments: Some(fragment),
},
}
}
}
#[derive(Serialize, Clone)]
struct ToolCallFunctionOut {
name: String,
arguments: String,
}
#[derive(Serialize)]
struct ChatCompletionResponse {
id: String,
request_id: String,
object: &'static str,
model: String,
choices: Vec<ChatCompletionChoice>,
usage: generate::Usage,
ferrox_cache: &'static str,
}
#[derive(Serialize)]
struct ChatCompletionChunkDelta {
#[serde(skip_serializing_if = "Option::is_none")]
role: Option<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
reasoning_content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_calls: Option<Vec<ToolCallDelta>>,
}
#[derive(Serialize)]
struct ChatCompletionChunkChoice {
index: usize,
delta: ChatCompletionChunkDelta,
finish_reason: Option<&'static str>,
}
#[derive(Serialize)]
struct ChatCompletionChunk {
id: String,
#[serde(skip_serializing_if = "Option::is_none")]
request_id: Option<String>,
object: &'static str,
model: String,
choices: Vec<ChatCompletionChunkChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
usage: Option<generate::Usage>,
}
async fn health(State(state): State<Arc<AppState>>) -> Response {
let snapshot = state.detection.snapshot();
let mut capabilities = snapshot.capabilities;
let active = state.active();
capabilities.push(match active.as_deref() {
None => ferrox_api::Capability::unavailable(
ferrox_api::health::capability::REAL_WEIGHTS,
ferrox_api::health::reason::MODEL_NOT_LOADED,
"No model is loaded. POST /admin/models/load with an id from GET /admin/models.",
),
Some(active) if active.is_synthetic() => ferrox_api::Capability::unavailable(
ferrox_api::health::capability::REAL_WEIGHTS,
ferrox_api::health::reason::MODEL_NOT_LOADED,
"Serving synthetic random weights: set FERROX_MODEL_PATH (or -m) to a real \
checkpoint. Output from this model is noise.",
),
Some(active) if active.encoder().is_some() => ferrox_api::Capability::available(
ferrox_api::health::capability::REAL_WEIGHTS,
format!(
"Serving the real embedding checkpoint '{}'. This is an ENCODER: {} serves \
it, generation endpoints refuse it.",
active.name(),
ferrox_api::routes::V1_EMBEDDINGS,
),
),
Some(active) => ferrox_api::Capability::available(
ferrox_api::health::capability::REAL_WEIGHTS,
format!("Serving the real checkpoint '{}'.", active.name()),
),
});
capabilities.push(if active.as_ref().is_some_and(|a| a.batcher.is_some()) {
ferrox_api::Capability::available(
ferrox_api::health::capability::CONTINUOUS_BATCHING,
"Concurrent requests share one batched decode step.",
)
} else {
ferrox_api::Capability::unavailable(
ferrox_api::health::capability::CONTINUOUS_BATCHING,
ferrox_api::health::reason::DISABLED,
"Off; set FERROX_CONTINUOUS_BATCHING=1 (incompatible with a KV pool or prefix cache).",
)
});
let last_request_ms = state
.last_request_ms
.load(std::sync::atomic::Ordering::Relaxed);
let uptime = state.started_at.elapsed();
let health_state = if active.is_none() {
ferrox_api::HealthState::Unavailable
} else {
snapshot.state
};
let body = ferrox_api::HealthResponse {
state: health_state,
reason: match health_state {
ferrox_api::HealthState::Ready => None,
ferrox_api::HealthState::Unavailable => {
Some(ferrox_api::health::reason::MODEL_NOT_LOADED.to_string())
}
ferrox_api::HealthState::Detecting => {
Some(ferrox_api::health::reason::DETECTING.to_string())
}
},
detail: match health_state {
ferrox_api::HealthState::Ready => None,
ferrox_api::HealthState::Unavailable => Some(
"No model is loaded. POST /admin/models/load with an id from GET /admin/models."
.to_string(),
),
ferrox_api::HealthState::Detecting => {
Some("Probing available compute backends.".to_string())
}
},
model: active
.as_deref()
.map(|active| ferrox_api::health::ModelSummary {
id: active.name().to_string(),
tokenizer: active.tokenizer_kind().to_string(),
synthetic_weights: active.is_synthetic(),
}),
capabilities,
version: env!("CARGO_PKG_VERSION").to_string(),
pid: std::process::id(),
uptime_seconds: uptime.as_secs_f64(),
server_time_unix_ms: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis().min(u64::MAX as u128) as u64)
.unwrap_or(0),
last_request_age_seconds: (last_request_ms > 0)
.then(|| uptime.as_secs_f64() - (last_request_ms as f64 / 1000.0))
.map(|age| age.max(0.0)),
};
let status =
StatusCode::from_u16(body.state.http_status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
(status, Json(body)).into_response()
}
async fn list_models(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
let Some(active) = state.active() else {
return Json(serde_json::json!({ "object": "list", "data": [] }));
};
let mut model_entry = serde_json::json!({
"id": active.name(),
"object": "model",
"ferrox_synthetic_weights": active.is_synthetic(),
"ferrox_tokenizer": active.tokenizer_kind(),
});
if let Some(encoder) = active.encoder() {
model_entry["ferrox_model_kind"] = serde_json::json!("embedding");
model_entry["ferrox_endpoints"] = serde_json::json!([ferrox_api::routes::V1_EMBEDDINGS]);
model_entry["ferrox_n_embd"] = serde_json::json!(encoder.n_embd());
model_entry["ferrox_pooling"] = serde_json::json!(encoder.pooling_type().name());
model_entry["ferrox_context_length"] = serde_json::json!(encoder.n_ctx_train());
}
if let Some(model) = active.generative_opt() {
let parser_configured =
crate::policy::parser::ReasoningFormat::infer(active.name()).is_some();
let gears = model.chat_template().think_gears(parser_configured);
if !gears.is_empty() {
model_entry["supported_reasoning_efforts"] = serde_json::json!(gears.supported);
if let Some(default) = &gears.default {
model_entry["default_reasoning_effort"] = serde_json::json!(default);
}
model_entry["reasoning_effort_kwargs"] = serde_json::json!(gears.kwargs);
}
}
if let Some(mcp) = &state.mcp {
model_entry["ferrox_mcp"] = mcp.models_metadata();
}
Json(serde_json::json!({
"object": "list",
"data": [model_entry]
}))
}
async fn serving_stats(State(state): State<Arc<AppState>>) -> Json<serde_json::Value> {
let now_ms = state.uptime().as_millis().min(u64::MAX as u128) as u64;
let mut serving = state.serving.lock().unwrap_or_else(|p| p.into_inner());
let active = state.active();
Json(serde_json::json!({
"model": active.as_ref().map(|a| a.name()),
"state": state
.maintenance
.lock()
.unwrap_or_else(|p| p.into_inner())
.state()
.as_str(),
"uptime_s": state.uptime().as_secs(),
"throughput": {
"decode_tps": (serving.decode_tokens_per_second(now_ms) * 10.0).round() / 10.0,
"prefill_tps": (serving.prefill_tokens_per_second(now_ms) * 10.0).round() / 10.0,
},
"requests": {
"active": state.cancels.live_count(),
"completed": state.stats.recorded_total(),
"p95_ms": state.stats.p95_duration_ms(),
"ttft_mean_ms": state.stats.ttft_mean_ms(),
"prompt_tokens_total": state.stats.tokens_prompt_total(),
"completion_tokens_total": state.stats.tokens_generated_total(),
},
"pools": cache_admin::pool_gauges(&state),
"memory": cache_admin::footprint_json(&state),
}))
}
#[derive(Deserialize)]
struct RequestsQuery {
#[serde(default)]
since: u64,
#[serde(default = "default_requests_limit")]
limit: usize,
}
fn default_requests_limit() -> usize {
stats::MAX_PAGE
}
async fn recent_requests(
State(state): State<Arc<AppState>>,
axum::extract::Query(q): axum::extract::Query<RequestsQuery>,
) -> Json<serde_json::Value> {
let (rows, cursor, missed) = state.stats.page(q.since, q.limit);
Json(serde_json::json!({
"requests": rows,
"next_cursor": cursor,
"missed": missed,
"total": state.stats.recorded_total(),
}))
}
#[derive(Serialize)]
struct CombinedCacheStats {
response_cache: response_cache::CacheStats,
prefix_cache: Option<ferrox_models::PrefixCacheStats>,
}
async fn cache_stats(State(state): State<Arc<AppState>>) -> Json<CombinedCacheStats> {
Json(CombinedCacheStats {
response_cache: lock_cache(&state.response_cache).stats(),
prefix_cache: state
.prefix_cache
.as_ref()
.map(|pc| pc.lock().unwrap_or_else(|p| p.into_inner()).stats()),
})
}
async fn metrics(State(state): State<Arc<AppState>>) -> Response {
use std::sync::atomic::Ordering;
let cache_stats = lock_cache(&state.response_cache).stats();
let active = state.active();
let requests_total = state.requests_total.load(Ordering::Relaxed);
let errors_total = state.request_errors_total.load(Ordering::Relaxed);
let uptime = state.started_at.elapsed().as_secs_f64();
let body = format!(
"# HELP ferrox_requests_total Total chat completion requests received.\n\
# TYPE ferrox_requests_total counter\n\
ferrox_requests_total {requests_total}\n\
# HELP ferrox_request_errors_total Total chat completion requests that returned an error.\n\
# TYPE ferrox_request_errors_total counter\n\
ferrox_request_errors_total {errors_total}\n\
# HELP ferrox_cache_hits_total Whole-response cache hits.\n\
# TYPE ferrox_cache_hits_total counter\n\
ferrox_cache_hits_total {}\n\
# HELP ferrox_cache_misses_total Whole-response cache misses.\n\
# TYPE ferrox_cache_misses_total counter\n\
ferrox_cache_misses_total {}\n\
# HELP ferrox_cache_entries Current whole-response cache entry count.\n\
# TYPE ferrox_cache_entries gauge\n\
ferrox_cache_entries {}\n\
# HELP ferrox_synthetic_weights 1 if serving synthetic random weights instead of a real checkpoint.\n\
# TYPE ferrox_synthetic_weights gauge\n\
ferrox_synthetic_weights {}\n\
# HELP ferrox_uptime_seconds Seconds since this server process started.\n\
# TYPE ferrox_uptime_seconds gauge\n\
ferrox_uptime_seconds {uptime}\n",
cache_stats.hits,
cache_stats.misses,
cache_stats.entries,
active
.as_ref()
.map(|a| a.is_synthetic() as u8)
.unwrap_or(0),
);
let body = match active
.as_ref()
.and_then(|a| a.expert_store_stats())
{
Some(es) => format!(
"{body}\
# HELP ferrox_expert_cache_hits_total Expert-store cache hits.\n\
# TYPE ferrox_expert_cache_hits_total counter\n\
ferrox_expert_cache_hits_total {}\n\
# HELP ferrox_expert_cache_misses_total Expert-store cache misses (source reads).\n\
# TYPE ferrox_expert_cache_misses_total counter\n\
ferrox_expert_cache_misses_total {}\n\
# HELP ferrox_expert_cache_evictions_total Expert-store LRU evictions.\n\
# TYPE ferrox_expert_cache_evictions_total counter\n\
ferrox_expert_cache_evictions_total {}\n\
# HELP ferrox_expert_cache_pass_throughs_total Acquires served uncached (entry could not fit the budget).\n\
# TYPE ferrox_expert_cache_pass_throughs_total counter\n\
ferrox_expert_cache_pass_throughs_total {}\n\
# HELP ferrox_expert_cache_bytes_read_total Bytes read from the checkpoint for expert misses.\n\
# TYPE ferrox_expert_cache_bytes_read_total counter\n\
ferrox_expert_cache_bytes_read_total {}\n\
# HELP ferrox_expert_cache_resident_bytes Current expert-cache footprint in bytes.\n\
# TYPE ferrox_expert_cache_resident_bytes gauge\n\
ferrox_expert_cache_resident_bytes {}\n",
es.hits, es.misses, es.evictions, es.pass_throughs, es.bytes_read, es.resident_bytes,
),
None => body,
};
let body = match active.as_ref().and_then(|a| a.batcher.as_ref()) {
Some(batcher) => {
let sched = batcher.stats();
format!(
"{body}\
# HELP ferrox_prefill_chunks_total Bounded prefill chunks the batch scheduler has run.\n\
# TYPE ferrox_prefill_chunks_total counter\n\
ferrox_prefill_chunks_total {}\n\
# HELP ferrox_prefill_tokens_total Prompt tokens run through chunked prefill.\n\
# TYPE ferrox_prefill_tokens_total counter\n\
ferrox_prefill_tokens_total {}\n\
# HELP ferrox_decode_steps_total Batched decode steps the batch scheduler has run.\n\
# TYPE ferrox_decode_steps_total counter\n\
ferrox_decode_steps_total {}\n\
# HELP ferrox_scheduler_queue_depth Requests waiting for admission to the batch scheduler.\n\
# TYPE ferrox_scheduler_queue_depth gauge\n\
ferrox_scheduler_queue_depth {}\n\
# HELP ferrox_scheduler_queue_rejected_total Requests refused with 503 because the admission queue was full.\n\
# TYPE ferrox_scheduler_queue_rejected_total counter\n\
ferrox_scheduler_queue_rejected_total {}\n\
# HELP ferrox_kv_blocks_total KV blocks in the scheduler's admission budget (0 when unconfigured).\n\
# TYPE ferrox_kv_blocks_total gauge\n\
ferrox_kv_blocks_total {}\n\
# HELP ferrox_kv_blocks_free KV blocks not reserved by an in-flight request.\n\
# TYPE ferrox_kv_blocks_free gauge\n\
ferrox_kv_blocks_free {}\n\
# HELP ferrox_kv_block_size Token positions per KV block.\n\
# TYPE ferrox_kv_block_size gauge\n\
ferrox_kv_block_size {}\n\
# HELP ferrox_kv_rejected_too_large_total Requests refused with 400 because they exceed the whole KV block budget.\n\
# TYPE ferrox_kv_rejected_too_large_total counter\n\
ferrox_kv_rejected_too_large_total {}\n\
# HELP ferrox_kv_rejected_context_length_total Requests refused with 400 for exceeding the per-request context ceiling.\n\
# TYPE ferrox_kv_rejected_context_length_total counter\n\
ferrox_kv_rejected_context_length_total {}\n\
# HELP ferrox_scheduler_aborted_total Requests the batch scheduler stopped because they were cancelled.\n\
# TYPE ferrox_scheduler_aborted_total counter\n\
ferrox_scheduler_aborted_total {}\n",
sched.prefill_chunks,
sched.prefill_tokens,
sched.decode_steps,
sched.queue_depth,
sched.queue_rejected,
sched.kv_blocks_total,
sched.kv_blocks_free,
sched.kv_block_size,
sched.kv_rejected_too_large,
sched.kv_rejected_context_length,
sched.aborted,
)
}
None => body,
};
(
[(
axum::http::header::CONTENT_TYPE,
"text/plain; version=0.0.4",
)],
body,
)
.into_response()
}
pub(crate) type ApiError = (StatusCode, Json<serde_json::Value>);
pub(crate) fn invalid_request(message: &str, param: &str) -> ApiError {
(
StatusCode::BAD_REQUEST,
Json(serde_json::json!({"error": {
"message": message,
"type": "invalid_request_error",
"param": param,
"code": null,
}})),
)
}
pub(crate) fn unsupported_feature(message: &str) -> ApiError {
(
StatusCode::NOT_IMPLEMENTED,
Json(serde_json::json!({"error": {"message": message, "type": "unsupported"}})),
)
}
pub(crate) fn decode_error_response(e: generate::DecodeError) -> ApiError {
let status = match e {
generate::DecodeError::TokenOutOfVocab { .. } => StatusCode::BAD_REQUEST,
generate::DecodeError::KvBudgetExceeded { .. } => StatusCode::BAD_REQUEST,
generate::DecodeError::KvPoolExhausted | generate::DecodeError::QueueFull { .. } => {
StatusCode::SERVICE_UNAVAILABLE
}
generate::DecodeError::GrammarConstraint { .. } => StatusCode::BAD_REQUEST,
};
tracing::warn!("decode error: {e}");
let mut body = serde_json::json!({"error": {"message": e.to_string()}});
if let generate::DecodeError::KvBudgetExceeded {
binding,
estimated_bytes,
limit_bytes,
positions,
positions_limit,
..
} = &e
{
body["error"]["type"] = serde_json::json!("invalid_request_error");
body["error"]["code"] = serde_json::json!(binding);
body["error"]["binding"] = serde_json::json!(binding);
body["error"]["estimated_bytes"] = serde_json::json!(estimated_bytes);
body["error"]["limit_bytes"] = serde_json::json!(limit_bytes);
body["error"]["positions"] = serde_json::json!(positions);
body["error"]["positions_limit"] = serde_json::json!(positions_limit);
}
if let Some(secs) = e.retry_after_secs() {
body["error"]["retry_after_seconds"] = serde_json::json!(secs);
}
(status, Json(body))
}
pub(crate) fn join_error_response(e: tokio::task::JoinError) -> ApiError {
tracing::error!("generation task panicked: {e}");
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({"error": {"message": "internal error during generation"}})),
)
}
#[allow(clippy::too_many_arguments)] fn run_generation_emit(
model: &Model,
prompt: &str,
params: &GenerationParams,
kv_pool: Option<&generate::KvPoolConfig>,
paged_kv: Option<&generate::PagedKvConfig>,
prefix_cache: Option<&Mutex<PrefixCache>>,
continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
ceiling: Option<&budget::ContextCeiling>,
mut emit: impl FnMut(&str),
) -> Result<(FinishReason, generate::Usage, String), generate::DecodeError> {
let synthetic = model.is_synthetic();
let mut chunks = Vec::new();
let params = &{
let mut resolved = params.clone();
resolved.stop_token_ids =
crate::stop::resolve_stop_tokens(&resolved.stop, |text| model.encode(text));
resolved
};
let used_batcher = matches!((model, continuous_batcher), (Model::Gguf(_), Some(_)));
let (finish, usage) = match model {
Model::Gguf(m) => {
if let Some(batcher) = continuous_batcher {
let mut tokens = m.tokenizer.encode(prompt);
ferrox_models::tokenizer::prepend_bos(&mut tokens, m.bos_id);
let (finish, _generated_ids, text, usage) =
batcher.generate(tokens, params.clone(), m.stop_tokens.clone())?;
if !text.is_empty() {
chunks.push(text);
}
(finish, usage)
} else {
generate::generate(
&m.decoder,
m.tokenizer.as_ref(),
&m.stop_tokens,
m.bos_id,
prompt,
params,
kv_pool,
paged_kv,
prefix_cache,
ceiling,
|chunk| {
chunks.push(chunk.to_string());
if !synthetic {
emit(chunk);
}
},
)?
}
}
Model::Kimi(m) => generate::generate_engine(
&m.engine,
&m.tokenizer,
&m.stop_tokens,
None,
prompt,
params,
|chunk| {
chunks.push(chunk.to_string());
if !synthetic {
emit(chunk);
}
},
)?,
Model::Mla(m) => generate::generate_engine(
&m.engine,
&m.tokenizer,
&m.stop_tokens,
m.bos_id,
prompt,
params,
|chunk| {
chunks.push(chunk.to_string());
if !synthetic {
emit(chunk);
}
},
)?,
Model::Gemma4(m) => generate::generate_engine(
&m.engine,
&m.tokenizer,
&m.stop_tokens,
m.bos_id,
prompt,
params,
|chunk| {
chunks.push(chunk.to_string());
if !synthetic {
emit(chunk);
}
},
)?,
Model::Glm52(m) => generate::generate_engine(
&m.engine,
&m.tokenizer,
&m.stop_tokens,
m.bos_id,
prompt,
params,
|chunk| {
chunks.push(chunk.to_string());
if !synthetic {
emit(chunk);
}
},
)?,
};
let mut full = chunks.concat();
if synthetic {
full = format!(
"[ferrox synthetic-weight demo: no real checkpoint loaded -- set FERROX_MODEL_PATH \
to serve a real model. Decoded ids -> {full:?}]"
);
emit(&full);
} else if used_batcher && !full.is_empty() {
emit(&full);
}
Ok((finish, usage, full))
}
#[allow(clippy::too_many_arguments)] pub(crate) fn run_generation(
model: &Model,
prompt: &str,
params: &GenerationParams,
kv_pool: Option<&generate::KvPoolConfig>,
paged_kv: Option<&generate::PagedKvConfig>,
prefix_cache: Option<&Mutex<PrefixCache>>,
continuous_batcher: Option<&serving::batch::ContinuousBatcher>,
ceiling: Option<&budget::ContextCeiling>,
) -> Result<(Vec<String>, FinishReason, generate::Usage), generate::DecodeError> {
let (finish, usage, full) = run_generation_emit(
model,
prompt,
params,
kv_pool,
paged_kv,
prefix_cache,
continuous_batcher,
ceiling,
|_| {},
)?;
Ok((
if full.is_empty() {
Vec::new()
} else {
vec![full]
},
finish,
usage,
))
}
pub(crate) fn prompt_from_messages(
messages: &[ChatMessage],
template: &chat_template::PromptTemplate,
tools: &[ToolDef],
extra: serde_json::Map<String, serde_json::Value>,
) -> Result<String, ApiError> {
let rendered = if tools.is_empty() || template.handles_tools() {
template.render(messages, tools, extra)
} else {
let mut with_preamble = Vec::with_capacity(messages.len() + 1);
with_preamble.push(ChatMessage {
role: "system".to_string(),
content: Some(MessageContent::Text(tool_preamble(tools))),
tool_calls: None,
tool_call_id: None,
reasoning_content: None,
});
with_preamble.extend_from_slice(messages);
template.render(&with_preamble, &[], extra)
};
rendered.map_err(template_error_response)
}
fn template_error_response(err: ferrox_models::chat_template::TemplateError) -> ApiError {
(
StatusCode::BAD_REQUEST,
Json(serde_json::json!({
"error": {
"message": format!("chat template failed to render: {err}"),
"type": "invalid_request_error",
"param": "messages",
"code": null,
}
})),
)
}
fn tool_preamble(tools: &[ToolDef]) -> String {
let mut out = String::from(
"You can call tools to help answer the user. To call a tool, respond with \
EXACTLY one line in this format and nothing else:\n\
<tool_call>{\"name\": \"<tool name>\", \"arguments\": {<arguments as a JSON \
object matching that tool's parameters>}}</tool_call>\n\n\
Available tools:\n",
);
for t in tools {
out.push_str(&format!(
"- {}: {}\n parameters (JSON schema): {}\n",
t.function.name,
t.function.description.as_deref().unwrap_or(""),
t.function
.parameters
.as_ref()
.map(|v| v.to_string())
.unwrap_or_else(|| "{}".to_string()),
));
}
out
}
fn tool_call_deltas(
events: Vec<crate::policy::parser::ToolCallEvent>,
opened: &std::cell::Cell<usize>,
) -> (String, Vec<ToolCallDelta>) {
let mut text = String::new();
let mut deltas = Vec::new();
for event in events {
match event {
crate::policy::parser::ToolCallEvent::Text(chunk) => text.push_str(&chunk),
crate::policy::parser::ToolCallEvent::CallStart { index, name } => {
opened.set(opened.get().max(index + 1));
deltas.push(ToolCallDelta::opening(index, name));
}
crate::policy::parser::ToolCallEvent::CallArguments { index, fragment } => {
if !fragment.is_empty() {
deltas.push(ToolCallDelta::arguments(index, fragment));
}
}
crate::policy::parser::ToolCallEvent::CallEnd { .. } => {}
}
}
(text, deltas)
}
fn build_response_message(
text: String,
tools: &[ToolDef],
posture: output::OutputPosture,
base_finish: &'static str,
) -> (ChatCompletionResponseMessage, &'static str) {
let parsed = output::parse_output(&text, tools, posture);
let calls: Vec<ToolCallOut> = parsed
.calls
.into_iter()
.enumerate()
.map(|(index, call)| ToolCallOut {
id: format!("call_{index}"),
kind: "function",
function: ToolCallFunctionOut {
name: call.name,
arguments: call.arguments,
},
})
.collect();
if !calls.is_empty() {
return (
ChatCompletionResponseMessage {
role: "assistant",
content: None,
reasoning_content: parsed.reasoning,
tool_calls: Some(calls),
},
"tool_calls",
);
}
(
ChatCompletionResponseMessage {
role: "assistant",
content: Some(parsed.content),
reasoning_content: parsed.reasoning,
tool_calls: None,
},
base_finish,
)
}
fn resolve_history(state: &AppState, req: &ChatCompletionRequest) -> Vec<ChatMessage> {
let mut history = match &req.session_id {
Some(id) => state.sessions.extend_and_get(id, &req.messages),
None => req.messages.clone(),
};
if req.json_object_mode() {
inject_json_object_system_hint(&mut history);
}
history
}
fn inject_json_object_system_hint(messages: &mut Vec<ChatMessage>) {
const HINT: &str =
"You must respond with valid JSON only (a single JSON object, no markdown fences).";
if let Some(sys) = messages.iter_mut().find(|m| m.role == "system") {
match &mut sys.content {
Some(MessageContent::Text(s)) if !s.contains("JSON") => {
s.push_str("\n\n");
s.push_str(HINT);
}
None => {
sys.content = Some(MessageContent::Text(HINT.to_string()));
}
_ => {}
}
} else {
messages.insert(
0,
ChatMessage {
role: "system".to_string(),
content: Some(MessageContent::Text(HINT.to_string())),
tool_calls: None,
tool_call_id: None,
reasoning_content: None,
},
);
}
}
async fn chat_completions(
State(state): State<Arc<AppState>>,
headers: axum::http::HeaderMap,
Json(req): Json<ChatCompletionRequest>,
) -> Response {
let attribution = attribution::Attribution::from_headers(&headers);
state
.requests_total
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let started = std::time::Instant::now();
let request_id = ferrox_api::next_request_id();
let stream = req.stream.unwrap_or(false);
let refusal = cache_admin::check_admission(&state)
.err()
.or_else(|| req.validate_supported_fields().err());
if let Some(err) = refusal {
state
.request_errors_total
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let response = err.into_response();
state.record_request(stats::Record {
request_id: &request_id,
route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
model: state.active_model_name(),
status: response.status().as_u16(),
stream,
duration_ms: started.elapsed().as_millis() as u64,
usage: None,
attribution: &attribution,
});
return response;
}
let response = if stream {
chat_completions_stream(
Arc::clone(&state),
req,
request_id.clone(),
started,
attribution.clone(),
)
.await
.into_response()
} else {
chat_completions_full(
Arc::clone(&state),
req,
request_id.clone(),
started,
attribution.clone(),
)
.await
.into_response()
};
if response.status().is_client_error() || response.status().is_server_error() {
state
.request_errors_total
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
state.record_request(stats::Record {
request_id: &request_id,
route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
model: state.active_model_name(),
status: response.status().as_u16(),
stream,
duration_ms: started.elapsed().as_millis() as u64,
usage: None,
attribution: &attribution,
});
}
state.mark_request_finished();
response
}
async fn chat_completions_full(
state: Arc<AppState>,
req: ChatCompletionRequest,
request_id: String,
started: std::time::Instant,
attribution: attribution::Attribution,
) -> Result<Json<ChatCompletionResponse>, ApiError> {
let tools_active = req.tools_active();
let active = state.require_active()?;
let history = resolve_history(&state, &req);
let template = active.generative()?.chat_template();
let kwargs = req.resolve_template_kwargs(&template);
let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
let key = req.is_cacheable().then(|| req.cache_key(&prompt));
let (completion, cache_status) = if let Some(cached) = key
.as_ref()
.and_then(|key| lock_cache(&state.response_cache).get(key))
{
tracing::debug!("cache hit for key {}", key.as_ref().unwrap().digest());
(cached, "hit")
} else {
let params = req.generation_params_for_template(&template, active.name())?;
let (chunks, finish, usage) = decode_task::buffered(
decode_task::DecodeHandles::take(&state, &active)?,
prompt.clone(),
params,
)
.await?;
let completion = response_cache::CachedCompletion {
content: chunks.concat(),
finish,
usage,
};
let cache_status = if let Some(key) = key {
tracing::debug!("cache miss for key {}", key.digest());
lock_cache(&state.response_cache).put(key, completion.clone());
"miss"
} else {
"skip"
};
(completion, cache_status)
};
let content = completion.content;
if req.json_object_mode() {
json_mode::validate_json_object_output(&content)?;
}
if let Some(id) = &req.session_id {
state.sessions.store_reply(
id,
ChatMessage {
role: "assistant".to_string(),
content: Some(MessageContent::Text(content.clone())),
tool_calls: None,
tool_call_id: None,
reasoning_content: None,
},
);
}
let (message, finish_reason) = build_response_message(
content,
if tools_active { &req.tools } else { &[] },
output::OutputPosture::resolve(active.name(), &prompt),
completion.finish.as_str(),
);
state.record_request(stats::Record {
request_id: &request_id,
route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
model: Some(active.name().to_string()),
status: 200,
stream: false,
duration_ms: started.elapsed().as_millis() as u64,
usage: Some(&completion.usage),
attribution: &attribution,
});
Ok(Json(ChatCompletionResponse {
id: request_id.clone(),
request_id,
object: "chat.completion",
model: req.model,
choices: vec![ChatCompletionChoice {
index: 0,
message,
finish_reason,
}],
usage: completion.usage,
ferrox_cache: cache_status,
}))
}
async fn chat_completions_stream(
state: Arc<AppState>,
req: ChatCompletionRequest,
request_id: String,
started: std::time::Instant,
attribution: attribution::Attribution,
) -> Result<Response, ApiError> {
let tools_active = req.tools_active();
let active = state.require_active()?;
let history = resolve_history(&state, &req);
let template = active.generative()?.chat_template();
let kwargs = req.resolve_template_kwargs(&template);
let prompt = prompt_from_messages(&history, &template, &req.tools, kwargs)?;
let model_name = req.model.clone();
let session_id = req.session_id.clone();
let sessions = state.sessions.clone();
let model = Arc::clone(active.generative()?);
let kv_pool = state.kv_pool.clone();
let paged_kv = state.paged_kv.clone();
let prefix_cache = state.prefix_cache.clone();
let batcher = active.batcher.clone();
let ceiling = active.ceiling.clone();
let mut params = req.generation_params_for_template(&template, active.name())?;
let stats_state = Arc::clone(&state);
let served_model = active.name().to_string();
let posture = output::OutputPosture::resolve(&served_model, &prompt);
let offered_tools: Vec<ToolDef> = if tools_active {
req.tools.clone()
} else {
Vec::new()
};
let (cancel_token, cancel_guard) = state.cancels.register(&request_id);
params.cancel = Some(cancel_token.clone());
let overlap = batcher.is_none();
let slot = req
.stream_resumable
.unwrap_or(false)
.then(|| state.streams.register(&request_id));
let emitter = resume::Emitter::new(slot);
let (tx, rx) = tokio::sync::mpsc::channel::<Result<Event, Infallible>>(64);
let keepalive = sse::keepalive_event(&ChatCompletionChunk {
id: request_id.clone(),
request_id: None,
object: "chat.completion.chunk",
model: model_name.clone(),
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: ChatCompletionChunkDelta {
role: None,
content: None,
reasoning_content: None,
tool_calls: None,
},
finish_reason: None,
}],
usage: None,
});
tokio::task::spawn_blocking(move || {
let _cancel_guard = cancel_guard;
let tx_chunks = tx.clone();
let orphan_timeout = sse::orphan_timeout_from_env();
let mut first = true;
let head_request_id = request_id.clone();
let stream_reasoning: Rc<RefCell<Option<crate::policy::parser::ReasoningParser>>> =
Rc::new(RefCell::new(posture.reasoning_parser()));
let emit_reasoning = Rc::clone(&stream_reasoning);
let stream_tools: Rc<RefCell<Option<crate::policy::parser::ToolCallParser>>> = Rc::new(
RefCell::new(tools_active.then(|| posture.tool_call_parser(&offered_tools))),
);
let emit_tools = Rc::clone(&stream_tools);
let streamed_calls = Rc::new(std::cell::Cell::new(0usize));
let emit_streamed_calls = Rc::clone(&streamed_calls);
let result = run_generation_emit(
&model,
&prompt,
¶ms,
kv_pool.as_ref(),
paged_kv.as_ref(),
prefix_cache.as_deref(),
batcher.as_ref(),
ceiling.as_deref(),
|chunk| {
if !overlap || chunk.is_empty() {
return;
}
let (reasoning, content) = match emit_reasoning.borrow_mut().as_mut() {
Some(parser) => {
let delta = parser.push(chunk);
(delta.reasoning, delta.content)
}
None => (String::new(), chunk.to_string()),
};
let (content, tool_calls) = match emit_tools.borrow_mut().as_mut() {
Some(parser) => {
let (text, calls) =
tool_call_deltas(parser.push(&content), &emit_streamed_calls);
(text, calls)
}
None => (content, Vec::new()),
};
if reasoning.is_empty() && content.is_empty() && tool_calls.is_empty() {
return;
}
let role = if first { Some("assistant") } else { None };
let request_id = first.then(|| head_request_id.clone());
first = false;
let payload = ChatCompletionChunk {
id: head_request_id.clone(),
request_id,
object: "chat.completion.chunk",
model: model_name.clone(),
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: ChatCompletionChunkDelta {
role,
content: (!content.is_empty()).then_some(content),
reasoning_content: (!reasoning.is_empty()).then_some(reasoning),
tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
},
finish_reason: None,
}],
usage: None,
};
if let Err(why) =
sse::send_or_orphan(&tx_chunks, Ok(emitter.event(&payload)), orphan_timeout)
{
if why == sse::SendFailure::Orphaned {
tracing::warn!(
"SSE stream {head_request_id} accepted nothing for the orphan \
deadline; treating it as abandoned"
);
}
if !emitter.is_resumable() {
cancel_token.cancel();
}
}
},
);
let mut pending_request_id = first.then(|| request_id.clone());
match result {
Ok((finish, usage, full_text)) => {
if let Some(id) = &session_id {
sessions.store_reply(
id,
ChatMessage {
role: "assistant".to_string(),
content: Some(MessageContent::Text(full_text.clone())),
tool_calls: None,
tool_call_id: None,
reasoning_content: None,
},
);
}
let mut streamed_finish: Option<&'static str> = None;
if overlap {
let tail = stream_reasoning
.borrow_mut()
.as_mut()
.map(|parser| parser.flush())
.unwrap_or_default();
let (mut content, mut tool_calls) = (tail.content, Vec::new());
if let Some(parser) = stream_tools.borrow_mut().as_mut() {
let mut events = parser.push(&content);
events.extend(parser.finish());
let (text, calls) = tool_call_deltas(events, &streamed_calls);
content = text;
tool_calls = calls;
}
if !content.is_empty() || !tail.reasoning.is_empty() || !tool_calls.is_empty() {
let payload = ChatCompletionChunk {
id: request_id.clone(),
request_id: pending_request_id.take(),
object: "chat.completion.chunk",
model: model_name.clone(),
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: ChatCompletionChunkDelta {
role: None,
content: (!content.is_empty()).then_some(content),
reasoning_content: (!tail.reasoning.is_empty())
.then_some(tail.reasoning),
tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
},
finish_reason: None,
}],
usage: None,
};
let _ =
sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
}
if streamed_calls.get() > 0 {
streamed_finish = Some("tool_calls");
}
} else {
let parsed = output::parse_output(&full_text, &offered_tools, posture);
let tool_calls: Vec<ToolCallDelta> = parsed
.calls
.iter()
.enumerate()
.map(|(index, call)| {
ToolCallDelta::whole(index, call.name.clone(), call.arguments.clone())
})
.collect();
if !tool_calls.is_empty() {
streamed_finish = Some("tool_calls");
}
if !tool_calls.is_empty()
|| !parsed.content.is_empty()
|| parsed.reasoning.is_some()
{
let payload = ChatCompletionChunk {
id: request_id.clone(),
request_id: pending_request_id.take(),
object: "chat.completion.chunk",
model: model_name.clone(),
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: ChatCompletionChunkDelta {
role: Some("assistant"),
content: (!parsed.content.is_empty() && tool_calls.is_empty())
.then(|| parsed.content.clone()),
reasoning_content: parsed.reasoning.clone(),
tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
},
finish_reason: None,
}],
usage: None,
};
let _ =
sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
}
}
let final_finish_reason = match streamed_finish {
Some(reason) if finish.as_str() != "length" => reason,
_ => finish.as_str(),
};
let final_payload = ChatCompletionChunk {
id: request_id.clone(),
request_id: pending_request_id.take(),
object: "chat.completion.chunk",
model: model_name,
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: ChatCompletionChunkDelta {
role: None,
content: None,
reasoning_content: None,
tool_calls: None,
},
finish_reason: Some(final_finish_reason),
}],
usage: Some(usage.clone()),
};
let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&final_payload)), orphan_timeout);
let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
stats_state.record_request(stats::Record {
request_id: &request_id,
route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
model: Some(served_model.clone()),
status: 200,
stream: true,
duration_ms: started.elapsed().as_millis() as u64,
usage: Some(&usage),
attribution: &attribution,
});
}
Err(e) => {
tracing::warn!("decode error on streamed request {request_id}: {e}");
stats_state.record_request(stats::Record {
request_id: &request_id,
route: ferrox_api::routes::V1_CHAT_COMPLETIONS,
model: Some(served_model.clone()),
status: 500,
stream: true,
duration_ms: started.elapsed().as_millis() as u64,
usage: None,
attribution: &attribution,
});
let payload = ChatCompletionChunk {
id: request_id.clone(),
request_id: pending_request_id.take(),
object: "chat.completion.chunk",
model: model_name,
choices: vec![ChatCompletionChunkChoice {
index: 0,
delta: ChatCompletionChunkDelta {
role: Some("assistant"),
content: Some(format!("[error: {e}]")),
reasoning_content: None,
tool_calls: None,
},
finish_reason: Some("stop"),
}],
usage: None,
};
let _ = sse::send_or_orphan(&tx, Ok(emitter.event(&payload)), orphan_timeout);
let _ = sse::send_or_orphan(&tx, Ok(emitter.done()), orphan_timeout);
}
}
drop(emitter);
});
let stream = sse::with_keepalive(rx, keepalive, sse::KEEPALIVE_INTERVAL);
Ok((
[(
axum::http::HeaderName::from_static("x-accel-buffering"),
axum::http::HeaderValue::from_static("no"),
)],
Sse::new(stream),
)
.into_response())
}
fn axum_path(template: &str) -> String {
let mut out = String::with_capacity(template.len());
let mut rest = template;
while let Some(open) = rest.find('{') {
let Some(close) = rest[open..].find('}').map(|c| open + c) else {
break;
};
out.push_str(&rest[..open]);
out.push(':');
out.push_str(&rest[open + 1..close]);
rest = &rest[close + 1..];
}
out.push_str(rest);
out
}
async fn cancel_generation(
State(state): State<Arc<AppState>>,
Json(req): Json<ferrox_api::CancelGenerationRequest>,
) -> Response {
let cancelled = state.cancels.cancel(&req.request_id);
let status = if cancelled {
StatusCode::OK
} else {
StatusCode::NOT_FOUND
};
let detail = if cancelled {
"the generation was asked to stop; it ends at its next token".to_string()
} else {
"no generation with that request_id is running -- it has already \
finished, was never issued, or was served by a path that does \
not register for cancellation"
.to_string()
};
(
status,
Json(ferrox_api::CancelGenerationResponse {
request_id: req.request_id,
cancelled,
detail,
}),
)
.into_response()
}
type Activated = (
Loaded,
Option<serving::batch::ContinuousBatcher>,
Option<Arc<budget::ContextCeiling>>,
);
fn price_batcher_config(path: Option<&str>) -> serving::batch::BatcherConfig {
let mut batcher = serving::batch::BatcherConfig::from_env();
if batcher.max_context.is_some() && batcher.kv_blocks.is_some() {
return batcher;
}
let Some(path) = path else {
return batcher;
};
let priced = budget::price_gguf(path, ferrox_models::KvElem::F32, 1, 1);
let Some((priced, gguf_ctx, source)) = priced else {
return batcher;
};
let Some(derived) = budget::derive_limits(&priced, gguf_ctx, batcher.kv_block_size) else {
tracing::warn!(
"this checkpoint's weights leave no room for KV inside the {source}: {} weight \
bytes against a {} byte budget. Serving with no derived context ceiling -- set \
FERROX_DEVICE_BUDGET_BYTES if the probe is wrong, or FERROX_CB_MAX_CONTEXT to \
admit on a number you choose.",
priced.weights_bytes,
priced.device_budget_bytes,
);
return batcher;
};
tracing::info!("{source}");
tracing::info!("{}", derived.fit);
let adopted = budget::apply_derived(&mut batcher, &derived);
if adopted.max_context {
tracing::info!(
"derived per-request context ceiling: {} token positions (prompt + max_tokens); \
override with FERROX_CB_MAX_CONTEXT",
derived.max_context
);
}
if adopted.kv_blocks {
tracing::info!(
"derived KV block budget: {} blocks x {} positions; override with FERROX_CB_KV_BLOCKS",
derived.kv_blocks,
batcher.kv_block_size
);
}
batcher
}
pub(crate) fn activate_loaded_model(
loaded: model::LoadedModel,
enable_continuous_batching: bool,
path: Option<&str>,
paged_kv: Option<&generate::PagedKvConfig>,
) -> Activated {
match loaded {
model::LoadedModel::Gguf(g) => {
let decoder = Arc::new(g.decoder);
let tokenizer = Arc::new(g.tokenizer);
let config = price_batcher_config(path);
let shape =
ferrox_models::KvShape::from_config(&decoder.config, ferrox_models::KvElem::F32, 1);
let ceiling = Arc::new(budget::ContextCeiling::new(config.max_context, shape));
let batcher = if enable_continuous_batching {
tracing::info!(
"continuous batching enabled: decode steps share Decoder::forward_multi_seq \
(stop sequences use the same pending-buffer trim as the private generate loop)"
);
let tok = Arc::clone(&tokenizer);
let decode = Arc::new(move |ids: &[usize]| tok.decode(ids));
Some(serving::batch::ContinuousBatcher::spawn_with_ceiling(
Arc::clone(&decoder),
decode,
config,
Arc::clone(&ceiling),
paged_kv.cloned(),
))
} else {
None
};
(
Loaded::Generative(Arc::new(Model::Gguf(GgufModel {
decoder,
tokenizer,
stop_tokens: g.stop_tokens,
bos_id: g.bos_id,
is_synthetic: g.is_synthetic,
chat_template: g.chat_template,
}))),
batcher,
Some(ceiling),
)
}
model::LoadedModel::Kimi(k) => (
Loaded::Generative(Arc::new(Model::Kimi(KimiModel {
engine: k.engine,
tokenizer: k.tokenizer,
stop_tokens: k.stop_tokens,
chat_template: k.chat_template,
}))),
None,
None,
),
model::LoadedModel::Mla(m) => (
Loaded::Generative(Arc::new(Model::Mla(MlaModel {
engine: m.engine,
tokenizer: m.tokenizer,
stop_tokens: m.stop_tokens,
bos_id: m.bos_id,
name: m.name,
chat_template: m.chat_template,
}))),
None,
None,
),
model::LoadedModel::Gemma4(m) => (
Loaded::Generative(Arc::new(Model::Gemma4(Gemma4Model {
engine: m.engine,
tokenizer: m.tokenizer,
stop_tokens: m.stop_tokens,
bos_id: m.bos_id,
name: m.name,
chat_template: m.chat_template,
}))),
None,
None,
),
model::LoadedModel::Glm52(g) => (
Loaded::Generative(Arc::new(Model::Glm52(Glm52Model {
engine: g.engine,
tokenizer: g.tokenizer,
stop_tokens: g.stop_tokens,
bos_id: g.bos_id,
name: g.name,
chat_template: g.chat_template,
}))),
None,
None,
),
model::LoadedModel::Encoder(e) => (Loaded::Encoder(e), None, None),
}
}
struct StartupModels {
loaded: model::LoadedModel,
embedding: Option<Arc<ferrox_models::EmbeddingModel>>,
}
fn build_app_state(
models: StartupModels,
kv_pool: Option<generate::KvPoolConfig>,
paged_kv: Option<generate::PagedKvConfig>,
prefix_cache: Option<Arc<Mutex<PrefixCache>>>,
enable_continuous_batching: bool,
mcp: Option<mcp::LoadedMcpConfig>,
detection: Arc<health::Detection>,
) -> AppState {
let StartupModels { loaded, embedding } = models;
let (loaded, batcher, ceiling) = activate_loaded_model(
loaded,
enable_continuous_batching,
std::env::var("FERROX_MODEL_PATH").ok().as_deref(),
paged_kv.as_ref(),
);
let id = startup_model_id();
AppState {
embedding,
active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
id,
loaded,
batcher,
ceiling,
}))),
paged_kv,
load_in_progress: std::sync::atomic::AtomicBool::new(false),
tasks: Arc::new(tasks::TaskRegistry::new()),
cancels: Arc::new(cancel::CancelRegistry::new()),
stats: stats::Stats::new(),
streams: resume::StreamRegistry::new(),
model_dir: admin::model_dirs().into_iter().next(),
response_cache: Mutex::new(ResponseCache::new(1000, Duration::from_secs(3600))),
kv_pool,
prefix_cache,
sessions: session::SessionStore::new(),
requests_total: std::sync::atomic::AtomicU64::new(0),
request_errors_total: std::sync::atomic::AtomicU64::new(0),
started_at: std::time::Instant::now(),
last_request_ms: std::sync::atomic::AtomicU64::new(0),
detection,
mcp,
continuous_batching_enabled: enable_continuous_batching,
loading_model: Mutex::new(None),
last_load_error: Mutex::new(None),
serving: Mutex::new(crate::stats::ServingStats::default()),
maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
started_unix: unix_now(),
}
}
fn load_embedding_model() -> anyhow::Result<Option<Arc<ferrox_models::EmbeddingModel>>> {
let Ok(path) = std::env::var("FERROX_EMBEDDING_MODEL_PATH") else {
return Ok(None);
};
let model = ferrox_models::EmbeddingModel::from_gguf_path(&path)
.map_err(|e| anyhow::anyhow!("FERROX_EMBEDDING_MODEL_PATH={path}: {e}"))?;
tracing::info!(
"loaded embedding model '{}' ({}, {} dims, pooling {}, max {} tokens)",
model.name(),
model.architecture(),
model.n_embd(),
model.pooling_type().name(),
model.n_ctx_train(),
);
Ok(Some(Arc::new(model)))
}
fn unix_now() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
fn startup_model_id() -> Option<String> {
let configured = std::env::var("FERROX_MODEL_PATH").ok()?;
let configured = std::fs::canonicalize(&configured).ok()?;
admin::discover(&admin::model_dirs())
.into_iter()
.find(|d| {
std::fs::canonicalize(&d.path)
.map(|p| p == configured)
.unwrap_or(false)
})
.map(|d| d.id)
}
fn init_cpu_pool() {
match ferrox_core::threads::init_cpu_pool() {
Some(n) => eprintln!(
"ferrox-server: rayon pool {n} threads (perf cores {}; override with FERROX_CPU_THREADS)",
ferrox_core::threads::perf_core_count()
),
None => eprintln!("ferrox-server: global rayon pool already built; leaving it alone"),
}
}
fn announce_ready(addr: SocketAddr, scheme: &str) {
use std::io::Write;
let ready =
ferrox_api::ServerReady::new(addr, scheme, env!("CARGO_PKG_VERSION"), std::process::id());
let mut stdout = std::io::stdout().lock();
let _ = writeln!(stdout, "{}", ready.to_line());
let _ = stdout.flush();
}
async fn shutdown_signal(exit_on_stdin_close: bool) {
if !exit_on_stdin_close {
std::future::pending::<()>().await;
return;
}
let _ = tokio::task::spawn_blocking(|| {
use std::io::Read;
let mut sink = [0u8; 256];
let mut stdin = std::io::stdin().lock();
loop {
match stdin.read(&mut sink) {
Ok(0) => break,
Ok(_) => continue,
Err(e) => {
tracing::warn!("stdin read failed ({e}); treating it as closed");
break;
}
}
}
})
.await;
tracing::info!("stdin closed; shutting down");
}
fn tokio_worker_threads() -> usize {
std::env::var("FERROX_TOKIO_WORKERS")
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.filter(|n| *n > 0)
.unwrap_or(2)
}
fn install_ring_crypto_provider() {
let _ = rustls::crypto::ring::default_provider().install_default();
}
pub fn run_server(args: ServerArgs) -> anyhow::Result<()> {
if args.list_devices {
print_available_devices();
return Ok(());
}
apply_cli_overrides(&args)?;
let _instance = {
use ferrox_core::instance::{register, InstancePolicy};
let policy = if args.allow_multiple_instances {
InstancePolicy::Multi
} else {
InstancePolicy::from_env_or(InstancePolicy::Single)
};
let model = std::env::var("FERROX_MODEL_PATH").ok();
register(
"server",
model.as_deref(),
ferrox_core::instance::current_backend(),
policy,
)
.map_err(|conflict| anyhow::anyhow!("{conflict}"))?
};
let journal = journal::Journal::from_env();
eprintln!(
"ferrox-server: process lifecycle journal at {:?} (override with FERROX_JOURNAL_PATH)",
journal.path()
);
journal.append(&journal::Record::session_start(
env!("CARGO_PKG_VERSION"),
std::process::id(),
));
journal::install_panic_hook(journal.clone());
let mcp_config_path = args.mcp_config.clone();
let exit_on_stdin_close = args.exit_on_stdin_close
|| std::env::var("FERROX_EXIT_ON_STDIN_CLOSE")
.map(|v| v == "1")
.unwrap_or(false);
unsafe { ferrox_core::weight_matrix::default_cpu_int_dot_on() };
init_cpu_pool();
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(tokio_worker_threads())
.enable_all()
.build()?;
let result = runtime.block_on(run(mcp_config_path, exit_on_stdin_close));
let reason = match &result {
Ok(()) => "normal".to_string(),
Err(e) => e.to_string(),
};
journal.append(&journal::Record::session_exit(reason));
runtime.shutdown_background();
result
}
async fn run(mcp_config_path: Option<PathBuf>, exit_on_stdin_close: bool) -> anyhow::Result<()> {
let _ = tracing_subscriber::fmt::try_init();
let addr = std::env::var("FERROX_ADDR").unwrap_or_else(|_| "127.0.0.1:8383".to_string());
let api_key_configured = std::env::var("FERROX_API_KEY").is_ok();
let allow_unauthenticated_remote = std::env::var("FERROX_ALLOW_UNAUTHENTICATED_REMOTE")
.map(|v| v == "1")
.unwrap_or(false);
if let Err(msg) =
security::check_bind_authorization(&addr, api_key_configured, allow_unauthenticated_remote)
{
anyhow::bail!(msg);
}
let embedding_model = load_embedding_model()?;
let mut loaded = model::load()?;
match &loaded {
model::LoadedModel::Gguf(g) => tracing::info!(
"loaded GGUF model '{}' (synthetic={}, tokenizer={})",
g.decoder.config.name,
g.is_synthetic,
g.tokenizer.kind()
),
model::LoadedModel::Kimi(k) => tracing::info!(
"loaded Kimi K3 checkpoint (tokenizer={} base tokens)",
k.tokenizer.vocab_size()
),
model::LoadedModel::Mla(m) => tracing::info!(
"loaded MLA GGUF '{}' (tokenizer={})",
m.name,
m.tokenizer.kind()
),
model::LoadedModel::Gemma4(m) => tracing::info!(
"loaded Gemma4 GGUF '{}' (tokenizer={})",
m.name,
m.tokenizer.kind()
),
model::LoadedModel::Glm52(g) => tracing::info!(
"loaded GLM-5.2 GGUF '{}' (tokenizer={})",
g.name,
g.tokenizer.kind()
),
model::LoadedModel::Encoder(_) => {}
}
let metal_default_moe_budget = {
#[cfg(feature = "metal")]
{
ferrox_core::metal_dense_enabled()
&& std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES").is_err()
}
#[cfg(not(feature = "metal"))]
{
false
}
};
if let Ok(budget_str) = std::env::var("FERROX_GPU_VRAM_BUDGET_BYTES") {
let budget: u64 = budget_str
.parse()
.expect("FERROX_GPU_VRAM_BUDGET_BYTES must be a non-negative integer");
match &mut loaded {
model::LoadedModel::Gguf(g) => {
tracing::info!(
"GPU expert placement enabled: {budget} byte VRAM budget for routed experts \
(CUDA and/or Metal matvecs when built with the matching feature)"
);
g.decoder.gpu_vram_budget_bytes = Some(budget);
}
model::LoadedModel::Kimi(_) => {
tracing::warn!(
"FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Kimi K3 -- not \
supported yet (its MoE stack isn't wired to PlacementPlan), ignoring"
);
}
model::LoadedModel::Mla(_) => {
tracing::warn!(
"FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is MLA -- dense \
FFN path only today; ignoring expert VRAM budget"
);
}
model::LoadedModel::Gemma4(_) => {
tracing::warn!(
"FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is Gemma4 -- \
ignoring expert VRAM budget"
);
}
model::LoadedModel::Glm52(_) => {
tracing::warn!(
"FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is GLM-5.2 DSA -- \
GPU expert placement not wired yet; ignoring"
);
}
model::LoadedModel::Encoder(_) => {
tracing::warn!(
"FERROX_GPU_VRAM_BUDGET_BYTES is set but the loaded model is an encoder -- \
it has no routed experts to place; ignoring"
);
}
}
} else if metal_default_moe_budget {
const METAL_DEFAULT_MOE_BUDGET: u64 = 64 * 1024 * 1024 * 1024;
if let model::LoadedModel::Gguf(g) = &mut loaded {
tracing::info!(
"Metal MoE expert placement default-on ({METAL_DEFAULT_MOE_BUDGET} byte budget); \
set FERROX_GPU_VRAM_BUDGET_BYTES=0 to force CPU experts"
);
g.decoder.gpu_vram_budget_bytes = Some(METAL_DEFAULT_MOE_BUDGET);
}
}
#[cfg(feature = "cuda")]
{
if ferrox_core::cuda_dense_enabled() {
tracing::info!(
"CUDA dense matvec enabled for WeightMatrix::apply \
(FERROX_CUDA=0|cpu forces CPU; weight buffers stay resident after first upload)"
);
} else {
tracing::info!(
"CUDA dense matvec disabled (FERROX_CUDA); dense decode uses CPU or Metal"
);
}
}
#[cfg(feature = "metal")]
{
if ferrox_core::metal_dense_enabled() {
tracing::info!(
"Metal dense matvec enabled for WeightMatrix::apply \
(FERROX_METAL=0|cpu forces CPU; weight buffers stay resident after first upload)"
);
match std::env::var("FERROX_METAL_ATTN").ok().as_deref() {
Some("1") | Some("true") | Some("on") | Some("attn") => {
tracing::info!(
"Metal fused attention requested (FERROX_METAL_ATTN): \
QKV→RoPE→GQA→O on-GPU for Norm/NeoX decode without QKV bias/QK-norm"
);
}
_ => {}
}
tracing::info!(
"Metal greedy GPU argmax: temperature<=0 folds \
final_norm+lm_head+argmax into the dense stack"
);
} else {
tracing::info!("Metal dense matvec disabled (FERROX_METAL); dense decode uses CPU");
}
}
let blocks_env = std::env::var("FERROX_KV_POOL_BLOCKS");
let block_size_env = std::env::var("FERROX_KV_POOL_BLOCK_SIZE");
let byte_budget_env = std::env::var("FERROX_KV_BYTE_BUDGET");
if blocks_env.is_ok() && byte_budget_env.is_ok() {
panic!(
"FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive \
(set one block-count source plus FERROX_KV_POOL_BLOCK_SIZE, or neither to disable)"
);
}
let kv_pool = match (blocks_env, block_size_env, byte_budget_env) {
(Ok(blocks), Ok(block_size), Err(_)) => {
let total_blocks: usize = blocks
.parse()
.expect("FERROX_KV_POOL_BLOCKS must be a positive integer");
let block_size: usize = block_size
.parse()
.expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
.ok()
.map(|v| {
v.parse()
.expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
})
.unwrap_or(0);
tracing::info!(
"KV cache block pool enabled: {total_blocks} blocks x {block_size} positions \
each, shared across all concurrent requests, {queue_wait_ms}ms admission queue wait"
);
Some(generate::KvPoolConfig {
pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
queue_wait: Duration::from_millis(queue_wait_ms),
})
}
(Err(_), Ok(block_size), Ok(byte_budget)) => {
let block_size: usize = block_size
.parse()
.expect("FERROX_KV_POOL_BLOCK_SIZE must be a positive integer");
let budget: u64 = byte_budget
.parse()
.expect("FERROX_KV_BYTE_BUDGET must be a positive integer");
let cfg = match &loaded {
model::LoadedModel::Gguf(g) => &g.decoder.config,
model::LoadedModel::Kimi(_)
| model::LoadedModel::Mla(_)
| model::LoadedModel::Gemma4(_)
| model::LoadedModel::Glm52(_)
| model::LoadedModel::Encoder(_) => {
panic!(
"FERROX_KV_BYTE_BUDGET requires a GGUF decoder model \
(set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
);
}
};
let bytes_per_block = block_size
* cfg.n_layers
* cfg.n_kv_heads
* cfg.head_dim
* 2
* std::mem::size_of::<f32>();
assert!(
bytes_per_block > 0,
"derived KV block byte size must be positive (check model config and block size)"
);
let total_blocks = (budget as usize / bytes_per_block).max(1);
let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
.ok()
.map(|v| {
v.parse()
.expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
})
.unwrap_or(0);
tracing::info!(
"KV cache block pool enabled from byte budget: {budget} bytes / \
{bytes_per_block} bytes per block ({block_size} positions x {} layers) -> \
{total_blocks} blocks, {queue_wait_ms}ms admission queue wait",
cfg.n_layers
);
Some(generate::KvPoolConfig {
pool: Arc::new(Mutex::new(KvBlockPool::new(block_size, total_blocks))),
queue_wait: Duration::from_millis(queue_wait_ms),
})
}
(Err(_), Err(_), Err(_)) => None,
(Err(_), Ok(_), Err(_)) => panic!(
"FERROX_KV_POOL_BLOCK_SIZE requires FERROX_KV_POOL_BLOCKS or FERROX_KV_BYTE_BUDGET \
(or unset all three to disable KV cache pooling)"
),
(Ok(_), Ok(_), Ok(_)) => {
unreachable!("FERROX_KV_POOL_BLOCKS and FERROX_KV_BYTE_BUDGET are mutually exclusive")
}
(Ok(_), Err(_), _) | (Err(_), Err(_), Ok(_)) => panic!(
"FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET and FERROX_KV_POOL_BLOCK_SIZE must be \
set together (or neither, to disable KV cache pooling)"
),
};
let paged_kv = match (
std::env::var("FERROX_PAGED_KV_BLOCKS"),
std::env::var("FERROX_PAGED_KV_BLOCK_SIZE"),
) {
(Ok(blocks), Ok(block_size)) => {
assert!(
kv_pool.is_none(),
"FERROX_PAGED_KV_BLOCKS and FERROX_KV_POOL_BLOCKS/FERROX_KV_BYTE_BUDGET are \
mutually exclusive: both bound the same KV memory, by different means. \
Set one."
);
let blocks_per_layer: usize = blocks
.parse()
.expect("FERROX_PAGED_KV_BLOCKS must be a positive integer");
let block_size: usize = block_size
.parse()
.expect("FERROX_PAGED_KV_BLOCK_SIZE must be a positive integer");
let gguf = match &loaded {
model::LoadedModel::Gguf(g) => g,
_ => panic!(
"FERROX_PAGED_KV_BLOCKS requires a GGUF decoder model \
(set FERROX_MODEL_PATH to a generic-decoder .gguf file)"
),
};
let cfg = &gguf.decoder.config;
let queue_wait_ms: u64 = std::env::var("FERROX_KV_POOL_QUEUE_TIMEOUT_MS")
.ok()
.map(|v| {
v.parse()
.expect("FERROX_KV_POOL_QUEUE_TIMEOUT_MS must be a non-negative integer")
})
.unwrap_or(0);
tracing::info!(
"Paged KV enabled: {blocks_per_layer} blocks x {block_size} positions per \
layer across {} layers, shared by all concurrent requests, \
{queue_wait_ms}ms admission queue wait",
cfg.n_layers
);
let radix = Some(Arc::new(Mutex::new(crate::policy::radix::RadixCache::new(
block_size,
))));
let anchor_token = crate::policy::anchor::resolve_anchor_token(
crate::policy::parser::ToolCallFormat::infer(
&std::env::var("FERROX_MODEL_PATH").unwrap_or_default(),
)
.opener(),
|text| {
gguf.tokenizer
.encode(text)
.into_iter()
.map(|t| t as u32)
.collect()
},
);
if let Some(id) = anchor_token {
tracing::info!(
"Paged KV window slide: tool-call anchor is token {id}, so a turn's \
window stops short of where its next turn rejoins"
);
}
let slide_interval: usize = std::env::var("FERROX_PAGED_KV_SLIDE_INTERVAL")
.ok()
.map(|v| {
v.parse()
.expect("FERROX_PAGED_KV_SLIDE_INTERVAL must be a positive integer")
})
.unwrap_or(crate::policy::pool_budget::DEFAULT_SWA_EVICTION_INTERVAL);
if let Some(window) = cfg.uniform_sliding_window() {
tracing::info!(
"Paged KV window slide enabled: every layer slides by {window} every \
{slide_interval} decode steps, so a request holds its prompt and a \
window rather than its whole context"
);
} else if cfg.kv_block_window().is_some() {
tracing::info!(
"Paged KV window slide NOT enabled: this model has full-attention layers, \
and a page group holds one block in every layer"
);
}
Some(generate::PagedKvConfig {
store: Arc::new(ferrox_core::cache::SharedPagedKv::new(
cfg.n_layers,
block_size,
blocks_per_layer,
cfg.n_kv_heads,
cfg.head_dim,
)),
queue_wait: Duration::from_millis(queue_wait_ms),
radix,
anchor_token,
slide_interval,
})
}
(Err(_), Err(_)) => None,
_ => panic!(
"FERROX_PAGED_KV_BLOCKS and FERROX_PAGED_KV_BLOCK_SIZE must be set together \
(or neither, to disable paged KV)"
),
};
let prefix_cache = std::env::var("FERROX_PREFIX_CACHE_ENTRIES").ok().map(|v| {
let max_entries: usize = v
.parse()
.expect("FERROX_PREFIX_CACHE_ENTRIES must be a positive integer");
if kv_pool.is_some() {
tracing::warn!(
"FERROX_PREFIX_CACHE_ENTRIES is set but so is the KV pool -- prefix \
caching will never be consulted while a KV pool is configured"
);
}
assert!(
paged_kv.is_none(),
"FERROX_PREFIX_CACHE_ENTRIES and FERROX_PAGED_KV_BLOCKS are mutually exclusive: \
the prefix cache stores contiguous KV snapshots, which a paged request does not \
produce, so the cache could never hit. Set one."
);
tracing::info!(
"KV-prefix cache enabled: up to {max_entries} stored prefixes, shared across \
all requests"
);
Arc::new(Mutex::new(PrefixCache::new(max_entries)))
});
if matches!(
loaded,
model::LoadedModel::Kimi(_) | model::LoadedModel::Mla(_) | model::LoadedModel::Glm52(_)
) && (kv_pool.is_some() || prefix_cache.is_some())
{
tracing::warn!(
"KV pool / prefix cache are configured but the loaded model is Kimi, MLA, or GLM-5.2 -- \
neither is consulted for those engines (state shapes differ from Decoder KV); see \
ferrox_models::engine's module docs"
);
}
let enable_cb = std::env::var("FERROX_CONTINUOUS_BATCHING")
.map(|v| v == "1")
.unwrap_or(false)
&& (paged_kv.is_some() || (kv_pool.is_none() && prefix_cache.is_none()))
&& matches!(loaded, model::LoadedModel::Gguf(_));
if std::env::var("FERROX_CONTINUOUS_BATCHING")
.map(|v| v == "1")
.unwrap_or(false)
&& (kv_pool.is_some() || prefix_cache.is_some())
{
tracing::warn!(
"FERROX_CONTINUOUS_BATCHING=1 ignored while KV pool or prefix cache is configured \
(those modes keep the private generate path)"
);
}
if let Ok(n) = std::env::var("FERROX_CHUNKED_PREFILL") {
if let Ok(chunk) = n.parse::<usize>() {
if chunk > 0 {
tracing::info!("chunked prefill enabled: {chunk} tokens per forward_batch chunk");
}
}
}
if matches!(
std::env::var("FERROX_CPU_KV_OFFLOAD").ok().as_deref(),
Some("1")
) {
tracing::warn!(
"FERROX_CPU_KV_OFFLOAD=1: syncing Metal KV to host after each decode step \
(minimal spill; full layer offload still planned)"
);
}
let mcp = match mcp_config_path {
Some(path) => {
let loaded = mcp::load_mcp_config(&path)?;
tracing::info!(
"MCP config loaded from {} ({} server(s); invocation not wired yet)",
loaded.path,
loaded.servers.len()
);
Some(loaded)
}
None => None,
};
let detection = health::Detection::spawn();
let state = Arc::new(build_app_state(
StartupModels {
loaded,
embedding: embedding_model,
},
kv_pool,
paged_kv,
prefix_cache,
enable_cb,
mcp,
detection,
));
use ferrox_api::routes;
let public = Router::new().route(routes::HEALTH, get(health));
let mut protected = Router::new()
.route(routes::V1_MODELS, get(list_models))
.route(routes::V1_RESPONSES, post(responses::responses))
.route(
&axum_path(routes::V1_RESPONSE),
get(responses::responses_get),
)
.route(
&axum_path(routes::V1_RESPONSE_CANCEL),
post(responses::responses_cancel),
)
.route(routes::V1_STATS, get(serving_stats))
.route(routes::V1_REQUESTS, get(recent_requests))
.route(routes::V1_CACHE_STATUS, get(cache_admin::cache_status))
.route(routes::V1_CACHE_REBUILD, post(cache_admin::cache_rebuild))
.route(routes::ADMIN_PREPARE_STOP, post(cache_admin::prepare_stop))
.route(routes::V1_CHAT_COMPLETIONS, post(chat_completions))
.route(routes::V1_CANCEL, post(cancel_generation))
.route(&axum_path(routes::V1_STREAM), get(resume::resume))
.route(&axum_path(routes::V1_STREAM_POLL), get(resume::poll))
.route(routes::V1_MESSAGES, post(anthropic::messages))
.route(
routes::V1_MESSAGES_COUNT_TOKENS,
post(anthropic::count_tokens),
)
.route(routes::V1_COMPLETIONS, post(openai_extra::completions))
.route(routes::COMPLETION, post(completion::completion))
.route(routes::COMPLETIONS, post(completion::completion))
.route(routes::V1_TOKENIZE, post(openai_extra::tokenize))
.route(routes::V1_DETOKENIZE, post(openai_extra::detokenize))
.route(routes::TOKENIZE, post(openai_extra::tokenize))
.route(routes::DETOKENIZE, post(openai_extra::detokenize))
.route(routes::V1_EMBEDDINGS, post(embeddings::embeddings))
.route(routes::CACHE_STATS, get(cache_stats))
.route(routes::METRICS, get(metrics))
.route(routes::ADMIN_MODELS, get(admin::models))
.route(routes::ADMIN_MODELS_LOAD, post(admin::load_model))
.route(routes::ADMIN_MODELS_UNLOAD, post(admin::unload_model))
.route(routes::ADMIN_DOWNLOAD, post(admin::download))
.route(routes::ADMIN_TASKS, get(admin::tasks))
.route(&admin::cancel_route(), post(admin::cancel_task))
.route(routes::ADMIN_STATS, get(admin::stats))
.merge(conversations::router());
if let Ok(key) = std::env::var("FERROX_API_KEY") {
tracing::info!("API key auth enabled");
let auth = limits::AuthConfig {
api_key: Arc::new(key),
};
protected = protected.route_layer(axum::middleware::from_fn_with_state(
auth,
limits::require_api_key,
));
}
if let Ok(rpm) = std::env::var("FERROX_RATE_LIMIT_PER_MINUTE") {
let rpm: u32 = rpm
.parse()
.expect("FERROX_RATE_LIMIT_PER_MINUTE must be a positive integer");
tracing::info!("rate limiting enabled: {rpm} requests/minute (global)");
let limiter = Arc::new(limits::RateLimiter::per_minute(rpm));
protected = protected.route_layer(axum::middleware::from_fn_with_state(
limiter,
limits::rate_limit,
));
}
if let Ok(spec) = std::env::var("FERROX_CORS_ORIGINS") {
let origins = security::parse_cors_origins(&spec)
.unwrap_or_else(|e| panic!("FERROX_CORS_ORIGINS: {e}"));
tracing::info!(
"CORS enabled: {} allow-listed origin(s) ({})",
origins.len(),
spec
);
let cors = tower_http::cors::CorsLayer::new()
.allow_origin(tower_http::cors::AllowOrigin::list(origins))
.allow_methods([axum::http::Method::GET, axum::http::Method::POST])
.allow_headers([
axum::http::header::CONTENT_TYPE,
axum::http::header::AUTHORIZATION,
axum::http::HeaderName::from_static(attribution::CLIENT_HEADER),
axum::http::HeaderName::from_static("last-event-id"),
]);
protected = protected.route_layer(cors);
}
let app = public
.merge(protected)
.layer(axum::middleware::from_fn(limits::retry_after))
.with_state(state);
let tls_paths = security::tls_paths_from_env().unwrap_or_else(|e| panic!("{e}"));
install_ring_crypto_provider();
match tls_paths {
Some(paths) => {
let config =
axum_server::tls_rustls::RustlsConfig::from_pem_file(&paths.cert, &paths.key)
.await
.map_err(|e| {
anyhow::anyhow!(
"failed to load TLS cert/key ({:?}, {:?}): {e}",
paths.cert,
paths.key
)
})?;
let socket_addr: std::net::SocketAddr = addr
.parse()
.map_err(|e| anyhow::anyhow!("invalid FERROX_ADDR {addr:?} for TLS: {e}"))?;
let listener = std::net::TcpListener::bind(socket_addr)?;
listener.set_nonblocking(true)?;
let bound = listener.local_addr()?;
tracing::info!("TLS enabled: ferrox-server listening on https://{bound}");
announce_ready(bound, "https");
let handle = axum_server::Handle::new();
let shutdown_handle = handle.clone();
tokio::spawn(async move {
shutdown_signal(exit_on_stdin_close).await;
shutdown_handle.graceful_shutdown(Some(Duration::from_secs(5)));
});
axum_server::from_tcp_rustls(listener, config)?
.handle(handle)
.serve(app.into_make_service())
.await?;
}
None => {
let listener = tokio::net::TcpListener::bind(&addr).await?;
let bound = listener.local_addr()?;
tracing::info!("ferrox-server listening on {bound}");
announce_ready(bound, "http");
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal(exit_on_stdin_close))
.await?;
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use ferrox_models::config::test_dense_fixture;
#[test]
fn parses_llama_server_style_options() {
let argv = [
"ferrox-server",
"-m",
"model.gguf",
"--host",
"::1",
"--port",
"9000",
"-t",
"4",
"-dev",
"Metal",
"-ngl",
"all",
]
.into_iter()
.map(String::from)
.collect();
let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
assert_eq!(args.model.as_deref(), Some("model.gguf"));
assert_eq!(args.host, Some(IpAddr::V6(std::net::Ipv6Addr::LOCALHOST)));
assert_eq!(args.port, Some(9000));
assert_eq!(args.threads, Some(4));
assert_eq!(args.device, Some(OffloadDevice::Metal));
assert_eq!(args.n_gpu_layers, Some(GpuLayers::All));
assert_eq!(
cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
Some("[::1]:9000")
);
}
#[test]
fn port_zero_survives_argument_parsing_as_a_real_request() {
let argv = ["ferrox-server", "--port", "0"]
.into_iter()
.map(String::from)
.collect();
let args = ServerArgs::try_parse_from(rewrite_llama_style_argv(argv)).unwrap();
assert_eq!(args.port, Some(0));
assert_eq!(
cli_bind_addr(&args, Some("127.0.0.1:8383")).as_deref(),
Some("127.0.0.1:0")
);
}
#[test]
fn stdin_close_exit_is_opt_in() {
let args =
ServerArgs::try_parse_from(["ferrox-server"].into_iter().map(String::from)).unwrap();
assert!(!args.exit_on_stdin_close);
let args = ServerArgs::try_parse_from(
["ferrox-server", "--exit-on-stdin-close"]
.into_iter()
.map(String::from),
)
.unwrap();
assert!(args.exit_on_stdin_close);
}
#[test]
fn the_ready_line_round_trips_through_a_parent_reading_stdout() {
let addr: SocketAddr = "127.0.0.1:51999".parse().unwrap();
let ready = ferrox_api::ServerReady::new(addr, "http", "0.5.0", std::process::id());
let parsed = ferrox_api::ServerReady::from_line(&ready.to_line()).unwrap();
assert_eq!(parsed.port, 51999);
assert_eq!(parsed.base_url(), "http://127.0.0.1:51999");
assert!(ferrox_api::ServerReady::from_line("INFO ferrox-server listening").is_none());
}
fn test_model() -> Model {
let cfg = test_dense_fixture();
Model::Gguf(GgufModel {
decoder: Arc::new(Decoder::new_random_small(cfg, 2, 32)),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: None,
is_synthetic: true,
chat_template: chat_template::PromptTemplate::plain(),
})
}
fn greedy_params(max_tokens: usize) -> GenerationParams {
GenerationParams {
max_tokens,
sampling: SamplingParams::default(),
seed: 1,
stop: Vec::new(),
stop_token_ids: Vec::new(),
json_object: false,
grammar: None,
cancel: None,
ignore_eos: false,
}
}
fn test_model_full_byte_vocab() -> Model {
let mut cfg = test_dense_fixture();
cfg.vocab_size = 256;
Model::Gguf(GgufModel {
decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: None,
is_synthetic: true,
chat_template: chat_template::PromptTemplate::plain(),
})
}
fn test_state(model: Model, response_cache: ResponseCache) -> AppState {
AppState {
embedding: None,
paged_kv: None,
active: std::sync::RwLock::new(Some(Arc::new(ActiveModel {
id: None,
loaded: Loaded::Generative(Arc::new(model)),
batcher: None,
ceiling: None,
}))),
load_in_progress: std::sync::atomic::AtomicBool::new(false),
tasks: Arc::new(tasks::TaskRegistry::new()),
cancels: Arc::new(cancel::CancelRegistry::new()),
stats: stats::Stats::new(),
streams: resume::StreamRegistry::new(),
model_dir: None,
response_cache: Mutex::new(response_cache),
kv_pool: None,
prefix_cache: None,
sessions: session::SessionStore::new(),
requests_total: std::sync::atomic::AtomicU64::new(0),
request_errors_total: std::sync::atomic::AtomicU64::new(0),
started_at: std::time::Instant::now(),
last_request_ms: std::sync::atomic::AtomicU64::new(0),
detection: Arc::new(health::Detection::ready(health::probe_backends())),
mcp: None,
continuous_batching_enabled: false,
loading_model: Mutex::new(None),
last_load_error: Mutex::new(None),
serving: Mutex::new(crate::stats::ServingStats::default()),
maintenance: Mutex::new(crate::policy::maintenance::MaintenanceGate::serving()),
footprint: Mutex::new(crate::policy::footprint::ProbeCache::new(FOOTPRINT_TTL_MS)),
started_unix: unix_now(),
}
}
fn test_app() -> Router {
test_app_with_state(Arc::new(test_state(
test_model_full_byte_vocab(),
ResponseCache::new(1000, Duration::from_secs(3600)),
)))
}
fn test_app_with_state(state: Arc<AppState>) -> Router {
Router::new()
.route(ferrox_api::routes::HEALTH, get(health))
.route(ferrox_api::routes::V1_MODELS, get(list_models))
.route(ferrox_api::routes::V1_RESPONSES, post(responses::responses))
.route(
&axum_path(ferrox_api::routes::V1_RESPONSE),
get(responses::responses_get),
)
.route(
&axum_path(ferrox_api::routes::V1_RESPONSE_CANCEL),
post(responses::responses_cancel),
)
.route(ferrox_api::routes::V1_STATS, get(serving_stats))
.route(ferrox_api::routes::V1_REQUESTS, get(recent_requests))
.route(
ferrox_api::routes::V1_CACHE_STATUS,
get(cache_admin::cache_status),
)
.route(
ferrox_api::routes::V1_CACHE_REBUILD,
post(cache_admin::cache_rebuild),
)
.route(
ferrox_api::routes::ADMIN_PREPARE_STOP,
post(cache_admin::prepare_stop),
)
.route("/v1/chat/completions", post(chat_completions))
.route(ferrox_api::routes::V1_MESSAGES, post(anthropic::messages))
.route(
ferrox_api::routes::V1_MESSAGES_COUNT_TOKENS,
post(anthropic::count_tokens),
)
.route("/v1/tokenize", post(openai_extra::tokenize))
.route("/v1/detokenize", post(openai_extra::detokenize))
.route(ferrox_api::routes::TOKENIZE, post(openai_extra::tokenize))
.route(
ferrox_api::routes::DETOKENIZE,
post(openai_extra::detokenize),
)
.route("/v1/embeddings", post(embeddings::embeddings))
.route("/v1/completions", post(openai_extra::completions))
.route(ferrox_api::routes::COMPLETION, post(completion::completion))
.route(
ferrox_api::routes::COMPLETIONS,
post(completion::completion),
)
.route(
ferrox_api::routes::ADMIN_MODELS_UNLOAD,
post(admin::unload_model),
)
.route(ferrox_api::routes::ADMIN_TASKS, get(admin::tasks))
.route(ferrox_api::routes::ADMIN_STATS, get(admin::stats))
.route(ferrox_api::routes::V1_CANCEL, post(cancel_generation))
.route(
&axum_path(ferrox_api::routes::V1_STREAM),
get(resume::resume),
)
.route(
&axum_path(ferrox_api::routes::V1_STREAM_POLL),
get(resume::poll),
)
.with_state(state)
}
fn named_test_model(name: &'static str, vocab_size: usize) -> Model {
let mut cfg = test_dense_fixture();
cfg.name = name;
cfg.vocab_size = vocab_size;
Model::Gguf(GgufModel {
decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: None,
is_synthetic: true,
chat_template: chat_template::PromptTemplate::plain(),
})
}
fn model_with_template(name: &'static str, source: &str) -> Model {
let mut cfg = test_dense_fixture();
cfg.name = name;
cfg.vocab_size = 256;
Model::Gguf(GgufModel {
decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: None,
is_synthetic: true,
chat_template: chat_template::PromptTemplate::from_gguf_metadata(
Some(source),
Some("qwen3"),
false,
true,
None,
None,
),
})
}
#[tokio::test]
async fn a_template_that_rejects_the_conversation_is_a_400_on_the_streaming_path() {
let strict = "{% if messages | length > 1 %}\
{{ raise_exception('this template takes one turn') }}\
{% endif %}{{ messages[0].content }}";
let state = Arc::new(test_state(
model_with_template("strict", strict),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(state);
let (status, body) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "strict",
"stream": true,
"messages": [
{"role": "user", "content": "one"},
{"role": "user", "content": "two"},
],
}),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["error"]["param"], serde_json::json!("messages"));
assert!(
body["error"]["message"]
.as_str()
.unwrap()
.contains("one turn"),
"the template's own message must reach the caller: {body}"
);
let (status, _) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "strict",
"stream": true,
"max_tokens": 1,
"messages": [{"role": "user", "content": "one"}],
}),
)
.await;
assert_eq!(status, StatusCode::OK);
}
#[tokio::test]
async fn models_advertises_the_gears_this_checkpoint_actually_has() {
let reasoning = "{% if enable_thinking %}<think>{% endif %}\
{% if reasoning_effort %}\
{% if reasoning_effort not in ['low','medium','high'] %}\
{{ raise_exception('bad effort') }}\
{% endif %}[{{ reasoning_effort }}]\
{% endif %}{{ messages[0].content }}";
let state = Arc::new(test_state(
model_with_template("thinker", reasoning),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(state);
let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
assert_eq!(status, StatusCode::OK);
let entry = &models["data"][0];
assert_eq!(
entry["supported_reasoning_efforts"],
serde_json::json!(["off", "low", "medium", "high"])
);
assert_eq!(entry["default_reasoning_effort"], serde_json::json!("off"));
}
#[tokio::test]
async fn a_checkpoint_with_no_thinking_controls_advertises_neither_field() {
let app = test_app();
let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
let entry = &models["data"][0];
assert!(entry.get("supported_reasoning_efforts").is_none());
assert!(entry.get("default_reasoning_effort").is_none());
}
fn active_model(state: &AppState, name: &'static str) -> Arc<ActiveModel> {
Arc::new(ActiveModel {
id: Some(name.to_string()),
loaded: Loaded::Generative(Arc::new(named_test_model(name, 256))),
batcher: None,
ceiling: None,
})
.tap_into(state)
}
trait TapInto {
fn tap_into(self, state: &AppState) -> Self;
}
impl TapInto for Arc<ActiveModel> {
fn tap_into(self, state: &AppState) -> Self {
state.swap_active(Some(Arc::clone(&self)));
self
}
}
#[test]
fn an_in_flight_request_keeps_the_model_it_started_on() {
let state = test_state(
named_test_model("model-a", 256),
ResponseCache::new(4, Duration::from_secs(60)),
);
let in_flight = state.active().expect("a model is loaded");
assert_eq!(in_flight.name(), "model-a");
active_model(&state, "model-b");
assert_eq!(state.active().unwrap().name(), "model-b");
assert_eq!(in_flight.name(), "model-a");
let (_chunks, finish, _usage) = run_generation(
in_flight.generative().unwrap(),
"hi",
&greedy_params(3),
None,
None,
None,
None,
None,
)
.expect("the old model must still decode after being swapped out");
assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
}
#[test]
fn a_swapped_out_model_lives_until_its_last_holder_releases_it() {
let state = test_state(
named_test_model("model-a", 256),
ResponseCache::new(4, Duration::from_secs(60)),
);
let in_flight = state.active().expect("a model is loaded");
let weights = Arc::clone(in_flight.generative().unwrap());
assert!(Arc::strong_count(&weights) >= 2);
let previous = state.swap_active(Some(Arc::new(ActiveModel {
id: Some("model-b".to_string()),
loaded: Loaded::Generative(Arc::new(named_test_model("model-b", 256))),
batcher: None,
ceiling: None,
})));
drop(previous);
assert!(Arc::strong_count(&weights) >= 2);
drop(in_flight);
assert_eq!(Arc::strong_count(&weights), 1);
}
#[tokio::test]
async fn unloading_answers_503_instead_of_serving_the_dropped_model() {
let state = Arc::new(test_state(
named_test_model("model-a", 256),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(Arc::clone(&state));
let (status, body) = post_json_uri(
&app,
ferrox_api::routes::ADMIN_MODELS_UNLOAD,
serde_json::json!({}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["ok"], true);
assert!(body["active"].is_null());
assert!(state.active().is_none());
let (status, _) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
assert_eq!(status, StatusCode::OK);
let (_, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
assert_eq!(models["data"].as_array().unwrap().len(), 0);
let (status, body) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "x",
"messages": [{"role": "user", "content": "hi"}]
}),
)
.await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(body["error"]["type"], "model_not_loaded");
}
#[tokio::test]
async fn health_reports_the_unloaded_state_rather_than_going_silent() {
let state = Arc::new(test_state(
named_test_model("model-a", 256),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(Arc::clone(&state));
state.swap_active(None);
let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(body["state"], "unavailable");
assert_eq!(body["reason"], "model_not_loaded");
assert!(body["model"].is_null());
let real_weights = body["capabilities"]
.as_array()
.unwrap()
.iter()
.find(|c| c["id"] == "real_weights")
.cloned()
.expect("real_weights is always reported");
assert_eq!(real_weights["available"], false);
assert_eq!(real_weights["reason"], "model_not_loaded");
}
#[tokio::test]
async fn a_finished_request_lands_in_the_stats_ring_with_both_durations() {
let app = test_app();
let (status, completion) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "x",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4
}),
)
.await;
assert_eq!(status, StatusCode::OK);
let request_id = completion["request_id"].as_str().unwrap().to_string();
let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
assert_eq!(status, StatusCode::OK);
let recent = stats["recent"].as_array().unwrap();
assert_eq!(recent.len(), 1);
let row = &recent[0];
assert_eq!(row["request_id"], request_id);
assert_eq!(row["route"], ferrox_api::routes::V1_CHAT_COMPLETIONS);
assert_eq!(row["status"], 200);
assert_eq!(row["stream"], false);
assert!(row["duration_ms"].is_number());
assert!(row["decode_ms"].is_number());
assert!(stats["tokens_generated_total"].as_u64().unwrap() > 0);
assert_eq!(
stats["tokens_prompt_total"].as_u64().unwrap(),
row["prompt_tokens"].as_u64().unwrap()
);
}
#[tokio::test]
async fn a_rejected_request_is_recorded_too() {
let state = Arc::new(test_state(
named_test_model("model-a", 256),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(Arc::clone(&state));
state.swap_active(None);
let (status, _) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({"model": "x", "messages": [{"role": "user", "content": "hi"}]}),
)
.await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let recent = stats["recent"].as_array().unwrap();
assert_eq!(recent.len(), 1);
assert_eq!(recent[0]["status"], 503);
assert_eq!(recent[0]["completion_tokens"], 0);
assert!(recent[0]["decode_ms"].is_null());
assert_eq!(stats["errors_total"], 1);
}
async fn post_json_with_headers(
app: &Router,
uri: &str,
body: serde_json::Value,
headers: &[(&str, &str)],
) -> (StatusCode, serde_json::Value) {
use http_body_util::BodyExt;
use tower::ServiceExt;
let mut builder = axum::http::Request::builder()
.method("POST")
.uri(uri)
.header("content-type", "application/json");
for (name, value) in headers {
builder = builder.header(*name, *value);
}
let response = app
.clone()
.oneshot(
builder
.body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let bytes = response.into_body().collect().await.unwrap().to_bytes();
let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
(status, json)
}
#[tokio::test]
async fn tokenize_detokenize_and_embeddings_all_land_in_the_ring() {
let app = test_app();
let (status, _) = post_json_uri(
&app,
ferrox_api::routes::V1_TOKENIZE,
serde_json::json!({"prompt": "hello"}),
)
.await;
assert_eq!(status, StatusCode::OK);
let (status, _) = post_json_uri(
&app,
ferrox_api::routes::V1_DETOKENIZE,
serde_json::json!({"tokens": [104, 105]}),
)
.await;
assert_eq!(status, StatusCode::OK);
let (status, _) = post_json_uri(
&app,
ferrox_api::routes::V1_EMBEDDINGS,
serde_json::json!({"input": "hello"}),
)
.await;
assert_eq!(status, StatusCode::OK);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let routes: Vec<&str> = stats["recent"]
.as_array()
.unwrap()
.iter()
.map(|row| row["route"].as_str().unwrap())
.collect();
for expected in [
ferrox_api::routes::V1_TOKENIZE,
ferrox_api::routes::V1_DETOKENIZE,
ferrox_api::routes::V1_EMBEDDINGS,
] {
assert!(
routes.contains(&expected),
"{expected} is missing: {routes:?}"
);
}
let row = |route: &str| {
stats["recent"]
.as_array()
.unwrap()
.iter()
.find(|r| r["route"] == route)
.cloned()
.unwrap()
};
let embed = row(ferrox_api::routes::V1_EMBEDDINGS);
assert!(embed["prompt_tokens"].as_u64().unwrap() > 0);
assert!(embed["decode_ms"].is_null());
assert_eq!(embed["completion_tokens"], 0);
assert_eq!(row(ferrox_api::routes::V1_TOKENIZE)["prompt_tokens"], 0);
assert_eq!(
stats["tokens_prompt_total"].as_u64().unwrap(),
embed["prompt_tokens"].as_u64().unwrap(),
"only the forward pass counted"
);
}
fn streaming_test_app() -> Router {
let mut cfg = test_dense_fixture();
cfg.vocab_size = 256;
let model = Model::Gguf(GgufModel {
decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: None,
is_synthetic: false,
chat_template: chat_template::PromptTemplate::plain(),
});
test_app_with_state(Arc::new(test_state(
model,
ResponseCache::new(1000, Duration::from_secs(3600)),
)))
}
#[tokio::test]
async fn the_native_completion_wire_is_not_the_openai_one() {
let app = test_app();
let (status, native) = post_json_uri(
&app,
ferrox_api::routes::COMPLETION,
serde_json::json!({"prompt": "hi", "n_predict": 4}),
)
.await;
assert_eq!(status, StatusCode::OK, "{native}");
assert!(native["content"].is_string(), "{native}");
assert_eq!(native["stop"], true);
assert_eq!(native["stop_type"], "limit");
assert_eq!(native["stopping_word"], "");
assert_eq!(native["truncated"], false);
assert_eq!(native["id_slot"], -1);
assert!(native["timings"]["prompt_n"].is_number(), "{native}");
assert!(native["generation_settings"]["n_predict"] == 4, "{native}");
assert!(
native.get("choices").is_none(),
"the native shape has no `choices`: {native}"
);
let (status, openai) = post_json_uri(
&app,
ferrox_api::routes::V1_COMPLETIONS,
serde_json::json!({"prompt": "hi", "max_tokens": 4}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert!(openai["choices"][0]["text"].is_string(), "{openai}");
assert!(
openai.get("content").is_none(),
"the OpenAI shape has no top-level `content`: {openai}"
);
}
#[tokio::test]
async fn both_native_spellings_reach_the_same_handler() {
let app = test_app();
for route in [
ferrox_api::routes::COMPLETION,
ferrox_api::routes::COMPLETIONS,
] {
let (status, body) = post_json_uri(
&app,
route,
serde_json::json!({"prompt": "hi", "n_predict": 2, "seed": 1}),
)
.await;
assert_eq!(status, StatusCode::OK, "{route}: {body}");
assert_eq!(body["stop"], true, "{route}");
assert!(body["content"].is_string(), "{route}");
}
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let routes: Vec<&str> = stats["recent"]
.as_array()
.unwrap()
.iter()
.map(|row| row["route"].as_str().unwrap())
.collect();
assert!(
routes.contains(&ferrox_api::routes::COMPLETION),
"{routes:?}"
);
assert!(
routes.contains(&ferrox_api::routes::COMPLETIONS),
"{routes:?}"
);
}
#[tokio::test]
async fn a_native_stream_ends_on_a_stop_frame_with_no_done_sentinel() {
let app = streaming_test_app();
let raw = post_sse_raw_uri(
&app,
ferrox_api::routes::COMPLETION,
serde_json::json!({"prompt": "hi", "n_predict": 6, "stream": true, "seed": 7}),
)
.await;
assert!(
!raw.contains("[DONE]"),
"llama.cpp's native stream has no sentinel: {raw}"
);
let frames: Vec<serde_json::Value> = raw
.lines()
.filter_map(|line| line.strip_prefix("data: "))
.map(|json| serde_json::from_str(json).expect("every frame is one JSON object"))
.collect();
assert!(frames.len() >= 2, "expected partials then a final: {raw}");
let (last, partials) = frames.split_last().unwrap();
assert_eq!(last["stop"], true, "the last frame closes the stream");
assert!(last["timings"].is_object(), "{last}");
assert!(last["stop_type"].is_string(), "{last}");
for partial in partials {
assert_eq!(partial["stop"], false, "{partial}");
assert!(partial["content"].is_string(), "{partial}");
assert!(partial.get("timings").is_none(), "{partial}");
assert!(partial.get("generation_settings").is_none(), "{partial}");
}
let streamed: String = partials
.iter()
.filter_map(|p| p["content"].as_str())
.collect();
assert_eq!(last["content"].as_str().unwrap(), streamed);
}
#[tokio::test]
async fn an_unbounded_n_predict_is_refused_rather_than_quietly_shrunk() {
let app = test_app();
for body in [
serde_json::json!({"prompt": "hi"}),
serde_json::json!({"prompt": "hi", "n_predict": -1}),
] {
let (status, refusal) =
post_json_uri(&app, ferrox_api::routes::COMPLETION, body.clone()).await;
assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}: {refusal}");
assert!(
refusal["error"]["message"]
.as_str()
.unwrap()
.contains("n_predict"),
"{refusal}"
);
}
let (status, _) = post_json_uri(
&app,
ferrox_api::routes::COMPLETION,
serde_json::json!({"prompt": "hi", "n_predict": 2}),
)
.await;
assert_eq!(status, StatusCode::OK);
}
#[tokio::test]
async fn a_stop_string_halts_the_answer_and_is_named_back() {
let app = streaming_test_app();
let ask = |stop: serde_json::Value| {
let app = app.clone();
async move {
post_json_uri(
&app,
ferrox_api::routes::COMPLETION,
serde_json::json!({
"prompt": "hi",
"n_predict": 64,
"ignore_eos": true,
"stop": stop,
}),
)
.await
.1
}
};
let baseline = ask(serde_json::json!([])).await;
assert_eq!(baseline["stop_type"], "limit");
assert_eq!(baseline["stopping_word"], "");
let text = baseline["content"].as_str().unwrap().to_string();
let sentinel: String = text.chars().skip(1).take(2).collect();
assert_eq!(
sentinel.chars().count(),
2,
"the fixture must produce enough output to cut: {text:?}"
);
let cut = text.find(&sentinel).expect("it came out of this text");
let stopped = ask(serde_json::json!([sentinel])).await;
assert_eq!(stopped["stop_type"], "word", "{stopped}");
assert_eq!(stopped["stopping_word"], sentinel);
assert_eq!(
stopped["content"].as_str().unwrap(),
&text[..cut],
"the answer must be cut at the sentinel, not run past it"
);
}
#[tokio::test]
async fn the_llama_cpp_spelling_of_tokenize_reaches_the_same_handler() {
let app = test_app();
let (v1_status, v1) = post_json_uri(
&app,
ferrox_api::routes::V1_TOKENIZE,
serde_json::json!({"prompt": "hello"}),
)
.await;
let (alias_status, alias) = post_json_uri(
&app,
ferrox_api::routes::TOKENIZE,
serde_json::json!({"content": "hello"}),
)
.await;
assert_eq!(v1_status, StatusCode::OK);
assert_eq!(alias_status, StatusCode::OK, "{alias}");
assert_eq!(v1["tokens"], alias["tokens"]);
assert!(!alias["tokens"].as_array().unwrap().is_empty());
let (status, both_ways) = post_json_uri(
&app,
ferrox_api::routes::TOKENIZE,
serde_json::json!({"prompt": "hello"}),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(both_ways["tokens"], v1["tokens"]);
}
#[tokio::test]
async fn detokenize_answers_under_both_dialects_keys() {
let app = test_app();
for route in [
ferrox_api::routes::DETOKENIZE,
ferrox_api::routes::V1_DETOKENIZE,
] {
let (status, body) =
post_json_uri(&app, route, serde_json::json!({"tokens": [104, 105]})).await;
assert_eq!(status, StatusCode::OK, "{route}");
assert_eq!(body["text"], "hi", "{route}");
assert_eq!(body["content"], body["text"], "{route}");
}
}
#[tokio::test]
async fn the_alias_is_recorded_under_the_path_the_client_called() {
let app = test_app();
let (status, _) = post_json_uri(
&app,
ferrox_api::routes::TOKENIZE,
serde_json::json!({"content": "hello"}),
)
.await;
assert_eq!(status, StatusCode::OK);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let routes: Vec<&str> = stats["recent"]
.as_array()
.unwrap()
.iter()
.map(|row| row["route"].as_str().unwrap())
.collect();
assert!(
routes.contains(&ferrox_api::routes::TOKENIZE),
"the alias must be its own row: {routes:?}"
);
assert!(
!routes.contains(&ferrox_api::routes::V1_TOKENIZE),
"nothing called /v1/tokenize: {routes:?}"
);
}
#[tokio::test]
async fn add_special_prepends_the_same_bos_the_decoder_would() {
let mut cfg = test_dense_fixture();
cfg.vocab_size = 256;
let model = Model::Gguf(GgufModel {
decoder: Arc::new(Decoder::new_random_small(cfg, 2, 256)),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: Some(7),
is_synthetic: true,
chat_template: chat_template::PromptTemplate::plain(),
});
let app = test_app_with_state(Arc::new(test_state(
model,
ResponseCache::new(1000, Duration::from_secs(3600)),
)));
let (_, plain) = post_json_uri(
&app,
ferrox_api::routes::TOKENIZE,
serde_json::json!({"content": "hi"}),
)
.await;
let (_, special) = post_json_uri(
&app,
ferrox_api::routes::TOKENIZE,
serde_json::json!({"content": "hi", "add_special": true}),
)
.await;
assert_eq!(plain["tokens"], serde_json::json!([104, 105]));
assert_eq!(special["tokens"], serde_json::json!([7, 104, 105]));
assert_eq!(special["count"], 3);
}
#[tokio::test]
async fn a_rejected_embeddings_request_is_recorded_with_its_status() {
let app = test_app();
let (status, _) = post_json_uri(
&app,
ferrox_api::routes::V1_EMBEDDINGS,
serde_json::json!({"input": "hi", "encoding_format": "base64"}),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let recent = stats["recent"].as_array().unwrap();
assert_eq!(recent.len(), 1);
assert_eq!(recent[0]["route"], ferrox_api::routes::V1_EMBEDDINGS);
assert_eq!(recent[0]["status"], 400);
assert_eq!(
recent[0]["prompt_tokens"], 0,
"a rejected call embedded nothing"
);
}
#[tokio::test]
async fn a_row_names_the_key_that_served_it_without_carrying_the_key() {
let app = test_app();
let key = "sk-monitor-secret";
let (status, _) = post_json_with_headers(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "x",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 2
}),
&[
("authorization", &format!("Bearer {key}")),
("x-ferrox-client", "ferrox-studio"),
],
)
.await;
assert_eq!(status, StatusCode::OK);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let row = stats["recent"].as_array().unwrap()[0].clone();
let fingerprint = row["via_api_key"]
.as_str()
.expect("the row names the key that served it")
.to_string();
assert_eq!(fingerprint, attribution::key_fingerprint(key));
assert!(!fingerprint.contains(key));
assert!(
!serde_json::to_string(&stats).unwrap().contains(key),
"the stats payload must not carry the key in any form"
);
assert_eq!(row["client"], "ferrox-studio");
}
#[tokio::test]
async fn different_keys_are_different_callers_and_no_key_is_null() {
let app = test_app();
let body = serde_json::json!({
"model": "x",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 1
});
for headers in [
vec![("authorization", "Bearer key-one")],
vec![("authorization", "Bearer key-two")],
vec![],
] {
let (status, _) =
post_json_with_headers(&app, "/v1/chat/completions", body.clone(), &headers).await;
assert_eq!(status, StatusCode::OK);
}
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let recent = stats["recent"].as_array().unwrap();
assert_eq!(recent.len(), 3);
let one = recent[0]["via_api_key"].as_str().unwrap();
let two = recent[1]["via_api_key"].as_str().unwrap();
assert_ne!(one, two, "two keys must not collapse into one caller");
assert!(
recent[2]["via_api_key"].is_null(),
"an unauthenticated call is null, not a fingerprint of nothing"
);
assert!(recent[2]["client"].is_null());
}
#[tokio::test]
async fn a_row_names_the_model_that_served_it_not_the_one_requested() {
let state = Arc::new(test_state(
named_test_model("really-loaded", 256),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(Arc::clone(&state));
let (status, _) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "gpt-4-turbo-that-is-not-here",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 2
}),
)
.await;
assert_eq!(status, StatusCode::OK);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
assert_eq!(stats["recent"][0]["model"], "really-loaded");
state.swap_active(None);
let (status, _) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "gpt-4-turbo-that-is-not-here",
"messages": [{"role": "user", "content": "hi"}]
}),
)
.await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
let recent = stats["recent"].as_array().unwrap();
assert!(recent[recent.len() - 1]["model"].is_null());
}
#[tokio::test]
async fn a_streamed_row_names_the_model_it_decoded_against() {
let state = Arc::new(test_state(
named_test_model("model-before", 256),
ResponseCache::new(4, Duration::from_secs(60)),
));
let app = test_app_with_state(Arc::clone(&state));
let _ = post_sse_raw(&app, resumable_request()).await;
active_model(&state, "model-after");
let (_, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
assert_eq!(stats["recent"][0]["model"], "model-before");
}
#[tokio::test]
async fn the_queue_gauge_is_null_when_nothing_can_queue() {
let app = test_app();
let (status, stats) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
assert_eq!(status, StatusCode::OK);
assert!(
stats["queue_depth"].is_null(),
"without continuous batching nothing queues, so there is nothing to measure"
);
assert!(stats["queue_rejected_total"].is_null());
assert_eq!(
stats["generating_now"], 0,
"work in progress is measured and really is zero here"
);
}
async fn post_sse_raw(app: &Router, body: serde_json::Value) -> String {
post_sse_raw_uri(app, ferrox_api::routes::V1_CHAT_COMPLETIONS, body).await
}
async fn post_sse_raw_uri(app: &Router, uri: &str, body: serde_json::Value) -> String {
use http_body_util::BodyExt;
use tower::ServiceExt;
let response = app
.clone()
.oneshot(
axum::http::Request::builder()
.method("POST")
.uri(uri)
.header("content-type", "application/json")
.body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
let bytes = response.into_body().collect().await.unwrap().to_bytes();
String::from_utf8(bytes.to_vec()).unwrap()
}
async fn get_json_with_headers(
app: &Router,
uri: &str,
headers: &[(&str, &str)],
) -> (StatusCode, serde_json::Value) {
use http_body_util::BodyExt;
use tower::ServiceExt;
let mut builder = axum::http::Request::builder().method("GET").uri(uri);
for (name, value) in headers {
builder = builder.header(*name, *value);
}
let response = app
.clone()
.oneshot(builder.body(axum::body::Body::empty()).unwrap())
.await
.unwrap();
let status = response.status();
let bytes = response.into_body().collect().await.unwrap().to_bytes();
(
status,
serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({})),
)
}
fn sse_field<'a>(body: &'a str, field: &str) -> Vec<&'a str> {
body.lines()
.filter_map(|line| line.strip_prefix(field))
.map(str::trim)
.collect()
}
fn resumable_request() -> serde_json::Value {
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 4,
"temperature": 0,
"stream": true,
"stream_resumable": true,
})
}
#[tokio::test]
async fn a_resumable_stream_numbers_every_event_and_states_retry_once() {
let app = test_app();
let body = post_sse_raw(&app, resumable_request()).await;
let request_id = body
.lines()
.find_map(|l| l.strip_prefix("data: "))
.and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
.and_then(|v| v["request_id"].as_str().map(str::to_string))
.expect("the first chunk names the request");
let ids = sse_field(&body, "id:");
let datas = sse_field(&body, "data:");
assert_eq!(
ids.len(),
datas.len(),
"every event carries an id, or a reconnect cannot name where it stopped"
);
for (i, id) in ids.iter().enumerate() {
assert_eq!(*id, format!("{request_id}:{i}"));
}
let retries = sse_field(&body, "retry:");
assert_eq!(
retries.len(),
1,
"the reconnect delay is stated once, not on every event"
);
assert_eq!(retries[0], "1500");
assert!(
body.contains("data: [DONE]"),
"the end of stream is still stated"
);
}
#[tokio::test]
async fn a_plain_stream_carries_no_id_because_nothing_could_replay_it() {
let app = test_app();
let mut request = resumable_request();
request["stream_resumable"] = serde_json::json!(false);
let body = post_sse_raw(&app, request).await;
assert!(!sse_field(&body, "data:").is_empty(), "it still streams");
assert!(
sse_field(&body, "id:").is_empty(),
"an id promises a replay this stream cannot serve"
);
assert!(sse_field(&body, "retry:").is_empty());
}
#[tokio::test]
async fn the_polling_fallback_serves_exactly_what_the_stream_delivered() {
let app = test_app();
let body = post_sse_raw(&app, resumable_request()).await;
let request_id = sse_field(&body, "id:")[0]
.rsplit_once(':')
.unwrap()
.0
.to_string();
let streamed: Vec<String> = sse_field(&body, "data:")
.iter()
.map(|d| d.to_string())
.collect();
let (status, polled) = get_json(
&app,
&format!("{}?from=0", ferrox_api::routes::v1_stream_poll(&request_id)),
)
.await;
assert_eq!(status, StatusCode::OK);
let events: Vec<String> = polled["events"]
.as_array()
.unwrap()
.iter()
.map(|e| e["data"].as_str().unwrap().to_string())
.collect();
assert_eq!(
events, streamed,
"the fallback must deliver the same answer, not a re-run of it"
);
assert_eq!(polled["request_id"], request_id);
assert_eq!(
polled["done"], false,
"events were still being handed out, so the client must ask again"
);
let next = polled["next_index"].as_u64().unwrap();
let (_, drained) = get_json(
&app,
&format!(
"{}?from={next}",
ferrox_api::routes::v1_stream_poll(&request_id)
),
)
.await;
assert_eq!(drained["done"], true);
assert_eq!(drained["events"].as_array().unwrap().len(), 0);
}
#[tokio::test]
async fn a_resume_continues_after_the_last_event_id_rather_than_repeating() {
let app = test_app();
let body = post_sse_raw(&app, resumable_request()).await;
let ids = sse_field(&body, "id:");
let datas: Vec<String> = sse_field(&body, "data:")
.iter()
.map(|d| d.to_string())
.collect();
assert!(
ids.len() >= 3,
"need a few events to resume into the middle"
);
let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
let (status, resumed) = get_json_with_headers(
&app,
&format!("{}/poll", ferrox_api::routes::v1_stream(&request_id)),
&[],
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(resumed["events"].as_array().unwrap().len(), datas.len());
let (_, tail) = get_json(
&app,
&format!("{}?from=2", ferrox_api::routes::v1_stream_poll(&request_id)),
)
.await;
let tail_events: Vec<String> = tail["events"]
.as_array()
.unwrap()
.iter()
.map(|e| e["data"].as_str().unwrap().to_string())
.collect();
assert_eq!(tail_events, datas[2..].to_vec());
}
#[tokio::test]
async fn an_sse_reconnect_resumes_from_the_last_event_id() {
use http_body_util::BodyExt;
use tower::ServiceExt;
let app = test_app();
let body = post_sse_raw(&app, resumable_request()).await;
let ids = sse_field(&body, "id:");
let datas: Vec<String> = sse_field(&body, "data:")
.iter()
.map(|d| d.to_string())
.collect();
let request_id = ids[0].rsplit_once(':').unwrap().0.to_string();
let response = app
.clone()
.oneshot(
axum::http::Request::builder()
.method("GET")
.uri(ferrox_api::routes::v1_stream(&request_id))
.header("last-event-id", format!("{request_id}:0"))
.body(axum::body::Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get("x-accel-buffering")
.and_then(|v| v.to_str().ok()),
Some("no"),
"the reconnect needs the same anti-buffering header as the stream"
);
let bytes = response.into_body().collect().await.unwrap().to_bytes();
let resumed = String::from_utf8(bytes.to_vec()).unwrap();
assert_eq!(
sse_field(&resumed, "data:")
.iter()
.map(|d| d.to_string())
.collect::<Vec<_>>(),
datas[1..].to_vec()
);
assert_eq!(sse_field(&resumed, "id:")[0], format!("{request_id}:1"));
}
#[tokio::test]
async fn a_last_event_id_from_another_stream_is_refused() {
let app = test_app();
let body = post_sse_raw(&app, resumable_request()).await;
let request_id = sse_field(&body, "id:")[0]
.rsplit_once(':')
.unwrap()
.0
.to_string();
let (status, err) = get_json_with_headers(
&app,
&ferrox_api::routes::v1_stream(&request_id),
&[("last-event-id", "chatcmpl-someone-else:3")],
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(err["error"]["code"], "bad_last_event_id");
}
#[tokio::test]
async fn resuming_a_stream_that_was_never_resumable_is_a_404_that_says_why() {
let app = test_app();
let mut request = resumable_request();
request["stream_resumable"] = serde_json::json!(false);
let body = post_sse_raw(&app, request).await;
let request_id = body
.lines()
.find_map(|l| l.strip_prefix("data: "))
.and_then(|d| serde_json::from_str::<serde_json::Value>(d).ok())
.and_then(|v| v["request_id"].as_str().map(str::to_string))
.unwrap();
let (status, err) = get_json(&app, &ferrox_api::routes::v1_stream_poll(&request_id)).await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert_eq!(err["error"]["code"], "stream_not_found");
assert!(err["error"]["message"]
.as_str()
.unwrap()
.contains("stream_resumable"));
}
#[test]
fn the_axum_stream_patterns_match_the_published_templates() {
assert_eq!(
axum_path(ferrox_api::routes::V1_STREAM),
"/v1/stream/:request_id"
);
assert_eq!(
axum_path(ferrox_api::routes::V1_STREAM_POLL),
"/v1/stream/:request_id/poll"
);
assert_eq!(
ferrox_api::routes::v1_stream("abc"),
axum_path(ferrox_api::routes::V1_STREAM).replace(":request_id", "abc")
);
}
#[test]
fn no_published_template_reaches_the_router_with_its_braces() {
for template in [
ferrox_api::routes::V1_STREAM,
ferrox_api::routes::V1_STREAM_POLL,
ferrox_api::routes::V1_RESPONSE,
ferrox_api::routes::V1_RESPONSE_CANCEL,
ferrox_api::routes::ADMIN_TASK_CANCEL,
] {
assert!(
template.contains('{'),
"{template} is in the template list but has no placeholder"
);
let mounted = axum_path(template);
assert!(
!mounted.contains('{') && !mounted.contains('}'),
"{template} would be mounted as {mounted}, whose braces axum reads as a literal segment"
);
assert!(
mounted.contains(':'),
"{template} lost its placeholder entirely and would match one path only"
);
}
}
#[tokio::test]
async fn an_unknown_response_id_gets_the_handler_not_a_bare_404() {
let app = test_app();
let (status, body) = get_json(&app, "/v1/responses/resp_nonexistent").await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert!(
!body.is_null(),
"empty body means axum never matched the route, so the id was read as a literal segment"
);
}
#[tokio::test]
async fn the_task_list_starts_empty_rather_than_absent() {
let app = test_app();
let (status, body) = get_json(&app, ferrox_api::routes::ADMIN_TASKS).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["tasks"].as_array().unwrap().len(), 0);
}
async fn post_json_uri(
app: &Router,
uri: &str,
body: serde_json::Value,
) -> (StatusCode, serde_json::Value) {
use http_body_util::BodyExt;
use tower::ServiceExt;
let response = app
.clone()
.oneshot(
axum::http::Request::builder()
.method("POST")
.uri(uri)
.header("content-type", "application/json")
.body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let bytes = response.into_body().collect().await.unwrap().to_bytes();
let json = serde_json::from_slice(&bytes).unwrap_or(serde_json::json!({}));
(status, json)
}
async fn post_json(app: &Router, body: serde_json::Value) -> serde_json::Value {
post_json_uri(app, "/v1/chat/completions", body).await.1
}
#[tokio::test]
async fn stats_says_what_the_engine_is_using_and_which_quantity_that_is() {
let app = test_app();
let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
assert_eq!(status, StatusCode::OK);
let memory = &body["memory"];
if memory.is_null() {
return;
}
assert!(
memory["bytes"].as_u64().is_some_and(|b| b > 0),
"a read that produced a zero is a broken read, not an idle \
engine: {memory}"
);
assert!(
["pss", "rss"].contains(&memory["kind"].as_str().unwrap_or("")),
"the quantity must travel with the number: {memory}"
);
}
#[tokio::test]
async fn stats_reports_a_pool_it_does_not_have_as_absent_and_not_as_zero() {
let app = test_app();
let (status, body) = get_json(&app, ferrox_api::routes::V1_STATS).await;
assert_eq!(status, StatusCode::OK);
for pool in ["kv_pages", "window_slots", "state_slots"] {
assert!(
body["pools"][pool].is_null(),
"{pool} must be null rather than a zero row: {}",
body["pools"]
);
}
}
#[tokio::test]
async fn a_messages_response_states_the_id_that_v1_cancel_takes() {
use http_body_util::BodyExt;
use tower::ServiceExt;
let app = test_app();
let send = |body: serde_json::Value| {
let app = app.clone();
async move {
app.oneshot(
axum::http::Request::builder()
.method("POST")
.uri(ferrox_api::routes::V1_MESSAGES)
.header("content-type", "application/json")
.body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap()
}
};
let ok = send(serde_json::json!({
"model": "test",
"max_tokens": 1,
"messages": [{"role": "user", "content": "hi"}],
}))
.await;
assert_eq!(ok.status(), StatusCode::OK);
let id = ok
.headers()
.get("request-id")
.expect("a served message names its id")
.to_str()
.unwrap()
.to_string();
assert!(!id.is_empty());
let bad = send(serde_json::json!({"model": "test"})).await;
assert!(bad.status().is_client_error());
let other = bad.headers().get("request-id").expect("errors too");
assert_ne!(other.to_str().unwrap(), id);
let _ = bad.into_body().collect().await.unwrap();
}
#[tokio::test]
async fn a_request_that_arrives_mid_rebuild_is_refused_and_admitted_again_after() {
let state = Arc::new(test_state(
test_model_full_byte_vocab(),
ResponseCache::new(1000, Duration::from_secs(3600)),
));
let app = test_app_with_state(Arc::clone(&state));
let body = serde_json::json!({
"model": "test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 1,
});
state
.maintenance
.lock()
.unwrap()
.begin_rebuild()
.expect("a fresh server is serving, so the rebuild starts");
let (status, refused) = post_json_uri(&app, "/v1/chat/completions", body.clone()).await;
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(refused["error"]["type"], "cache_rebuilding");
state.maintenance.lock().unwrap().finish_rebuild(true);
let (status, _) = post_json_uri(&app, "/v1/chat/completions", body).await;
assert_eq!(
status,
StatusCode::OK,
"the gate reopens; a rebuild is not a latch"
);
}
#[tokio::test]
async fn cancelling_an_id_that_is_not_generating_is_a_404_that_says_so() {
let app = test_app();
let (status, body) = post_json_uri(
&app,
ferrox_api::routes::V1_CANCEL,
serde_json::json!({ "request_id": "chatcmpl-never-issued" }),
)
.await;
assert_eq!(status, StatusCode::NOT_FOUND);
assert_eq!(body["cancelled"], serde_json::json!(false));
assert_eq!(body["request_id"], "chatcmpl-never-issued");
assert!(
body["detail"].as_str().is_some_and(|d| !d.is_empty()),
"the verdict must carry a human reason: {body}"
);
}
#[tokio::test]
async fn cancelling_a_live_generation_signals_its_token_and_answers_200() {
let state = Arc::new(test_state(
test_model_full_byte_vocab(),
ResponseCache::new(1000, Duration::from_secs(3600)),
));
let app = test_app_with_state(Arc::clone(&state));
let (token, _guard) = state.cancels.register("chatcmpl-live");
let (status, before) = get_json(&app, ferrox_api::routes::ADMIN_STATS).await;
assert_eq!(status, StatusCode::OK);
assert_eq!(before["generating_now"], serde_json::json!(1));
let (status, body) = post_json_uri(
&app,
ferrox_api::routes::V1_CANCEL,
serde_json::json!({ "request_id": "chatcmpl-live" }),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(body["cancelled"], serde_json::json!(true));
assert!(
token.is_cancelled(),
"the endpoint answered ok without setting the flag the decode loop reads"
);
}
#[tokio::test]
async fn tokenize_detokenize_roundtrip_and_embeddings_mean() {
let app = test_app();
let (status, tok) =
post_json_uri(&app, "/v1/tokenize", serde_json::json!({ "prompt": "Hi" })).await;
assert_eq!(status, StatusCode::OK);
let tokens = tok["tokens"].as_array().unwrap();
assert_eq!(tok["count"], tokens.len());
assert!(!tokens.is_empty());
let (status, detok) = post_json_uri(
&app,
"/v1/detokenize",
serde_json::json!({ "tokens": tokens }),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(detok["text"], "Hi");
let (status, emb) = post_json_uri(
&app,
"/v1/embeddings",
serde_json::json!({
"input": "Hi",
"embedding_type": "mean"
}),
)
.await;
assert_eq!(status, StatusCode::OK);
let vec = emb["data"][0]["embedding"].as_array().unwrap();
assert!(!vec.is_empty());
assert!(vec.iter().all(|v| v.as_f64().is_some()));
}
#[tokio::test]
async fn the_decoder_path_still_refuses_a_pooling_it_cannot_mean() {
let app = test_app();
let (status, body) = post_json_uri(
&app,
"/v1/embeddings",
serde_json::json!({ "input": "Hi", "embedding_type": "cls" }),
)
.await;
assert_eq!(status, StatusCode::BAD_REQUEST);
let msg = body["error"]["message"].as_str().unwrap();
assert!(msg.contains("mean") && msg.contains("last"), "{msg}");
}
#[tokio::test]
#[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
async fn a_real_embedding_model_serves_v1_embeddings() {
let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../../models/bge-small-en-v1.5-q8_0.gguf");
if !path.exists() {
eprintln!("SKIP: {} not present", path.display());
return;
}
let encoder = ferrox_models::EmbeddingModel::from_gguf_path(&path).expect("load bge");
let mut state = test_state(
test_model_full_byte_vocab(),
ResponseCache::new(1000, Duration::from_secs(3600)),
);
state.embedding = Some(Arc::new(encoder));
let app = test_app_with_state(Arc::new(state));
let (status, body) = post_json_uri(
&app,
"/v1/embeddings",
serde_json::json!({ "input": ["Hello world", "a second input"] }),
)
.await;
assert_eq!(status, StatusCode::OK, "{body}");
assert_eq!(body["model"], "bge-small-en-v1.5");
let data = body["data"].as_array().unwrap();
assert_eq!(data.len(), 2);
for (i, row) in data.iter().enumerate() {
assert_eq!(row["index"], i);
let v: Vec<f64> = row["embedding"]
.as_array()
.unwrap()
.iter()
.map(|x| x.as_f64().unwrap())
.collect();
assert_eq!(v.len(), 384, "the encoder\'s width, not the decoder\'s");
let norm = v.iter().map(|x| x * x).sum::<f64>().sqrt();
assert!((norm - 1.0).abs() < 1e-4, "not L2-normalized: {norm}");
}
assert!(body["usage"]["prompt_tokens"].as_u64().unwrap() >= 4 + 2);
let (status, mean) = post_json_uri(
&app,
"/v1/embeddings",
serde_json::json!({ "input": "Hello world", "embedding_type": "mean" }),
)
.await;
assert_eq!(status, StatusCode::OK);
assert_ne!(mean["data"][0]["embedding"], data[0]["embedding"]);
}
#[tokio::test]
#[ignore = "needs models/bge-small-en-v1.5-q8_0.gguf"]
async fn an_encoder_can_be_the_loaded_model() {
let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../../models/bge-small-en-v1.5-q8_0.gguf");
if !path.exists() {
eprintln!("SKIP: {} not present", path.display());
return;
}
let loaded = model::load_from_path(path.to_str().unwrap()).expect("load bge as the model");
assert!(
matches!(loaded, model::LoadedModel::Encoder(_)),
"an encoder-only GGUF reached a decoder loader"
);
let (loaded, batcher, ceiling) = activate_loaded_model(loaded, true, None, None);
assert!(
matches!(loaded, Loaded::Encoder(_)),
"the encoder did not stay an encoder through activation"
);
assert!(
batcher.is_none() && ceiling.is_none(),
"an encoder was given a decode batcher or a KV ceiling it has no use for"
);
let state = test_state(
test_model_full_byte_vocab(),
ResponseCache::new(1000, Duration::from_secs(3600)),
);
state.swap_active(Some(Arc::new(ActiveModel {
id: None,
loaded,
batcher,
ceiling,
})));
let app = test_app_with_state(Arc::new(state));
let (status, body) = post_json_uri(
&app,
"/v1/embeddings",
serde_json::json!({ "input": "Hello world" }),
)
.await;
assert_eq!(status, StatusCode::OK, "{body}");
assert_eq!(body["model"], "bge-small-en-v1.5");
let v = body["data"][0]["embedding"].as_array().unwrap();
assert_eq!(v.len(), 384, "the encoder's width, not the decoder's");
let (status, body) = post_json_uri(
&app,
"/v1/chat/completions",
serde_json::json!({
"model": "bge-small-en-v1.5",
"messages": [{"role": "user", "content": "hi"}],
}),
)
.await;
assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}");
let msg = body["error"]["message"].as_str().unwrap();
for fact in [
"bge-small-en-v1.5",
"bert",
"embedding model",
"/v1/embeddings",
] {
assert!(msg.contains(fact), "the refusal does not say {fact}: {msg}");
}
let (status, models) = get_json(&app, ferrox_api::routes::V1_MODELS).await;
assert_eq!(status, StatusCode::OK);
let entry = &models["data"][0];
assert_eq!(entry["id"], "bge-small-en-v1.5");
assert_eq!(entry["ferrox_model_kind"], "embedding");
assert_eq!(entry["ferrox_tokenizer"], "gguf-wordpiece");
assert_eq!(entry["ferrox_n_embd"], 384);
assert_eq!(entry["ferrox_pooling"], "CLS");
assert_eq!(
entry["ferrox_endpoints"],
serde_json::json!(["/v1/embeddings"])
);
assert!(entry.get("supported_reasoning_efforts").is_none());
let (status, health) = get_json(&app, ferrox_api::routes::HEALTH).await;
assert_eq!(status, StatusCode::OK, "an encoder is a loaded model");
assert_eq!(health["model"]["id"], "bge-small-en-v1.5");
assert_eq!(health["model"]["synthetic_weights"], false);
let weights = health["capabilities"]
.as_array()
.unwrap()
.iter()
.find(|c| c["id"] == ferrox_api::health::capability::REAL_WEIGHTS)
.expect("a real-weights capability row");
let detail = weights["detail"].as_str().unwrap_or_default();
assert!(detail.contains("ENCODER"), "{detail}");
}
#[tokio::test]
async fn metrics_exposes_expert_store_counters_when_streaming_is_active() {
use http_body_util::BodyExt;
use tower::ServiceExt;
let fixture = concat!(
"../ferrox-models/tests/fixtures/",
"ferrox_real_moe_test.gguf"
);
let fixture = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(fixture);
let decoder = Decoder::from_gguf_with_expert_cache(
&fixture,
ferrox_models::config::test_moe_fixture(),
Some(1024 * 1024),
)
.expect("MoE fixture must load store-backed");
let mut caches: Vec<ferrox_core::cache::KvCache> = decoder
.layers
.iter()
.map(|_| {
ferrox_core::cache::KvCache::new(decoder.config.n_kv_heads, decoder.config.head_dim)
})
.collect();
decoder.forward_token(1, 0, &mut caches);
let model = Model::Gguf(GgufModel {
decoder: Arc::new(decoder),
tokenizer: Arc::new(ServerTokenizer::Byte),
stop_tokens: StopTokens::default(),
bos_id: None,
is_synthetic: false,
chat_template: chat_template::PromptTemplate::plain(),
});
let state = Arc::new(test_state(
model,
ResponseCache::new(16, Duration::from_secs(60)),
));
let app = Router::new()
.route("/metrics", axum::routing::get(metrics))
.route("/v1/chat/completions", post(chat_completions))
.with_state(state);
let fetch_metrics = |app: Router| async move {
let resp = app
.oneshot(
axum::http::Request::builder()
.method("GET")
.uri("/metrics")
.body(axum::body::Body::empty())
.unwrap(),
)
.await
.unwrap();
let bytes = resp.into_body().collect().await.unwrap().to_bytes();
String::from_utf8(bytes.to_vec()).unwrap()
};
let after = fetch_metrics(app.clone()).await;
assert!(
after.contains("ferrox_expert_cache_misses_total"),
"streaming model must expose expert-cache metrics: {after}"
);
let misses: u64 = after
.lines()
.find(|l| l.starts_with("ferrox_expert_cache_misses_total"))
.and_then(|l| l.split_whitespace().nth(1))
.and_then(|v| v.parse().ok())
.expect("misses metric line must parse");
assert!(
misses > 0,
"decode must have read experts through the store: {after}"
);
}
fn weather_tool() -> serde_json::Value {
serde_json::json!({
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather for a location.",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"]
}
}
})
}
fn weather_tool_def() -> ToolDef {
ToolDef {
kind: "function".to_string(),
function: ToolFunctionDef {
name: "get_weather".to_string(),
description: Some("Get the current weather for a location.".to_string()),
parameters: Some(serde_json::json!({
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"]
})),
},
}
}
#[test]
fn tool_preamble_mentions_every_tool_name_and_description() {
let preamble = tool_preamble(&[weather_tool_def()]);
assert!(preamble.contains("get_weather"));
assert!(preamble.contains("Get the current weather for a location."));
assert!(preamble.contains("<tool_call>"));
assert!(preamble.contains("</tool_call>"));
}
#[test]
fn a_real_marker_becomes_a_structured_tool_call() {
let text = "sure, let me check.<tool_call>{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Paris\"}}</tool_call>";
let (message, finish) = build_response_message(
text.to_string(),
&[weather_tool_def()],
output::OutputPosture::for_model("test-model"),
"stop",
);
assert_eq!(finish, "tool_calls");
let calls = message.tool_calls.expect("must carry a tool call");
assert_eq!(calls[0].function.name, "get_weather");
let parsed: serde_json::Value = serde_json::from_str(&calls[0].function.arguments).unwrap();
assert_eq!(parsed["location"], "Paris");
}
#[test]
fn a_plain_answer_is_not_promoted_to_a_tool_call() {
let (message, finish) = build_response_message(
"just an answer".to_string(),
&[weather_tool_def()],
output::OutputPosture::for_model("test-model"),
"stop",
);
assert_eq!(finish, "stop");
assert!(message.tool_calls.is_none());
assert_eq!(message.content.as_deref(), Some("just an answer"));
}
#[test]
fn a_malformed_payload_is_not_a_tool_call() {
let (message, finish) = build_response_message(
"<tool_call>not valid json at all</tool_call>".to_string(),
&[weather_tool_def()],
output::OutputPosture::for_model("test-model"),
"stop",
);
assert_eq!(finish, "stop");
assert!(message.tool_calls.is_none());
}
#[test]
fn a_tool_that_was_never_offered_is_not_returned() {
let (message, finish) = build_response_message(
"<tool_call>{\"name\": \"ping\", \"arguments\": {}}</tool_call>".to_string(),
&[weather_tool_def()],
output::OutputPosture::for_model("test-model"),
"stop",
);
assert_eq!(finish, "stop");
assert!(message.tool_calls.is_none());
}
#[test]
fn marker_text_with_no_tools_offered_stays_content() {
let (message, finish) = build_response_message(
"<tool_call>{\"name\": \"get_weather\", \"arguments\": {}}</tool_call>".to_string(),
&[],
output::OutputPosture::for_model("test-model"),
"stop",
);
assert_eq!(finish, "stop");
assert!(message.tool_calls.is_none());
assert!(message.content.is_some());
}
#[test]
fn a_streamed_call_opens_then_delivers_its_arguments_in_pieces() {
let opened = std::cell::Cell::new(0usize);
let mut parser = crate::policy::parser::ToolCallParser::new(
crate::policy::parser::ToolCallFormat::Qwen3Coder,
vec![
crate::policy::parser::tool_call::ToolSchema::with_parameters(
"write_file",
serde_json::json!({"type": "object", "properties": {
"path": {"type": "string"},
"contents": {"type": "string"}
}}),
),
],
);
let wire = "<tool_call><function=write_file>\
<parameter=path>\n/tmp/x\n</parameter>\
<parameter=contents>\nhello world\n</parameter>\
</function></tool_call>";
let mut deltas = Vec::new();
let mut text = String::new();
for piece in wire.as_bytes().chunks(7) {
let chunk = String::from_utf8_lossy(piece).into_owned();
let (more_text, more) = tool_call_deltas(parser.push(&chunk), &opened);
text.push_str(&more_text);
deltas.extend(more);
}
let (more_text, more) = tool_call_deltas(parser.finish(), &opened);
text.push_str(&more_text);
deltas.extend(more);
assert_eq!(opened.get(), 1, "one call opened");
assert!(text.is_empty(), "the markers are not content: {text:?}");
let first = &deltas[0];
assert_eq!(first.index, 0);
assert_eq!(first.id.as_deref(), Some("call_0"));
assert_eq!(first.kind, Some("function"));
assert_eq!(first.function.name.as_deref(), Some("write_file"));
let joined: String = deltas
.iter()
.filter_map(|d| d.function.arguments.clone())
.collect();
let parsed: serde_json::Value =
serde_json::from_str(&joined).expect("the fragments concatenate to valid JSON");
assert_eq!(parsed["path"], serde_json::json!("/tmp/x"));
assert_eq!(parsed["contents"], serde_json::json!("hello world"));
assert!(
deltas.len() >= 3,
"the arguments arrived in pieces, not whole: {}",
deltas.len()
);
assert!(
deltas[1..].iter().all(|d| d.function.name.is_none()),
"only the opening delta carries identity"
);
}
#[test]
fn text_around_a_streamed_call_is_still_content() {
let opened = std::cell::Cell::new(0usize);
let mut parser = crate::policy::parser::ToolCallParser::new(
crate::policy::parser::ToolCallFormat::Qwen25,
vec![crate::policy::parser::tool_call::ToolSchema::new(
"get_weather",
)],
);
let wire = "let me check. <tool_call>{\"name\": \"get_weather\", \
\"arguments\": {}}</tool_call> done";
let mut text = String::new();
for piece in wire.as_bytes().chunks(5) {
let chunk = String::from_utf8_lossy(piece).into_owned();
let (more, _) = tool_call_deltas(parser.push(&chunk), &opened);
text.push_str(&more);
}
let (more, _) = tool_call_deltas(parser.finish(), &opened);
text.push_str(&more);
assert_eq!(opened.get(), 1);
assert!(text.starts_with("let me check. "), "{text:?}");
assert!(text.ends_with(" done"), "{text:?}");
assert!(!text.contains("<tool_call>"), "markers leaked: {text:?}");
}
#[test]
fn a_reasoning_block_is_split_out_of_the_answer() {
let (message, finish) = build_response_message(
"<think>weighing it up</think>The answer is 4.".to_string(),
&[],
output::OutputPosture::for_model("Qwen3-8B"),
"stop",
);
assert_eq!(finish, "stop");
assert_eq!(message.content.as_deref(), Some("The answer is 4."));
assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
}
#[test]
fn a_non_reasoning_model_keeps_a_literal_marker_in_its_answer() {
let (message, _) = build_response_message(
"Use the <think> tag like this.".to_string(),
&[],
output::OutputPosture::for_model("llama-3.1-8b"),
"stop",
);
assert_eq!(
message.content.as_deref(),
Some("Use the <think> tag like this.")
);
assert!(message.reasoning_content.is_none());
}
#[tokio::test]
async fn a_request_with_no_tools_or_session_behaves_exactly_as_before() {
let app = test_app();
let body = serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 4,
"temperature": 0,
});
let resp = post_json(&app, body).await;
let message = &resp["choices"][0]["message"];
assert!(message["content"].is_string());
assert!(message.get("tool_calls").is_none());
assert_eq!(resp["choices"][0]["finish_reason"], "length");
assert_eq!(resp["usage"]["completion_tokens"], 4);
assert_eq!(
resp["usage"]["total_tokens"],
resp["usage"]["prompt_tokens"].as_u64().unwrap() + 4
);
}
async fn get_json(app: &Router, uri: &str) -> (StatusCode, serde_json::Value) {
use http_body_util::BodyExt;
use tower::ServiceExt;
let response = app
.clone()
.oneshot(
axum::http::Request::builder()
.method("GET")
.uri(uri)
.body(axum::body::Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let bytes = response.into_body().collect().await.unwrap().to_bytes();
(status, serde_json::from_slice(&bytes).unwrap())
}
#[tokio::test]
async fn health_answers_a_capability_handshake_not_a_boolean() {
let app = test_app();
let (status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
assert_eq!(status, StatusCode::OK);
let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
assert_eq!(health.state, ferrox_api::HealthState::Ready);
assert!(health.pid > 0);
assert!(health.server_time_unix_ms > 0);
assert_eq!(health.last_request_age_seconds, None);
for id in [
ferrox_api::health::capability::CPU,
ferrox_api::health::capability::METAL,
ferrox_api::health::capability::CUDA,
ferrox_api::health::capability::REAL_WEIGHTS,
ferrox_api::health::capability::CONTINUOUS_BATCHING,
] {
let cap = health
.capability(id)
.unwrap_or_else(|| panic!("{id} missing"));
assert!(!cap.reason.is_empty(), "{cap:?}");
assert!(!cap.detail.is_empty(), "{cap:?}");
}
let weights = health
.capability(ferrox_api::health::capability::REAL_WEIGHTS)
.unwrap();
assert!(!weights.available);
assert_eq!(weights.reason, ferrox_api::health::reason::MODEL_NOT_LOADED);
assert!(health.model.as_ref().unwrap().synthetic_weights);
}
#[tokio::test]
async fn health_vouches_for_liveness_after_a_request_has_been_served() {
let app = test_app();
let _ = post_json(
&app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}"}],
"max_tokens": 1,
"temperature": 0,
}),
)
.await;
let (_status, body) = get_json(&app, ferrox_api::routes::HEALTH).await;
let health: ferrox_api::HealthResponse = serde_json::from_value(body).unwrap();
let age = health
.last_request_age_seconds
.expect("a served request is evidence of liveness");
assert!((0.0..5.0).contains(&age), "implausible age {age}");
}
async fn post_sse_chunks(app: &Router, body: serde_json::Value) -> Vec<serde_json::Value> {
use http_body_util::BodyExt;
use tower::ServiceExt;
let response = app
.clone()
.oneshot(
axum::http::Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("content-type", "application/json")
.body(axum::body::Body::from(serde_json::to_vec(&body).unwrap()))
.unwrap(),
)
.await
.unwrap();
let bytes = response.into_body().collect().await.unwrap().to_bytes();
String::from_utf8(bytes.to_vec())
.unwrap()
.lines()
.filter_map(|line| line.strip_prefix("data: "))
.filter(|payload| *payload != "[DONE]")
.map(|payload| serde_json::from_str(payload).unwrap())
.collect()
}
#[tokio::test]
async fn a_stream_states_its_request_id_once_in_the_first_chunk() {
let app = test_app();
let chunks = post_sse_chunks(
&app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 4,
"temperature": 0,
"stream": true,
}),
)
.await;
assert!(!chunks.is_empty());
let request_id = chunks[0]["request_id"]
.as_str()
.expect("the first chunk names the request")
.to_string();
assert!(request_id.starts_with("chatcmpl-"), "{request_id}");
for (i, chunk) in chunks.iter().enumerate().skip(1) {
assert!(
chunk.get("request_id").is_none(),
"chunk {i} repeats request_id"
);
}
for chunk in &chunks {
assert_eq!(chunk["id"], serde_json::json!(request_id));
}
let other = post_sse_chunks(
&app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 4,
"temperature": 0,
"stream": true,
}),
)
.await;
assert_ne!(
other[0]["request_id"].as_str().unwrap(),
request_id,
"two concurrent chats must not share an id"
);
}
#[tokio::test]
async fn a_non_streamed_response_names_the_same_request_id_as_its_completion_id() {
let app = test_app();
let resp = post_json(
&app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 2,
"temperature": 0,
}),
)
.await;
assert_eq!(resp["id"], resp["request_id"]);
assert!(resp["request_id"]
.as_str()
.unwrap()
.starts_with("chatcmpl-"));
}
#[tokio::test]
async fn usage_carries_separate_prefill_and_decode_timings() {
let app = test_app();
let resp = post_json(
&app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 4,
"temperature": 0,
}),
)
.await;
let usage = &resp["usage"];
assert!(usage["prompt_eval_duration_ms"].is_number(), "{usage}");
assert!(usage["generation_duration_ms"].is_number(), "{usage}");
assert!(usage["time_to_first_token_ms"].is_number(), "{usage}");
assert!(usage["predicted_per_second"].is_number(), "{usage}");
assert!(usage.get("cached_tokens").is_none(), "{usage}");
}
#[tokio::test]
async fn a_tools_request_with_no_marker_in_the_output_falls_back_to_plain_content() {
let app = test_app();
let body = serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 4,
"temperature": 0,
"tools": [weather_tool()],
});
let resp = post_json(&app, body).await;
let message = &resp["choices"][0]["message"];
assert!(
message["content"].is_string(),
"must fall back to plain content when no real tool-call marker is present: {resp:?}"
);
assert!(message.get("tool_calls").is_none());
assert_eq!(resp["choices"][0]["finish_reason"], "length");
}
#[tokio::test]
async fn a_cache_hit_reports_the_original_finish_reason_and_usage() {
let app = test_app();
let body = serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}"}],
"max_tokens": 3,
"temperature": 0,
});
let first = post_json(&app, body.clone()).await;
assert_eq!(first["ferrox_cache"], "miss");
let second = post_json(&app, body).await;
assert_eq!(second["ferrox_cache"], "hit");
assert_eq!(
first["choices"][0]["message"]["content"],
second["choices"][0]["message"]["content"]
);
assert_eq!(
first["choices"][0]["finish_reason"],
second["choices"][0]["finish_reason"]
);
assert_eq!(first["usage"], second["usage"]);
assert_eq!(second["usage"]["completion_tokens"], 3);
}
#[tokio::test]
async fn session_reuse_produces_the_same_output_as_manually_resending_full_history() {
let session_app = test_app();
let manual_app = test_app();
let turn1 = post_json(
&session_app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"session_id": "s1",
"max_tokens": 5,
"temperature": 0,
}),
)
.await;
let reply1 = turn1["choices"][0]["message"]["content"]
.as_str()
.unwrap()
.to_string();
let manual_turn1 = post_json(
&manual_app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{1}\u{2}\u{3}"}],
"max_tokens": 5,
"temperature": 0,
}),
)
.await;
assert_eq!(
manual_turn1["choices"][0]["message"]["content"]
.as_str()
.unwrap(),
reply1
);
let turn2 = post_json(
&session_app,
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "\u{4}\u{5}"}],
"session_id": "s1",
"max_tokens": 5,
"temperature": 0,
}),
)
.await;
let reply2 = turn2["choices"][0]["message"]["content"]
.as_str()
.unwrap()
.to_string();
let manual_turn2 = post_json(
&manual_app,
serde_json::json!({
"model": "m",
"messages": [
{"role": "user", "content": "\u{1}\u{2}\u{3}"},
{"role": "assistant", "content": reply1},
{"role": "user", "content": "\u{4}\u{5}"},
],
"max_tokens": 5,
"temperature": 0,
}),
)
.await;
assert_eq!(
manual_turn2["choices"][0]["message"]["content"]
.as_str()
.unwrap(),
reply2,
"resuming a session must produce identical output to manually resending the full history"
);
}
#[test]
fn lock_cache_recovers_from_a_poisoned_mutex() {
let cache = Arc::new(Mutex::new(ResponseCache::new(10, Duration::from_secs(60))));
let poison_cache = Arc::clone(&cache);
let _ = std::thread::spawn(move || {
let _guard = poison_cache.lock().unwrap();
panic!("simulated panic while holding the lock");
})
.join();
let recovered = lock_cache(&cache);
assert_eq!(recovered.stats().entries, 0);
}
#[test]
fn is_cacheable_true_for_greedy_or_seeded_requests() {
let mut req_body = serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
});
let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
assert!(
req.is_cacheable(),
"default (temperature 0) must be cacheable"
);
req_body["temperature"] = serde_json::json!(0.8);
let req: ChatCompletionRequest = serde_json::from_value(req_body.clone()).unwrap();
assert!(
!req.is_cacheable(),
"unseeded sampling must never be cacheable"
);
req_body["seed"] = serde_json::json!(42);
let req: ChatCompletionRequest = serde_json::from_value(req_body).unwrap();
assert!(
req.is_cacheable(),
"sampling with an explicit seed is deterministic and must be cacheable"
);
}
const GRADED: &str = "{% if reasoning_effort %}\
{% if reasoning_effort not in ['low','medium','high'] %}\
{{ raise_exception('unsupported effort') }}\
{% endif %}E:{{ reasoning_effort }}|{% endif %}\
{% if enable_thinking %}THINK|{% endif %}{{ messages[0].content }}";
fn graded_template() -> chat_template::PromptTemplate {
chat_template::PromptTemplate::from_gguf_metadata(
Some(GRADED),
Some("qwen3"),
false,
true,
None,
None,
)
}
fn chat_request(value: serde_json::Value) -> ChatCompletionRequest {
serde_json::from_value(value).expect("request")
}
#[test]
fn a_grammar_on_the_chat_wire_reaches_the_generation_params() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"grammar": "root ::= \"a\"+",
}));
req.validate_supported_fields()
.expect("a valid grammar is a valid request");
let params = req
.generation_params()
.expect("a valid grammar compiles at params time too");
assert!(
params.grammar.is_some(),
"the grammar was dropped between the wire and the sampler"
);
assert!(
params.needs_vocab_logits(),
"a grammar request that may fold lm_head into a GPU argmax is \
a grammar request served unconstrained"
);
let plain = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
}));
assert!(plain.generation_params().unwrap().grammar.is_none());
}
fn tool_request(tool_choice: serde_json::Value) -> ChatCompletionRequest {
chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "weather in Rome?"}],
"tools": [weather_tool()],
"tool_choice": tool_choice,
}))
}
#[test]
fn a_forced_tool_choice_puts_a_grammar_on_the_generation_params() {
for choice in [
serde_json::json!("required"),
serde_json::json!({"type": "function", "function": {"name": "get_weather"}}),
] {
let req = tool_request(choice.clone());
req.validate_supported_fields()
.unwrap_or_else(|e| panic!("{choice} is a valid request: {e:?}"));
let params = req
.generation_params_for_template(&graded_template(), "Qwen3-8B")
.unwrap_or_else(|e| panic!("{choice} compiles: {e:?}"));
let grammar = params
.grammar
.as_ref()
.unwrap_or_else(|| panic!("{choice} was accepted and then not enforced"));
assert!(
grammar.is_awaiting_trigger(),
"the model must be free to think before it calls"
);
assert!(
!grammar.allows_eog(),
"{choice} must not be able to end the turn without a call"
);
assert!(
params.needs_vocab_logits(),
"{choice} would let a backend return a token id instead of logits"
);
assert!(
!generate::greedy_gpu_fold_allowed(¶ms),
"{choice} at temperature 0 must still refuse the greedy GPU fold"
);
}
}
#[test]
fn an_unforced_tool_choice_leaves_the_generation_unconstrained() {
for choice in [serde_json::json!("auto"), serde_json::json!("none")] {
let req = tool_request(choice.clone());
req.validate_supported_fields().expect("still supported");
let params = match req.generation_params_for_template(&graded_template(), "Qwen3-8B") {
Ok(p) => p,
Err((status, _)) => panic!("{choice} has no constraint to compile: {status}"),
};
assert!(
params.grammar.is_none(),
"{choice} does not force a call and must not be constrained"
);
}
}
#[test]
fn a_forced_tool_choice_refuses_rather_than_quietly_not_forcing() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"tool_choice": "required",
}));
let (status, _) = req
.validate_supported_fields()
.expect_err("nothing to call");
assert_eq!(status, StatusCode::BAD_REQUEST);
let req =
tool_request(serde_json::json!({"type": "function", "function": {"name": "nope"}}));
let (status, Json(body)) = req.validate_supported_fields().expect_err("no such tool");
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["error"]["param"], "tool_choice");
let req = tool_request(serde_json::json!({"type": "function"}));
let (status, _) = req.validate_supported_fields().expect_err("names nothing");
assert_eq!(status, StatusCode::BAD_REQUEST);
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"tools": [weather_tool()],
"tool_choice": "required",
"grammar": "root ::= \"a\"+",
}));
let (status, _) = req
.validate_supported_fields()
.expect_err("a grammar and a forced call are two constraints");
assert_eq!(status, StatusCode::BAD_REQUEST);
let req = tool_request(serde_json::json!("required"));
let (status, Json(body)) =
match req.generation_params_for_template(&graded_template(), "GLM-4.7") {
Err(e) => e,
Ok(_) => panic!("glm calls are not JSON behind a marker"),
};
assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
assert!(
body["error"]["message"].as_str().unwrap().contains("glm47"),
"{body}"
);
}
#[test]
fn an_unparseable_grammar_on_the_chat_wire_is_a_400() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"grammar": "root ::= \"a",
}));
let (status, Json(body)) = req
.validate_supported_fields()
.expect_err("this does not parse");
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["error"]["param"], "grammar");
assert!(req.generation_params().is_err(), "and again at params time");
}
#[test]
fn response_format_json_schema_names_the_missing_converter() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"response_format": {
"type": "json_schema",
"json_schema": {"name": "x", "schema": {"type": "object"}},
},
}));
let (status, Json(body)) = req
.validate_supported_fields()
.expect_err("not implemented yet");
assert_eq!(status, StatusCode::NOT_IMPLEMENTED);
let message = body["error"]["message"].as_str().expect("a message");
assert!(
message.contains("JSON schema") && message.contains("grammar"),
"the refusal must say what is missing: {message}"
);
}
#[test]
fn an_omitted_output_budget_is_a_whole_answer_not_sixteen_tokens() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
}));
assert_eq!(req.max_tokens, DEFAULT_CHAT_MAX_TOKENS);
}
#[test]
fn min_p_reaches_the_sampler_from_the_chat_wire() {
let asked = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"min_p": 0.07,
}));
assert_eq!(asked.sampling_params().min_p, 0.07);
let silent = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
}));
assert_eq!(
silent.sampling_params().min_p,
0.0,
"an unset min_p must be off, not llama.cpp's CLI default"
);
}
#[test]
fn no_sampler_knob_is_missing_from_the_cache_key() {
let base = serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"seed": 1,
});
let baseline = chat_request(base.clone()).cache_key("prompt");
for (knob, value) in [
("temperature", serde_json::json!(0.5)),
("top_p", serde_json::json!(0.9)),
("min_p", serde_json::json!(0.05)),
("top_k", serde_json::json!(40)),
("repetition_penalty", serde_json::json!(1.1)),
("presence_penalty", serde_json::json!(0.3)),
("frequency_penalty", serde_json::json!(0.3)),
] {
let mut body = base.clone();
body[knob] = value;
assert_ne!(
chat_request(body).cache_key("prompt"),
baseline,
"`{knob}` is not in the cache key: two requests differing \
only in it would share one cached answer"
);
}
}
#[test]
fn an_explicit_zero_output_budget_is_a_client_error() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 0,
}));
let (status, body) = req.validate_supported_fields().expect_err("rejected");
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["error"]["param"], serde_json::json!("max_tokens"));
}
#[test]
fn a_request_can_turn_thinking_off() {
let template = graded_template();
for body in [
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"reasoning_effort": "none",
}),
serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"thinking": {"type": "disabled"},
}),
] {
let kwargs = chat_request(body).resolve_template_kwargs(&template);
assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
assert_eq!(kwargs["thinking_mode"], serde_json::json!("disabled"));
assert!(!kwargs.contains_key("reasoning_effort"));
}
}
#[test]
fn a_disabled_switch_beats_an_effort_in_the_same_request() {
let template = graded_template();
let kwargs = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"reasoning_effort": "high",
"thinking": {"type": "disabled"},
}))
.resolve_template_kwargs(&template);
assert_eq!(kwargs["enable_thinking"], serde_json::json!(false));
assert!(!kwargs.contains_key("reasoning_effort"));
}
#[test]
fn an_unrecognized_thinking_switch_is_refused_rather_than_read_as_on() {
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"thinking": {"type": "disable"},
}));
let (status, _) = req.validate_supported_fields().expect_err("rejected");
assert_eq!(status, StatusCode::BAD_REQUEST);
}
#[test]
fn an_explicit_template_kwarg_stands_the_protocol_knobs_down() {
let template = graded_template();
let kwargs = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"reasoning_effort": "none",
"chat_template_kwargs": {"enable_thinking": true},
}))
.resolve_template_kwargs(&template);
assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
}
#[test]
fn an_off_vocabulary_reasoning_effort_is_quantized_rather_than_interpolated() {
let template = graded_template();
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"reasoning_effort": "minimal",
}));
let kwargs = req.resolve_template_kwargs(&template);
assert_eq!(kwargs["reasoning_effort"], serde_json::json!("low"));
let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
assert!(prompt.starts_with("E:low|"), "{prompt}");
}
#[test]
fn an_effort_with_no_near_gear_is_dropped_so_the_template_default_applies() {
let template = graded_template();
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": {"reasoning_effort": "none"},
}));
let kwargs = req.resolve_template_kwargs(&template);
assert!(!kwargs.contains_key("reasoning_effort"));
let prompt = prompt_from_messages(&req.messages, &template, &[], kwargs).expect("renders");
assert_eq!(prompt, "hi");
}
#[test]
fn chat_template_kwargs_wins_over_the_top_level_reasoning_effort() {
let template = graded_template();
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"reasoning_effort": "low",
"chat_template_kwargs": {"reasoning_effort": "high"},
}));
assert_eq!(
req.resolve_template_kwargs(&template)["reasoning_effort"],
serde_json::json!("high")
);
}
#[test]
fn offering_tools_turns_thinking_on_by_itself() {
let template = graded_template();
let quiet = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
}));
assert!(!quiet
.resolve_template_kwargs(&template)
.contains_key("enable_thinking"));
let with_tools = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"tools": [{"type": "function", "function": {"name": "get_weather"}}],
}));
let kwargs = with_tools.resolve_template_kwargs(&template);
assert_eq!(kwargs["enable_thinking"], serde_json::json!(true));
let prompt =
prompt_from_messages(&with_tools.messages, &template, &[], kwargs).expect("renders");
assert!(prompt.starts_with("THINK|"), "{prompt}");
}
#[test]
fn a_prompt_that_opens_the_reasoning_block_makes_the_first_token_reasoning() {
let opener = chat_template::PromptTemplate::from_gguf_metadata(
Some("{{ messages[0].content }}{% if enable_thinking %}<think>{% endif %}"),
Some("qwen3"),
false,
true,
None,
None,
);
let req = chat_request(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"chat_template_kwargs": {"enable_thinking": true},
}));
let kwargs = req.resolve_template_kwargs(&opener);
let prompt = prompt_from_messages(&req.messages, &opener, &[], kwargs).expect("renders");
assert!(prompt.ends_with("<think>"), "{prompt}");
let posture = output::OutputPosture::resolve("Qwen3-8B", &prompt);
let (message, _) = build_response_message(
"weighing it up</think>Paris.".to_string(),
&[],
posture,
"stop",
);
assert_eq!(message.reasoning_content.as_deref(), Some("weighing it up"));
assert_eq!(message.content.as_deref(), Some("Paris."));
let closed = output::OutputPosture::resolve("Qwen3-8B", "<|im_start|>assistant\n");
let (message, _) = build_response_message(
"weighing it up</think>Paris.".to_string(),
&[],
closed,
"stop",
);
assert_eq!(message.reasoning_content, None);
}
#[test]
fn stop_param_accepts_both_single_string_and_array() {
let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"stop": "END",
}))
.unwrap();
assert_eq!(req.stop_sequences(), vec!["END".to_string()]);
let req: ChatCompletionRequest = serde_json::from_value(serde_json::json!({
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"stop": ["A", "B"],
}))
.unwrap();
assert_eq!(req.stop_sequences(), vec!["A".to_string(), "B".to_string()]);
}
#[test]
fn run_generation_rejects_out_of_vocab_tokens_instead_of_panicking() {
let model = test_model();
let result = run_generation(
&model,
"hello",
&greedy_params(4),
None,
None,
None,
None,
None,
);
assert!(matches!(
result,
Err(generate::DecodeError::TokenOutOfVocab { .. })
));
}
#[test]
fn run_generation_honors_an_exhausted_kv_pool_and_maps_it_to_a_503() {
let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
let holder_pool = Arc::clone(&pool);
let holder = std::thread::spawn(move || {
let mut held = ferrox_core::cache::KvCache::with_pool(1, 1, holder_pool, 0).unwrap();
held.push(&[0.0], &[0.0]).unwrap(); std::thread::sleep(Duration::from_millis(200));
drop(held);
});
std::thread::sleep(Duration::from_millis(15));
let config = generate::KvPoolConfig {
pool,
queue_wait: Duration::ZERO,
};
let result = run_generation(
&model,
&prompt,
&greedy_params(4),
Some(&config),
None,
None,
None,
None,
);
assert!(matches!(
result,
Err(generate::DecodeError::KvPoolExhausted)
));
let (status, _body) = decode_error_response(result.unwrap_err());
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
holder.join().unwrap();
}
#[test]
fn a_request_too_big_for_the_whole_pool_is_a_400_not_a_retryable_503() {
let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 1)));
let config = generate::KvPoolConfig {
pool,
queue_wait: Duration::ZERO,
};
let result = run_generation(
&model,
&prompt,
&greedy_params(4),
Some(&config),
None,
None,
None,
None,
);
let err = result.expect_err("one block cannot hold two layers' caches");
assert!(
matches!(
&err,
generate::DecodeError::KvBudgetExceeded { binding, .. }
if *binding == ferrox_models::Ceiling::DeviceMemory.code()
),
"expected an immovable device-memory refusal, got {err:?}"
);
let (status, _body) = decode_error_response(err);
assert_eq!(status, StatusCode::BAD_REQUEST);
}
#[test]
fn decode_error_response_maps_a_full_queue_to_a_retryable_503() {
let (status, Json(body)) = decode_error_response(generate::DecodeError::QueueFull {
queued: 512,
cap: 512,
});
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(body["error"]["retry_after_seconds"], 1);
let message = body["error"]["message"].as_str().expect("message");
assert!(message.contains("512"), "{message}");
}
#[test]
fn decode_error_response_omits_a_retry_hint_for_an_unretryable_error() {
let (_status, Json(body)) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
token: 99,
vocab_size: 32,
});
assert!(
body["error"]["retry_after_seconds"].is_null(),
"retrying a prompt this model cannot tokenize never helps"
);
}
#[test]
fn decode_error_response_maps_token_out_of_vocab_to_bad_request() {
let (status, _body) = decode_error_response(generate::DecodeError::TokenOutOfVocab {
token: 99,
vocab_size: 32,
});
assert_eq!(status, StatusCode::BAD_REQUEST);
}
#[test]
fn run_generation_succeeds_and_releases_blocks_when_the_pool_has_room() {
let model = test_model(); let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 2)));
let config = generate::KvPoolConfig {
pool: pool.clone(),
queue_wait: Duration::ZERO,
};
let (_, finish, _usage) = run_generation(
&model,
&prompt,
&greedy_params(4),
Some(&config),
None,
None,
None,
None,
)
.unwrap();
assert_eq!(finish, FinishReason::Length);
assert_eq!(
pool.lock().unwrap().free_blocks(),
2,
"a completed request must return its blocks to the pool"
);
}
#[tokio::test]
async fn concurrent_requests_against_the_same_model_do_not_interfere() {
let model = Arc::new(test_model());
let prompt = String::from_utf8(vec![1u8, 2]).unwrap();
let mut handles = Vec::new();
for _ in 0..8 {
let model = Arc::clone(&model);
let prompt = prompt.clone();
handles.push(tokio::task::spawn_blocking(move || {
run_generation(
&model,
&prompt,
&greedy_params(6),
None,
None,
None,
None,
None,
)
.unwrap()
}));
}
let mut results = Vec::new();
for h in handles {
results.push(h.await.unwrap());
}
for r in &results[1..] {
assert_eq!(r.0, results[0].0, "decoded chunks must match");
assert_eq!(r.1, results[0].1, "finish reason must match");
assert_eq!(
r.2.prompt_tokens, results[0].2.prompt_tokens,
"prompt token count must match"
);
assert_eq!(
r.2.completion_tokens, results[0].2.completion_tokens,
"completion token count must match"
);
}
}
fn write_safetensors_shard(tensors: &[(String, Vec<usize>, Vec<f32>)]) -> Vec<u8> {
let mut header_entries = Vec::new();
let mut data = Vec::new();
for (name, shape, values) in tensors {
let start = data.len();
for v in values {
data.extend_from_slice(&v.to_le_bytes());
}
let end = data.len();
let shape_str = shape
.iter()
.map(|d| d.to_string())
.collect::<Vec<_>>()
.join(",");
header_entries.push(format!(
"\"{name}\":{{\"dtype\":\"F32\",\"shape\":[{shape_str}],\"data_offsets\":[{start},{end}]}}"
));
}
let header = format!("{{{}}}", header_entries.join(","));
let header_bytes = header.as_bytes();
let mut out = Vec::with_capacity(8 + header_bytes.len() + data.len());
out.extend_from_slice(&(header_bytes.len() as u64).to_le_bytes());
out.extend_from_slice(header_bytes);
out.extend_from_slice(&data);
out
}
fn build_synthetic_kimi_loaded() -> model::KimiLoaded {
use ferrox_models::config::{AttentionKind, KdaConfig, KimiHybridAttention, MlaConfig};
use ferrox_models::kimi_loader::KimiRealHparams;
use ferrox_moe::{GatingFunction, MoeLayerConfig};
let hidden_dim = 8;
let kda_num_heads = 2;
let kda_head_dim = 3;
let kda_proj = kda_num_heads * kda_head_dim;
let conv_kernel = 4;
let dense_intermediate = 5;
let vocab_size = 256;
let mla_num_heads = 1;
let mla_q_lora_rank = 2;
let mla_kv_lora_rank = 2;
let mla_qk_nope_head_dim = 2;
let mla_qk_rope_head_dim = 2;
let mla_v_head_dim = 2;
let model_cfg = ferrox_models::ModelConfig {
name: "synthetic-kimi-server-test",
n_layers: 1,
hidden_dim,
n_heads: 1,
n_kv_heads: 1,
head_dim: 4,
vocab_size,
rope_theta: 10000.0,
rms_norm_eps: 1e-5,
sliding_window: None,
moe: MoeLayerConfig {
expert_weights_scale: 1.0,
n_experts: 1,
n_experts_active: 1,
n_shared_experts: 0,
hidden_dim,
expert_ffn_dim: 4,
gating: GatingFunction::Sigmoid,
norm_topk_prob: true,
expert_group_count: None,
expert_group_used_count: None,
},
n_dense_leading_layers: 1,
attention: AttentionKind::KimiHybrid(KimiHybridAttention {
kda_layers: vec![1],
full_attn_layers: vec![],
mla: MlaConfig {
num_heads: mla_num_heads,
q_lora_rank: mla_q_lora_rank,
kv_lora_rank: mla_kv_lora_rank,
qk_nope_head_dim: mla_qk_nope_head_dim,
qk_rope_head_dim: mla_qk_rope_head_dim,
v_head_dim: mla_v_head_dim,
use_output_gate: true,
rope: None,
},
kda: KdaConfig {
num_heads: kda_num_heads,
head_dim: kda_head_dim,
short_conv_kernel_size: conv_kernel,
gate_lower_bound: -5.0,
use_full_rank_gate: true,
},
}),
rope_freqs: None,
rope_attn_factor: 1.0,
rope_dim: None,
rope_freqs_long: None,
rope_freqs_short: None,
rope_orig_ctx: None,
rope_layout: ferrox_models::config::RopeLayout::Neox,
qk_norm_style: ferrox_models::capability::QkNormStyle::WholeVector,
swa_pattern: None,
swa_dense_first: false,
attn_logit_softcap: None,
final_logit_softcap: None,
embedding_scale: None,
attention_scale: None,
rope_theta_swa: None,
ffn_activation: ferrox_models::config::FfnActivation::Swiglu,
best_effort_fields: &["synthetic test config, not a real preset"],
};
let hp = KimiRealHparams {
hidden_dim,
kda_num_heads,
kda_head_dim,
mla_num_heads,
mla_q_lora_rank,
mla_kv_lora_rank,
mla_qk_nope_head_dim,
mla_qk_rope_head_dim,
mla_v_head_dim,
dense_intermediate_dim: dense_intermediate,
moe_hidden_dim: hidden_dim,
moe_intermediate_dim: 4,
n_experts: 1,
num_shared_experts: 0,
};
let prefix = "language_model.model.layers.0";
let mut tensors: Vec<(String, Vec<usize>, Vec<f32>)> = Vec::new();
let push = |tensors: &mut Vec<(String, Vec<usize>, Vec<f32>)>,
name: String,
shape: Vec<usize>,
n: usize| {
tensors.push((name, shape, vec![0.01f32; n]));
};
push(
&mut tensors,
format!("{prefix}.input_layernorm.weight"),
vec![hidden_dim],
hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.post_attention_layernorm.weight"),
vec![hidden_dim],
hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attention_res_norm.weight"),
vec![hidden_dim],
hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attention_res_proj.weight"),
vec![1, hidden_dim],
hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.mlp_res_norm.weight"),
vec![hidden_dim],
hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.mlp_res_proj.weight"),
vec![1, hidden_dim],
hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.q_proj.weight"),
vec![kda_proj, hidden_dim],
kda_proj * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.k_proj.weight"),
vec![kda_proj, hidden_dim],
kda_proj * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.v_proj.weight"),
vec![kda_proj, hidden_dim],
kda_proj * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.q_conv1d.weight"),
vec![kda_proj, 1, conv_kernel],
kda_proj * conv_kernel,
);
push(
&mut tensors,
format!("{prefix}.self_attn.k_conv1d.weight"),
vec![kda_proj, 1, conv_kernel],
kda_proj * conv_kernel,
);
push(
&mut tensors,
format!("{prefix}.self_attn.v_conv1d.weight"),
vec![kda_proj, 1, conv_kernel],
kda_proj * conv_kernel,
);
push(
&mut tensors,
format!("{prefix}.self_attn.A_log"),
vec![kda_num_heads],
kda_num_heads,
);
push(
&mut tensors,
format!("{prefix}.self_attn.f_a_proj.weight"),
vec![kda_head_dim, hidden_dim],
kda_head_dim * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.f_b_proj.weight"),
vec![kda_proj, kda_head_dim],
kda_proj * kda_head_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.dt_bias"),
vec![kda_proj],
kda_proj,
);
push(
&mut tensors,
format!("{prefix}.self_attn.b_proj.weight"),
vec![kda_num_heads, hidden_dim],
kda_num_heads * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.g_proj.weight"),
vec![kda_proj, hidden_dim],
kda_proj * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.o_norm.weight"),
vec![kda_head_dim],
kda_head_dim,
);
push(
&mut tensors,
format!("{prefix}.self_attn.o_proj.weight"),
vec![hidden_dim, kda_proj],
hidden_dim * kda_proj,
);
push(
&mut tensors,
format!("{prefix}.mlp.gate_proj.weight"),
vec![dense_intermediate, hidden_dim],
dense_intermediate * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.mlp.up_proj.weight"),
vec![dense_intermediate, hidden_dim],
dense_intermediate * hidden_dim,
);
push(
&mut tensors,
format!("{prefix}.mlp.down_proj.weight"),
vec![hidden_dim, dense_intermediate],
hidden_dim * dense_intermediate,
);
push(
&mut tensors,
"language_model.model.embed_tokens.weight".to_string(),
vec![vocab_size, hidden_dim],
vocab_size * hidden_dim,
);
push(
&mut tensors,
"language_model.lm_head.weight".to_string(),
vec![vocab_size, hidden_dim],
vocab_size * hidden_dim,
);
push(
&mut tensors,
"language_model.model.norm.weight".to_string(),
vec![hidden_dim],
hidden_dim,
);
push(
&mut tensors,
"language_model.model.output_attn_res_norm.weight".to_string(),
vec![hidden_dim],
hidden_dim,
);
push(
&mut tensors,
"language_model.model.output_attn_res_proj.weight".to_string(),
vec![1, hidden_dim],
hidden_dim,
);
static FIXTURE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let dir = std::env::temp_dir().join(format!(
"ferrox_server_kimi_e2e_test_{}_{}",
std::process::id(),
FIXTURE.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
));
std::fs::create_dir_all(&dir).unwrap();
let shard_bytes = write_safetensors_shard(&tensors);
std::fs::write(dir.join("shard0.safetensors"), &shard_bytes).unwrap();
let map_entries: Vec<String> = tensors
.iter()
.map(|(name, ..)| format!("\"{name}\":\"shard0.safetensors\""))
.collect();
let index = format!("{{\"weight_map\":{{{}}}}}", map_entries.join(","));
std::fs::write(dir.join("model.safetensors.index.json"), &index).unwrap();
use base64::Engine;
let vocab_lines: Vec<String> = (0..vocab_size as u32)
.map(|b| {
let b64 = base64::engine::general_purpose::STANDARD.encode([b as u8]);
format!("{b64} {b}")
})
.collect();
std::fs::write(dir.join("tiktoken.model"), vocab_lines.join("\n")).unwrap();
let loaded = model::load_kimi_checkpoint_with_config(dir.to_str().unwrap(), model_cfg, hp)
.expect("must load the synthetic Kimi checkpoint end to end");
std::fs::remove_dir_all(&dir).ok();
loaded
}
#[test]
fn kimi_model_serves_real_text_end_to_end_via_run_generation() {
let loaded = build_synthetic_kimi_loaded();
let state = build_app_state(
StartupModels {
loaded: model::LoadedModel::Kimi(loaded),
embedding: None,
},
None,
None,
None,
false,
None,
Arc::new(health::Detection::ready(health::probe_backends())),
);
let active = state.active().expect("a freshly built state has a model");
assert_eq!(active.tokenizer_kind(), "kimi-tiktoken-bpe");
assert!(!active.is_synthetic());
let (_chunks, finish, _usage) = run_generation(
active.generative().unwrap(),
"hi",
&greedy_params(5),
None,
None,
None,
None,
None,
)
.expect("a real Kimi checkpoint must generate without error");
assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
}
#[test]
fn a_grammar_constrains_the_engine_decode_path() {
let loaded = build_synthetic_kimi_loaded();
let state = build_app_state(
StartupModels {
loaded: model::LoadedModel::Kimi(loaded),
embedding: None,
},
None,
None,
None,
false,
None,
Arc::new(health::Detection::ready(health::probe_backends())),
);
let active = state.active().expect("a freshly built state has a model");
let run = |grammar: Option<&str>| {
let mut params = greedy_params(6);
params.grammar = grammar.map(|src| {
Arc::new(
ferrox_models::grammar::Grammar::from_str_with_root(src, "root")
.expect("test grammar parses"),
)
});
run_generation(
active.generative().unwrap(),
"hi",
¶ms,
None,
None,
None,
None,
None,
)
};
let (chunks, _, _) = run(None).expect("the unconstrained run must serve");
let unconstrained = chunks.concat();
assert!(
unconstrained.chars().any(|c| c != 'a'),
"the unconstrained run produced only `a` ({unconstrained:?}), so the \
constrained run below would prove nothing"
);
let (chunks, finish, _) =
run(Some(r#"root ::= "a"+"#)).expect("a grammar this vocabulary can spell must serve");
let constrained = chunks.concat();
assert!(
!constrained.is_empty() && constrained.chars().all(|c| c == 'a'),
"the engine decode path served text its grammar forbids ({constrained:?}): \
the constraint was dropped between `generate_engine` and the sampler"
);
assert!(matches!(finish, FinishReason::Length | FinishReason::Stop));
}
#[test]
fn kv_pool_and_prefix_cache_are_never_consulted_for_a_kimi_model() {
let loaded = build_synthetic_kimi_loaded();
let state = build_app_state(
StartupModels {
loaded: model::LoadedModel::Kimi(loaded),
embedding: None,
},
None,
None,
None,
false,
None,
Arc::new(health::Detection::ready(health::probe_backends())),
);
let pool = Arc::new(Mutex::new(ferrox_core::cache::KvBlockPool::new(64, 4)));
let kv_pool_config = generate::KvPoolConfig {
pool: pool.clone(),
queue_wait: Duration::ZERO,
};
let pc = Mutex::new(PrefixCache::new(4));
run_generation(
state
.active()
.expect("a freshly built state has a model")
.generative()
.unwrap(),
"hi",
&greedy_params(5),
Some(&kv_pool_config),
None,
Some(&pc),
None,
None,
)
.expect("a real Kimi checkpoint must generate without error");
assert_eq!(
pool.lock().unwrap().free_blocks(),
4,
"the KV pool must be completely untouched by a Kimi request"
);
let stats = pc.lock().unwrap().stats();
assert_eq!(
stats.hits + stats.misses,
0,
"the prefix cache must never be consulted for a Kimi request"
);
}
}