Skip to main content

mold_server/
state.rs

1use mold_core::Config;
2use mold_inference::InferenceEngine;
3use std::path::PathBuf;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::{Arc, RwLock};
6use std::time::Instant;
7use tokio::sync::Mutex;
8
9use mold_inference::shared_pool::SharedPool;
10
11use crate::downloads::DownloadQueue;
12use crate::events::EventBroadcaster;
13use crate::gpu_pool::GpuPool;
14use crate::job_registry::{JobRegistry, SharedJobRegistry};
15use crate::model_cache::ModelCache;
16use crate::resources::ResourceBroadcaster;
17
18#[derive(Debug, Clone, Default)]
19pub struct EngineSnapshot {
20    /// Currently GPU-loaded model (None if no model on GPU).
21    pub model_name: Option<String>,
22    pub is_loaded: bool,
23    /// All models in the cache (loaded + unloaded), for status display.
24    pub cached_models: Vec<String>,
25}
26
27#[derive(Debug, Clone)]
28pub struct ActiveGenerationSnapshot {
29    pub model: String,
30    pub prompt_sha256: String,
31    pub started_at_unix_ms: u64,
32    pub started_at: Instant,
33}
34
35// ── Generation queue types ──────────────────────────────────────────────────
36
37/// Internal SSE message type used by both the queue worker and SSE streams.
38pub enum SseMessage {
39    Progress(mold_core::SseProgressEvent),
40    Complete(mold_core::SseCompleteEvent),
41    UpscaleComplete(mold_core::SseUpscaleCompleteEvent),
42    Error(mold_core::SseErrorEvent),
43}
44
45/// A generation job submitted to the queue worker.
46pub struct GenerationJob {
47    /// Server-assigned UUIDv4. Echoed back to the client in the initial
48    /// `SseProgressEvent::Queued` event so the SPA can later reconcile
49    /// a persisted "running" card against `GET /api/queue` — anything
50    /// the registry doesn't know about is a zombie left over from a
51    /// dropped SSE stream and gets dead-lettered client-side. Always
52    /// non-empty for jobs created through the public API; tests that
53    /// construct `GenerationJob` directly may leave it empty.
54    pub id: String,
55    pub request: mold_core::GenerateRequest,
56    /// Channel to send SSE progress/complete/error events (None for non-streaming).
57    pub progress_tx: Option<tokio::sync::mpsc::UnboundedSender<SseMessage>>,
58    /// Oneshot to return the final result for non-streaming callers.
59    pub result_tx: tokio::sync::oneshot::Sender<Result<GenerationJobResult, String>>,
60    /// Pre-resolved output directory for server-side image saving.
61    pub output_dir: Option<PathBuf>,
62}
63
64pub struct GenerationJobResult {
65    pub response: mold_core::GenerateResponse,
66    pub image: mold_core::ImageData,
67}
68
69/// Handle for submitting jobs to the generation queue.
70#[derive(Clone)]
71pub struct QueueHandle {
72    job_tx: tokio::sync::mpsc::Sender<GenerationJob>,
73    pending_count: Arc<AtomicUsize>,
74}
75
76/// Reason a `QueueHandle::submit` attempt failed.
77#[derive(Debug)]
78pub enum SubmitError {
79    /// Queue is at capacity — caller should return 503 with `Retry-After`.
80    Full { pending: usize, capacity: usize },
81    /// Receiving end is gone (server shutting down).
82    Shutdown,
83}
84
85impl QueueHandle {
86    pub fn new(job_tx: tokio::sync::mpsc::Sender<GenerationJob>) -> Self {
87        Self {
88            job_tx,
89            pending_count: Arc::new(AtomicUsize::new(0)),
90        }
91    }
92
93    /// Submit a generation job.
94    ///
95    /// Atomically reserves a slot against `capacity` using fetch_add, so a
96    /// burst of concurrent callers cannot all slip past a separate pending()
97    /// pre-check (TOCTOU).  Returns the queue position on success.
98    pub async fn submit(&self, job: GenerationJob, capacity: usize) -> Result<usize, SubmitError> {
99        let prev = self.pending_count.fetch_add(1, Ordering::SeqCst);
100        if prev >= capacity {
101            self.pending_count.fetch_sub(1, Ordering::SeqCst);
102            return Err(SubmitError::Full {
103                pending: prev,
104                capacity,
105            });
106        }
107        if self.job_tx.send(job).await.is_err() {
108            self.pending_count.fetch_sub(1, Ordering::SeqCst);
109            return Err(SubmitError::Shutdown);
110        }
111        #[cfg(feature = "metrics")]
112        {
113            crate::metrics::record_queue_submit();
114            crate::metrics::record_queue_depth(self.pending_count.load(Ordering::SeqCst));
115        }
116        Ok(prev)
117    }
118
119    pub fn decrement(&self) {
120        self.pending_count.fetch_sub(1, Ordering::SeqCst);
121    }
122
123    pub fn pending(&self) -> usize {
124        self.pending_count.load(Ordering::SeqCst)
125    }
126}
127
128// ── AppState ────────────────────────────────────────────────────────────────
129
130#[derive(Clone)]
131pub struct AppState {
132    // ── Multi-GPU fields ────────────────────────────────────────────────────
133    /// GPU worker pool for multi-GPU dispatch.
134    pub gpu_pool: Arc<GpuPool>,
135    /// Maximum queue capacity (for status reporting and 503 responses).
136    pub queue_capacity: usize,
137
138    // ── Legacy single-GPU fields (retained during migration) ────────────────
139    pub model_cache: Arc<Mutex<ModelCache>>,
140    /// Uses std::sync::RwLock (not tokio) because it's only accessed from
141    /// synchronous contexts (inside spawn_blocking closures and brief reads).
142    /// Must never be held across an .await point.
143    pub active_generation: Arc<RwLock<Option<ActiveGenerationSnapshot>>>,
144    pub config: Arc<tokio::sync::RwLock<Config>>,
145    pub start_time: Instant,
146    /// Guards concurrent model loads and hot-swaps.
147    pub model_load_lock: Arc<Mutex<()>>,
148    /// Guards concurrent pulls — only one download at a time.
149    pub pull_lock: Arc<Mutex<()>>,
150    /// Generation request queue.
151    pub queue: QueueHandle,
152    /// Authoritative registry of in-flight jobs (queued + running). Powers
153    /// `GET /api/queue` so the SPA can reconcile persisted "running" cards
154    /// against server reality and dead-letter zombies whose SSE stream
155    /// silently dropped.
156    pub job_registry: SharedJobRegistry,
157    /// Dispatch pause gate toggled by `POST /api/queue/pause` /
158    /// `POST /api/queue/resume`. The dispatch loops park on this at the top of
159    /// each iteration; a job already running on a worker finishes untouched.
160    pub queue_pause: Arc<crate::queue::QueuePause>,
161    /// Shared tokenizer pool for cross-engine caching.
162    pub shared_pool: Arc<std::sync::Mutex<SharedPool>>,
163    /// Shutdown trigger for graceful shutdown via `/api/shutdown` endpoint.
164    pub shutdown_tx: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
165    /// Cached upscaler engine to avoid recreating per request. Small models (2-64MB), single slot.
166    pub upscaler_cache: Arc<std::sync::Mutex<Option<Box<dyn mold_inference::UpscaleEngine>>>>,
167    /// SQLite-backed gallery metadata store. `None` when MOLD_DB_DISABLE=1 or
168    /// when MOLD_HOME could not be resolved — callers must fall back to the
169    /// filesystem walk in `routes::scan_gallery_dir`.
170    pub metadata_db: Arc<Option<mold_db::MetadataDb>>,
171    /// Durable chain-job runner handle. `None` when DB-backed chain jobs are
172    /// unavailable; chain-job API handlers return 503 in that state.
173    pub chain_jobs: Option<Arc<crate::chain_job_runner::ChainJobRunnerHandle>>,
174    // ── Downloads UI (Agent A) ──────────────────────────────────────────────
175    /// Bounded-parallel download queue.
176    pub downloads: Arc<DownloadQueue>,
177    /// Always-on resource telemetry (Agent B).
178    pub resources: Arc<ResourceBroadcaster>,
179    /// Server-wide lifecycle broadcast backing `GET /api/events` — job
180    /// queued/started/ended (mirrored by `job_registry`) plus gallery
181    /// added/removed. One SSE connection observes the whole server.
182    pub events: Arc<EventBroadcaster>,
183    // ── Catalog (live HF + Civitai) ─────────────────────────────────────────
184    /// In-process TTL cache backing `/api/catalog/search`.
185    pub catalog_live_cache: mold_catalog::live::LiveCache,
186    /// Upstream base URL for live Civitai search. Production runs with
187    /// the public host; tests override via `with_civitai_base`.
188    pub catalog_live_civitai_base: Arc<String>,
189    /// Cache of pure synthesis intents for installed catalog (`cv:*`/`hf:*`)
190    /// IDs, keyed by ID. Resolution into a `ModelConfig` is deferred to
191    /// engine-load time so disk-state races (file present vs not) can't
192    /// seal a wrong config; see `model_manager::resolve_intent_to_paths`.
193    pub catalog_intents: Arc<
194        tokio::sync::RwLock<
195            std::collections::HashMap<String, mold_catalog::synthesis::CatalogModelIntent>,
196        >,
197    >,
198}
199
200/// Default TTL for the live-search cache. Five minutes is long enough
201/// to absorb SPA re-mounts and quick paging without serving stale
202/// rankings.
203const CATALOG_LIVE_CACHE_TTL_SECS: u64 = 300;
204/// Cap on cached query keys. The working set per user is small —
205/// caps too high just retain stale rows past their TTL.
206const CATALOG_LIVE_CACHE_MAX_KEYS: usize = 256;
207/// Production Civitai endpoint. Tests inject a wiremock URI via
208/// [`AppState::with_civitai_base`].
209pub const CATALOG_LIVE_CIVITAI_BASE: &str = "https://civitai.com";
210
211fn default_live_cache() -> mold_catalog::live::LiveCache {
212    mold_catalog::live::LiveCache::new(
213        std::time::Duration::from_secs(CATALOG_LIVE_CACHE_TTL_SECS),
214        CATALOG_LIVE_CACHE_MAX_KEYS,
215    )
216}
217
218/// Default maximum number of cached models (GPU-resident + parked engine structs).
219pub const DEFAULT_MAX_CACHED_MODELS: usize = 3;
220/// Lower / upper bounds applied to env-overridden cache caps. Below 1 the
221/// cache can't hold the active engine; above 16 the OOM risk dwarfs the
222/// hit-rate gains for a typical local server.
223const MAX_CACHED_MODELS_LOWER: usize = 1;
224const MAX_CACHED_MODELS_UPPER: usize = 16;
225/// Env var that overrides `DEFAULT_MAX_CACHED_MODELS` at runtime.
226pub const MAX_CACHED_MODELS_ENV: &str = "MOLD_MAX_CACHED_MODELS";
227
228/// Default idle-TTL for parked cache entries — 30 minutes. Tuned for a
229/// local-first workflow: long enough that a user toggling between two
230/// models inside a session never pays a reload, short enough that
231/// background memory pressure doesn't accumulate overnight.
232pub const DEFAULT_CACHE_IDLE_TTL_SECS: u64 = 1800;
233const CACHE_IDLE_TTL_LOWER_SECS: u64 = 60;
234const CACHE_IDLE_TTL_UPPER_SECS: u64 = 86_400;
235/// Env var that overrides `DEFAULT_CACHE_IDLE_TTL_SECS`.
236pub const CACHE_IDLE_TTL_ENV: &str = "MOLD_CACHE_IDLE_TTL_SECS";
237
238/// Resolve the cache idle-TTL from env, falling back to the default.
239/// Out-of-range or unparseable values log a warning and use the default.
240pub fn resolve_cache_idle_ttl_secs() -> u64 {
241    match std::env::var(CACHE_IDLE_TTL_ENV) {
242        Ok(raw) => match raw.trim().parse::<u64>() {
243            Ok(n) if (CACHE_IDLE_TTL_LOWER_SECS..=CACHE_IDLE_TTL_UPPER_SECS).contains(&n) => n,
244            Ok(n) => {
245                tracing::warn!(
246                    env = CACHE_IDLE_TTL_ENV,
247                    value = n,
248                    lower = CACHE_IDLE_TTL_LOWER_SECS,
249                    upper = CACHE_IDLE_TTL_UPPER_SECS,
250                    "ignoring out-of-range cache idle-TTL; using default"
251                );
252                DEFAULT_CACHE_IDLE_TTL_SECS
253            }
254            Err(e) => {
255                tracing::warn!(
256                    env = CACHE_IDLE_TTL_ENV,
257                    raw = %raw,
258                    error = %e,
259                    "ignoring unparseable cache idle-TTL; using default"
260                );
261                DEFAULT_CACHE_IDLE_TTL_SECS
262            }
263        },
264        Err(_) => DEFAULT_CACHE_IDLE_TTL_SECS,
265    }
266}
267
268/// Resolve the model-cache capacity from env, falling back to the default.
269/// Out-of-range or unparseable values log a warning and use the default so
270/// a typo in the env never silently shrinks the cache to an unusable size.
271pub fn resolve_max_cached_models() -> usize {
272    match std::env::var(MAX_CACHED_MODELS_ENV) {
273        Ok(raw) => match raw.trim().parse::<usize>() {
274            Ok(n) if (MAX_CACHED_MODELS_LOWER..=MAX_CACHED_MODELS_UPPER).contains(&n) => n,
275            Ok(n) => {
276                tracing::warn!(
277                    env = MAX_CACHED_MODELS_ENV,
278                    value = n,
279                    lower = MAX_CACHED_MODELS_LOWER,
280                    upper = MAX_CACHED_MODELS_UPPER,
281                    "ignoring out-of-range cache cap; using default"
282                );
283                DEFAULT_MAX_CACHED_MODELS
284            }
285            Err(e) => {
286                tracing::warn!(
287                    env = MAX_CACHED_MODELS_ENV,
288                    raw = %raw,
289                    error = %e,
290                    "ignoring unparseable cache cap; using default"
291                );
292                DEFAULT_MAX_CACHED_MODELS
293            }
294        },
295        Err(_) => DEFAULT_MAX_CACHED_MODELS,
296    }
297}
298
299impl AppState {
300    /// Create state with a pre-loaded engine (server starts with a configured model).
301    pub fn new(
302        engine: Box<dyn InferenceEngine>,
303        config: Config,
304        queue: QueueHandle,
305        gpu_pool: Arc<GpuPool>,
306        queue_capacity: usize,
307    ) -> Self {
308        let mut cache = ModelCache::new(resolve_max_cached_models());
309        cache.insert(engine, 0);
310        let events = EventBroadcaster::new();
311        Self {
312            gpu_pool,
313            queue_capacity,
314            model_cache: Arc::new(Mutex::new(cache)),
315            active_generation: Arc::new(RwLock::new(None)),
316            config: Arc::new(tokio::sync::RwLock::new(config)),
317            start_time: Instant::now(),
318            model_load_lock: Arc::new(Mutex::new(())),
319            pull_lock: Arc::new(Mutex::new(())),
320            queue,
321            job_registry: JobRegistry::with_events(events.clone()),
322            queue_pause: crate::queue::QueuePause::new(),
323            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
324            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
325            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
326            metadata_db: Arc::new(None),
327            chain_jobs: None,
328            downloads: DownloadQueue::new(),
329            resources: ResourceBroadcaster::new(),
330            events,
331            catalog_live_cache: default_live_cache(),
332            catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
333            catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
334        }
335    }
336
337    /// Create state with no engine (zero-config startup, models pulled on demand).
338    pub fn empty(
339        config: Config,
340        queue: QueueHandle,
341        gpu_pool: Arc<GpuPool>,
342        queue_capacity: usize,
343    ) -> Self {
344        let events = EventBroadcaster::new();
345        Self {
346            gpu_pool,
347            queue_capacity,
348            model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
349            active_generation: Arc::new(RwLock::new(None)),
350            config: Arc::new(tokio::sync::RwLock::new(config)),
351            start_time: Instant::now(),
352            model_load_lock: Arc::new(Mutex::new(())),
353            pull_lock: Arc::new(Mutex::new(())),
354            queue,
355            job_registry: JobRegistry::with_events(events.clone()),
356            queue_pause: crate::queue::QueuePause::new(),
357            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
358            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
359            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
360            metadata_db: Arc::new(None),
361            chain_jobs: None,
362            downloads: DownloadQueue::new(),
363            resources: ResourceBroadcaster::new(),
364            events,
365            catalog_live_cache: default_live_cache(),
366            catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
367            catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
368        }
369    }
370
371    /// Create an empty GpuPool for testing (no GPU workers).
372    #[cfg(test)]
373    pub(crate) fn empty_gpu_pool() -> Arc<GpuPool> {
374        Arc::new(GpuPool {
375            workers: Vec::new(),
376        })
377    }
378
379    /// Alias for `empty_gpu_pool` — exposed for tests in sibling modules
380    /// (routes_test.rs, downloads_test.rs) that live in the crate but not
381    /// in the same file.
382    #[cfg(test)]
383    pub(crate) fn empty_gpu_pool_for_test() -> Arc<GpuPool> {
384        Self::empty_gpu_pool()
385    }
386
387    #[cfg(test)]
388    pub fn with_engine(engine: impl InferenceEngine + 'static) -> Self {
389        let (tx, _rx) = tokio::sync::mpsc::channel(16);
390        let queue = QueueHandle::new(tx);
391        let mut cache = ModelCache::new(resolve_max_cached_models());
392        cache.insert(Box::new(engine), 0);
393        let events = EventBroadcaster::new();
394        Self {
395            gpu_pool: Self::empty_gpu_pool(),
396            queue_capacity: 200,
397            model_cache: Arc::new(Mutex::new(cache)),
398            active_generation: Arc::new(RwLock::new(None)),
399            config: Arc::new(tokio::sync::RwLock::new(Config::default())),
400            start_time: Instant::now(),
401            model_load_lock: Arc::new(Mutex::new(())),
402            pull_lock: Arc::new(Mutex::new(())),
403            queue,
404            job_registry: JobRegistry::with_events(events.clone()),
405            queue_pause: crate::queue::QueuePause::new(),
406            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
407            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
408            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
409            metadata_db: Arc::new(None),
410            chain_jobs: None,
411            downloads: DownloadQueue::new(),
412            resources: ResourceBroadcaster::new(),
413            events,
414            catalog_live_cache: default_live_cache(),
415            catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
416            catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
417        }
418    }
419
420    /// Create state with a queue whose receiver is returned for testing.
421    #[cfg(test)]
422    pub fn with_engine_and_queue(
423        engine: impl InferenceEngine + 'static,
424    ) -> (Self, tokio::sync::mpsc::Receiver<GenerationJob>) {
425        let (tx, rx) = tokio::sync::mpsc::channel(16);
426        let queue = QueueHandle::new(tx);
427        let mut cache = ModelCache::new(resolve_max_cached_models());
428        cache.insert(Box::new(engine), 0);
429        let events = EventBroadcaster::new();
430        let state = Self {
431            gpu_pool: Self::empty_gpu_pool(),
432            queue_capacity: 200,
433            model_cache: Arc::new(Mutex::new(cache)),
434            active_generation: Arc::new(RwLock::new(None)),
435            config: Arc::new(tokio::sync::RwLock::new(Config::default())),
436            start_time: Instant::now(),
437            model_load_lock: Arc::new(Mutex::new(())),
438            pull_lock: Arc::new(Mutex::new(())),
439            queue,
440            job_registry: JobRegistry::with_events(events.clone()),
441            queue_pause: crate::queue::QueuePause::new(),
442            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
443            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
444            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
445            metadata_db: Arc::new(None),
446            chain_jobs: None,
447            downloads: DownloadQueue::new(),
448            resources: ResourceBroadcaster::new(),
449            events,
450            catalog_live_cache: default_live_cache(),
451            catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
452            catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
453        };
454        (state, rx)
455    }
456
457    /// Construct an empty AppState for integration tests. Catalog
458    /// surfaces hit live HF/Civitai (test callers point those at a
459    /// wiremock instance via `with_civitai_base`).
460    pub fn for_tests() -> Self {
461        let (tx, _rx) = tokio::sync::mpsc::channel(16);
462        let queue = QueueHandle::new(tx);
463        let events = EventBroadcaster::new();
464        Self {
465            gpu_pool: Arc::new(GpuPool {
466                workers: Vec::new(),
467            }),
468            queue_capacity: 200,
469            model_cache: Arc::new(Mutex::new(ModelCache::new(resolve_max_cached_models()))),
470            active_generation: Arc::new(RwLock::new(None)),
471            config: Arc::new(tokio::sync::RwLock::new(Config::default())),
472            start_time: Instant::now(),
473            model_load_lock: Arc::new(Mutex::new(())),
474            pull_lock: Arc::new(Mutex::new(())),
475            queue,
476            job_registry: JobRegistry::with_events(events.clone()),
477            queue_pause: crate::queue::QueuePause::new(),
478            shared_pool: Arc::new(std::sync::Mutex::new(SharedPool::new())),
479            shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
480            upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
481            metadata_db: Arc::new(None),
482            chain_jobs: None,
483            downloads: DownloadQueue::new(),
484            resources: ResourceBroadcaster::new(),
485            events,
486            catalog_live_cache: default_live_cache(),
487            catalog_live_civitai_base: Arc::new(CATALOG_LIVE_CIVITAI_BASE.to_string()),
488            catalog_intents: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
489        }
490    }
491
492    /// Override the live-search Civitai base URL on an existing state.
493    /// Tests point this at a wiremock instance; production never calls
494    /// it (the `new` / `empty` constructors set the public host).
495    pub fn with_civitai_base(mut self, base: impl Into<String>) -> Self {
496        self.catalog_live_civitai_base = Arc::new(base.into());
497        self
498    }
499}
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504
505    #[test]
506    fn engine_snapshot_default_is_unloaded() {
507        let snap = EngineSnapshot::default();
508        assert!(snap.model_name.is_none());
509        assert!(!snap.is_loaded);
510    }
511
512    #[test]
513    fn active_generation_snapshot_stores_fields() {
514        let snap = ActiveGenerationSnapshot {
515            model: "flux-dev:q8".to_string(),
516            prompt_sha256: "abc123".to_string(),
517            started_at_unix_ms: 1700000000000,
518            started_at: std::time::Instant::now(),
519        };
520        assert_eq!(snap.model, "flux-dev:q8");
521        assert_eq!(snap.prompt_sha256, "abc123");
522        assert_eq!(snap.started_at_unix_ms, 1700000000000);
523    }
524
525    #[test]
526    fn queue_handle_pending_starts_at_zero() {
527        let (tx, _rx) = tokio::sync::mpsc::channel::<GenerationJob>(16);
528        let handle = QueueHandle::new(tx);
529        assert_eq!(handle.pending(), 0);
530    }
531
532    #[test]
533    fn upscaler_cache_starts_empty() {
534        let config = mold_core::Config::default();
535        let state = AppState::empty(
536            config,
537            QueueHandle::new(tokio::sync::mpsc::channel(1).0),
538            AppState::empty_gpu_pool(),
539            200,
540        );
541        let cache = state.upscaler_cache.lock().unwrap();
542        assert!(cache.is_none());
543    }
544
545    #[test]
546    fn upscaler_cache_cleared_by_setting_none() {
547        let config = mold_core::Config::default();
548        let state = AppState::empty(
549            config,
550            QueueHandle::new(tokio::sync::mpsc::channel(1).0),
551            AppState::empty_gpu_pool(),
552            200,
553        );
554        {
555            let mut cache = state.upscaler_cache.lock().unwrap();
556            *cache = None;
557        }
558        let cache = state.upscaler_cache.lock().unwrap();
559        assert!(cache.is_none());
560    }
561
562    #[test]
563    fn app_state_exposes_resources_broadcaster() {
564        let config = mold_core::Config::default();
565        let state = AppState::empty(
566            config,
567            QueueHandle::new(tokio::sync::mpsc::channel(1).0),
568            AppState::empty_gpu_pool(),
569            200,
570        );
571        // The broadcaster must exist and return None before any aggregator tick.
572        assert!(state.resources.latest().is_none());
573        // Subscribing must succeed (no panics).
574        let _rx = state.resources.subscribe();
575    }
576
577    /// Serializes every test that touches a process-wide env var via
578    /// `std::env::set_var` — mutating env is global state, so without this
579    /// guard parallel tests would race on the env table.
580    static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
581
582    /// Set `name` to `value` (or remove it when `value` is `None`), invoke
583    /// `f`, then restore the original value. Lock-serialized so concurrent
584    /// tests don't race on the env table.
585    fn with_env<R>(name: &str, value: Option<&str>, f: impl FnOnce() -> R) -> R {
586        let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
587        let prev = std::env::var(name).ok();
588        match value {
589            Some(v) => std::env::set_var(name, v),
590            None => std::env::remove_var(name),
591        }
592        let out = f();
593        match prev {
594            Some(v) => std::env::set_var(name, v),
595            None => std::env::remove_var(name),
596        }
597        out
598    }
599
600    #[test]
601    fn resolve_max_cached_uses_default_when_env_missing() {
602        let n = with_env(MAX_CACHED_MODELS_ENV, None, resolve_max_cached_models);
603        assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
604    }
605
606    #[test]
607    fn resolve_max_cached_honors_env_within_range() {
608        let n = with_env(MAX_CACHED_MODELS_ENV, Some("8"), resolve_max_cached_models);
609        assert_eq!(n, 8);
610    }
611
612    #[test]
613    fn resolve_max_cached_clamps_zero_back_to_default() {
614        let n = with_env(MAX_CACHED_MODELS_ENV, Some("0"), resolve_max_cached_models);
615        assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
616    }
617
618    #[test]
619    fn resolve_max_cached_clamps_overflow_back_to_default() {
620        let n = with_env(
621            MAX_CACHED_MODELS_ENV,
622            Some("999"),
623            resolve_max_cached_models,
624        );
625        assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
626    }
627
628    #[test]
629    fn resolve_max_cached_falls_back_when_env_unparseable() {
630        let n = with_env(
631            MAX_CACHED_MODELS_ENV,
632            Some("not-a-number"),
633            resolve_max_cached_models,
634        );
635        assert_eq!(n, DEFAULT_MAX_CACHED_MODELS);
636    }
637
638    #[test]
639    fn resolve_cache_idle_ttl_uses_default_when_env_missing() {
640        let n = with_env(CACHE_IDLE_TTL_ENV, None, resolve_cache_idle_ttl_secs);
641        assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
642    }
643
644    #[test]
645    fn resolve_cache_idle_ttl_honors_env_within_range() {
646        // 600s is comfortably inside [60, 86_400] — the resolver must echo it.
647        let n = with_env(CACHE_IDLE_TTL_ENV, Some("600"), resolve_cache_idle_ttl_secs);
648        assert_eq!(n, 600);
649    }
650
651    #[test]
652    fn resolve_cache_idle_ttl_clamps_zero_back_to_default() {
653        // 0 is below the 60s lower bound; falls back to the default with a warn log.
654        let n = with_env(CACHE_IDLE_TTL_ENV, Some("0"), resolve_cache_idle_ttl_secs);
655        assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
656    }
657
658    #[test]
659    fn resolve_cache_idle_ttl_clamps_overflow_back_to_default() {
660        // 100_000 is above the 86_400s upper bound; falls back to the default.
661        let n = with_env(
662            CACHE_IDLE_TTL_ENV,
663            Some("100000"),
664            resolve_cache_idle_ttl_secs,
665        );
666        assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
667    }
668
669    #[test]
670    fn resolve_cache_idle_ttl_falls_back_when_env_unparseable() {
671        let n = with_env(
672            CACHE_IDLE_TTL_ENV,
673            Some("not-a-number"),
674            resolve_cache_idle_ttl_secs,
675        );
676        assert_eq!(n, DEFAULT_CACHE_IDLE_TTL_SECS);
677    }
678}