mold_server/
lib.rs

1pub mod model_manager;
2pub mod routes;
3pub mod state;
4
5#[cfg(test)]
6mod routes_test;
7
8use anyhow::Result;
9use mold_core::{Config, ModelPaths};
10use std::net::SocketAddr;
11use std::path::PathBuf;
12use tokio::net::TcpListener;
13use tower_http::cors::CorsLayer;
14use tower_http::trace::TraceLayer;
15use tracing::info;
16
17pub async fn run_server(bind: &str, port: u16, models_dir: PathBuf) -> Result<()> {
18    Config::install_runtime_models_dir_override(models_dir.clone());
19
20    let mut config = Config::load_or_default();
21    config.models_dir = models_dir.to_string_lossy().into_owned();
22    let model_name = config.default_model.clone();
23
24    let state = match ModelPaths::resolve(&model_name, &config) {
25        Some(paths) => {
26            info!(model = %model_name, "configured model");
27            info!(transformer = %paths.transformer.display());
28            info!(vae = %paths.vae.display());
29            if let Some(t5) = &paths.t5_encoder {
30                info!(t5 = %t5.display());
31            }
32            if let Some(clip) = &paths.clip_encoder {
33                info!(clip = %clip.display());
34            }
35            if let Some(t5_tok) = &paths.t5_tokenizer {
36                info!(t5_tok = %t5_tok.display());
37            }
38            if let Some(clip_tok) = &paths.clip_tokenizer {
39                info!(clip_tok = %clip_tok.display());
40            }
41            if let Some(clip2) = &paths.clip_encoder_2 {
42                info!(clip2 = %clip2.display());
43            }
44            if let Some(clip2_tok) = &paths.clip_tokenizer_2 {
45                info!(clip2_tok = %clip2_tok.display());
46            }
47            for (i, te) in paths.text_encoder_files.iter().enumerate() {
48                info!(text_encoder_shard = i, path = %te.display());
49            }
50            if let Some(text_tok) = &paths.text_tokenizer {
51                info!(text_tok = %text_tok.display());
52            }
53
54            let engine = mold_inference::create_engine(
55                model_name,
56                paths,
57                &config,
58                mold_inference::LoadStrategy::Eager,
59            )?;
60            state::AppState::new(engine, config)
61        }
62        None => {
63            info!("no default model configured — models will be pulled on first request");
64            state::AppState::empty(config)
65        }
66    };
67
68    let cors = build_cors_layer()?;
69
70    let app = routes::create_router(state)
71        .layer(TraceLayer::new_for_http())
72        .layer(cors);
73
74    let addr: SocketAddr = format!("{bind}:{port}").parse()?;
75    info!(%addr, "starting mold server");
76
77    let listener = TcpListener::bind(addr).await?;
78    axum::serve(listener, app).await?;
79
80    Ok(())
81}
82
83fn build_cors_layer() -> Result<CorsLayer> {
84    let cors = match std::env::var("MOLD_CORS_ORIGIN") {
85        Ok(origin) if !origin.is_empty() => {
86            let origin = origin
87                .parse::<axum::http::HeaderValue>()
88                .map_err(|_| anyhow::anyhow!("invalid MOLD_CORS_ORIGIN value: {origin}"))?;
89            CorsLayer::new()
90                .allow_origin(origin)
91                .allow_methods([
92                    axum::http::Method::GET,
93                    axum::http::Method::POST,
94                    axum::http::Method::DELETE,
95                ])
96                .allow_headers(tower_http::cors::Any)
97                .expose_headers([axum::http::header::HeaderName::from_static(
98                    "x-mold-seed-used",
99                )])
100        }
101        _ => CorsLayer::permissive(),
102    };
103    Ok(cors)
104}
105
106#[cfg(test)]
107mod tests {
108    use super::build_cors_layer;
109    use std::sync::Mutex;
110
111    static ENV_LOCK: Mutex<()> = Mutex::new(());
112
113    #[test]
114    fn invalid_cors_origin_returns_error() {
115        let _lock = ENV_LOCK.lock().unwrap();
116        std::env::set_var("MOLD_CORS_ORIGIN", "\nnot-a-header");
117        let result = build_cors_layer();
118        std::env::remove_var("MOLD_CORS_ORIGIN");
119        assert!(result.is_err());
120    }
121
122    #[test]
123    fn valid_cors_origin_builds_layer() {
124        let _lock = ENV_LOCK.lock().unwrap();
125        std::env::set_var("MOLD_CORS_ORIGIN", "https://example.com");
126        let result = build_cors_layer();
127        std::env::remove_var("MOLD_CORS_ORIGIN");
128        assert!(result.is_ok());
129    }
130}
mold_server/lib.rs

mold_server/
lib.rs
