Skip to main content

mold_server/
lib.rs

1pub mod auth;
2pub mod catalog_api;
3pub mod chain_job_runner;
4pub mod chain_limits;
5pub mod test_support;
6// Agent A (downloads)
7pub mod downloads;
8pub mod gpu_pool;
9pub mod gpu_worker;
10pub mod job_registry;
11pub mod logging;
12mod memory_preflight;
13#[cfg(feature = "metrics")]
14pub mod metrics;
15pub mod model_cache;
16pub mod model_manager;
17pub mod queue;
18pub mod rate_limit;
19pub mod request_id;
20pub mod resources;
21pub mod routes;
22pub mod routes_chain;
23pub mod routes_chain_jobs;
24mod signals;
25pub mod state;
26pub mod web_ui;
27
28#[cfg(all(test, feature = "metrics"))]
29mod metrics_test;
30#[cfg(test)]
31mod resources_test;
32#[cfg(test)]
33mod routes_test;
34
35use anyhow::Result;
36use axum::{extract::DefaultBodyLimit, middleware};
37use mold_core::types::GpuSelection;
38use mold_core::{Config, ModelPaths};
39use std::net::SocketAddr;
40use std::path::PathBuf;
41use std::sync::atomic::AtomicUsize;
42use tokio::net::TcpListener;
43use tower_http::cors::CorsLayer;
44use tower_http::trace::TraceLayer;
45use tracing::info;
46
47use state::QueueHandle;
48
49const MAX_REQUEST_BODY_BYTES: usize = 64 * 1024 * 1024;
50
51pub async fn run_server(
52    bind: &str,
53    port: u16,
54    models_dir: PathBuf,
55    gpu_selection: GpuSelection,
56    queue_size: usize,
57) -> Result<()> {
58    // Re-arm SIG_IGN for SIGPIPE. The CLI resets it to SIG_DFL in main() for
59    // clean piping of short-lived commands, but for this long-running server
60    // that is fatal — a single client dropping mid-write would kill the whole
61    // process (issue #342). With SIG_IGN, such writes surface as EPIPE and are
62    // handled per-request by hyper/axum.
63    signals::ignore_sigpipe();
64
65    Config::install_runtime_models_dir_override(models_dir.clone());
66
67    let mut config = Config::load_or_default();
68    config.models_dir = models_dir.to_string_lossy().into_owned();
69    let model_name = config.resolved_default_model();
70
71    // ── Discover and initialize GPU workers ────────────────────────────────
72    let shared_pool = std::sync::Arc::new(std::sync::Mutex::new(
73        mold_inference::shared_pool::SharedPool::new(),
74    ));
75
76    let discovered = mold_inference::device::discover_gpus();
77    let selected = mold_inference::device::filter_gpus(&discovered, &gpu_selection);
78
79    if selected.is_empty() && !discovered.is_empty() {
80        anyhow::bail!(
81            "No GPUs matched selection {:?} (discovered: {:?})",
82            gpu_selection,
83            discovered.iter().map(|g| g.ordinal).collect::<Vec<_>>()
84        );
85    }
86
87    let mut workers = Vec::new();
88    let mut _gpu_thread_handles = Vec::new();
89
90    // Per-worker channel is a tiny buffer: one in-flight plus one immediate
91    // handoff. The global queue cap is enforced by `QueueHandle` against
92    // `queue_size`; per-worker overflow triggers the dispatcher's cross-worker
93    // retry path in `run_queue_dispatcher`.
94    const PER_WORKER_CHANNEL_SIZE: usize = 2;
95
96    let max_cached = state::resolve_max_cached_models();
97    for gpu in &selected {
98        let (job_tx, job_rx) = std::sync::mpsc::sync_channel(PER_WORKER_CHANNEL_SIZE);
99        let worker = std::sync::Arc::new(gpu_pool::GpuWorker {
100            gpu: gpu.clone(),
101            model_cache: std::sync::Arc::new(std::sync::Mutex::new(model_cache::ModelCache::new(
102                max_cached,
103            ))),
104            active_generation: std::sync::Arc::new(std::sync::RwLock::new(None)),
105            model_load_lock: std::sync::Arc::new(std::sync::Mutex::new(())),
106            shared_pool: shared_pool.clone(),
107            in_flight: AtomicUsize::new(0),
108            consecutive_failures: AtomicUsize::new(0),
109            degraded_until: std::sync::RwLock::new(None),
110            job_tx,
111        });
112
113        let handle = gpu_worker::spawn_gpu_thread(worker.clone(), job_rx);
114        _gpu_thread_handles.push(handle);
115        workers.push(worker);
116    }
117
118    let gpu_pool = std::sync::Arc::new(gpu_pool::GpuPool { workers });
119
120    // Log discovered GPUs.
121    for status in gpu_pool.gpu_status() {
122        info!(
123            gpu = status.ordinal,
124            name = %status.name,
125            vram_mb = status.vram_total_bytes / 1_000_000,
126            "GPU worker ready"
127        );
128    }
129
130    if selected.is_empty() {
131        info!("no GPUs discovered — server will operate in CPU/pull-only mode");
132    }
133
134    // ── Create generation queue ────────────────────────────────────────────
135    let (job_tx, job_rx) = tokio::sync::mpsc::channel(queue_size.max(1));
136    let queue_handle = QueueHandle::new(job_tx);
137
138    // ── Create AppState ────────────────────────────────────────────────────
139    let mut state = if gpu_pool.worker_count() > 0 {
140        if let Some(paths) = ModelPaths::resolve(&model_name, &config) {
141            info!(model = %model_name, "configured model");
142            info!(transformer = %paths.transformer.display());
143            info!(vae = %paths.vae.display());
144            if let Some(spatial_upscaler) = &paths.spatial_upscaler {
145                info!(spatial_upscaler = %spatial_upscaler.display());
146            }
147            if let Some(t5) = &paths.t5_encoder {
148                info!(t5 = %t5.display());
149            }
150            if let Some(clip) = &paths.clip_encoder {
151                info!(clip = %clip.display());
152            }
153            if let Some(t5_tok) = &paths.t5_tokenizer {
154                info!(t5_tok = %t5_tok.display());
155            }
156            if let Some(clip_tok) = &paths.clip_tokenizer {
157                info!(clip_tok = %clip_tok.display());
158            }
159            if let Some(clip2) = &paths.clip_encoder_2 {
160                info!(clip2 = %clip2.display());
161            }
162            if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
163                info!(clip2_tok = %clip2_tok.display());
164            }
165            for (i, te) in paths.text_encoder_files.iter().enumerate() {
166                info!(text_encoder_shard = i, path = %te.display());
167            }
168            if let Some(text_tok) = &paths.text_tokenizer {
169                info!(text_tok = %text_tok.display());
170            }
171            info!("multi-GPU mode defers model loading to per-GPU workers");
172        } else {
173            info!("no default model configured — models will be pulled on first request");
174        }
175        let mut state = state::AppState::empty(config, queue_handle, gpu_pool.clone(), queue_size);
176        state.shared_pool = shared_pool;
177        state
178    } else {
179        match ModelPaths::resolve(&model_name, &config) {
180            Some(paths) => {
181                info!(model = %model_name, "configured model");
182                info!(transformer = %paths.transformer.display());
183                info!(vae = %paths.vae.display());
184                if let Some(spatial_upscaler) = &paths.spatial_upscaler {
185                    info!(spatial_upscaler = %spatial_upscaler.display());
186                }
187                if let Some(t5) = &paths.t5_encoder {
188                    info!(t5 = %t5.display());
189                }
190                if let Some(clip) = &paths.clip_encoder {
191                    info!(clip = %clip.display());
192                }
193                if let Some(t5_tok) = &paths.t5_tokenizer {
194                    info!(t5_tok = %t5_tok.display());
195                }
196                if let Some(clip_tok) = &paths.clip_tokenizer {
197                    info!(clip_tok = %clip_tok.display());
198                }
199                if let Some(clip2) = &paths.clip_encoder_2 {
200                    info!(clip2 = %clip2.display());
201                }
202                if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
203                    info!(clip2_tok = %clip2_tok.display());
204                }
205                for (i, te) in paths.text_encoder_files.iter().enumerate() {
206                    info!(text_encoder_shard = i, path = %te.display());
207                }
208                if let Some(text_tok) = &paths.text_tokenizer {
209                    info!(text_tok = %text_tok.display());
210                }
211
212                let offload = std::env::var("MOLD_OFFLOAD").is_ok_and(|v| v == "1");
213                let engine = mold_inference::create_engine_with_pool(
214                    model_name,
215                    paths,
216                    &config,
217                    mold_inference::LoadStrategy::Eager,
218                    0,
219                    offload,
220                    Some(shared_pool.clone()),
221                )?;
222                let mut state = state::AppState::new(
223                    engine,
224                    config,
225                    queue_handle,
226                    gpu_pool.clone(),
227                    queue_size,
228                );
229                state.shared_pool = shared_pool;
230                state
231            }
232            None => {
233                info!("no default model configured — models will be pulled on first request");
234                state::AppState::empty(config, queue_handle, gpu_pool.clone(), queue_size)
235            }
236        }
237    };
238
239    // Open the gallery metadata DB (best-effort — server still runs without it).
240    match mold_db::open_default() {
241        Ok(Some(db)) => {
242            info!(db = %db.path().display(), "metadata DB opened");
243            state.metadata_db = std::sync::Arc::new(Some(db));
244        }
245        Ok(None) => {
246            tracing::info!("metadata DB disabled (MOLD_DB_DISABLE set or MOLD_HOME unresolved)");
247        }
248        Err(e) => {
249            tracing::warn!(
250                "failed to open metadata DB: {e:#} — gallery falls back to filesystem scan"
251            );
252        }
253    }
254
255    if state.metadata_db.is_some() {
256        let Some(jobs_root) = Config::mold_dir().map(|dir| dir.join("jobs")) else {
257            anyhow::bail!("metadata DB opened but MOLD_HOME could not be resolved for chain jobs");
258        };
259        std::fs::create_dir_all(&jobs_root)?;
260        let db_arc = state.metadata_db.clone();
261        let reconcile_root = jobs_root.clone();
262        let (flipped, repaired) = tokio::task::spawn_blocking(move || {
263            let Some(db) = db_arc.as_ref().as_ref() else {
264                anyhow::bail!("metadata DB disappeared before chain reconcile");
265            };
266            chain_job_runner::startup_reconcile(db, &reconcile_root)
267        })
268        .await??;
269        tracing::info!(
270            flipped,
271            repaired,
272            jobs_root = %jobs_root.display(),
273            "chain job startup reconcile complete"
274        );
275
276        let gc_db = state.metadata_db.clone();
277        let gc_root = jobs_root.clone();
278        let startup_gc = tokio::task::spawn_blocking(move || {
279            let Some(db) = gc_db.as_ref().as_ref() else {
280                anyhow::bail!("metadata DB disappeared before chain startup GC");
281            };
282            chain_job_runner::startup_gc_sweep(db, &gc_root)
283        })
284        .await??;
285        tracing::info!(
286            swept_ephemeral_jobs = startup_gc.swept_ephemeral_jobs,
287            pruned_artifact_dirs = startup_gc.pruned_artifact_dirs,
288            jobs_root = %jobs_root.display(),
289            "chain job startup GC complete"
290        );
291
292        let config_snapshot = state.config.read().await.clone();
293        let output_dir = if config_snapshot.is_output_disabled() {
294            None
295        } else {
296            Some(config_snapshot.effective_output_dir())
297        };
298        let deps = chain_job_runner::RunnerDeps {
299            db: state.metadata_db.clone(),
300            jobs_root,
301            executor: std::sync::Arc::new(chain_job_runner::ProductionStageExecutor::new(
302                state.gpu_pool.clone(),
303                config_snapshot,
304            )),
305            queue_probe: std::sync::Arc::new(chain_job_runner::ProductionQueueProbe::new(
306                state.queue.clone(),
307                state.gpu_pool.clone(),
308            )),
309            events: std::sync::Arc::new(chain_job_runner::JobEventBus::new()),
310            cancel: std::sync::Arc::new(chain_job_runner::CancelRegistry::new()),
311            job_locks: std::sync::Arc::new(chain_job_runner::JobMutationLocks::new()),
312            claims: std::sync::Arc::new(chain_job_runner::EphemeralClaims::default()),
313            output_dir,
314        };
315        state.chain_jobs = Some(std::sync::Arc::new(chain_job_runner::spawn_runner(deps)));
316    }
317
318    // Spawn the generation queue worker — processes jobs sequentially (single GPU).
319    // Spawn queue worker: use multi-GPU dispatcher if GPUs are available,
320    // otherwise fall back to the single-threaded queue worker.
321    let worker_state = state.clone();
322    if gpu_pool.worker_count() > 0 {
323        tokio::spawn(queue::run_queue_dispatcher(job_rx, worker_state));
324    } else {
325        tokio::spawn(queue::run_queue_worker(job_rx, worker_state));
326    }
327
328    // Background idle-TTL sweeper: reclaims parked engines that haven't been
329    // touched for `MOLD_CACHE_IDLE_TTL_SECS` seconds. Abort handle bound to
330    // graceful shutdown like every other long-running task in this fn.
331    let idle_evict_handle = spawn_cache_idle_evictor(
332        state.model_cache.clone(),
333        state.model_load_lock.clone(),
334        gpu_pool.clone(),
335        std::time::Duration::from_secs(state::resolve_cache_idle_ttl_secs()),
336    );
337
338    // ── Downloads UI (Agent A) ──────────────────────────────────────────────
339    // Single-writer download queue driver. Bind the `JoinHandle` so we can
340    // `.abort()` it when `axum::serve` returns — same pattern as the resource
341    // telemetry aggregator (see commit 5e43886). Without this the task would
342    // outlive graceful shutdown and keep polling its cancellation token until
343    // process exit.
344    let downloads_shutdown = tokio_util::sync::CancellationToken::new();
345    let downloads_models_dir = state.config.read().await.resolved_models_dir();
346    let downloads_driver = crate::downloads::spawn_driver(
347        state.downloads.clone(),
348        std::sync::Arc::new(crate::downloads::HfPullDriver),
349        std::sync::Arc::new(crate::downloads::CivitaiRecipeDriver),
350        downloads_models_dir,
351        downloads_shutdown.clone(),
352    );
353
354    // Ensure output directory exists and pre-generate thumbnails.
355    {
356        let config = state.config.read().await;
357        if config.is_output_disabled() {
358            tracing::warn!(
359                "image output is disabled (output_dir is empty) — \
360                 generated images will not be saved and the TUI gallery will be empty"
361            );
362        } else {
363            let output_dir = config.effective_output_dir();
364            let _ = std::fs::create_dir_all(&output_dir);
365            info!(output_dir = %output_dir.display(), "gallery output directory");
366            routes::spawn_thumbnail_warmup(&config);
367
368            // Async reconcile: import any existing files into the DB and
369            // drop rows whose backing files are missing. Runs on a blocking
370            // worker so it never stalls the request path even on large dirs.
371            if state.metadata_db.is_some() {
372                let db_arc = state.metadata_db.clone();
373                let dir = output_dir.clone();
374                tokio::spawn(async move {
375                    let join = tokio::task::spawn_blocking(move || {
376                        if let Some(db) = db_arc.as_ref() {
377                            db.reconcile(&dir)
378                        } else {
379                            Ok(mold_db::ReconcileStats::default())
380                        }
381                    })
382                    .await;
383                    match join {
384                        Ok(Ok(stats)) => tracing::info!(
385                            imported = stats.imported,
386                            updated = stats.updated,
387                            removed = stats.removed,
388                            kept = stats.kept,
389                            "metadata DB reconciled with gallery directory"
390                        ),
391                        Ok(Err(e)) => tracing::warn!("metadata DB reconcile failed: {e:#}"),
392                        Err(e) => tracing::warn!("metadata DB reconcile task join error: {e}"),
393                    }
394                });
395            }
396        }
397    }
398
399    // Load optional auth and rate-limit configuration from env vars.
400    let auth_state = auth::load_api_keys()?;
401    let rl_config = rate_limit::load_rate_limit_config()?;
402
403    let cors = build_cors_layer()?;
404
405    // Install the Prometheus metrics recorder (when feature-enabled).
406    // Must happen before any middleware or handler that records metrics.
407    #[cfg(feature = "metrics")]
408    let prometheus_handle = metrics::install_recorder();
409
410    // Build the router with middleware layers.
411    // Order (outermost → innermost): CORS → Trace → RequestID → Metrics → Auth → RateLimit → routes
412    // All inject + enforce layers use .layer() (not .route_layer()) so they run on
413    // ALL requests, including unmatched 404 paths — preventing auth/rate-limit bypass.
414    // Set up graceful shutdown: fires on SIGTERM or POST /api/shutdown.
415    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
416    *state.shutdown_tx.lock().await = Some(shutdown_tx);
417
418    #[cfg(unix)]
419    {
420        let sigterm_state = state.clone();
421        tokio::spawn(async move {
422            if let Ok(mut sig) =
423                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
424            {
425                sig.recv().await;
426                tracing::info!("received SIGTERM, initiating graceful shutdown");
427                if let Some(tx) = sigterm_state.shutdown_tx.lock().await.take() {
428                    let _ = tx.send(());
429                }
430            }
431        });
432    }
433
434    // Spawn the resource telemetry aggregator (1 Hz). Keep the `JoinHandle`
435    // bound so we can `.abort()` it when `axum::serve` returns — otherwise
436    // the task outlives server shutdown and keeps ticking until process exit.
437    let resources_aggregator = resources::spawn_aggregator(state.resources.clone());
438
439    // Save start_time before state is moved into the router (needed for metrics).
440    #[cfg(feature = "metrics")]
441    let server_start_time = state.start_time;
442
443    // The /metrics endpoint is mounted outside the auth/rate-limit stack so it
444    // is always accessible for monitoring scrapers (Prometheus, Grafana Agent, etc.).
445    #[allow(unused_mut)]
446    let mut app = routes::create_router(state)
447        .merge(web_ui::router())
448        .layer(DefaultBodyLimit::max(MAX_REQUEST_BODY_BYTES))
449        .layer(middleware::from_fn(rate_limit::rate_limit_middleware))
450        .layer(middleware::from_fn_with_state(
451            rl_config,
452            rate_limit::inject_rate_limit_state,
453        ))
454        .layer(middleware::from_fn(auth::require_api_key))
455        .layer(middleware::from_fn_with_state(
456            auth_state,
457            auth::inject_auth_state,
458        ));
459
460    // HTTP metrics middleware sits outside auth so it observes all requests
461    // (including auth failures and rate-limited responses).
462    #[cfg(feature = "metrics")]
463    {
464        app = app.layer(middleware::from_fn(metrics::http_metrics_middleware));
465    }
466
467    #[cfg(feature = "metrics")]
468    {
469        let metrics_state = metrics::MetricsState {
470            handle: prometheus_handle,
471            start_time: server_start_time,
472        };
473        app = app.route(
474            "/metrics",
475            axum::routing::get(metrics::metrics_endpoint).with_state(metrics_state),
476        );
477    }
478
479    let app = app
480        .layer(middleware::from_fn(request_id::request_id_middleware))
481        .layer(TraceLayer::new_for_http())
482        .layer(cors);
483
484    let addr: SocketAddr = format!("{bind}:{port}").parse()?;
485    let version = mold_core::build_info::version_string();
486    info!(%addr, %version, "starting mold server");
487
488    let listener = TcpListener::bind(addr).await?;
489    axum::serve(
490        listener,
491        app.into_make_service_with_connect_info::<SocketAddr>(),
492    )
493    .with_graceful_shutdown(async {
494        let _ = shutdown_rx.await;
495        tracing::info!("shutting down");
496    })
497    .await?;
498
499    // Server has stopped accepting requests — cancel the downloads token so the
500    // driver's `wait_for_work` arm returns, then abort the JoinHandle to ensure
501    // the task is cleaned up on the same shutdown path as the HTTP server.
502    // Matches the aggregator handle pattern from commit 5e43886.
503    downloads_shutdown.cancel();
504    downloads_driver.abort();
505    idle_evict_handle.abort();
506    // Server has stopped accepting requests — stop the telemetry aggregator
507    // so it doesn't outlive the server loop.
508    resources_aggregator.abort();
509
510    Ok(())
511}
512
513/// Spawn a tokio task that wakes every 60s and drops any cache entry whose
514/// `last_used` is older than `ttl` (and that isn't actively GPU-resident).
515/// Sweeps the legacy single-GPU cache and every per-worker cache in the
516/// multi-GPU pool. Returns the `JoinHandle` so the caller can `.abort()` on
517/// shutdown.
518///
519/// After dropping evicted engines, calls `reclaim_gpu_memory` on a thread
520/// bound to the relevant GPU ordinal so the freed memory actually returns to
521/// the OS rather than sitting in CUDA's per-context caching allocator. Without
522/// this step, `nvidia-smi` shows VRAM as still-allocated to the process even
523/// after the engine struct is dropped, which is the proximate cause of "model
524/// B OOMs even though model A finished and the queue went idle."
525fn spawn_cache_idle_evictor(
526    legacy_cache: std::sync::Arc<tokio::sync::Mutex<model_cache::ModelCache>>,
527    legacy_load_lock: std::sync::Arc<tokio::sync::Mutex<()>>,
528    gpu_pool: std::sync::Arc<gpu_pool::GpuPool>,
529    ttl: std::time::Duration,
530) -> tokio::task::JoinHandle<()> {
531    use tokio::time::{interval, MissedTickBehavior};
532    tokio::spawn(async move {
533        let mut tick = interval(std::time::Duration::from_secs(60));
534        tick.set_missed_tick_behavior(MissedTickBehavior::Skip);
535        // First tick fires immediately; skip it so a freshly-loaded model
536        // doesn't get reaped on boot before it's even been used.
537        tick.tick().await;
538        loop {
539            tick.tick().await;
540
541            // ── Legacy single-GPU cache ─────────────────────────────────────
542            //
543            // Take the legacy load lock for the full eviction+reclaim window
544            // so a generation request can't race in between us evicting and
545            // reclaiming, which would slot a fresh model load into a context
546            // we're about to reset.
547            {
548                let _load_guard = legacy_load_lock.lock().await;
549                let evicted = {
550                    let mut cache = legacy_cache.lock().await;
551                    cache.evict_idle(ttl)
552                };
553                let evicted_count = evicted.len();
554                // Drop engines OUTSIDE the cache lock — `cuMemFree` and
555                // safetensor unmap during drop can block other cache users.
556                drop(evicted);
557
558                // Only reclaim when something was evicted (nothing to flush
559                // otherwise) and when no GPU-resident engine remains
560                // (`cuDevicePrimaryCtxReset` would corrupt it). The load
561                // lock above guarantees no concurrent load can sneak one in
562                // between this check and the reclaim.
563                let legacy_active = legacy_cache.lock().await.active_model().is_some();
564                if evicted_count > 0 && !legacy_active {
565                    tokio::task::spawn_blocking(|| {
566                        mold_inference::reclaim_gpu_memory(0);
567                    })
568                    .await
569                    .ok();
570                }
571            }
572
573            // ── Multi-GPU per-worker caches ─────────────────────────────────
574            //
575            // Same pattern but per-worker: hold `worker.model_load_lock` for
576            // evict + drop + reclaim. Done under spawn_blocking because the
577            // worker locks are std mutexes and the reclaim itself is sync.
578            for worker in &gpu_pool.workers {
579                let worker = worker.clone();
580                let ttl = ttl;
581                tokio::task::spawn_blocking(move || {
582                    let _load_guard = match worker.model_load_lock.lock() {
583                        Ok(g) => g,
584                        Err(poisoned) => poisoned.into_inner(),
585                    };
586                    let evicted = {
587                        let mut cache =
588                            worker.model_cache.lock().unwrap_or_else(|e| e.into_inner());
589                        cache.evict_idle(ttl)
590                    };
591                    let evicted_count = evicted.len();
592                    drop(evicted);
593
594                    let active = worker
595                        .model_cache
596                        .lock()
597                        .map(|c| c.active_model().is_some())
598                        .unwrap_or(true);
599                    if evicted_count > 0 && !active {
600                        // Bind the thread to the worker's ordinal so
601                        // reclaim_gpu_memory's debug-assert is satisfied,
602                        // then clear so the spawn_blocking thread (which
603                        // returns to the tokio pool) doesn't carry a stale
604                        // binding.
605                        mold_inference::device::init_thread_gpu_ordinal(worker.gpu.ordinal);
606                        mold_inference::reclaim_gpu_memory(worker.gpu.ordinal);
607                        mold_inference::device::clear_thread_gpu_ordinal();
608                    }
609                })
610                .await
611                .ok();
612            }
613        }
614    })
615}
616
617fn build_cors_layer() -> Result<CorsLayer> {
618    let cors = match std::env::var("MOLD_CORS_ORIGIN") {
619        Ok(origin) if !origin.is_empty() => {
620            let origin = origin
621                .parse::<axum::http::HeaderValue>()
622                .map_err(|_| anyhow::anyhow!("invalid MOLD_CORS_ORIGIN value: {origin}"))?;
623            CorsLayer::new()
624                .allow_origin(origin)
625                .allow_methods([
626                    axum::http::Method::GET,
627                    axum::http::Method::POST,
628                    axum::http::Method::DELETE,
629                ])
630                .allow_headers(tower_http::cors::Any)
631                .expose_headers([
632                    axum::http::header::HeaderName::from_static("x-mold-seed-used"),
633                    axum::http::header::HeaderName::from_static("x-request-id"),
634                    axum::http::header::HeaderName::from_static("retry-after"),
635                    axum::http::header::HeaderName::from_static("x-mold-video-frames"),
636                    axum::http::header::HeaderName::from_static("x-mold-video-fps"),
637                    axum::http::header::HeaderName::from_static("x-mold-video-width"),
638                    axum::http::header::HeaderName::from_static("x-mold-video-height"),
639                    axum::http::header::HeaderName::from_static("x-mold-video-has-audio"),
640                    axum::http::header::HeaderName::from_static("x-mold-video-duration-ms"),
641                    axum::http::header::HeaderName::from_static("x-mold-video-audio-sample-rate"),
642                    axum::http::header::HeaderName::from_static("x-mold-video-audio-channels"),
643                    axum::http::header::HeaderName::from_static("x-mold-dimension-warning"),
644                ])
645        }
646        _ => CorsLayer::permissive(),
647    };
648    Ok(cors)
649}
650
651#[cfg(test)]
652mod tests {
653    use super::build_cors_layer;
654    use std::sync::Mutex;
655
656    static ENV_LOCK: Mutex<()> = Mutex::new(());
657
658    #[test]
659    fn invalid_cors_origin_returns_error() {
660        let _lock = ENV_LOCK.lock().unwrap();
661        std::env::set_var("MOLD_CORS_ORIGIN", "\nnot-a-header");
662        let result = build_cors_layer();
663        std::env::remove_var("MOLD_CORS_ORIGIN");
664        assert!(result.is_err());
665    }
666
667    #[test]
668    fn valid_cors_origin_builds_layer() {
669        let _lock = ENV_LOCK.lock().unwrap();
670        std::env::set_var("MOLD_CORS_ORIGIN", "https://example.com");
671        let result = build_cors_layer();
672        std::env::remove_var("MOLD_CORS_ORIGIN");
673        assert!(result.is_ok());
674    }
675}