1use crate::inference::{RequestQos, SamplingOverrides, TokenLimits};
4use crate::{Daemon, DaemonError};
5use axum::http::{HeaderMap, StatusCode};
6use axum::response::{IntoResponse, Response};
7use axum::Json;
8use std::sync::Arc;
9
10#[derive(Clone)]
11pub(crate) struct AppState {
12 pub(crate) registry: Arc<crate::registry::ModelRegistry>,
13 pub(crate) limits: TokenLimits,
14 pub(crate) sampling: SamplingOverrides,
15 pub(crate) qos: HttpQosPolicy,
16 pub(crate) keys: Option<Arc<crate::keypolicy::KeyTable>>,
17 pub(crate) fim_buckets: Arc<crate::completion_bucket::ClientBuckets>,
18 pub(crate) nonstream_keepalive: Option<std::time::Duration>,
19 pub(crate) model_naming: crate::openai::ModelNaming,
20 pub(crate) identity: Arc<crate::openai::ServerIdentity>,
21}
22
23impl AppState {
24 pub(crate) fn requires_model(&self) -> bool {
26 self.model_naming == crate::openai::ModelNaming::Exact
27 }
28 pub(crate) fn default_daemon(&self) -> Arc<Daemon> {
29 self.registry
30 .resolve(None)
31 .map(|(_, d)| d)
32 .expect("a default model is always loaded")
33 }
34 pub(crate) fn window_of(&self, id: &str) -> u32 {
37 let fleet = self.registry.is_loaded(id).then(|| self.registry.resolve(Some(id))).flatten();
38 match fleet.and_then(|(_, d)| d.fleet_window()) {
39 Some(n) => n.min(u32::MAX as u64) as u32,
40 None => self.limits.max_context,
41 }
42 }
43
44 pub(crate) fn default_name(&self) -> String {
45 self.registry.default_name()
46 }
47 #[allow(clippy::result_large_err)]
48 pub(crate) fn sync_class(
49 &self,
50 headers: &HeaderMap,
51 key: &Option<axum::Extension<crate::keypolicy::ClientKey>>,
52 ) -> Result<Option<u8>, Response> {
53 self.qos
54 .resolve(headers, key.as_ref().map(|k| &*k.0 .0))
55 .map(|q| Some(q.class))
56 }
57 #[allow(clippy::result_large_err)]
58 pub(crate) fn resolve(&self, model: Option<&str>) -> Result<(String, Arc<Daemon>), Response> {
59 let loose = match model {
60 Some(m) if self.model_naming == crate::openai::ModelNaming::Loose => self.registry.loose_name(m),
61 _ => None,
62 };
63 let model = loose.as_deref().or(model);
64 let needs_load = model.is_some_and(|m| !self.registry.is_loaded(m) && self.registry.is_known(m))
65 || self.registry.reload_pending(model);
66 let many_threads = || {
67 tokio::runtime::Handle::try_current().is_ok_and(|h| h.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread)
68 };
69 let hit = if needs_load && many_threads() {
70 tokio::task::block_in_place(|| self.registry.try_resolve(model))
71 } else {
72 self.registry.try_resolve(model)
73 };
74 let (name, daemon) = hit.map_err(|e| match e {
75 crate::registry::ResolveError::Load(e) => load_failed(model.unwrap_or(""), e),
76 crate::registry::ResolveError::NotFound => error_response(
77 StatusCode::NOT_FOUND,
78 format!(
79 "model '{}' not found; loaded: {}",
80 model.unwrap_or(""),
81 self.registry.names().join(", ")
82 ),
83 ),
84 })?;
85 if daemon.fleet_serving().is_some_and(|(up, _)| up == 0) {
87 return Err(no_node_response(format!("fleet: no node serving model '{name}' is up")));
88 }
89 Ok((name, daemon))
90 }
91}
92
93pub const QOS_HEADER: &str = "x-superfluid-qos";
94pub const BATCH_INVARIANT_HEADER: &str = "x-superfluid-batch-invariant";
95
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub struct HttpQosPolicy {
98 pub default: RequestQos,
99 pub honor_headers: bool,
100 pub allow_batch_invariant: bool,
101}
102
103impl Default for HttpQosPolicy {
104 fn default() -> HttpQosPolicy {
105 HttpQosPolicy {
106 default: RequestQos::AGENT,
107 honor_headers: true,
108 allow_batch_invariant: false,
109 }
110 }
111}
112
113impl HttpQosPolicy {
114 #[allow(clippy::result_large_err)]
115 pub(crate) fn resolve(
116 &self,
117 headers: &HeaderMap,
118 key: Option<&crate::keypolicy::KeyPolicy>,
119 ) -> Result<RequestQos, Response> {
120 if let Some(key) = key {
121 let q = HttpQosPolicy {
122 default: RequestQos { class: key.class, batch_invariant: false },
123 honor_headers: true,
124 allow_batch_invariant: key.allow_batch_invariant,
125 }
126 .resolve(headers, None)
127 .map_err(|r| {
128 if r.status() != StatusCode::FORBIDDEN {
129 return r;
130 }
131 key.note_qos_reject();
132 error_response(
133 StatusCode::FORBIDDEN,
134 format!(
135 "API key '{}' may not request a batch-invariant lane ({BATCH_INVARIANT_HEADER}); \
136 grant it with \"batch_invariant\": true on this key in the --key-policy file",
137 key.name
138 ),
139 )
140 })?;
141 if q.class < key.max_class {
142 key.note_qos_reject();
143 return Err(error_response(
144 StatusCode::FORBIDDEN,
145 format!(
146 "API key '{}' may request at most the {} class ({QOS_HEADER}: {} asked for more)",
147 key.name,
148 crate::qos::NAMES[key.max_class as usize],
149 crate::qos::NAMES[q.class as usize],
150 ),
151 ));
152 }
153 return Ok(q);
154 }
155 let mut q = self.default;
156 if !self.honor_headers {
157 return Ok(q);
158 }
159 if let Some(v) = headers.get(QOS_HEADER) {
160 q.class = v.to_str().ok().and_then(crate::qos::parse).ok_or_else(|| {
161 error_response(
162 StatusCode::BAD_REQUEST,
163 format!(
164 "invalid {QOS_HEADER} header {:?}: expected one of {}",
165 String::from_utf8_lossy(v.as_bytes()),
166 crate::qos::NAMES.join(", ")
167 ),
168 )
169 })?;
170 }
171 if let Some(v) = headers.get(BATCH_INVARIANT_HEADER) {
172 q.batch_invariant = match v.to_str().map(|s| s.trim().to_ascii_lowercase()).as_deref() {
173 Ok("1" | "true" | "yes" | "on") => true,
174 Ok("0" | "false" | "no" | "off") => false,
175 _ => {
176 return Err(error_response(
177 StatusCode::BAD_REQUEST,
178 format!(
179 "invalid {BATCH_INVARIANT_HEADER} header {:?}: expected true or false",
180 String::from_utf8_lossy(v.as_bytes())
181 ),
182 ))
183 }
184 };
185 if q.batch_invariant && !self.allow_batch_invariant {
186 return Err(error_response(
187 StatusCode::FORBIDDEN,
188 format!(
189 "{BATCH_INVARIANT_HEADER} is disabled on this server: a batch-invariant lane \
190 runs alone and holds back every other request (start superfluid with \
191 --http-allow-batch-invariant to permit it)"
192 ),
193 ));
194 }
195 }
196 Ok(q)
197 }
198}
199
200#[derive(Clone)]
201pub(crate) struct AuthenticatedCredential(pub(crate) Arc<str>);
202
203impl AuthenticatedCredential {
204 pub(crate) fn global_key(key: &str) -> AuthenticatedCredential {
205 AuthenticatedCredential(Arc::from(format!("api-key:{key}")))
206 }
207
208 pub(crate) fn policy_key(name: &str) -> AuthenticatedCredential {
209 AuthenticatedCredential(Arc::from(format!("policy:{name}")))
210 }
211}
212
213pub(crate) fn refusal_response(status: StatusCode, r: &crate::Refusal) -> Response {
214 (
215 status,
216 Json(serde_json::json!({
217 "error": {"message": r.message, "type": "invalid_request_error", "param": r.param, "code": r.code}
218 })),
219 )
220 .into_response()
221}
222
223pub(crate) fn error_response(status: StatusCode, message: String) -> Response {
224 error_response_with_code(status, message, serde_json::Value::Null)
225}
226
227pub(crate) fn error_response_code(status: StatusCode, message: String, code: &str) -> Response {
228 error_response_with_code(status, message, serde_json::Value::String(code.into()))
229}
230
231pub(crate) fn error_response_typed(
232 status: StatusCode,
233 message: String,
234 kind: &str,
235 code: &str,
236) -> Response {
237 (
238 status,
239 Json(serde_json::json!({
240 "error": {"message": message, "type": kind, "param": null, "code": code}
241 })),
242 )
243 .into_response()
244}
245
246fn error_response_with_code(
247 status: StatusCode,
248 message: String,
249 code: serde_json::Value,
250) -> Response {
251 let kind = if status.is_client_error() {
252 "invalid_request_error"
253 } else {
254 "server_error"
255 };
256 (
257 status,
258 Json(serde_json::json!({
259 "error": {"message": message, "type": kind, "param": null, "code": code}
260 })),
261 )
262 .into_response()
263}
264
265fn load_failed(model: &str, e: DaemonError) -> Response {
266 use DaemonError as E;
267 let status = match e {
268 E::Config(_) | E::Constraint(_) | E::TemplateRefused(_) | E::Unsupported(_) => StatusCode::BAD_REQUEST,
269 E::Io(_)
270 | E::Codec(_)
271 | E::WalCorrupt
272 | E::EngineLoad(_)
273 | E::Agent(_)
274 | E::Fleet(_)
275 | E::FleetNode { .. }
276 | E::Rejected(_)
277 | E::Generation(_)
278 | E::Protocol(_)
279 | E::NoNode(_)
280 | E::CompletionThrottled { .. } => StatusCode::SERVICE_UNAVAILABLE,
281 E::UnknownSession(_)
282 | E::EmptySession(_)
283 | E::StreamTooLong { .. }
284 | E::SessionBusy(_)
285 | E::UnknownToolCall(_)
286 | E::UnknownRole(_)
287 | E::InvalidTool { .. }
288 | E::SystemNotFirst(_)
289 | E::ForkPoint { .. }
290 | E::InheritedToolCall(_)
291 | E::MetaConflict { .. }
292 | E::GenerationConflict { .. }
293 | E::Purged(_)
294 | E::FilesQuota { .. }
295 | E::RebaseSpec(_)
296 | E::PurgeTimeout(_)
297 | E::UnknownPermission(_)
298 | E::PermissionPending(_)
299 | E::ToolLeaseExpired(_)
300 | E::ToolLateEffect(_)
301 | E::UnknownBlockKind(_)
302 | E::UnknownMedia(_)
303 | E::MediaTooLarge(_)
304 | E::NoFimDialect => StatusCode::SERVICE_UNAVAILABLE,
305 };
306 let message = format!("model '{model}' could not be loaded: {e}");
307 if let E::Unsupported(r) = &e {
308 return refusal_response(status, &crate::Refusal { message, ..r.clone() });
309 }
310 error_response(status, message)
311}
312
313pub(crate) fn daemon_error(e: DaemonError) -> Response {
314 if matches!(e, DaemonError::StreamTooLong { .. }) {
315 return error_response_typed(
316 StatusCode::BAD_REQUEST,
317 e.to_string(),
318 "invalid_request_error",
319 "context_length_exceeded",
320 );
321 }
322 if matches!(e, DaemonError::NoFimDialect) {
323 return error_response_code(
324 StatusCode::BAD_REQUEST,
325 "this model has no fill-in-middle dialect and no --fim-model satellite is configured"
326 .into(),
327 "fim_unsupported",
328 );
329 }
330 if let DaemonError::CompletionThrottled { retry_after_ms } = e {
331 return throttled_response(std::time::Duration::from_millis(retry_after_ms));
332 }
333 if let DaemonError::Unsupported(r) = &e {
334 return refusal_response(StatusCode::BAD_REQUEST, r);
335 }
336 if matches!(e, DaemonError::NoNode(_)) {
337 return no_node_response(e.to_string());
338 }
339 let status = match &e {
340 DaemonError::UnknownSession(_)
341 | DaemonError::EmptySession(_)
342 | DaemonError::UnknownToolCall(_)
343 | DaemonError::SessionBusy(_)
344 | DaemonError::Protocol(_)
345 | DaemonError::Constraint(_)
346 | DaemonError::TemplateRefused(_)
347 | DaemonError::Unsupported(_)
348 | DaemonError::Config(_) => StatusCode::BAD_REQUEST,
349 _ => StatusCode::INTERNAL_SERVER_ERROR,
350 };
351 error_response(status, e.to_string())
352}
353
354pub(crate) fn no_node_response(message: String) -> Response {
356 error_response_typed(StatusCode::SERVICE_UNAVAILABLE, message, "server_error", "no_node_up")
357}
358
359pub(crate) fn throttled_response(wait: std::time::Duration) -> Response {
360 let secs = crate::completion_bucket::retry_after_secs(wait);
361 let retry_after_ms = crate::completion_bucket::retry_after_ms(wait);
362 let mut resp = error_response_typed(
363 StatusCode::TOO_MANY_REQUESTS,
364 format!("completion rate limit exceeded; retry in {retry_after_ms} ms"),
365 "rate_limit_error",
366 "completion_rate_limited",
367 );
368 if let Ok(v) = axum::http::HeaderValue::from_str(&secs.to_string()) {
369 resp.headers_mut()
370 .insert(axum::http::header::RETRY_AFTER, v);
371 }
372 resp
373}
374
375#[cfg(test)]
376mod stream_failure_tests {
377 use crate::generation_output::StreamFailure;
378 use crate::{DaemonError, Refusal};
379
380 #[test]
381 fn a_stream_failure_carries_them_too() {
382 let long = StreamFailure::new(DaemonError::StreamTooLong { len: 9000, max: 8192 }, true);
383 assert_eq!((long.code, long.param.as_deref()), (Some("context_length_exceeded"), None));
384
385 let refusal = Refusal::new("unsupported_input", Some("messages[0].content"), "model 'm' on runtime r does not accept images: no media path");
386 let refused = StreamFailure::new(DaemonError::Unsupported(refusal), true);
387 assert_eq!((refused.code, refused.param.as_deref()), (Some("unsupported_input"), Some("messages[0].content")));
388 assert_eq!(refused.message, "model 'm' on runtime r does not accept images: no media path");
389
390 let other = StreamFailure::new(DaemonError::Protocol("sse client gone"), false);
391 assert_eq!((other.code, other.param), (None, None));
392 }
393}