use anyhow::Result;
use base64::{engine::general_purpose, Engine as _};
use mold_core::{
Config, GalleryImage, GenerateRequest, GenerateResponse, LoraInfo, LoraWeight, MoldClient,
OutputFormat, SseProgressEvent,
};
use serde::Deserialize;
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::sync::Mutex;
const MCP_PROTOCOL_VERSION: &str = "2025-06-18";
const MAX_ASYNC_JOBS: usize = 32;
pub async fn run(host: Option<String>) -> Result<()> {
let server = McpServer::new(host);
let stdin = BufReader::new(tokio::io::stdin());
let mut lines = stdin.lines();
let mut stdout = BufWriter::new(tokio::io::stdout());
while let Some(line) = lines.next_line().await? {
if line.trim().is_empty() {
continue;
}
let response = match serde_json::from_str::<Value>(&line) {
Ok(message) => server.handle_message(message).await,
Err(e) => Some(error_response(
Value::Null,
-32700,
format!("parse error: {e}"),
)),
};
if let Some(response) = response {
let encoded = serde_json::to_string(&response)?;
stdout.write_all(encoded.as_bytes()).await?;
stdout.write_all(b"\n").await?;
stdout.flush().await?;
}
}
Ok(())
}
struct McpServer {
client: MoldClient,
jobs: AsyncJobRegistry,
}
impl McpServer {
fn new(host: Option<String>) -> Self {
let client = match host {
Some(host) => match std::env::var("MOLD_API_KEY").ok().filter(|k| !k.is_empty()) {
Some(api_key) => MoldClient::with_api_key(&host, api_key),
None => MoldClient::new(&host),
},
None => MoldClient::from_env(),
};
Self {
client,
jobs: AsyncJobRegistry::default(),
}
}
async fn handle_message(&self, message: Value) -> Option<Value> {
match message {
Value::Array(items) => {
if items.is_empty() {
return Some(error_response(
Value::Null,
-32600,
"invalid request: empty batch",
));
}
let mut responses = Vec::new();
for item in items {
if let Some(response) = self.handle_single(item).await {
responses.push(response);
}
}
(!responses.is_empty()).then_some(Value::Array(responses))
}
other => self.handle_single(other).await,
}
}
async fn handle_single(&self, message: Value) -> Option<Value> {
let method = match message.get("method").and_then(Value::as_str) {
Some(method) => method,
None => {
return Some(error_response(
message.get("id").cloned().unwrap_or(Value::Null),
-32600,
"invalid request: missing method",
));
}
};
if method == "tools/call" {
return self.handle_tool_call(&message).await;
}
handle_protocol_message(message)
}
async fn handle_tool_call(&self, message: &Value) -> Option<Value> {
let id = message.get("id").cloned()?;
let params = message.get("params").cloned().unwrap_or_else(|| json!({}));
let Some(name) = params.get("name").and_then(Value::as_str) else {
return Some(error_response(id, -32602, "tools/call missing params.name"));
};
let arguments = params
.get("arguments")
.cloned()
.unwrap_or_else(|| json!({}));
let result = match name {
"generate_image" => self.tool_generate_image(arguments).await,
"generate_image_async" => self.tool_generate_image_async(arguments).await,
"generation_status" => self.tool_generation_status(arguments).await,
"list_gallery" => self.tool_list_gallery(arguments).await,
"get_gallery_image" => self.tool_get_gallery_image(arguments).await,
"list_models" => self.tool_list_models(arguments).await,
"list_loras" => self.tool_list_loras(arguments).await,
"server_status" => self.tool_server_status().await,
other => Err(format!("unknown tool: {other}")),
};
Some(match result {
Ok(result) => response(id, result),
Err(err) => response(
id,
json!({
"content": [{ "type": "text", "text": err }],
"isError": true
}),
),
})
}
async fn tool_generate_image(&self, arguments: Value) -> std::result::Result<Value, String> {
let mut args: GenerateImageArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let loras = self.resolve_loras(args.loras.take()).await?;
let req = build_generate_request(args, loras)?;
let response = self
.client
.generate(req)
.await
.map_err(|e| format!("mold generation failed: {e}"))?;
let image = response
.images
.first()
.ok_or_else(|| "mold did not return an image".to_string())?;
let encoded = general_purpose::STANDARD.encode(&image.data);
let mut details = format!(
"Generated {}x{} {} image with {} in {:.1}s; seed {}",
image.width,
image.height,
image.format.extension(),
response.model,
response.generation_time_ms as f64 / 1000.0,
response.seed_used
);
if let Some(gpu) = response.gpu {
details.push_str(&format!("; gpu {gpu}"));
}
Ok(json!({
"content": [
{ "type": "text", "text": details },
{
"type": "image",
"data": encoded,
"mimeType": image.format.content_type()
}
]
}))
}
async fn tool_generate_image_async(
&self,
arguments: Value,
) -> std::result::Result<Value, String> {
let mut args: GenerateImageArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let loras = self.resolve_loras(args.loras.take()).await?;
let req = build_generate_request(args, loras)?;
let job_id = self.jobs.create(&req).await;
let client = self.client.clone();
let jobs = self.jobs.clone();
let task_job_id = job_id.clone();
tokio::spawn(async move {
jobs.mark_running(&task_job_id).await;
let (progress_tx, mut progress_rx) =
tokio::sync::mpsc::unbounded_channel::<SseProgressEvent>();
let progress_jobs = jobs.clone();
let progress_job_id = task_job_id.clone();
let progress_task = tokio::spawn(async move {
while let Some(event) = progress_rx.recv().await {
progress_jobs.record_progress(&progress_job_id, event).await;
}
});
let result = match client.generate_stream(&req, progress_tx).await {
Ok(Some(response)) => Ok(response),
Ok(None) => client.generate(req).await.map_err(|e| {
format!("mold generation failed after non-streaming fallback: {e}")
}),
Err(e) => Err(format!("mold generation failed: {e}")),
};
let _ = progress_task.await;
jobs.finish(&task_job_id, result).await;
});
Ok(json!({
"content": [{
"type": "text",
"text": format!(
"Started async generation job {job_id}. Poll generation_status with {{\"job_id\":\"{job_id}\"}}."
)
}],
"structuredContent": {
"job_id": job_id,
"status": "queued"
}
}))
}
async fn tool_generation_status(&self, arguments: Value) -> std::result::Result<Value, String> {
let args: GenerationStatusArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let include_result = args.include_result.unwrap_or(true);
if let Some(job_id) = args.job_id {
let Some(job) = self.jobs.get(&job_id).await else {
return Err(format!("unknown async generation job: {job_id}"));
};
return Ok(job_tool_result(&job, include_result));
}
let jobs = self.jobs.list().await;
if jobs.is_empty() {
return Ok(text_result("No async generation jobs are tracked."));
}
let lines = jobs.iter().map(job_status_line).collect::<Vec<_>>();
Ok(json!({
"content": [{ "type": "text", "text": lines.join("\n") }],
"structuredContent": {
"jobs": jobs.iter().map(job_summary_json).collect::<Vec<_>>()
}
}))
}
async fn tool_list_gallery(&self, arguments: Value) -> std::result::Result<Value, String> {
let args: ListGalleryArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let images = self
.client
.list_gallery()
.await
.map_err(|e| format!("failed to list gallery: {e}"))?;
let images = filter_sort_gallery(images, &args.filter())?;
let total = images.len();
let limit = args.limit.unwrap_or(25).clamp(1, 200);
let offset = args.offset.unwrap_or(0);
let page = images
.into_iter()
.skip(offset)
.take(limit)
.collect::<Vec<_>>();
if page.is_empty() {
return Ok(text_result("No gallery images matched."));
}
let lines = page.iter().map(gallery_line).collect::<Vec<_>>();
Ok(json!({
"content": [{
"type": "text",
"text": format!(
"Showing {} of {total} matched gallery images:\n{}",
lines.len(),
lines.join("\n")
)
}],
"structuredContent": {
"total": total,
"offset": offset,
"limit": limit,
"images": page.iter().map(gallery_summary_json).collect::<Vec<_>>()
}
}))
}
async fn tool_get_gallery_image(&self, arguments: Value) -> std::result::Result<Value, String> {
let args: GetGalleryImageArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let gallery = self
.client
.list_gallery()
.await
.map_err(|e| format!("failed to list gallery: {e}"))?;
let selected = if let Some(filename) = args.filename.as_deref() {
gallery
.iter()
.find(|image| image.filename == filename)
.cloned()
.ok_or_else(|| format!("gallery image not found: {filename}"))?
} else {
let matches = filter_sort_gallery(gallery, &args.filter())?;
matches
.into_iter()
.next()
.ok_or_else(|| "no gallery image matched".to_string())?
};
let format = gallery_format(&selected);
let use_thumbnail = args.thumbnail.unwrap_or(format == Some(OutputFormat::Mp4));
let (data, mime_type, returned_thumbnail) = if use_thumbnail {
(
self.client
.get_gallery_thumbnail(&selected.filename)
.await
.map_err(|e| format!("failed to fetch gallery thumbnail: {e}"))?,
OutputFormat::Png.content_type(),
true,
)
} else {
if format == Some(OutputFormat::Mp4) {
return Err(
"selected gallery item is an MP4 video; set thumbnail=true to fetch a preview"
.to_string(),
);
}
let data = self
.client
.get_gallery_image(&selected.filename)
.await
.map_err(|e| format!("failed to fetch gallery image: {e}"))?;
let mime_type = format
.map(|f| f.content_type())
.unwrap_or_else(|| content_type_for_filename(&selected.filename));
(data, mime_type, false)
};
Ok(json!({
"content": [
{
"type": "text",
"text": if returned_thumbnail {
format!("Fetched thumbnail for {}", gallery_line(&selected))
} else {
format!("Fetched {}", gallery_line(&selected))
}
},
{
"type": "image",
"data": general_purpose::STANDARD.encode(data),
"mimeType": mime_type
}
],
"structuredContent": {
"image": gallery_summary_json(&selected),
"thumbnail": returned_thumbnail
}
}))
}
async fn tool_list_models(&self, arguments: Value) -> std::result::Result<Value, String> {
let args: ListModelsArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let models = self
.client
.list_models_extended()
.await
.map_err(|e| format!("failed to list models: {e}"))?;
let limit = args.limit.unwrap_or(50).clamp(1, 200);
let mut lines = Vec::new();
for model in models
.into_iter()
.filter(|m| !args.downloaded_only.unwrap_or(true) || m.downloaded)
.filter(|m| !args.generation_only.unwrap_or(true) || m.is_generation_model())
.take(limit)
{
let status = if model.downloaded {
"downloaded"
} else {
"not downloaded"
};
lines.push(format!(
"- {} [{}] {status}; default {}x{}, steps {}, guidance {}",
model.name,
model.family,
model.defaults.default_width,
model.defaults.default_height,
model.defaults.default_steps,
model.defaults.default_guidance
));
}
if lines.is_empty() {
lines.push("No matching models found.".to_string());
}
Ok(text_result(lines.join("\n")))
}
async fn tool_list_loras(&self, arguments: Value) -> std::result::Result<Value, String> {
let args: ListLorasArgs =
serde_json::from_value(arguments).map_err(|e| format!("invalid arguments: {e}"))?;
let loras = self
.client
.list_loras(args.model.as_deref())
.await
.map_err(|e| format!("failed to list LoRAs: {e}"))?;
if loras.is_empty() {
let scope = args
.model
.as_deref()
.map(|model| format!(" compatible with {model}"))
.unwrap_or_default();
return Ok(json!({
"content": [{ "type": "text", "text": format!("No installed LoRAs{scope} found.") }],
"structuredContent": { "loras": [] }
}));
}
let scope = args
.model
.as_deref()
.map(|model| format!(" compatible with {model}"))
.unwrap_or_default();
let lines = loras.iter().map(lora_line).collect::<Vec<_>>();
Ok(json!({
"content": [{
"type": "text",
"text": format!(
"Showing {} installed LoRAs{scope}:\n{}",
loras.len(),
lines.join("\n")
)
}],
"structuredContent": { "loras": loras }
}))
}
async fn tool_server_status(&self) -> std::result::Result<Value, String> {
let status = self
.client
.server_status()
.await
.map_err(|e| format!("failed to read server status: {e}"))?;
let mut lines = vec![
format!("mold server {}", status.version),
format!("busy: {}", status.busy),
format!("loaded models: {}", status.models_loaded.join(", ")),
format!("uptime: {}s", status.uptime_secs),
];
if let Some(hostname) = status.hostname {
lines.push(format!("host: {hostname}"));
}
if let Some(memory) = status.memory_status {
lines.push(memory);
}
if let Some(gpus) = status.gpus {
for gpu in gpus {
lines.push(format!(
"gpu {}: {} ({:?}){}",
gpu.ordinal,
gpu.name,
gpu.state,
gpu.loaded_model
.map(|m| format!(", loaded {m}"))
.unwrap_or_default()
));
}
}
if let Some(depth) = status.queue_depth {
lines.push(format!("queue depth: {depth}"));
}
Ok(text_result(lines.join("\n")))
}
async fn resolve_loras(
&self,
loras: Option<Vec<McpLoraArg>>,
) -> std::result::Result<Option<Vec<LoraWeight>>, String> {
let Some(loras) = loras else {
return Ok(None);
};
if loras.is_empty() {
return Ok(None);
}
let refs = loras
.into_iter()
.map(normalise_mcp_lora_arg)
.collect::<std::result::Result<Vec<_>, _>>()?;
let catalog = if refs.iter().any(|lora| is_lora_catalog_id(&lora.value)) {
self.client
.list_loras(None)
.await
.map_err(|e| format!("failed to resolve LoRA ids: {e}"))?
} else {
Vec::new()
};
Ok(Some(resolve_mcp_lora_refs(&refs, &catalog)?))
}
}
#[derive(Debug, Deserialize)]
struct GenerateImageArgs {
prompt: String,
model: Option<String>,
width: Option<u32>,
height: Option<u32>,
steps: Option<u32>,
guidance: Option<f64>,
seed: Option<u64>,
negative_prompt: Option<String>,
output_format: Option<String>,
expand: Option<bool>,
loras: Option<Vec<McpLoraArg>>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
enum McpLoraArg {
Ref(String),
Object(McpLoraObject),
}
#[derive(Debug, Clone, Deserialize)]
struct McpLoraObject {
id: Option<String>,
path: Option<String>,
scale: Option<f64>,
}
#[derive(Debug, Clone, PartialEq)]
struct McpLoraRef {
value: String,
scale: f64,
}
#[derive(Debug, Deserialize)]
struct ListModelsArgs {
downloaded_only: Option<bool>,
generation_only: Option<bool>,
limit: Option<usize>,
}
#[derive(Debug, Deserialize)]
struct ListLorasArgs {
model: Option<String>,
}
#[derive(Debug, Deserialize)]
struct GenerationStatusArgs {
job_id: Option<String>,
include_result: Option<bool>,
}
#[derive(Debug, Deserialize, Default)]
struct ListGalleryArgs {
query: Option<String>,
filename: Option<String>,
model: Option<String>,
format: Option<String>,
sort_by: Option<String>,
sort_order: Option<String>,
limit: Option<usize>,
offset: Option<usize>,
}
impl ListGalleryArgs {
fn filter(&self) -> GalleryFilter {
GalleryFilter {
query: self.query.clone(),
filename: self.filename.clone(),
model: self.model.clone(),
format: self.format.clone(),
sort_by: self.sort_by.clone(),
sort_order: self.sort_order.clone(),
}
}
}
#[derive(Debug, Deserialize, Default)]
struct GetGalleryImageArgs {
filename: Option<String>,
latest: Option<bool>,
query: Option<String>,
model: Option<String>,
format: Option<String>,
sort_by: Option<String>,
sort_order: Option<String>,
thumbnail: Option<bool>,
}
impl GetGalleryImageArgs {
fn filter(&self) -> GalleryFilter {
GalleryFilter {
query: self.query.clone(),
filename: None,
model: self.model.clone(),
format: self.format.clone(),
sort_by: self.sort_by.clone().or_else(|| {
if self.latest.unwrap_or(true) {
Some("timestamp".to_string())
} else {
None
}
}),
sort_order: self.sort_order.clone().or_else(|| {
if self.latest.unwrap_or(true) {
Some("desc".to_string())
} else {
None
}
}),
}
}
}
#[derive(Debug, Clone, Default)]
struct GalleryFilter {
query: Option<String>,
filename: Option<String>,
model: Option<String>,
format: Option<String>,
sort_by: Option<String>,
sort_order: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AsyncJobStatus {
Queued,
Running,
Succeeded,
Failed,
}
impl AsyncJobStatus {
fn as_str(self) -> &'static str {
match self {
Self::Queued => "queued",
Self::Running => "running",
Self::Succeeded => "succeeded",
Self::Failed => "failed",
}
}
fn is_terminal(self) -> bool {
matches!(self, Self::Succeeded | Self::Failed)
}
}
#[derive(Debug, Clone)]
struct AsyncGenerationJob {
id: String,
status: AsyncJobStatus,
model: String,
prompt_preview: String,
created_at_ms: u64,
updated_at_ms: u64,
started_at_ms: Option<u64>,
finished_at_ms: Option<u64>,
latest_progress: Option<SseProgressEvent>,
error: Option<String>,
result: Option<GenerateResponse>,
}
impl AsyncGenerationJob {
fn new(id: String, req: &GenerateRequest) -> Self {
let now = now_ms();
Self {
id,
status: AsyncJobStatus::Queued,
model: req.model.clone(),
prompt_preview: prompt_preview(&req.prompt),
created_at_ms: now,
updated_at_ms: now,
started_at_ms: None,
finished_at_ms: None,
latest_progress: None,
error: None,
result: None,
}
}
}
#[derive(Debug, Default)]
struct AsyncJobRegistryInner {
next_id: AtomicU64,
jobs: Mutex<HashMap<String, AsyncGenerationJob>>,
}
#[derive(Debug, Clone, Default)]
struct AsyncJobRegistry {
inner: Arc<AsyncJobRegistryInner>,
}
impl AsyncJobRegistry {
async fn create(&self, req: &GenerateRequest) -> String {
let id = format!(
"gen-{}",
self.inner.next_id.fetch_add(1, Ordering::Relaxed) + 1
);
let mut jobs = self.inner.jobs.lock().await;
jobs.insert(id.clone(), AsyncGenerationJob::new(id.clone(), req));
prune_completed_jobs(&mut jobs);
id
}
async fn mark_running(&self, id: &str) {
let mut jobs = self.inner.jobs.lock().await;
if let Some(job) = jobs.get_mut(id) {
let now = now_ms();
job.status = AsyncJobStatus::Running;
job.started_at_ms.get_or_insert(now);
job.updated_at_ms = now;
}
}
async fn record_progress(&self, id: &str, event: SseProgressEvent) {
let mut jobs = self.inner.jobs.lock().await;
if let Some(job) = jobs.get_mut(id) {
job.status = match event {
SseProgressEvent::Queued { .. } => AsyncJobStatus::Queued,
_ => AsyncJobStatus::Running,
};
job.latest_progress = Some(event);
job.updated_at_ms = now_ms();
}
}
async fn finish(&self, id: &str, result: std::result::Result<GenerateResponse, String>) {
let mut jobs = self.inner.jobs.lock().await;
if let Some(job) = jobs.get_mut(id) {
let now = now_ms();
job.finished_at_ms = Some(now);
job.updated_at_ms = now;
match result {
Ok(response) => {
job.status = AsyncJobStatus::Succeeded;
job.result = Some(response);
job.error = None;
}
Err(err) => {
job.status = AsyncJobStatus::Failed;
job.error = Some(err);
job.result = None;
}
}
}
prune_completed_jobs(&mut jobs);
}
async fn get(&self, id: &str) -> Option<AsyncGenerationJob> {
self.inner.jobs.lock().await.get(id).cloned()
}
async fn list(&self) -> Vec<AsyncGenerationJob> {
let mut jobs = self
.inner
.jobs
.lock()
.await
.values()
.cloned()
.collect::<Vec<_>>();
jobs.sort_by_key(|job| job.created_at_ms);
jobs
}
}
fn prune_completed_jobs(jobs: &mut HashMap<String, AsyncGenerationJob>) {
if jobs.len() <= MAX_ASYNC_JOBS {
return;
}
let mut removable = jobs
.values()
.filter(|job| job.status.is_terminal())
.map(|job| {
(
job.finished_at_ms.unwrap_or(job.updated_at_ms),
job.id.clone(),
)
})
.collect::<Vec<_>>();
removable.sort_by_key(|(ts, _)| *ts);
let remove_count = jobs.len().saturating_sub(MAX_ASYNC_JOBS);
for (_, id) in removable.into_iter().take(remove_count) {
jobs.remove(&id);
}
}
fn job_tool_result(job: &AsyncGenerationJob, include_result: bool) -> Value {
let mut content = vec![json!({ "type": "text", "text": job_status_line(job) })];
if include_result {
if let Some(response) = &job.result {
if let Some(image) = response.images.first() {
content.push(json!({
"type": "image",
"data": general_purpose::STANDARD.encode(&image.data),
"mimeType": image.format.content_type()
}));
}
}
}
json!({
"content": content,
"structuredContent": job_summary_json(job)
})
}
fn job_status_line(job: &AsyncGenerationJob) -> String {
let mut line = format!("{}: {} model {}", job.id, job.status.as_str(), job.model);
if let Some(progress) = &job.latest_progress {
line.push_str(&format!("; {}", progress_summary(progress)));
}
if let Some(error) = &job.error {
line.push_str(&format!("; error: {error}"));
}
if let Some(response) = &job.result {
if let Some(image) = response.images.first() {
line.push_str(&format!(
"; generated {}x{} {} in {:.1}s; seed {}",
image.width,
image.height,
image.format.extension(),
response.generation_time_ms as f64 / 1000.0,
response.seed_used
));
if let Some(gpu) = response.gpu {
line.push_str(&format!("; gpu {gpu}"));
}
}
}
line
}
fn job_summary_json(job: &AsyncGenerationJob) -> Value {
let result = job.result.as_ref().map(|response| {
let image = response.images.first();
json!({
"model": response.model,
"seed_used": response.seed_used,
"generation_time_ms": response.generation_time_ms,
"gpu": response.gpu,
"image": image.map(|image| json!({
"width": image.width,
"height": image.height,
"format": image.format.extension(),
"index": image.index
}))
})
});
json!({
"job_id": job.id,
"status": job.status.as_str(),
"model": job.model,
"prompt_preview": job.prompt_preview,
"created_at_ms": job.created_at_ms,
"updated_at_ms": job.updated_at_ms,
"started_at_ms": job.started_at_ms,
"finished_at_ms": job.finished_at_ms,
"latest_progress": job.latest_progress,
"error": job.error,
"result": result,
})
}
fn progress_summary(event: &SseProgressEvent) -> String {
match event {
SseProgressEvent::StageStart { name } => format!("stage started: {name}"),
SseProgressEvent::StageDone { name, elapsed_ms } => {
format!("stage done: {name} in {:.1}s", *elapsed_ms as f64 / 1000.0)
}
SseProgressEvent::Info { message } => message.clone(),
SseProgressEvent::CacheHit { resource } => format!("cache hit: {resource}"),
SseProgressEvent::DenoiseStep {
step,
total,
elapsed_ms,
} => format!(
"denoise step {step}/{total} after {:.1}s",
*elapsed_ms as f64 / 1000.0
),
SseProgressEvent::DownloadProgress {
filename,
file_index,
total_files,
bytes_downloaded,
bytes_total,
..
} => format!(
"downloading {filename} ({file_index}/{total_files}, {} / {})",
human_bytes(*bytes_downloaded),
human_bytes(*bytes_total)
),
SseProgressEvent::DownloadDone {
filename,
file_index,
total_files,
..
} => format!("downloaded {filename} ({file_index}/{total_files})"),
SseProgressEvent::PullComplete { model } => format!("pull complete: {model}"),
SseProgressEvent::Queued { position, id } => {
if id.is_empty() {
format!("queued at position {position}")
} else {
format!("queued at position {position}; server job {id}")
}
}
SseProgressEvent::WeightLoad {
bytes_loaded,
bytes_total,
component,
} => format!(
"loading {component} ({} / {})",
human_bytes(*bytes_loaded),
human_bytes(*bytes_total)
),
}
}
fn human_bytes(bytes: u64) -> String {
const MIB: f64 = 1024.0 * 1024.0;
const GIB: f64 = 1024.0 * 1024.0 * 1024.0;
let bytes = bytes as f64;
if bytes >= GIB {
format!("{:.1} GiB", bytes / GIB)
} else if bytes >= MIB {
format!("{:.1} MiB", bytes / MIB)
} else {
format!("{} B", bytes as u64)
}
}
fn prompt_preview(prompt: &str) -> String {
const MAX_CHARS: usize = 120;
let mut chars = prompt.chars();
let preview = chars.by_ref().take(MAX_CHARS).collect::<String>();
if chars.next().is_some() {
format!("{preview}...")
} else {
preview
}
}
fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
fn filter_sort_gallery(
mut images: Vec<GalleryImage>,
filter: &GalleryFilter,
) -> std::result::Result<Vec<GalleryImage>, String> {
let format = filter
.format
.as_deref()
.map(parse_gallery_format)
.transpose()?;
images.retain(|image| {
filter
.query
.as_deref()
.is_none_or(|query| gallery_matches_query(image, query))
&& filter
.filename
.as_deref()
.is_none_or(|name| contains_case_insensitive(&image.filename, name))
&& filter
.model
.as_deref()
.is_none_or(|model| contains_case_insensitive(&image.metadata.model, model))
&& format.is_none_or(|format| gallery_format(image) == Some(format))
});
let sort_by = filter.sort_by.as_deref().unwrap_or("timestamp");
match sort_by {
"timestamp" | "date" | "time" | "created_at" => {
images.sort_by_key(|image| image.timestamp);
}
"filename" | "name" => {
images.sort_by(|a, b| a.filename.cmp(&b.filename));
}
"model" => {
images.sort_by(|a, b| a.metadata.model.cmp(&b.metadata.model));
}
"prompt" | "description" => {
images.sort_by(|a, b| a.metadata.prompt.cmp(&b.metadata.prompt));
}
"size" | "size_bytes" => {
images.sort_by_key(|image| image.size_bytes.unwrap_or(0));
}
other => {
return Err(format!(
"unsupported sort_by '{other}'; use timestamp, filename, model, prompt, or size"
));
}
}
let descending = match filter.sort_order.as_deref().unwrap_or("desc") {
"desc" | "descending" => true,
"asc" | "ascending" => false,
other => {
return Err(format!("unsupported sort_order '{other}'; use asc or desc"));
}
};
if descending {
images.reverse();
}
Ok(images)
}
fn gallery_matches_query(image: &GalleryImage, query: &str) -> bool {
contains_case_insensitive(&image.filename, query)
|| contains_case_insensitive(&image.metadata.prompt, query)
|| image
.metadata
.original_prompt
.as_deref()
.is_some_and(|text| contains_case_insensitive(text, query))
|| image
.metadata
.negative_prompt
.as_deref()
.is_some_and(|text| contains_case_insensitive(text, query))
|| contains_case_insensitive(&image.metadata.model, query)
|| image
.metadata
.lora
.as_deref()
.is_some_and(|text| contains_case_insensitive(text, query))
|| contains_case_insensitive(&image.metadata.version, query)
}
fn contains_case_insensitive(haystack: &str, needle: &str) -> bool {
haystack.to_lowercase().contains(&needle.to_lowercase())
}
fn parse_gallery_format(format: &str) -> std::result::Result<OutputFormat, String> {
match format
.trim()
.trim_start_matches('.')
.to_lowercase()
.as_str()
{
"png" => Ok(OutputFormat::Png),
"jpg" | "jpeg" => Ok(OutputFormat::Jpeg),
"gif" => Ok(OutputFormat::Gif),
"apng" => Ok(OutputFormat::Apng),
"webp" => Ok(OutputFormat::Webp),
"mp4" => Ok(OutputFormat::Mp4),
other => Err(format!(
"unsupported format '{other}'; use png, jpeg, gif, apng, webp, or mp4"
)),
}
}
fn gallery_format(image: &GalleryImage) -> Option<OutputFormat> {
image.format.or_else(|| {
extension_for_filename(&image.filename).and_then(|ext| parse_gallery_format(ext).ok())
})
}
fn extension_for_filename(filename: &str) -> Option<&str> {
filename.rsplit_once('.').map(|(_, ext)| ext)
}
fn content_type_for_filename(filename: &str) -> &'static str {
extension_for_filename(filename)
.and_then(|ext| parse_gallery_format(ext).ok())
.map(|format| format.content_type())
.unwrap_or("application/octet-stream")
}
fn gallery_line(image: &GalleryImage) -> String {
let format = gallery_format(image)
.map(|format| format.extension().to_string())
.unwrap_or_else(|| "unknown".to_string());
let size = image
.size_bytes
.map(human_bytes)
.unwrap_or_else(|| "unknown size".to_string());
format!(
"{} [{}; {size}] {} model {}; {}x{}, steps {}, seed {}; prompt: {}",
image.filename,
format,
format_gallery_timestamp(image.timestamp),
image.metadata.model,
image.metadata.width,
image.metadata.height,
image.metadata.steps,
image.metadata.seed,
prompt_preview(&image.metadata.prompt)
)
}
fn gallery_summary_json(image: &GalleryImage) -> Value {
json!({
"filename": image.filename,
"timestamp": image.timestamp,
"datetime": format_gallery_timestamp(image.timestamp),
"format": gallery_format(image).map(|format| format.extension()),
"size_bytes": image.size_bytes,
"metadata_synthetic": image.metadata_synthetic,
"metadata": image.metadata,
})
}
fn lora_line(lora: &LoraInfo) -> String {
let author = lora
.author
.as_deref()
.map(|author| format!(" by {author}"))
.unwrap_or_default();
let triggers = if lora.trained_words.is_empty() {
String::new()
} else {
format!("; triggers: {}", lora.trained_words.join(", "))
};
format!(
"- {}{} [{}] {}; path: {}{}",
lora.name, author, lora.family, lora.id, lora.path, triggers
)
}
fn format_gallery_timestamp(timestamp: u64) -> String {
let secs = if timestamp > 1_000_000_000_000 {
timestamp / 1000
} else {
timestamp
};
time::OffsetDateTime::from_unix_timestamp(secs as i64)
.ok()
.and_then(|dt| {
dt.format(&time::format_description::well_known::Rfc3339)
.ok()
})
.unwrap_or_else(|| secs.to_string())
}
fn normalise_mcp_lora_arg(arg: McpLoraArg) -> std::result::Result<McpLoraRef, String> {
match arg {
McpLoraArg::Ref(value) => {
let value = value.trim().to_string();
if value.is_empty() {
return Err("loras entries must not be empty".to_string());
}
Ok(McpLoraRef { value, scale: 1.0 })
}
McpLoraArg::Object(object) => {
let id = non_empty_lora_ref(object.id);
let path = non_empty_lora_ref(object.path);
let value = match (id, path) {
(Some(_), Some(_)) => {
return Err(
"each loras object must provide either id or path, not both".to_string()
);
}
(Some(id), None) => id,
(None, Some(path)) => path,
(None, None) => {
return Err("each loras object must provide an id or path".to_string());
}
};
let scale = object.scale.unwrap_or(1.0);
if !scale.is_finite() {
return Err("LoRA scale must be a finite number".to_string());
}
Ok(McpLoraRef { value, scale })
}
}
}
fn non_empty_lora_ref(value: Option<String>) -> Option<String> {
value.and_then(|value| {
let value = value.trim().to_string();
(!value.is_empty()).then_some(value)
})
}
fn is_lora_catalog_id(value: &str) -> bool {
value.starts_with("cv:") || value.starts_with("hf:")
}
fn resolve_mcp_lora_refs(
refs: &[McpLoraRef],
catalog: &[LoraInfo],
) -> std::result::Result<Vec<LoraWeight>, String> {
refs.iter()
.map(|lora| {
let path = if is_lora_catalog_id(&lora.value) {
catalog
.iter()
.find(|candidate| candidate.id == lora.value)
.map(|candidate| candidate.path.clone())
.ok_or_else(|| {
format!(
"installed LoRA id '{}' was not found; use list_loras to see installed ids",
lora.value
)
})?
} else {
lora.value.clone()
};
Ok(LoraWeight {
path,
scale: lora.scale,
})
})
.collect()
}
fn build_generate_request(
args: GenerateImageArgs,
loras: Option<Vec<LoraWeight>>,
) -> std::result::Result<GenerateRequest, String> {
if args.prompt.trim().is_empty() {
return Err("prompt must not be empty".to_string());
}
let output_format = match args.output_format.as_deref().unwrap_or("png") {
"png" => OutputFormat::Png,
"jpeg" | "jpg" => OutputFormat::Jpeg,
other => {
return Err(format!(
"unsupported output_format '{other}'; use png or jpeg"
))
}
};
let mut config = Config::load_or_default();
let model = args
.model
.unwrap_or_else(|| config.resolved_default_model());
if crate::catalog_bridge::looks_like_catalog_id(&model) {
crate::catalog_bridge::install_catalog_model_from_installed_sidecar(&mut config, &model)
.map_err(|e| format!("failed to resolve installed catalog model '{model}': {e}"))?;
}
let model_cfg = config.resolved_model_config(&model);
let width = args
.width
.unwrap_or_else(|| model_cfg.effective_width(&config));
let height = args
.height
.unwrap_or_else(|| model_cfg.effective_height(&config));
if width == 0 || height == 0 {
return Err("width and height must be greater than zero".to_string());
}
if width & 15 != 0 || height & 15 != 0 {
return Err("width and height must be multiples of 16".to_string());
}
Ok(GenerateRequest {
prompt: args.prompt,
negative_prompt: args.negative_prompt,
model,
width,
height,
steps: args
.steps
.unwrap_or_else(|| model_cfg.effective_steps(&config)),
guidance: args
.guidance
.unwrap_or_else(|| model_cfg.effective_guidance()),
seed: args.seed,
batch_size: 1,
output_format: Some(output_format),
embed_metadata: Some(config.effective_embed_metadata(None)),
scheduler: None,
cfg_plus: None,
source_image: None,
edit_images: None,
strength: 0.75,
mask_image: None,
control_image: None,
control_model: None,
control_scale: 1.0,
expand: args.expand,
original_prompt: None,
lora: None,
frames: None,
fps: None,
upscale_model: None,
gif_preview: false,
enable_audio: None,
audio_file: None,
audio_file_path: None,
source_video: None,
source_video_path: None,
keyframes: None,
pipeline: None,
loras,
retake_range: None,
spatial_upscale: None,
temporal_upscale: None,
placement: None,
})
}
fn handle_protocol_message(message: Value) -> Option<Value> {
let id = message.get("id").cloned();
let method = message.get("method").and_then(Value::as_str)?;
match (method, id) {
("notifications/initialized", _) => None,
(_, None) => None,
("initialize", Some(id)) => Some(response(
id,
json!({
"protocolVersion": MCP_PROTOCOL_VERSION,
"capabilities": { "tools": {} },
"serverInfo": {
"name": "mold",
"version": mold_core::build_info::version_string()
}
}),
)),
("ping", Some(id)) => Some(response(id, json!({}))),
("tools/list", Some(id)) => Some(response(id, json!({ "tools": tool_definitions() }))),
("prompts/list", Some(id)) => Some(response(id, json!({ "prompts": [] }))),
("resources/list", Some(id)) => Some(response(id, json!({ "resources": [] }))),
(other, Some(id)) => Some(error_response(
id,
-32601,
format!("method not found: {other}"),
)),
}
}
fn tool_definitions() -> Value {
json!([
{
"name": "generate_image",
"description": "Generate one image with mold. Requires a running mold serve process, unless MOLD_HOST points at a remote mold server.",
"inputSchema": {
"type": "object",
"properties": {
"prompt": {
"type": "string",
"description": "The image prompt."
},
"model": {
"type": "string",
"description": "Optional mold model name, such as flux2-klein:q8. Defaults to mold's configured default model."
},
"width": {
"type": "integer",
"minimum": 16,
"multipleOf": 16,
"description": "Image width in pixels."
},
"height": {
"type": "integer",
"minimum": 16,
"multipleOf": 16,
"description": "Image height in pixels."
},
"steps": {
"type": "integer",
"minimum": 1,
"description": "Inference step count."
},
"guidance": {
"type": "number",
"description": "Guidance scale."
},
"seed": {
"type": "integer",
"minimum": 0,
"description": "Optional deterministic seed."
},
"negative_prompt": {
"type": "string",
"description": "Optional negative prompt for CFG-capable model families."
},
"output_format": {
"type": "string",
"enum": ["png", "jpeg", "jpg"],
"description": "Output image format. Defaults to png."
},
"expand": {
"type": "boolean",
"description": "Ask the mold server to expand the prompt before generation."
},
"loras": {
"type": "array",
"description": "Optional LoRA stack to apply. Each item can be an installed LoRA id from list_loras, a server-side path, or an object with id/path and optional scale.",
"items": {
"oneOf": [
{
"type": "string",
"description": "Installed LoRA id such as cv:827325, or a server-side .safetensors path. Scale defaults to 1.0."
},
{
"type": "object",
"properties": {
"id": {
"type": "string",
"description": "Installed LoRA catalog id, such as cv:827325."
},
"path": {
"type": "string",
"description": "Server-side path to a LoRA .safetensors file."
},
"scale": {
"type": "number",
"minimum": 0.0,
"maximum": 2.0,
"description": "LoRA strength. Defaults to 1.0."
}
},
"additionalProperties": false
}
]
}
}
},
"required": ["prompt"],
"additionalProperties": false
}
},
{
"name": "generate_image_async",
"description": "Start an image generation job and return immediately with a job id. Use generation_status to poll progress and fetch the completed image.",
"inputSchema": {
"type": "object",
"properties": {
"prompt": {
"type": "string",
"description": "The image prompt."
},
"model": {
"type": "string",
"description": "Optional mold model name, such as flux2-klein:q8. Defaults to mold's configured default model."
},
"width": {
"type": "integer",
"minimum": 16,
"multipleOf": 16,
"description": "Image width in pixels."
},
"height": {
"type": "integer",
"minimum": 16,
"multipleOf": 16,
"description": "Image height in pixels."
},
"steps": {
"type": "integer",
"minimum": 1,
"description": "Inference step count."
},
"guidance": {
"type": "number",
"description": "Guidance scale."
},
"seed": {
"type": "integer",
"minimum": 0,
"description": "Optional deterministic seed."
},
"negative_prompt": {
"type": "string",
"description": "Optional negative prompt for CFG-capable model families."
},
"output_format": {
"type": "string",
"enum": ["png", "jpeg", "jpg"],
"description": "Output image format. Defaults to png."
},
"expand": {
"type": "boolean",
"description": "Ask the mold server to expand the prompt before generation."
},
"loras": {
"type": "array",
"description": "Optional LoRA stack to apply. Each item can be an installed LoRA id from list_loras, a server-side path, or an object with id/path and optional scale.",
"items": {
"oneOf": [
{
"type": "string",
"description": "Installed LoRA id such as cv:827325, or a server-side .safetensors path. Scale defaults to 1.0."
},
{
"type": "object",
"properties": {
"id": {
"type": "string",
"description": "Installed LoRA catalog id, such as cv:827325."
},
"path": {
"type": "string",
"description": "Server-side path to a LoRA .safetensors file."
},
"scale": {
"type": "number",
"minimum": 0.0,
"maximum": 2.0,
"description": "LoRA strength. Defaults to 1.0."
}
},
"additionalProperties": false
}
]
}
}
},
"required": ["prompt"],
"additionalProperties": false
}
},
{
"name": "generation_status",
"description": "List async generation jobs or return status for one job. Completed image jobs include image content by default.",
"inputSchema": {
"type": "object",
"properties": {
"job_id": {
"type": "string",
"description": "Async generation job id returned by generate_image_async. Omit to list tracked jobs."
},
"include_result": {
"type": "boolean",
"description": "Include completed image content when job_id is provided. Defaults to true."
}
},
"additionalProperties": false
}
},
{
"name": "list_gallery",
"description": "List, search, filter, and sort saved mold gallery items.",
"inputSchema": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Case-insensitive search across filename, prompt/description, model, LoRA, and version."
},
"filename": {
"type": "string",
"description": "Case-insensitive filename substring filter."
},
"model": {
"type": "string",
"description": "Case-insensitive model substring filter."
},
"format": {
"type": "string",
"enum": ["png", "jpeg", "jpg", "gif", "apng", "webp", "mp4"],
"description": "Filter by output format."
},
"sort_by": {
"type": "string",
"enum": ["timestamp", "date", "time", "created_at", "filename", "name", "model", "prompt", "description", "size", "size_bytes"],
"description": "Sort field. Defaults to timestamp."
},
"sort_order": {
"type": "string",
"enum": ["desc", "descending", "asc", "ascending"],
"description": "Sort direction. Defaults to desc."
},
"limit": {
"type": "integer",
"minimum": 1,
"maximum": 200,
"description": "Maximum rows to return. Defaults to 25."
},
"offset": {
"type": "integer",
"minimum": 0,
"description": "Rows to skip after filtering/sorting. Defaults to 0."
}
},
"additionalProperties": false
}
},
{
"name": "get_gallery_image",
"description": "Fetch a saved gallery image by filename, or fetch the latest item matching search/filter arguments.",
"inputSchema": {
"type": "object",
"properties": {
"filename": {
"type": "string",
"description": "Exact gallery filename to fetch. If omitted, the latest matching item is fetched."
},
"latest": {
"type": "boolean",
"description": "When filename is omitted, choose the latest matching item. Defaults to true."
},
"query": {
"type": "string",
"description": "Case-insensitive search across filename, prompt/description, model, LoRA, and version."
},
"model": {
"type": "string",
"description": "Case-insensitive model substring filter."
},
"format": {
"type": "string",
"enum": ["png", "jpeg", "jpg", "gif", "apng", "webp", "mp4"],
"description": "Filter by output format."
},
"sort_by": {
"type": "string",
"enum": ["timestamp", "date", "time", "created_at", "filename", "name", "model", "prompt", "description", "size", "size_bytes"],
"description": "Sort field used when filename is omitted. Defaults to timestamp."
},
"sort_order": {
"type": "string",
"enum": ["desc", "descending", "asc", "ascending"],
"description": "Sort direction used when filename is omitted. Defaults to desc."
},
"thumbnail": {
"type": "boolean",
"description": "Return a PNG thumbnail instead of the full image. Defaults to true for MP4 gallery items and false otherwise."
}
},
"additionalProperties": false
}
},
{
"name": "list_models",
"description": "List mold models visible to the running mold server.",
"inputSchema": {
"type": "object",
"properties": {
"downloaded_only": {
"type": "boolean",
"description": "Only show downloaded models. Defaults to true."
},
"generation_only": {
"type": "boolean",
"description": "Hide upscalers, utility models, and auxiliary models. Defaults to true."
},
"limit": {
"type": "integer",
"minimum": 1,
"maximum": 200,
"description": "Maximum number of models to return. Defaults to 50."
}
},
"additionalProperties": false
}
},
{
"name": "list_loras",
"description": "List installed LoRA adapters. When model is provided, returns LoRAs compatible with that model family.",
"inputSchema": {
"type": "object",
"properties": {
"model": {
"type": "string",
"description": "Optional mold model name, such as flux-dev:q8 or realistic-vision-v5:fp16. Filters LoRAs to the model's compatible family."
}
},
"additionalProperties": false
}
},
{
"name": "server_status",
"description": "Show mold server health, queue, loaded model, and GPU status.",
"inputSchema": {
"type": "object",
"properties": {},
"additionalProperties": false
}
}
])
}
fn response(id: Value, result: Value) -> Value {
json!({
"jsonrpc": "2.0",
"id": id,
"result": result
})
}
fn error_response(id: Value, code: i64, message: impl Into<String>) -> Value {
json!({
"jsonrpc": "2.0",
"id": id,
"error": {
"code": code,
"message": message.into()
}
})
}
fn text_result(text: impl Into<String>) -> Value {
json!({
"content": [{ "type": "text", "text": text.into() }]
})
}
#[cfg(test)]
fn handle_protocol_message_for_test(message: Value) -> Option<Value> {
handle_protocol_message(message)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::ENV_LOCK;
use serde_json::json;
#[test]
fn initialize_declares_tools_capability() {
let response = handle_protocol_message_for_test(json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": { "protocolVersion": "2025-06-18" }
}))
.expect("initialize should produce a response");
assert_eq!(response["id"], 1);
assert_eq!(response["result"]["capabilities"]["tools"], json!({}));
assert_eq!(response["result"]["serverInfo"]["name"], "mold");
}
#[test]
fn tools_list_exposes_generate_image() {
let response = handle_protocol_message_for_test(json!({
"jsonrpc": "2.0",
"id": "tools",
"method": "tools/list"
}))
.expect("tools/list should produce a response");
let tools = response["result"]["tools"].as_array().unwrap();
assert!(tools.iter().any(|tool| tool["name"] == "generate_image"));
}
#[test]
fn tools_list_exposes_async_generation_tools() {
let response = handle_protocol_message_for_test(json!({
"jsonrpc": "2.0",
"id": "tools",
"method": "tools/list"
}))
.expect("tools/list should produce a response");
let tools = response["result"]["tools"].as_array().unwrap();
assert!(tools
.iter()
.any(|tool| tool["name"] == "generate_image_async"));
assert!(tools.iter().any(|tool| tool["name"] == "generation_status"));
}
#[test]
fn tools_list_exposes_gallery_tools() {
let response = handle_protocol_message_for_test(json!({
"jsonrpc": "2.0",
"id": "tools",
"method": "tools/list"
}))
.expect("tools/list should produce a response");
let tools = response["result"]["tools"].as_array().unwrap();
assert!(tools.iter().any(|tool| tool["name"] == "list_gallery"));
assert!(tools.iter().any(|tool| tool["name"] == "get_gallery_image"));
}
#[test]
fn tools_list_exposes_lora_stack_and_listing() {
let response = handle_protocol_message_for_test(json!({
"jsonrpc": "2.0",
"id": "tools",
"method": "tools/list"
}))
.expect("tools/list should produce a response");
let tools = response["result"]["tools"].as_array().unwrap();
let generate = tools
.iter()
.find(|tool| tool["name"] == "generate_image")
.expect("generate_image tool");
assert!(generate["inputSchema"]["properties"]["loras"].is_object());
assert!(
generate["inputSchema"]["properties"]["loras"]["items"]["oneOf"]
.as_array()
.is_some_and(|variants| variants.len() == 2)
);
assert!(tools.iter().any(|tool| tool["name"] == "list_loras"));
}
#[test]
fn normalise_lora_arg_accepts_ids_paths_and_default_scale() {
let parsed: GenerateImageArgs = serde_json::from_value(json!({
"prompt": "a cat",
"loras": [
"cv:827325",
{ "path": "/models/a.safetensors" },
{ "id": "cv:123", "scale": 0.65 }
]
}))
.unwrap();
assert_eq!(parsed.loras.as_ref().unwrap().len(), 3);
assert_eq!(
normalise_mcp_lora_arg(McpLoraArg::Ref("cv:827325".into())).unwrap(),
McpLoraRef {
value: "cv:827325".into(),
scale: 1.0,
}
);
assert_eq!(
normalise_mcp_lora_arg(McpLoraArg::Ref("/models/a.safetensors".into())).unwrap(),
McpLoraRef {
value: "/models/a.safetensors".into(),
scale: 1.0,
}
);
assert_eq!(
normalise_mcp_lora_arg(McpLoraArg::Object(McpLoraObject {
id: Some(" cv:123 ".into()),
path: None,
scale: Some(0.65),
}))
.unwrap(),
McpLoraRef {
value: "cv:123".into(),
scale: 0.65,
}
);
assert!(normalise_mcp_lora_arg(McpLoraArg::Object(McpLoraObject {
id: Some("cv:1".into()),
path: Some("/models/a.safetensors".into()),
scale: None,
}))
.is_err());
}
#[test]
fn resolve_lora_refs_maps_catalog_ids_and_keeps_paths() {
let catalog = vec![LoraInfo {
id: "cv:827325".into(),
name: "Detail".into(),
family: "flux".into(),
author: None,
path: "/models/cv-827325/detail.safetensors".into(),
trained_words: vec![],
size_bytes: Some(123),
thumbnail_url: None,
added_at: 10,
}];
let loras = resolve_mcp_lora_refs(
&[
McpLoraRef {
value: "cv:827325".into(),
scale: 0.7,
},
McpLoraRef {
value: "/models/manual.safetensors".into(),
scale: 1.0,
},
],
&catalog,
)
.unwrap();
assert_eq!(loras[0].path, "/models/cv-827325/detail.safetensors");
assert_eq!(loras[0].scale, 0.7);
assert_eq!(loras[1].path, "/models/manual.safetensors");
assert!(resolve_mcp_lora_refs(
&[McpLoraRef {
value: "cv:missing".into(),
scale: 1.0,
}],
&catalog,
)
.is_err());
}
#[test]
fn build_generate_request_preserves_lora_stack() {
let req = build_generate_request(
GenerateImageArgs {
prompt: "a cat".into(),
model: Some("flux-dev:q8".into()),
width: Some(1024),
height: Some(1024),
steps: Some(4),
guidance: Some(1.0),
seed: Some(42),
negative_prompt: None,
output_format: Some("png".into()),
expand: None,
loras: None,
},
Some(vec![
LoraWeight {
path: "/models/a.safetensors".into(),
scale: 0.8,
},
LoraWeight {
path: "/models/b.safetensors".into(),
scale: 0.4,
},
]),
)
.expect("valid generate request");
let loras = req.loras.expect("lora stack");
assert_eq!(loras.len(), 2);
assert_eq!(loras[0].path, "/models/a.safetensors");
assert_eq!(loras[1].scale, 0.4);
}
#[test]
fn build_generate_request_uses_installed_catalog_sidecar_defaults() {
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior_home = std::env::var_os("MOLD_HOME");
let prior_models_dir = std::env::var_os("MOLD_MODELS_DIR");
let dir = tempfile::tempdir().expect("tempdir");
unsafe {
std::env::set_var("MOLD_HOME", dir.path());
std::env::remove_var("MOLD_MODELS_DIR");
}
let install_dir = dir.path().join("models/cv-12345");
let primary_rel = "sdxl/civitai/12345/model.safetensors";
let primary = install_dir.join(primary_rel);
std::fs::create_dir_all(primary.parent().unwrap()).unwrap();
std::fs::write(&primary, b"fake").unwrap();
let sidecar = mold_catalog::sidecar::CatalogSidecar {
schema: mold_catalog::sidecar::SIDECAR_SCHEMA,
id: "cv:12345".to_string(),
source: "civitai".to_string(),
source_id: "12345".to_string(),
name: "SDXL Fine Tune".to_string(),
author: Some("tester".to_string()),
family: "sdxl".to_string(),
family_role: "finetune".to_string(),
sub_family: None,
kind: "checkpoint".to_string(),
modality: "image".to_string(),
thumbnail_url: None,
size_bytes: Some(4),
engine_phase: 1,
trained_words: Vec::new(),
primary_filename_rel: primary_rel.to_string(),
written_at: 0,
};
mold_catalog::sidecar::write_sidecar(
&install_dir.join(mold_catalog::sidecar::SIDECAR_FILENAME),
&sidecar,
)
.unwrap();
let req = build_generate_request(
GenerateImageArgs {
prompt: "a catalog default test".into(),
model: Some("cv:12345".into()),
width: None,
height: None,
steps: None,
guidance: None,
seed: None,
negative_prompt: None,
output_format: None,
expand: None,
loras: None,
},
None,
)
.expect("catalog sidecar should supply model defaults");
assert_eq!(req.width, 1024);
assert_eq!(req.height, 1024);
assert_eq!(req.steps, 25);
assert_eq!(req.guidance, 7.5);
unsafe {
match prior_home {
Some(v) => std::env::set_var("MOLD_HOME", v),
None => std::env::remove_var("MOLD_HOME"),
}
match prior_models_dir {
Some(v) => std::env::set_var("MOLD_MODELS_DIR", v),
None => std::env::remove_var("MOLD_MODELS_DIR"),
}
}
}
#[tokio::test]
async fn async_job_registry_tracks_completed_image() {
let jobs = AsyncJobRegistry::default();
let req = GenerateRequest {
prompt: "a cat".into(),
negative_prompt: None,
model: "flux2-klein:q8".into(),
width: 16,
height: 16,
steps: 1,
guidance: 0.0,
seed: Some(7),
batch_size: 1,
output_format: Some(OutputFormat::Png),
embed_metadata: Some(false),
scheduler: None,
cfg_plus: None,
source_image: None,
edit_images: None,
strength: 0.75,
mask_image: None,
control_image: None,
control_model: None,
control_scale: 1.0,
expand: None,
original_prompt: None,
lora: None,
frames: None,
fps: None,
upscale_model: None,
gif_preview: false,
enable_audio: None,
audio_file: None,
audio_file_path: None,
source_video: None,
source_video_path: None,
keyframes: None,
pipeline: None,
loras: None,
retake_range: None,
spatial_upscale: None,
temporal_upscale: None,
placement: None,
};
let id = jobs.create(&req).await;
jobs.mark_running(&id).await;
jobs.finish(
&id,
Ok(GenerateResponse {
images: vec![mold_core::ImageData {
data: vec![1, 2, 3],
format: OutputFormat::Png,
width: 16,
height: 16,
index: 0,
}],
video: None,
generation_time_ms: 42,
model: "flux2-klein:q8".into(),
seed_used: 7,
gpu: Some(0),
}),
)
.await;
let job = jobs.get(&id).await.expect("job should exist");
assert_eq!(job.status, AsyncJobStatus::Succeeded);
assert_eq!(job.result.unwrap().images[0].data, vec![1, 2, 3]);
}
#[test]
fn gallery_filter_searches_prompt_and_sorts_by_timestamp_desc() {
let images = vec![
test_gallery_image("older.png", "flux-dev:q8", "a red chair", 10),
test_gallery_image("newer.png", "sdxl:fp16", "a blue chair", 20),
test_gallery_image("other.png", "flux-dev:q8", "a yellow lamp", 30),
];
let filter = GalleryFilter {
query: Some("chair".into()),
sort_by: Some("timestamp".into()),
sort_order: Some("desc".into()),
..Default::default()
};
let filtered = filter_sort_gallery(images, &filter).unwrap();
let names = filtered
.iter()
.map(|image| image.filename.as_str())
.collect::<Vec<_>>();
assert_eq!(names, vec!["newer.png", "older.png"]);
}
fn test_gallery_image(
filename: &str,
model: &str,
prompt: &str,
timestamp: u64,
) -> GalleryImage {
GalleryImage {
filename: filename.into(),
metadata: mold_core::OutputMetadata {
prompt: prompt.into(),
negative_prompt: None,
original_prompt: None,
model: model.into(),
seed: 1,
steps: 4,
guidance: 0.0,
width: 768,
height: 768,
strength: None,
scheduler: None,
output_format: Some(OutputFormat::Png),
cfg_plus: None,
lora: None,
lora_scale: None,
loras: None,
control_model: None,
control_scale: None,
upscale_model: None,
gif_preview: None,
enable_audio: None,
audio_file_path: None,
source_video_path: None,
pipeline: None,
retake_range: None,
spatial_upscale: None,
temporal_upscale: None,
frames: None,
fps: None,
version: "test".into(),
},
timestamp,
format: Some(OutputFormat::Png),
size_bytes: Some(123),
metadata_synthetic: false,
}
}
}