Skip to main content

studio_worker/
local_api.rs

1//! Always-on local HTTP API for image generation (127.0.0.1 only).
2//!
3//! Synchronous: `POST /image` blocks until the engine finishes and returns the
4//! image bytes. Models come from the local [`Catalog`], which the operator can
5//! extend at runtime via `POST /models` — the same `ModelSource` shape the
6//! studio uses. Every job is recorded into the in-app local queue.
7//!
8//! ## Security model
9//!
10//! Binding loopback alone is **not** enough:
11//!
12//! * Browsers happily fire cross-site `text/plain` POSTs at
13//!   `127.0.0.1` without a CORS preflight, and this API parses bodies
14//!   as JSON regardless of content type — so without a guard any web
15//!   page could inject catalog models (with attacker-controlled
16//!   download URLs) or burn the GPU.  DNS rebinding additionally lets
17//!   a page *read* responses.
18//! * Any other local OS user can reach the port.
19//!
20//! Defence, checked in order on every route except `GET /healthz`:
21//!
22//! 1. **Host allow-list** — must be loopback (DNS-rebinding guard).
23//! 2. **Origin allow-list** — when present, must be a loopback origin
24//!    (CSRF guard; absent means a non-browser client and is allowed).
25//! 3. **Bearer token** — `Authorization: Bearer <token>`, compared in
26//!    constant time.  The token is generated per install, persisted in
27//!    `config.toml`, and published to local clients via the owner-only
28//!    `local-api.json` discovery file in the worker's config dir.
29
30use std::net::SocketAddr;
31use std::path::PathBuf;
32use std::sync::atomic::{AtomicBool, Ordering};
33use std::sync::Arc;
34use std::time::Duration;
35
36use parking_lot::Mutex;
37use serde::Deserialize;
38use tiny_http::{Header, Method, Request, Response, Server};
39
40use crate::catalog::{Catalog, CatalogModel};
41use crate::engine::Engine;
42use crate::host::{HostError, ModelHost, ModelStatus};
43use crate::job_gate::JobGate;
44use crate::lifecycle::ModelState;
45use crate::local::{chat_on_lane, run_image, run_kind, LocalError, LocalImageRequest};
46use crate::runtime::{JobOutcome, WorkerObservers};
47use crate::stt_stream::tokens::StreamTokens;
48use crate::types::{
49    AudioSttParams, AudioTtsParams, ChatMessage, LlmParams, Task, TaskKind, TaskResult, VideoParams,
50};
51
52const TRACE_TARGET: &str = "studio_worker::local_api";
53const POLL: Duration = Duration::from_millis(200);
54
55/// Maximum accepted request-body size.  `read_body` used to read to
56/// string unbounded, so a single request could OOM the worker.
57pub const MAX_BODY_BYTES: usize = 1024 * 1024;
58
59/// Why a request was denied before reaching its handler.
60#[derive(Debug, PartialEq, Eq)]
61enum Denial {
62    /// `Host` header present but not loopback — DNS rebinding.
63    Host(String),
64    /// `Origin` header present but not a loopback origin — CSRF.
65    Origin(String),
66    /// Missing or wrong bearer token.
67    Token,
68}
69
70/// True when `host` (an HTTP `Host` header value, optionally with a
71/// `:port` suffix) names the loopback interface this API binds.
72fn host_is_loopback(host: &str) -> bool {
73    // `[::1]:port` — bracketed IPv6 keeps its colons, so strip the
74    // port only after the closing bracket.
75    let bare = if let Some(rest) = host.strip_prefix('[') {
76        match rest.split_once(']') {
77            Some((addr, _port)) => addr,
78            None => return false,
79        }
80    } else {
81        host.rsplit_once(':').map(|(h, _)| h).unwrap_or(host)
82    };
83    bare.eq_ignore_ascii_case("localhost") || bare == "127.0.0.1" || bare == "::1"
84}
85
86/// True when an `Origin` header value is a loopback origin.  `null`
87/// (sandboxed iframes / redirects) and every remote origin are
88/// rejected; only `http(s)://<loopback>[:port]` passes.
89fn origin_is_loopback(origin: &str) -> bool {
90    let rest = origin
91        .strip_prefix("https://")
92        .or_else(|| origin.strip_prefix("http://"));
93    match rest {
94        Some(host) => host_is_loopback(host),
95        None => false,
96    }
97}
98
99/// Constant-time bearer-token comparison so a local attacker can't
100/// binary-search the token through response timing.
101fn token_matches(presented: &str, expected: &str) -> bool {
102    let (a, b) = (presented.as_bytes(), expected.as_bytes());
103    if a.len() != b.len() {
104        return false;
105    }
106    a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
107}
108
109/// The request gate, pure over extracted header values so every
110/// branch is unit-testable without a socket.  Checks run in
111/// cheapest-and-broadest-first order: Host (rebinding), Origin
112/// (CSRF), then the token.
113fn deny_reason(
114    host: Option<&str>,
115    origin: Option<&str>,
116    authorization: Option<&str>,
117    token: &str,
118) -> Option<Denial> {
119    if let Some(host) = host {
120        if !host_is_loopback(host) {
121            return Some(Denial::Host(host.to_string()));
122        }
123    }
124    if let Some(origin) = origin {
125        if !origin_is_loopback(origin) {
126            return Some(Denial::Origin(origin.to_string()));
127        }
128    }
129    let presented = authorization.and_then(|a| {
130        a.strip_prefix("Bearer ")
131            .or_else(|| a.strip_prefix("bearer "))
132    });
133    match presented {
134        Some(presented) if token_matches(presented, token) => None,
135        _ => Some(Denial::Token),
136    }
137}
138
139/// The local image API server, bound but not yet serving.
140pub struct LocalApi {
141    engine: Arc<dyn Engine>,
142    catalog: Arc<Mutex<Catalog>>,
143    catalog_path: Option<PathBuf>,
144    observers: WorkerObservers,
145    server: Server,
146    addr: SocketAddr,
147    /// Bearer token every route except `GET /healthz` requires.
148    token: String,
149    /// Shared one-job-at-a-time gate.  A local generation reserves it
150    /// so it can't run concurrently with a studio job on the same GPU.
151    gate: JobGate,
152    /// Root the engine downloads models into.  Reported (with its free
153    /// space) on `/healthz` so a stuck first-use download is visible.
154    models_root: Option<PathBuf>,
155    /// The model host and streaming tokens.
156    services: ModelServices,
157    /// The daemon's runtime handles, for the `/daemon/*` routes.  `None`
158    /// answers them `503 daemon_control_unavailable`.
159    control: Option<crate::control::DaemonControl>,
160}
161
162/// What the local API offers besides one-off jobs: the model host, and the
163/// tokens that open streaming sessions on the LAN listener.
164#[derive(Clone)]
165pub struct ModelServices {
166    pub host: ModelHost,
167    pub tokens: Arc<StreamTokens>,
168    /// Port of the running LAN stream listener; 0 while it is not listening.
169    pub stream_port: Arc<std::sync::atomic::AtomicU16>,
170}
171
172impl ModelServices {
173    pub fn new(host: ModelHost) -> Self {
174        Self {
175            host,
176            tokens: Arc::new(StreamTokens::default()),
177            stream_port: Arc::new(std::sync::atomic::AtomicU16::new(0)),
178        }
179    }
180}
181
182#[derive(Deserialize)]
183#[serde(rename_all = "camelCase")]
184struct StreamTokenBody {
185    model: String,
186    #[serde(default)]
187    ttl_secs: Option<i64>,
188}
189
190/// A stream token lives this long unless the request says otherwise.
191/// The holder (Runa) refreshes well before expiry.  Safe range: within
192/// the token store's clamp (30 s..=1 h).
193const DEFAULT_STREAM_TOKEN_TTL_SECS: i64 = 600;
194
195#[derive(Deserialize)]
196#[serde(rename_all = "camelCase")]
197struct ImageBody {
198    prompt: String,
199    #[serde(default)]
200    model: Option<String>,
201    #[serde(default)]
202    negative_prompt: Option<String>,
203    #[serde(default)]
204    width: Option<u32>,
205    #[serde(default)]
206    height: Option<u32>,
207    #[serde(default)]
208    steps: Option<u32>,
209    #[serde(default)]
210    seed: Option<u64>,
211    #[serde(default)]
212    ext: Option<String>,
213}
214
215/// OpenAI-compatible chat-completions request (non-streaming subset).
216#[derive(Deserialize)]
217struct ChatBody {
218    #[serde(default)]
219    model: Option<String>,
220    messages: Vec<ChatMessageBody>,
221    #[serde(default)]
222    max_tokens: Option<u32>,
223    #[serde(default)]
224    temperature: Option<f32>,
225    #[serde(default)]
226    top_p: Option<f32>,
227    #[serde(default)]
228    stop: Option<Vec<String>>,
229    /// llama-server compatible template switches (e.g. `enable_thinking`).
230    #[serde(default)]
231    chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
232}
233
234#[derive(Deserialize)]
235struct ChatMessageBody {
236    role: String,
237    content: String,
238}
239
240#[derive(Deserialize)]
241#[serde(rename_all = "camelCase")]
242struct TtsBody {
243    text: String,
244    #[serde(default)]
245    model: Option<String>,
246    #[serde(default)]
247    voice: Option<String>,
248    #[serde(default)]
249    speed: Option<f32>,
250    #[serde(default)]
251    language: Option<String>,
252    #[serde(default)]
253    ext: Option<String>,
254}
255
256#[derive(Deserialize)]
257#[serde(rename_all = "camelCase")]
258struct SttBody {
259    input_url: String,
260    #[serde(default)]
261    model: Option<String>,
262    #[serde(default)]
263    language: Option<String>,
264}
265
266#[derive(Deserialize)]
267#[serde(rename_all = "camelCase")]
268struct VideoBody {
269    prompt: String,
270    #[serde(default)]
271    model: Option<String>,
272    #[serde(default)]
273    negative_prompt: Option<String>,
274    #[serde(default)]
275    seconds: Option<f32>,
276    #[serde(default)]
277    width: Option<u32>,
278    #[serde(default)]
279    height: Option<u32>,
280    #[serde(default)]
281    ext: Option<String>,
282}
283
284impl LocalApi {
285    /// Bind to `addr` (e.g. `127.0.0.1:0` for an ephemeral port).
286    // A constructor wiring the API's collaborators (engine, catalog,
287    // observers, auth, gate, models-root); grouping them into a struct
288    // would only move the argument list, not reduce it.
289    #[allow(clippy::too_many_arguments)]
290    pub fn bind(
291        addr: &str,
292        engine: Arc<dyn Engine>,
293        catalog: Arc<Mutex<Catalog>>,
294        catalog_path: Option<PathBuf>,
295        observers: WorkerObservers,
296        token: String,
297        gate: JobGate,
298        models_root: Option<PathBuf>,
299        services: ModelServices,
300    ) -> anyhow::Result<Self> {
301        anyhow::ensure!(
302            !token.is_empty(),
303            "local api: refusing to serve with an empty token"
304        );
305        let server =
306            Server::http(addr).map_err(|e| anyhow::anyhow!("local api bind {addr}: {e}"))?;
307        let addr = server
308            .server_addr()
309            .to_ip()
310            .ok_or_else(|| anyhow::anyhow!("local api: non-ip listen address"))?;
311        Ok(Self {
312            engine,
313            catalog,
314            catalog_path,
315            observers,
316            server,
317            addr,
318            token,
319            gate,
320            models_root,
321            services,
322            control: None,
323        })
324    }
325
326    /// Serve the `/daemon/*` routes with `control`.
327    pub fn with_control(mut self, control: crate::control::DaemonControl) -> Self {
328        self.control = Some(control);
329        self
330    }
331
332    /// The bound socket address.
333    pub fn local_addr(&self) -> SocketAddr {
334        self.addr
335    }
336
337    /// The base URL the API is reachable at.
338    pub fn url(&self) -> String {
339        format!("http://{}", self.addr)
340    }
341
342    /// Serve requests until `stop` is set.
343    ///
344    /// Requests are handled on a small pool of worker threads (tiny_http's
345    /// `Server` is `Sync`, so several threads can `recv_timeout`
346    /// concurrently).  A single generation can take ~10 s, and a
347    /// first-use model download minutes — on the old single-threaded
348    /// loop that blocked `/healthz`, `/models`, and every other caller.
349    /// The pool keeps the cheap routes responsive while a job runs.
350    /// `std::thread::scope` lets the workers borrow `&self` + `stop`
351    /// without an `Arc`, so the public signature is unchanged.
352    pub fn serve(&self, stop: &AtomicBool) {
353        // Long generations each hold a thread while the tray UI polls once a
354        // second; eight keeps the cheap routes answering alongside them.
355        const WORKERS: usize = 8;
356        std::thread::scope(|scope| {
357            for _ in 0..WORKERS {
358                scope.spawn(|| {
359                    while !stop.load(Ordering::Relaxed) {
360                        match self.server.recv_timeout(POLL) {
361                            Ok(Some(request)) => self.route(request),
362                            Ok(None) => {}
363                            Err(err) => {
364                                tracing::warn!(target: TRACE_TARGET, error = %err, "local api recv error");
365                                break;
366                            }
367                        }
368                    }
369                });
370            }
371        });
372    }
373
374    fn route(&self, request: Request) {
375        let method = request.method().clone();
376        let url = request.url().to_string();
377        let path = url.split('?').next().unwrap_or("/");
378
379        // `GET /healthz` stays open: liveness only, no secrets.  Every
380        // other route passes the Host / Origin / token gate first.
381        if !(method == Method::Get && path == "/healthz") {
382            let header = |name: &'static str| {
383                request
384                    .headers()
385                    .iter()
386                    .find(|h| h.field.equiv(name))
387                    .map(|h| h.value.as_str().to_string())
388            };
389            let denial = deny_reason(
390                header("host").as_deref(),
391                header("origin").as_deref(),
392                header("authorization").as_deref(),
393                &self.token,
394            );
395            if let Some(denial) = denial {
396                let (status, body) = match &denial {
397                    Denial::Host(host) => {
398                        (403, format!("forbidden: non-loopback Host header {host:?}"))
399                    }
400                    Denial::Origin(origin) => (
401                        403,
402                        format!("forbidden: cross-site request from Origin {origin:?}"),
403                    ),
404                    Denial::Token => (
405                        401,
406                        "missing or invalid Authorization bearer token; local clients \
407                         can read the current token from the local-api.json discovery \
408                         file in the worker's config directory"
409                            .to_string(),
410                    ),
411                };
412                tracing::warn!(
413                    target: TRACE_TARGET,
414                    op = "deny",
415                    method = %method,
416                    path,
417                    status,
418                    reason = ?denial,
419                    "local api request denied"
420                );
421                if let Err(err) = respond(request, status, "text/plain", body.as_bytes()) {
422                    tracing::warn!(target: TRACE_TARGET, error = %err, "local api respond error");
423                }
424                return;
425            }
426        }
427
428        let outcome = match (&method, path) {
429            (Method::Get, "/healthz") => self.handle_healthz(request),
430            (Method::Post, "/image") => self.handle_image(request),
431            (Method::Post, "/v1/chat/completions") => self.handle_chat(request),
432            (Method::Post, "/tts") => self.handle_tts(request),
433            (Method::Post, "/stt") => self.handle_stt(request),
434            (Method::Post, "/video") => self.handle_video(request),
435            (Method::Get, "/models") => self.handle_list_models(request),
436            (Method::Post, "/models") => self.handle_add_model(request),
437            (Method::Get, "/jobs") => self.handle_jobs(request),
438            (Method::Post, "/stream-tokens") => self.handle_stream_token(request),
439            (_, p) if p.starts_with("/daemon/") => self.handle_daemon(request, &method, &url),
440            (Method::Get, p) if job_route(p, "/log").is_some() => {
441                let id = job_route(p, "/log").unwrap_or_default().to_string();
442                self.handle_job_log(request, &id)
443            }
444            (Method::Get, p) if job_route(p, "/thumbnail").is_some() => {
445                let id = job_route(p, "/thumbnail").unwrap_or_default().to_string();
446                self.handle_job_thumbnail(request, &id)
447            }
448            (Method::Get, p) if lifecycle_route(p, "/state").is_some() => {
449                let id = lifecycle_route(p, "/state").unwrap_or_default().to_string();
450                self.respond_lifecycle(request, self.services.host.status(&id), 200)
451            }
452            (Method::Post, p) if lifecycle_route(p, "/load").is_some() => {
453                let id = lifecycle_route(p, "/load").unwrap_or_default().to_string();
454                self.respond_lifecycle(request, self.services.host.load(&id), 202)
455            }
456            (Method::Post, p) if lifecycle_route(p, "/unload").is_some() => {
457                let id = lifecycle_route(p, "/unload")
458                    .unwrap_or_default()
459                    .to_string();
460                self.respond_lifecycle(request, self.services.host.unload(&id), 202)
461            }
462            (Method::Delete, p) if p.starts_with("/models/") => {
463                let id = p.trim_start_matches("/models/").to_string();
464                self.handle_delete_model(request, &id)
465            }
466            _ => respond(request, 404, "text/plain", b"not found"),
467        };
468        if let Err(err) = outcome {
469            tracing::warn!(target: TRACE_TARGET, error = %err, "local api respond error");
470        }
471    }
472
473    /// Liveness + a read-only runtime snapshot for operators and local
474    /// tooling.  Unauthenticated (no secrets, no prompts): the worker
475    /// version, whether a job is in flight, the engine name, and the
476    /// models-root free space so a stuck first-use download is
477    /// diagnosable without shelling into the box.
478    fn handle_healthz(&self, request: Request) -> std::io::Result<()> {
479        let free_bytes = self
480            .models_root
481            .as_deref()
482            .and_then(|root| fs4::available_space(root).ok());
483        let gpu = self
484            .observers
485            .gpu_runtime
486            .lock()
487            .clone()
488            .map(|g| serde_json::json!({ "ok": g.ok, "detail": g.detail }));
489        let body = serde_json::json!({
490            "ok": true,
491            "version": crate::AGENT_VERSION,
492            "busy": self.gate.is_busy(),
493            "engine": self.engine.name(),
494            "modelsRoot": self.models_root.as_ref().map(|p| p.display().to_string()),
495            "modelsRootFreeBytes": free_bytes,
496            "gpuRuntime": gpu,
497        });
498        match serde_json::to_vec(&body) {
499            Ok(bytes) => respond(request, 200, "application/json", &bytes),
500            // Never let a serialisation slip break liveness.
501            Err(_) => respond(request, 200, "application/json", b"{\"ok\":true}"),
502        }
503    }
504
505    fn handle_image(&self, mut request: Request) -> std::io::Result<()> {
506        let body = match read_body(&mut request)? {
507            BodyOutcome::Ok(body) => body,
508            BodyOutcome::TooLarge => return respond_too_large(request),
509        };
510        let parsed: ImageBody = match serde_json::from_str(&body) {
511            Ok(parsed) => parsed,
512            Err(err) => {
513                return respond(
514                    request,
515                    400,
516                    "text/plain",
517                    format!("bad json: {err}").as_bytes(),
518                )
519            }
520        };
521        let req = LocalImageRequest {
522            prompt: parsed.prompt,
523            model: parsed.model,
524            negative_prompt: parsed.negative_prompt,
525            width: parsed.width,
526            height: parsed.height,
527            steps: parsed.steps,
528            seed: parsed.seed,
529            ext: parsed.ext,
530        };
531
532        // One GPU, one job: reserve the shared gate so a local
533        // generation never runs alongside a studio job.  Busy → 503 +
534        // Retry-After so the caller backs off instead of OOMing.
535        let Some(_reservation) = self.gate.try_reserve() else {
536            return respond_busy(request);
537        };
538
539        let catalog = self.catalog.lock().clone();
540        match run_image(self.engine.as_ref(), &catalog, &self.observers, &req) {
541            Ok(TaskResult::Image { bytes, ext }) => {
542                respond(request, 200, content_type_for(&ext), &bytes)
543            }
544            Ok(_) => respond(request, 500, "text/plain", b"unexpected non-image result"),
545            Err(err) => respond_local_err(request, err),
546        }
547    }
548
549    /// OpenAI-compatible chat completions.  Resolves an LLM model from
550    /// the catalog (explicit `model` or the default), dispatches, and
551    /// returns the engine's JSON verbatim (the synthetic + llama
552    /// engines already emit a `chat.completion`-shaped body).
553    fn handle_chat(&self, mut request: Request) -> std::io::Result<()> {
554        let body = match read_body(&mut request)? {
555            BodyOutcome::Ok(body) => body,
556            BodyOutcome::TooLarge => return respond_too_large(request),
557        };
558        let parsed: ChatBody = match serde_json::from_str(&body) {
559            Ok(p) => p,
560            Err(err) => {
561                return respond(
562                    request,
563                    400,
564                    "text/plain",
565                    format!("bad json: {err}").as_bytes(),
566                )
567            }
568        };
569        let prompt_preview = parsed
570            .messages
571            .last()
572            .map(|m| m.content.clone())
573            .unwrap_or_default();
574        let params = LlmParams {
575            messages: parsed
576                .messages
577                .into_iter()
578                .map(|m| ChatMessage {
579                    role: m.role,
580                    content: m.content,
581                })
582                .collect(),
583            max_tokens: parsed.max_tokens.unwrap_or(512),
584            temperature: parsed.temperature.unwrap_or(0.7),
585            top_p: parsed.top_p,
586            stop: parsed.stop,
587            chat_template_kwargs: parsed.chat_template_kwargs,
588            ..Default::default()
589        };
590        let catalog = self.catalog.lock().clone();
591        // A loaded model answers on its own lane, outside the job gate.
592        if let Some(result) = chat_on_lane(
593            &self.services.host,
594            &catalog,
595            &self.observers,
596            parsed.model.as_deref(),
597            &prompt_preview,
598            params.clone(),
599        ) {
600            return respond_llm(request, result);
601        }
602        let Some(_reservation) = self.gate.try_reserve() else {
603            return respond_busy(request);
604        };
605        let outcome = run_kind(
606            self.engine.as_ref(),
607            &catalog,
608            &self.observers,
609            TaskKind::Llm,
610            parsed.model.as_deref(),
611            &prompt_preview,
612            Task::Llm(params),
613        );
614        respond_llm(request, outcome)
615    }
616
617    fn handle_tts(&self, mut request: Request) -> std::io::Result<()> {
618        let body = match read_body(&mut request)? {
619            BodyOutcome::Ok(body) => body,
620            BodyOutcome::TooLarge => return respond_too_large(request),
621        };
622        let parsed: TtsBody = match serde_json::from_str(&body) {
623            Ok(p) => p,
624            Err(err) => {
625                return respond(
626                    request,
627                    400,
628                    "text/plain",
629                    format!("bad json: {err}").as_bytes(),
630                )
631            }
632        };
633        let preview = parsed.text.clone();
634        let params = AudioTtsParams {
635            text: parsed.text,
636            voice: parsed.voice.unwrap_or_else(|| "default".into()),
637            speed: parsed.speed,
638            language: parsed.language,
639            ext: parsed.ext.unwrap_or_else(|| "wav".into()),
640        };
641        let Some(_reservation) = self.gate.try_reserve() else {
642            return respond_busy(request);
643        };
644        let catalog = self.catalog.lock().clone();
645        match run_kind(
646            self.engine.as_ref(),
647            &catalog,
648            &self.observers,
649            TaskKind::AudioTts,
650            parsed.model.as_deref(),
651            &preview,
652            Task::AudioTts(params),
653        ) {
654            Ok(TaskResult::AudioTts { bytes, ext }) => {
655                respond(request, 200, content_type_for(&ext), &bytes)
656            }
657            Ok(_) => respond(request, 500, "text/plain", b"unexpected non-audio result"),
658            Err(err) => respond_local_err(request, err),
659        }
660    }
661
662    fn handle_stt(&self, mut request: Request) -> std::io::Result<()> {
663        let body = match read_body(&mut request)? {
664            BodyOutcome::Ok(body) => body,
665            BodyOutcome::TooLarge => return respond_too_large(request),
666        };
667        let parsed: SttBody = match serde_json::from_str(&body) {
668            Ok(p) => p,
669            Err(err) => {
670                return respond(
671                    request,
672                    400,
673                    "text/plain",
674                    format!("bad json: {err}").as_bytes(),
675                )
676            }
677        };
678        let preview = parsed.input_url.clone();
679        let params = AudioSttParams {
680            input_url: parsed.input_url,
681            language: parsed.language,
682            ..Default::default()
683        };
684        let Some(_reservation) = self.gate.try_reserve() else {
685            return respond_busy(request);
686        };
687        let catalog = self.catalog.lock().clone();
688        match run_kind(
689            self.engine.as_ref(),
690            &catalog,
691            &self.observers,
692            TaskKind::AudioStt,
693            parsed.model.as_deref(),
694            &preview,
695            Task::AudioStt(params),
696        ) {
697            Ok(TaskResult::AudioStt { json }) => match serde_json::to_vec(&json) {
698                Ok(bytes) => respond(request, 200, "application/json", &bytes),
699                Err(e) => respond(request, 500, "text/plain", e.to_string().as_bytes()),
700            },
701            Ok(_) => respond(
702                request,
703                500,
704                "text/plain",
705                b"unexpected non-transcript result",
706            ),
707            Err(err) => respond_local_err(request, err),
708        }
709    }
710
711    fn handle_video(&self, mut request: Request) -> std::io::Result<()> {
712        let body = match read_body(&mut request)? {
713            BodyOutcome::Ok(body) => body,
714            BodyOutcome::TooLarge => return respond_too_large(request),
715        };
716        let parsed: VideoBody = match serde_json::from_str(&body) {
717            Ok(p) => p,
718            Err(err) => {
719                return respond(
720                    request,
721                    400,
722                    "text/plain",
723                    format!("bad json: {err}").as_bytes(),
724                )
725            }
726        };
727        let preview = parsed.prompt.clone();
728        let params = VideoParams {
729            prompt: parsed.prompt,
730            negative_prompt: parsed.negative_prompt,
731            seconds: parsed.seconds.unwrap_or(2.0),
732            width: parsed.width.unwrap_or(256),
733            height: parsed.height.unwrap_or(256),
734            ext: parsed.ext.unwrap_or_else(|| "mp4".into()),
735            ..Default::default()
736        };
737        let Some(_reservation) = self.gate.try_reserve() else {
738            return respond_busy(request);
739        };
740        let catalog = self.catalog.lock().clone();
741        match run_kind(
742            self.engine.as_ref(),
743            &catalog,
744            &self.observers,
745            TaskKind::Video,
746            parsed.model.as_deref(),
747            &preview,
748            Task::Video(params),
749        ) {
750            Ok(TaskResult::Video { bytes, ext }) => {
751                respond(request, 200, content_type_for(&ext), &bytes)
752            }
753            Ok(_) => respond(request, 500, "text/plain", b"unexpected non-video result"),
754            Err(err) => respond_local_err(request, err),
755        }
756    }
757
758    fn handle_list_models(&self, request: Request) -> std::io::Result<()> {
759        let models = self.catalog.lock().models.clone();
760        let statuses = self.services.host.statuses();
761        let listed: Vec<serde_json::Value> = models
762            .iter()
763            .map(|model| {
764                let mut value = serde_json::to_value(model).unwrap_or_default();
765                if let (Some(obj), Some(status)) = (
766                    value.as_object_mut(),
767                    statuses.iter().find(|s| s.id == model.id),
768                ) {
769                    obj.insert("state".into(), status.state.name().into());
770                    obj.insert("resident".into(), status.resident.into());
771                    obj.insert("since".into(), status.since.to_rfc3339().into());
772                    obj.insert("loadable".into(), self.services.host.can_load(model).into());
773                    if let ModelState::Failed { reason } = &status.state {
774                        obj.insert("error".into(), reason.clone().into());
775                    }
776                }
777                value
778            })
779            .collect();
780        match serde_json::to_vec(&listed) {
781            Ok(body) => respond(request, 200, "application/json", &body),
782            Err(err) => respond(request, 500, "text/plain", err.to_string().as_bytes()),
783        }
784    }
785
786    fn handle_add_model(&self, mut request: Request) -> std::io::Result<()> {
787        let body = match read_body(&mut request)? {
788            BodyOutcome::Ok(body) => body,
789            BodyOutcome::TooLarge => return respond_too_large(request),
790        };
791        let model: CatalogModel = match serde_json::from_str(&body) {
792            Ok(model) => model,
793            Err(err) => {
794                return respond(
795                    request,
796                    400,
797                    "text/plain",
798                    format!("bad model: {err}").as_bytes(),
799                )
800            }
801        };
802        let saved = {
803            let mut catalog = self.catalog.lock();
804            catalog.upsert(model);
805            self.persist(&catalog)
806        };
807        match saved {
808            Ok(()) => respond(request, 200, "application/json", b"{\"ok\":true}"),
809            Err(err) => respond(request, 500, "text/plain", err.to_string().as_bytes()),
810        }
811    }
812
813    fn handle_delete_model(&self, request: Request, id: &str) -> std::io::Result<()> {
814        // Free the weights and drop the residency before the entry goes.
815        if let Err(err) = self.services.host.unload(id) {
816            if !matches!(err, HostError::UnknownModel(_)) {
817                return respond_json(
818                    request,
819                    500,
820                    &serde_json::json!({ "error": "unload_failed", "message": err.to_string() }),
821                );
822            }
823        }
824        let (existed, saved) = {
825            let mut catalog = self.catalog.lock();
826            let existed = catalog.remove(id);
827            (existed, self.persist(&catalog))
828        };
829        if !existed {
830            return respond(request, 404, "text/plain", b"no such model");
831        }
832        match saved {
833            Ok(()) => respond(request, 200, "application/json", b"{\"ok\":true}"),
834            Err(err) => respond(request, 500, "text/plain", err.to_string().as_bytes()),
835        }
836    }
837
838    /// Answer a lifecycle call: the model's status, or a named error.
839    /// `pending_status` is used while the model is still transitioning.
840    fn respond_lifecycle(
841        &self,
842        request: Request,
843        outcome: Result<ModelStatus, HostError>,
844        pending_status: u16,
845    ) -> std::io::Result<()> {
846        match outcome {
847            Ok(status) => {
848                let code = match status.state {
849                    ModelState::Loading | ModelState::Unloading => pending_status,
850                    _ => 200,
851                };
852                respond_json(request, code, &status_json(&status))
853            }
854            Err(err) => {
855                let (code, body) = match &err {
856                    HostError::UnknownModel(_) => {
857                        (404, serde_json::json!({ "error": "unknown_model" }))
858                    }
859                    HostError::Disabled(_) => {
860                        (400, serde_json::json!({ "error": "model_disabled" }))
861                    }
862                    HostError::Refused(r) => (
863                        409,
864                        serde_json::json!({
865                            "error": "insufficient_memory",
866                            "neededGib": r.needed_gib,
867                            "freeGib": r.free_gib,
868                            "marginGib": r.margin_gib,
869                        }),
870                    ),
871                    HostError::NotLoaded { state, .. } => (
872                        409,
873                        serde_json::json!({ "error": "model_not_loaded", "state": state }),
874                    ),
875                    HostError::Persist(_) => {
876                        (500, serde_json::json!({ "error": "residency_not_saved" }))
877                    }
878                    HostError::LaneBusy(_) => (409, serde_json::json!({ "error": "model_busy" })),
879                };
880                let mut body = body;
881                body["message"] = err.to_string().into();
882                tracing::warn!(
883                    target: TRACE_TARGET,
884                    op = "lifecycle",
885                    status = code,
886                    error = %err,
887                    "lifecycle request refused"
888                );
889                respond_json(request, code, &body)
890            }
891        }
892    }
893
894    /// Mint a short-lived token that opens one streaming model on the LAN
895    /// listener (\`ws://<host>:<port>/transcribe?token=...\`).
896    fn handle_stream_token(&self, mut request: Request) -> std::io::Result<()> {
897        let body = match read_body(&mut request)? {
898            BodyOutcome::Ok(body) => body,
899            BodyOutcome::TooLarge => return respond_too_large(request),
900        };
901        let parsed: StreamTokenBody = match serde_json::from_str(&body) {
902            Ok(p) => p,
903            Err(err) => {
904                return respond_json(
905                    request,
906                    400,
907                    &serde_json::json!({ "error": "bad_request", "message": err.to_string() }),
908                )
909            }
910        };
911        let model = self.catalog.lock().get(&parsed.model).cloned();
912        let Some(model) = model else {
913            return respond_json(
914                request,
915                404,
916                &serde_json::json!({ "error": "unknown_model" }),
917            );
918        };
919        if model.source.engine != crate::types::ModelEngine::Parakeet {
920            return respond_json(
921                request,
922                400,
923                &serde_json::json!({ "error": "not_a_stream_model" }),
924            );
925        }
926        let port = self.services.stream_port.load(Ordering::SeqCst);
927        if port == 0 {
928            return respond_json(
929                request,
930                503,
931                &serde_json::json!({ "error": "stream_listener_down" }),
932            );
933        }
934        let ttl =
935            chrono::Duration::seconds(parsed.ttl_secs.unwrap_or(DEFAULT_STREAM_TOKEN_TTL_SECS));
936        let grant = self
937            .services
938            .tokens
939            .mint(&model.id, ttl, chrono::Utc::now());
940        tracing::info!(
941            target: TRACE_TARGET,
942            op = "stream_token",
943            model = %model.id,
944            expires_at = %grant.expires_at,
945            "stream token minted"
946        );
947        respond_json(
948            request,
949            200,
950            &serde_json::json!({
951                "token": grant.token,
952                "model": grant.model,
953                "expiresAt": grant.expires_at.to_rfc3339(),
954                "port": port,
955                "path": crate::stt_stream::server::STREAM_PATH,
956            }),
957        )
958    }
959
960    fn handle_jobs(&self, request: Request) -> std::io::Result<()> {
961        let jobs: Vec<serde_json::Value> = self
962            .observers
963            .local_jobs
964            .lock()
965            .iter()
966            .map(|job| {
967                let (status, reason) = match &job.outcome {
968                    JobOutcome::Completed => ("completed", None),
969                    JobOutcome::Failed { reason } => ("failed", Some(reason.clone())),
970                };
971                serde_json::json!({
972                    "jobId": job.job_id,
973                    "kind": job.kind.as_str(),
974                    "model": job.model,
975                    "prompt": job.prompt,
976                    "status": status,
977                    "reason": reason,
978                    "startedAt": job.started_at.to_rfc3339(),
979                    "finishedAt": job.finished_at.to_rfc3339(),
980                })
981            })
982            .collect();
983        match serde_json::to_vec(&jobs) {
984            Ok(body) => respond(request, 200, "application/json", &body),
985            Err(err) => respond(request, 500, "text/plain", err.to_string().as_bytes()),
986        }
987    }
988
989    /// `/daemon/*`: what the tray UI sees and does.
990    fn handle_daemon(
991        &self,
992        mut request: Request,
993        method: &Method,
994        url: &str,
995    ) -> std::io::Result<()> {
996        let Some(control) = &self.control else {
997            return respond_json(
998                request,
999                503,
1000                &serde_json::json!({ "error": "daemon_control_unavailable" }),
1001            );
1002        };
1003        let path = url.split('?').next().unwrap_or("/");
1004        match (method, path) {
1005            (Method::Get, "/daemon/status") => {
1006                let status = control.status(&self.observers, self.gate.is_busy());
1007                respond_serialised(request, 200, &status)
1008            }
1009            (Method::Get, "/daemon/logs") => {
1010                let after = query_param(url, "after")
1011                    .and_then(|v| v.parse::<u64>().ok())
1012                    .unwrap_or(0);
1013                let (entries, seq) = crate::runtime::recent_logs_after(&self.observers, after);
1014                respond_serialised(request, 200, &crate::daemon_api::LogsPage { entries, seq })
1015            }
1016            (Method::Post, "/daemon/pause") => {
1017                let paused = control.set_paused(true);
1018                respond_json(request, 200, &serde_json::json!({ "paused": paused }))
1019            }
1020            (Method::Post, "/daemon/resume") => {
1021                let paused = control.set_paused(false);
1022                respond_json(request, 200, &serde_json::json!({ "paused": paused }))
1023            }
1024            (Method::Get, "/daemon/config") => {
1025                respond_serialised(request, 200, &control.editable_config())
1026            }
1027            (Method::Put, "/daemon/config") => {
1028                let body = match read_body(&mut request)? {
1029                    BodyOutcome::Ok(body) => body,
1030                    BodyOutcome::TooLarge => return respond_too_large(request),
1031                };
1032                let edit: crate::daemon_api::EditableConfig = match serde_json::from_str(&body) {
1033                    Ok(edit) => edit,
1034                    Err(err) => {
1035                        return respond_error(request, 400, "bad_request", &err.to_string())
1036                    }
1037                };
1038                match control.update_config(edit) {
1039                    Ok(saved) => respond_serialised(request, 200, &saved),
1040                    Err(err @ crate::control::ControlError::Invalid(_)) => {
1041                        respond_error(request, 400, "invalid_config", &err.to_string())
1042                    }
1043                    Err(err) => respond_error(request, 500, "config_not_saved", &err.to_string()),
1044                }
1045            }
1046            (Method::Post, "/daemon/registration/reset") => {
1047                match control.request_registration_reset() {
1048                    Ok(()) => respond_json(request, 202, &serde_json::json!({ "ok": true })),
1049                    Err(err) => respond_error(request, 409, "not_rejected", &err.to_string()),
1050                }
1051            }
1052            (Method::Post, "/daemon/shutdown") => {
1053                control.shutdown();
1054                respond_json(request, 202, &serde_json::json!({ "ok": true }))
1055            }
1056            _ => respond_error(request, 404, "not_found", "no such daemon route"),
1057        }
1058    }
1059
1060    fn handle_job_log(&self, request: Request, id: &str) -> std::io::Result<()> {
1061        match crate::job_log::global().get(id) {
1062            Some(log) => respond_serialised(request, 200, &log),
1063            None => respond_error(request, 404, "unknown_job", "no log captured for that job"),
1064        }
1065    }
1066
1067    fn handle_job_thumbnail(&self, request: Request, id: &str) -> std::io::Result<()> {
1068        match self.observers.thumbnails.get(id) {
1069            Some(png) => respond(request, 200, "image/png", &png),
1070            None => respond_error(request, 404, "no_thumbnail", "no thumbnail for that job"),
1071        }
1072    }
1073
1074    fn persist(&self, catalog: &Catalog) -> std::io::Result<()> {
1075        match &self.catalog_path {
1076            Some(path) => catalog.save(path),
1077            None => Ok(()),
1078        }
1079    }
1080}
1081
1082/// A request body, or a refusal to read one past [`MAX_BODY_BYTES`].
1083enum BodyOutcome {
1084    Ok(String),
1085    TooLarge,
1086}
1087
1088fn read_body(request: &mut Request) -> std::io::Result<BodyOutcome> {
1089    // Declared length first — reject without reading a byte.
1090    if matches!(request.body_length(), Some(len) if len > MAX_BODY_BYTES) {
1091        return Ok(BodyOutcome::TooLarge);
1092    }
1093    // Then a hard cap on the reader for chunked / lying senders: read
1094    // at most one byte past the cap so overflow is detectable.
1095    let mut body = String::new();
1096    use std::io::Read as _;
1097    request
1098        .as_reader()
1099        .take(MAX_BODY_BYTES as u64 + 1)
1100        .read_to_string(&mut body)?;
1101    if body.len() > MAX_BODY_BYTES {
1102        return Ok(BodyOutcome::TooLarge);
1103    }
1104    Ok(BodyOutcome::Ok(body))
1105}
1106
1107fn respond_too_large(request: Request) -> std::io::Result<()> {
1108    respond(
1109        request,
1110        413,
1111        "text/plain",
1112        format!("request body exceeds {MAX_BODY_BYTES} bytes").as_bytes(),
1113    )
1114}
1115
1116/// 503 when the single job slot is taken by another job (studio or
1117/// local).  Carries `Retry-After: 2` so a client polls back rather
1118/// than hammering.
1119fn respond_busy(request: Request) -> std::io::Result<()> {
1120    let retry = Header::from_bytes(b"Retry-After".as_slice(), b"2".as_slice())
1121        .expect("static Retry-After header is valid");
1122    let response = Response::from_data(
1123        b"worker is busy with another job (studio or local); retry shortly".to_vec(),
1124    )
1125    .with_status_code(503)
1126    .with_header(retry)
1127    .with_header(
1128        Header::from_bytes(b"Content-Type".as_slice(), b"text/plain".as_slice())
1129            .expect("static content-type header is valid"),
1130    );
1131    request.respond(response)
1132}
1133
1134/// Publish the bound URL + bearer token for local clients, atomically
1135/// and owner-only (the file carries the token).  Written on every
1136/// successful bind; removed again by [`remove_discovery_file`] on
1137/// clean shutdown so stale files can't point at a dead port.
1138pub fn write_discovery_file(path: &std::path::Path, url: &str, token: &str) -> anyhow::Result<()> {
1139    let body = serde_json::json!({ "url": url, "token": token });
1140    let text = serde_json::to_string_pretty(&body)?;
1141    crate::config::write_atomic(path, text.as_bytes())?;
1142    tracing::info!(
1143        target: TRACE_TARGET,
1144        op = "discovery",
1145        path = %path.display(),
1146        url,
1147        "local api discovery file written"
1148    );
1149    Ok(())
1150}
1151
1152/// Best-effort removal of the discovery file on shutdown.  A missing
1153/// file is the desired end state; any other failure is warn-logged so
1154/// a stale token file never vanishes silently *and* never lingers
1155/// silently.
1156pub fn remove_discovery_file(path: &std::path::Path) {
1157    if let Err(e) = std::fs::remove_file(path) {
1158        if e.kind() != std::io::ErrorKind::NotFound {
1159            tracing::warn!(
1160                target: TRACE_TARGET,
1161                op = "discovery",
1162                path = %path.display(),
1163                error = %e,
1164                "failed to remove local api discovery file"
1165            );
1166        }
1167    }
1168}
1169
1170fn content_type_for(ext: &str) -> &'static str {
1171    match ext.to_ascii_lowercase().as_str() {
1172        "webp" => "image/webp",
1173        "png" => "image/png",
1174        "jpg" | "jpeg" => "image/jpeg",
1175        "gif" => "image/gif",
1176        "wav" => "audio/wav",
1177        "mp3" => "audio/mpeg",
1178        "ogg" | "opus" => "audio/ogg",
1179        "flac" => "audio/flac",
1180        "mp4" => "video/mp4",
1181        "webm" => "video/webm",
1182        _ => "application/octet-stream",
1183    }
1184}
1185
1186/// Map a [`LocalError`] onto an HTTP response: catalog/contract errors
1187/// (unknown or wrong-kind model, none configured) are the caller's
1188/// fault → 400 with the message; an engine failure is 500.
1189fn respond_local_err(request: Request, err: LocalError) -> std::io::Result<()> {
1190    let status = match err {
1191        LocalError::Engine(_) => 500,
1192        _ => 400,
1193    };
1194    respond(request, status, "text/plain", err.to_string().as_bytes())
1195}
1196
1197fn respond_llm(request: Request, outcome: Result<TaskResult, LocalError>) -> std::io::Result<()> {
1198    match outcome {
1199        Ok(TaskResult::Llm { json }) => match serde_json::to_vec(&json) {
1200            Ok(bytes) => respond(request, 200, "application/json", &bytes),
1201            Err(e) => respond(request, 500, "text/plain", e.to_string().as_bytes()),
1202        },
1203        Ok(_) => respond(request, 500, "text/plain", b"unexpected non-llm result"),
1204        Err(err) => respond_local_err(request, err),
1205    }
1206}
1207
1208/// `/models/<id><suffix>` -> `Some(id)` for a non-empty id without slashes.
1209fn lifecycle_route<'a>(path: &'a str, suffix: &str) -> Option<&'a str> {
1210    let id = path.strip_prefix("/models/")?.strip_suffix(suffix)?;
1211    (!id.is_empty() && !id.contains('/')).then_some(id)
1212}
1213
1214/// `/jobs/<id><suffix>` -> `Some(id)` for a non-empty id without slashes.
1215fn job_route<'a>(path: &'a str, suffix: &str) -> Option<&'a str> {
1216    let id = path.strip_prefix("/jobs/")?.strip_suffix(suffix)?;
1217    (!id.is_empty() && !id.contains('/')).then_some(id)
1218}
1219
1220/// The value of query parameter `name` in `url`, if present.
1221fn query_param<'a>(url: &'a str, name: &str) -> Option<&'a str> {
1222    url.split_once('?')?
1223        .1
1224        .split('&')
1225        .find_map(|pair| pair.strip_prefix(name)?.strip_prefix('='))
1226}
1227
1228/// The wire shape of a model's status.
1229fn status_json(status: &ModelStatus) -> serde_json::Value {
1230    let mut body = serde_json::json!({
1231        "id": status.id,
1232        "state": status.state.name(),
1233        "resident": status.resident,
1234        "since": status.since.to_rfc3339(),
1235    });
1236    if let ModelState::Failed { reason } = &status.state {
1237        body["error"] = reason.clone().into();
1238    }
1239    body
1240}
1241
1242fn respond_serialised<T: serde::Serialize>(
1243    request: Request,
1244    status: u16,
1245    body: &T,
1246) -> std::io::Result<()> {
1247    match serde_json::to_vec(body) {
1248        Ok(bytes) => respond(request, status, "application/json", &bytes),
1249        Err(err) => respond(request, 500, "text/plain", err.to_string().as_bytes()),
1250    }
1251}
1252
1253fn respond_error(request: Request, status: u16, code: &str, message: &str) -> std::io::Result<()> {
1254    respond_serialised(
1255        request,
1256        status,
1257        &crate::daemon_api::ErrorBody {
1258            error: code.to_string(),
1259            message: Some(message.to_string()),
1260        },
1261    )
1262}
1263
1264fn respond_json(request: Request, status: u16, body: &serde_json::Value) -> std::io::Result<()> {
1265    let bytes = serde_json::to_vec(body).unwrap_or_else(|_| b"{}".to_vec());
1266    respond(request, status, "application/json", &bytes)
1267}
1268
1269fn respond(request: Request, status: u16, content_type: &str, body: &[u8]) -> std::io::Result<()> {
1270    let header = Header::from_bytes(b"Content-Type".as_slice(), content_type.as_bytes())
1271        .expect("static content-type header is valid");
1272    let response = Response::from_data(body)
1273        .with_status_code(status)
1274        .with_header(header);
1275    request.respond(response)
1276}
1277
1278#[cfg(test)]
1279mod tests {
1280    use super::*;
1281    use crate::catalog::CatalogModel;
1282    use crate::engine::{EngineCapabilities, SyntheticEngine};
1283    use crate::types::{ModelCliDefaults, ModelEngine, ModelSource, Task, TaskKind};
1284
1285    /// An engine that sleeps in `dispatch` so a generation stays
1286    /// in-flight long enough to prove the pool keeps `/healthz`
1287    /// answering while a job runs.
1288    struct SlowEngine {
1289        inner: SyntheticEngine,
1290        delay: std::time::Duration,
1291    }
1292
1293    impl Engine for SlowEngine {
1294        fn name(&self) -> &'static str {
1295            "slow"
1296        }
1297        fn capabilities(&self) -> EngineCapabilities {
1298            self.inner.capabilities()
1299        }
1300        fn dispatch(&self, model: &str, task: Task) -> anyhow::Result<TaskResult> {
1301            std::thread::sleep(self.delay);
1302            self.inner.dispatch(model, task)
1303        }
1304    }
1305
1306    fn synthetic_model_of(id: &str, kind: TaskKind) -> CatalogModel {
1307        CatalogModel {
1308            kind,
1309            ..synthetic_model(id)
1310        }
1311    }
1312
1313    /// A catalog with one synthetic model per kind, so every local
1314    /// endpoint has a default to resolve.
1315    fn multi_kind_catalog() -> Catalog {
1316        Catalog {
1317            models: vec![
1318                synthetic_model_of("img", TaskKind::Image),
1319                synthetic_model_of("chat", TaskKind::Llm),
1320                synthetic_model_of("tts", TaskKind::AudioTts),
1321                synthetic_model_of("stt", TaskKind::AudioStt),
1322                synthetic_model_of("vid", TaskKind::Video),
1323            ],
1324            ..Default::default()
1325        }
1326    }
1327
1328    fn synthetic_model(id: &str) -> CatalogModel {
1329        CatalogModel {
1330            id: id.into(),
1331            display_name: id.into(),
1332            kind: TaskKind::Image,
1333            vram_gb_estimate: 0.0,
1334            description: None,
1335            source: ModelSource {
1336                engine: ModelEngine::Synthetic,
1337                files: vec![],
1338                cli_defaults: ModelCliDefaults {
1339                    cfg_scale: 1.0,
1340                    steps: 4,
1341                    width: 64,
1342                    height: 64,
1343                    ..Default::default()
1344                },
1345            },
1346            enabled: true,
1347            origin: "local".into(),
1348            exclusive_group: None,
1349        }
1350    }
1351
1352    const TEST_TOKEN: &str = "test-token-0123456789abcdef";
1353
1354    struct Harness {
1355        url: String,
1356        observers: WorkerObservers,
1357        host: crate::host::ModelHost,
1358        services: ModelServices,
1359        stop: Arc<AtomicBool>,
1360        handle: Option<std::thread::JoinHandle<()>>,
1361    }
1362
1363    impl Harness {
1364        fn start(catalog: Catalog) -> Self {
1365            Self::start_with_gate(catalog, JobGate::new())
1366        }
1367
1368        fn start_with_gate(catalog: Catalog, gate: JobGate) -> Self {
1369            Self::start_full(catalog, gate, 20.0)
1370        }
1371
1372        /// Start with a device reporting `free_gib` of free memory.
1373        fn start_with_free(catalog: Catalog, free_gib: f32) -> Self {
1374            Self::start_full(catalog, JobGate::new(), free_gib)
1375        }
1376
1377        fn start_full(catalog: Catalog, gate: JobGate, free_gib: f32) -> Self {
1378            let engine: Arc<dyn Engine> = Arc::new(SyntheticEngine::new());
1379            let observers = WorkerObservers::default();
1380            let catalog = Arc::new(Mutex::new(catalog));
1381            let host = crate::host::ModelHost::new(
1382                catalog.clone(),
1383                Arc::new(crate::test_support::InstantRuntime),
1384                Arc::new(crate::test_support::FixedProbe(free_gib)),
1385                crate::residency::Residency::load_for_serving(None),
1386            );
1387            let services = ModelServices::new(host.clone());
1388            services
1389                .stream_port
1390                .store(4798, std::sync::atomic::Ordering::SeqCst);
1391            let api = LocalApi::bind(
1392                "127.0.0.1:0",
1393                engine,
1394                catalog,
1395                None,
1396                observers.clone(),
1397                TEST_TOKEN.to_string(),
1398                gate.clone(),
1399                None,
1400                services.clone(),
1401            )
1402            .unwrap();
1403            let url = api.url();
1404            let stop = Arc::new(AtomicBool::new(false));
1405            let stop_thread = stop.clone();
1406            let handle = std::thread::spawn(move || api.serve(&stop_thread));
1407            Harness {
1408                url,
1409                observers,
1410                host,
1411                services,
1412                stop,
1413                handle: Some(handle),
1414            }
1415        }
1416
1417        /// Authed POST builder — what a legitimate local client sends.
1418        fn post(&self, path: &str) -> reqwest::blocking::RequestBuilder {
1419            reqwest::blocking::Client::new()
1420                .post(format!("{}{}", self.url, path))
1421                .bearer_auth(TEST_TOKEN)
1422        }
1423
1424        /// Authed GET.
1425        fn get(&self, path: &str) -> reqwest::blocking::RequestBuilder {
1426            reqwest::blocking::Client::new()
1427                .get(format!("{}{}", self.url, path))
1428                .bearer_auth(TEST_TOKEN)
1429        }
1430    }
1431
1432    impl Drop for Harness {
1433        fn drop(&mut self) {
1434            self.stop.store(true, Ordering::Relaxed);
1435            if let Some(handle) = self.handle.take() {
1436                let _ = handle.join();
1437            }
1438        }
1439    }
1440
1441    fn test_host(catalog: &Arc<Mutex<Catalog>>) -> crate::host::ModelHost {
1442        crate::host::ModelHost::new(
1443            catalog.clone(),
1444            Arc::new(crate::test_support::InstantRuntime),
1445            Arc::new(crate::test_support::FixedProbe(20.0)),
1446            crate::residency::Residency::load_for_serving(None),
1447        )
1448    }
1449
1450    fn seeded_catalog() -> Catalog {
1451        Catalog {
1452            models: vec![synthetic_model("synthetic-img")],
1453            ..Default::default()
1454        }
1455    }
1456
1457    #[test]
1458    fn post_image_returns_image_bytes_and_records_job() {
1459        let h = Harness::start(seeded_catalog());
1460
1461        let res = h
1462            .post("/image")
1463            .json(&serde_json::json!({ "prompt": "a blue bird" }))
1464            .send()
1465            .unwrap();
1466        assert_eq!(res.status(), 200);
1467        assert_eq!(res.headers()["content-type"], "image/webp");
1468        let bytes = res.bytes().unwrap();
1469        assert!(!bytes.is_empty());
1470
1471        assert_eq!(h.observers.local_jobs.lock().len(), 1);
1472    }
1473
1474    #[test]
1475    fn post_image_honours_requested_ext() {
1476        let h = Harness::start(seeded_catalog());
1477        let res = h
1478            .post("/image")
1479            .json(&serde_json::json!({ "prompt": "x", "ext": "png" }))
1480            .send()
1481            .unwrap();
1482        assert_eq!(res.status(), 200);
1483        assert_eq!(res.headers()["content-type"], "image/png");
1484    }
1485
1486    #[test]
1487    fn get_models_lists_catalog() {
1488        let h = Harness::start(seeded_catalog());
1489        let body = h.get("/models").send().unwrap().text().unwrap();
1490        assert!(body.contains("synthetic-img"));
1491    }
1492
1493    fn json(res: reqwest::blocking::Response) -> (u16, serde_json::Value) {
1494        let status = res.status().as_u16();
1495        (status, res.json().unwrap())
1496    }
1497
1498    fn wait_state(h: &Harness, id: &str, want: &str) -> serde_json::Value {
1499        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
1500        loop {
1501            let (status, body) = json(h.get(&format!("/models/{id}/state")).send().unwrap());
1502            assert_eq!(status, 200, "{body}");
1503            if body["state"] == want {
1504                return body;
1505            }
1506            assert!(
1507                std::time::Instant::now() < deadline,
1508                "never reached {want}: {body}"
1509            );
1510            std::thread::sleep(std::time::Duration::from_millis(10));
1511        }
1512    }
1513
1514    #[test]
1515    fn get_models_carries_state_and_residency() {
1516        let h = Harness::start(seeded_catalog());
1517        let (status, body) = json(h.get("/models").send().unwrap());
1518        assert_eq!(status, 200);
1519        let first = &body.as_array().unwrap()[0];
1520        assert_eq!(first["state"], "unloaded");
1521        assert_eq!(first["resident"], false);
1522        assert!(first["id"].is_string(), "catalogue fields stay: {first}");
1523    }
1524
1525    #[test]
1526    fn load_reaches_loaded_and_marks_resident() {
1527        let h = Harness::start(seeded_catalog());
1528        let (status, body) = json(h.post("/models/synthetic-img/load").send().unwrap());
1529        assert!(status == 202 || status == 200, "{status} {body}");
1530        assert_eq!(body["id"], "synthetic-img");
1531        assert_eq!(body["resident"], true);
1532        let body = wait_state(&h, "synthetic-img", "loaded");
1533        assert!(body["since"].is_string());
1534        let (status, _) = json(h.post("/models/synthetic-img/load").send().unwrap());
1535        assert_eq!(status, 200, "loading a loaded model is a no-op");
1536    }
1537
1538    #[test]
1539    fn unload_frees_and_clears_residency() {
1540        let h = Harness::start(seeded_catalog());
1541        h.post("/models/synthetic-img/load").send().unwrap();
1542        wait_state(&h, "synthetic-img", "loaded");
1543        let (status, body) = json(h.post("/models/synthetic-img/unload").send().unwrap());
1544        assert!(status == 202 || status == 200, "{status} {body}");
1545        assert_eq!(body["resident"], false);
1546        wait_state(&h, "synthetic-img", "unloaded");
1547        let (status, _) = json(h.post("/models/synthetic-img/unload").send().unwrap());
1548        assert_eq!(status, 200, "unloading an unloaded model is a no-op");
1549    }
1550
1551    #[test]
1552    fn a_load_that_does_not_fit_is_a_409_with_the_numbers() {
1553        let mut catalog = seeded_catalog();
1554        catalog.models[0].vram_gb_estimate = 8.0;
1555        let h = Harness::start_with_free(catalog, 4.0);
1556        let (status, body) = json(h.post("/models/synthetic-img/load").send().unwrap());
1557        assert_eq!(status, 409, "{body}");
1558        assert_eq!(body["error"], "insufficient_memory");
1559        assert_eq!(body["neededGib"], 8.0);
1560        assert_eq!(body["freeGib"], 4.0);
1561        assert!(body["marginGib"].is_number());
1562        wait_state(&h, "synthetic-img", "unloaded");
1563    }
1564
1565    #[test]
1566    fn unknown_models_are_404_on_every_lifecycle_route() {
1567        let h = Harness::start(seeded_catalog());
1568        for res in [
1569            h.get("/models/nope/state").send().unwrap(),
1570            h.post("/models/nope/load").send().unwrap(),
1571            h.post("/models/nope/unload").send().unwrap(),
1572        ] {
1573            let (status, body) = json(res);
1574            assert_eq!(status, 404);
1575            assert_eq!(body["error"], "unknown_model");
1576        }
1577    }
1578
1579    #[test]
1580    fn a_disabled_model_cannot_be_loaded() {
1581        let mut catalog = seeded_catalog();
1582        catalog.models[0].enabled = false;
1583        let h = Harness::start(catalog);
1584        let (status, body) = json(h.post("/models/synthetic-img/load").send().unwrap());
1585        assert_eq!(status, 400);
1586        assert_eq!(body["error"], "model_disabled");
1587    }
1588
1589    #[test]
1590    fn lifecycle_routes_need_the_token() {
1591        let h = Harness::start(seeded_catalog());
1592        let res = reqwest::blocking::Client::new()
1593            .post(format!("{}/models/synthetic-img/load", h.url))
1594            .send()
1595            .unwrap();
1596        assert_eq!(res.status(), 401);
1597        wait_state(&h, "synthetic-img", "unloaded");
1598    }
1599
1600    #[test]
1601    fn deleting_a_loaded_model_unloads_it_first() {
1602        let h = Harness::start(seeded_catalog());
1603        h.post("/models/synthetic-img/load").send().unwrap();
1604        wait_state(&h, "synthetic-img", "loaded");
1605        let res = reqwest::blocking::Client::new()
1606            .delete(format!("{}/models/synthetic-img", h.url))
1607            .bearer_auth(TEST_TOKEN)
1608            .send()
1609            .unwrap();
1610        assert_eq!(res.status(), 200);
1611        let unloaded = h.host.wait_for(
1612            "synthetic-img",
1613            |s| *s == crate::lifecycle::ModelState::Unloaded,
1614            std::time::Duration::from_secs(5),
1615        );
1616        assert!(unloaded.is_some(), "weights freed after delete");
1617        assert_eq!(h.host.loaded_gib(), 0.0);
1618    }
1619
1620    fn llm_catalog() -> Catalog {
1621        Catalog {
1622            models: vec![synthetic_model_of("chat-llm", TaskKind::Llm)],
1623            ..Default::default()
1624        }
1625    }
1626
1627    fn chat(h: &Harness, body: serde_json::Value) -> (u16, serde_json::Value) {
1628        json(h.post("/v1/chat/completions").json(&body).send().unwrap())
1629    }
1630
1631    #[test]
1632    fn chat_is_served_on_the_lane_of_a_loaded_model() {
1633        let h = Harness::start(llm_catalog());
1634        h.post("/models/chat-llm/load").send().unwrap();
1635        wait_state(&h, "chat-llm", "loaded");
1636        let (status, body) = chat(
1637            &h,
1638            serde_json::json!({ "model": "chat-llm", "messages": [{ "role": "user", "content": "hi" }] }),
1639        );
1640        assert_eq!(status, 200, "{body}");
1641        assert_eq!(body["choices"][0]["message"]["content"], "resident:hi");
1642    }
1643
1644    #[test]
1645    fn chat_uses_the_default_llm_when_it_is_loaded() {
1646        let h = Harness::start(llm_catalog());
1647        h.post("/models/chat-llm/load").send().unwrap();
1648        wait_state(&h, "chat-llm", "loaded");
1649        let (_, body) = chat(
1650            &h,
1651            serde_json::json!({ "messages": [{ "role": "user", "content": "yo" }] }),
1652        );
1653        assert_eq!(body["choices"][0]["message"]["content"], "resident:yo");
1654    }
1655
1656    #[test]
1657    fn chat_template_kwargs_reach_the_model() {
1658        let h = Harness::start(llm_catalog());
1659        h.post("/models/chat-llm/load").send().unwrap();
1660        wait_state(&h, "chat-llm", "loaded");
1661        let (_, body) = chat(
1662            &h,
1663            serde_json::json!({
1664                "messages": [{ "role": "user", "content": "hi" }],
1665                "chat_template_kwargs": { "enable_thinking": false },
1666            }),
1667        );
1668        assert_eq!(
1669            body["kwargs"],
1670            serde_json::json!({ "enable_thinking": false })
1671        );
1672    }
1673
1674    #[test]
1675    fn chat_on_an_unloaded_model_runs_as_a_transient_job() {
1676        let h = Harness::start(llm_catalog());
1677        let (status, body) = chat(
1678            &h,
1679            serde_json::json!({ "model": "chat-llm", "messages": [{ "role": "user", "content": "hi" }] }),
1680        );
1681        assert_eq!(status, 200, "{body}");
1682        let content = body["choices"][0]["message"]["content"].as_str().unwrap();
1683        assert!(!content.starts_with("resident:"), "{content}");
1684    }
1685
1686    #[test]
1687    fn a_resident_chat_is_recorded_as_a_local_job_and_skips_the_job_gate() {
1688        let gate = JobGate::new();
1689        let h = Harness::start_with_gate(llm_catalog(), gate.clone());
1690        h.post("/models/chat-llm/load").send().unwrap();
1691        wait_state(&h, "chat-llm", "loaded");
1692        let _held = gate.try_reserve().expect("a transient job holds the gate");
1693        let (status, _) = chat(
1694            &h,
1695            serde_json::json!({ "model": "chat-llm", "messages": [{ "role": "user", "content": "lane" }] }),
1696        );
1697        assert_eq!(status, 200, "a loaded model serves on its own lane");
1698        let jobs = h.observers.local_jobs.lock().clone();
1699        let last = jobs.front().expect("recorded");
1700        assert_eq!(last.model, "chat-llm");
1701        assert_eq!(last.prompt, "lane");
1702    }
1703
1704    fn stream_catalog() -> Catalog {
1705        let mut stt = synthetic_model_of("stt-a", TaskKind::AudioStt);
1706        stt.source.engine = crate::types::ModelEngine::Parakeet;
1707        Catalog {
1708            models: vec![stt, synthetic_model_of("chat-llm", TaskKind::Llm)],
1709            ..Default::default()
1710        }
1711    }
1712
1713    #[test]
1714    fn stream_tokens_are_minted_for_streaming_models() {
1715        let h = Harness::start(stream_catalog());
1716        let (status, body) = json(
1717            h.post("/stream-tokens")
1718                .json(&serde_json::json!({ "model": "stt-a", "ttlSecs": 600 }))
1719                .send()
1720                .unwrap(),
1721        );
1722        assert_eq!(status, 200, "{body}");
1723        let token = body["token"].as_str().unwrap();
1724        assert_eq!(token.len(), 64);
1725        assert_eq!(body["model"], "stt-a");
1726        assert_eq!(body["port"], 4798);
1727        assert_eq!(body["path"], "/transcribe");
1728        assert!(body["expiresAt"].is_string());
1729        assert_eq!(
1730            h.services.tokens.check(token, chrono::Utc::now()),
1731            Ok("stt-a".to_string()),
1732            "the listener accepts it"
1733        );
1734    }
1735
1736    #[test]
1737    fn stream_tokens_are_refused_for_other_models() {
1738        let h = Harness::start(stream_catalog());
1739        let (status, body) = json(
1740            h.post("/stream-tokens")
1741                .json(&serde_json::json!({ "model": "chat-llm" }))
1742                .send()
1743                .unwrap(),
1744        );
1745        assert_eq!(status, 400);
1746        assert_eq!(body["error"], "not_a_stream_model");
1747        let (status, body) = json(
1748            h.post("/stream-tokens")
1749                .json(&serde_json::json!({ "model": "nope" }))
1750                .send()
1751                .unwrap(),
1752        );
1753        assert_eq!(status, 404);
1754        assert_eq!(body["error"], "unknown_model");
1755    }
1756
1757    #[test]
1758    fn stream_tokens_need_a_running_listener() {
1759        let mut h = Harness::start(stream_catalog());
1760        h.services
1761            .stream_port
1762            .store(0, std::sync::atomic::Ordering::SeqCst);
1763        let (status, body) = json(
1764            h.post("/stream-tokens")
1765                .json(&serde_json::json!({ "model": "stt-a" }))
1766                .send()
1767                .unwrap(),
1768        );
1769        assert_eq!(status, 503);
1770        assert_eq!(body["error"], "stream_listener_down");
1771        let _ = &mut h;
1772    }
1773
1774    #[test]
1775    fn stream_tokens_need_the_install_token() {
1776        let h = Harness::start(stream_catalog());
1777        let res = reqwest::blocking::Client::new()
1778            .post(format!("{}/stream-tokens", h.url))
1779            .json(&serde_json::json!({ "model": "stt-a" }))
1780            .send()
1781            .unwrap();
1782        assert_eq!(res.status(), 401);
1783    }
1784
1785    #[test]
1786    fn post_models_adds_a_model_then_lists_it() {
1787        let h = Harness::start(seeded_catalog());
1788        let res = h
1789            .post("/models")
1790            .json(&synthetic_model("added-model"))
1791            .send()
1792            .unwrap();
1793        assert_eq!(res.status(), 200);
1794
1795        let body = h.get("/models").send().unwrap().text().unwrap();
1796        assert!(body.contains("added-model"));
1797    }
1798
1799    #[test]
1800    fn unknown_model_is_a_400() {
1801        let h = Harness::start(seeded_catalog());
1802        let res = h
1803            .post("/image")
1804            .json(&serde_json::json!({ "prompt": "x", "model": "nope" }))
1805            .send()
1806            .unwrap();
1807        assert_eq!(res.status(), 400);
1808    }
1809
1810    #[test]
1811    fn invalid_json_is_a_400() {
1812        let h = Harness::start(seeded_catalog());
1813        let res = h
1814            .post("/image")
1815            .body("not json")
1816            .header("content-type", "application/json")
1817            .send()
1818            .unwrap();
1819        assert_eq!(res.status(), 400);
1820    }
1821
1822    #[test]
1823    fn healthz_reports_a_runtime_snapshot() {
1824        let h = Harness::start(seeded_catalog());
1825        let body: serde_json::Value = reqwest::blocking::get(format!("{}/healthz", h.url))
1826            .unwrap()
1827            .json()
1828            .unwrap();
1829        assert_eq!(body["ok"], true);
1830        assert_eq!(body["version"], crate::AGENT_VERSION);
1831        assert_eq!(body["busy"], false);
1832        assert_eq!(body["engine"], "synthetic");
1833        // No secrets / prompts leak into the unauthenticated snapshot.
1834        let raw = serde_json::to_string(&body).unwrap();
1835        assert!(
1836            !raw.contains(TEST_TOKEN),
1837            "healthz must not carry the token"
1838        );
1839    }
1840
1841    #[test]
1842    fn healthz_surfaces_gpu_runtime_when_probed() {
1843        let h = Harness::start(seeded_catalog());
1844        // Simulate the startup probe having found a missing runtime.
1845        crate::runtime::set_gpu_runtime_status(
1846            &h.observers,
1847            Err(anyhow::anyhow!(
1848                "Vulkan runtime not available: install libvulkan1"
1849            )),
1850        );
1851        let body: serde_json::Value = reqwest::blocking::get(format!("{}/healthz", h.url))
1852            .unwrap()
1853            .json()
1854            .unwrap();
1855        assert_eq!(body["gpuRuntime"]["ok"], false);
1856        assert!(body["gpuRuntime"]["detail"]
1857            .as_str()
1858            .unwrap()
1859            .contains("libvulkan1"));
1860    }
1861
1862    // -----------------------------------------------------------------
1863    // Generic per-kind endpoints (chat / tts / stt / video) — the local
1864    // API serves every modality the worker's engines support, not just
1865    // image.
1866    // -----------------------------------------------------------------
1867
1868    #[test]
1869    fn chat_completions_returns_an_openai_shaped_body() {
1870        let h = Harness::start(multi_kind_catalog());
1871        let res = h
1872            .post("/v1/chat/completions")
1873            .json(&serde_json::json!({
1874                "messages": [{"role": "user", "content": "hello there"}],
1875                "max_tokens": 16
1876            }))
1877            .send()
1878            .unwrap();
1879        assert_eq!(res.status(), 200);
1880        assert_eq!(res.headers()["content-type"], "application/json");
1881        let body: serde_json::Value = res.json().unwrap();
1882        // The synthetic engine emits a chat.completion-shaped object.
1883        assert!(
1884            body.get("choices").is_some() || body.get("object").is_some(),
1885            "expected an OpenAI-ish body, got: {body}"
1886        );
1887        // The local job was recorded.
1888        assert!(h
1889            .observers
1890            .local_jobs
1891            .lock()
1892            .iter()
1893            .any(|j| j.kind == TaskKind::Llm));
1894    }
1895
1896    #[test]
1897    fn tts_returns_audio_bytes() {
1898        let h = Harness::start(multi_kind_catalog());
1899        let res = h
1900            .post("/tts")
1901            .json(&serde_json::json!({ "text": "read this aloud" }))
1902            .send()
1903            .unwrap();
1904        assert_eq!(res.status(), 200);
1905        assert_eq!(res.headers()["content-type"], "audio/wav");
1906        assert!(!res.bytes().unwrap().is_empty());
1907    }
1908
1909    #[test]
1910    fn stt_returns_a_transcript_json() {
1911        let h = Harness::start(multi_kind_catalog());
1912        let res = h
1913            .post("/stt")
1914            .json(&serde_json::json!({ "inputUrl": "https://example.com/a.wav" }))
1915            .send()
1916            .unwrap();
1917        assert_eq!(res.status(), 200);
1918        assert_eq!(res.headers()["content-type"], "application/json");
1919    }
1920
1921    #[test]
1922    fn video_returns_bytes() {
1923        let h = Harness::start(multi_kind_catalog());
1924        let res = h
1925            .post("/video")
1926            .json(&serde_json::json!({ "prompt": "a tiny dragon" }))
1927            .send()
1928            .unwrap();
1929        assert_eq!(res.status(), 200);
1930        assert!(!res.bytes().unwrap().is_empty());
1931    }
1932
1933    #[test]
1934    fn chat_without_an_llm_model_is_a_400() {
1935        // Only an image model in the catalog: a chat request has no
1936        // model to resolve and must say so (400), not 500.
1937        let h = Harness::start(seeded_catalog());
1938        let res = h
1939            .post("/v1/chat/completions")
1940            .json(&serde_json::json!({
1941                "messages": [{"role": "user", "content": "hi"}]
1942            }))
1943            .send()
1944            .unwrap();
1945        assert_eq!(res.status(), 400);
1946        assert!(res.text().unwrap().contains("llm"));
1947    }
1948
1949    #[test]
1950    fn chat_endpoint_respects_the_busy_gate() {
1951        let gate = JobGate::new();
1952        let h = Harness::start_with_gate(multi_kind_catalog(), gate.clone());
1953        let _held = gate.try_reserve().unwrap();
1954        let res = h
1955            .post("/v1/chat/completions")
1956            .json(&serde_json::json!({ "messages": [{"role":"user","content":"x"}] }))
1957            .send()
1958            .unwrap();
1959        assert_eq!(res.status(), 503);
1960    }
1961
1962    #[test]
1963    fn jobs_endpoint_reports_after_generation() {
1964        let h = Harness::start(seeded_catalog());
1965        h.post("/image")
1966            .json(&serde_json::json!({ "prompt": "x" }))
1967            .send()
1968            .unwrap();
1969        let body = h.get("/jobs").send().unwrap().text().unwrap();
1970        assert!(body.contains("\"completed\""));
1971        assert!(body.contains("synthetic-img"));
1972    }
1973
1974    // -----------------------------------------------------------------
1975    // Auth gate: every route except GET /healthz requires the bearer
1976    // token; Host / Origin headers must be loopback.  These pin the
1977    // CSRF / DNS-rebinding / local-user defences end-to-end through a
1978    // real socket.
1979    // -----------------------------------------------------------------
1980
1981    #[test]
1982    fn routes_reject_requests_without_a_token() {
1983        let h = Harness::start(seeded_catalog());
1984        let client = reqwest::blocking::Client::new();
1985        let cases: Vec<(reqwest::blocking::RequestBuilder, &str)> = vec![
1986            (
1987                client
1988                    .post(format!("{}/image", h.url))
1989                    .json(&serde_json::json!({ "prompt": "x" })),
1990                "POST /image",
1991            ),
1992            (client.get(format!("{}/models", h.url)), "GET /models"),
1993            (
1994                client
1995                    .post(format!("{}/models", h.url))
1996                    .json(&synthetic_model("evil")),
1997                "POST /models",
1998            ),
1999            (
2000                client.delete(format!("{}/models/synthetic-img", h.url)),
2001                "DELETE /models",
2002            ),
2003            (client.get(format!("{}/jobs", h.url)), "GET /jobs"),
2004        ];
2005        for (req, name) in cases {
2006            let res = req.send().unwrap();
2007            assert_eq!(res.status(), 401, "{name} must require the token");
2008            let body = res.text().unwrap();
2009            assert!(
2010                body.contains("local-api.json"),
2011                "{name}: the 401 must point at the discovery file, got: {body}"
2012            );
2013        }
2014        // Nothing was mutated by the unauthenticated attempts.
2015        let body = h.get("/models").send().unwrap().text().unwrap();
2016        assert!(!body.contains("evil"));
2017        assert!(body.contains("synthetic-img"));
2018    }
2019
2020    #[test]
2021    fn routes_reject_a_wrong_token() {
2022        let h = Harness::start(seeded_catalog());
2023        let res = reqwest::blocking::Client::new()
2024            .get(format!("{}/models", h.url))
2025            .bearer_auth("wrong-token")
2026            .send()
2027            .unwrap();
2028        assert_eq!(res.status(), 401);
2029    }
2030
2031    #[test]
2032    fn daemon_routes_need_daemon_control() {
2033        let h = Harness::start(multi_kind_catalog());
2034        let resp = h.get("/daemon/status").send().unwrap();
2035        assert_eq!(resp.status(), 503);
2036        let body: serde_json::Value = resp.json().unwrap();
2037        assert_eq!(body["error"], "daemon_control_unavailable");
2038    }
2039
2040    #[test]
2041    fn daemon_routes_need_the_token() {
2042        let h = Harness::start(multi_kind_catalog());
2043        let resp = reqwest::blocking::Client::new()
2044            .get(format!("{}/daemon/status", h.url))
2045            .send()
2046            .unwrap();
2047        assert_eq!(resp.status(), 401);
2048    }
2049
2050    #[test]
2051    fn an_unknown_daemon_route_is_not_found() {
2052        let daemon = crate::test_support::DaemonHarness::start();
2053        let resp = reqwest::blocking::Client::new()
2054            .get(format!("{}/daemon/nope", daemon.url))
2055            .bearer_auth(crate::test_support::HARNESS_TOKEN)
2056            .send()
2057            .unwrap();
2058        assert_eq!(resp.status(), 404);
2059    }
2060
2061    #[test]
2062    fn a_malformed_config_body_is_a_bad_request() {
2063        let daemon = crate::test_support::DaemonHarness::start();
2064        let resp = reqwest::blocking::Client::new()
2065            .put(format!("{}/daemon/config", daemon.url))
2066            .bearer_auth(crate::test_support::HARNESS_TOKEN)
2067            .body("{")
2068            .send()
2069            .unwrap();
2070        assert_eq!(resp.status(), 400);
2071    }
2072
2073    #[test]
2074    fn the_models_listing_carries_since_and_the_failure() {
2075        let daemon = crate::test_support::DaemonHarness::start();
2076        let models = daemon.client().models().unwrap();
2077        assert!(models.iter().all(|m| m.since.is_some() && m.loadable));
2078        assert!(models.iter().all(|m| m.error.is_none()));
2079    }
2080
2081    #[test]
2082    fn job_routes_and_query_params_parse() {
2083        assert_eq!(job_route("/jobs/local-1/log", "/log"), Some("local-1"));
2084        assert_eq!(job_route("/jobs//log", "/log"), None);
2085        assert_eq!(job_route("/jobs/a/b/log", "/log"), None);
2086        assert_eq!(query_param("/daemon/logs?after=12", "after"), Some("12"));
2087        assert_eq!(query_param("/daemon/logs?x=1&after=3", "after"), Some("3"));
2088        assert_eq!(query_param("/daemon/logs?afterx=3", "after"), None);
2089        assert_eq!(query_param("/daemon/logs", "after"), None);
2090    }
2091
2092    #[test]
2093    fn healthz_needs_no_token() {
2094        let h = Harness::start(seeded_catalog());
2095        let res = reqwest::blocking::get(format!("{}/healthz", h.url)).unwrap();
2096        assert_eq!(res.status(), 200);
2097    }
2098
2099    #[test]
2100    fn healthz_answers_while_a_generation_is_in_flight() {
2101        // The whole point of the worker pool: a slow (~400 ms) job must
2102        // not block liveness / cheap routes.  On the old
2103        // single-threaded loop this `/healthz` would queue behind the
2104        // generation and only answer after it finished.
2105        let engine: Arc<dyn Engine> = Arc::new(SlowEngine {
2106            inner: SyntheticEngine::new(),
2107            delay: std::time::Duration::from_millis(400),
2108        });
2109        let observers = WorkerObservers::default();
2110        let catalog = Arc::new(Mutex::new(seeded_catalog()));
2111        let api = LocalApi::bind(
2112            "127.0.0.1:0",
2113            engine,
2114            catalog.clone(),
2115            None,
2116            observers,
2117            TEST_TOKEN.to_string(),
2118            JobGate::new(),
2119            None,
2120            ModelServices::new(test_host(&catalog)),
2121        )
2122        .unwrap();
2123        let url = api.url();
2124        let stop = Arc::new(AtomicBool::new(false));
2125        let stop_thread = stop.clone();
2126        let handle = std::thread::spawn(move || api.serve(&stop_thread));
2127
2128        // Kick off the slow generation on a background thread.
2129        let gen_url = url.clone();
2130        let gen = std::thread::spawn(move || {
2131            reqwest::blocking::Client::new()
2132                .post(format!("{gen_url}/image"))
2133                .bearer_auth(TEST_TOKEN)
2134                .json(&serde_json::json!({ "prompt": "slow" }))
2135                .timeout(std::time::Duration::from_secs(5))
2136                .send()
2137                .unwrap()
2138                .status()
2139                .as_u16()
2140        });
2141
2142        // Give the generation time to occupy a worker, then time a
2143        // /healthz: it must answer well within the generation's 400 ms.
2144        std::thread::sleep(std::time::Duration::from_millis(100));
2145        let start = std::time::Instant::now();
2146        let health = reqwest::blocking::get(format!("{url}/healthz")).unwrap();
2147        let elapsed = start.elapsed();
2148        assert_eq!(health.status(), 200);
2149        assert!(
2150            elapsed < std::time::Duration::from_millis(250),
2151            "healthz blocked behind the generation ({elapsed:?}); the pool isn't concurrent"
2152        );
2153        // While the job runs, healthz reports busy=true.
2154        let body: serde_json::Value = health.json().unwrap();
2155        assert_eq!(body["busy"], true, "a running job must show busy=true");
2156
2157        assert_eq!(gen.join().unwrap(), 200, "the generation still succeeds");
2158        stop.store(true, Ordering::Relaxed);
2159        let _ = handle.join();
2160    }
2161
2162    #[test]
2163    fn non_loopback_host_header_is_forbidden_even_with_a_token() {
2164        // The DNS-rebinding shape: the TCP connection reaches loopback
2165        // but the browser's Host header names the attacker's domain.
2166        let h = Harness::start(seeded_catalog());
2167        let res = h
2168            .get("/models")
2169            .header("host", "evil.example:4787")
2170            .send()
2171            .unwrap();
2172        assert_eq!(res.status(), 403);
2173        assert!(res.text().unwrap().contains("Host"));
2174    }
2175
2176    #[test]
2177    fn cross_site_origin_is_forbidden_even_with_a_token() {
2178        // The CSRF shape: a browser always attaches the page's Origin
2179        // to cross-site POSTs.
2180        let h = Harness::start(seeded_catalog());
2181        let res = h
2182            .post("/image")
2183            .header("origin", "https://evil.example")
2184            .json(&serde_json::json!({ "prompt": "x" }))
2185            .send()
2186            .unwrap();
2187        assert_eq!(res.status(), 403);
2188        assert!(res.text().unwrap().contains("Origin"));
2189    }
2190
2191    #[test]
2192    fn loopback_origin_is_allowed() {
2193        // A local web app (e.g. a dashboard on localhost:5173) is a
2194        // legitimate browser client.
2195        let h = Harness::start(seeded_catalog());
2196        let res = h
2197            .get("/models")
2198            .header("origin", "http://localhost:5173")
2199            .send()
2200            .unwrap();
2201        assert_eq!(res.status(), 200);
2202    }
2203
2204    #[test]
2205    fn oversized_body_is_a_413() {
2206        let h = Harness::start(seeded_catalog());
2207        let big = "x".repeat(MAX_BODY_BYTES + 1);
2208        let res = h.post("/image").body(big).send().unwrap();
2209        assert_eq!(res.status(), 413);
2210    }
2211
2212    #[test]
2213    fn body_at_the_cap_is_still_read() {
2214        // Boundary: exactly MAX_BODY_BYTES must not be rejected as too
2215        // large (it fails later as bad JSON, which is the point — the
2216        // size gate stayed out of the way).
2217        let h = Harness::start(seeded_catalog());
2218        let exact = "x".repeat(MAX_BODY_BYTES);
2219        let res = h.post("/image").body(exact).send().unwrap();
2220        assert_eq!(res.status(), 400);
2221    }
2222
2223    #[test]
2224    fn bind_refuses_an_empty_token() {
2225        let engine: Arc<dyn Engine> = Arc::new(SyntheticEngine::new());
2226        let catalog = Arc::new(Mutex::new(seeded_catalog()));
2227        let err = LocalApi::bind(
2228            "127.0.0.1:0",
2229            engine,
2230            catalog.clone(),
2231            None,
2232            WorkerObservers::default(),
2233            String::new(),
2234            JobGate::new(),
2235            None,
2236            ModelServices::new(test_host(&catalog)),
2237        )
2238        .err()
2239        .expect("empty token must be refused")
2240        .to_string();
2241        assert!(err.contains("empty token"), "got: {err}");
2242    }
2243
2244    #[test]
2245    fn post_image_returns_503_when_the_shared_gate_is_held() {
2246        // A studio job (or another local job) holds the one-job gate;
2247        // a concurrent local generation must be refused with 503 +
2248        // Retry-After rather than run a second job on the same GPU.
2249        let gate = JobGate::new();
2250        let h = Harness::start_with_gate(seeded_catalog(), gate.clone());
2251        let reservation = gate.try_reserve().expect("pre-hold the slot");
2252        let res = h
2253            .post("/image")
2254            .json(&serde_json::json!({ "prompt": "x" }))
2255            .send()
2256            .unwrap();
2257        assert_eq!(res.status(), 503);
2258        assert_eq!(res.headers()["retry-after"], "2");
2259
2260        // Once the holder releases, the same request succeeds — proving
2261        // the 503 was the gate, not a broken engine.
2262        drop(reservation);
2263        let res = h
2264            .post("/image")
2265            .json(&serde_json::json!({ "prompt": "x" }))
2266            .send()
2267            .unwrap();
2268        assert_eq!(res.status(), 200);
2269    }
2270
2271    // -----------------------------------------------------------------
2272    // Pure guards.
2273    // -----------------------------------------------------------------
2274
2275    #[test]
2276    fn host_is_loopback_accepts_only_loopback_shapes() {
2277        for ok in [
2278            "127.0.0.1",
2279            "127.0.0.1:4787",
2280            "localhost",
2281            "LOCALHOST:80",
2282            "[::1]",
2283            "[::1]:4787",
2284        ] {
2285            assert!(host_is_loopback(ok), "{ok} should be loopback");
2286        }
2287        for bad in [
2288            "evil.example",
2289            "evil.example:4787",
2290            "127.0.0.1.evil.example",
2291            "192.168.1.10:4787",
2292            "[::2]:4787",
2293            "[::1",
2294            "",
2295        ] {
2296            assert!(!host_is_loopback(bad), "{bad} should be rejected");
2297        }
2298    }
2299
2300    #[test]
2301    fn origin_is_loopback_accepts_only_loopback_origins() {
2302        for ok in [
2303            "http://127.0.0.1:4787",
2304            "http://localhost:5173",
2305            "https://localhost",
2306            "http://[::1]:3000",
2307        ] {
2308            assert!(origin_is_loopback(ok), "{ok} should be allowed");
2309        }
2310        for bad in [
2311            "https://evil.example",
2312            "http://192.168.1.10",
2313            "null",
2314            "file://",
2315            "chrome-extension://abc",
2316            "",
2317        ] {
2318            assert!(!origin_is_loopback(bad), "{bad} should be rejected");
2319        }
2320    }
2321
2322    #[test]
2323    fn deny_reason_orders_host_origin_then_token() {
2324        let t = "tok";
2325        // Bad host wins even when everything else is bad too.
2326        assert!(matches!(
2327            deny_reason(Some("evil.example"), Some("https://evil.example"), None, t),
2328            Some(Denial::Host(_))
2329        ));
2330        // Good host, bad origin.
2331        assert!(matches!(
2332            deny_reason(Some("127.0.0.1"), Some("https://evil.example"), None, t),
2333            Some(Denial::Origin(_))
2334        ));
2335        // Good host + origin, missing token.
2336        assert_eq!(
2337            deny_reason(Some("127.0.0.1"), None, None, t),
2338            Some(Denial::Token)
2339        );
2340        // Malformed authorization schemes are a token failure.
2341        assert_eq!(
2342            deny_reason(None, None, Some("Basic dXNlcjpwdw=="), t),
2343            Some(Denial::Token)
2344        );
2345        // Absent host + origin (curl-style) with the right token passes.
2346        assert_eq!(deny_reason(None, None, Some("Bearer tok"), t), None);
2347        // Lowercase scheme is tolerated.
2348        assert_eq!(deny_reason(None, None, Some("bearer tok"), t), None);
2349    }
2350
2351    #[test]
2352    fn token_matches_is_exact() {
2353        assert!(token_matches("abc", "abc"));
2354        assert!(!token_matches("abd", "abc"));
2355        assert!(!token_matches("ab", "abc"));
2356        assert!(!token_matches("", "abc"));
2357    }
2358
2359    // -----------------------------------------------------------------
2360    // Discovery file.
2361    // -----------------------------------------------------------------
2362
2363    #[test]
2364    fn discovery_file_round_trips_and_is_owner_only() {
2365        let dir = tempfile::tempdir().unwrap();
2366        let path = dir.path().join("local-api.json");
2367        write_discovery_file(&path, "http://127.0.0.1:4787", "tok-123").unwrap();
2368
2369        let parsed: serde_json::Value =
2370            serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
2371        assert_eq!(parsed["url"], "http://127.0.0.1:4787");
2372        assert_eq!(parsed["token"], "tok-123");
2373
2374        #[cfg(unix)]
2375        {
2376            use std::os::unix::fs::PermissionsExt;
2377            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
2378            assert_eq!(
2379                mode & 0o077,
2380                0,
2381                "discovery file carries the token and must be owner-only, got {mode:o}"
2382            );
2383        }
2384
2385        remove_discovery_file(&path);
2386        assert!(!path.exists());
2387        // Idempotent: removing a missing file is quiet.
2388        remove_discovery_file(&path);
2389    }
2390}