1use axum::{
2 extract::{Path, Request, State},
3 http::{header, HeaderMap, HeaderValue, StatusCode},
4 response::{
5 sse::{Event as SseEvent, KeepAlive, Sse},
6 IntoResponse,
7 },
8 routing::{delete, get, patch, post},
9 Json, Router,
10};
11use base64::Engine as _;
12use mold_core::{
13 types::GpuSelection, ActiveGenerationStatus, GenerateRequest, GpuInfo, GpuWorkerState,
14 ModelInfoExtended, ResourceSnapshot, ServerStatus, SseErrorEvent, SseProgressEvent,
15};
16use serde::{Deserialize, Serialize};
17use std::cmp::Reverse;
18use std::convert::Infallible;
19use std::sync::atomic::Ordering;
20use tokio_stream::StreamExt as _;
21use utoipa::OpenApi;
22
23use crate::model_manager;
24use crate::state::{AppState, GenerationJob, SseMessage, SubmitError};
25
26fn submit_error_to_api(e: SubmitError) -> ApiError {
27 match e {
28 SubmitError::Full { pending, capacity } => {
29 ApiError::queue_full(format!("generation queue is full ({pending}/{capacity})"))
30 }
31 SubmitError::Shutdown => ApiError::internal("generation queue shut down"),
32 }
33}
34
35#[derive(Debug, Serialize)]
38pub struct ApiError {
39 pub error: String,
40 pub code: String,
41 #[serde(skip)]
42 status: StatusCode,
43}
44
45impl ApiError {
46 pub fn validation(msg: impl Into<String>) -> Self {
47 Self {
48 error: msg.into(),
49 code: "VALIDATION_ERROR".to_string(),
50 status: StatusCode::UNPROCESSABLE_ENTITY,
51 }
52 }
53
54 pub fn not_found(msg: impl Into<String>) -> Self {
55 Self {
56 error: msg.into(),
57 code: "MODEL_NOT_FOUND".to_string(),
58 status: StatusCode::NOT_FOUND,
59 }
60 }
61
62 pub fn unknown_model(msg: impl Into<String>) -> Self {
63 Self {
64 error: msg.into(),
65 code: "UNKNOWN_MODEL".to_string(),
66 status: StatusCode::BAD_REQUEST,
67 }
68 }
69
70 pub fn inference(msg: impl Into<String>) -> Self {
71 Self {
72 error: msg.into(),
73 code: "INFERENCE_ERROR".to_string(),
74 status: StatusCode::INTERNAL_SERVER_ERROR,
75 }
76 }
77
78 pub fn internal(msg: impl Into<String>) -> Self {
79 Self {
80 error: msg.into(),
81 code: "INTERNAL_ERROR".to_string(),
82 status: StatusCode::INTERNAL_SERVER_ERROR,
83 }
84 }
85
86 pub fn internal_with_status(msg: impl Into<String>, status: StatusCode) -> Self {
87 Self {
88 error: msg.into(),
89 code: "INTERNAL_ERROR".to_string(),
90 status,
91 }
92 }
93
94 pub fn with_code(msg: impl Into<String>, code: impl Into<String>, status: StatusCode) -> Self {
95 Self {
96 error: msg.into(),
97 code: code.into(),
98 status,
99 }
100 }
101
102 pub fn queue_job_not_found(msg: impl Into<String>) -> Self {
103 Self {
104 error: msg.into(),
105 code: "QUEUE_JOB_NOT_FOUND".to_string(),
106 status: StatusCode::NOT_FOUND,
107 }
108 }
109
110 pub fn queue_job_running(msg: impl Into<String>) -> Self {
111 Self {
112 error: msg.into(),
113 code: "QUEUE_JOB_RUNNING".to_string(),
114 status: StatusCode::CONFLICT,
115 }
116 }
117
118 pub fn cancelled(msg: impl Into<String>) -> Self {
123 Self {
124 error: msg.into(),
125 code: "CANCELLED".to_string(),
126 status: StatusCode::from_u16(499).expect("499 is a valid status code"),
127 }
128 }
129
130 pub fn insufficient_memory(msg: impl Into<String>) -> Self {
131 Self {
132 error: msg.into(),
133 code: "INSUFFICIENT_MEMORY".to_string(),
134 status: StatusCode::SERVICE_UNAVAILABLE,
135 }
136 }
137
138 pub fn forbidden(msg: impl Into<String>) -> Self {
139 Self {
140 error: msg.into(),
141 code: "FORBIDDEN".to_string(),
142 status: StatusCode::FORBIDDEN,
143 }
144 }
145
146 pub fn queue_full(msg: impl Into<String>) -> Self {
147 Self {
148 error: msg.into(),
149 code: "QUEUE_FULL".to_string(),
150 status: StatusCode::SERVICE_UNAVAILABLE,
151 }
152 }
153}
154
155impl IntoResponse for ApiError {
156 fn into_response(self) -> axum::response::Response {
157 let status = self.status;
158 if self.code == "QUEUE_FULL" {
160 let mut headers = HeaderMap::new();
161 headers.insert(header::RETRY_AFTER, HeaderValue::from_static("1"));
162 return (status, headers, Json(self)).into_response();
163 }
164 (status, Json(self)).into_response()
165 }
166}
167
168#[cfg(test)]
170use crate::queue::clean_error_message;
171
172#[derive(OpenApi)]
173#[openapi(
174 paths(
175 generate,
176 generate_stream,
177 expand_prompt,
178 list_models,
179 crate::catalog_api::list_loras,
180 load_model,
181 pull_model_endpoint,
182 unload_model,
183 delete_model,
184 server_status,
185 list_queue,
186 patch_queue_job,
187 cancel_queue_job,
188 list_history,
189 delete_history,
190 crate::routes_config::list_config,
191 crate::routes_config::get_config_key,
192 crate::routes_config::put_config_key,
193 crate::routes_config::delete_config_key,
194 crate::routes_config::list_config_profiles,
195 crate::routes_config::put_config_profile,
196 health,
197 capabilities_chain_limits,
198 crate::routes_chain::generate_chain,
199 crate::routes_chain::generate_chain_stream,
200 crate::routes_chain_jobs::create_chain_job,
201 crate::routes_chain_jobs::list_chain_jobs,
202 crate::routes_chain_jobs::get_chain_job,
203 crate::routes_chain_jobs::chain_job_events,
204 crate::routes_chain_jobs::resume_chain_job,
205 crate::routes_chain_jobs::retake_chain_job,
206 crate::routes_chain_jobs::cancel_chain_job,
207 crate::routes_chain_jobs::delete_chain_job,
208 crate::routes_chain_jobs::gc_chain_jobs,
209 crate::routes_chain_jobs::chain_job_stage_preview,
210 ),
211 components(schemas(
212 mold_core::GenerateRequest,
213 mold_core::GenerateResponse,
214 mold_core::ExpandRequest,
215 mold_core::ExpandResponse,
216 mold_core::ImageData,
217 mold_core::OutputFormat,
218 mold_core::ModelInfo,
219 mold_core::LoraInfo,
220 mold_core::ServerStatus,
221 mold_core::ActiveGenerationStatus,
222 mold_core::GpuInfo,
223 mold_core::SseProgressEvent,
224 mold_core::SseCompleteEvent,
225 mold_core::SseErrorEvent,
226 mold_core::ChainRequest,
227 mold_core::ChainResponse,
228 mold_core::ChainStage,
229 mold_core::ChainProgressEvent,
230 mold_core::SseChainCompleteEvent,
231 mold_core::chain_job::ChainJobSummary,
232 mold_core::chain_job::ChainJobStageDetail,
233 mold_core::chain_job::ChainJobDetail,
234 mold_core::chain_job::ChainJobListing,
235 mold_core::chain_job::CreateChainJobResponse,
236 mold_core::chain_job::RetakeRequest,
237 mold_core::chain_job::ChainJobEvent,
238 mold_core::chain_job::FinalizeRecord,
239 mold_core::chain_job::RetakeAmendment,
240 mold_core::chain_job::ChainJobState,
241 mold_core::chain_job::StageState,
242 mold_core::chain_job::RetakeMode,
243 mold_core::chain_job::GcOutcome,
244 ModelInfoExtended,
245 LoadModelBody,
246 UnloadRequest,
247 mold_core::ModelRemovalResponse,
248 mold_core::KeptComponent,
249 QueuePatchRequest,
250 mold_core::HistoryEntry,
251 mold_core::HistoryListing,
252 mold_core::ConfigEntry,
253 mold_core::ConfigListing,
254 mold_core::ConfigProfiles,
255 crate::routes_config::ConfigSetRequest,
256 crate::routes_config::ProfileSetRequest,
257 crate::job_registry::JobEntry,
258 crate::job_registry::QueueListing,
259 crate::chain_limits::ChainLimits,
260 )),
261 tags(
262 (name = "generation", description = "Image generation"),
263 (name = "models", description = "Model management"),
264 (name = "server", description = "Server status and health"),
265 (name = "chain-jobs", description = "Durable chained video jobs"),
266 ),
267 info(
268 title = "mold",
269 description = "Local AI image generation server — FLUX, SD3.5, SD1.5, SDXL, Z-Image, Flux.2, Qwen-Image",
270 version = env!("CARGO_PKG_VERSION"),
271 )
272)]
273pub struct ApiDoc;
274
275pub fn create_router(state: AppState) -> Router {
276 Router::new()
279 .route("/api/generate", post(generate))
280 .route("/api/generate/estimate", post(generate_estimate))
281 .route("/api/generate/stream", post(generate_stream))
282 .route(
283 "/api/generate/chain",
284 post(crate::routes_chain::generate_chain),
285 )
286 .route(
287 "/api/generate/chain/stream",
288 post(crate::routes_chain::generate_chain_stream),
289 )
290 .route(
291 "/api/chain-jobs",
292 post(crate::routes_chain_jobs::create_chain_job)
293 .get(crate::routes_chain_jobs::list_chain_jobs),
294 )
295 .route(
296 "/api/chain-jobs/:id",
297 get(crate::routes_chain_jobs::get_chain_job)
298 .delete(crate::routes_chain_jobs::delete_chain_job),
299 )
300 .route(
301 "/api/chain-jobs/:id/events",
302 get(crate::routes_chain_jobs::chain_job_events),
303 )
304 .route(
305 "/api/chain-jobs/:id/resume",
306 post(crate::routes_chain_jobs::resume_chain_job),
307 )
308 .route(
309 "/api/chain-jobs/:id/retake",
310 post(crate::routes_chain_jobs::retake_chain_job),
311 )
312 .route(
313 "/api/chain-jobs/:id/cancel",
314 post(crate::routes_chain_jobs::cancel_chain_job),
315 )
316 .route(
317 "/api/chain-jobs/gc",
318 post(crate::routes_chain_jobs::gc_chain_jobs),
319 )
320 .route(
321 "/api/chain-jobs/:id/stages/:idx/preview",
322 get(crate::routes_chain_jobs::chain_job_stage_preview),
323 )
324 .route("/api/expand", post(expand_prompt))
325 .route("/api/models", get(list_models))
326 .route("/api/models/:model", delete(delete_model))
327 .route("/api/models/:model/components", get(model_components))
328 .route("/api/loras", get(crate::catalog_api::list_loras))
329 .route("/api/models/load", post(load_model))
330 .route("/api/models/pull", post(pull_model_endpoint))
331 .route("/api/models/unload", delete(unload_model))
332 .route("/api/gallery", get(list_gallery))
333 .route(
334 "/api/gallery/image/:filename",
335 get(get_gallery_image).delete(delete_gallery_image),
336 )
337 .route(
338 "/api/gallery/thumbnail/:filename",
339 get(get_gallery_thumbnail),
340 )
341 .route("/api/gallery/preview/:filename", get(get_gallery_preview))
342 .route("/api/downloads", get(list_downloads).post(create_download))
344 .route("/api/downloads/:id", delete(delete_download))
345 .route("/api/downloads/stream", get(stream_downloads))
346 .route(
348 "/api/catalog/families",
349 get(crate::catalog_api::list_families),
350 )
351 .route(
352 "/api/catalog/search",
353 get(crate::catalog_api::live_search_catalog),
354 )
355 .route(
356 "/api/catalog/installed",
357 get(crate::catalog_api::list_installed_catalog),
358 )
359 .route(
360 "/api/catalog/*id",
361 get(crate::catalog_api::get_catalog_entry)
362 .post(crate::catalog_api::post_catalog_dispatch),
363 )
364 .route("/api/upscale", post(upscale))
365 .route("/api/upscale/stream", post(upscale_stream))
366 .route("/api/resources", get(get_resources))
367 .route("/api/resources/stream", get(get_resources_stream))
368 .route("/api/status", get(server_status))
369 .route("/api/queue", get(list_queue))
370 .route(
371 "/api/queue/:id",
372 patch(patch_queue_job).delete(cancel_queue_job),
373 )
374 .route("/api/history", get(list_history).delete(delete_history))
375 .route("/api/capabilities", get(server_capabilities))
376 .route(
377 "/api/capabilities/chain-limits",
378 get(capabilities_chain_limits),
379 )
380 .route("/api/shutdown", post(shutdown_server))
381 .route("/api/config", get(crate::routes_config::list_config))
383 .route(
384 "/api/config/profiles",
385 get(crate::routes_config::list_config_profiles),
386 )
387 .route(
388 "/api/config/profile",
389 axum::routing::put(crate::routes_config::put_config_profile),
390 )
391 .route(
392 "/api/config/:key",
393 get(crate::routes_config::get_config_key)
394 .put(crate::routes_config::put_config_key)
395 .delete(crate::routes_config::delete_config_key),
396 )
397 .route(
399 "/api/config/model/:name/placement",
400 get(get_model_placement)
401 .put(put_model_placement)
402 .delete(delete_model_placement),
403 )
404 .route("/health", get(health))
405 .with_state(state)
406 .route("/api/openapi.json", get(openapi_json))
407 .route("/api/docs", get(scalar_docs))
408}
409
410fn sse_message_to_event(msg: SseMessage) -> SseEvent {
413 fn serialize_event<T: Serialize>(event_name: &str, payload: &T) -> SseEvent {
414 match serde_json::to_string(payload) {
415 Ok(data) => SseEvent::default().event(event_name).data(data),
416 Err(err) => SseEvent::default().event("error").data(
417 serde_json::json!({
418 "message": format!("failed to serialize SSE payload: {err}")
419 })
420 .to_string(),
421 ),
422 }
423 }
424
425 match msg {
426 SseMessage::Progress(payload) => serialize_event("progress", &payload),
427 SseMessage::Complete(payload) => serialize_event("complete", &payload),
428 SseMessage::UpscaleComplete(payload) => serialize_event("complete", &payload),
429 SseMessage::Error(payload) => serialize_event("error", &payload),
430 }
431}
432
433#[cfg(test)]
434fn save_image_to_dir(
435 dir: &std::path::Path,
436 img: &mold_core::ImageData,
437 model: &str,
438 batch_size: u32,
439) {
440 if let Err(e) = std::fs::create_dir_all(dir) {
441 tracing::warn!("failed to create output dir {}: {e}", dir.display());
442 return;
443 }
444 let timestamp_ms = std::time::SystemTime::now()
447 .duration_since(std::time::UNIX_EPOCH)
448 .unwrap_or_default()
449 .as_millis() as u64;
450 let ext = img.format.to_string();
451 let filename =
452 mold_core::default_output_filename(model, timestamp_ms, &ext, batch_size, img.index);
453 let path = dir.join(&filename);
454 match std::fs::write(&path, &img.data) {
455 Ok(()) => tracing::info!("saved image to {}", path.display()),
456 Err(e) => tracing::warn!("failed to save image to {}: {e}", path.display()),
457 }
458}
459
460async fn prepare_generation(
468 state: &AppState,
469 request: &mut mold_core::GenerateRequest,
470) -> Result<(Option<std::path::PathBuf>, Option<String>, Option<usize>), ApiError> {
471 apply_default_metadata_setting(state, request).await;
476
477 let preferred_gpu = validate_multi_gpu_placement(state, request.placement.as_ref())?;
478
479 maybe_expand_prompt(state, request, preferred_gpu).await?;
481
482 if let Err(e) = model_manager::install_catalog_model(state, &request.model).await {
492 return Err(model_manager::install_error_to_api_error(&e));
493 }
494 let family_hint = model_manager::catalog_family_for(state, &request.model).await;
495
496 let resolved_family = model_manager::family_for_model(state, &request.model).await;
500 request.normalise_output_format(resolved_family.as_deref());
504
505 if let Err(e) = validate_generate_request(request, family_hint.as_deref()) {
506 return Err(ApiError::validation(e));
507 }
508
509 resolve_server_local_media_paths(state, request).await?;
510
511 let _ = model_manager::check_model_available(state, &request.model).await?;
512
513 let (output_dir, dim_warning) = {
514 let config = state.config.read().await;
515 let output_dir = if config.is_output_disabled() {
516 None
517 } else {
518 Some(config.effective_output_dir())
519 };
520 let family = config.resolved_model_config(&request.model).family;
521 let dim_warning = family
522 .as_deref()
523 .and_then(|f| mold_core::dimension_warning(request.width, request.height, f));
524 (output_dir, dim_warning)
525 };
526
527 Ok((output_dir, dim_warning, preferred_gpu))
528}
529
530fn record_prompt_history(state: &AppState, prompt: &str, negative: Option<&str>, model: &str) {
536 let Some(db) = state.metadata_db.as_ref().as_ref() else {
537 return;
538 };
539 let history = mold_db::PromptHistory::new(db);
540 if let Ok(rows) = history.recent(1) {
541 if rows.first().is_some_and(|latest| {
542 latest.prompt == prompt
543 && latest.model == model
544 && latest.negative.as_deref() == negative
545 }) {
546 return;
547 }
548 }
549 if let Err(e) = history.push(&mold_db::HistoryEntry {
550 prompt: prompt.to_string(),
551 negative: negative.map(str::to_string),
552 model: model.to_string(),
553 created_at_ms: 0, }) {
555 tracing::warn!("failed to record prompt history: {e:#}");
556 }
557}
558
559pub(crate) async fn resolve_server_local_media_paths(
560 state: &AppState,
561 request: &mut mold_core::GenerateRequest,
562) -> Result<(), ApiError> {
563 if request.audio_file_path.is_none() && request.source_video_path.is_none() {
564 return Ok(());
565 }
566
567 let roots = state.config.read().await.resolved_media_roots();
568 if let Some(path) = request.audio_file_path.as_deref() {
569 let resolved = mold_core::resolve_server_media_path(path, &roots)
570 .map_err(|e| ApiError::validation(format!("audio_file_path: {e}")))?;
571 request.audio_file_path = Some(resolved.to_string_lossy().to_string());
572 }
573 if let Some(path) = request.source_video_path.as_deref() {
574 let resolved = mold_core::resolve_server_media_path(path, &roots)
575 .map_err(|e| ApiError::validation(format!("source_video_path: {e}")))?;
576 request.source_video_path = Some(resolved.to_string_lossy().to_string());
577 }
578
579 Ok(())
580}
581
582fn active_gpu_selection(state: &AppState) -> GpuSelection {
583 let ordinals: Vec<usize> = state
584 .gpu_pool
585 .workers
586 .iter()
587 .map(|w| w.gpu.ordinal)
588 .collect();
589 if ordinals.is_empty() {
590 GpuSelection::All
591 } else {
592 GpuSelection::Specific(ordinals)
593 }
594}
595
596fn validate_multi_gpu_placement(
597 state: &AppState,
598 placement: Option<&mold_core::types::DevicePlacement>,
599) -> Result<Option<usize>, ApiError> {
600 state
601 .gpu_pool
602 .resolve_explicit_placement_gpu(placement)
603 .map_err(ApiError::validation)
604}
605
606fn select_aux_worker(
607 state: &AppState,
608) -> Result<std::sync::Arc<crate::gpu_pool::GpuWorker>, ApiError> {
609 let mut workers: Vec<_> = state
610 .gpu_pool
611 .workers
612 .iter()
613 .filter(|w| !w.is_degraded())
614 .cloned()
615 .collect();
616 workers.sort_by_key(|w| {
617 (
618 w.in_flight.load(Ordering::SeqCst),
619 Reverse(w.gpu.total_vram_bytes),
620 )
621 });
622 workers
623 .into_iter()
624 .next()
625 .ok_or_else(|| ApiError::internal("no GPU worker available for auxiliary workload"))
626}
627
628fn clear_global_upscaler_cache(state: &AppState) {
629 if let Ok(mut cache) = state.upscaler_cache.try_lock() {
630 if cache.is_some() {
631 *cache = None;
632 tracing::info!("upscaler cache cleared");
633 }
634 }
635}
636
637#[utoipa::path(
640 post,
641 path = "/api/generate",
642 tag = "generation",
643 request_body = mold_core::GenerateRequest,
644 responses(
645 (status = 200, description = "Generated image bytes", content_type = "image/png"),
646 (status = 404, description = "Model not downloaded"),
647 (status = 422, description = "Invalid request parameters"),
648 (status = 500, description = "Inference error"),
649 (status = 503, description = "Generation queue full"),
650 )
651)]
652async fn generate(
655 State(state): State<AppState>,
656 Json(mut req): Json<mold_core::GenerateRequest>,
657) -> Result<impl IntoResponse, ApiError> {
658 let typed = (
661 req.prompt.clone(),
662 req.negative_prompt.clone(),
663 req.model.clone(),
664 );
665 let (output_dir, dim_warning, preferred_gpu) = prepare_generation(&state, &mut req).await?;
666 record_prompt_history(&state, &typed.0, typed.1.as_deref(), &typed.2);
667
668 tracing::info!(
669 model = %req.model,
670 prompt = %req.prompt,
671 width = req.width,
672 height = req.height,
673 steps = req.steps,
674 guidance = req.guidance,
675 seed = ?req.seed,
676 format = %req.resolved_output_format(),
677 lora = ?req.lora.as_ref().map(|l| &l.path),
678 lora_scale = ?req.lora.as_ref().map(|l| l.scale),
679 loras = ?req.loras.as_ref().map(|v| v.iter().map(|l| &l.path).collect::<Vec<_>>()),
680 "generate request"
681 );
682
683 let (result_tx, result_rx) = tokio::sync::oneshot::channel();
685 let job_id = uuid::Uuid::new_v4().to_string();
690 let cancel = state
691 .job_registry
692 .register_with_target_gpu(&job_id, &req.model, preferred_gpu);
693 let job = GenerationJob {
694 id: job_id.clone(),
695 request: req,
696 progress_tx: None,
697 result_tx,
698 output_dir,
699 };
700
701 let _position = state
702 .queue
703 .submit(job, state.queue_capacity)
704 .await
705 .inspect_err(|_| state.job_registry.remove(&job_id))
706 .map_err(submit_error_to_api)?;
707
708 let result = tokio::select! {
713 result = result_rx => result
714 .map_err(|_| ApiError::internal("generation queue worker dropped the job"))?,
715 _ = cancel.notified() => {
716 return Err(ApiError::cancelled(format!(
717 "generation job {job_id} was cancelled while queued"
718 )));
719 }
720 };
721
722 match result {
723 Ok(job_result) => {
724 let img = job_result.image;
725 let response = job_result.response;
726 let content_type = HeaderValue::from_static(img.format.content_type());
727 let mut headers = HeaderMap::new();
728 headers.insert(header::CONTENT_TYPE, content_type);
729 headers.insert(
730 "x-mold-seed-used",
731 HeaderValue::from_str(&response.seed_used.to_string()).map_err(|e| {
732 ApiError::internal(format!("failed to serialize seed header: {e}"))
733 })?,
734 );
735 if let Some(ordinal) = response.gpu {
736 headers.insert(
737 "x-mold-gpu",
738 HeaderValue::from_str(&ordinal.to_string()).map_err(|e| {
739 ApiError::internal(format!("failed to serialize gpu header: {e}"))
740 })?,
741 );
742 }
743 if let Some(warning) = dim_warning {
744 match HeaderValue::from_str(&warning.replace('\n', " ")) {
745 Ok(val) => {
746 headers.insert("x-mold-dimension-warning", val);
747 }
748 Err(e) => {
749 tracing::warn!("dimension warning could not be encoded as header: {e}");
750 }
751 }
752 }
753 let output_data = if let Some(ref video) = response.video {
756 let ct = HeaderValue::from_static(video.format.content_type());
757 headers.insert(header::CONTENT_TYPE, ct);
758 if let Ok(v) = HeaderValue::from_str(&video.frames.to_string()) {
759 headers.insert("x-mold-video-frames", v);
760 }
761 if let Ok(v) = HeaderValue::from_str(&video.fps.to_string()) {
762 headers.insert("x-mold-video-fps", v);
763 }
764 if let Ok(v) = HeaderValue::from_str(&video.width.to_string()) {
765 headers.insert("x-mold-video-width", v);
766 }
767 if let Ok(v) = HeaderValue::from_str(&video.height.to_string()) {
768 headers.insert("x-mold-video-height", v);
769 }
770 if video.has_audio {
771 headers.insert("x-mold-video-has-audio", HeaderValue::from_static("1"));
772 }
773 if let Some(dur) = video.duration_ms {
774 if let Ok(v) = HeaderValue::from_str(&dur.to_string()) {
775 headers.insert("x-mold-video-duration-ms", v);
776 }
777 }
778 if let Some(sr) = video.audio_sample_rate {
779 if let Ok(v) = HeaderValue::from_str(&sr.to_string()) {
780 headers.insert("x-mold-video-audio-sample-rate", v);
781 }
782 }
783 if let Some(ch) = video.audio_channels {
784 if let Ok(v) = HeaderValue::from_str(&ch.to_string()) {
785 headers.insert("x-mold-video-audio-channels", v);
786 }
787 }
788 video.data.clone()
789 } else {
790 img.data
791 };
792 Ok((headers, output_data))
793 }
794 Err(err_msg) => {
795 if err_msg.contains("queue is full") {
799 Err(ApiError::queue_full(err_msg))
800 } else {
801 Err(ApiError::inference(err_msg))
802 }
803 }
804 }
805}
806
807fn validate_generate_request(
808 req: &mold_core::GenerateRequest,
809 family_hint: Option<&str>,
810) -> Result<(), String> {
811 mold_core::validate_generate_request_with_family(req, family_hint)
812}
813
814async fn apply_default_metadata_setting(state: &AppState, req: &mut mold_core::GenerateRequest) {
815 if req.embed_metadata.is_some() {
816 return;
817 }
818
819 let config = state.config.read().await;
820 req.embed_metadata = Some(config.effective_embed_metadata(None));
821}
822
823async fn maybe_expand_prompt(
825 state: &AppState,
826 req: &mut mold_core::GenerateRequest,
827 preferred_gpu: Option<usize>,
828) -> Result<(), ApiError> {
829 if req.expand != Some(true) {
830 return Ok(());
831 }
832
833 let config = state.config.read().await;
834 let config_snapshot = config.clone();
835 let expand_settings = config.expand.clone().with_env_overrides();
836
837 let model_family = config
839 .resolved_model_config(&req.model)
840 .family
841 .or_else(|| mold_core::manifest::find_manifest(&req.model).map(|m| m.family.clone()))
842 .unwrap_or_else(|| {
843 tracing::warn!(
844 model = %req.model,
845 "could not resolve model family for prompt expansion, defaulting to \"flux\""
846 );
847 "flux".to_string()
848 });
849
850 let expand_config = expand_settings.to_expand_config(&model_family, 1);
851 let original_prompt = req.prompt.clone();
852
853 drop(config);
855
856 let expander = create_server_expander(
857 &config_snapshot,
858 &expand_settings,
859 active_gpu_selection(state),
860 preferred_gpu,
861 )?;
862 let result =
863 tokio::task::spawn_blocking(move || expander.expand(&original_prompt, &expand_config))
864 .await
865 .map_err(|e| ApiError::internal(format!("expand task failed: {e}")))?
866 .map_err(|e| ApiError::internal(format!("prompt expansion failed: {e}")))?;
867
868 if let Some(expanded) = result.expanded.first() {
869 req.original_prompt = Some(req.prompt.clone());
870 req.prompt = expanded.clone();
871 }
872
873 Ok(())
874}
875
876fn create_server_expander(
878 _config: &mold_core::Config,
879 settings: &mold_core::ExpandSettings,
880 _gpu_selection: GpuSelection,
881 _preferred_gpu: Option<usize>,
882) -> Result<Box<dyn mold_core::PromptExpander>, ApiError> {
883 if let Some(api_expander) = settings.create_api_expander() {
884 return Ok(Box::new(api_expander));
885 }
886
887 #[cfg(feature = "expand")]
888 {
889 if let Some(local) =
890 mold_inference::expand::LocalExpander::from_config(_config, Some(&settings.model))
891 {
892 return Ok(Box::new(
893 local
894 .with_gpu_selection(_gpu_selection)
895 .with_preferred_gpu(_preferred_gpu),
896 ));
897 }
898 return Err(ApiError::validation(
899 "local expand model not found — run: mold pull qwen3-expand".to_string(),
900 ));
901 }
902
903 #[cfg(not(feature = "expand"))]
904 {
905 Err(ApiError::validation(
906 "local prompt expansion not available — built without expand feature. \
907 Configure an API backend in [expand] settings."
908 .to_string(),
909 ))
910 }
911}
912
913#[utoipa::path(
916 post,
917 path = "/api/expand",
918 tag = "generation",
919 request_body = mold_core::ExpandRequest,
920 responses(
921 (status = 200, description = "Expanded prompt(s)", body = mold_core::ExpandResponse),
922 (status = 422, description = "Invalid request parameters"),
923 (status = 500, description = "Expansion failed"),
924 )
925)]
926async fn expand_prompt(
927 State(state): State<AppState>,
928 Json(req): Json<mold_core::ExpandRequest>,
929) -> Result<Json<mold_core::ExpandResponse>, ApiError> {
930 if req.variations == 0 || req.variations > mold_core::expand::MAX_VARIATIONS {
931 return Err(ApiError::validation(format!(
932 "variations must be between 1 and {}",
933 mold_core::expand::MAX_VARIATIONS,
934 )));
935 }
936
937 let config = state.config.read().await;
938 let expand_settings = config.expand.clone().with_env_overrides();
939 let expand_config = expand_settings.to_expand_config(&req.model_family, req.variations);
940 let prompt = req.prompt.clone();
941 let config_snapshot = config.clone();
942 drop(config);
943
944 let expander = create_server_expander(
945 &config_snapshot,
946 &expand_settings,
947 active_gpu_selection(&state),
948 None,
949 )?;
950 let result = tokio::task::spawn_blocking(move || expander.expand(&prompt, &expand_config))
951 .await
952 .map_err(|e| ApiError::internal(format!("expand task failed: {e}")))?
953 .map_err(|e| ApiError::internal(format!("prompt expansion failed: {e}")))?;
954
955 Ok(Json(mold_core::ExpandResponse {
956 original: req.prompt,
957 expanded: result.expanded,
958 }))
959}
960
961async fn upscale(
964 State(state): State<AppState>,
965 Json(req): Json<mold_core::UpscaleRequest>,
966) -> Result<Json<mold_core::UpscaleResponse>, ApiError> {
967 if let Err(msg) = mold_core::validate_upscale_request(&req) {
968 return Err(ApiError::validation(msg));
969 }
970
971 let model_name = mold_core::manifest::resolve_model_name(&req.model);
972
973 let needs_pull = {
975 let config = state.config.read().await;
976 config
977 .models
978 .get(&model_name)
979 .and_then(|c| c.transformer.as_ref())
980 .is_none()
981 };
982 if needs_pull {
983 if mold_core::manifest::find_manifest(&model_name).is_none() {
984 return Err(ApiError::not_found(format!(
985 "unknown upscaler model '{}'. Run 'mold list' to see available models.",
986 model_name
987 )));
988 }
989 model_manager::pull_model(&state, &model_name, None).await?;
990 }
991
992 let config = state.config.read().await;
993 let weights_path = config
994 .models
995 .get(&model_name)
996 .and_then(|c| c.transformer.as_ref())
997 .ok_or_else(|| {
998 ApiError::not_found(format!(
999 "upscaler model '{}' not configured after pull",
1000 model_name
1001 ))
1002 })?;
1003 let weights_path = std::path::PathBuf::from(weights_path);
1004 let model_name_owned = model_name.clone();
1005 drop(config);
1006
1007 let resp = if state.gpu_pool.worker_count() > 0 {
1008 let worker = select_aux_worker(&state)?;
1009 worker.in_flight.fetch_add(1, Ordering::SeqCst);
1010 let worker_clone = worker.clone();
1011 let result =
1012 tokio::task::spawn_blocking(move || -> anyhow::Result<mold_core::UpscaleResponse> {
1013 struct ThreadGpuGuard;
1014 impl Drop for ThreadGpuGuard {
1015 fn drop(&mut self) {
1016 mold_inference::device::clear_thread_gpu_ordinal();
1017 }
1018 }
1019
1020 mold_inference::device::init_thread_gpu_ordinal(worker_clone.gpu.ordinal);
1021 let _thread_gpu = ThreadGpuGuard;
1022 let _load_lock = worker_clone.model_load_lock.lock().unwrap();
1023 let mut engine = mold_inference::create_upscale_engine(
1024 model_name_owned,
1025 weights_path,
1026 mold_inference::LoadStrategy::Eager,
1027 worker_clone.gpu.ordinal,
1028 )?;
1029 engine.upscale(&req)
1030 })
1031 .await
1032 .map_err(|e| ApiError::internal(format!("upscale task panicked: {e}")));
1033 worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1034 result?.map_err(|e| ApiError::internal(format!("upscale failed: {e}")))?
1035 } else {
1036 let upscaler_cache = state.upscaler_cache.clone();
1037 tokio::task::spawn_blocking(move || -> anyhow::Result<mold_core::UpscaleResponse> {
1038 let mut cache = upscaler_cache.lock().unwrap_or_else(|e| e.into_inner());
1039
1040 let needs_new = cache
1042 .as_ref()
1043 .is_none_or(|e| e.model_name() != model_name_owned);
1044 if needs_new {
1045 let new_engine = mold_inference::create_upscale_engine(
1046 model_name_owned,
1047 weights_path,
1048 mold_inference::LoadStrategy::Eager,
1049 0,
1050 )?;
1051 *cache = Some(new_engine);
1052 }
1053
1054 cache.as_mut().unwrap().upscale(&req)
1055 })
1056 .await
1057 .map_err(|e| ApiError::internal(format!("upscale task panicked: {e}")))?
1058 .map_err(|e| ApiError::internal(format!("upscale failed: {e}")))?
1059 };
1060
1061 Ok(Json(resp))
1062}
1063
1064async fn upscale_stream(
1067 State(state): State<AppState>,
1068 Json(req): Json<mold_core::UpscaleRequest>,
1069) -> Result<Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
1070 if let Err(msg) = mold_core::validate_upscale_request(&req) {
1071 return Err(ApiError::validation(msg));
1072 }
1073
1074 let model_name = mold_core::manifest::resolve_model_name(&req.model);
1075
1076 let needs_pull = {
1078 let config = state.config.read().await;
1079 config
1080 .models
1081 .get(&model_name)
1082 .and_then(|c| c.transformer.as_ref())
1083 .is_none()
1084 };
1085
1086 if needs_pull && mold_core::manifest::find_manifest(&model_name).is_none() {
1088 return Err(ApiError::not_found(format!(
1089 "unknown upscaler model '{}'. Run 'mold list' to see available models.",
1090 model_name
1091 )));
1092 }
1093
1094 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseMessage>();
1095 let model_name_owned = model_name.clone();
1096 let state_clone = state.clone();
1097 let upscaler_cache = state.upscaler_cache.clone();
1098
1099 tokio::spawn(async move {
1100 if needs_pull {
1102 let progress_tx = tx.clone();
1103 let callback =
1104 std::sync::Arc::new(move |event: mold_core::download::DownloadProgressEvent| {
1105 let sse_event = match event {
1106 mold_core::download::DownloadProgressEvent::Status { message } => {
1107 SseProgressEvent::Info { message }
1108 }
1109 mold_core::download::DownloadProgressEvent::FileStart {
1110 filename,
1111 file_index,
1112 total_files,
1113 size_bytes,
1114 batch_bytes_downloaded,
1115 batch_bytes_total,
1116 batch_elapsed_ms,
1117 } => SseProgressEvent::DownloadProgress {
1118 filename,
1119 file_index,
1120 total_files,
1121 bytes_downloaded: 0,
1122 bytes_total: size_bytes,
1123 batch_bytes_downloaded,
1124 batch_bytes_total,
1125 batch_elapsed_ms,
1126 },
1127 mold_core::download::DownloadProgressEvent::FileProgress {
1128 filename,
1129 file_index,
1130 bytes_downloaded,
1131 bytes_total,
1132 batch_bytes_downloaded,
1133 batch_bytes_total,
1134 batch_elapsed_ms,
1135 } => SseProgressEvent::DownloadProgress {
1136 filename,
1137 file_index,
1138 total_files: 0,
1139 bytes_downloaded,
1140 bytes_total,
1141 batch_bytes_downloaded,
1142 batch_bytes_total,
1143 batch_elapsed_ms,
1144 },
1145 mold_core::download::DownloadProgressEvent::FileDone {
1146 filename,
1147 file_index,
1148 total_files,
1149 batch_bytes_downloaded,
1150 batch_bytes_total,
1151 batch_elapsed_ms,
1152 } => SseProgressEvent::DownloadDone {
1153 filename,
1154 file_index,
1155 total_files,
1156 batch_bytes_downloaded,
1157 batch_bytes_total,
1158 batch_elapsed_ms,
1159 },
1160 };
1161 let _ = progress_tx.send(SseMessage::Progress(sse_event));
1162 });
1163
1164 match model_manager::pull_model(&state_clone, &model_name_owned, Some(callback)).await {
1165 Ok(_) => {
1166 let _ = tx.send(SseMessage::Progress(SseProgressEvent::PullComplete {
1167 model: model_name_owned.clone(),
1168 }));
1169 }
1170 Err(e) => {
1171 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1172 message: format!("failed to pull upscaler model: {}", e.error),
1173 }));
1174 return;
1175 }
1176 }
1177 }
1178
1179 let weights_path = {
1181 let config = state_clone.config.read().await;
1182 config
1183 .models
1184 .get(&model_name_owned)
1185 .and_then(|c| c.transformer.as_ref())
1186 .map(std::path::PathBuf::from)
1187 };
1188
1189 let Some(weights_path) = weights_path else {
1190 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1191 message: format!(
1192 "upscaler model '{}' not configured after pull",
1193 model_name_owned
1194 ),
1195 }));
1196 return;
1197 };
1198
1199 let result = if state_clone.gpu_pool.worker_count() > 0 {
1200 match select_aux_worker(&state_clone) {
1201 Ok(worker) => {
1202 worker.in_flight.fetch_add(1, Ordering::SeqCst);
1203 let worker_clone = worker.clone();
1204 let tx_for_worker = tx.clone();
1205 let model_name_for_worker = model_name_owned.clone();
1206 let weights_path_for_worker = weights_path.clone();
1207 let req_for_worker = req.clone();
1208 let result = tokio::task::spawn_blocking(move || {
1209 struct ThreadGpuGuard;
1210 impl Drop for ThreadGpuGuard {
1211 fn drop(&mut self) {
1212 mold_inference::device::clear_thread_gpu_ordinal();
1213 }
1214 }
1215
1216 mold_inference::device::init_thread_gpu_ordinal(worker_clone.gpu.ordinal);
1217 let _thread_gpu = ThreadGpuGuard;
1218 let _load_lock = worker_clone.model_load_lock.lock().unwrap();
1219 let _ = tx_for_worker.send(SseMessage::Progress(
1220 mold_core::SseProgressEvent::StageStart {
1221 name: format!(
1222 "Loading upscaler model on GPU {}",
1223 worker_clone.gpu.ordinal
1224 ),
1225 },
1226 ));
1227 let mut engine = match mold_inference::create_upscale_engine(
1228 model_name_for_worker,
1229 weights_path_for_worker,
1230 mold_inference::LoadStrategy::Eager,
1231 worker_clone.gpu.ordinal,
1232 ) {
1233 Ok(engine) => engine,
1234 Err(e) => {
1235 let _ = tx_for_worker.send(SseMessage::Error(
1236 mold_core::SseErrorEvent {
1237 message: format!("failed to load upscaler: {e}"),
1238 },
1239 ));
1240 return;
1241 }
1242 };
1243
1244 let tx_progress = tx_for_worker.clone();
1245 engine.set_on_progress(Box::new(move |event| {
1246 let sse_event: mold_core::SseProgressEvent = event.into();
1247 let _ = tx_progress.send(SseMessage::Progress(sse_event));
1248 }));
1249
1250 match engine.upscale(&req_for_worker) {
1251 Ok(resp) => {
1252 let image_b64 = base64::engine::general_purpose::STANDARD
1253 .encode(&resp.image.data);
1254 let _ = tx_for_worker.send(SseMessage::UpscaleComplete(
1255 mold_core::SseUpscaleCompleteEvent {
1256 image: image_b64,
1257 format: resp.image.format,
1258 model: resp.model,
1259 scale_factor: resp.scale_factor,
1260 original_width: resp.original_width,
1261 original_height: resp.original_height,
1262 upscale_time_ms: resp.upscale_time_ms,
1263 },
1264 ));
1265 }
1266 Err(e) => {
1267 let _ = tx_for_worker.send(SseMessage::Error(
1268 mold_core::SseErrorEvent {
1269 message: format!("upscale failed: {e}"),
1270 },
1271 ));
1272 }
1273 }
1274
1275 engine.clear_on_progress();
1276 })
1277 .await;
1278 worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1279 result
1280 }
1281 Err(e) => {
1282 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1283 message: e.error,
1284 }));
1285 return;
1286 }
1287 }
1288 } else {
1289 let model_name_for_cache = model_name_owned.clone();
1290 let weights_path_for_cache = weights_path.clone();
1291 let req_for_cache = req.clone();
1292 tokio::task::spawn_blocking(move || {
1293 let mut cache = upscaler_cache.lock().unwrap();
1294
1295 let needs_new = cache
1296 .as_ref()
1297 .is_none_or(|e| e.model_name() != model_name_for_cache);
1298 if needs_new {
1299 let _ = tx.send(SseMessage::Progress(
1300 mold_core::SseProgressEvent::StageStart {
1301 name: "Loading upscaler model".to_string(),
1302 },
1303 ));
1304 match mold_inference::create_upscale_engine(
1305 model_name_for_cache,
1306 weights_path_for_cache,
1307 mold_inference::LoadStrategy::Eager,
1308 0,
1309 ) {
1310 Ok(new_engine) => {
1311 *cache = Some(new_engine);
1312 }
1313 Err(e) => {
1314 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1315 message: format!("failed to load upscaler: {e}"),
1316 }));
1317 return;
1318 }
1319 }
1320 }
1321
1322 let engine = cache.as_mut().unwrap();
1323 let tx_progress = tx.clone();
1324 engine.set_on_progress(Box::new(move |event| {
1325 let sse_event: mold_core::SseProgressEvent = event.into();
1326 let _ = tx_progress.send(SseMessage::Progress(sse_event));
1327 }));
1328
1329 match engine.upscale(&req_for_cache) {
1330 Ok(resp) => {
1331 let image_b64 =
1332 base64::engine::general_purpose::STANDARD.encode(&resp.image.data);
1333 let _ = tx.send(SseMessage::UpscaleComplete(
1334 mold_core::SseUpscaleCompleteEvent {
1335 image: image_b64,
1336 format: resp.image.format,
1337 model: resp.model,
1338 scale_factor: resp.scale_factor,
1339 original_width: resp.original_width,
1340 original_height: resp.original_height,
1341 upscale_time_ms: resp.upscale_time_ms,
1342 },
1343 ));
1344 }
1345 Err(e) => {
1346 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1347 message: format!("upscale failed: {e}"),
1348 }));
1349 }
1350 }
1351
1352 engine.clear_on_progress();
1353 })
1354 .await
1355 };
1356
1357 if let Err(e) = result {
1358 tracing::error!("upscale task panicked: {e}");
1359 }
1360 });
1361
1362 let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx)
1363 .map(|msg| Ok::<_, Infallible>(sse_message_to_event(msg)));
1364
1365 Ok(Sse::new(stream).keep_alive(
1366 KeepAlive::new()
1367 .interval(std::time::Duration::from_secs(15))
1368 .text("ping"),
1369 ))
1370}
1371
1372#[utoipa::path(
1375 post,
1376 path = "/api/generate/stream",
1377 tag = "generation",
1378 request_body = mold_core::GenerateRequest,
1379 responses(
1380 (status = 200, description = "SSE event stream with progress and result"),
1381 (status = 404, description = "Model not downloaded"),
1382 (status = 422, description = "Invalid request parameters"),
1383 (status = 500, description = "Inference error"),
1384 (status = 503, description = "Generation queue full"),
1385 )
1386)]
1387async fn generate_stream(
1388 State(state): State<AppState>,
1389 Json(mut req): Json<mold_core::GenerateRequest>,
1390) -> Result<Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
1391 let typed = (
1394 req.prompt.clone(),
1395 req.negative_prompt.clone(),
1396 req.model.clone(),
1397 );
1398 let (output_dir, dim_warning, preferred_gpu) = prepare_generation(&state, &mut req).await?;
1399 record_prompt_history(&state, &typed.0, typed.1.as_deref(), &typed.2);
1400
1401 tracing::info!(
1402 model = %req.model,
1403 prompt = %req.prompt,
1404 "generate/stream request"
1405 );
1406
1407 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseMessage>();
1409
1410 if let Some(warning) = dim_warning {
1412 let _ = tx.send(SseMessage::Progress(SseProgressEvent::Info {
1413 message: warning,
1414 }));
1415 }
1416
1417 let (result_tx, result_rx) = tokio::sync::oneshot::channel();
1418 let job_id = uuid::Uuid::new_v4().to_string();
1421 let cancel = state
1422 .job_registry
1423 .register_with_target_gpu(&job_id, &req.model, preferred_gpu);
1424 let job = GenerationJob {
1425 id: job_id.clone(),
1426 request: req,
1427 progress_tx: Some(tx.clone()),
1428 result_tx,
1429 output_dir,
1430 };
1431
1432 let position = state
1433 .queue
1434 .submit(job, state.queue_capacity)
1435 .await
1436 .inspect_err(|_| state.job_registry.remove(&job_id))
1437 .map_err(submit_error_to_api)?;
1438
1439 let _ = tx.send(SseMessage::Progress(SseProgressEvent::Queued {
1442 position,
1443 id: job_id,
1444 }));
1445
1446 tokio::spawn(async move {
1453 tokio::select! {
1454 _ = result_rx => {}
1455 _ = cancel.notified() => {
1456 let _ = tx.send(SseMessage::Error(SseErrorEvent {
1457 message: "cancelled".to_string(),
1458 }));
1459 }
1460 }
1461 drop(tx); });
1463
1464 let stream = async_stream::stream! {
1470 let mut rx = rx;
1471 while let Some(msg) = rx.recv().await {
1472 let is_error = matches!(msg, SseMessage::Error(_));
1473 yield Ok::<_, Infallible>(sse_message_to_event(msg));
1474 if is_error {
1475 break;
1476 }
1477 }
1478 };
1479
1480 Ok(Sse::new(stream).keep_alive(
1481 KeepAlive::new()
1482 .interval(std::time::Duration::from_secs(15))
1483 .text("ping"),
1484 ))
1485}
1486
1487#[utoipa::path(
1490 get,
1491 path = "/api/models",
1492 tag = "models",
1493 responses(
1494 (status = 200, description = "List of available models", body = Vec<ModelInfoExtended>),
1495 )
1496)]
1497async fn list_models(State(state): State<AppState>) -> Json<Vec<ModelInfoExtended>> {
1498 Json(model_manager::list_models(&state).await)
1499}
1500
1501async fn generate_estimate(
1502 State(state): State<AppState>,
1503 Json(req): Json<GenerateRequest>,
1504) -> Result<Json<mold_core::GenerationMemoryEstimate>, ApiError> {
1505 Ok(Json(
1506 model_manager::estimate_generation_memory(&state, &req).await?,
1507 ))
1508}
1509
1510async fn model_components(
1511 State(state): State<AppState>,
1512 Path(model): Path<String>,
1513) -> Result<Json<mold_core::ModelComponentsResponse>, ApiError> {
1514 Ok(Json(
1515 model_manager::model_component_status(&state, &model).await?,
1516 ))
1517}
1518
1519#[derive(Debug, Deserialize, utoipa::ToSchema)]
1522pub struct LoadModelBody {
1523 #[schema(example = "flux-schnell:q8")]
1524 pub model: String,
1525 #[serde(default, skip_serializing_if = "Option::is_none")]
1528 pub gpu: Option<usize>,
1529}
1530
1531#[utoipa::path(
1532 post,
1533 path = "/api/models/load",
1534 tag = "models",
1535 request_body = LoadModelBody,
1536 responses(
1537 (status = 200, description = "Model loaded successfully"),
1538 (status = 404, description = "Model not downloaded"),
1539 (status = 400, description = "Unknown model"),
1540 (status = 500, description = "Failed to load model"),
1541 )
1542)]
1543async fn load_model(
1544 State(state): State<AppState>,
1545 Json(body): Json<LoadModelBody>,
1546) -> Result<impl IntoResponse, ApiError> {
1547 if let Err(e) = model_manager::install_catalog_model(&state, &body.model).await {
1548 return Err(model_manager::install_error_to_api_error(&e));
1549 }
1550 let _ = model_manager::check_model_available(&state, &body.model).await?;
1551
1552 if state.gpu_pool.worker_count() > 0 {
1554 let worker = match body.gpu {
1555 Some(ordinal) => state
1556 .gpu_pool
1557 .workers
1558 .iter()
1559 .find(|w| w.gpu.ordinal == ordinal)
1560 .cloned()
1561 .ok_or_else(|| {
1562 ApiError::not_found(format!("no GPU worker with ordinal {ordinal}"))
1563 })?,
1564 None => {
1565 let est = crate::queue::estimate_model_vram(&body.model);
1566 state
1567 .gpu_pool
1568 .select_worker(&body.model, est)
1569 .ok_or_else(|| {
1570 ApiError::internal(format!(
1571 "no GPU available to load model '{}'",
1572 body.model
1573 ))
1574 })?
1575 }
1576 };
1577 let config_snapshot = state.config.read().await.clone();
1578 let model_name = body.model.clone();
1579 let worker_clone = worker.clone();
1580 tokio::task::spawn_blocking(move || {
1581 crate::gpu_worker::load_blocking(&worker_clone, &model_name, &config_snapshot)
1582 })
1583 .await
1584 .map_err(|e| ApiError::internal(format!("model load task failed: {e}")))?
1585 .map_err(|e| ApiError::internal(format!("model load error: {e}")))?;
1586
1587 tracing::info!(
1588 model = %body.model,
1589 gpu = worker.gpu.ordinal,
1590 "model loaded via API"
1591 );
1592 return Ok(StatusCode::OK);
1593 }
1594
1595 model_manager::ensure_model_ready(&state, &body.model, None, None, false).await?;
1598 tracing::info!(model = %body.model, gpu = ?body.gpu, "model loaded via API (legacy)");
1599 Ok(StatusCode::OK)
1600}
1601
1602#[utoipa::path(
1605 post,
1606 path = "/api/models/pull",
1607 tag = "models",
1608 request_body = LoadModelBody,
1609 responses(
1610 (status = 200, description = "Model pulled (SSE stream or plain text)"),
1611 (status = 400, description = "Unknown model"),
1612 (status = 500, description = "Download failed"),
1613 )
1614)]
1615async fn pull_model_endpoint(
1616 State(state): State<AppState>,
1617 headers: HeaderMap,
1618 Json(body): Json<LoadModelBody>,
1619) -> Result<impl IntoResponse, ApiError> {
1620 let wants_sse = headers
1621 .get(header::ACCEPT)
1622 .and_then(|v| v.to_str().ok())
1623 .is_some_and(|v| v.contains("text/event-stream"));
1624
1625 if body.model.starts_with("cv:") || body.model.starts_with("hf:") {
1626 if let Err(e) = model_manager::install_catalog_model(&state, &body.model).await {
1627 return Err(model_manager::install_error_to_api_error(&e));
1628 }
1629 if model_manager::check_model_available(&state, &body.model)
1630 .await
1631 .is_ok()
1632 {
1633 return Ok(PullResponse::Text(format!(
1634 "model '{}' is already present",
1635 body.model
1636 )));
1637 }
1638 let companion_names = {
1639 let intents = state.catalog_intents.read().await;
1640 intents
1641 .get(&body.model)
1642 .map(|intent| {
1643 intent
1644 .companions
1645 .iter()
1646 .map(|companion| companion.name.clone())
1647 .collect::<Vec<_>>()
1648 })
1649 .unwrap_or_default()
1650 };
1651 let models_dir = state.config.read().await.resolved_models_dir();
1652 let companion_jobs = crate::catalog_api::enqueue_missing_companions(
1653 &companion_names,
1654 &models_dir,
1655 &state.downloads,
1656 Some(&body.model),
1657 )
1658 .await;
1659 let primary_job = crate::catalog_api::enqueue_catalog_primary_repair(&state, &body.model)
1660 .await
1661 .map_err(|(status, msg)| ApiError::internal_with_status(msg, status))?;
1662 if !companion_jobs.is_empty() || primary_job.is_some() {
1663 let primary_count = usize::from(primary_job.is_some());
1664 return Ok(PullResponse::Text(format!(
1665 "queued repair for model '{}' ({} primary job(s), {} companion job(s))",
1666 body.model,
1667 primary_count,
1668 companion_jobs.len()
1669 )));
1670 }
1671 model_manager::check_model_available(&state, &body.model).await?;
1672 return Ok(PullResponse::Text(format!(
1673 "model '{}' is already present",
1674 body.model
1675 )));
1676 }
1677
1678 let (job_id, _position) = match state.downloads.enqueue(body.model.clone()).await {
1680 Ok((id, pos, _)) => (id, pos),
1681 Err(crate::downloads::EnqueueError::UnknownModel(_)) => {
1682 return Err(ApiError::unknown_model(format!(
1683 "unknown model '{}'. Run 'mold list' to see available models.",
1684 body.model
1685 )));
1686 }
1687 Err(crate::downloads::EnqueueError::LockPoisoned) => {
1688 return Err(ApiError::internal("download queue state is corrupt"));
1689 }
1690 };
1691
1692 if !wants_sse {
1693 let mut rx = state.downloads.subscribe();
1695 loop {
1696 match rx.recv().await {
1697 Ok(mold_core::types::DownloadEvent::JobDone { id, model }) if id == job_id => {
1698 return Ok(PullResponse::Text(format!(
1699 "model '{model}' pulled successfully"
1700 )));
1701 }
1702 Ok(mold_core::types::DownloadEvent::JobFailed { id, error }) if id == job_id => {
1703 return Err(ApiError::internal(format!(
1704 "failed to pull model '{}': {error}",
1705 body.model
1706 )));
1707 }
1708 Ok(mold_core::types::DownloadEvent::JobCancelled { id }) if id == job_id => {
1709 return Err(ApiError::internal(format!(
1710 "pull of '{}' was cancelled",
1711 body.model
1712 )));
1713 }
1714 Ok(_) => continue,
1715 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
1716 Err(tokio::sync::broadcast::error::RecvError::Closed) => {
1717 return Err(ApiError::internal("download queue channel closed"));
1718 }
1719 }
1720 }
1721 }
1722
1723 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseMessage>();
1726 let mut events = state.downloads.subscribe();
1727 let model_for_cb = body.model.clone();
1728 tokio::spawn(async move {
1729 loop {
1730 match events.recv().await {
1731 Ok(mold_core::types::DownloadEvent::Started {
1732 id,
1733 files_total,
1734 bytes_total,
1735 }) if id == job_id => {
1736 let _ = tx.send(SseMessage::Progress(SseProgressEvent::DownloadProgress {
1737 filename: String::new(),
1738 file_index: 0,
1739 total_files: files_total,
1740 bytes_downloaded: 0,
1741 bytes_total,
1742 batch_bytes_downloaded: 0,
1743 batch_bytes_total: bytes_total,
1744 batch_elapsed_ms: 0,
1745 }));
1746 }
1747 Ok(mold_core::types::DownloadEvent::Progress {
1748 id,
1749 files_done,
1750 bytes_done,
1751 current_file,
1752 }) if id == job_id => {
1753 let _ = tx.send(SseMessage::Progress(SseProgressEvent::DownloadProgress {
1754 filename: current_file.unwrap_or_default(),
1755 file_index: files_done,
1756 total_files: 0,
1757 bytes_downloaded: bytes_done,
1758 bytes_total: 0,
1759 batch_bytes_downloaded: bytes_done,
1760 batch_bytes_total: 0,
1761 batch_elapsed_ms: 0,
1762 }));
1763 }
1764 Ok(mold_core::types::DownloadEvent::JobDone { id, .. }) if id == job_id => {
1765 let _ = tx.send(SseMessage::Progress(SseProgressEvent::PullComplete {
1766 model: model_for_cb.clone(),
1767 }));
1768 break;
1769 }
1770 Ok(mold_core::types::DownloadEvent::JobFailed { id, error }) if id == job_id => {
1771 let _ = tx.send(SseMessage::Error(SseErrorEvent { message: error }));
1772 break;
1773 }
1774 Ok(mold_core::types::DownloadEvent::JobCancelled { id }) if id == job_id => {
1775 let _ = tx.send(SseMessage::Error(SseErrorEvent {
1776 message: "pull cancelled".into(),
1777 }));
1778 break;
1779 }
1780 Ok(_) => continue,
1781 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
1782 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
1783 }
1784 }
1785 });
1786
1787 let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx)
1788 .map(|msg| Ok::<_, Infallible>(sse_message_to_event(msg)));
1789
1790 Ok(PullResponse::Sse(
1791 Sse::new(stream)
1792 .keep_alive(
1793 KeepAlive::new()
1794 .interval(std::time::Duration::from_secs(15))
1795 .text("ping"),
1796 )
1797 .into_response(),
1798 ))
1799}
1800
1801enum PullResponse {
1803 Sse(axum::response::Response),
1804 Text(String),
1805}
1806
1807impl IntoResponse for PullResponse {
1808 fn into_response(self) -> axum::response::Response {
1809 match self {
1810 PullResponse::Sse(resp) => resp,
1811 PullResponse::Text(text) => text.into_response(),
1812 }
1813 }
1814}
1815
1816#[derive(Debug, Default, Deserialize, utoipa::ToSchema)]
1821pub struct UnloadRequest {
1822 #[serde(default, skip_serializing_if = "Option::is_none")]
1824 pub model: Option<String>,
1825 #[serde(default, skip_serializing_if = "Option::is_none")]
1827 pub gpu: Option<usize>,
1828}
1829
1830#[utoipa::path(
1831 delete,
1832 path = "/api/models/unload",
1833 tag = "models",
1834 request_body(content = Option<UnloadRequest>, content_type = "application/json"),
1835 responses(
1836 (status = 200, description = "Model unloaded or no model was loaded", body = String),
1837 )
1838)]
1839async fn unload_model(
1840 State(state): State<AppState>,
1841 body: Option<Json<UnloadRequest>>,
1842) -> Result<impl IntoResponse, ApiError> {
1843 let req = body.map(|b| b.0).unwrap_or_default();
1844 tracing::debug!(model = ?req.model, gpu = ?req.gpu, "unload request");
1845 clear_global_upscaler_cache(&state);
1846
1847 if state.gpu_pool.worker_count() > 0 {
1849 let targets: Vec<_> = match (req.gpu, req.model.as_deref()) {
1851 (Some(ordinal), _) => state
1852 .gpu_pool
1853 .workers
1854 .iter()
1855 .filter(|w| w.gpu.ordinal == ordinal)
1856 .cloned()
1857 .collect(),
1858 (None, Some(model)) => state
1859 .gpu_pool
1860 .workers
1861 .iter()
1862 .filter(|w| {
1863 let cache = w.model_cache.lock().unwrap();
1864 cache
1865 .get(model)
1866 .map(|e| e.residency == crate::model_cache::ModelResidency::Gpu)
1867 .unwrap_or(false)
1868 })
1869 .cloned()
1870 .collect(),
1871 (None, None) => state.gpu_pool.workers.clone(),
1872 };
1873
1874 if targets.is_empty() {
1875 return Ok((StatusCode::OK, "no model loaded".to_string()));
1876 }
1877
1878 let mut unloaded_pairs: Vec<(usize, String)> = Vec::new();
1879 for worker in targets {
1880 let worker_clone = worker.clone();
1881 let result = tokio::task::spawn_blocking(move || {
1882 crate::gpu_worker::unload_blocking(&worker_clone)
1883 })
1884 .await
1885 .map_err(|e| ApiError::internal(format!("unload task failed: {e}")))?;
1886 if let Some(name) = result {
1887 unloaded_pairs.push((worker.gpu.ordinal, name));
1888 }
1889 }
1890
1891 let msg = if unloaded_pairs.is_empty() {
1892 "no model loaded".to_string()
1893 } else {
1894 let joined: Vec<String> = unloaded_pairs
1895 .iter()
1896 .map(|(o, m)| format!("gpu{o}:{m}"))
1897 .collect();
1898 format!("unloaded {}", joined.join(", "))
1899 };
1900 return Ok((StatusCode::OK, msg));
1901 }
1902
1903 Ok((StatusCode::OK, model_manager::unload_model(&state).await))
1905}
1906
1907async fn model_is_gpu_resident(state: &AppState, canonical: &str) -> bool {
1913 {
1914 let cache = state.model_cache.lock().await;
1915 if cache.active_model() == Some(canonical) {
1916 return true;
1917 }
1918 }
1919 if state
1920 .active_generation
1921 .read()
1922 .unwrap_or_else(|e| e.into_inner())
1923 .as_ref()
1924 .is_some_and(|g| g.model == canonical)
1925 {
1926 return true;
1927 }
1928 for worker in &state.gpu_pool.workers {
1929 if let Ok(active) = worker.active_generation.read() {
1930 if active.as_ref().is_some_and(|g| g.model == canonical) {
1931 return true;
1932 }
1933 }
1934 if let Ok(cache) = worker.model_cache.lock() {
1935 if cache.active_model() == Some(canonical) {
1936 return true;
1937 }
1938 }
1939 }
1940 false
1941}
1942
1943#[utoipa::path(
1952 delete,
1953 path = "/api/models/{model}",
1954 tag = "models",
1955 params(("model" = String, Path, description = "Model name (e.g. flux-schnell:q8)")),
1956 responses(
1957 (status = 200, description = "Model removed", body = mold_core::ModelRemovalResponse),
1958 (status = 404, description = "Model not installed"),
1959 (status = 409, description = "Model is currently loaded — unload it first"),
1960 )
1961)]
1962async fn delete_model(
1963 State(state): State<AppState>,
1964 Path(model): Path<String>,
1965) -> Result<Json<mold_core::ModelRemovalResponse>, ApiError> {
1966 let canonical = mold_core::manifest::resolve_model_name(&model);
1967
1968 if model_is_gpu_resident(&state, &canonical).await
1974 || (model != canonical && model_is_gpu_resident(&state, &model).await)
1975 {
1976 return Err(ApiError::with_code(
1977 format!(
1978 "model '{canonical}' is currently loaded; unload it first (DELETE /api/models/unload)"
1979 ),
1980 "MODEL_LOADED",
1981 StatusCode::CONFLICT,
1982 ));
1983 }
1984
1985 let mut config = state.config.write().await;
1988 let in_config = config.models.contains_key(&canonical);
1989 if !in_config && !config.manifest_model_is_downloaded(&canonical) {
1990 return Err(ApiError::with_code(
1991 format!("model '{canonical}' is not installed"),
1992 "UNKNOWN_MODEL",
1993 StatusCode::NOT_FOUND,
1994 ));
1995 }
1996
1997 tracing::info!(model = %canonical, "model removal requested");
1998 let plan = mold_core::removal::plan_removal(&config, &canonical);
1999 let outcome = mold_core::removal::execute_removal(&config, &plan);
2000 for warning in &outcome.warnings {
2001 tracing::warn!(model = %canonical, "model removal: {warning}");
2002 }
2003
2004 mold_core::download::remove_pulling_marker(&canonical);
2005 if in_config {
2006 config.remove_model(&canonical);
2007 if let Err(e) = config.save() {
2008 tracing::warn!("failed to persist model removal to config.toml: {e}");
2009 }
2010 }
2011 drop(config);
2012
2013 {
2016 let mut cache = state.model_cache.lock().await;
2017 let _ = cache.remove(&canonical);
2018 }
2019 for worker in &state.gpu_pool.workers {
2020 if let Ok(mut cache) = worker.model_cache.lock() {
2021 let _ = cache.remove(&canonical);
2022 }
2023 }
2024
2025 let kept = plan
2026 .shared_files
2027 .iter()
2028 .map(|(path, used_by)| mold_core::KeptComponent {
2029 component: path.clone(),
2030 used_by: used_by.clone(),
2031 })
2032 .collect();
2033
2034 Ok(Json(mold_core::ModelRemovalResponse {
2035 removed: outcome.removed,
2036 kept,
2037 freed_bytes: outcome.freed_bytes,
2038 }))
2039}
2040
2041#[utoipa::path(
2044 get,
2045 path = "/api/status",
2046 tag = "server",
2047 responses(
2048 (status = 200, description = "Server status", body = ServerStatus),
2049 )
2050)]
2051async fn server_status(State(state): State<AppState>) -> Json<ServerStatus> {
2052 let gpu_statuses = state.gpu_pool.gpu_status();
2054 let has_gpus = !gpu_statuses.is_empty();
2055
2056 let gpu_models_loaded: Vec<String> = gpu_statuses
2058 .iter()
2059 .filter_map(|g| g.loaded_model.clone())
2060 .collect();
2061 let gpu_busy = gpu_statuses
2062 .iter()
2063 .any(|g| g.state == GpuWorkerState::Generating);
2064
2065 let multi_gpu_current_gen = if has_gpus {
2068 state.gpu_pool.workers.iter().find_map(|w| {
2069 let gen = w.active_generation.read().ok()?;
2070 gen.as_ref().map(|g| ActiveGenerationStatus {
2071 model: g.model.clone(),
2072 prompt_sha256: g.prompt_sha256.clone(),
2073 started_at_unix_ms: g.started_at_unix_ms,
2074 elapsed_ms: g.started_at.elapsed().as_millis() as u64,
2075 })
2076 })
2077 } else {
2078 None
2079 };
2080
2081 let (models_loaded, busy, current_generation) = if has_gpus {
2083 (gpu_models_loaded, gpu_busy, multi_gpu_current_gen)
2084 } else {
2085 let snapshot = state.model_cache.lock().await.snapshot();
2086 let models = match (snapshot.model_name, snapshot.is_loaded) {
2087 (Some(model_name), true) => vec![model_name],
2088 _ => vec![],
2089 };
2090 let gen = state
2091 .active_generation
2092 .read()
2093 .unwrap_or_else(|e| e.into_inner())
2094 .as_ref()
2095 .map(|active| ActiveGenerationStatus {
2096 model: active.model.clone(),
2097 prompt_sha256: active.prompt_sha256.clone(),
2098 started_at_unix_ms: active.started_at_unix_ms,
2099 elapsed_ms: active.started_at.elapsed().as_millis() as u64,
2100 });
2101 let is_busy = gen.is_some();
2102 (models, is_busy, gen)
2103 };
2104
2105 Json(ServerStatus {
2106 version: env!("CARGO_PKG_VERSION").to_string(),
2107 git_sha: if mold_core::build_info::GIT_SHA == "unknown" {
2108 None
2109 } else {
2110 Some(mold_core::build_info::GIT_SHA.to_string())
2111 },
2112 build_date: if mold_core::build_info::BUILD_DATE == "unknown" {
2113 None
2114 } else {
2115 Some(mold_core::build_info::BUILD_DATE.to_string())
2116 },
2117 models_loaded,
2118 busy,
2119 current_generation,
2120 gpu_info: query_gpu_info(),
2121 uptime_secs: state.start_time.elapsed().as_secs(),
2122 hostname: hostname::get().ok().and_then(|h| h.into_string().ok()),
2123 memory_status: mold_inference::device::memory_status_string(),
2124 gpus: if has_gpus { Some(gpu_statuses) } else { None },
2125 queue_depth: Some(state.queue.pending()),
2126 queue_capacity: Some(state.queue_capacity),
2127 })
2128}
2129
2130#[utoipa::path(
2133 get,
2134 path = "/health",
2135 tag = "server",
2136 responses(
2137 (status = 200, description = "Server is healthy"),
2138 )
2139)]
2140async fn health() -> impl IntoResponse {
2141 StatusCode::OK
2142}
2143
2144#[utoipa::path(
2151 get,
2152 path = "/api/queue",
2153 tag = "queue",
2154 responses(
2155 (status = 200, description = "Queue snapshot", body = crate::job_registry::QueueListing),
2156 )
2157)]
2158async fn list_queue(State(state): State<AppState>) -> Json<crate::job_registry::QueueListing> {
2159 Json(state.job_registry.snapshot())
2160}
2161
2162#[derive(Debug, Deserialize, utoipa::ToSchema)]
2163struct QueuePatchRequest {
2164 target_gpu: Option<usize>,
2165}
2166
2167#[utoipa::path(
2168 patch,
2169 path = "/api/queue/{id}",
2170 tag = "queue",
2171 request_body = QueuePatchRequest,
2172 responses(
2173 (status = 200, description = "Updated queue entry", body = crate::job_registry::JobEntry),
2174 (status = 404, description = "Queue job not found"),
2175 (status = 409, description = "Queue job is already running"),
2176 (status = 422, description = "Invalid GPU target"),
2177 )
2178)]
2179async fn patch_queue_job(
2180 State(state): State<AppState>,
2181 Path(id): Path<String>,
2182 Json(req): Json<QueuePatchRequest>,
2183) -> Result<Json<crate::job_registry::JobEntry>, ApiError> {
2184 if let Some(target) = req.target_gpu {
2185 let available = state
2186 .gpu_pool
2187 .workers
2188 .iter()
2189 .any(|w| w.gpu.ordinal == target);
2190 if !available {
2191 return Err(ApiError::validation(format!(
2192 "gpu:{target} is not available in this server's worker pool"
2193 )));
2194 }
2195 }
2196
2197 state
2198 .job_registry
2199 .set_target_gpu(&id, req.target_gpu)
2200 .map_err(|e| match e {
2201 crate::job_registry::TargetGpuUpdateError::NotFound => {
2202 ApiError::queue_job_not_found(format!("queue job {id} not found"))
2203 }
2204 crate::job_registry::TargetGpuUpdateError::AlreadyRunning => {
2205 ApiError::queue_job_running(format!(
2206 "queue job {id} is already running; lane changes only apply to queued jobs"
2207 ))
2208 }
2209 })?;
2210
2211 let entry = state
2212 .job_registry
2213 .entry(&id)
2214 .ok_or_else(|| ApiError::queue_job_not_found(format!("queue job {id} not found")))?;
2215 Ok(Json(entry))
2216}
2217
2218#[utoipa::path(
2224 delete,
2225 path = "/api/queue/{id}",
2226 tag = "queue",
2227 params(("id" = String, Path, description = "Queue job id")),
2228 responses(
2229 (status = 204, description = "Queued job cancelled"),
2230 (status = 404, description = "Queue job not found"),
2231 (status = 409, description = "Queue job is already running"),
2232 )
2233)]
2234async fn cancel_queue_job(
2235 State(state): State<AppState>,
2236 Path(id): Path<String>,
2237) -> Result<StatusCode, ApiError> {
2238 state.job_registry.cancel_queued(&id).map_err(|e| match e {
2239 crate::job_registry::QueuedJobCancelError::NotFound => {
2240 ApiError::queue_job_not_found(format!("queue job {id} not found"))
2241 }
2242 crate::job_registry::QueuedJobCancelError::AlreadyRunning => ApiError::queue_job_running(
2243 format!("queue job {id} is already running; only queued jobs can be cancelled"),
2244 ),
2245 })?;
2246 Ok(StatusCode::NO_CONTENT)
2247}
2248
2249const HISTORY_DEFAULT_LIMIT: usize = 50;
2253const HISTORY_MAX_LIMIT: usize = 500;
2255
2256const HISTORY_UNAVAILABLE: &str = "HISTORY_UNAVAILABLE";
2258
2259fn history_db(state: &AppState) -> Result<&mold_db::MetadataDb, ApiError> {
2260 state.metadata_db.as_ref().as_ref().ok_or_else(|| {
2261 ApiError::with_code(
2262 "prompt history is unavailable because the metadata DB is disabled",
2263 HISTORY_UNAVAILABLE,
2264 StatusCode::SERVICE_UNAVAILABLE,
2265 )
2266 })
2267}
2268
2269#[derive(Debug, Deserialize)]
2270struct HistoryListQuery {
2271 query: Option<String>,
2273 limit: Option<usize>,
2275}
2276
2277#[utoipa::path(
2278 get,
2279 path = "/api/history",
2280 tag = "server",
2281 params(
2282 ("query" = Option<String>, Query, description = "Substring filter over prompt text (case-insensitive)"),
2283 ("limit" = Option<usize>, Query, description = "Max rows to return (default 50, max 500)"),
2284 ),
2285 responses(
2286 (status = 200, description = "Prompt history, newest first", body = mold_core::HistoryListing),
2287 (status = 503, description = "Metadata DB disabled"),
2288 )
2289)]
2290async fn list_history(
2291 State(state): State<AppState>,
2292 axum::extract::Query(params): axum::extract::Query<HistoryListQuery>,
2293) -> Result<Json<mold_core::HistoryListing>, ApiError> {
2294 let db = history_db(&state)?;
2295 let limit = params
2296 .limit
2297 .unwrap_or(HISTORY_DEFAULT_LIMIT)
2298 .min(HISTORY_MAX_LIMIT);
2299 let history = mold_db::PromptHistory::new(db);
2300 let rows = match params
2301 .query
2302 .as_deref()
2303 .map(str::trim)
2304 .filter(|q| !q.is_empty())
2305 {
2306 Some(query) => history.search(query, limit),
2307 None => history.recent(limit),
2308 }
2309 .map_err(|e| ApiError::internal(format!("failed to read prompt history: {e:#}")))?;
2310 let entries = rows
2311 .into_iter()
2312 .map(|e| mold_core::HistoryEntry {
2313 prompt: e.prompt,
2314 model: e.model,
2315 used_at: e.created_at_ms,
2316 })
2317 .collect();
2318 Ok(Json(mold_core::HistoryListing { entries }))
2319}
2320
2321#[derive(Debug, Deserialize)]
2322struct HistoryDeleteQuery {
2323 keep: Option<usize>,
2325}
2326
2327#[utoipa::path(
2328 delete,
2329 path = "/api/history",
2330 tag = "server",
2331 params(
2332 ("keep" = Option<usize>, Query, description = "Keep only the most recent N entries instead of clearing everything"),
2333 ),
2334 responses(
2335 (status = 204, description = "Prompt history cleared (or trimmed)"),
2336 (status = 503, description = "Metadata DB disabled"),
2337 )
2338)]
2339async fn delete_history(
2340 State(state): State<AppState>,
2341 axum::extract::Query(params): axum::extract::Query<HistoryDeleteQuery>,
2342) -> Result<StatusCode, ApiError> {
2343 let db = history_db(&state)?;
2344 let history = mold_db::PromptHistory::new(db);
2345 match params.keep {
2346 Some(keep) => history.trim_to(keep),
2347 None => history.clear(),
2348 }
2349 .map_err(|e| ApiError::internal(format!("failed to clear prompt history: {e:#}")))?;
2350 Ok(StatusCode::NO_CONTENT)
2351}
2352
2353async fn server_capabilities() -> Json<mold_core::ServerCapabilities> {
2359 let catalog_available = std::env::var("MOLD_CATALOG_DISABLE")
2360 .map(|v| v != "1" && !v.eq_ignore_ascii_case("true"))
2361 .unwrap_or(true);
2362
2363 Json(mold_core::ServerCapabilities {
2364 gallery: mold_core::GalleryCapabilities { can_delete: true },
2365 catalog: mold_core::CatalogCapabilities {
2366 available: catalog_available,
2367 families: mold_catalog::families::ALL_FAMILIES
2368 .iter()
2369 .map(|f| f.as_str().to_string())
2370 .collect::<Vec<_>>(),
2371 },
2372 })
2373}
2374
2375#[utoipa::path(
2378 get,
2379 path = "/api/capabilities/chain-limits",
2380 tag = "server",
2381 params(
2382 ("model" = String, Query, description = "Model name (e.g. ltx-2-19b-distilled:fp8)")
2383 ),
2384 responses(
2385 (status = 200, description = "Chain limits for the requested model",
2386 body = crate::chain_limits::ChainLimits),
2387 (status = 400, description = "Missing required 'model' query parameter"),
2388 (status = 404, description = "Unknown or unsupported model"),
2389 )
2390)]
2391async fn capabilities_chain_limits(
2392 axum::extract::Query(params): axum::extract::Query<std::collections::HashMap<String, String>>,
2393) -> axum::response::Response {
2394 let raw_model = match params.get("model") {
2395 Some(m) => m.clone(),
2396 None => {
2397 return (
2398 StatusCode::BAD_REQUEST,
2399 "missing required 'model' query parameter\n",
2400 )
2401 .into_response();
2402 }
2403 };
2404
2405 let resolved = mold_core::manifest::resolve_model_name(&raw_model);
2406 let Some(manifest) = mold_core::manifest::find_manifest(&resolved) else {
2407 return (StatusCode::NOT_FOUND, "unknown model\n").into_response();
2408 };
2409 let family = manifest.family.clone();
2410
2411 if crate::chain_limits::family_cap(&family).is_none() {
2412 return (StatusCode::NOT_FOUND, "model is not chain-capable\n").into_response();
2413 }
2414
2415 let quant = resolved
2416 .split_once(':')
2417 .map(|(_, tag)| tag.to_string())
2418 .unwrap_or_default();
2419
2420 let limits = crate::chain_limits::compute_limits(&resolved, &family, &quant, 0);
2422 Json(limits).into_response()
2423}
2424
2425#[utoipa::path(
2433 post,
2434 path = "/api/shutdown",
2435 tag = "server",
2436 responses(
2437 (status = 200, description = "Shutdown initiated"),
2438 (status = 403, description = "Forbidden — remote shutdown requires API key auth"),
2439 )
2440)]
2441async fn shutdown_server(State(state): State<AppState>, request: Request) -> impl IntoResponse {
2442 let auth_enabled = request
2445 .extensions()
2446 .get::<crate::auth::AuthState>()
2447 .is_some_and(|s| s.is_some());
2448
2449 if !auth_enabled {
2450 let is_loopback = request
2451 .extensions()
2452 .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
2453 .map(|ci| ci.0.ip().is_loopback())
2454 .unwrap_or(false);
2455 if !is_loopback {
2456 return (
2457 StatusCode::FORBIDDEN,
2458 "shutdown requires API key auth or localhost access\n",
2459 );
2460 }
2461 }
2462
2463 tracing::info!("shutdown requested via API");
2464 if let Some(tx) = state.shutdown_tx.lock().await.take() {
2465 let _ = tx.send(());
2466 }
2467 (StatusCode::OK, "shutdown initiated\n")
2468}
2469
2470async fn list_gallery(
2480 State(state): State<AppState>,
2481) -> Result<Json<Vec<mold_core::GalleryImage>>, ApiError> {
2482 let config = state.config.read().await;
2483 if config.is_output_disabled() {
2484 return Ok(Json(Vec::new()));
2485 }
2486 let output_dir = config.effective_output_dir();
2487 drop(config);
2488
2489 if !output_dir.is_dir() {
2490 return Ok(Json(Vec::new()));
2491 }
2492
2493 if state.metadata_db.is_some() {
2494 let db_arc = state.metadata_db.clone();
2495 let dir = output_dir.clone();
2496 let listed = tokio::task::spawn_blocking(move || {
2497 db_arc
2498 .as_ref()
2499 .as_ref()
2500 .map(|db| db.list(Some(&dir)))
2501 .transpose()
2502 })
2503 .await
2504 .map_err(|e| ApiError::internal(format!("gallery DB query failed: {e}")))?
2505 .map_err(|e| ApiError::internal(format!("gallery DB query failed: {e:#}")))?;
2506 if let Some(rows) = listed {
2507 if !rows.is_empty() {
2508 let images = rows.iter().map(|r| r.to_gallery_image()).collect();
2509 return Ok(Json(images));
2510 }
2511 }
2512 }
2513
2514 let images = tokio::task::spawn_blocking(move || scan_gallery_dir(&output_dir))
2515 .await
2516 .map_err(|e| ApiError::internal(format!("gallery scan failed: {e}")))?;
2517
2518 Ok(Json(images))
2519}
2520
2521async fn get_gallery_image(
2531 State(state): State<AppState>,
2532 headers: HeaderMap,
2533 axum::extract::Path(filename): axum::extract::Path<String>,
2534) -> Result<axum::response::Response, ApiError> {
2535 let config = state.config.read().await;
2536 if config.is_output_disabled() {
2537 return Err(ApiError::not_found("image output is disabled"));
2538 }
2539 let output_dir = config.effective_output_dir();
2540 drop(config);
2541
2542 let clean_name = std::path::Path::new(&filename)
2544 .file_name()
2545 .map(|f| f.to_string_lossy().to_string())
2546 .unwrap_or_default();
2547
2548 if clean_name.is_empty() || clean_name != filename {
2549 return Err(ApiError::validation("invalid filename"));
2550 }
2551
2552 let path = output_dir.join(&clean_name);
2553 let meta = match tokio::fs::metadata(&path).await {
2554 Ok(m) if m.is_file() => m,
2555 _ => {
2556 return Err(ApiError::not_found(format!(
2557 "image not found: {clean_name}"
2558 )));
2559 }
2560 };
2561 let total_len = meta.len();
2562 let content_type = content_type_for_filename(&clean_name);
2563
2564 let range_header = headers
2565 .get(header::RANGE)
2566 .and_then(|v| v.to_str().ok())
2567 .map(|s| s.to_string());
2568
2569 let file = tokio::fs::File::open(&path)
2570 .await
2571 .map_err(|e| ApiError::internal(format!("failed to open file: {e}")))?;
2572
2573 if let Some(raw) = range_header {
2574 if let Some((start, end)) = parse_byte_range(&raw, total_len) {
2575 return serve_range(file, start, end, total_len, content_type).await;
2576 } else {
2577 return Ok(axum::response::Response::builder()
2579 .status(StatusCode::RANGE_NOT_SATISFIABLE)
2580 .header(header::CONTENT_RANGE, format!("bytes */{total_len}"))
2581 .body(axum::body::Body::empty())
2582 .unwrap());
2583 }
2584 }
2585
2586 let stream = tokio_util::io::ReaderStream::new(file);
2588 let body = axum::body::Body::from_stream(stream);
2589 Ok(axum::response::Response::builder()
2590 .status(StatusCode::OK)
2591 .header(header::CONTENT_TYPE, content_type)
2592 .header(header::ACCEPT_RANGES, "bytes")
2593 .header(header::CONTENT_LENGTH, total_len)
2594 .header(header::CACHE_CONTROL, "public, max-age=3600")
2595 .body(body)
2596 .unwrap())
2597}
2598
2599fn parse_byte_range(header: &str, total_len: u64) -> Option<(u64, u64)> {
2607 let spec = header.strip_prefix("bytes=")?;
2608 if spec.contains(',') {
2609 return None;
2610 }
2611 let (start_s, end_s) = spec.split_once('-')?;
2612 let start_s = start_s.trim();
2613 let end_s = end_s.trim();
2614
2615 if total_len == 0 {
2616 return None;
2617 }
2618
2619 if start_s.is_empty() {
2620 let suffix: u64 = end_s.parse().ok()?;
2622 if suffix == 0 {
2623 return None;
2624 }
2625 let start = total_len.saturating_sub(suffix);
2626 return Some((start, total_len - 1));
2627 }
2628
2629 let start: u64 = start_s.parse().ok()?;
2630 if start >= total_len {
2631 return None;
2632 }
2633 let end: u64 = if end_s.is_empty() {
2634 total_len - 1
2635 } else {
2636 end_s.parse().ok()?
2637 };
2638 let end = end.min(total_len - 1);
2639 if end < start {
2640 return None;
2641 }
2642 Some((start, end))
2643}
2644
2645async fn serve_range(
2649 mut file: tokio::fs::File,
2650 start: u64,
2651 end: u64,
2652 total_len: u64,
2653 content_type: &'static str,
2654) -> Result<axum::response::Response, ApiError> {
2655 use tokio::io::{AsyncReadExt, AsyncSeekExt};
2656 file.seek(std::io::SeekFrom::Start(start))
2657 .await
2658 .map_err(|e| ApiError::internal(format!("seek failed: {e}")))?;
2659 let len = end - start + 1;
2660 let stream = tokio_util::io::ReaderStream::new(file.take(len));
2661 let body = axum::body::Body::from_stream(stream);
2662 Ok(axum::response::Response::builder()
2663 .status(StatusCode::PARTIAL_CONTENT)
2664 .header(header::CONTENT_TYPE, content_type)
2665 .header(header::ACCEPT_RANGES, "bytes")
2666 .header(header::CONTENT_LENGTH, len)
2667 .header(
2668 header::CONTENT_RANGE,
2669 format!("bytes {start}-{end}/{total_len}"),
2670 )
2671 .header(header::CACHE_CONTROL, "public, max-age=300")
2674 .body(body)
2675 .unwrap())
2676}
2677
2678fn content_type_for_filename(name: &str) -> &'static str {
2681 let lower = name.to_ascii_lowercase();
2682 if lower.ends_with(".png") {
2683 "image/png"
2684 } else if lower.ends_with(".jpg") || lower.ends_with(".jpeg") {
2685 "image/jpeg"
2686 } else if lower.ends_with(".gif") {
2687 "image/gif"
2688 } else if lower.ends_with(".webp") {
2689 "image/webp"
2690 } else if lower.ends_with(".apng") {
2691 "image/apng"
2692 } else if lower.ends_with(".mp4") {
2693 "video/mp4"
2694 } else {
2695 "application/octet-stream"
2696 }
2697}
2698
2699async fn delete_gallery_image(
2704 State(state): State<AppState>,
2705 axum::extract::Path(filename): axum::extract::Path<String>,
2706) -> Result<impl IntoResponse, ApiError> {
2707 let config = state.config.read().await;
2708 if config.is_output_disabled() {
2709 return Err(ApiError::not_found("image output is disabled"));
2710 }
2711 let output_dir = config.effective_output_dir();
2712 drop(config);
2713
2714 let clean_name = std::path::Path::new(&filename)
2715 .file_name()
2716 .map(|f| f.to_string_lossy().to_string())
2717 .unwrap_or_default();
2718
2719 if clean_name.is_empty() || clean_name != filename {
2720 return Err(ApiError::validation("invalid filename"));
2721 }
2722
2723 let path = output_dir.join(&clean_name);
2724 if path.is_file() {
2725 std::fs::remove_file(&path)
2726 .map_err(|e| ApiError::internal(format!("failed to delete image: {e}")))?;
2727 }
2728
2729 let thumb_dir = server_thumbnail_dir();
2734 let _ = std::fs::remove_file(thumb_dir.join(&clean_name));
2735 let _ = std::fs::remove_file(thumb_dir.join(format!("{clean_name}.png")));
2736 let _ = std::fs::remove_file(
2737 server_preview_gif_dir().join(mold_core::media_paths::preview_gif_filename(&clean_name)),
2738 );
2739
2740 if let Some(db) = state.metadata_db.as_ref().as_ref() {
2744 match db.delete(&output_dir, &clean_name) {
2745 Ok(true) => {}
2746 Ok(false) => tracing::debug!(
2747 "delete: no metadata row for {}",
2748 output_dir.join(&clean_name).display()
2749 ),
2750 Err(e) => tracing::warn!(
2751 "metadata DB delete failed for {}: {e:#}",
2752 output_dir.join(&clean_name).display()
2753 ),
2754 }
2755 }
2756
2757 Ok(StatusCode::NO_CONTENT)
2758}
2759
2760async fn get_gallery_thumbnail(
2763 State(state): State<AppState>,
2764 axum::extract::Path(filename): axum::extract::Path<String>,
2765) -> Result<impl IntoResponse, ApiError> {
2766 let config = state.config.read().await;
2767 if config.is_output_disabled() {
2768 return Err(ApiError::not_found("image output is disabled"));
2769 }
2770 let output_dir = config.effective_output_dir();
2771 drop(config);
2772
2773 let clean_name = std::path::Path::new(&filename)
2774 .file_name()
2775 .map(|f| f.to_string_lossy().to_string())
2776 .unwrap_or_default();
2777
2778 if clean_name.is_empty() || clean_name != filename {
2779 return Err(ApiError::validation("invalid filename"));
2780 }
2781
2782 let source_path = output_dir.join(&clean_name);
2783 if !source_path.is_file() {
2784 return Err(ApiError::not_found(format!(
2785 "image not found: {clean_name}"
2786 )));
2787 }
2788
2789 let thumb_dir = server_thumbnail_dir();
2793 let thumb_path = thumb_dir.join(format!("{clean_name}.png"));
2794 let lower = clean_name.to_ascii_lowercase();
2795 let is_video = lower.ends_with(".mp4");
2796
2797 if !thumb_path.is_file() {
2798 let source = source_path.clone();
2805 let dest = thumb_path.clone();
2806 let gen_result = tokio::task::spawn_blocking(move || {
2807 if is_video {
2808 generate_video_thumbnail(&source, &dest)
2809 } else {
2810 generate_server_thumbnail(&source, &dest)
2811 }
2812 })
2813 .await
2814 .map_err(|e| ApiError::internal(format!("thumbnail generation failed: {e}")))?;
2815
2816 if let Err(err) = gen_result {
2817 tracing::warn!(
2818 file = %clean_name,
2819 error = %err,
2820 "thumbnail decode failed; falling back to source bytes"
2821 );
2822 if is_video {
2826 let mut headers = HeaderMap::new();
2827 headers.insert(
2828 header::CONTENT_TYPE,
2829 HeaderValue::from_static("image/svg+xml"),
2830 );
2831 headers.insert(
2832 header::CACHE_CONTROL,
2833 HeaderValue::from_static("public, max-age=300"),
2834 );
2835 return Ok((headers, VIDEO_PLACEHOLDER_SVG.as_bytes().to_vec()));
2836 }
2837 let raw = tokio::fs::read(&source_path)
2838 .await
2839 .map_err(|e| ApiError::internal(format!("failed to read source: {e}")))?;
2840 let mut headers = HeaderMap::new();
2841 headers.insert(
2842 header::CONTENT_TYPE,
2843 HeaderValue::from_static(content_type_for_filename(&clean_name)),
2844 );
2845 headers.insert(
2846 header::CACHE_CONTROL,
2847 HeaderValue::from_static("public, max-age=300"),
2848 );
2849 return Ok((headers, raw));
2850 }
2851 }
2852
2853 let data = tokio::fs::read(&thumb_path)
2854 .await
2855 .map_err(|e| ApiError::internal(format!("failed to read thumbnail: {e}")))?;
2856
2857 let mut headers = HeaderMap::new();
2858 headers.insert(header::CONTENT_TYPE, HeaderValue::from_static("image/png"));
2859 headers.insert(
2860 header::CACHE_CONTROL,
2861 HeaderValue::from_static("public, max-age=3600"),
2862 );
2863
2864 Ok((headers, data))
2865}
2866
2867const VIDEO_PLACEHOLDER_SVG: &str = r##"<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 256 256" width="256" height="256"><defs><linearGradient id="g" x1="0" y1="0" x2="1" y2="1"><stop offset="0" stop-color="#1e293b"/><stop offset="1" stop-color="#0f172a"/></linearGradient></defs><rect width="256" height="256" fill="url(#g)"/><circle cx="128" cy="128" r="52" fill="rgba(255,255,255,0.08)"/><polygon points="112,100 112,156 160,128" fill="rgba(226,232,240,0.85)"/></svg>"##;
2868
2869async fn get_gallery_preview(
2879 State(state): State<AppState>,
2880 axum::extract::Path(filename): axum::extract::Path<String>,
2881) -> Result<axum::response::Response, ApiError> {
2882 let config = state.config.read().await;
2883 if config.is_output_disabled() {
2884 return Err(ApiError::not_found("image output is disabled"));
2885 }
2886 let output_dir = config.effective_output_dir();
2887 drop(config);
2888
2889 let clean_name = std::path::Path::new(&filename)
2892 .file_name()
2893 .map(|f| f.to_string_lossy().to_string())
2894 .unwrap_or_default();
2895 if clean_name.is_empty() || clean_name != filename {
2896 return Err(ApiError::validation("invalid filename"));
2897 }
2898
2899 let source_path = output_dir.join(&clean_name);
2906 if !tokio::fs::metadata(&source_path)
2907 .await
2908 .map(|m| m.is_file())
2909 .unwrap_or(false)
2910 {
2911 return Err(ApiError::not_found(format!(
2912 "image not found: {clean_name}"
2913 )));
2914 }
2915
2916 let preview_path =
2917 server_preview_gif_dir().join(mold_core::media_paths::preview_gif_filename(&clean_name));
2918 let meta = match tokio::fs::metadata(&preview_path).await {
2919 Ok(m) if m.is_file() => m,
2920 _ => {
2921 return Err(ApiError::not_found(format!(
2922 "preview not found: {clean_name}"
2923 )));
2924 }
2925 };
2926 let total_len = meta.len();
2927
2928 let file = tokio::fs::File::open(&preview_path)
2929 .await
2930 .map_err(|e| ApiError::internal(format!("failed to open preview: {e}")))?;
2931 let stream = tokio_util::io::ReaderStream::new(file);
2932 let body = axum::body::Body::from_stream(stream);
2933 Ok(axum::response::Response::builder()
2934 .status(StatusCode::OK)
2935 .header(header::CONTENT_TYPE, "image/gif")
2936 .header(header::CONTENT_LENGTH, total_len)
2937 .header(header::CACHE_CONTROL, "public, max-age=3600")
2938 .body(body)
2939 .unwrap())
2940}
2941
2942fn server_preview_gif_dir() -> std::path::PathBuf {
2947 mold_core::Config::mold_dir()
2948 .unwrap_or_else(|| std::path::PathBuf::from(".mold"))
2949 .join("cache")
2950 .join("previews")
2951}
2952
2953fn server_thumbnail_dir() -> std::path::PathBuf {
2955 mold_core::Config::mold_dir()
2956 .unwrap_or_else(|| std::path::PathBuf::from(".mold"))
2957 .join("cache")
2958 .join("thumbnails")
2959}
2960
2961fn generate_server_thumbnail(
2965 source: &std::path::Path,
2966 dest: &std::path::Path,
2967) -> anyhow::Result<()> {
2968 let img = image::open(source)?;
2969 let thumb = img.thumbnail(256, 256);
2970 if let Some(parent) = dest.parent() {
2971 std::fs::create_dir_all(parent)?;
2972 }
2973 thumb.save_with_format(dest, image::ImageFormat::Png)?;
2974 Ok(())
2975}
2976
2977fn generate_video_thumbnail(
2985 source: &std::path::Path,
2986 dest: &std::path::Path,
2987) -> anyhow::Result<()> {
2988 if let Some(parent) = dest.parent() {
2989 std::fs::create_dir_all(parent)?;
2990 }
2991 let tmp = dest.with_extension("firstframe.png");
2996 mold_inference::ltx2::media::extract_thumbnail(source, &tmp)?;
2997 let decode_result = (|| -> anyhow::Result<()> {
2998 let img = image::open(&tmp)?;
2999 let thumb = img.thumbnail(256, 256);
3000 thumb.save_with_format(dest, image::ImageFormat::Png)?;
3001 Ok(())
3002 })();
3003 let _ = std::fs::remove_file(&tmp);
3004 decode_result
3005}
3006
3007pub fn spawn_thumbnail_warmup(config: &mold_core::Config) {
3009 if !thumbnail_warmup_enabled() {
3010 tracing::info!("thumbnail warmup disabled; thumbnails will be generated on demand");
3011 return;
3012 }
3013
3014 let output_dir = config.effective_output_dir();
3015 std::thread::spawn(move || {
3016 if !output_dir.is_dir() {
3017 return;
3018 }
3019 let thumb_dir = server_thumbnail_dir();
3020 let walker = walkdir::WalkDir::new(&output_dir).max_depth(1).into_iter();
3021 for entry in walker.filter_map(|e| e.ok()) {
3022 let path = entry.path();
3023 if !path.is_file() {
3024 continue;
3025 }
3026 let ext = path
3027 .extension()
3028 .and_then(|e| e.to_str())
3029 .map(|e| e.to_lowercase());
3030 let is_raster = matches!(
3031 ext.as_deref(),
3032 Some("png" | "jpg" | "jpeg" | "gif" | "apng" | "webp")
3033 );
3034 let is_video = matches!(ext.as_deref(), Some("mp4"));
3035 if !is_raster && !is_video {
3036 continue;
3037 }
3038 let filename = path
3039 .file_name()
3040 .map(|f| f.to_string_lossy().to_string())
3041 .unwrap_or_default();
3042 let thumb_path = thumb_dir.join(format!("{filename}.png"));
3043 if thumb_path.is_file() {
3044 continue;
3045 }
3046 let result = if is_video {
3047 generate_video_thumbnail(path, &thumb_path)
3048 } else {
3049 generate_server_thumbnail(path, &thumb_path)
3050 };
3051 if let Err(e) = result {
3052 tracing::warn!("failed to generate thumbnail for {}: {e}", path.display());
3053 }
3054 }
3055 tracing::info!("thumbnail warmup complete");
3056 });
3057}
3058
3059fn thumbnail_warmup_enabled() -> bool {
3060 std::env::var("MOLD_THUMBNAIL_WARMUP")
3061 .map(|v| matches!(v.as_str(), "1" | "true" | "TRUE" | "yes" | "YES"))
3062 .unwrap_or(false)
3063}
3064
3065fn scan_gallery_dir(dir: &std::path::Path) -> Vec<mold_core::GalleryImage> {
3084 let mut images: Vec<mold_core::GalleryImage> = mold_db::scan::scan_output_dir(dir)
3085 .filter_map(|item| match item {
3086 mold_db::scan::ScanItem::Valid(file) => Some(file),
3087 _ => None,
3088 })
3089 .map(|file| {
3090 let timestamp = file.timestamp_secs();
3091 let size_bytes = file.size_u64();
3092 let (metadata, synthetic) = mold_db::metadata_io::read_or_synthesize(
3093 &file.path,
3094 file.format,
3095 &file.filename,
3096 timestamp,
3097 );
3098 mold_core::GalleryImage {
3099 filename: file.filename,
3100 metadata,
3101 timestamp,
3102 format: Some(file.format),
3103 size_bytes: Some(size_bytes),
3104 metadata_synthetic: synthetic,
3105 }
3106 })
3107 .collect();
3108
3109 images.sort_by_key(|img| std::cmp::Reverse(img.timestamp));
3110 images
3111}
3112
3113async fn get_model_placement(
3121 State(state): State<AppState>,
3122 axum::extract::Path(name): axum::extract::Path<String>,
3123) -> Result<Json<mold_core::types::DevicePlacement>, ApiError> {
3124 let cfg = state.config.read().await;
3125 match cfg.models.get(&name).and_then(|mc| mc.placement.clone()) {
3126 Some(placement) => Ok(Json(placement)),
3127 None => Err(ApiError::not_found(format!(
3128 "no placement saved for model '{name}'"
3129 ))),
3130 }
3131}
3132
3133async fn put_model_placement(
3134 State(state): State<AppState>,
3135 axum::extract::Path(name): axum::extract::Path<String>,
3136 Json(placement): Json<mold_core::types::DevicePlacement>,
3137) -> Result<Json<serde_json::Value>, ApiError> {
3138 validate_multi_gpu_placement(&state, Some(&placement))?;
3139 {
3140 let mut cfg = state.config.write().await;
3141 cfg.set_model_placement(&name, Some(placement.clone()));
3142 cfg.save().map_err(|e| {
3143 tracing::warn!("failed to persist placement to config.toml: {e}");
3144 ApiError::internal(format!("failed to persist placement to config.toml: {e}"))
3145 })?;
3146 }
3147 Ok(Json(serde_json::json!({
3148 "ok": true,
3149 "model": name,
3150 })))
3151}
3152
3153async fn delete_model_placement(
3154 State(state): State<AppState>,
3155 axum::extract::Path(name): axum::extract::Path<String>,
3156) -> Result<Json<serde_json::Value>, ApiError> {
3157 let mut cfg = state.config.write().await;
3158 cfg.set_model_placement(&name, None);
3159 cfg.save().map_err(|e| {
3160 tracing::warn!("failed to persist placement removal to config.toml: {e}");
3161 ApiError::internal(format!(
3162 "failed to persist placement removal to config.toml: {e}"
3163 ))
3164 })?;
3165 Ok(Json(serde_json::json!({ "ok": true })))
3166}
3167
3168async fn openapi_json() -> impl IntoResponse {
3171 Json(ApiDoc::openapi())
3172}
3173
3174async fn scalar_docs() -> impl IntoResponse {
3177 (
3178 [(header::CONTENT_TYPE, "text/html")],
3179 r#"<!DOCTYPE html>
3180<html>
3181<head>
3182 <title>mold API</title>
3183 <meta charset="utf-8" />
3184 <meta name="viewport" content="width=device-width, initial-scale=1" />
3185</head>
3186<body>
3187 <script id="api-reference" data-url="/api/openapi.json"></script>
3188 <script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference"></script>
3189</body>
3190</html>"#,
3191 )
3192}
3193
3194fn query_gpu_info() -> Option<GpuInfo> {
3197 let nvidia_smi = if std::path::Path::new("/run/current-system/sw/bin/nvidia-smi").exists() {
3198 "/run/current-system/sw/bin/nvidia-smi"
3199 } else {
3200 "nvidia-smi"
3201 };
3202
3203 let output = std::process::Command::new(nvidia_smi)
3204 .args([
3205 "--query-gpu=name,memory.total,memory.used",
3206 "--format=csv,noheader,nounits",
3207 ])
3208 .output()
3209 .ok()?;
3210
3211 if !output.status.success() {
3212 return None;
3213 }
3214
3215 let text = String::from_utf8(output.stdout).ok()?;
3216 let line = text.lines().next()?;
3217 let parts: Vec<&str> = line.split(',').map(str::trim).collect();
3218 if parts.len() < 3 {
3219 return None;
3220 }
3221
3222 Some(GpuInfo {
3223 name: parts[0].to_string(),
3224 vram_total_mb: parts[1].parse().ok()?,
3225 vram_used_mb: parts[2].parse().ok()?,
3226 })
3227}
3228
3229#[derive(serde::Deserialize, utoipa::ToSchema)]
3232pub struct CreateDownloadBody {
3233 pub model: String,
3234}
3235
3236#[derive(serde::Serialize, utoipa::ToSchema)]
3237pub struct CreateDownloadResponse {
3238 pub id: String,
3239 pub position: usize,
3240}
3241
3242#[utoipa::path(
3243 post,
3244 path = "/api/downloads",
3245 tag = "downloads",
3246 request_body = CreateDownloadBody,
3247 responses(
3248 (status = 200, description = "Enqueued; position 0 = will start immediately", body = CreateDownloadResponse),
3249 (status = 400, description = "Unknown model"),
3250 (status = 409, description = "Already active or queued; body contains existing id", body = CreateDownloadResponse),
3251 )
3252)]
3253pub async fn create_download(
3254 State(state): State<AppState>,
3255 Json(body): Json<CreateDownloadBody>,
3256) -> axum::response::Response {
3257 use crate::downloads::{EnqueueError, EnqueueOutcome};
3258 match state.downloads.enqueue(body.model.clone()).await {
3259 Ok((id, position, EnqueueOutcome::Created)) => (
3260 StatusCode::OK,
3261 Json(CreateDownloadResponse { id, position }),
3262 )
3263 .into_response(),
3264 Ok((id, position, EnqueueOutcome::AlreadyPresent)) => (
3265 StatusCode::CONFLICT,
3266 Json(CreateDownloadResponse { id, position }),
3267 )
3268 .into_response(),
3269 Err(EnqueueError::UnknownModel(_)) => (
3270 StatusCode::BAD_REQUEST,
3271 Json(serde_json::json!({
3272 "error": format!("unknown model '{}'. Run 'mold list' to see available models.", body.model)
3273 })),
3274 )
3275 .into_response(),
3276 Err(EnqueueError::LockPoisoned) => (
3277 StatusCode::INTERNAL_SERVER_ERROR,
3278 Json(serde_json::json!({ "error": "download queue state is corrupt" })),
3279 )
3280 .into_response(),
3281 }
3282}
3283
3284#[utoipa::path(
3285 delete,
3286 path = "/api/downloads/{id}",
3287 tag = "downloads",
3288 params(("id" = String, Path, description = "Job id")),
3289 responses(
3290 (status = 204, description = "Cancelled"),
3291 (status = 404, description = "Unknown id"),
3292 )
3293)]
3294pub async fn delete_download(
3295 State(state): State<AppState>,
3296 axum::extract::Path(id): axum::extract::Path<String>,
3297) -> axum::response::Response {
3298 if state.downloads.cancel(&id).await {
3299 StatusCode::NO_CONTENT.into_response()
3300 } else {
3301 (
3302 StatusCode::NOT_FOUND,
3303 Json(serde_json::json!({ "error": format!("unknown download id '{id}'") })),
3304 )
3305 .into_response()
3306 }
3307}
3308
3309#[utoipa::path(
3310 get,
3311 path = "/api/downloads",
3312 tag = "downloads",
3313 responses((status = 200, description = "Current queue state"))
3314)]
3315pub async fn list_downloads(State(state): State<AppState>) -> axum::response::Response {
3316 Json(state.downloads.listing().await).into_response()
3317}
3318
3319#[utoipa::path(
3320 get,
3321 path = "/api/downloads/stream",
3322 tag = "downloads",
3323 responses((status = 200, description = "SSE stream of DownloadEvent JSON")),
3324)]
3325pub async fn stream_downloads(
3326 State(state): State<AppState>,
3327) -> Sse<
3328 impl futures_core::Stream<Item = Result<axum::response::sse::Event, std::convert::Infallible>>,
3329> {
3330 use axum::response::sse::Event;
3331 use tokio_stream::wrappers::BroadcastStream;
3332 use tokio_stream::StreamExt as _;
3333
3334 let rx = state.downloads.subscribe();
3341 let initial = state.downloads.listing().await;
3342 let snapshot_event = mold_core::types::DownloadEvent::Snapshot { listing: initial };
3343
3344 let stream = async_stream::stream! {
3345 let data = serde_json::to_string(&snapshot_event).unwrap_or_else(|_| "{}".to_string());
3346 yield Ok::<_, std::convert::Infallible>(Event::default().event("download").data(data));
3347
3348 let mut bs = BroadcastStream::new(rx);
3349 while let Some(item) = bs.next().await {
3350 match item {
3351 Ok(event) => {
3352 let data = serde_json::to_string(&event)
3353 .unwrap_or_else(|_| "{}".to_string());
3354 yield Ok(Event::default().event("download").data(data));
3355 }
3356 Err(_lagged) => continue,
3360 }
3361 }
3362 };
3363
3364 Sse::new(stream).keep_alive(
3365 KeepAlive::new()
3366 .interval(std::time::Duration::from_secs(15))
3367 .text("ping"),
3368 )
3369}
3370
3371async fn get_resources(State(state): State<AppState>) -> Result<Json<ResourceSnapshot>, ApiError> {
3377 match state.resources.latest() {
3378 Some(snap) => Ok(Json(snap)),
3379 None => Err(ApiError::internal_with_status(
3380 "resource telemetry not ready",
3381 StatusCode::SERVICE_UNAVAILABLE,
3382 )),
3383 }
3384}
3385
3386async fn get_resources_stream(
3389 State(state): State<AppState>,
3390) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>> {
3391 use tokio_stream::wrappers::BroadcastStream;
3392
3393 let rx = state.resources.subscribe();
3394 let initial = state.resources.latest();
3397
3398 let stream = async_stream::stream! {
3399 if let Some(snap) = initial {
3400 yield Ok::<_, Infallible>(snapshot_to_sse(&snap));
3401 }
3402 let mut bs = BroadcastStream::new(rx);
3403 while let Some(item) = bs.next().await {
3404 match item {
3405 Ok(snap) => yield Ok(snapshot_to_sse(&snap)),
3406 Err(_lagged) => continue,
3409 }
3410 }
3411 };
3412
3413 Sse::new(stream).keep_alive(
3414 KeepAlive::new()
3415 .interval(std::time::Duration::from_secs(15))
3416 .text("ping"),
3417 )
3418}
3419
3420fn snapshot_to_sse(snap: &ResourceSnapshot) -> SseEvent {
3421 match serde_json::to_string(snap) {
3422 Ok(data) => SseEvent::default().event("snapshot").data(data),
3423 Err(e) => SseEvent::default()
3424 .event("error")
3425 .data(format!("{{\"message\":\"serialize failed: {e}\"}}")),
3426 }
3427}
3428
3429#[cfg(test)]
3430mod tests {
3431 use super::*;
3432
3433 fn env_lock() -> &'static std::sync::Mutex<()> {
3434 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
3435 &ENV_LOCK
3436 }
3437
3438 #[test]
3439 fn clean_error_message_strips_backtrace() {
3440 let err = anyhow::anyhow!(
3441 "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")\n\
3442 \x20 0: candle_core::error::Error::bt\n\
3443 \x20 at /home/user/.cargo/git/candle/src/error.rs:264:25\n\
3444 \x20 1: <core::result::Result<O,E> as candle_core::cuda_backend::error::WrapErr<O>>::w\n\
3445 \x20 at /home/user/.cargo/git/candle/src/cuda_backend/error.rs:60:65"
3446 );
3447 let msg = clean_error_message(&err);
3448 assert_eq!(
3449 msg,
3450 "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
3451 );
3452 }
3453
3454 #[test]
3455 fn clean_error_message_preserves_simple_error() {
3456 let err = anyhow::anyhow!("model not found: flux-dev:q4");
3457 let msg = clean_error_message(&err);
3458 assert_eq!(msg, "model not found: flux-dev:q4");
3459 }
3460
3461 #[test]
3462 fn clean_error_message_preserves_multiline_without_backtrace() {
3463 let err = anyhow::anyhow!("validation failed\nprompt is empty\nsteps must be > 0");
3464 let msg = clean_error_message(&err);
3465 assert_eq!(msg, "validation failed\nprompt is empty\nsteps must be > 0");
3466 }
3467
3468 #[test]
3469 fn clean_error_message_strips_high_numbered_frames() {
3470 let err = anyhow::anyhow!(
3471 "some error\n\
3472 \x20 10: tokio::runtime::task::core::Core<T,S>::poll at /home/user/.cargo/tokio/src/core.rs:375\n\
3473 \x20 11: std::panicking::catch_unwind at /nix/store/rust/src/panicking.rs:544"
3474 );
3475 let msg = clean_error_message(&err);
3476 assert_eq!(msg, "some error");
3477 }
3478
3479 #[test]
3480 fn clean_error_message_empty_fallback() {
3481 let err = anyhow::anyhow!("0: candle_core::error::Error::bt at /some/path.rs:10:5");
3483 let msg = clean_error_message(&err);
3484 assert!(!msg.is_empty());
3486 }
3487
3488 #[test]
3489 fn clean_error_message_renders_full_anyhow_chain() {
3490 let root = std::io::Error::new(std::io::ErrorKind::InvalidData, "bytes past end");
3495 let err: anyhow::Error = anyhow::Error::new(root)
3496 .context("validate single-file checkpoint at /tmp/foo.safetensors");
3497 let msg = clean_error_message(&err);
3498 assert!(
3499 msg.contains("validate single-file checkpoint") && msg.contains("bytes past end"),
3500 "expected both context layers in the rendered chain, got: {msg}",
3501 );
3502 }
3503
3504 #[test]
3505 fn save_image_to_dir_creates_directory_and_writes_file() {
3506 let dir = std::env::temp_dir().join(format!(
3507 "mold-save-test-{}",
3508 std::time::SystemTime::now()
3509 .duration_since(std::time::UNIX_EPOCH)
3510 .unwrap()
3511 .as_nanos()
3512 ));
3513 assert!(!dir.exists());
3514
3515 let img = mold_core::ImageData {
3516 data: vec![0x89, 0x50, 0x4E, 0x47], format: mold_core::OutputFormat::Png,
3518 width: 64,
3519 height: 64,
3520 index: 0,
3521 };
3522
3523 save_image_to_dir(&dir, &img, "test-model:q8", 1);
3524
3525 assert!(dir.exists(), "directory should be created");
3526 let files: Vec<_> = std::fs::read_dir(&dir).unwrap().collect();
3527 assert_eq!(files.len(), 1, "should have exactly one file");
3528 let file = files[0].as_ref().unwrap();
3529 let filename = file.file_name().to_str().unwrap().to_string();
3530 assert!(filename.starts_with("mold-test-model-q8-"), "{filename}");
3531 assert!(filename.ends_with(".png"), "{filename}");
3532 let contents = std::fs::read(file.path()).unwrap();
3533 assert_eq!(contents, vec![0x89, 0x50, 0x4E, 0x47]);
3534
3535 std::fs::remove_dir_all(&dir).ok();
3536 }
3537
3538 #[test]
3539 fn save_image_to_dir_batch_includes_index() {
3540 let dir = std::env::temp_dir().join(format!(
3541 "mold-save-batch-{}",
3542 std::time::SystemTime::now()
3543 .duration_since(std::time::UNIX_EPOCH)
3544 .unwrap()
3545 .as_nanos()
3546 ));
3547
3548 let img = mold_core::ImageData {
3549 data: vec![0xFF, 0xD8], format: mold_core::OutputFormat::Jpeg,
3551 width: 64,
3552 height: 64,
3553 index: 2,
3554 };
3555
3556 save_image_to_dir(&dir, &img, "flux-dev", 4);
3557
3558 let files: Vec<_> = std::fs::read_dir(&dir).unwrap().collect();
3559 assert_eq!(files.len(), 1);
3560 let filename = files[0]
3561 .as_ref()
3562 .unwrap()
3563 .file_name()
3564 .to_str()
3565 .unwrap()
3566 .to_string();
3567 assert!(
3568 filename.contains("-2.jpeg"),
3569 "batch index in name: {filename}"
3570 );
3571
3572 std::fs::remove_dir_all(&dir).ok();
3573 }
3574
3575 #[test]
3576 fn save_image_to_dir_invalid_path_logs_warning_no_panic() {
3577 let img = mold_core::ImageData {
3579 data: vec![0x00],
3580 format: mold_core::OutputFormat::Png,
3581 width: 1,
3582 height: 1,
3583 index: 0,
3584 };
3585 save_image_to_dir(
3587 std::path::Path::new("/dev/null/impossible"),
3588 &img,
3589 "test",
3590 1,
3591 );
3592 }
3594
3595 #[test]
3596 fn thumbnail_warmup_is_disabled_by_default() {
3597 let _guard = env_lock().lock().unwrap();
3598 unsafe {
3599 std::env::remove_var("MOLD_THUMBNAIL_WARMUP");
3600 }
3601 assert!(!thumbnail_warmup_enabled());
3602 }
3603
3604 #[test]
3605 fn thumbnail_warmup_accepts_truthy_env_values() {
3606 let _guard = env_lock().lock().unwrap();
3607 unsafe {
3608 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "1");
3609 }
3610 assert!(thumbnail_warmup_enabled());
3611 unsafe {
3612 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "true");
3613 }
3614 assert!(thumbnail_warmup_enabled());
3615 unsafe {
3616 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "YES");
3617 }
3618 assert!(thumbnail_warmup_enabled());
3619 unsafe {
3620 std::env::remove_var("MOLD_THUMBNAIL_WARMUP");
3621 }
3622 }
3623
3624 #[test]
3625 fn thumbnail_warmup_rejects_falsey_env_values() {
3626 let _guard = env_lock().lock().unwrap();
3627 unsafe {
3628 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "0");
3629 }
3630 assert!(!thumbnail_warmup_enabled());
3631 unsafe {
3632 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "false");
3633 }
3634 assert!(!thumbnail_warmup_enabled());
3635 unsafe {
3636 std::env::remove_var("MOLD_THUMBNAIL_WARMUP");
3637 }
3638 }
3639
3640 #[test]
3641 fn content_type_covers_every_output_format() {
3642 assert_eq!(content_type_for_filename("a.png"), "image/png");
3643 assert_eq!(content_type_for_filename("a.PNG"), "image/png");
3644 assert_eq!(content_type_for_filename("a.jpg"), "image/jpeg");
3645 assert_eq!(content_type_for_filename("a.jpeg"), "image/jpeg");
3646 assert_eq!(content_type_for_filename("a.gif"), "image/gif");
3647 assert_eq!(content_type_for_filename("a.webp"), "image/webp");
3648 assert_eq!(content_type_for_filename("a.apng"), "image/apng");
3649 assert_eq!(content_type_for_filename("a.mp4"), "video/mp4");
3650 assert_eq!(
3651 content_type_for_filename("a.unknown"),
3652 "application/octet-stream"
3653 );
3654 }
3655 struct TempDir(std::path::PathBuf);
3664 impl TempDir {
3665 fn new(tag: &str) -> Self {
3666 let mut p = std::env::temp_dir();
3667 p.push(format!("mold-gallery-test-{tag}-{}", uuid::Uuid::new_v4()));
3668 std::fs::create_dir_all(&p).expect("create tempdir");
3669 Self(p)
3670 }
3671 fn path(&self) -> &std::path::Path {
3672 &self.0
3673 }
3674 }
3675 impl Drop for TempDir {
3676 fn drop(&mut self) {
3677 let _ = std::fs::remove_dir_all(&self.0);
3678 }
3679 }
3680
3681 fn make_png_bytes(width: u32, height: u32) -> Vec<u8> {
3686 let img = image::RgbImage::from_fn(width, height, |x, y| {
3687 let n = (x.wrapping_mul(37) ^ y.wrapping_mul(131)) as u8;
3688 image::Rgb([n, n.wrapping_add(85), n.wrapping_sub(17)])
3689 });
3690 let mut buf = Vec::new();
3691 image::DynamicImage::ImageRgb8(img)
3692 .write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
3693 .expect("encode png");
3694 buf
3695 }
3696
3697 #[test]
3698 fn scan_gallery_dir_filters_invalid_and_keeps_valid() {
3699 let td = TempDir::new("scan");
3700 let dir = td.path();
3701
3702 std::fs::write(dir.join("mold-model-1000.png"), make_png_bytes(32, 32)).unwrap();
3704
3705 let mut junk = vec![0u8; 512];
3707 junk[..4].copy_from_slice(b"JUNK");
3708 std::fs::write(dir.join("mold-broken-2000.png"), &junk).unwrap();
3709
3710 std::fs::write(
3712 dir.join("mold-tiny-3000.png"),
3713 b"\x89PNG\r\n\x1a\n", )
3715 .unwrap();
3716
3717 let mut mp4 = Vec::new();
3719 mp4.extend_from_slice(&[0x00, 0x00, 0x00, 0x20]);
3720 mp4.extend_from_slice(b"ftyp");
3721 mp4.extend_from_slice(b"isom\x00\x00\x02\x00");
3722 mp4.resize(8192, 0);
3725 std::fs::write(dir.join("mold-ltx-4000.mp4"), &mp4).unwrap();
3726
3727 let bad_mp4 = vec![0u8; 8192];
3729 std::fs::write(dir.join("mold-no-ftyp-5000.mp4"), &bad_mp4).unwrap();
3730
3731 std::fs::write(dir.join("random.txt"), b"not an output").unwrap();
3733
3734 let results = scan_gallery_dir(dir);
3735 let names: Vec<&str> = results.iter().map(|i| i.filename.as_str()).collect();
3736 assert!(
3737 names.contains(&"mold-model-1000.png"),
3738 "valid PNG should survive: {names:?}"
3739 );
3740 assert!(
3741 names.contains(&"mold-ltx-4000.mp4"),
3742 "valid MP4 with ftyp should survive: {names:?}"
3743 );
3744 assert!(
3745 !names.contains(&"mold-broken-2000.png"),
3746 "PNG with no valid header should be filtered: {names:?}"
3747 );
3748 assert!(
3749 !names.contains(&"mold-tiny-3000.png"),
3750 "under-size PNG stub should be filtered: {names:?}"
3751 );
3752 assert!(
3753 !names.contains(&"mold-no-ftyp-5000.mp4"),
3754 "MP4 without ftyp should be filtered: {names:?}"
3755 );
3756 assert_eq!(names.len(), 2, "only the 2 valid fixtures remain");
3757 }
3758
3759 #[test]
3760 fn solid_black_png_is_filtered_at_scan_time() {
3761 let td = TempDir::new("black");
3762 let dir = td.path();
3763
3764 let black = image::RgbImage::from_pixel(256, 256, image::Rgb([0, 0, 0]));
3768 let mut buf = Vec::new();
3769 image::DynamicImage::ImageRgb8(black)
3770 .write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
3771 .unwrap();
3772 std::fs::write(dir.join("mold-noisy-1000.png"), &buf).unwrap();
3773
3774 std::fs::write(dir.join("mold-valid-2000.png"), make_png_bytes(256, 256)).unwrap();
3776
3777 let results = scan_gallery_dir(dir);
3778 let names: Vec<&str> = results.iter().map(|i| i.filename.as_str()).collect();
3779 assert!(
3780 !names.contains(&"mold-noisy-1000.png"),
3781 "solid-black PNG should be filtered: {names:?}"
3782 );
3783 assert!(
3784 names.contains(&"mold-valid-2000.png"),
3785 "noisy PNG should survive: {names:?}"
3786 );
3787 }
3788 #[test]
3789 fn parse_byte_range_handles_common_forms() {
3790 assert_eq!(parse_byte_range("bytes=0-499", 2000), Some((0, 499)));
3792 assert_eq!(parse_byte_range("bytes=100-", 2000), Some((100, 1999)));
3794 assert_eq!(parse_byte_range("bytes=-500", 2000), Some((1500, 1999)));
3796 assert_eq!(parse_byte_range("bytes=0-9999", 2000), Some((0, 1999)));
3798 assert_eq!(parse_byte_range("bytes=0-1999", 2000), Some((0, 1999)));
3800 }
3801
3802 #[test]
3803 fn parse_byte_range_rejects_malformed_and_unsatisfiable() {
3804 assert_eq!(parse_byte_range("bytes=", 1000), None);
3805 assert_eq!(parse_byte_range("bytes=abc-100", 1000), None);
3806 assert_eq!(parse_byte_range("bytes=2000-", 1000), None);
3808 assert_eq!(parse_byte_range("bytes=500-100", 1000), None);
3810 assert_eq!(parse_byte_range("bytes=0-10,20-30", 1000), None);
3812 assert_eq!(parse_byte_range("bytes=-0", 1000), None);
3814 assert_eq!(parse_byte_range("bytes=0-10", 0), None);
3816 assert_eq!(parse_byte_range("items=0-10", 1000), None);
3818 }
3819
3820 #[test]
3821 fn scan_populates_real_dimensions_for_synthesized_metadata() {
3822 let td = TempDir::new("dims");
3826 let dir = td.path();
3827 std::fs::write(dir.join("mold-nometa-1000.png"), make_png_bytes(128, 96)).unwrap();
3828
3829 let results = scan_gallery_dir(dir);
3830 assert_eq!(results.len(), 1);
3831 let entry = &results[0];
3832 assert!(entry.metadata_synthetic);
3833 assert_eq!(entry.metadata.width, 128);
3834 assert_eq!(entry.metadata.height, 96);
3835 }
3836}