1use axum::{
2 extract::{Extension, 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, DiskUsage, GenerateRequest, GpuBackend, GpuInfo,
14 GpuWorkerState, ModelInfoExtended, ResourceSnapshot, ServerStatus, SseErrorEvent,
15 SseProgressEvent,
16};
17use serde::{Deserialize, Serialize};
18use std::cmp::Reverse;
19use std::convert::Infallible;
20use std::sync::atomic::Ordering;
21use tokio_stream::StreamExt as _;
22use utoipa::OpenApi;
23
24use crate::model_manager;
25use crate::state::{AppState, GenerationJob, SseCompletionPayload, SseMessage, SubmitError};
26
27fn submit_error_to_api(e: SubmitError) -> ApiError {
28 match e {
29 SubmitError::Full { pending, capacity } => {
30 ApiError::queue_full(format!("generation queue is full ({pending}/{capacity})"))
31 }
32 SubmitError::Shutdown => ApiError::internal("generation queue shut down"),
33 }
34}
35
36#[derive(Debug, Serialize)]
39pub struct ApiError {
40 pub error: String,
41 pub code: String,
42 #[serde(skip)]
43 status: StatusCode,
44}
45
46impl ApiError {
47 pub fn validation(msg: impl Into<String>) -> Self {
48 Self {
49 error: msg.into(),
50 code: "VALIDATION_ERROR".to_string(),
51 status: StatusCode::UNPROCESSABLE_ENTITY,
52 }
53 }
54
55 pub fn not_found(msg: impl Into<String>) -> Self {
56 Self {
57 error: msg.into(),
58 code: "MODEL_NOT_FOUND".to_string(),
59 status: StatusCode::NOT_FOUND,
60 }
61 }
62
63 pub fn unknown_model(msg: impl Into<String>) -> Self {
64 Self {
65 error: msg.into(),
66 code: "UNKNOWN_MODEL".to_string(),
67 status: StatusCode::BAD_REQUEST,
68 }
69 }
70
71 pub fn inference(msg: impl Into<String>) -> Self {
72 Self {
73 error: msg.into(),
74 code: "INFERENCE_ERROR".to_string(),
75 status: StatusCode::INTERNAL_SERVER_ERROR,
76 }
77 }
78
79 pub fn internal(msg: impl Into<String>) -> Self {
80 Self {
81 error: msg.into(),
82 code: "INTERNAL_ERROR".to_string(),
83 status: StatusCode::INTERNAL_SERVER_ERROR,
84 }
85 }
86
87 pub fn internal_with_status(msg: impl Into<String>, status: StatusCode) -> Self {
88 Self {
89 error: msg.into(),
90 code: "INTERNAL_ERROR".to_string(),
91 status,
92 }
93 }
94
95 pub fn with_code(msg: impl Into<String>, code: impl Into<String>, status: StatusCode) -> Self {
96 Self {
97 error: msg.into(),
98 code: code.into(),
99 status,
100 }
101 }
102
103 pub fn queue_job_not_found(msg: impl Into<String>) -> Self {
104 Self {
105 error: msg.into(),
106 code: "QUEUE_JOB_NOT_FOUND".to_string(),
107 status: StatusCode::NOT_FOUND,
108 }
109 }
110
111 pub fn queue_job_running(msg: impl Into<String>) -> Self {
112 Self {
113 error: msg.into(),
114 code: "QUEUE_JOB_RUNNING".to_string(),
115 status: StatusCode::CONFLICT,
116 }
117 }
118
119 pub fn cancelled(msg: impl Into<String>) -> Self {
124 Self {
125 error: msg.into(),
126 code: "CANCELLED".to_string(),
127 status: StatusCode::from_u16(499).expect("499 is a valid status code"),
128 }
129 }
130
131 pub fn insufficient_memory(msg: impl Into<String>) -> Self {
132 Self {
133 error: msg.into(),
134 code: "INSUFFICIENT_MEMORY".to_string(),
135 status: StatusCode::SERVICE_UNAVAILABLE,
136 }
137 }
138
139 pub fn forbidden(msg: impl Into<String>) -> Self {
140 Self {
141 error: msg.into(),
142 code: "FORBIDDEN".to_string(),
143 status: StatusCode::FORBIDDEN,
144 }
145 }
146
147 pub fn queue_full(msg: impl Into<String>) -> Self {
148 Self {
149 error: msg.into(),
150 code: "QUEUE_FULL".to_string(),
151 status: StatusCode::SERVICE_UNAVAILABLE,
152 }
153 }
154}
155
156impl IntoResponse for ApiError {
157 fn into_response(self) -> axum::response::Response {
158 let status = self.status;
159 if self.code == "QUEUE_FULL" {
161 let mut headers = HeaderMap::new();
162 headers.insert(header::RETRY_AFTER, HeaderValue::from_static("1"));
163 return (status, headers, Json(self)).into_response();
164 }
165 (status, Json(self)).into_response()
166 }
167}
168
169#[cfg(test)]
171use crate::queue::clean_error_message;
172
173#[derive(OpenApi)]
174#[openapi(
175 paths(
176 generate,
177 generate_stream,
178 expand_prompt,
179 list_models,
180 crate::catalog_api::list_loras,
181 load_model,
182 pull_model_endpoint,
183 unload_model,
184 delete_model,
185 create_gallery_media_token,
186 server_status,
187 list_queue,
188 patch_queue_job,
189 cancel_queue_job,
190 pause_queue,
191 resume_queue,
192 cancel_all_queue,
193 list_history,
194 delete_history,
195 crate::routes_config::list_config,
196 crate::routes_config::get_config_key,
197 crate::routes_config::put_config_key,
198 crate::routes_config::delete_config_key,
199 crate::routes_config::list_config_profiles,
200 crate::routes_config::put_config_profile,
201 health,
202 capabilities_chain_limits,
203 stream_events,
204 crate::routes_chain::generate_chain,
205 crate::routes_chain::generate_chain_stream,
206 crate::routes_chain_jobs::create_chain_job,
207 crate::routes_chain_jobs::list_chain_jobs,
208 crate::routes_chain_jobs::get_chain_job,
209 crate::routes_chain_jobs::chain_job_events,
210 crate::routes_chain_jobs::resume_chain_job,
211 crate::routes_chain_jobs::retake_chain_job,
212 crate::routes_chain_jobs::cancel_chain_job,
213 crate::routes_chain_jobs::delete_chain_job,
214 crate::routes_chain_jobs::gc_chain_jobs,
215 crate::routes_chain_jobs::chain_job_stage_preview,
216 ),
217 components(schemas(
218 mold_core::GenerateRequest,
219 mold_core::GenerateResponse,
220 mold_core::ExpandRequest,
221 mold_core::ExpandResponse,
222 mold_core::ImageData,
223 mold_core::OutputFormat,
224 mold_core::ModelInfo,
225 mold_core::LoraInfo,
226 mold_core::ServerStatus,
227 mold_core::ActiveGenerationStatus,
228 mold_core::GpuInfo,
229 mold_core::DiskUsage,
230 mold_core::SseProgressEvent,
231 mold_core::SseCompleteEvent,
232 mold_core::SseErrorEvent,
233 mold_core::ChainRequest,
234 mold_core::ChainResponse,
235 mold_core::ChainStage,
236 mold_core::ChainProgressEvent,
237 mold_core::SseChainCompleteEvent,
238 mold_core::chain_job::ChainJobSummary,
239 mold_core::chain_job::ChainJobStageDetail,
240 mold_core::chain_job::ChainJobDetail,
241 mold_core::chain_job::ChainJobListing,
242 mold_core::chain_job::CreateChainJobResponse,
243 mold_core::chain_job::RetakeRequest,
244 mold_core::chain_job::ChainJobEvent,
245 mold_core::chain_job::FinalizeRecord,
246 mold_core::chain_job::RetakeAmendment,
247 mold_core::chain_job::ChainJobState,
248 mold_core::chain_job::StageState,
249 mold_core::chain_job::RetakeMode,
250 mold_core::chain_job::GcOutcome,
251 ModelInfoExtended,
252 LoadModelBody,
253 UnloadRequest,
254 mold_core::ModelRemovalResponse,
255 mold_core::KeptComponent,
256 QueuePatchRequest,
257 QueuePauseResponse,
258 QueueCancelAllResponse,
259 mold_core::HistoryEntry,
260 mold_core::HistoryListing,
261 mold_core::ConfigEntry,
262 mold_core::ConfigListing,
263 mold_core::ConfigProfiles,
264 crate::routes_config::ConfigSetRequest,
265 crate::routes_config::ProfileSetRequest,
266 crate::job_registry::JobEntry,
267 crate::job_registry::QueueListing,
268 crate::chain_limits::ChainLimits,
269 GalleryMediaTokenRequest,
270 GalleryMediaTokenResponse,
271 )),
272 tags(
273 (name = "generation", description = "Image generation"),
274 (name = "models", description = "Model management"),
275 (name = "server", description = "Server status and health"),
276 (name = "chain-jobs", description = "Durable chained video jobs"),
277 ),
278 info(
279 title = "mold",
280 description = "Local AI image generation server — FLUX, SD3.5, SD1.5, SDXL, Z-Image, Flux.2, Qwen-Image",
281 version = env!("CARGO_PKG_VERSION"),
282 )
283)]
284pub struct ApiDoc;
285
286pub fn create_router(state: AppState) -> Router {
287 Router::new()
290 .route("/api/generate", post(generate))
291 .route("/api/generate/estimate", post(generate_estimate))
292 .route("/api/generate/stream", post(generate_stream))
293 .route(
294 "/api/generate/chain",
295 post(crate::routes_chain::generate_chain),
296 )
297 .route(
298 "/api/generate/chain/stream",
299 post(crate::routes_chain::generate_chain_stream),
300 )
301 .route(
302 "/api/chain-jobs",
303 post(crate::routes_chain_jobs::create_chain_job)
304 .get(crate::routes_chain_jobs::list_chain_jobs),
305 )
306 .route(
307 "/api/chain-jobs/:id",
308 get(crate::routes_chain_jobs::get_chain_job)
309 .delete(crate::routes_chain_jobs::delete_chain_job),
310 )
311 .route(
312 "/api/chain-jobs/:id/events",
313 get(crate::routes_chain_jobs::chain_job_events),
314 )
315 .route(
316 "/api/chain-jobs/:id/resume",
317 post(crate::routes_chain_jobs::resume_chain_job),
318 )
319 .route(
320 "/api/chain-jobs/:id/retake",
321 post(crate::routes_chain_jobs::retake_chain_job),
322 )
323 .route(
324 "/api/chain-jobs/:id/cancel",
325 post(crate::routes_chain_jobs::cancel_chain_job),
326 )
327 .route(
328 "/api/chain-jobs/gc",
329 post(crate::routes_chain_jobs::gc_chain_jobs),
330 )
331 .route(
332 "/api/chain-jobs/:id/stages/:idx/preview",
333 get(crate::routes_chain_jobs::chain_job_stage_preview),
334 )
335 .route("/api/expand", post(expand_prompt))
336 .route("/api/models", get(list_models))
337 .route("/api/models/:model", delete(delete_model))
338 .route("/api/models/:model/components", get(model_components))
339 .route("/api/loras", get(crate::catalog_api::list_loras))
340 .route("/api/models/load", post(load_model))
341 .route("/api/models/pull", post(pull_model_endpoint))
342 .route("/api/models/unload", delete(unload_model))
343 .route("/api/gallery", get(list_gallery))
344 .route("/api/gallery/media-token", post(create_gallery_media_token))
345 .route(
346 "/api/gallery/image/:filename",
347 get(get_gallery_image).delete(delete_gallery_image),
348 )
349 .route(
350 "/api/gallery/thumbnail/:filename",
351 get(get_gallery_thumbnail),
352 )
353 .route("/api/gallery/preview/:filename", get(get_gallery_preview))
354 .route("/api/downloads", get(list_downloads).post(create_download))
356 .route("/api/downloads/:id", delete(delete_download))
357 .route("/api/downloads/stream", get(stream_downloads))
358 .route(
360 "/api/catalog/families",
361 get(crate::catalog_api::list_families),
362 )
363 .route(
364 "/api/catalog/search",
365 get(crate::catalog_api::live_search_catalog),
366 )
367 .route(
368 "/api/catalog/installed",
369 get(crate::catalog_api::list_installed_catalog),
370 )
371 .route(
372 "/api/catalog/*id",
373 get(crate::catalog_api::get_catalog_entry)
374 .post(crate::catalog_api::post_catalog_dispatch),
375 )
376 .route("/api/upscale", post(upscale))
377 .route("/api/upscale/stream", post(upscale_stream))
378 .route("/api/resources", get(get_resources))
379 .route("/api/resources/stream", get(get_resources_stream))
380 .route("/api/events", get(stream_events))
381 .route("/api/status", get(server_status))
382 .route("/api/queue", get(list_queue).delete(cancel_all_queue))
383 .route("/api/queue/pause", post(pause_queue))
384 .route("/api/queue/resume", post(resume_queue))
385 .route(
386 "/api/queue/:id",
387 patch(patch_queue_job).delete(cancel_queue_job),
388 )
389 .route("/api/history", get(list_history).delete(delete_history))
390 .route("/api/capabilities", get(server_capabilities))
391 .route(
392 "/api/capabilities/chain-limits",
393 get(capabilities_chain_limits),
394 )
395 .route("/api/shutdown", post(shutdown_server))
396 .route("/api/config", get(crate::routes_config::list_config))
398 .route(
399 "/api/config/profiles",
400 get(crate::routes_config::list_config_profiles),
401 )
402 .route(
403 "/api/config/profile",
404 axum::routing::put(crate::routes_config::put_config_profile),
405 )
406 .route(
407 "/api/config/:key",
408 get(crate::routes_config::get_config_key)
409 .put(crate::routes_config::put_config_key)
410 .delete(crate::routes_config::delete_config_key),
411 )
412 .route(
414 "/api/config/model/:name/placement",
415 get(get_model_placement)
416 .put(put_model_placement)
417 .delete(delete_model_placement),
418 )
419 .route("/health", get(health))
420 .with_state(state)
421 .route("/api/openapi.json", get(openapi_json))
422 .route("/api/docs", get(scalar_docs))
423}
424
425fn sse_message_to_event(msg: SseMessage) -> SseEvent {
428 fn serialize_event<T: Serialize>(event_name: &str, payload: &T) -> SseEvent {
429 match serde_json::to_string(payload) {
430 Ok(data) => SseEvent::default().event(event_name).data(data),
431 Err(err) => SseEvent::default().event("error").data(
432 serde_json::json!({
433 "message": format!("failed to serialize SSE payload: {err}")
434 })
435 .to_string(),
436 ),
437 }
438 }
439
440 match msg {
441 SseMessage::Progress(payload) => serialize_event("progress", &payload),
442 SseMessage::Complete(payload) => serialize_event("complete", &payload),
443 SseMessage::UpscaleComplete(payload) => serialize_event("complete", &payload),
444 SseMessage::Error(payload) => serialize_event("error", &payload),
445 }
446}
447
448#[cfg(test)]
449fn save_image_to_dir(
450 dir: &std::path::Path,
451 img: &mold_core::ImageData,
452 model: &str,
453 batch_size: u32,
454) {
455 if let Err(e) = std::fs::create_dir_all(dir) {
456 tracing::warn!("failed to create output dir {}: {e}", dir.display());
457 return;
458 }
459 let timestamp_ms = std::time::SystemTime::now()
462 .duration_since(std::time::UNIX_EPOCH)
463 .unwrap_or_default()
464 .as_millis() as u64;
465 let ext = img.format.to_string();
466 let filename =
467 mold_core::default_output_filename(model, timestamp_ms, &ext, batch_size, img.index);
468 let path = dir.join(&filename);
469 match std::fs::write(&path, &img.data) {
470 Ok(()) => tracing::info!("saved image to {}", path.display()),
471 Err(e) => tracing::warn!("failed to save image to {}: {e}", path.display()),
472 }
473}
474
475async fn prepare_generation(
483 state: &AppState,
484 request: &mut mold_core::GenerateRequest,
485) -> Result<(Option<std::path::PathBuf>, Option<String>, Option<usize>), ApiError> {
486 apply_default_metadata_setting(state, request).await;
491
492 let preferred_gpu = validate_multi_gpu_placement(state, request.placement.as_ref())?;
493
494 maybe_expand_prompt(state, request, preferred_gpu).await?;
496
497 if let Err(e) = model_manager::install_catalog_model(state, &request.model).await {
507 return Err(model_manager::install_error_to_api_error(&e));
508 }
509 let family_hint = model_manager::catalog_family_for(state, &request.model).await;
510
511 let resolved_family = model_manager::family_for_model(state, &request.model).await;
515 request.normalise_output_format(resolved_family.as_deref());
519
520 if let Err(e) = validate_generate_request(request, family_hint.as_deref()) {
521 return Err(ApiError::validation(e));
522 }
523
524 resolve_server_local_media_paths(state, request).await?;
525
526 let _ = model_manager::check_model_available(state, &request.model).await?;
527
528 let (output_dir, dim_warning) = {
529 let config = state.config.read().await;
530 let output_dir = if config.is_output_disabled() {
531 None
532 } else {
533 Some(config.effective_output_dir())
534 };
535 let family = config.resolved_model_config(&request.model).family;
536 let dim_warning = family
537 .as_deref()
538 .and_then(|f| mold_core::dimension_warning(request.width, request.height, f));
539 (output_dir, dim_warning)
540 };
541
542 Ok((output_dir, dim_warning, preferred_gpu))
543}
544
545fn record_prompt_history(state: &AppState, prompt: &str, negative: Option<&str>, model: &str) {
551 let Some(db) = state.metadata_db.as_ref().as_ref() else {
552 return;
553 };
554 let history = mold_db::PromptHistory::new(db);
555 if let Ok(rows) = history.recent(1) {
556 if rows.first().is_some_and(|latest| {
557 latest.prompt == prompt
558 && latest.model == model
559 && latest.negative.as_deref() == negative
560 }) {
561 return;
562 }
563 }
564 if let Err(e) = history.push(&mold_db::HistoryEntry {
565 prompt: prompt.to_string(),
566 negative: negative.map(str::to_string),
567 model: model.to_string(),
568 created_at_ms: 0, }) {
570 tracing::warn!("failed to record prompt history: {e:#}");
571 }
572}
573
574pub(crate) async fn resolve_server_local_media_paths(
575 state: &AppState,
576 request: &mut mold_core::GenerateRequest,
577) -> Result<(), ApiError> {
578 if request.audio_file_path.is_none() && request.source_video_path.is_none() {
579 return Ok(());
580 }
581
582 let roots = state.config.read().await.resolved_media_roots();
583 if let Some(path) = request.audio_file_path.as_deref() {
584 let resolved = mold_core::resolve_server_media_path(path, &roots)
585 .map_err(|e| ApiError::validation(format!("audio_file_path: {e}")))?;
586 request.audio_file_path = Some(resolved.to_string_lossy().to_string());
587 }
588 if let Some(path) = request.source_video_path.as_deref() {
589 let resolved = mold_core::resolve_server_media_path(path, &roots)
590 .map_err(|e| ApiError::validation(format!("source_video_path: {e}")))?;
591 request.source_video_path = Some(resolved.to_string_lossy().to_string());
592 }
593
594 Ok(())
595}
596
597fn active_gpu_selection(state: &AppState) -> GpuSelection {
598 let ordinals: Vec<usize> = state
599 .gpu_pool
600 .workers
601 .iter()
602 .map(|w| w.gpu.ordinal)
603 .collect();
604 if ordinals.is_empty() {
605 GpuSelection::All
606 } else {
607 GpuSelection::Specific(ordinals)
608 }
609}
610
611fn validate_multi_gpu_placement(
612 state: &AppState,
613 placement: Option<&mold_core::types::DevicePlacement>,
614) -> Result<Option<usize>, ApiError> {
615 state
616 .gpu_pool
617 .resolve_explicit_placement_gpu(placement)
618 .map_err(ApiError::validation)
619}
620
621fn select_aux_worker(
622 state: &AppState,
623) -> Result<std::sync::Arc<crate::gpu_pool::GpuWorker>, ApiError> {
624 let mut workers: Vec<_> = state
625 .gpu_pool
626 .workers
627 .iter()
628 .filter(|w| !w.is_degraded())
629 .cloned()
630 .collect();
631 workers.sort_by_key(|w| {
632 (
633 w.in_flight.load(Ordering::SeqCst),
634 Reverse(w.gpu.total_vram_bytes),
635 )
636 });
637 workers
638 .into_iter()
639 .next()
640 .ok_or_else(|| ApiError::internal("no GPU worker available for auxiliary workload"))
641}
642
643fn clear_global_upscaler_cache(state: &AppState) {
644 if let Ok(mut cache) = state.upscaler_cache.try_lock() {
645 if cache.is_some() {
646 *cache = None;
647 tracing::info!("upscaler cache cleared");
648 }
649 }
650}
651
652#[utoipa::path(
655 post,
656 path = "/api/generate",
657 tag = "generation",
658 request_body = mold_core::GenerateRequest,
659 responses(
660 (status = 200, description = "Generated image bytes", content_type = "image/png"),
661 (status = 404, description = "Model not downloaded"),
662 (status = 422, description = "Invalid request parameters"),
663 (status = 500, description = "Inference error"),
664 (status = 503, description = "Generation queue full"),
665 )
666)]
667async fn generate(
670 State(state): State<AppState>,
671 Json(mut req): Json<mold_core::GenerateRequest>,
672) -> Result<impl IntoResponse, ApiError> {
673 let typed = (
676 req.prompt.clone(),
677 req.negative_prompt.clone(),
678 req.model.clone(),
679 );
680 let (output_dir, dim_warning, preferred_gpu) = prepare_generation(&state, &mut req).await?;
681 record_prompt_history(&state, &typed.0, typed.1.as_deref(), &typed.2);
682
683 tracing::info!(
684 model = %req.model,
685 prompt = %req.prompt,
686 width = req.width,
687 height = req.height,
688 steps = req.steps,
689 guidance = req.guidance,
690 seed = ?req.seed,
691 format = %req.resolved_output_format(),
692 lora = ?req.lora.as_ref().map(|l| &l.path),
693 lora_scale = ?req.lora.as_ref().map(|l| l.scale),
694 loras = ?req.loras.as_ref().map(|v| v.iter().map(|l| &l.path).collect::<Vec<_>>()),
695 "generate request"
696 );
697
698 let (result_tx, result_rx) = tokio::sync::oneshot::channel();
700 let job_id = uuid::Uuid::new_v4().to_string();
705 let queue_metadata = Box::new(mold_core::OutputMetadata::from_generate_request(
710 &req,
711 req.seed.unwrap_or(0),
712 req.scheduler,
713 mold_core::build_info::version_string(),
714 ));
715 let cancel = state.job_registry.register_job(
716 &job_id,
717 &req.model,
718 preferred_gpu,
719 Some(req.seed.is_some()),
720 Some(queue_metadata),
721 );
722 let job = GenerationJob {
723 id: job_id.clone(),
724 request: req,
725 completion_payload: SseCompletionPayload::Full,
726 progress_tx: None,
727 result_tx,
728 output_dir,
729 };
730
731 let _position = state
732 .queue
733 .submit(job, state.queue_capacity)
734 .await
735 .inspect_err(|_| state.job_registry.remove(&job_id))
736 .map_err(submit_error_to_api)?;
737
738 let result = tokio::select! {
743 result = result_rx => result
744 .map_err(|_| ApiError::internal("generation queue worker dropped the job"))?,
745 _ = cancel.notified() => {
746 return Err(ApiError::cancelled(format!(
747 "generation job {job_id} was cancelled while queued"
748 )));
749 }
750 };
751
752 match result {
753 Ok(job_result) => {
754 let img = job_result.image;
755 let response = job_result.response;
756 let content_type = HeaderValue::from_static(img.format.content_type());
757 let mut headers = HeaderMap::new();
758 headers.insert(header::CONTENT_TYPE, content_type);
759 headers.insert(
760 "x-mold-seed-used",
761 HeaderValue::from_str(&response.seed_used.to_string()).map_err(|e| {
762 ApiError::internal(format!("failed to serialize seed header: {e}"))
763 })?,
764 );
765 if let Some(ordinal) = response.gpu {
766 headers.insert(
767 "x-mold-gpu",
768 HeaderValue::from_str(&ordinal.to_string()).map_err(|e| {
769 ApiError::internal(format!("failed to serialize gpu header: {e}"))
770 })?,
771 );
772 }
773 if let Some(warning) = dim_warning {
774 match HeaderValue::from_str(&warning.replace('\n', " ")) {
775 Ok(val) => {
776 headers.insert("x-mold-dimension-warning", val);
777 }
778 Err(e) => {
779 tracing::warn!("dimension warning could not be encoded as header: {e}");
780 }
781 }
782 }
783 let output_data = if let Some(ref video) = response.video {
786 let ct = HeaderValue::from_static(video.format.content_type());
787 headers.insert(header::CONTENT_TYPE, ct);
788 if let Ok(v) = HeaderValue::from_str(&video.frames.to_string()) {
789 headers.insert("x-mold-video-frames", v);
790 }
791 if let Ok(v) = HeaderValue::from_str(&video.fps.to_string()) {
792 headers.insert("x-mold-video-fps", v);
793 }
794 if let Ok(v) = HeaderValue::from_str(&video.width.to_string()) {
795 headers.insert("x-mold-video-width", v);
796 }
797 if let Ok(v) = HeaderValue::from_str(&video.height.to_string()) {
798 headers.insert("x-mold-video-height", v);
799 }
800 if video.has_audio {
801 headers.insert("x-mold-video-has-audio", HeaderValue::from_static("1"));
802 }
803 if let Some(dur) = video.duration_ms {
804 if let Ok(v) = HeaderValue::from_str(&dur.to_string()) {
805 headers.insert("x-mold-video-duration-ms", v);
806 }
807 }
808 if let Some(sr) = video.audio_sample_rate {
809 if let Ok(v) = HeaderValue::from_str(&sr.to_string()) {
810 headers.insert("x-mold-video-audio-sample-rate", v);
811 }
812 }
813 if let Some(ch) = video.audio_channels {
814 if let Ok(v) = HeaderValue::from_str(&ch.to_string()) {
815 headers.insert("x-mold-video-audio-channels", v);
816 }
817 }
818 video.data.clone()
819 } else {
820 img.data
821 };
822 Ok((headers, output_data))
823 }
824 Err(err_msg) => {
825 if err_msg.contains("queue is full") {
829 Err(ApiError::queue_full(err_msg))
830 } else {
831 Err(ApiError::inference(err_msg))
832 }
833 }
834 }
835}
836
837fn validate_generate_request(
838 req: &mold_core::GenerateRequest,
839 family_hint: Option<&str>,
840) -> Result<(), String> {
841 mold_core::validate_generate_request_with_family(req, family_hint)
842}
843
844async fn apply_default_metadata_setting(state: &AppState, req: &mut mold_core::GenerateRequest) {
845 if req.embed_metadata.is_some() {
846 return;
847 }
848
849 let config = state.config.read().await;
850 req.embed_metadata = Some(config.effective_embed_metadata(None));
851}
852
853async fn maybe_expand_prompt(
855 state: &AppState,
856 req: &mut mold_core::GenerateRequest,
857 preferred_gpu: Option<usize>,
858) -> Result<(), ApiError> {
859 if req.expand != Some(true) {
860 return Ok(());
861 }
862
863 let config = state.config.read().await;
864 let config_snapshot = config.clone();
865 let expand_settings = config.expand.clone().with_env_overrides();
866
867 let model_family = config
869 .resolved_model_config(&req.model)
870 .family
871 .or_else(|| mold_core::manifest::find_manifest(&req.model).map(|m| m.family.clone()))
872 .unwrap_or_else(|| {
873 tracing::warn!(
874 model = %req.model,
875 "could not resolve model family for prompt expansion, defaulting to \"flux\""
876 );
877 "flux".to_string()
878 });
879
880 let expand_config = expand_settings.to_expand_config(&model_family, 1);
881 let original_prompt = req.prompt.clone();
882
883 drop(config);
885
886 let expander = create_server_expander(
887 &config_snapshot,
888 &expand_settings,
889 active_gpu_selection(state),
890 preferred_gpu,
891 )?;
892 let result =
893 tokio::task::spawn_blocking(move || expander.expand(&original_prompt, &expand_config))
894 .await
895 .map_err(|e| ApiError::internal(format!("expand task failed: {e}")))?
896 .map_err(|e| ApiError::internal(format!("prompt expansion failed: {e}")))?;
897
898 if let Some(expanded) = result.expanded.first() {
899 req.original_prompt = Some(req.prompt.clone());
900 req.prompt = expanded.clone();
901 }
902
903 Ok(())
904}
905
906fn create_server_expander(
908 _config: &mold_core::Config,
909 settings: &mold_core::ExpandSettings,
910 _gpu_selection: GpuSelection,
911 _preferred_gpu: Option<usize>,
912) -> Result<Box<dyn mold_core::PromptExpander>, ApiError> {
913 if let Some(api_expander) = settings.create_api_expander() {
914 return Ok(Box::new(api_expander));
915 }
916
917 #[cfg(feature = "expand")]
918 {
919 if let Some(local) =
920 mold_inference::expand::LocalExpander::from_config(_config, Some(&settings.model))
921 {
922 return Ok(Box::new(
923 local
924 .with_gpu_selection(_gpu_selection)
925 .with_preferred_gpu(_preferred_gpu),
926 ));
927 }
928 return Err(ApiError::validation(
929 "local expand model not found — run: mold pull qwen3-expand".to_string(),
930 ));
931 }
932
933 #[cfg(not(feature = "expand"))]
934 {
935 Err(ApiError::validation(
936 "local prompt expansion not available — built without expand feature. \
937 Configure an API backend in [expand] settings."
938 .to_string(),
939 ))
940 }
941}
942
943fn expand_config_for_request(
947 settings: &mold_core::ExpandSettings,
948 req: &mold_core::ExpandRequest,
949) -> mold_core::ExpandConfig {
950 let mut config = settings.to_expand_config(&req.model_family, req.variations);
951 config.style = req.style.clone();
952 config
953}
954
955#[utoipa::path(
958 post,
959 path = "/api/expand",
960 tag = "generation",
961 request_body = mold_core::ExpandRequest,
962 responses(
963 (status = 200, description = "Expanded prompt(s)", body = mold_core::ExpandResponse),
964 (status = 422, description = "Invalid request parameters"),
965 (status = 500, description = "Expansion failed"),
966 )
967)]
968async fn expand_prompt(
969 State(state): State<AppState>,
970 Json(req): Json<mold_core::ExpandRequest>,
971) -> Result<Json<mold_core::ExpandResponse>, ApiError> {
972 if req.variations == 0 || req.variations > mold_core::expand::MAX_VARIATIONS {
973 return Err(ApiError::validation(format!(
974 "variations must be between 1 and {}",
975 mold_core::expand::MAX_VARIATIONS,
976 )));
977 }
978
979 let config = state.config.read().await;
980 let expand_settings = config.expand.clone().with_env_overrides();
981 let expand_config = expand_config_for_request(&expand_settings, &req);
982 let prompt = req.prompt.clone();
983 let config_snapshot = config.clone();
984 drop(config);
985
986 let expander = create_server_expander(
987 &config_snapshot,
988 &expand_settings,
989 active_gpu_selection(&state),
990 None,
991 )?;
992 let result = tokio::task::spawn_blocking(move || expander.expand(&prompt, &expand_config))
993 .await
994 .map_err(|e| ApiError::internal(format!("expand task failed: {e}")))?
995 .map_err(|e| ApiError::internal(format!("prompt expansion failed: {e}")))?;
996
997 Ok(Json(mold_core::ExpandResponse {
998 original: req.prompt,
999 expanded: result.expanded,
1000 }))
1001}
1002
1003async fn upscale(
1006 State(state): State<AppState>,
1007 Json(req): Json<mold_core::UpscaleRequest>,
1008) -> Result<Json<mold_core::UpscaleResponse>, ApiError> {
1009 if let Err(msg) = mold_core::validate_upscale_request(&req) {
1010 return Err(ApiError::validation(msg));
1011 }
1012
1013 let model_name = mold_core::manifest::resolve_model_name(&req.model);
1014
1015 let needs_pull = {
1017 let config = state.config.read().await;
1018 config
1019 .models
1020 .get(&model_name)
1021 .and_then(|c| c.transformer.as_ref())
1022 .is_none()
1023 };
1024 if needs_pull {
1025 if mold_core::manifest::find_manifest(&model_name).is_none() {
1026 return Err(ApiError::not_found(format!(
1027 "unknown upscaler model '{}'. Run 'mold list' to see available models.",
1028 model_name
1029 )));
1030 }
1031 model_manager::pull_model(&state, &model_name, None).await?;
1032 }
1033
1034 let config = state.config.read().await;
1035 let weights_path = config
1036 .models
1037 .get(&model_name)
1038 .and_then(|c| c.transformer.as_ref())
1039 .ok_or_else(|| {
1040 ApiError::not_found(format!(
1041 "upscaler model '{}' not configured after pull",
1042 model_name
1043 ))
1044 })?;
1045 let weights_path = std::path::PathBuf::from(weights_path);
1046 let model_name_owned = model_name.clone();
1047 drop(config);
1048
1049 let resp = if state.gpu_pool.worker_count() > 0 {
1050 let worker = select_aux_worker(&state)?;
1051 worker.in_flight.fetch_add(1, Ordering::SeqCst);
1052 let worker_clone = worker.clone();
1053 let joined =
1054 tokio::task::spawn_blocking(move || -> anyhow::Result<mold_core::UpscaleResponse> {
1055 struct ThreadGpuGuard;
1056 impl Drop for ThreadGpuGuard {
1057 fn drop(&mut self) {
1058 mold_inference::device::clear_thread_gpu_ordinal();
1059 }
1060 }
1061
1062 mold_inference::device::init_thread_gpu_ordinal(worker_clone.gpu.ordinal);
1063 let _thread_gpu = ThreadGpuGuard;
1064 let _load_lock = worker_clone.model_load_lock.lock().unwrap();
1065 crate::gpu_worker::ensure_worker_not_poisoned(&worker_clone, &model_name_owned)?;
1066 let mut engine = mold_inference::create_upscale_engine(
1067 model_name_owned,
1068 weights_path,
1069 mold_inference::LoadStrategy::Eager,
1070 worker_clone.gpu.ordinal,
1071 )?;
1072 engine.upscale(&req)
1073 })
1074 .await;
1075 worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1076 let result =
1077 joined.map_err(|e| ApiError::internal(format!("upscale task panicked: {e}")))?;
1078 if let Err(error) = &result {
1079 crate::gpu_worker::quarantine_if_fatal_cuda_error(&worker, error);
1080 }
1081 result.map_err(|e| ApiError::internal(format!("upscale failed: {e}")))?
1082 } else {
1083 let upscaler_cache = state.upscaler_cache.clone();
1084 tokio::task::spawn_blocking(move || -> anyhow::Result<mold_core::UpscaleResponse> {
1085 let mut cache = upscaler_cache.lock().unwrap_or_else(|e| e.into_inner());
1086
1087 let needs_new = cache
1089 .as_ref()
1090 .is_none_or(|e| e.model_name() != model_name_owned);
1091 if needs_new {
1092 let new_engine = mold_inference::create_upscale_engine(
1093 model_name_owned,
1094 weights_path,
1095 mold_inference::LoadStrategy::Eager,
1096 0,
1097 )?;
1098 *cache = Some(new_engine);
1099 }
1100
1101 cache.as_mut().unwrap().upscale(&req)
1102 })
1103 .await
1104 .map_err(|e| ApiError::internal(format!("upscale task panicked: {e}")))?
1105 .map_err(|e| ApiError::internal(format!("upscale failed: {e}")))?
1106 };
1107
1108 Ok(Json(resp))
1109}
1110
1111async fn upscale_stream(
1114 State(state): State<AppState>,
1115 Json(req): Json<mold_core::UpscaleRequest>,
1116) -> Result<Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
1117 if let Err(msg) = mold_core::validate_upscale_request(&req) {
1118 return Err(ApiError::validation(msg));
1119 }
1120
1121 let model_name = mold_core::manifest::resolve_model_name(&req.model);
1122
1123 let needs_pull = {
1125 let config = state.config.read().await;
1126 config
1127 .models
1128 .get(&model_name)
1129 .and_then(|c| c.transformer.as_ref())
1130 .is_none()
1131 };
1132
1133 if needs_pull && mold_core::manifest::find_manifest(&model_name).is_none() {
1135 return Err(ApiError::not_found(format!(
1136 "unknown upscaler model '{}'. Run 'mold list' to see available models.",
1137 model_name
1138 )));
1139 }
1140
1141 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseMessage>();
1142 let model_name_owned = model_name.clone();
1143 let state_clone = state.clone();
1144 let upscaler_cache = state.upscaler_cache.clone();
1145
1146 tokio::spawn(async move {
1147 if needs_pull {
1149 let progress_tx = tx.clone();
1150 let callback =
1151 std::sync::Arc::new(move |event: mold_core::download::DownloadProgressEvent| {
1152 let sse_event = match event {
1153 mold_core::download::DownloadProgressEvent::Status { message } => {
1154 SseProgressEvent::Info { message }
1155 }
1156 mold_core::download::DownloadProgressEvent::FileStart {
1157 filename,
1158 file_index,
1159 total_files,
1160 size_bytes,
1161 batch_bytes_downloaded,
1162 batch_bytes_total,
1163 batch_elapsed_ms,
1164 } => SseProgressEvent::DownloadProgress {
1165 filename,
1166 file_index,
1167 total_files,
1168 bytes_downloaded: 0,
1169 bytes_total: size_bytes,
1170 batch_bytes_downloaded,
1171 batch_bytes_total,
1172 batch_elapsed_ms,
1173 },
1174 mold_core::download::DownloadProgressEvent::FileProgress {
1175 filename,
1176 file_index,
1177 bytes_downloaded,
1178 bytes_total,
1179 batch_bytes_downloaded,
1180 batch_bytes_total,
1181 batch_elapsed_ms,
1182 } => SseProgressEvent::DownloadProgress {
1183 filename,
1184 file_index,
1185 total_files: 0,
1186 bytes_downloaded,
1187 bytes_total,
1188 batch_bytes_downloaded,
1189 batch_bytes_total,
1190 batch_elapsed_ms,
1191 },
1192 mold_core::download::DownloadProgressEvent::FileDone {
1193 filename,
1194 file_index,
1195 total_files,
1196 batch_bytes_downloaded,
1197 batch_bytes_total,
1198 batch_elapsed_ms,
1199 } => SseProgressEvent::DownloadDone {
1200 filename,
1201 file_index,
1202 total_files,
1203 batch_bytes_downloaded,
1204 batch_bytes_total,
1205 batch_elapsed_ms,
1206 },
1207 };
1208 let _ = progress_tx.send(SseMessage::Progress(sse_event));
1209 });
1210
1211 match model_manager::pull_model(&state_clone, &model_name_owned, Some(callback)).await {
1212 Ok(_) => {
1213 let _ = tx.send(SseMessage::Progress(SseProgressEvent::PullComplete {
1214 model: model_name_owned.clone(),
1215 }));
1216 }
1217 Err(e) => {
1218 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1219 message: format!("failed to pull upscaler model: {}", e.error),
1220 }));
1221 return;
1222 }
1223 }
1224 }
1225
1226 let weights_path = {
1228 let config = state_clone.config.read().await;
1229 config
1230 .models
1231 .get(&model_name_owned)
1232 .and_then(|c| c.transformer.as_ref())
1233 .map(std::path::PathBuf::from)
1234 };
1235
1236 let Some(weights_path) = weights_path else {
1237 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1238 message: format!(
1239 "upscaler model '{}' not configured after pull",
1240 model_name_owned
1241 ),
1242 }));
1243 return;
1244 };
1245
1246 let result = if state_clone.gpu_pool.worker_count() > 0 {
1247 match select_aux_worker(&state_clone) {
1248 Ok(worker) => {
1249 worker.in_flight.fetch_add(1, Ordering::SeqCst);
1250 let worker_clone = worker.clone();
1251 let tx_for_worker = tx.clone();
1252 let model_name_for_worker = model_name_owned.clone();
1253 let weights_path_for_worker = weights_path.clone();
1254 let req_for_worker = req.clone();
1255 let result = tokio::task::spawn_blocking(move || {
1256 struct ThreadGpuGuard;
1257 impl Drop for ThreadGpuGuard {
1258 fn drop(&mut self) {
1259 mold_inference::device::clear_thread_gpu_ordinal();
1260 }
1261 }
1262
1263 mold_inference::device::init_thread_gpu_ordinal(worker_clone.gpu.ordinal);
1264 let _thread_gpu = ThreadGpuGuard;
1265 let _load_lock = worker_clone.model_load_lock.lock().unwrap();
1266 if let Err(error) = crate::gpu_worker::ensure_worker_not_poisoned(
1267 &worker_clone,
1268 &model_name_for_worker,
1269 ) {
1270 let _ =
1271 tx_for_worker.send(SseMessage::Error(mold_core::SseErrorEvent {
1272 message: error.to_string(),
1273 }));
1274 return;
1275 }
1276 let _ = tx_for_worker.send(SseMessage::Progress(
1277 mold_core::SseProgressEvent::StageStart {
1278 name: format!(
1279 "Loading upscaler model on GPU {}",
1280 worker_clone.gpu.ordinal
1281 ),
1282 },
1283 ));
1284 let mut engine = match mold_inference::create_upscale_engine(
1285 model_name_for_worker,
1286 weights_path_for_worker,
1287 mold_inference::LoadStrategy::Eager,
1288 worker_clone.gpu.ordinal,
1289 ) {
1290 Ok(engine) => engine,
1291 Err(e) => {
1292 crate::gpu_worker::quarantine_if_fatal_cuda_error(
1293 &worker_clone,
1294 &e,
1295 );
1296 let _ = tx_for_worker.send(SseMessage::Error(
1297 mold_core::SseErrorEvent {
1298 message: format!("failed to load upscaler: {e}"),
1299 },
1300 ));
1301 return;
1302 }
1303 };
1304
1305 let tx_progress = tx_for_worker.clone();
1306 engine.set_on_progress(Box::new(move |event| {
1307 let sse_event: mold_core::SseProgressEvent = event.into();
1308 let _ = tx_progress.send(SseMessage::Progress(sse_event));
1309 }));
1310
1311 match engine.upscale(&req_for_worker) {
1312 Ok(resp) => {
1313 let image_b64 = base64::engine::general_purpose::STANDARD
1314 .encode(&resp.image.data);
1315 let _ = tx_for_worker.send(SseMessage::UpscaleComplete(
1316 mold_core::SseUpscaleCompleteEvent {
1317 image: image_b64,
1318 format: resp.image.format,
1319 model: resp.model,
1320 scale_factor: resp.scale_factor,
1321 original_width: resp.original_width,
1322 original_height: resp.original_height,
1323 upscale_time_ms: resp.upscale_time_ms,
1324 },
1325 ));
1326 }
1327 Err(e) => {
1328 crate::gpu_worker::quarantine_if_fatal_cuda_error(
1329 &worker_clone,
1330 &e,
1331 );
1332 let _ = tx_for_worker.send(SseMessage::Error(
1333 mold_core::SseErrorEvent {
1334 message: format!("upscale failed: {e}"),
1335 },
1336 ));
1337 }
1338 }
1339
1340 engine.clear_on_progress();
1341 })
1342 .await;
1343 worker.in_flight.fetch_sub(1, Ordering::SeqCst);
1344 result
1345 }
1346 Err(e) => {
1347 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1348 message: e.error,
1349 }));
1350 return;
1351 }
1352 }
1353 } else {
1354 let model_name_for_cache = model_name_owned.clone();
1355 let weights_path_for_cache = weights_path.clone();
1356 let req_for_cache = req.clone();
1357 tokio::task::spawn_blocking(move || {
1358 let mut cache = upscaler_cache.lock().unwrap();
1359
1360 let needs_new = cache
1361 .as_ref()
1362 .is_none_or(|e| e.model_name() != model_name_for_cache);
1363 if needs_new {
1364 let _ = tx.send(SseMessage::Progress(
1365 mold_core::SseProgressEvent::StageStart {
1366 name: "Loading upscaler model".to_string(),
1367 },
1368 ));
1369 match mold_inference::create_upscale_engine(
1370 model_name_for_cache,
1371 weights_path_for_cache,
1372 mold_inference::LoadStrategy::Eager,
1373 0,
1374 ) {
1375 Ok(new_engine) => {
1376 *cache = Some(new_engine);
1377 }
1378 Err(e) => {
1379 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1380 message: format!("failed to load upscaler: {e}"),
1381 }));
1382 return;
1383 }
1384 }
1385 }
1386
1387 let engine = cache.as_mut().unwrap();
1388 let tx_progress = tx.clone();
1389 engine.set_on_progress(Box::new(move |event| {
1390 let sse_event: mold_core::SseProgressEvent = event.into();
1391 let _ = tx_progress.send(SseMessage::Progress(sse_event));
1392 }));
1393
1394 match engine.upscale(&req_for_cache) {
1395 Ok(resp) => {
1396 let image_b64 =
1397 base64::engine::general_purpose::STANDARD.encode(&resp.image.data);
1398 let _ = tx.send(SseMessage::UpscaleComplete(
1399 mold_core::SseUpscaleCompleteEvent {
1400 image: image_b64,
1401 format: resp.image.format,
1402 model: resp.model,
1403 scale_factor: resp.scale_factor,
1404 original_width: resp.original_width,
1405 original_height: resp.original_height,
1406 upscale_time_ms: resp.upscale_time_ms,
1407 },
1408 ));
1409 }
1410 Err(e) => {
1411 let _ = tx.send(SseMessage::Error(mold_core::SseErrorEvent {
1412 message: format!("upscale failed: {e}"),
1413 }));
1414 }
1415 }
1416
1417 engine.clear_on_progress();
1418 })
1419 .await
1420 };
1421
1422 if let Err(e) = result {
1423 tracing::error!("upscale task panicked: {e}");
1424 }
1425 });
1426
1427 let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx)
1428 .map(|msg| Ok::<_, Infallible>(sse_message_to_event(msg)));
1429
1430 Ok(Sse::new(stream).keep_alive(
1431 KeepAlive::new()
1432 .interval(std::time::Duration::from_secs(15))
1433 .text("ping"),
1434 ))
1435}
1436
1437const SSE_PAYLOAD_HEADER: &str = "x-mold-sse-payload";
1440
1441pub(crate) fn requested_sse_completion_payload(
1442 headers: &HeaderMap,
1443) -> Result<SseCompletionPayload, ApiError> {
1444 let Some(value) = headers.get(SSE_PAYLOAD_HEADER) else {
1445 return Ok(SseCompletionPayload::Full);
1446 };
1447 match value.to_str().ok() {
1448 Some("metadata-only") => Ok(SseCompletionPayload::MetadataOnly),
1449 _ => Err(ApiError::validation(
1450 "X-Mold-SSE-Payload must be 'metadata-only' when provided",
1451 )),
1452 }
1453}
1454
1455#[utoipa::path(
1456 post,
1457 path = "/api/generate/stream",
1458 tag = "generation",
1459 request_body = mold_core::GenerateRequest,
1460 params((
1461 "X-Mold-SSE-Payload" = Option<String>,
1462 Header,
1463 description = "Set to metadata-only to omit encoded media and return the saved gallery filename"
1464 )),
1465 responses(
1466 (status = 200, description = "SSE event stream with progress and result"),
1467 (status = 404, description = "Model not downloaded"),
1468 (status = 422, description = "Invalid request parameters"),
1469 (status = 500, description = "Inference error"),
1470 (status = 503, description = "Generation queue full"),
1471 )
1472)]
1473async fn generate_stream(
1474 State(state): State<AppState>,
1475 headers: HeaderMap,
1476 Json(mut req): Json<mold_core::GenerateRequest>,
1477) -> Result<Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>>, ApiError> {
1478 let completion_payload = requested_sse_completion_payload(&headers)?;
1479 let typed = (
1482 req.prompt.clone(),
1483 req.negative_prompt.clone(),
1484 req.model.clone(),
1485 );
1486 let (output_dir, dim_warning, preferred_gpu) = prepare_generation(&state, &mut req).await?;
1487 if completion_payload == SseCompletionPayload::MetadataOnly && output_dir.is_none() {
1488 return Err(ApiError::validation(
1489 "metadata-only SSE completions require server gallery output to be enabled",
1490 ));
1491 }
1492 record_prompt_history(&state, &typed.0, typed.1.as_deref(), &typed.2);
1493
1494 tracing::info!(
1495 model = %req.model,
1496 prompt = %req.prompt,
1497 "generate/stream request"
1498 );
1499
1500 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseMessage>();
1502
1503 if let Some(warning) = dim_warning {
1505 let _ = tx.send(SseMessage::Progress(SseProgressEvent::Info {
1506 message: warning,
1507 }));
1508 }
1509
1510 let (result_tx, result_rx) = tokio::sync::oneshot::channel();
1511 let job_id = uuid::Uuid::new_v4().to_string();
1514 let queue_metadata = Box::new(mold_core::OutputMetadata::from_generate_request(
1519 &req,
1520 req.seed.unwrap_or(0),
1521 req.scheduler,
1522 mold_core::build_info::version_string(),
1523 ));
1524 let cancel = state.job_registry.register_job(
1525 &job_id,
1526 &req.model,
1527 preferred_gpu,
1528 Some(req.seed.is_some()),
1529 Some(queue_metadata),
1530 );
1531 let job = GenerationJob {
1532 id: job_id.clone(),
1533 request: req,
1534 completion_payload,
1535 progress_tx: Some(tx.clone()),
1536 result_tx,
1537 output_dir,
1538 };
1539
1540 let position = state
1541 .queue
1542 .submit(job, state.queue_capacity)
1543 .await
1544 .inspect_err(|_| state.job_registry.remove(&job_id))
1545 .map_err(submit_error_to_api)?;
1546
1547 let _ = tx.send(SseMessage::Progress(SseProgressEvent::Queued {
1550 position,
1551 id: job_id,
1552 }));
1553
1554 tokio::spawn(async move {
1561 tokio::select! {
1562 _ = result_rx => {}
1563 _ = cancel.notified() => {
1564 let _ = tx.send(SseMessage::Error(SseErrorEvent {
1565 message: "cancelled".to_string(),
1566 }));
1567 }
1568 }
1569 drop(tx); });
1571
1572 let stream = async_stream::stream! {
1578 let mut rx = rx;
1579 while let Some(msg) = rx.recv().await {
1580 let is_error = matches!(msg, SseMessage::Error(_));
1581 yield Ok::<_, Infallible>(sse_message_to_event(msg));
1582 if is_error {
1583 break;
1584 }
1585 }
1586 };
1587
1588 Ok(Sse::new(stream).keep_alive(
1589 KeepAlive::new()
1590 .interval(std::time::Duration::from_secs(15))
1591 .text("ping"),
1592 ))
1593}
1594
1595#[utoipa::path(
1598 get,
1599 path = "/api/models",
1600 tag = "models",
1601 responses(
1602 (status = 200, description = "List of available models", body = Vec<ModelInfoExtended>),
1603 )
1604)]
1605async fn list_models(State(state): State<AppState>) -> Json<Vec<ModelInfoExtended>> {
1606 Json(model_manager::list_models(&state).await)
1607}
1608
1609async fn generate_estimate(
1610 State(state): State<AppState>,
1611 Json(req): Json<GenerateRequest>,
1612) -> Result<Json<mold_core::GenerationMemoryEstimate>, ApiError> {
1613 Ok(Json(
1614 model_manager::estimate_generation_memory(&state, &req).await?,
1615 ))
1616}
1617
1618async fn model_components(
1619 State(state): State<AppState>,
1620 Path(model): Path<String>,
1621) -> Result<Json<mold_core::ModelComponentsResponse>, ApiError> {
1622 Ok(Json(
1623 model_manager::model_component_status(&state, &model).await?,
1624 ))
1625}
1626
1627#[derive(Debug, Deserialize, utoipa::ToSchema)]
1630pub struct LoadModelBody {
1631 #[schema(example = "flux-schnell:q8")]
1632 pub model: String,
1633 #[serde(default, skip_serializing_if = "Option::is_none")]
1636 pub gpu: Option<usize>,
1637}
1638
1639#[utoipa::path(
1640 post,
1641 path = "/api/models/load",
1642 tag = "models",
1643 request_body = LoadModelBody,
1644 responses(
1645 (status = 200, description = "Model loaded successfully"),
1646 (status = 404, description = "Model not downloaded"),
1647 (status = 400, description = "Unknown model"),
1648 (status = 500, description = "Failed to load model"),
1649 )
1650)]
1651async fn load_model(
1652 State(state): State<AppState>,
1653 Json(body): Json<LoadModelBody>,
1654) -> Result<impl IntoResponse, ApiError> {
1655 if let Err(e) = model_manager::install_catalog_model(&state, &body.model).await {
1656 return Err(model_manager::install_error_to_api_error(&e));
1657 }
1658 let _ = model_manager::check_model_available(&state, &body.model).await?;
1659
1660 if state.gpu_pool.worker_count() > 0 {
1662 let worker = match body.gpu {
1663 Some(ordinal) => state
1664 .gpu_pool
1665 .workers
1666 .iter()
1667 .find(|w| w.gpu.ordinal == ordinal)
1668 .cloned()
1669 .ok_or_else(|| {
1670 ApiError::not_found(format!("no GPU worker with ordinal {ordinal}"))
1671 })?,
1672 None => {
1673 let est = crate::queue::estimate_model_vram(&body.model);
1674 state
1675 .gpu_pool
1676 .select_worker(&body.model, est)
1677 .ok_or_else(|| {
1678 ApiError::internal(format!(
1679 "no GPU available to load model '{}'",
1680 body.model
1681 ))
1682 })?
1683 }
1684 };
1685 let config_snapshot = state.config.read().await.clone();
1686 let model_name = body.model.clone();
1687 let worker_clone = worker.clone();
1688 tokio::task::spawn_blocking(move || {
1689 crate::gpu_worker::load_blocking(&worker_clone, &model_name, &config_snapshot)
1690 })
1691 .await
1692 .map_err(|e| ApiError::internal(format!("model load task failed: {e}")))?
1693 .map_err(|e| ApiError::internal(format!("model load error: {e}")))?;
1694
1695 tracing::info!(
1696 model = %body.model,
1697 gpu = worker.gpu.ordinal,
1698 "model loaded via API"
1699 );
1700 return Ok(StatusCode::OK);
1701 }
1702
1703 model_manager::ensure_model_ready(&state, &body.model, None, None, false).await?;
1706 tracing::info!(model = %body.model, gpu = ?body.gpu, "model loaded via API (legacy)");
1707 Ok(StatusCode::OK)
1708}
1709
1710#[utoipa::path(
1713 post,
1714 path = "/api/models/pull",
1715 tag = "models",
1716 request_body = LoadModelBody,
1717 responses(
1718 (status = 200, description = "Model pulled (SSE stream or plain text)"),
1719 (status = 400, description = "Unknown model"),
1720 (status = 500, description = "Download failed"),
1721 )
1722)]
1723async fn pull_model_endpoint(
1724 State(state): State<AppState>,
1725 headers: HeaderMap,
1726 Json(body): Json<LoadModelBody>,
1727) -> Result<impl IntoResponse, ApiError> {
1728 let wants_sse = headers
1729 .get(header::ACCEPT)
1730 .and_then(|v| v.to_str().ok())
1731 .is_some_and(|v| v.contains("text/event-stream"));
1732
1733 if body.model.starts_with("cv:") || body.model.starts_with("hf:") {
1734 if let Err(e) = model_manager::install_catalog_model(&state, &body.model).await {
1735 return Err(model_manager::install_error_to_api_error(&e));
1736 }
1737 if model_manager::check_model_available(&state, &body.model)
1738 .await
1739 .is_ok()
1740 {
1741 return Ok(PullResponse::Text(format!(
1742 "model '{}' is already present",
1743 body.model
1744 )));
1745 }
1746 let companion_names = {
1747 let intents = state.catalog_intents.read().await;
1748 intents
1749 .get(&body.model)
1750 .map(|intent| {
1751 intent
1752 .companions
1753 .iter()
1754 .map(|companion| companion.name.clone())
1755 .collect::<Vec<_>>()
1756 })
1757 .unwrap_or_default()
1758 };
1759 let models_dir = state.config.read().await.resolved_models_dir();
1760 let companion_jobs = crate::catalog_api::enqueue_missing_companions(
1761 &companion_names,
1762 &models_dir,
1763 &state.downloads,
1764 Some(&body.model),
1765 None,
1766 )
1767 .await;
1768 let primary_job = crate::catalog_api::enqueue_catalog_primary_repair(&state, &body.model)
1769 .await
1770 .map_err(|(status, msg)| ApiError::internal_with_status(msg, status))?;
1771 if !companion_jobs.is_empty() || primary_job.is_some() {
1772 let primary_count = usize::from(primary_job.is_some());
1773 return Ok(PullResponse::Text(format!(
1774 "queued repair for model '{}' ({} primary job(s), {} companion job(s))",
1775 body.model,
1776 primary_count,
1777 companion_jobs.len()
1778 )));
1779 }
1780 model_manager::check_model_available(&state, &body.model).await?;
1781 return Ok(PullResponse::Text(format!(
1782 "model '{}' is already present",
1783 body.model
1784 )));
1785 }
1786
1787 let (job_id, _position) = match state.downloads.enqueue(body.model.clone()).await {
1789 Ok((id, pos, _)) => (id, pos),
1790 Err(crate::downloads::EnqueueError::UnknownModel(_)) => {
1791 return Err(ApiError::unknown_model(format!(
1792 "unknown model '{}'. Run 'mold list' to see available models.",
1793 body.model
1794 )));
1795 }
1796 Err(crate::downloads::EnqueueError::LockPoisoned) => {
1797 return Err(ApiError::internal("download queue state is corrupt"));
1798 }
1799 };
1800
1801 if !wants_sse {
1802 let mut rx = state.downloads.subscribe();
1804 loop {
1805 match rx.recv().await {
1806 Ok(mold_core::types::DownloadEvent::JobDone { id, model }) if id == job_id => {
1807 return Ok(PullResponse::Text(format!(
1808 "model '{model}' pulled successfully"
1809 )));
1810 }
1811 Ok(mold_core::types::DownloadEvent::JobFailed { id, error }) if id == job_id => {
1812 return Err(ApiError::internal(format!(
1813 "failed to pull model '{}': {error}",
1814 body.model
1815 )));
1816 }
1817 Ok(mold_core::types::DownloadEvent::JobCancelled { id }) if id == job_id => {
1818 return Err(ApiError::internal(format!(
1819 "pull of '{}' was cancelled",
1820 body.model
1821 )));
1822 }
1823 Ok(_) => continue,
1824 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
1825 Err(tokio::sync::broadcast::error::RecvError::Closed) => {
1826 return Err(ApiError::internal("download queue channel closed"));
1827 }
1828 }
1829 }
1830 }
1831
1832 let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<SseMessage>();
1835 let mut events = state.downloads.subscribe();
1836 let model_for_cb = body.model.clone();
1837 tokio::spawn(async move {
1838 loop {
1839 match events.recv().await {
1840 Ok(mold_core::types::DownloadEvent::Started {
1841 id,
1842 files_total,
1843 bytes_total,
1844 }) if id == job_id => {
1845 let _ = tx.send(SseMessage::Progress(SseProgressEvent::DownloadProgress {
1846 filename: String::new(),
1847 file_index: 0,
1848 total_files: files_total,
1849 bytes_downloaded: 0,
1850 bytes_total,
1851 batch_bytes_downloaded: 0,
1852 batch_bytes_total: bytes_total,
1853 batch_elapsed_ms: 0,
1854 }));
1855 }
1856 Ok(mold_core::types::DownloadEvent::Progress {
1857 id,
1858 files_done,
1859 bytes_done,
1860 current_file,
1861 }) if id == job_id => {
1862 let _ = tx.send(SseMessage::Progress(SseProgressEvent::DownloadProgress {
1863 filename: current_file.unwrap_or_default(),
1864 file_index: files_done,
1865 total_files: 0,
1866 bytes_downloaded: bytes_done,
1867 bytes_total: 0,
1868 batch_bytes_downloaded: bytes_done,
1869 batch_bytes_total: 0,
1870 batch_elapsed_ms: 0,
1871 }));
1872 }
1873 Ok(mold_core::types::DownloadEvent::JobDone { id, .. }) if id == job_id => {
1874 let _ = tx.send(SseMessage::Progress(SseProgressEvent::PullComplete {
1875 model: model_for_cb.clone(),
1876 }));
1877 break;
1878 }
1879 Ok(mold_core::types::DownloadEvent::JobFailed { id, error }) if id == job_id => {
1880 let _ = tx.send(SseMessage::Error(SseErrorEvent { message: error }));
1881 break;
1882 }
1883 Ok(mold_core::types::DownloadEvent::JobCancelled { id }) if id == job_id => {
1884 let _ = tx.send(SseMessage::Error(SseErrorEvent {
1885 message: "pull cancelled".into(),
1886 }));
1887 break;
1888 }
1889 Ok(_) => continue,
1890 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
1891 Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
1892 }
1893 }
1894 });
1895
1896 let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(rx)
1897 .map(|msg| Ok::<_, Infallible>(sse_message_to_event(msg)));
1898
1899 Ok(PullResponse::Sse(
1900 Sse::new(stream)
1901 .keep_alive(
1902 KeepAlive::new()
1903 .interval(std::time::Duration::from_secs(15))
1904 .text("ping"),
1905 )
1906 .into_response(),
1907 ))
1908}
1909
1910enum PullResponse {
1912 Sse(axum::response::Response),
1913 Text(String),
1914}
1915
1916impl IntoResponse for PullResponse {
1917 fn into_response(self) -> axum::response::Response {
1918 match self {
1919 PullResponse::Sse(resp) => resp,
1920 PullResponse::Text(text) => text.into_response(),
1921 }
1922 }
1923}
1924
1925#[derive(Debug, Default, Deserialize, utoipa::ToSchema)]
1930pub struct UnloadRequest {
1931 #[serde(default, skip_serializing_if = "Option::is_none")]
1933 pub model: Option<String>,
1934 #[serde(default, skip_serializing_if = "Option::is_none")]
1936 pub gpu: Option<usize>,
1937}
1938
1939#[utoipa::path(
1940 delete,
1941 path = "/api/models/unload",
1942 tag = "models",
1943 request_body(content = Option<UnloadRequest>, content_type = "application/json"),
1944 responses(
1945 (status = 200, description = "Model unloaded or no model was loaded", body = String),
1946 )
1947)]
1948async fn unload_model(
1949 State(state): State<AppState>,
1950 body: Option<Json<UnloadRequest>>,
1951) -> Result<impl IntoResponse, ApiError> {
1952 let req = body.map(|b| b.0).unwrap_or_default();
1953 tracing::debug!(model = ?req.model, gpu = ?req.gpu, "unload request");
1954 clear_global_upscaler_cache(&state);
1955
1956 if state.gpu_pool.worker_count() > 0 {
1958 let targets: Vec<_> = match (req.gpu, req.model.as_deref()) {
1960 (Some(ordinal), _) => state
1961 .gpu_pool
1962 .workers
1963 .iter()
1964 .filter(|w| w.gpu.ordinal == ordinal)
1965 .cloned()
1966 .collect(),
1967 (None, Some(model)) => state
1968 .gpu_pool
1969 .workers
1970 .iter()
1971 .filter(|w| {
1972 let cache = w.model_cache.lock().unwrap();
1973 cache
1974 .get(model)
1975 .map(|e| e.residency == crate::model_cache::ModelResidency::Gpu)
1976 .unwrap_or(false)
1977 })
1978 .cloned()
1979 .collect(),
1980 (None, None) => state.gpu_pool.workers.clone(),
1981 };
1982
1983 if targets.is_empty() {
1984 return Ok((StatusCode::OK, "no model loaded".to_string()));
1985 }
1986
1987 let mut unloaded_pairs: Vec<(usize, String)> = Vec::new();
1988 for worker in targets {
1989 let worker_clone = worker.clone();
1990 let result = tokio::task::spawn_blocking(move || {
1991 crate::gpu_worker::unload_blocking(&worker_clone)
1992 })
1993 .await
1994 .map_err(|e| ApiError::internal(format!("unload task failed: {e}")))?;
1995 if let Some(name) = result {
1996 unloaded_pairs.push((worker.gpu.ordinal, name));
1997 }
1998 }
1999
2000 let msg = if unloaded_pairs.is_empty() {
2001 "no model loaded".to_string()
2002 } else {
2003 let joined: Vec<String> = unloaded_pairs
2004 .iter()
2005 .map(|(o, m)| format!("gpu{o}:{m}"))
2006 .collect();
2007 format!("unloaded {}", joined.join(", "))
2008 };
2009 return Ok((StatusCode::OK, msg));
2010 }
2011
2012 Ok((StatusCode::OK, model_manager::unload_model(&state).await))
2014}
2015
2016async fn model_is_gpu_resident(state: &AppState, canonical: &str) -> bool {
2022 {
2023 let cache = state.model_cache.lock().await;
2024 if cache.active_model() == Some(canonical) {
2025 return true;
2026 }
2027 }
2028 if state
2029 .active_generation
2030 .read()
2031 .unwrap_or_else(|e| e.into_inner())
2032 .as_ref()
2033 .is_some_and(|g| g.model == canonical)
2034 {
2035 return true;
2036 }
2037 for worker in &state.gpu_pool.workers {
2038 if let Ok(active) = worker.active_generation.read() {
2039 if active.as_ref().is_some_and(|g| g.model == canonical) {
2040 return true;
2041 }
2042 }
2043 if let Ok(cache) = worker.model_cache.lock() {
2044 if cache.active_model() == Some(canonical) {
2045 return true;
2046 }
2047 }
2048 }
2049 false
2050}
2051
2052#[utoipa::path(
2061 delete,
2062 path = "/api/models/{model}",
2063 tag = "models",
2064 params(("model" = String, Path, description = "Model name (e.g. flux-schnell:q8)")),
2065 responses(
2066 (status = 200, description = "Model removed", body = mold_core::ModelRemovalResponse),
2067 (status = 404, description = "Model not installed"),
2068 (status = 409, description = "Model is currently loaded — unload it first"),
2069 )
2070)]
2071async fn delete_model(
2072 State(state): State<AppState>,
2073 Path(model): Path<String>,
2074) -> Result<Json<mold_core::ModelRemovalResponse>, ApiError> {
2075 let canonical = mold_core::manifest::resolve_model_name(&model);
2076
2077 if model_is_gpu_resident(&state, &canonical).await
2083 || (model != canonical && model_is_gpu_resident(&state, &model).await)
2084 {
2085 return Err(ApiError::with_code(
2086 format!(
2087 "model '{canonical}' is currently loaded; unload it first (DELETE /api/models/unload)"
2088 ),
2089 "MODEL_LOADED",
2090 StatusCode::CONFLICT,
2091 ));
2092 }
2093
2094 let mut config = state.config.write().await;
2097 let in_config = config.models.contains_key(&canonical);
2098 if !in_config && !config.manifest_model_is_downloaded(&canonical) {
2099 return Err(ApiError::with_code(
2100 format!("model '{canonical}' is not installed"),
2101 "UNKNOWN_MODEL",
2102 StatusCode::NOT_FOUND,
2103 ));
2104 }
2105
2106 tracing::info!(model = %canonical, "model removal requested");
2107 let plan = mold_core::removal::plan_removal(&config, &canonical);
2108 let outcome = mold_core::removal::execute_removal(&config, &plan);
2109 for warning in &outcome.warnings {
2110 tracing::warn!(model = %canonical, "model removal: {warning}");
2111 }
2112
2113 mold_core::download::remove_pulling_marker(&canonical);
2114 if in_config {
2115 config.remove_model(&canonical);
2116 if let Err(e) = config.save() {
2117 tracing::warn!("failed to persist model removal to config.toml: {e}");
2118 }
2119 }
2120 drop(config);
2121
2122 {
2125 let mut cache = state.model_cache.lock().await;
2126 let _ = cache.remove(&canonical);
2127 }
2128 for worker in &state.gpu_pool.workers {
2129 if let Ok(mut cache) = worker.model_cache.lock() {
2130 let _ = cache.remove(&canonical);
2131 }
2132 }
2133
2134 let kept = plan
2135 .shared_files
2136 .iter()
2137 .map(|(path, used_by)| mold_core::KeptComponent {
2138 component: path.clone(),
2139 used_by: used_by.clone(),
2140 })
2141 .collect();
2142
2143 Ok(Json(mold_core::ModelRemovalResponse {
2144 removed: outcome.removed,
2145 kept,
2146 freed_bytes: outcome.freed_bytes,
2147 }))
2148}
2149
2150fn disk_usage_for_path(
2156 disks: &[(std::path::PathBuf, u64, u64)],
2157 dir: &std::path::Path,
2158) -> Option<DiskUsage> {
2159 disks
2160 .iter()
2161 .filter(|(mount, _, _)| dir.starts_with(mount))
2162 .max_by_key(|(mount, _, _)| mount.as_os_str().len())
2163 .map(|&(_, total_bytes, free_bytes)| DiskUsage {
2164 total_bytes,
2165 free_bytes,
2166 })
2167}
2168
2169fn canonical_models_dir(dir: &std::path::Path) -> std::path::PathBuf {
2174 std::fs::canonicalize(dir).unwrap_or_else(|_| dir.to_path_buf())
2175}
2176
2177fn models_disk_usage(dir: &std::path::Path) -> Option<DiskUsage> {
2182 let dir = canonical_models_dir(dir);
2183 let disks = sysinfo::Disks::new_with_refreshed_list();
2184 let entries: Vec<(std::path::PathBuf, u64, u64)> = disks
2185 .list()
2186 .iter()
2187 .map(|d| {
2188 (
2189 d.mount_point().to_path_buf(),
2190 d.total_space(),
2191 d.available_space(),
2192 )
2193 })
2194 .collect();
2195 disk_usage_for_path(&entries, &dir)
2196}
2197
2198#[utoipa::path(
2199 get,
2200 path = "/api/status",
2201 tag = "server",
2202 responses(
2203 (status = 200, description = "Server status", body = ServerStatus),
2204 )
2205)]
2206async fn server_status(State(state): State<AppState>) -> Json<ServerStatus> {
2207 let (models_disk, needs_disk_refresh) = state.models_disk_cache.read();
2213 if needs_disk_refresh {
2214 let cache = state.models_disk_cache.clone();
2215 let models_dir = state.config.read().await.resolved_models_dir();
2216 tokio::task::spawn_blocking(move || cache.store(models_disk_usage(&models_dir)));
2217 }
2218
2219 let mut gpu_statuses = state.gpu_pool.gpu_status();
2224 if let Some(resources) = state.resources.latest() {
2225 let visible_devices = std::env::var("CUDA_VISIBLE_DEVICES").ok();
2226 for status in &mut gpu_statuses {
2227 let Some(physical_ordinal) = crate::resources::physical_ordinal_for_worker(
2228 status.ordinal,
2229 visible_devices.as_deref(),
2230 ) else {
2231 continue;
2232 };
2233 if let Some(snapshot) = resources
2234 .gpus
2235 .iter()
2236 .find(|gpu| gpu.ordinal == physical_ordinal)
2237 {
2238 status.vram_total_bytes = snapshot.vram_total;
2239 status.vram_used_bytes = snapshot.vram_used;
2240 }
2241 }
2242 }
2243 let has_gpus = !gpu_statuses.is_empty();
2244
2245 let gpu_models_loaded: Vec<String> = gpu_statuses
2247 .iter()
2248 .filter_map(|g| g.loaded_model.clone())
2249 .collect();
2250 let gpu_busy = gpu_statuses
2251 .iter()
2252 .any(|g| g.state == GpuWorkerState::Generating);
2253
2254 let multi_gpu_current_gen = if has_gpus {
2257 state.gpu_pool.workers.iter().find_map(|w| {
2258 let gen = w.active_generation.read().ok()?;
2259 gen.as_ref().map(|g| ActiveGenerationStatus {
2260 model: g.model.clone(),
2261 prompt_sha256: g.prompt_sha256.clone(),
2262 started_at_unix_ms: g.started_at_unix_ms,
2263 elapsed_ms: g.started_at.elapsed().as_millis() as u64,
2264 })
2265 })
2266 } else {
2267 None
2268 };
2269
2270 let (models_loaded, busy, current_generation) = if has_gpus {
2272 (gpu_models_loaded, gpu_busy, multi_gpu_current_gen)
2273 } else {
2274 let snapshot = state.model_cache.lock().await.snapshot();
2275 let models = match (snapshot.model_name, snapshot.is_loaded) {
2276 (Some(model_name), true) => vec![model_name],
2277 _ => vec![],
2278 };
2279 let gen = state
2280 .active_generation
2281 .read()
2282 .unwrap_or_else(|e| e.into_inner())
2283 .as_ref()
2284 .map(|active| ActiveGenerationStatus {
2285 model: active.model.clone(),
2286 prompt_sha256: active.prompt_sha256.clone(),
2287 started_at_unix_ms: active.started_at_unix_ms,
2288 elapsed_ms: active.started_at.elapsed().as_millis() as u64,
2289 });
2290 let is_busy = gen.is_some();
2291 (models, is_busy, gen)
2292 };
2293
2294 Json(ServerStatus {
2295 version: env!("CARGO_PKG_VERSION").to_string(),
2296 git_sha: if mold_core::build_info::GIT_SHA == "unknown" {
2297 None
2298 } else {
2299 Some(mold_core::build_info::GIT_SHA.to_string())
2300 },
2301 build_date: if mold_core::build_info::BUILD_DATE == "unknown" {
2302 None
2303 } else {
2304 Some(mold_core::build_info::BUILD_DATE.to_string())
2305 },
2306 models_loaded,
2307 busy,
2308 current_generation,
2309 gpu_info: query_gpu_info(),
2310 uptime_secs: state.start_time.elapsed().as_secs(),
2311 hostname: hostname::get().ok().and_then(|h| h.into_string().ok()),
2312 memory_status: mold_inference::device::memory_status_string(),
2313 gpus: if has_gpus { Some(gpu_statuses) } else { None },
2314 queue_depth: Some(state.queue.pending()),
2315 queue_capacity: Some(state.queue_capacity),
2316 queue_paused: Some(state.queue_pause.is_paused()),
2317 instance_id: Some(state.instance_id.as_ref().clone()),
2318 models_disk,
2319 })
2320}
2321
2322#[utoipa::path(
2325 get,
2326 path = "/health",
2327 tag = "server",
2328 responses(
2329 (status = 200, description = "Server is healthy"),
2330 )
2331)]
2332async fn health() -> impl IntoResponse {
2333 StatusCode::OK
2334}
2335
2336#[utoipa::path(
2343 get,
2344 path = "/api/queue",
2345 tag = "queue",
2346 responses(
2347 (status = 200, description = "Queue snapshot", body = crate::job_registry::QueueListing),
2348 )
2349)]
2350async fn list_queue(State(state): State<AppState>) -> Json<crate::job_registry::QueueListing> {
2351 Json(state.job_registry.snapshot())
2352}
2353
2354fn deserialize_some<'de, T, D>(deserializer: D) -> Result<Option<T>, D::Error>
2359where
2360 T: Deserialize<'de>,
2361 D: serde::Deserializer<'de>,
2362{
2363 T::deserialize(deserializer).map(Some)
2364}
2365
2366#[derive(Debug, Deserialize, utoipa::ToSchema)]
2367struct QueuePatchRequest {
2368 #[serde(default, deserialize_with = "deserialize_some")]
2372 #[schema(value_type = Option<usize>)]
2373 target_gpu: Option<Option<usize>>,
2374 #[serde(default)]
2377 position: Option<usize>,
2378}
2379
2380#[utoipa::path(
2381 patch,
2382 path = "/api/queue/{id}",
2383 tag = "queue",
2384 request_body = QueuePatchRequest,
2385 responses(
2386 (status = 200, description = "Updated queue entry", body = crate::job_registry::JobEntry),
2387 (status = 404, description = "Queue job not found"),
2388 (status = 409, description = "Queue job is already running"),
2389 (status = 422, description = "Invalid GPU target"),
2390 )
2391)]
2392async fn patch_queue_job(
2393 State(state): State<AppState>,
2394 Path(id): Path<String>,
2395 Json(req): Json<QueuePatchRequest>,
2396) -> Result<Json<crate::job_registry::JobEntry>, ApiError> {
2397 if let Some(Some(target)) = req.target_gpu {
2398 let available = state
2399 .gpu_pool
2400 .workers
2401 .iter()
2402 .any(|w| w.gpu.ordinal == target);
2403 if !available {
2404 return Err(ApiError::validation(format!(
2405 "gpu:{target} is not available in this server's worker pool"
2406 )));
2407 }
2408 }
2409
2410 if let Some(target_gpu) = req.target_gpu {
2414 state
2415 .job_registry
2416 .set_target_gpu(&id, target_gpu)
2417 .map_err(|e| match e {
2418 crate::job_registry::TargetGpuUpdateError::NotFound => {
2419 ApiError::queue_job_not_found(format!("queue job {id} not found"))
2420 }
2421 crate::job_registry::TargetGpuUpdateError::AlreadyRunning => {
2422 ApiError::queue_job_running(format!(
2423 "queue job {id} is already running; lane changes only apply to queued jobs"
2424 ))
2425 }
2426 })?;
2427 }
2428
2429 if let Some(position) = req.position {
2430 state
2431 .job_registry
2432 .reorder_queued(&id, position)
2433 .map_err(|e| match e {
2434 crate::job_registry::QueueReorderError::NotFound => {
2435 ApiError::queue_job_not_found(format!("queue job {id} not found"))
2436 }
2437 crate::job_registry::QueueReorderError::AlreadyRunning => {
2438 ApiError::queue_job_running(format!(
2439 "queue job {id} is already running; only queued jobs can be reordered"
2440 ))
2441 }
2442 })?;
2443 }
2444
2445 let entry = state
2446 .job_registry
2447 .entry(&id)
2448 .ok_or_else(|| ApiError::queue_job_not_found(format!("queue job {id} not found")))?;
2449 Ok(Json(entry))
2450}
2451
2452#[utoipa::path(
2458 delete,
2459 path = "/api/queue/{id}",
2460 tag = "queue",
2461 params(("id" = String, Path, description = "Queue job id")),
2462 responses(
2463 (status = 204, description = "Queued job cancelled"),
2464 (status = 404, description = "Queue job not found"),
2465 (status = 409, description = "Queue job is already running"),
2466 )
2467)]
2468async fn cancel_queue_job(
2469 State(state): State<AppState>,
2470 Path(id): Path<String>,
2471) -> Result<StatusCode, ApiError> {
2472 state.job_registry.cancel_queued(&id).map_err(|e| match e {
2473 crate::job_registry::QueuedJobCancelError::NotFound => {
2474 ApiError::queue_job_not_found(format!("queue job {id} not found"))
2475 }
2476 crate::job_registry::QueuedJobCancelError::AlreadyRunning => ApiError::queue_job_running(
2477 format!("queue job {id} is already running; only queued jobs can be cancelled"),
2478 ),
2479 })?;
2480 Ok(StatusCode::NO_CONTENT)
2481}
2482
2483#[derive(Debug, Serialize, utoipa::ToSchema)]
2486struct QueuePauseResponse {
2487 paused: bool,
2488}
2489
2490#[derive(Debug, Serialize, utoipa::ToSchema)]
2492struct QueueCancelAllResponse {
2493 cancelled: usize,
2494}
2495
2496#[utoipa::path(
2500 post,
2501 path = "/api/queue/pause",
2502 tag = "queue",
2503 responses(
2504 (status = 200, description = "Queue dispatch paused", body = QueuePauseResponse),
2505 )
2506)]
2507async fn pause_queue(State(state): State<AppState>) -> Json<QueuePauseResponse> {
2508 if state.queue_pause.pause() {
2509 state.events.publish(mold_core::ServerEvent::QueuePaused);
2510 }
2511 Json(QueuePauseResponse { paused: true })
2512}
2513
2514#[utoipa::path(
2517 post,
2518 path = "/api/queue/resume",
2519 tag = "queue",
2520 responses(
2521 (status = 200, description = "Queue dispatch resumed", body = QueuePauseResponse),
2522 )
2523)]
2524async fn resume_queue(State(state): State<AppState>) -> Json<QueuePauseResponse> {
2525 if state.queue_pause.resume() {
2526 state.events.publish(mold_core::ServerEvent::QueueResumed);
2527 }
2528 Json(QueuePauseResponse { paused: false })
2529}
2530
2531#[utoipa::path(
2536 delete,
2537 path = "/api/queue",
2538 tag = "queue",
2539 responses(
2540 (status = 200, description = "Queued jobs cancelled", body = QueueCancelAllResponse),
2541 )
2542)]
2543async fn cancel_all_queue(State(state): State<AppState>) -> Json<QueueCancelAllResponse> {
2544 let cancelled = state.job_registry.cancel_all_queued();
2545 Json(QueueCancelAllResponse { cancelled })
2546}
2547
2548const HISTORY_DEFAULT_LIMIT: usize = 50;
2552const HISTORY_MAX_LIMIT: usize = 500;
2554
2555const HISTORY_UNAVAILABLE: &str = "HISTORY_UNAVAILABLE";
2557
2558fn history_db(state: &AppState) -> Result<&mold_db::MetadataDb, ApiError> {
2559 state.metadata_db.as_ref().as_ref().ok_or_else(|| {
2560 ApiError::with_code(
2561 "prompt history is unavailable because the metadata DB is disabled",
2562 HISTORY_UNAVAILABLE,
2563 StatusCode::SERVICE_UNAVAILABLE,
2564 )
2565 })
2566}
2567
2568#[derive(Debug, Deserialize)]
2569struct HistoryListQuery {
2570 query: Option<String>,
2572 limit: Option<usize>,
2574}
2575
2576#[utoipa::path(
2577 get,
2578 path = "/api/history",
2579 tag = "server",
2580 params(
2581 ("query" = Option<String>, Query, description = "Substring filter over prompt text (case-insensitive)"),
2582 ("limit" = Option<usize>, Query, description = "Max rows to return (default 50, max 500)"),
2583 ),
2584 responses(
2585 (status = 200, description = "Prompt history, newest first", body = mold_core::HistoryListing),
2586 (status = 503, description = "Metadata DB disabled"),
2587 )
2588)]
2589async fn list_history(
2590 State(state): State<AppState>,
2591 axum::extract::Query(params): axum::extract::Query<HistoryListQuery>,
2592) -> Result<Json<mold_core::HistoryListing>, ApiError> {
2593 let db = history_db(&state)?;
2594 let limit = params
2595 .limit
2596 .unwrap_or(HISTORY_DEFAULT_LIMIT)
2597 .min(HISTORY_MAX_LIMIT);
2598 let history = mold_db::PromptHistory::new(db);
2599 let rows = match params
2600 .query
2601 .as_deref()
2602 .map(str::trim)
2603 .filter(|q| !q.is_empty())
2604 {
2605 Some(query) => history.search(query, limit),
2606 None => history.recent(limit),
2607 }
2608 .map_err(|e| ApiError::internal(format!("failed to read prompt history: {e:#}")))?;
2609 let entries = rows
2610 .into_iter()
2611 .map(|e| mold_core::HistoryEntry {
2612 prompt: e.prompt,
2613 model: e.model,
2614 used_at: e.created_at_ms,
2615 })
2616 .collect();
2617 Ok(Json(mold_core::HistoryListing { entries }))
2618}
2619
2620#[derive(Debug, Deserialize)]
2621struct HistoryDeleteQuery {
2622 keep: Option<usize>,
2624}
2625
2626#[utoipa::path(
2627 delete,
2628 path = "/api/history",
2629 tag = "server",
2630 params(
2631 ("keep" = Option<usize>, Query, description = "Keep only the most recent N entries instead of clearing everything"),
2632 ),
2633 responses(
2634 (status = 204, description = "Prompt history cleared (or trimmed)"),
2635 (status = 503, description = "Metadata DB disabled"),
2636 )
2637)]
2638async fn delete_history(
2639 State(state): State<AppState>,
2640 axum::extract::Query(params): axum::extract::Query<HistoryDeleteQuery>,
2641) -> Result<StatusCode, ApiError> {
2642 let db = history_db(&state)?;
2643 let history = mold_db::PromptHistory::new(db);
2644 match params.keep {
2645 Some(keep) => history.trim_to(keep),
2646 None => history.clear(),
2647 }
2648 .map_err(|e| ApiError::internal(format!("failed to clear prompt history: {e:#}")))?;
2649 Ok(StatusCode::NO_CONTENT)
2650}
2651
2652async fn server_capabilities(State(state): State<AppState>) -> Json<mold_core::ServerCapabilities> {
2658 let catalog_available = std::env::var("MOLD_CATALOG_DISABLE")
2659 .map(|v| v != "1" && !v.eq_ignore_ascii_case("true"))
2660 .unwrap_or(true);
2661 let config = state.config.read().await;
2662 let expand_settings = config.expand.clone().with_env_overrides();
2663 let expand_model_present =
2664 expand_settings.is_local() && config.manifest_model_is_downloaded(&expand_settings.model);
2665 let expand = expand_capabilities(&expand_settings, expand_model_present);
2666
2667 Json(mold_core::ServerCapabilities {
2668 gallery: mold_core::GalleryCapabilities { can_delete: true },
2669 catalog: mold_core::CatalogCapabilities {
2670 available: catalog_available,
2671 families: mold_catalog::families::ALL_FAMILIES
2672 .iter()
2673 .map(|f| f.as_str().to_string())
2674 .collect::<Vec<_>>(),
2675 sort: mold_catalog::live::CatalogSort::WIRE_VALUES
2676 .iter()
2677 .map(|s| s.to_string())
2678 .collect::<Vec<_>>(),
2679 },
2680 events: mold_core::EventsCapabilities { available: true },
2681 queue: mold_core::QueueCapabilities {
2682 can_pause: true,
2683 can_cancel_all: true,
2684 can_reorder: true,
2685 },
2686 expand: Some(expand),
2687 })
2688}
2689
2690fn expand_capabilities(
2694 settings: &mold_core::expand::ExpandSettings,
2695 model_present: bool,
2696) -> mold_core::ExpandCapabilities {
2697 if settings.is_local() {
2698 mold_core::ExpandCapabilities {
2699 configured: cfg!(feature = "expand"),
2700 model_present: Some(model_present),
2701 backend: mold_core::ExpandBackend::Local,
2702 }
2703 } else {
2704 mold_core::ExpandCapabilities {
2705 configured: !settings.backend.trim().is_empty(),
2706 model_present: None,
2707 backend: mold_core::ExpandBackend::Api,
2708 }
2709 }
2710}
2711
2712#[utoipa::path(
2715 get,
2716 path = "/api/capabilities/chain-limits",
2717 tag = "server",
2718 params(
2719 ("model" = String, Query, description = "Model name (e.g. ltx-2-19b-distilled:fp8)")
2720 ),
2721 responses(
2722 (status = 200, description = "Chain limits for the requested model",
2723 body = crate::chain_limits::ChainLimits),
2724 (status = 400, description = "Missing required 'model' query parameter"),
2725 (status = 404, description = "Unknown or unsupported model"),
2726 )
2727)]
2728async fn capabilities_chain_limits(
2729 axum::extract::Query(params): axum::extract::Query<std::collections::HashMap<String, String>>,
2730) -> axum::response::Response {
2731 let raw_model = match params.get("model") {
2732 Some(m) => m.clone(),
2733 None => {
2734 return (
2735 StatusCode::BAD_REQUEST,
2736 "missing required 'model' query parameter\n",
2737 )
2738 .into_response();
2739 }
2740 };
2741
2742 let resolved = mold_core::manifest::resolve_model_name(&raw_model);
2743 let Some(manifest) = mold_core::manifest::find_manifest(&resolved) else {
2744 return (StatusCode::NOT_FOUND, "unknown model\n").into_response();
2745 };
2746 let family = manifest.family.clone();
2747
2748 if crate::chain_limits::family_cap(&family).is_none() {
2749 return (StatusCode::NOT_FOUND, "model is not chain-capable\n").into_response();
2750 }
2751
2752 let quant = resolved
2753 .split_once(':')
2754 .map(|(_, tag)| tag.to_string())
2755 .unwrap_or_default();
2756
2757 let limits = crate::chain_limits::compute_limits(&resolved, &family, &quant, 0);
2759 Json(limits).into_response()
2760}
2761
2762#[utoipa::path(
2770 post,
2771 path = "/api/shutdown",
2772 tag = "server",
2773 responses(
2774 (status = 200, description = "Shutdown initiated"),
2775 (status = 403, description = "Forbidden — remote shutdown requires API key auth"),
2776 )
2777)]
2778async fn shutdown_server(State(state): State<AppState>, request: Request) -> impl IntoResponse {
2779 let auth_enabled = request
2782 .extensions()
2783 .get::<crate::auth::AuthState>()
2784 .is_some_and(|s| s.is_some());
2785
2786 if !auth_enabled {
2787 let is_loopback = request
2788 .extensions()
2789 .get::<axum::extract::ConnectInfo<std::net::SocketAddr>>()
2790 .map(|ci| ci.0.ip().is_loopback())
2791 .unwrap_or(false);
2792 if !is_loopback {
2793 return (
2794 StatusCode::FORBIDDEN,
2795 "shutdown requires API key auth or localhost access\n",
2796 );
2797 }
2798 }
2799
2800 tracing::info!("shutdown requested via API");
2801 if let Some(tx) = state.shutdown_tx.lock().await.take() {
2802 let _ = tx.send(());
2803 }
2804 (StatusCode::OK, "shutdown initiated\n")
2805}
2806
2807#[derive(Debug, Deserialize, utoipa::ToSchema)]
2810pub struct GalleryMediaTokenRequest {
2811 path: String,
2812}
2813
2814#[derive(Debug, Serialize, utoipa::ToSchema)]
2815pub struct GalleryMediaTokenResponse {
2816 token: Option<String>,
2817 expires_at: Option<u64>,
2818 auth_required: bool,
2819}
2820
2821#[utoipa::path(
2830 post,
2831 path = "/api/gallery/media-token",
2832 tag = "server",
2833 request_body = GalleryMediaTokenRequest,
2834 responses(
2835 (status = 200, description = "Short-lived gallery media ticket", body = GalleryMediaTokenResponse),
2836 (status = 401, description = "API key authentication is required"),
2837 (status = 422, description = "Requested path is not a gallery media path"),
2838 )
2839)]
2840async fn create_gallery_media_token(
2841 auth_state: Option<Extension<crate::auth::AuthState>>,
2842 authenticated: Option<Extension<crate::auth::ApiKeyAuthenticated>>,
2843 Json(request): Json<GalleryMediaTokenRequest>,
2844) -> Result<impl IntoResponse, ApiError> {
2845 if !crate::auth::is_gallery_image_path(&request.path) {
2846 return Err(ApiError::validation(
2847 "media token path must match /api/gallery/image/:filename",
2848 ));
2849 }
2850
2851 let key_set = auth_state.and_then(|Extension(state)| state);
2852 let (token, expires_at, auth_required) = match key_set {
2853 Some(key_set) => {
2854 let _authenticated = authenticated.ok_or_else(|| {
2855 ApiError::with_code(
2856 "API key authentication is required to issue a media token",
2857 "UNAUTHORIZED",
2858 StatusCode::UNAUTHORIZED,
2859 )
2860 })?;
2861 let (token, expires_at) = key_set.issue_gallery_media_token(&request.path);
2862 (Some(token), Some(expires_at), true)
2863 }
2864 None => (None, None, false),
2867 };
2868
2869 let mut headers = HeaderMap::new();
2870 headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
2871 Ok((
2872 headers,
2873 Json(GalleryMediaTokenResponse {
2874 token,
2875 expires_at,
2876 auth_required,
2877 }),
2878 ))
2879}
2880
2881async fn list_gallery(
2889 State(state): State<AppState>,
2890) -> Result<Json<Vec<mold_core::GalleryImage>>, ApiError> {
2891 let config = state.config.read().await;
2892 if config.is_output_disabled() {
2893 return Ok(Json(Vec::new()));
2894 }
2895 let output_dir = config.effective_output_dir();
2896 drop(config);
2897
2898 if !output_dir.is_dir() {
2899 return Ok(Json(Vec::new()));
2900 }
2901
2902 if state.metadata_db.is_some() {
2903 let db_arc = state.metadata_db.clone();
2904 let dir = output_dir.clone();
2905 let listed = tokio::task::spawn_blocking(move || {
2906 db_arc
2907 .as_ref()
2908 .as_ref()
2909 .map(|db| db.list(Some(&dir)))
2910 .transpose()
2911 })
2912 .await
2913 .map_err(|e| ApiError::internal(format!("gallery DB query failed: {e}")))?
2914 .map_err(|e| ApiError::internal(format!("gallery DB query failed: {e:#}")))?;
2915 if let Some(rows) = listed {
2916 if !rows.is_empty() {
2917 let images = rows.iter().map(|r| r.to_gallery_image()).collect();
2918 return Ok(Json(images));
2919 }
2920 }
2921 }
2922
2923 let images = tokio::task::spawn_blocking(move || scan_gallery_dir(&output_dir))
2924 .await
2925 .map_err(|e| ApiError::internal(format!("gallery scan failed: {e}")))?;
2926
2927 Ok(Json(images))
2928}
2929
2930async fn get_gallery_image(
2940 State(state): State<AppState>,
2941 headers: HeaderMap,
2942 axum::extract::Path(filename): axum::extract::Path<String>,
2943) -> Result<axum::response::Response, ApiError> {
2944 let config = state.config.read().await;
2945 if config.is_output_disabled() {
2946 return Err(ApiError::not_found("image output is disabled"));
2947 }
2948 let output_dir = config.effective_output_dir();
2949 drop(config);
2950
2951 let clean_name = std::path::Path::new(&filename)
2953 .file_name()
2954 .map(|f| f.to_string_lossy().to_string())
2955 .unwrap_or_default();
2956
2957 if clean_name.is_empty() || clean_name != filename {
2958 return Err(ApiError::validation("invalid filename"));
2959 }
2960
2961 let path = output_dir.join(&clean_name);
2962 let meta = match tokio::fs::metadata(&path).await {
2963 Ok(m) if m.is_file() => m,
2964 _ => {
2965 return Err(ApiError::not_found(format!(
2966 "image not found: {clean_name}"
2967 )));
2968 }
2969 };
2970 let total_len = meta.len();
2971 let content_type = content_type_for_filename(&clean_name);
2972
2973 let range_header = headers
2974 .get(header::RANGE)
2975 .and_then(|v| v.to_str().ok())
2976 .map(|s| s.to_string());
2977
2978 let file = tokio::fs::File::open(&path)
2979 .await
2980 .map_err(|e| ApiError::internal(format!("failed to open file: {e}")))?;
2981
2982 if let Some(raw) = range_header {
2983 if let Some((start, end)) = parse_byte_range(&raw, total_len) {
2984 return serve_range(file, start, end, total_len, content_type).await;
2985 } else {
2986 return Ok(axum::response::Response::builder()
2988 .status(StatusCode::RANGE_NOT_SATISFIABLE)
2989 .header(header::CONTENT_RANGE, format!("bytes */{total_len}"))
2990 .body(axum::body::Body::empty())
2991 .unwrap());
2992 }
2993 }
2994
2995 let stream = tokio_util::io::ReaderStream::new(file);
2997 let body = axum::body::Body::from_stream(stream);
2998 Ok(axum::response::Response::builder()
2999 .status(StatusCode::OK)
3000 .header(header::CONTENT_TYPE, content_type)
3001 .header(header::ACCEPT_RANGES, "bytes")
3002 .header(header::CONTENT_LENGTH, total_len)
3003 .header(header::CACHE_CONTROL, "private, no-store")
3004 .body(body)
3005 .unwrap())
3006}
3007
3008fn parse_byte_range(header: &str, total_len: u64) -> Option<(u64, u64)> {
3016 let spec = header.strip_prefix("bytes=")?;
3017 if spec.contains(',') {
3018 return None;
3019 }
3020 let (start_s, end_s) = spec.split_once('-')?;
3021 let start_s = start_s.trim();
3022 let end_s = end_s.trim();
3023
3024 if total_len == 0 {
3025 return None;
3026 }
3027
3028 if start_s.is_empty() {
3029 let suffix: u64 = end_s.parse().ok()?;
3031 if suffix == 0 {
3032 return None;
3033 }
3034 let start = total_len.saturating_sub(suffix);
3035 return Some((start, total_len - 1));
3036 }
3037
3038 let start: u64 = start_s.parse().ok()?;
3039 if start >= total_len {
3040 return None;
3041 }
3042 let end: u64 = if end_s.is_empty() {
3043 total_len - 1
3044 } else {
3045 end_s.parse().ok()?
3046 };
3047 let end = end.min(total_len - 1);
3048 if end < start {
3049 return None;
3050 }
3051 Some((start, end))
3052}
3053
3054async fn serve_range(
3058 mut file: tokio::fs::File,
3059 start: u64,
3060 end: u64,
3061 total_len: u64,
3062 content_type: &'static str,
3063) -> Result<axum::response::Response, ApiError> {
3064 use tokio::io::{AsyncReadExt, AsyncSeekExt};
3065 file.seek(std::io::SeekFrom::Start(start))
3066 .await
3067 .map_err(|e| ApiError::internal(format!("seek failed: {e}")))?;
3068 let len = end - start + 1;
3069 let stream = tokio_util::io::ReaderStream::new(file.take(len));
3070 let body = axum::body::Body::from_stream(stream);
3071 Ok(axum::response::Response::builder()
3072 .status(StatusCode::PARTIAL_CONTENT)
3073 .header(header::CONTENT_TYPE, content_type)
3074 .header(header::ACCEPT_RANGES, "bytes")
3075 .header(header::CONTENT_LENGTH, len)
3076 .header(
3077 header::CONTENT_RANGE,
3078 format!("bytes {start}-{end}/{total_len}"),
3079 )
3080 .header(header::CACHE_CONTROL, "private, no-store")
3083 .body(body)
3084 .unwrap())
3085}
3086
3087fn content_type_for_filename(name: &str) -> &'static str {
3090 let lower = name.to_ascii_lowercase();
3091 if lower.ends_with(".png") {
3092 "image/png"
3093 } else if lower.ends_with(".jpg") || lower.ends_with(".jpeg") {
3094 "image/jpeg"
3095 } else if lower.ends_with(".gif") {
3096 "image/gif"
3097 } else if lower.ends_with(".webp") {
3098 "image/webp"
3099 } else if lower.ends_with(".apng") {
3100 "image/apng"
3101 } else if lower.ends_with(".mp4") {
3102 "video/mp4"
3103 } else {
3104 "application/octet-stream"
3105 }
3106}
3107
3108async fn delete_gallery_image(
3113 State(state): State<AppState>,
3114 axum::extract::Path(filename): axum::extract::Path<String>,
3115) -> Result<impl IntoResponse, ApiError> {
3116 let config = state.config.read().await;
3117 if config.is_output_disabled() {
3118 return Err(ApiError::not_found("image output is disabled"));
3119 }
3120 let output_dir = config.effective_output_dir();
3121 drop(config);
3122
3123 let clean_name = std::path::Path::new(&filename)
3124 .file_name()
3125 .map(|f| f.to_string_lossy().to_string())
3126 .unwrap_or_default();
3127
3128 if clean_name.is_empty() || clean_name != filename {
3129 return Err(ApiError::validation("invalid filename"));
3130 }
3131
3132 let db = state.metadata_db.clone();
3136 let name = clean_name.clone();
3137 let dir = output_dir.clone();
3138 tokio::task::spawn_blocking(move || -> Result<(), ApiError> {
3139 let path = dir.join(&name);
3140 if path.is_file() {
3141 std::fs::remove_file(&path)
3142 .map_err(|e| ApiError::internal(format!("failed to delete image: {e}")))?;
3143 }
3144
3145 let thumb_dir = server_thumbnail_dir();
3150 let _ = std::fs::remove_file(thumb_dir.join(&name));
3151 let _ = std::fs::remove_file(thumb_dir.join(format!("{name}.png")));
3152 let _ = std::fs::remove_file(
3153 server_preview_gif_dir().join(mold_core::media_paths::preview_gif_filename(&name)),
3154 );
3155
3156 if let Some(db) = db.as_ref().as_ref() {
3160 match db.delete(&dir, &name) {
3161 Ok(true) => {}
3162 Ok(false) => {
3163 tracing::debug!("delete: no metadata row for {}", dir.join(&name).display())
3164 }
3165 Err(e) => tracing::warn!(
3166 "metadata DB delete failed for {}: {e:#}",
3167 dir.join(&name).display()
3168 ),
3169 }
3170 }
3171 Ok(())
3172 })
3173 .await
3174 .map_err(|e| ApiError::internal(format!("gallery delete task failed: {e}")))??;
3175
3176 state
3177 .events
3178 .publish(mold_core::ServerEvent::GalleryRemoved {
3179 filename: clean_name,
3180 });
3181
3182 Ok(StatusCode::NO_CONTENT)
3183}
3184
3185async fn get_gallery_thumbnail(
3188 State(state): State<AppState>,
3189 axum::extract::Path(filename): axum::extract::Path<String>,
3190) -> Result<impl IntoResponse, ApiError> {
3191 let config = state.config.read().await;
3192 if config.is_output_disabled() {
3193 return Err(ApiError::not_found("image output is disabled"));
3194 }
3195 let output_dir = config.effective_output_dir();
3196 drop(config);
3197
3198 let clean_name = std::path::Path::new(&filename)
3199 .file_name()
3200 .map(|f| f.to_string_lossy().to_string())
3201 .unwrap_or_default();
3202
3203 if clean_name.is_empty() || clean_name != filename {
3204 return Err(ApiError::validation("invalid filename"));
3205 }
3206
3207 let source_path = output_dir.join(&clean_name);
3208 if !source_path.is_file() {
3209 return Err(ApiError::not_found(format!(
3210 "image not found: {clean_name}"
3211 )));
3212 }
3213
3214 let thumb_dir = server_thumbnail_dir();
3218 let thumb_path = thumb_dir.join(format!("{clean_name}.png"));
3219 let lower = clean_name.to_ascii_lowercase();
3220 let is_video = lower.ends_with(".mp4");
3221
3222 if !thumb_path.is_file() {
3223 let source = source_path.clone();
3230 let dest = thumb_path.clone();
3231 let gen_result = tokio::task::spawn_blocking(move || {
3232 if is_video {
3233 generate_video_thumbnail(&source, &dest)
3234 } else {
3235 generate_server_thumbnail(&source, &dest)
3236 }
3237 })
3238 .await
3239 .map_err(|e| ApiError::internal(format!("thumbnail generation failed: {e}")))?;
3240
3241 if let Err(err) = gen_result {
3242 tracing::warn!(
3243 file = %clean_name,
3244 error = %err,
3245 "thumbnail decode failed; falling back to source bytes"
3246 );
3247 if is_video {
3251 let mut headers = HeaderMap::new();
3252 headers.insert(
3253 header::CONTENT_TYPE,
3254 HeaderValue::from_static("image/svg+xml"),
3255 );
3256 headers.insert(
3257 header::CACHE_CONTROL,
3258 HeaderValue::from_static("public, max-age=300"),
3259 );
3260 return Ok((headers, VIDEO_PLACEHOLDER_SVG.as_bytes().to_vec()));
3261 }
3262 let raw = tokio::fs::read(&source_path)
3263 .await
3264 .map_err(|e| ApiError::internal(format!("failed to read source: {e}")))?;
3265 let mut headers = HeaderMap::new();
3266 headers.insert(
3267 header::CONTENT_TYPE,
3268 HeaderValue::from_static(content_type_for_filename(&clean_name)),
3269 );
3270 headers.insert(
3271 header::CACHE_CONTROL,
3272 HeaderValue::from_static("public, max-age=300"),
3273 );
3274 return Ok((headers, raw));
3275 }
3276 }
3277
3278 let data = tokio::fs::read(&thumb_path)
3279 .await
3280 .map_err(|e| ApiError::internal(format!("failed to read thumbnail: {e}")))?;
3281
3282 let mut headers = HeaderMap::new();
3283 headers.insert(header::CONTENT_TYPE, HeaderValue::from_static("image/png"));
3284 headers.insert(
3285 header::CACHE_CONTROL,
3286 HeaderValue::from_static("public, max-age=3600"),
3287 );
3288
3289 Ok((headers, data))
3290}
3291
3292const 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>"##;
3293
3294async fn get_gallery_preview(
3304 State(state): State<AppState>,
3305 axum::extract::Path(filename): axum::extract::Path<String>,
3306) -> Result<axum::response::Response, ApiError> {
3307 let config = state.config.read().await;
3308 if config.is_output_disabled() {
3309 return Err(ApiError::not_found("image output is disabled"));
3310 }
3311 let output_dir = config.effective_output_dir();
3312 drop(config);
3313
3314 let clean_name = std::path::Path::new(&filename)
3317 .file_name()
3318 .map(|f| f.to_string_lossy().to_string())
3319 .unwrap_or_default();
3320 if clean_name.is_empty() || clean_name != filename {
3321 return Err(ApiError::validation("invalid filename"));
3322 }
3323
3324 let source_path = output_dir.join(&clean_name);
3331 if !tokio::fs::metadata(&source_path)
3332 .await
3333 .map(|m| m.is_file())
3334 .unwrap_or(false)
3335 {
3336 return Err(ApiError::not_found(format!(
3337 "image not found: {clean_name}"
3338 )));
3339 }
3340
3341 let preview_path =
3342 server_preview_gif_dir().join(mold_core::media_paths::preview_gif_filename(&clean_name));
3343 let meta = match tokio::fs::metadata(&preview_path).await {
3344 Ok(m) if m.is_file() => m,
3345 _ => {
3346 return Err(ApiError::not_found(format!(
3347 "preview not found: {clean_name}"
3348 )));
3349 }
3350 };
3351 let total_len = meta.len();
3352
3353 let file = tokio::fs::File::open(&preview_path)
3354 .await
3355 .map_err(|e| ApiError::internal(format!("failed to open preview: {e}")))?;
3356 let stream = tokio_util::io::ReaderStream::new(file);
3357 let body = axum::body::Body::from_stream(stream);
3358 Ok(axum::response::Response::builder()
3359 .status(StatusCode::OK)
3360 .header(header::CONTENT_TYPE, "image/gif")
3361 .header(header::CONTENT_LENGTH, total_len)
3362 .header(header::CACHE_CONTROL, "public, max-age=3600")
3363 .body(body)
3364 .unwrap())
3365}
3366
3367fn server_preview_gif_dir() -> std::path::PathBuf {
3372 mold_core::Config::mold_dir()
3373 .unwrap_or_else(|| std::path::PathBuf::from(".mold"))
3374 .join("cache")
3375 .join("previews")
3376}
3377
3378fn server_thumbnail_dir() -> std::path::PathBuf {
3380 mold_core::Config::mold_dir()
3381 .unwrap_or_else(|| std::path::PathBuf::from(".mold"))
3382 .join("cache")
3383 .join("thumbnails")
3384}
3385
3386fn generate_server_thumbnail(
3390 source: &std::path::Path,
3391 dest: &std::path::Path,
3392) -> anyhow::Result<()> {
3393 let img = image::open(source)?;
3394 let thumb = img.thumbnail(256, 256);
3395 if let Some(parent) = dest.parent() {
3396 std::fs::create_dir_all(parent)?;
3397 }
3398 thumb.save_with_format(dest, image::ImageFormat::Png)?;
3399 Ok(())
3400}
3401
3402fn generate_video_thumbnail(
3410 source: &std::path::Path,
3411 dest: &std::path::Path,
3412) -> anyhow::Result<()> {
3413 if let Some(parent) = dest.parent() {
3414 std::fs::create_dir_all(parent)?;
3415 }
3416 let tmp = dest.with_extension("firstframe.png");
3421 mold_inference::ltx2::media::extract_thumbnail(source, &tmp)?;
3422 let decode_result = (|| -> anyhow::Result<()> {
3423 let img = image::open(&tmp)?;
3424 let thumb = img.thumbnail(256, 256);
3425 thumb.save_with_format(dest, image::ImageFormat::Png)?;
3426 Ok(())
3427 })();
3428 let _ = std::fs::remove_file(&tmp);
3429 decode_result
3430}
3431
3432pub fn spawn_thumbnail_warmup(config: &mold_core::Config) {
3434 if !thumbnail_warmup_enabled() {
3435 tracing::info!("thumbnail warmup disabled; thumbnails will be generated on demand");
3436 return;
3437 }
3438
3439 let output_dir = config.effective_output_dir();
3440 std::thread::spawn(move || {
3441 if !output_dir.is_dir() {
3442 return;
3443 }
3444 let thumb_dir = server_thumbnail_dir();
3445 let walker = walkdir::WalkDir::new(&output_dir).max_depth(1).into_iter();
3446 for entry in walker.filter_map(|e| e.ok()) {
3447 let path = entry.path();
3448 if !path.is_file() {
3449 continue;
3450 }
3451 let ext = path
3452 .extension()
3453 .and_then(|e| e.to_str())
3454 .map(|e| e.to_lowercase());
3455 let is_raster = matches!(
3456 ext.as_deref(),
3457 Some("png" | "jpg" | "jpeg" | "gif" | "apng" | "webp")
3458 );
3459 let is_video = matches!(ext.as_deref(), Some("mp4"));
3460 if !is_raster && !is_video {
3461 continue;
3462 }
3463 let filename = path
3464 .file_name()
3465 .map(|f| f.to_string_lossy().to_string())
3466 .unwrap_or_default();
3467 let thumb_path = thumb_dir.join(format!("{filename}.png"));
3468 if thumb_path.is_file() {
3469 continue;
3470 }
3471 let result = if is_video {
3472 generate_video_thumbnail(path, &thumb_path)
3473 } else {
3474 generate_server_thumbnail(path, &thumb_path)
3475 };
3476 if let Err(e) = result {
3477 tracing::warn!("failed to generate thumbnail for {}: {e}", path.display());
3478 }
3479 }
3480 tracing::info!("thumbnail warmup complete");
3481 });
3482}
3483
3484fn thumbnail_warmup_enabled() -> bool {
3485 std::env::var("MOLD_THUMBNAIL_WARMUP")
3486 .map(|v| matches!(v.as_str(), "1" | "true" | "TRUE" | "yes" | "YES"))
3487 .unwrap_or(false)
3488}
3489
3490fn scan_gallery_dir(dir: &std::path::Path) -> Vec<mold_core::GalleryImage> {
3509 let mut images: Vec<mold_core::GalleryImage> = mold_db::scan::scan_output_dir(dir)
3510 .filter_map(|item| match item {
3511 mold_db::scan::ScanItem::Valid(file) => Some(file),
3512 _ => None,
3513 })
3514 .map(|file| {
3515 let timestamp = file.timestamp_secs();
3516 let size_bytes = file.size_u64();
3517 let (metadata, synthetic) = mold_db::metadata_io::read_or_synthesize(
3518 &file.path,
3519 file.format,
3520 &file.filename,
3521 timestamp,
3522 );
3523 mold_core::GalleryImage {
3524 filename: file.filename,
3525 metadata,
3526 timestamp,
3527 format: Some(file.format),
3528 size_bytes: Some(size_bytes),
3529 metadata_synthetic: synthetic,
3530 }
3531 })
3532 .collect();
3533
3534 images.sort_by_key(|img| std::cmp::Reverse(img.timestamp));
3535 images
3536}
3537
3538async fn get_model_placement(
3546 State(state): State<AppState>,
3547 axum::extract::Path(name): axum::extract::Path<String>,
3548) -> Result<Json<mold_core::types::DevicePlacement>, ApiError> {
3549 let cfg = state.config.read().await;
3550 match cfg.models.get(&name).and_then(|mc| mc.placement.clone()) {
3551 Some(placement) => Ok(Json(placement)),
3552 None => Err(ApiError::not_found(format!(
3553 "no placement saved for model '{name}'"
3554 ))),
3555 }
3556}
3557
3558async fn put_model_placement(
3559 State(state): State<AppState>,
3560 axum::extract::Path(name): axum::extract::Path<String>,
3561 Json(placement): Json<mold_core::types::DevicePlacement>,
3562) -> Result<Json<serde_json::Value>, ApiError> {
3563 validate_multi_gpu_placement(&state, Some(&placement))?;
3564 {
3565 let mut cfg = state.config.write().await;
3566 cfg.set_model_placement(&name, Some(placement.clone()));
3567 cfg.save().map_err(|e| {
3568 tracing::warn!("failed to persist placement to config.toml: {e}");
3569 ApiError::internal(format!("failed to persist placement to config.toml: {e}"))
3570 })?;
3571 }
3572 Ok(Json(serde_json::json!({
3573 "ok": true,
3574 "model": name,
3575 })))
3576}
3577
3578async fn delete_model_placement(
3579 State(state): State<AppState>,
3580 axum::extract::Path(name): axum::extract::Path<String>,
3581) -> Result<Json<serde_json::Value>, ApiError> {
3582 let mut cfg = state.config.write().await;
3583 cfg.set_model_placement(&name, None);
3584 cfg.save().map_err(|e| {
3585 tracing::warn!("failed to persist placement removal to config.toml: {e}");
3586 ApiError::internal(format!(
3587 "failed to persist placement removal to config.toml: {e}"
3588 ))
3589 })?;
3590 Ok(Json(serde_json::json!({ "ok": true })))
3591}
3592
3593async fn openapi_json() -> impl IntoResponse {
3596 Json(ApiDoc::openapi())
3597}
3598
3599async fn scalar_docs() -> impl IntoResponse {
3602 (
3603 [(header::CONTENT_TYPE, "text/html")],
3604 r#"<!DOCTYPE html>
3605<html>
3606<head>
3607 <title>mold API</title>
3608 <meta charset="utf-8" />
3609 <meta name="viewport" content="width=device-width, initial-scale=1" />
3610</head>
3611<body>
3612 <script id="api-reference" data-url="/api/openapi.json"></script>
3613 <script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference"></script>
3614</body>
3615</html>"#,
3616 )
3617}
3618
3619fn query_gpu_info() -> Option<GpuInfo> {
3622 query_cuda_gpu_info().or_else(query_metal_gpu_info)
3623}
3624
3625fn query_cuda_gpu_info() -> Option<GpuInfo> {
3626 let nvidia_smi = if std::path::Path::new("/run/current-system/sw/bin/nvidia-smi").exists() {
3627 "/run/current-system/sw/bin/nvidia-smi"
3628 } else {
3629 "nvidia-smi"
3630 };
3631
3632 let output = std::process::Command::new(nvidia_smi)
3633 .args([
3634 "--query-gpu=name,memory.total,memory.used",
3635 "--format=csv,noheader,nounits",
3636 ])
3637 .output()
3638 .ok()?;
3639
3640 if !output.status.success() {
3641 return None;
3642 }
3643
3644 let text = String::from_utf8(output.stdout).ok()?;
3645 let line = text.lines().next()?;
3646 let parts: Vec<&str> = line.split(',').map(str::trim).collect();
3647 if parts.len() < 3 {
3648 return None;
3649 }
3650
3651 Some(GpuInfo {
3652 name: parts[0].to_string(),
3653 vram_total_mb: parts[1].parse().ok()?,
3654 vram_used_mb: parts[2].parse().ok()?,
3655 backend: Some(GpuBackend::Cuda),
3656 })
3657}
3658
3659fn query_metal_gpu_info() -> Option<GpuInfo> {
3664 let gpu = crate::resources::metal_snapshot().into_iter().next()?;
3665 Some(GpuInfo {
3666 name: gpu.name,
3667 vram_total_mb: gpu.vram_total / 1_000_000,
3670 vram_used_mb: gpu.vram_used / 1_000_000,
3671 backend: Some(GpuBackend::Metal),
3672 })
3673}
3674
3675#[derive(serde::Deserialize, utoipa::ToSchema)]
3678pub struct CreateDownloadBody {
3679 pub model: String,
3680}
3681
3682#[derive(serde::Serialize, utoipa::ToSchema)]
3683pub struct CreateDownloadResponse {
3684 pub id: String,
3685 pub position: usize,
3686}
3687
3688#[utoipa::path(
3689 post,
3690 path = "/api/downloads",
3691 tag = "downloads",
3692 request_body = CreateDownloadBody,
3693 responses(
3694 (status = 200, description = "Enqueued; position 0 = will start immediately", body = CreateDownloadResponse),
3695 (status = 400, description = "Unknown model"),
3696 (status = 409, description = "Already active or queued; body contains existing id", body = CreateDownloadResponse),
3697 )
3698)]
3699pub async fn create_download(
3700 State(state): State<AppState>,
3701 Json(body): Json<CreateDownloadBody>,
3702) -> axum::response::Response {
3703 use crate::downloads::{EnqueueError, EnqueueOutcome};
3704 match state.downloads.enqueue(body.model.clone()).await {
3705 Ok((id, position, EnqueueOutcome::Created)) => (
3706 StatusCode::OK,
3707 Json(CreateDownloadResponse { id, position }),
3708 )
3709 .into_response(),
3710 Ok((id, position, EnqueueOutcome::AlreadyPresent)) => (
3711 StatusCode::CONFLICT,
3712 Json(CreateDownloadResponse { id, position }),
3713 )
3714 .into_response(),
3715 Err(EnqueueError::UnknownModel(_)) => (
3716 StatusCode::BAD_REQUEST,
3717 Json(serde_json::json!({
3718 "error": format!("unknown model '{}'. Run 'mold list' to see available models.", body.model)
3719 })),
3720 )
3721 .into_response(),
3722 Err(EnqueueError::LockPoisoned) => (
3723 StatusCode::INTERNAL_SERVER_ERROR,
3724 Json(serde_json::json!({ "error": "download queue state is corrupt" })),
3725 )
3726 .into_response(),
3727 }
3728}
3729
3730#[utoipa::path(
3731 delete,
3732 path = "/api/downloads/{id}",
3733 tag = "downloads",
3734 params(("id" = String, Path, description = "Job id")),
3735 responses(
3736 (status = 204, description = "Cancelled"),
3737 (status = 404, description = "Unknown id"),
3738 )
3739)]
3740pub async fn delete_download(
3741 State(state): State<AppState>,
3742 axum::extract::Path(id): axum::extract::Path<String>,
3743) -> axum::response::Response {
3744 if state.downloads.cancel(&id).await {
3745 StatusCode::NO_CONTENT.into_response()
3746 } else {
3747 (
3748 StatusCode::NOT_FOUND,
3749 Json(serde_json::json!({ "error": format!("unknown download id '{id}'") })),
3750 )
3751 .into_response()
3752 }
3753}
3754
3755#[utoipa::path(
3756 get,
3757 path = "/api/downloads",
3758 tag = "downloads",
3759 responses((status = 200, description = "Current queue state"))
3760)]
3761pub async fn list_downloads(State(state): State<AppState>) -> axum::response::Response {
3762 Json(state.downloads.listing().await).into_response()
3763}
3764
3765#[utoipa::path(
3766 get,
3767 path = "/api/downloads/stream",
3768 tag = "downloads",
3769 responses((status = 200, description = "SSE stream of DownloadEvent JSON")),
3770)]
3771pub async fn stream_downloads(
3772 State(state): State<AppState>,
3773) -> Sse<
3774 impl futures_core::Stream<Item = Result<axum::response::sse::Event, std::convert::Infallible>>,
3775> {
3776 use axum::response::sse::Event;
3777 use tokio_stream::wrappers::BroadcastStream;
3778 use tokio_stream::StreamExt as _;
3779
3780 let rx = state.downloads.subscribe();
3787 let initial = state.downloads.listing().await;
3788 let snapshot_event = mold_core::types::DownloadEvent::Snapshot { listing: initial };
3789
3790 let stream = async_stream::stream! {
3791 let data = serde_json::to_string(&snapshot_event).unwrap_or_else(|_| "{}".to_string());
3792 yield Ok::<_, std::convert::Infallible>(Event::default().event("download").data(data));
3793
3794 let mut bs = BroadcastStream::new(rx);
3795 while let Some(item) = bs.next().await {
3796 match item {
3797 Ok(event) => {
3798 let data = serde_json::to_string(&event)
3799 .unwrap_or_else(|_| "{}".to_string());
3800 yield Ok(Event::default().event("download").data(data));
3801 }
3802 Err(_lagged) => continue,
3806 }
3807 }
3808 };
3809
3810 Sse::new(stream).keep_alive(
3811 KeepAlive::new()
3812 .interval(std::time::Duration::from_secs(15))
3813 .text("ping"),
3814 )
3815}
3816
3817async fn get_resources(State(state): State<AppState>) -> Result<Json<ResourceSnapshot>, ApiError> {
3823 match state.resources.latest() {
3824 Some(snap) => Ok(Json(snap)),
3825 None => Err(ApiError::internal_with_status(
3826 "resource telemetry not ready",
3827 StatusCode::SERVICE_UNAVAILABLE,
3828 )),
3829 }
3830}
3831
3832async fn get_resources_stream(
3835 State(state): State<AppState>,
3836) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>> {
3837 use tokio_stream::wrappers::BroadcastStream;
3838
3839 let rx = state.resources.subscribe();
3840 let initial = state.resources.latest();
3843
3844 let stream = async_stream::stream! {
3845 if let Some(snap) = initial {
3846 yield Ok::<_, Infallible>(snapshot_to_sse(&snap));
3847 }
3848 let mut bs = BroadcastStream::new(rx);
3849 while let Some(item) = bs.next().await {
3850 match item {
3851 Ok(snap) => yield Ok(snapshot_to_sse(&snap)),
3852 Err(_lagged) => continue,
3855 }
3856 }
3857 };
3858
3859 Sse::new(stream).keep_alive(
3860 KeepAlive::new()
3861 .interval(std::time::Duration::from_secs(15))
3862 .text("ping"),
3863 )
3864}
3865
3866fn snapshot_to_sse(snap: &ResourceSnapshot) -> SseEvent {
3867 match serde_json::to_string(snap) {
3868 Ok(data) => SseEvent::default().event("snapshot").data(data),
3869 Err(e) => SseEvent::default()
3870 .event("error")
3871 .data(format!("{{\"message\":\"serialize failed: {e}\"}}")),
3872 }
3873}
3874
3875#[utoipa::path(
3883 get,
3884 path = "/api/events",
3885 tag = "server",
3886 responses(
3887 (status = 200, description = "SSE stream of server lifecycle events", content_type = "text/event-stream")
3888 )
3889)]
3890async fn stream_events(
3891 State(state): State<AppState>,
3892) -> Sse<impl futures_core::Stream<Item = Result<SseEvent, Infallible>>> {
3893 use tokio_stream::wrappers::BroadcastStream;
3894
3895 let rx = state.events.subscribe();
3896 let stream = async_stream::stream! {
3897 let mut bs = BroadcastStream::new(rx);
3898 while let Some(item) = bs.next().await {
3899 match item {
3900 Ok(ev) => yield Ok::<_, Infallible>(server_event_to_sse(&ev)),
3901 Err(_lagged) => continue,
3904 }
3905 }
3906 };
3907
3908 Sse::new(stream).keep_alive(
3909 KeepAlive::new()
3910 .interval(std::time::Duration::from_secs(15))
3911 .text("ping"),
3912 )
3913}
3914
3915fn server_event_to_sse(ev: &mold_core::ServerEvent) -> SseEvent {
3916 match serde_json::to_string(ev) {
3917 Ok(data) => SseEvent::default().event("event").data(data),
3918 Err(e) => SseEvent::default()
3921 .event("error")
3922 .data(serde_json::json!({ "message": format!("serialize failed: {e}") }).to_string()),
3923 }
3924}
3925
3926#[cfg(test)]
3927mod tests {
3928 use super::*;
3929
3930 #[test]
3931 fn expand_config_for_request_threads_style() {
3932 let settings = mold_core::expand::ExpandSettings::default();
3933 let req = mold_core::ExpandRequest {
3934 prompt: "a cat".to_string(),
3935 model_family: "flux".to_string(),
3936 variations: 2,
3937 style: Some("oil painting".to_string()),
3938 };
3939 let config = expand_config_for_request(&settings, &req);
3940 assert_eq!(config.style.as_deref(), Some("oil painting"));
3941 assert_eq!(config.model_family, "flux");
3942 assert_eq!(config.variations, 2);
3943
3944 let bare = mold_core::ExpandRequest {
3945 prompt: "a cat".to_string(),
3946 model_family: "flux".to_string(),
3947 variations: 1,
3948 style: None,
3949 };
3950 assert!(expand_config_for_request(&settings, &bare).style.is_none());
3951 }
3952
3953 #[test]
3954 fn local_expand_capability_reports_feature_and_model_facts_separately() {
3955 let settings = mold_core::expand::ExpandSettings::default();
3956 let missing = expand_capabilities(&settings, false);
3957 let present = expand_capabilities(&settings, true);
3958
3959 assert_eq!(missing.backend, mold_core::ExpandBackend::Local);
3960 assert_eq!(missing.configured, cfg!(feature = "expand"));
3961 assert_eq!(missing.model_present, Some(false));
3962 assert_eq!(present.configured, cfg!(feature = "expand"));
3963 assert_eq!(present.model_present, Some(true));
3964 }
3965
3966 #[test]
3967 fn api_expand_capability_does_not_claim_external_reachability() {
3968 let settings = mold_core::expand::ExpandSettings {
3969 backend: "http://localhost:11434".into(),
3970 ..Default::default()
3971 };
3972 let capability = expand_capabilities(&settings, false);
3973
3974 assert_eq!(capability.backend, mold_core::ExpandBackend::Api);
3975 assert!(capability.configured);
3976 assert_eq!(capability.model_present, None);
3977
3978 let unconfigured = mold_core::expand::ExpandSettings {
3979 backend: " ".into(),
3980 ..Default::default()
3981 };
3982 let capability = expand_capabilities(&unconfigured, true);
3983 assert_eq!(capability.backend, mold_core::ExpandBackend::Api);
3984 assert!(!capability.configured);
3985 assert_eq!(capability.model_present, None);
3986 }
3987
3988 fn env_lock() -> &'static std::sync::Mutex<()> {
3989 static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
3990 &ENV_LOCK
3991 }
3992
3993 #[test]
3994 fn disk_usage_for_path_picks_longest_matching_mount() {
3995 let disks = vec![
3996 (std::path::PathBuf::from("/"), 100, 10),
3997 (std::path::PathBuf::from("/data"), 500, 50),
3998 ];
3999 let usage = disk_usage_for_path(&disks, std::path::Path::new("/data/models")).unwrap();
4000 assert_eq!(usage.total_bytes, 500);
4001 assert_eq!(usage.free_bytes, 50);
4002 let usage = disk_usage_for_path(&disks, std::path::Path::new("/home/u/models")).unwrap();
4004 assert_eq!(usage.total_bytes, 100);
4005 }
4006
4007 #[test]
4008 fn disk_usage_for_path_component_boundaries_and_no_match() {
4009 let disks = vec![(std::path::PathBuf::from("/data"), 500, 50)];
4010 assert_eq!(
4013 disk_usage_for_path(&disks, std::path::Path::new("/database/models")),
4014 None
4015 );
4016 assert_eq!(
4018 disk_usage_for_path(&disks, std::path::Path::new("relative/models")),
4019 None
4020 );
4021 }
4022
4023 #[cfg(unix)]
4024 #[test]
4025 fn canonical_models_dir_resolves_symlinks_for_mount_matching() {
4026 let tmp = tempfile::tempdir().unwrap();
4031 let root = std::fs::canonicalize(tmp.path()).unwrap();
4034 let target = root.join("big-volume").join("models");
4035 std::fs::create_dir_all(&target).unwrap();
4036 let link = root.join("link-models");
4037 std::os::unix::fs::symlink(&target, &link).unwrap();
4038
4039 let resolved = canonical_models_dir(&link);
4040 assert_eq!(resolved, target);
4041
4042 let disks = vec![
4045 (std::path::PathBuf::from("/"), 100, 10),
4046 (root.join("big-volume"), 500, 50),
4047 ];
4048 let usage = disk_usage_for_path(&disks, &resolved).unwrap();
4049 assert_eq!(usage.total_bytes, 500);
4050 assert_eq!(usage.free_bytes, 50);
4051 }
4052
4053 #[test]
4054 fn canonical_models_dir_falls_back_to_the_literal_path() {
4055 let missing = std::path::Path::new("/definitely/not/a/real/mold/models/dir");
4058 assert_eq!(canonical_models_dir(missing), missing.to_path_buf());
4059 }
4060
4061 #[test]
4062 fn clean_error_message_strips_backtrace() {
4063 let err = anyhow::anyhow!(
4064 "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")\n\
4065 \x20 0: candle_core::error::Error::bt\n\
4066 \x20 at /home/user/.cargo/git/candle/src/error.rs:264:25\n\
4067 \x20 1: <core::result::Result<O,E> as candle_core::cuda_backend::error::WrapErr<O>>::w\n\
4068 \x20 at /home/user/.cargo/git/candle/src/cuda_backend/error.rs:60:65"
4069 );
4070 let msg = clean_error_message(&err);
4071 assert_eq!(
4072 msg,
4073 "DriverError(CUDA_ERROR_OUT_OF_MEMORY, \"out of memory\")"
4074 );
4075 }
4076
4077 #[test]
4078 fn clean_error_message_preserves_simple_error() {
4079 let err = anyhow::anyhow!("model not found: flux-dev:q4");
4080 let msg = clean_error_message(&err);
4081 assert_eq!(msg, "model not found: flux-dev:q4");
4082 }
4083
4084 #[test]
4085 fn clean_error_message_preserves_multiline_without_backtrace() {
4086 let err = anyhow::anyhow!("validation failed\nprompt is empty\nsteps must be > 0");
4087 let msg = clean_error_message(&err);
4088 assert_eq!(msg, "validation failed\nprompt is empty\nsteps must be > 0");
4089 }
4090
4091 #[test]
4092 fn clean_error_message_strips_high_numbered_frames() {
4093 let err = anyhow::anyhow!(
4094 "some error\n\
4095 \x20 10: tokio::runtime::task::core::Core<T,S>::poll at /home/user/.cargo/tokio/src/core.rs:375\n\
4096 \x20 11: std::panicking::catch_unwind at /nix/store/rust/src/panicking.rs:544"
4097 );
4098 let msg = clean_error_message(&err);
4099 assert_eq!(msg, "some error");
4100 }
4101
4102 #[test]
4103 fn clean_error_message_empty_fallback() {
4104 let err = anyhow::anyhow!("0: candle_core::error::Error::bt at /some/path.rs:10:5");
4106 let msg = clean_error_message(&err);
4107 assert!(!msg.is_empty());
4109 }
4110
4111 #[test]
4112 fn clean_error_message_renders_full_anyhow_chain() {
4113 let root = std::io::Error::new(std::io::ErrorKind::InvalidData, "bytes past end");
4118 let err: anyhow::Error = anyhow::Error::new(root)
4119 .context("validate single-file checkpoint at /tmp/foo.safetensors");
4120 let msg = clean_error_message(&err);
4121 assert!(
4122 msg.contains("validate single-file checkpoint") && msg.contains("bytes past end"),
4123 "expected both context layers in the rendered chain, got: {msg}",
4124 );
4125 }
4126
4127 #[test]
4128 fn save_image_to_dir_creates_directory_and_writes_file() {
4129 let dir = std::env::temp_dir().join(format!(
4130 "mold-save-test-{}",
4131 std::time::SystemTime::now()
4132 .duration_since(std::time::UNIX_EPOCH)
4133 .unwrap()
4134 .as_nanos()
4135 ));
4136 assert!(!dir.exists());
4137
4138 let img = mold_core::ImageData {
4139 data: vec![0x89, 0x50, 0x4E, 0x47], format: mold_core::OutputFormat::Png,
4141 width: 64,
4142 height: 64,
4143 index: 0,
4144 };
4145
4146 save_image_to_dir(&dir, &img, "test-model:q8", 1);
4147
4148 assert!(dir.exists(), "directory should be created");
4149 let files: Vec<_> = std::fs::read_dir(&dir).unwrap().collect();
4150 assert_eq!(files.len(), 1, "should have exactly one file");
4151 let file = files[0].as_ref().unwrap();
4152 let filename = file.file_name().to_str().unwrap().to_string();
4153 assert!(filename.starts_with("mold-test-model-q8-"), "{filename}");
4154 assert!(filename.ends_with(".png"), "{filename}");
4155 let contents = std::fs::read(file.path()).unwrap();
4156 assert_eq!(contents, vec![0x89, 0x50, 0x4E, 0x47]);
4157
4158 std::fs::remove_dir_all(&dir).ok();
4159 }
4160
4161 #[test]
4162 fn save_image_to_dir_batch_includes_index() {
4163 let dir = std::env::temp_dir().join(format!(
4164 "mold-save-batch-{}",
4165 std::time::SystemTime::now()
4166 .duration_since(std::time::UNIX_EPOCH)
4167 .unwrap()
4168 .as_nanos()
4169 ));
4170
4171 let img = mold_core::ImageData {
4172 data: vec![0xFF, 0xD8], format: mold_core::OutputFormat::Jpeg,
4174 width: 64,
4175 height: 64,
4176 index: 2,
4177 };
4178
4179 save_image_to_dir(&dir, &img, "flux-dev", 4);
4180
4181 let files: Vec<_> = std::fs::read_dir(&dir).unwrap().collect();
4182 assert_eq!(files.len(), 1);
4183 let filename = files[0]
4184 .as_ref()
4185 .unwrap()
4186 .file_name()
4187 .to_str()
4188 .unwrap()
4189 .to_string();
4190 assert!(
4191 filename.contains("-2.jpeg"),
4192 "batch index in name: {filename}"
4193 );
4194
4195 std::fs::remove_dir_all(&dir).ok();
4196 }
4197
4198 #[test]
4199 fn save_image_to_dir_invalid_path_logs_warning_no_panic() {
4200 let img = mold_core::ImageData {
4202 data: vec![0x00],
4203 format: mold_core::OutputFormat::Png,
4204 width: 1,
4205 height: 1,
4206 index: 0,
4207 };
4208 save_image_to_dir(
4210 std::path::Path::new("/dev/null/impossible"),
4211 &img,
4212 "test",
4213 1,
4214 );
4215 }
4217
4218 #[test]
4219 fn thumbnail_warmup_is_disabled_by_default() {
4220 let _guard = env_lock().lock().unwrap();
4221 unsafe {
4222 std::env::remove_var("MOLD_THUMBNAIL_WARMUP");
4223 }
4224 assert!(!thumbnail_warmup_enabled());
4225 }
4226
4227 #[test]
4228 fn thumbnail_warmup_accepts_truthy_env_values() {
4229 let _guard = env_lock().lock().unwrap();
4230 unsafe {
4231 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "1");
4232 }
4233 assert!(thumbnail_warmup_enabled());
4234 unsafe {
4235 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "true");
4236 }
4237 assert!(thumbnail_warmup_enabled());
4238 unsafe {
4239 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "YES");
4240 }
4241 assert!(thumbnail_warmup_enabled());
4242 unsafe {
4243 std::env::remove_var("MOLD_THUMBNAIL_WARMUP");
4244 }
4245 }
4246
4247 #[test]
4248 fn thumbnail_warmup_rejects_falsey_env_values() {
4249 let _guard = env_lock().lock().unwrap();
4250 unsafe {
4251 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "0");
4252 }
4253 assert!(!thumbnail_warmup_enabled());
4254 unsafe {
4255 std::env::set_var("MOLD_THUMBNAIL_WARMUP", "false");
4256 }
4257 assert!(!thumbnail_warmup_enabled());
4258 unsafe {
4259 std::env::remove_var("MOLD_THUMBNAIL_WARMUP");
4260 }
4261 }
4262
4263 #[test]
4264 fn content_type_covers_every_output_format() {
4265 assert_eq!(content_type_for_filename("a.png"), "image/png");
4266 assert_eq!(content_type_for_filename("a.PNG"), "image/png");
4267 assert_eq!(content_type_for_filename("a.jpg"), "image/jpeg");
4268 assert_eq!(content_type_for_filename("a.jpeg"), "image/jpeg");
4269 assert_eq!(content_type_for_filename("a.gif"), "image/gif");
4270 assert_eq!(content_type_for_filename("a.webp"), "image/webp");
4271 assert_eq!(content_type_for_filename("a.apng"), "image/apng");
4272 assert_eq!(content_type_for_filename("a.mp4"), "video/mp4");
4273 assert_eq!(
4274 content_type_for_filename("a.unknown"),
4275 "application/octet-stream"
4276 );
4277 }
4278 struct TempDir(std::path::PathBuf);
4287 impl TempDir {
4288 fn new(tag: &str) -> Self {
4289 let mut p = std::env::temp_dir();
4290 p.push(format!("mold-gallery-test-{tag}-{}", uuid::Uuid::new_v4()));
4291 std::fs::create_dir_all(&p).expect("create tempdir");
4292 Self(p)
4293 }
4294 fn path(&self) -> &std::path::Path {
4295 &self.0
4296 }
4297 }
4298 impl Drop for TempDir {
4299 fn drop(&mut self) {
4300 let _ = std::fs::remove_dir_all(&self.0);
4301 }
4302 }
4303
4304 fn make_png_bytes(width: u32, height: u32) -> Vec<u8> {
4309 let img = image::RgbImage::from_fn(width, height, |x, y| {
4310 let n = (x.wrapping_mul(37) ^ y.wrapping_mul(131)) as u8;
4311 image::Rgb([n, n.wrapping_add(85), n.wrapping_sub(17)])
4312 });
4313 let mut buf = Vec::new();
4314 image::DynamicImage::ImageRgb8(img)
4315 .write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
4316 .expect("encode png");
4317 buf
4318 }
4319
4320 #[test]
4321 fn scan_gallery_dir_filters_invalid_and_keeps_valid() {
4322 let td = TempDir::new("scan");
4323 let dir = td.path();
4324
4325 std::fs::write(dir.join("mold-model-1000.png"), make_png_bytes(32, 32)).unwrap();
4327
4328 let mut junk = vec![0u8; 512];
4330 junk[..4].copy_from_slice(b"JUNK");
4331 std::fs::write(dir.join("mold-broken-2000.png"), &junk).unwrap();
4332
4333 std::fs::write(
4335 dir.join("mold-tiny-3000.png"),
4336 b"\x89PNG\r\n\x1a\n", )
4338 .unwrap();
4339
4340 let mut mp4 = Vec::new();
4342 mp4.extend_from_slice(&[0x00, 0x00, 0x00, 0x20]);
4343 mp4.extend_from_slice(b"ftyp");
4344 mp4.extend_from_slice(b"isom\x00\x00\x02\x00");
4345 mp4.resize(8192, 0);
4348 std::fs::write(dir.join("mold-ltx-4000.mp4"), &mp4).unwrap();
4349
4350 let bad_mp4 = vec![0u8; 8192];
4352 std::fs::write(dir.join("mold-no-ftyp-5000.mp4"), &bad_mp4).unwrap();
4353
4354 std::fs::write(dir.join("random.txt"), b"not an output").unwrap();
4356
4357 let results = scan_gallery_dir(dir);
4358 let names: Vec<&str> = results.iter().map(|i| i.filename.as_str()).collect();
4359 assert!(
4360 names.contains(&"mold-model-1000.png"),
4361 "valid PNG should survive: {names:?}"
4362 );
4363 assert!(
4364 names.contains(&"mold-ltx-4000.mp4"),
4365 "valid MP4 with ftyp should survive: {names:?}"
4366 );
4367 assert!(
4368 !names.contains(&"mold-broken-2000.png"),
4369 "PNG with no valid header should be filtered: {names:?}"
4370 );
4371 assert!(
4372 !names.contains(&"mold-tiny-3000.png"),
4373 "under-size PNG stub should be filtered: {names:?}"
4374 );
4375 assert!(
4376 !names.contains(&"mold-no-ftyp-5000.mp4"),
4377 "MP4 without ftyp should be filtered: {names:?}"
4378 );
4379 assert_eq!(names.len(), 2, "only the 2 valid fixtures remain");
4380 }
4381
4382 #[test]
4383 fn solid_black_png_is_filtered_at_scan_time() {
4384 let td = TempDir::new("black");
4385 let dir = td.path();
4386
4387 let black = image::RgbImage::from_pixel(256, 256, image::Rgb([0, 0, 0]));
4391 let mut buf = Vec::new();
4392 image::DynamicImage::ImageRgb8(black)
4393 .write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png)
4394 .unwrap();
4395 std::fs::write(dir.join("mold-noisy-1000.png"), &buf).unwrap();
4396
4397 std::fs::write(dir.join("mold-valid-2000.png"), make_png_bytes(256, 256)).unwrap();
4399
4400 let results = scan_gallery_dir(dir);
4401 let names: Vec<&str> = results.iter().map(|i| i.filename.as_str()).collect();
4402 assert!(
4403 !names.contains(&"mold-noisy-1000.png"),
4404 "solid-black PNG should be filtered: {names:?}"
4405 );
4406 assert!(
4407 names.contains(&"mold-valid-2000.png"),
4408 "noisy PNG should survive: {names:?}"
4409 );
4410 }
4411 #[test]
4412 fn parse_byte_range_handles_common_forms() {
4413 assert_eq!(parse_byte_range("bytes=0-499", 2000), Some((0, 499)));
4415 assert_eq!(parse_byte_range("bytes=100-", 2000), Some((100, 1999)));
4417 assert_eq!(parse_byte_range("bytes=-500", 2000), Some((1500, 1999)));
4419 assert_eq!(parse_byte_range("bytes=0-9999", 2000), Some((0, 1999)));
4421 assert_eq!(parse_byte_range("bytes=0-1999", 2000), Some((0, 1999)));
4423 }
4424
4425 #[test]
4426 fn parse_byte_range_rejects_malformed_and_unsatisfiable() {
4427 assert_eq!(parse_byte_range("bytes=", 1000), None);
4428 assert_eq!(parse_byte_range("bytes=abc-100", 1000), None);
4429 assert_eq!(parse_byte_range("bytes=2000-", 1000), None);
4431 assert_eq!(parse_byte_range("bytes=500-100", 1000), None);
4433 assert_eq!(parse_byte_range("bytes=0-10,20-30", 1000), None);
4435 assert_eq!(parse_byte_range("bytes=-0", 1000), None);
4437 assert_eq!(parse_byte_range("bytes=0-10", 0), None);
4439 assert_eq!(parse_byte_range("items=0-10", 1000), None);
4441 }
4442
4443 #[test]
4444 fn scan_populates_real_dimensions_for_synthesized_metadata() {
4445 let td = TempDir::new("dims");
4449 let dir = td.path();
4450 std::fs::write(dir.join("mold-nometa-1000.png"), make_png_bytes(128, 96)).unwrap();
4451
4452 let results = scan_gallery_dir(dir);
4453 assert_eq!(results.len(), 1);
4454 let entry = &results[0];
4455 assert!(entry.metadata_synthetic);
4456 assert_eq!(entry.metadata.width, 128);
4457 assert_eq!(entry.metadata.height, 96);
4458 }
4459}