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    /// The window `/v1/models` gives for `id`: on a fleet head, the largest a node serving it
35    /// has, as it is the nodes that hold the context.
36    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        // Refused before a stream starts, so a streamed request gets the status too.
86        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
354/// A fleet head with no node up for the model: 503, so a client retries rather than gives up.
355pub(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}