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