Skip to main content

superfluid_daemon/
http_state.rs

1//! Shared HTTP admission state, scheduling policy and error responses.
2
3use 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    /// Whether a request that names no model is refused.
25    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}