Skip to main content

mold_server/
routes.rs

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