1use 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
55pub const MAX_BODY_BYTES: usize = 1024 * 1024;
58
59#[derive(Debug, PartialEq, Eq)]
61enum Denial {
62 Host(String),
64 Origin(String),
66 Token,
68}
69
70fn host_is_loopback(host: &str) -> bool {
73 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
86fn 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
99fn 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
109fn 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
139pub 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 token: String,
149 gate: JobGate,
152 models_root: Option<PathBuf>,
155 services: ModelServices,
157 control: Option<crate::control::DaemonControl>,
160}
161
162#[derive(Clone)]
165pub struct ModelServices {
166 pub host: ModelHost,
167 pub tokens: Arc<StreamTokens>,
168 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
190const 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#[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 #[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 #[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 pub fn with_control(mut self, control: crate::control::DaemonControl) -> Self {
328 self.control = Some(control);
329 self
330 }
331
332 pub fn local_addr(&self) -> SocketAddr {
334 self.addr
335 }
336
337 pub fn url(&self) -> String {
339 format!("http://{}", self.addr)
340 }
341
342 pub fn serve(&self, stop: &AtomicBool) {
353 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 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 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 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 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 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 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 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 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 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 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
1082enum BodyOutcome {
1084 Ok(String),
1085 TooLarge,
1086}
1087
1088fn read_body(request: &mut Request) -> std::io::Result<BodyOutcome> {
1089 if matches!(request.body_length(), Some(len) if len > MAX_BODY_BYTES) {
1091 return Ok(BodyOutcome::TooLarge);
1092 }
1093 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
1116fn 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
1134pub 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
1152pub 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
1186fn 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
1208fn 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
1214fn 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
1220fn 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
1228fn 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 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 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 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 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 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 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 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 #[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 assert!(
1884 body.get("choices").is_some() || body.get("object").is_some(),
1885 "expected an OpenAI-ish body, got: {body}"
1886 );
1887 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 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 #[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 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 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 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 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 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 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 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 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 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 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 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 #[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 assert!(matches!(
2327 deny_reason(Some("evil.example"), Some("https://evil.example"), None, t),
2328 Some(Denial::Host(_))
2329 ));
2330 assert!(matches!(
2332 deny_reason(Some("127.0.0.1"), Some("https://evil.example"), None, t),
2333 Some(Denial::Origin(_))
2334 ));
2335 assert_eq!(
2337 deny_reason(Some("127.0.0.1"), None, None, t),
2338 Some(Denial::Token)
2339 );
2340 assert_eq!(
2342 deny_reason(None, None, Some("Basic dXNlcjpwdw=="), t),
2343 Some(Denial::Token)
2344 );
2345 assert_eq!(deny_reason(None, None, Some("Bearer tok"), t), None);
2347 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 #[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 remove_discovery_file(&path);
2389 }
2390}