From 97d9301a7154e5e5264f387bffe47ae9f70cdda2 Mon Sep 17 00:00:00 2001 From: Cameron Cordes Date: Fri, 30 Oct 2026 15:33:23 -0400 Subject: [PATCH] feat(enhance): discard, tests against a mock ComfyUI, docs (#192) Co-Authored-By: Claude Opus 5.5 --- .gitignore | 1 + CLAUDE.md | 89 +++ src/database/enhance_dao.rs | 13 +- src/enhance/comfy.rs | 9 +- src/enhance/config.rs | 4 +- src/enhance/handlers.rs | 34 +- src/enhance/integration_tests.rs | 985 +++++++++++++++++++++++++++++++ src/enhance/path_map.rs | 3 +- src/enhance/service.rs | 45 +- src/enhance/template.rs | 30 +- 10 files changed, 1165 insertions(+), 48 deletions(-) diff --git a/.gitignore b/.gitignore index 969edff..9083d82 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ database/target .env # Server-local LLM config (llm.example.toml is the template) /llm.toml +/enhance.toml # Server-local TTS pronunciation overrides (tts_pronunciations.example.json is the template) /tts_pronunciations.json /tmp diff --git a/CLAUDE.md b/CLAUDE.md index 7320b69..b059cd0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -424,6 +424,94 @@ for the `link_versions` binary). against the current hash; `unresolved` rows are listed, never guessed. `--apply` takes a `VACUUM INTO` snapshot first. +### Enhance on demand (#192) + +An **Enhance** action runs a ComfyUI workflow (a *template*) on one photo; the +user previews the result in the app and **approves** it, and it lands in the +library as a photo version (#191). Nothing reaches the library without an +approval. Code: `src/enhance/` (`config.rs`, `path_map.rs`, `template.rs`, +`comfy.rs`, `service.rs`, `handlers.rs`), table `enhance_jobs` +(`src/database/enhance_dao.rs`). The apps only talk to ImageApi; ComfyUI has no +auth and stays LAN-only, so ImageApi relays previews too. + +- **Config: `enhance.toml`** (path from `ENHANCE_CONFIG_PATH`, else `./enhance.toml`; + gitignored; `enhance.example.toml` is the annotated schema and a test parses it). + No secrets live in it. Without `url` (or without the file) Enhance is off: + `GET /enhance/templates` answers `available: false`. A malformed file fails startup; + a *template* that doesn't validate is logged and skipped, the rest load. + Keys: `url`, `timeout_secs`, `probe_ttl_secs`, `poll_secs`, `idle_free_secs`, + `unreachable_grace_secs`, `default_template`, `warn_above_megapixels`, + `[[path_map]]` (`from`, `to`), `[[template]]` (`name`, `label`, `description`, + `file`, `prompt_mode`, `prompt_join`, `output_megapixels`). +- **Templates** are API-format workflow exports, read-only (never edit them from + here). Nodes are found by `_meta.title`: `imageapi:source` (required, gets the + photo path), `imageapi:save` (required, forced to `require_approval = true`), + `imageapi:prompt` (optional; fills the first of `text` / `prompt` / `value`), + `imageapi:reference` (optional). `prompt_mode = "append"`: the user's text is joined + after the template's own prompt with `prompt_join` (default one space); empty text + leaves the prompt as is. `"replace"`: the text *is* the prompt and is required (empty + = 400); it needs a prompt node. With no `reference_path` the reference node **and + every `["", n]` input linking to it** are removed (its leftover export path must + never be used); if ComfyUI then rejects the prompt with `node_errors` the error says + "this template needs a reference photo". Every `seed`, `noise_seed` and `*.seed` + input gets a fresh random value per submit, or a retry would return ComfyUI's cached + result. A stray `SaveImage` is only a startup warning. +- **Path map**: library prefix <-> ComfyUI prefix (`/media/nas/` <-> `Z:\`), longest + prefix wins, separators converted, case-insensitive on the ComfyUI side; used on + submit and to recognise the path `/imageapi/approve` returns. No entries = identity. + A mapped drive only exists in the logged-in session; use UNC if ComfyUI becomes a + service. +- **Jobs** (`enhance_jobs`, keyed by the submitted file's `content_hash`): status + `queued | running | ready | approved | failed | expired | unavailable`. The poller + (`EnhanceService::spawn`, every `poll_secs`, active jobs only) reads `/queue` then + `/history/{id}`: in queue -> queued/running with `queue_position`; finished -> + `ready` (staged results in `pending`) or `failed` (ComfyUI's message); known to + neither -> **`expired`** (ComfyUI restarted and cleared its temp folder). An + unreachable ComfyUI for `unreachable_grace_secs` -> `unavailable`. A preview or + approve that finds the staged result gone also expires the job (HTTP 410). +- **Idle `/free`**: every 30 s, if the queue is empty, none of our jobs is active and + ComfyUI's newest `/history` entry has not changed for `idle_free_secs` (it includes + hand-run workflows), `POST /free {unload_models, free_memory}` is sent once per quiet + spell. "Not changed for" is measured on ImageApi's clock (no cross-host skew); the + first observation after ImageApi starts counts as activity. `idle_free_secs = 0` = off. +- **Approve** (`POST /enhance/jobs/{id}/approve`): ComfyUI's `/imageapi/approve` writes + the file (ImageApi writes nothing under library roots). The returned path is mapped + back, then a **targeted ingest** (`watcher::register_new_files` + `versions::link_new_files`) + registers and links just that file, so it is in `/image/versions` before the response. + Idempotent: an already-saved result returns its version without calling ComfyUI. +- **Output size**: Turbo output is about 1 MP regardless of the original, so + `output_megapixels` / `warn_above_megapixels` let clients warn above ~1.5 MP. + +Endpoints (all need a bearer token; errors are `{"error": msg, "code": ...}`; `library` +is accepted as a query or body field like the versions endpoints): + +``` +GET /enhance/templates + -> {"available": true, "reason"?: "ComfyUI isn't running" | "Enhance isn't configured", + "default": "restore", "warn_above_megapixels": 1.5, + "templates": [{"name","label","description","accepts_prompt", + "prompt_mode": "append"|"replace", "prompt_required", + "accepts_reference","default","output_megapixels", + "warn_above_megapixels"}]} +POST /enhance {"path","template","prompt"?,"reference_path"?,"library"?} -> job +GET /enhance/jobs/{id} -> job +GET /enhance/jobs?path=&library= -> {"jobs": [job, ...]} newest first, discarded hidden +GET /enhance/jobs/{id}/preview/{n} -> the staged PNG (relayed from ComfyUI /view) +POST /enhance/jobs/{id}/approve {"n"} -> a /image/versions `versions[]` entry +POST /enhance/jobs/{id}/discard -> job (hidden from the list; ComfyUI untouched) + +job = {"id","status","queue_position"?,"error"?,"path","library_id","template", + "prompt"?,"reference_path"?,"created_at","updated_at", + "pending": [{"n","preview": "/enhance/jobs/{id}/preview/{n}","approved"}], + "saved": [{"n","path","library_id"}]} +``` + +Error statuses: 400 bad request / empty required prompt / job not ready; 404 unknown +template, photo or job; 410 `expired`; 422 ComfyUI rejected the workflow; 502 upstream +or an approved file outside every library/path_map; 503 `unavailable` (ComfyUI down) or +`disabled`. `pending` is empty unless the job is `ready`/`approved`. Approve answers 202 +`{"linked": false, ...}` if the file was saved but couldn't be ingested synchronously. + ### File Processing Pipeline **Thumbnail Generation:** @@ -632,6 +720,7 @@ GET /image/versions?path=...&library=... (path of the original or any version { original:{path,library_id,hash}, versions:[{path, library_id,hash,kind,source,created_at,references, width,height}], display_hash }) +POST /enhance ... (Enhance on demand; see "Enhance on demand") PUT /image/versions/display (body { path, display_hash|null, library? }; global choice) // /photos hides versions behind their original (include_versions=true lists them) and adds // parallel arrays photo_display_paths / photo_version_counts; memories items, search hits and diff --git a/src/database/enhance_dao.rs b/src/database/enhance_dao.rs index aa90569..67633c3 100644 --- a/src/database/enhance_dao.rs +++ b/src/database/enhance_dao.rs @@ -93,8 +93,11 @@ pub trait EnhanceJobDao: Sync + Send { fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError>; fn get(&mut self, id: &str) -> Result, DbError>; /// Newest first. - fn list_for_hash(&mut self, original_hash: &str, limit: i64) - -> Result, DbError>; + fn list_for_hash( + &mut self, + original_hash: &str, + limit: i64, + ) -> Result, DbError>; /// Jobs the poller still has to advance (`queued` / `running`). fn list_active(&mut self) -> Result, DbError>; fn update(&mut self, id: &str, update: &JobUpdate) -> Result<(), DbError>; @@ -325,7 +328,11 @@ mod tests { assert!(a.updated_at >= 100); d.discard("b").unwrap(); - assert_eq!(d.list_for_hash("h1", 10).unwrap().len(), 1, "discarded hidden"); + assert_eq!( + d.list_for_hash("h1", 10).unwrap().len(), + 1, + "discarded hidden" + ); // Prune drops only old, finished jobs (updated_at was just bumped // for `a`; `c` is active). diff --git a/src/enhance/comfy.rs b/src/enhance/comfy.rs index 4102410..3189f65 100644 --- a/src/enhance/comfy.rs +++ b/src/enhance/comfy.rs @@ -64,7 +64,9 @@ impl std::fmt::Display for SubmitError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { SubmitError::Comfy(e) => e.fmt(f), - SubmitError::Rejected { message, .. } => write!(f, "ComfyUI rejected the workflow: {message}"), + SubmitError::Rejected { message, .. } => { + write!(f, "ComfyUI rejected the workflow: {message}") + } } } } @@ -375,7 +377,10 @@ fn rejection_message(json: &Value, raw: &str) -> String { pub(crate) fn parse_history_entry(entry: &Value) -> HistoryEntry { let status = &entry["status"]; - let status_str = status["status_str"].as_str().unwrap_or_default().to_string(); + let status_str = status["status_str"] + .as_str() + .unwrap_or_default() + .to_string(); let completed = status["completed"].as_bool().unwrap_or(false); let messages = status["messages"].as_array().cloned().unwrap_or_default(); let error = messages.iter().find_map(|m| { diff --git a/src/enhance/config.rs b/src/enhance/config.rs index dd1f34f..1cac49f 100644 --- a/src/enhance/config.rs +++ b/src/enhance/config.rs @@ -127,7 +127,7 @@ impl EnhanceConfig { bail!("{key} must be greater than 0"); } } - if !(self.warn_above_megapixels > 0.0) { + if self.warn_above_megapixels.is_nan() || self.warn_above_megapixels <= 0.0 { bail!("warn_above_megapixels must be greater than 0"); } let mut seen = std::collections::HashSet::new(); @@ -148,7 +148,7 @@ impl EnhanceConfig { if t.label.trim().is_empty() { bail!("template {:?} needs a label", t.name); } - if !(t.output_megapixels > 0.0) { + if t.output_megapixels.is_nan() || t.output_megapixels <= 0.0 { bail!("template {:?}: output_megapixels must be > 0", t.name); } } diff --git a/src/enhance/handlers.rs b/src/enhance/handlers.rs index 3f91989..1757b68 100644 --- a/src/enhance/handlers.rs +++ b/src/enhance/handlers.rs @@ -180,11 +180,7 @@ fn resolve_file( /// Library and relative path of an absolute file ComfyUI wrote. The job's own /// library wins when roots overlap. -fn place_in_libraries( - libs: &[Library], - prefer: i32, - abs: &Path, -) -> Option<(i32, String)> { +fn place_in_libraries(libs: &[Library], prefer: i32, abs: &Path) -> Option<(i32, String)> { let mut ordered: Vec<&Library> = libs.iter().collect(); ordered.sort_by_key(|l| l.id != prefer); for lib in ordered { @@ -292,10 +288,16 @@ pub async fn submit( Ok(f) => f, Err(r) => return r, }; - let reference = match body.reference_path.as_deref().filter(|p| !p.trim().is_empty()) { + let reference = match body + .reference_path + .as_deref() + .filter(|p| !p.trim().is_empty()) + { Some(p) => { // The reference lives in the photo's own library. - let lib = app_state.library_by_id(photo.library_id).map(|l| l.name.clone()); + let lib = app_state + .library_by_id(photo.library_id) + .map(|l| l.name.clone()); match resolve_file(&app_state, &exif_dao, p, lib.as_deref()) { Ok(f) => Some(f), Err(r) => return r, @@ -304,7 +306,12 @@ pub async fn submit( None => None, }; match svc - .submit(&photo, &body.template, body.prompt.as_deref(), reference.as_ref()) + .submit( + &photo, + &body.template, + body.prompt.as_deref(), + reference.as_ref(), + ) .await { Ok(job) => HttpResponse::Ok().json(job_view(&job)), @@ -442,10 +449,8 @@ pub async fn approve( let version_dao = app_state.version_dao.clone(); let exif_arc: Arc>> = exif_dao.clone().into_inner(); let (lib2, rel2, abs2) = (library.clone(), rel.clone(), abs.clone()); - let ingested = web::block(move || { - ingest_saved(&lib2, &rel2, &abs2, &exif_arc, &version_dao) - }) - .await; + let ingested = + web::block(move || ingest_saved(&lib2, &rel2, &abs2, &exif_arc, &version_dao)).await; let hash = match ingested { Ok(Some(h)) => h, _ => { @@ -461,7 +466,10 @@ pub async fn approve( // Which version row did that become? let entry = { - let mut dao = app_state.version_dao.lock().expect("Unable to lock VersionDao"); + let mut dao = app_state + .version_dao + .lock() + .expect("Unable to lock VersionDao"); let original = match dao.get_link(&hash) { Ok(Some(link)) => link.original_hash, _ => job.original_hash.clone(), diff --git a/src/enhance/integration_tests.rs b/src/enhance/integration_tests.rs index e69de29..6abe693 100644 --- a/src/enhance/integration_tests.rs +++ b/src/enhance/integration_tests.rs @@ -0,0 +1,985 @@ +//! End-to-end checks for Enhance (#192) against a mock ComfyUI: nothing here +//! touches a real ComfyUI, a real library or the prod DB. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use actix_web::test as awtest; +use actix_web::web::Data; +use actix_web::{App, HttpResponse, HttpServer, web}; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::Text; +use diesel::sqlite::SqliteConnection; +use serde_json::{Value, json}; + +use super::config::{EnhanceConfig, PathMapping, PromptMode, TemplateConfig}; +use super::service::EnhanceService; +use super::template::Template; +use crate::data::Claims; +use crate::database::enhance_dao::{EnhanceJobDao, SqliteEnhanceJobDao}; +use crate::database::test::in_memory_db_connection; +use crate::database::{ExifDao, SqliteExifDao, SqliteVersionDao}; +use crate::state::AppState; +use crate::versions::xmp::test_png::{producer_xmp, write_png}; + +const RESTORE: &str = include_str!("../../tests/fixtures/enhance/restore_turbo.json"); +const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json"); + +// ---- mock ComfyUI --------------------------------------------------------- + +#[derive(Default)] +struct Mock { + submitted: Mutex>, + queue_running: Mutex>, + queue_pending: Mutex>, + /// prompt_id -> history entry + history: Mutex>, + newest: Mutex>, + approve_calls: Mutex>, + approve_path: Mutex, + approve_gone: Mutex, + free_calls: AtomicUsize, + reject_prompts: Mutex>, + view_404: Mutex, +} + +async fn start_mock() -> (Arc, String) { + let mock = Arc::new(Mock::default()); + let m = mock.clone(); + let server = HttpServer::new(move || { + App::new() + .app_data(Data::from(m.clone())) + .route("/system_stats", web::get().to(|| async { HttpResponse::Ok().json(json!({})) })) + .route( + "/prompt", + web::post().to(|m: Data, body: web::Json| async move { + if let Some(err) = m.reject_prompts.lock().unwrap().clone() { + return HttpResponse::BadRequest().json(err); + } + let mut s = m.submitted.lock().unwrap(); + s.push(body.into_inner()); + let id = format!("prompt-{}", s.len()); + m.queue_pending.lock().unwrap().push(id.clone()); + HttpResponse::Ok().json(json!({"prompt_id": id, "number": s.len(), "node_errors": {}})) + }), + ) + .route( + "/queue", + web::get().to(|m: Data| async move { + let item = |id: &String| json!([1, id, {}, {}, []]); + HttpResponse::Ok().json(json!({ + "queue_running": m.queue_running.lock().unwrap().iter().map(item).collect::>(), + "queue_pending": m.queue_pending.lock().unwrap().iter().map(item).collect::>(), + })) + }), + ) + .route( + "/history", + web::get().to(|m: Data| async move { + match m.newest.lock().unwrap().clone() { + Some(e) => HttpResponse::Ok().json(json!({"newest": e})), + None => HttpResponse::Ok().json(json!({})), + } + }), + ) + .route( + "/history/{id}", + web::get().to(|m: Data, id: web::Path| async move { + match m.history.lock().unwrap().get(id.as_str()) { + Some(e) => HttpResponse::Ok().json(json!({ id.as_str(): e })), + None => HttpResponse::Ok().json(json!({})), + } + }), + ) + .route( + "/view", + web::get().to(|m: Data, q: web::Query>| async move { + if *m.view_404.lock().unwrap() { + return HttpResponse::NotFound().finish(); + } + assert_eq!(q.get("type").map(String::as_str), Some("temp")); + HttpResponse::Ok() + .content_type("image/png") + .body(format!("PNG:{}", q["filename"])) + }), + ) + .route( + "/imageapi/approve", + web::post().to(|m: Data, body: web::Json| async move { + if *m.approve_gone.lock().unwrap() { + return HttpResponse::NotFound().json(json!({"error": "no longer pending"})); + } + m.approve_calls + .lock() + .unwrap() + .push(body["id"].as_str().unwrap().to_string()); + HttpResponse::Ok().json(json!({"path": m.approve_path.lock().unwrap().clone()})) + }), + ) + .route( + "/free", + web::post().to(|m: Data, body: web::Json| async move { + assert_eq!(body["unload_models"], true); + m.free_calls.fetch_add(1, Ordering::SeqCst); + HttpResponse::Ok().finish() + }), + ) + }) + .workers(1) + .bind(("127.0.0.1", 0)) + .unwrap(); + let addr = server.addrs()[0]; + actix_rt::spawn(server.run()); + (mock, format!("http://{addr}")) +} + +fn finished_entry(pending_id: &str, ts: i64) -> Value { + json!({ + "status": {"status_str": "success", "completed": true, + "messages": [["execution_start", {"timestamp": ts - 3000}], ["execution_success", {"timestamp": ts}]]}, + "outputs": {"36": { + "images": [{"filename": format!("imageapi_pending_{pending_id}.png"), "subfolder": "", "type": "temp"}], + "imageapi_versions": [{"id": pending_id, "saved_path": null}], + }}, + }) +} + +// ---- fixture -------------------------------------------------------------- + +struct Fx { + mock: Arc, + svc: Data, + state: Data, + exif: Data>>, + conn: Arc>, + root: PathBuf, +} + +fn template(name: &str, mode: PromptMode, json: &str) -> Template { + Template::from_value( + TemplateConfig { + name: name.into(), + label: name.to_uppercase(), + description: format!("the {name} template"), + file: "x.json".into(), + prompt_mode: mode, + prompt_join: " ".into(), + output_megapixels: 1.0, + }, + serde_json::from_str(json).unwrap(), + ) + .unwrap() +} + +async fn fx_with(url: Option, tweak: impl FnOnce(&mut EnhanceConfig)) -> Fx { + let (mock, mock_url) = start_mock().await; + let mut state = AppState::test_state(); + let conn = Arc::new(Mutex::new(in_memory_db_connection())); + state.version_dao = Arc::new(Mutex::new(Box::new(SqliteVersionDao::from_shared( + conn.clone(), + )))); + let library = state.primary_library().clone(); + let root = PathBuf::from(&library.root_path); + sql_query("UPDATE libraries SET root_path = ? WHERE id = 1") + .bind::(&library.root_path) + .execute(&mut *conn.lock().unwrap()) + .unwrap(); + + let mut cfg = EnhanceConfig { + url: Some(url.unwrap_or(mock_url)), + path_map: vec![PathMapping { + from: format!("{}/", library.root_path), + to: "Z:\\".into(), + }], + ..EnhanceConfig::default() + }; + tweak(&mut cfg); + let dao: Arc>> = Arc::new(Mutex::new(Box::new( + SqliteEnhanceJobDao::from_shared(conn.clone()), + ))); + let svc = EnhanceService::new( + cfg, + vec![ + template("restore", PromptMode::Append, RESTORE), + template("edit", PromptMode::Replace, EDIT), + ], + dao, + ) + .unwrap(); + Fx { + mock, + svc: Data::new(svc), + exif: Data::new(Mutex::new( + Box::new(SqliteExifDao::from_shared(conn.clone())) as Box, + )), + state: Data::new(state), + conn, + root, + } +} + +async fn fx() -> Fx { + fx_with(None, |_| {}).await +} + +impl Fx { + fn put(&self, rel: &str, bytes: &[u8]) -> String { + let abs = self.root.join(rel); + std::fs::create_dir_all(abs.parent().unwrap()).unwrap(); + std::fs::write(&abs, bytes).unwrap(); + self.register(rel) + } + + fn register(&self, rel: &str) -> String { + let hash = crate::content_hash::compute(&self.root.join(rel)) + .unwrap() + .content_hash; + sql_query( + "INSERT INTO image_exif (library_id, rel_path, created_time, last_modified, content_hash) \ + VALUES (1, ?, 0, 0, ?)", + ) + .bind::(rel) + .bind::(&hash) + .execute(&mut *self.conn.lock().unwrap()) + .unwrap(); + hash + } + + fn token() -> String { + jsonwebtoken::encode( + &jsonwebtoken::Header::default(), + &Claims::valid_user("1".to_string()), + &jsonwebtoken::EncodingKey::from_secret(b"test_key"), + ) + .unwrap() + } + + async fn call(&self, method: &str, uri: &str, body: Option) -> (u16, Value, Vec) { + let app = awtest::init_service( + App::new() + .service(super::get_templates) + .service(super::submit) + .service(super::list_jobs) + .service(super::get_job) + .service(super::preview) + .service(super::approve) + .service(super::discard) + .service(crate::versions::get_versions) + .app_data(self.svc.clone()) + .app_data(self.state.clone()) + .app_data(self.exif.clone()), + ) + .await; + let mut req = match method { + "POST" => awtest::TestRequest::post(), + _ => awtest::TestRequest::get(), + } + .uri(uri) + .insert_header(("Authorization", format!("Bearer {}", Fx::token()))); + if let Some(b) = body { + req = req.set_json(b); + } + let resp = awtest::call_service(&app, req.to_request()).await; + let status = resp.status().as_u16(); + let bytes = awtest::read_body(resp).await.to_vec(); + let json = serde_json::from_slice(&bytes).unwrap_or(Value::Null); + (status, json, bytes) + } +} + +async fn call(fx: &Fx, method: &str, uri: &str, body: Option) -> (u16, Value, Vec) { + fx.call(method, uri, body).await +} + +const PHOTO: &str = "2009_200/019_19.JPG"; + +// ---- tests ---------------------------------------------------------------- + +#[actix_rt::test] +async fn templates_report_availability_and_shape() { + let fx = fx().await; + let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await; + assert_eq!(st, 200); + assert_eq!(body["available"], true); + assert!(body.get("reason").is_none()); + assert_eq!(body["default"], "restore"); + let t = &body["templates"]; + assert_eq!(t[0]["name"], "restore"); + assert_eq!(t[0]["label"], "RESTORE"); + assert_eq!(t[0]["accepts_prompt"], true); + assert_eq!(t[0]["prompt_mode"], "append"); + assert_eq!(t[0]["prompt_required"], false); + assert_eq!(t[0]["accepts_reference"], false); + assert_eq!(t[0]["default"], true); + assert_eq!(t[0]["output_megapixels"], 1.0); + assert_eq!(t[0]["warn_above_megapixels"], 1.5); + assert_eq!(t[1]["name"], "edit"); + assert_eq!(t[1]["prompt_mode"], "replace"); + assert_eq!(t[1]["prompt_required"], true); + assert_eq!(t[1]["accepts_reference"], true); + assert_eq!(t[1]["default"], false); +} + +#[actix_rt::test] +async fn comfy_down_means_unavailable_and_a_clear_submit_error() { + // Nothing listens on port 9 (discard) of loopback. + let fx = fx_with(Some("http://127.0.0.1:9".into()), |_| {}).await; + fx.put(PHOTO, b"photo"); + let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await; + assert_eq!(st, 200); + assert_eq!(body["available"], false); + assert_eq!(body["reason"], "ComfyUI isn't running"); + assert_eq!(body["templates"].as_array().unwrap().len(), 2); + + let (st, body, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + assert_eq!(st, 503); + assert_eq!(body["error"], "ComfyUI isn't running"); + assert_eq!(body["code"], "unavailable"); +} + +#[actix_rt::test] +async fn unconfigured_enhance_is_off_not_an_error() { + let fx = fx_with(None, |c| c.url = None).await; + let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await; + assert_eq!((st, &body["available"]), (200, &json!(false))); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + assert_eq!(st, 503); +} + +#[actix_rt::test] +async fn submit_maps_the_path_and_uses_a_fresh_seed_each_time() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let body = json!({"path": PHOTO, "template": "restore", "prompt": " warmer colours "}); + let (st, job1, _) = call(&fx, "POST", "/enhance", Some(body.clone())).await; + assert_eq!(st, 200, "{job1}"); + let (_, job2, _) = call(&fx, "POST", "/enhance", Some(body)).await; + + assert_eq!(job1["status"], "queued"); + assert_eq!(job1["path"], PHOTO); + assert_eq!(job1["template"], "restore"); + assert_eq!(job1["prompt"], "warmer colours"); + assert_eq!(job1["pending"], json!([])); + assert!(job1["created_at"].as_i64().unwrap() > 0); + assert_ne!(job1["id"], job2["id"]); + + let sent = fx.mock.submitted.lock().unwrap(); + assert_eq!(sent.len(), 2); + let p1 = &sent[0]["prompt"]; + let p2 = &sent[1]["prompt"]; + assert_eq!(p1["37"]["inputs"]["path"], "Z:\\2009_200\\019_19.JPG"); + assert_eq!(p1["36"]["inputs"]["require_approval"], true); + assert!( + p1["13"]["inputs"]["value"] + .as_str() + .unwrap() + .ends_with("overlays. warmer colours") + ); + assert_ne!(p1["5"]["inputs"]["seed"], p2["5"]["inputs"]["seed"]); + assert_ne!( + p1["14"]["inputs"]["sampling_mode.seed"], + p2["14"]["inputs"]["sampling_mode.seed"] + ); + assert_eq!(sent[0]["client_id"], "imageapi"); +} + +#[actix_rt::test] +async fn submit_validation() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + // Edit needs an instruction. + let (st, body, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "edit"})), + ) + .await; + assert_eq!(st, 400, "{body}"); + assert!(body["error"].as_str().unwrap().contains("instruction")); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "edit", "prompt": " "})), + ) + .await; + assert_eq!(st, 400); + // Unknown template / photo / traversal / reference on a template without one. + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "nope"})), + ) + .await; + assert_eq!(st, 404); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": "missing.jpg", "template": "restore"})), + ) + .await; + assert_eq!(st, 404); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": "../etc/passwd", "template": "restore"})), + ) + .await; + assert_eq!(st, 400); + fx.put("ref.jpg", b"other"); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore", "reference_path": "ref.jpg"})), + ) + .await; + assert_eq!(st, 400); + assert!( + fx.mock.submitted.lock().unwrap().is_empty(), + "nothing reached ComfyUI" + ); +} + +#[actix_rt::test] +async fn edit_without_reference_drops_the_node_and_with_one_keeps_it() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + fx.put("ref/other.jpg", b"other view"); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "edit", "prompt": "make it b&w"})), + ) + .await; + assert_eq!(st, 200); + let (st, job, _) = call( + &fx, "POST", "/enhance", + Some(json!({"path": PHOTO, "template": "edit", "prompt": "x", "reference_path": "ref/other.jpg"})), + ) + .await; + assert_eq!(st, 200); + assert_eq!(job["reference_path"], "ref/other.jpg"); + let sent = fx.mock.submitted.lock().unwrap(); + let without = &sent[0]["prompt"]; + assert!(without.get("38").is_none()); + assert!(!without.to_string().contains("[\"38\"")); + assert_eq!(without["13"]["inputs"]["value"], "make it b&w"); + let with = &sent[1]["prompt"]; + assert_eq!(with["38"]["inputs"]["path"], "Z:\\ref\\other.jpg"); +} + +#[actix_rt::test] +async fn comfy_rejection_is_reported_and_a_missing_reference_is_named() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + *fx.mock.reject_prompts.lock().unwrap() = Some(json!({ + "error": {"type": "prompt_outputs_failed_validation", "message": "Prompt outputs failed validation", "details": ""}, + "node_errors": {"36": {"class_type": "ImageApiSaveImageVersion", "errors": [{"message": "Required input is missing"}]}}, + })); + let (st, body, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "edit", "prompt": "x"})), + ) + .await; + assert_eq!(st, 422); + assert!( + body["error"] + .as_str() + .unwrap() + .contains("needs a reference photo"), + "{body}" + ); + // Restore never dropped a reference, so the message is ComfyUI's own. + let (st, body, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + assert_eq!(st, 422); + assert!(!body["error"].as_str().unwrap().contains("reference")); + assert!( + body["error"] + .as_str() + .unwrap() + .contains("Required input is missing") + ); +} + +#[actix_rt::test] +async fn job_moves_queued_running_ready_and_the_preview_is_relayed() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let id = job["id"].as_str().unwrap().to_string(); + + // Behind someone else's run: queued with a position. + fx.mock.queue_running.lock().unwrap().push("other".into()); + fx.svc.poll_once().await; + let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; + assert_eq!( + (j["status"].as_str().unwrap(), j["queue_position"].as_i64()), + ("queued", Some(1)) + ); + + // Now it runs. + *fx.mock.queue_running.lock().unwrap() = vec!["prompt-1".into()]; + fx.mock.queue_pending.lock().unwrap().clear(); + fx.svc.poll_once().await; + let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; + assert_eq!(j["status"], "running"); + + // And finishes. + fx.mock.queue_running.lock().unwrap().clear(); + fx.mock + .history + .lock() + .unwrap() + .insert("prompt-1".into(), finished_entry("abc", 10_000)); + fx.svc.poll_once().await; + let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; + assert_eq!(j["status"], "ready"); + assert!(j.get("queue_position").is_none()); + assert_eq!(j["pending"].as_array().unwrap().len(), 1); + assert_eq!(j["pending"][0]["n"], 0); + assert_eq!(j["pending"][0]["approved"], false); + // The staged file name is ComfyUI-internal and not exposed. + assert!(!j.to_string().contains("imageapi_pending")); + + let uri = j["pending"][0]["preview"].as_str().unwrap().to_string(); + let (st, _, bytes) = call(&fx, "GET", &uri, None).await; + assert_eq!(st, 200); + assert_eq!(bytes, b"PNG:imageapi_pending_abc.png"); + let (st, _, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}/preview/5"), None).await; + assert_eq!(st, 404); + + // The photo's job list shows it (and a failed one too). + let (_, list, _) = call(&fx, "GET", &format!("/enhance/jobs?path={PHOTO}"), None).await; + assert_eq!(list["jobs"].as_array().unwrap().len(), 1); + assert_eq!(list["jobs"][0]["id"], id.as_str()); +} + +#[actix_rt::test] +async fn failed_run_reports_comfys_message() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + fx.mock.queue_pending.lock().unwrap().clear(); + fx.mock.history.lock().unwrap().insert( + "prompt-1".into(), + json!({"status": {"status_str": "error", "completed": false, "messages": [ + ["execution_error", {"exception_message": "CUDA out of memory\n", "timestamp": 5}]]}, "outputs": {}}), + ); + fx.svc.poll_once().await; + let (_, j, _) = call( + &fx, + "GET", + &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(j["status"], "failed"); + assert_eq!(j["error"], "CUDA out of memory"); +} + +#[actix_rt::test] +async fn a_job_comfy_forgot_expires() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + // ComfyUI restarted: empty queue, empty history. + fx.mock.queue_pending.lock().unwrap().clear(); + fx.svc.poll_once().await; + let (_, j, _) = call( + &fx, + "GET", + &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(j["status"], "expired"); + assert!(j["error"].as_str().unwrap().contains("restarted")); + assert_eq!(j["pending"], json!([])); +} + +#[actix_rt::test] +async fn ready_job_expires_when_the_preview_is_gone_and_approve_says_so() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let id = job["id"].as_str().unwrap().to_string(); + fx.mock.queue_pending.lock().unwrap().clear(); + fx.mock + .history + .lock() + .unwrap() + .insert("prompt-1".into(), finished_entry("abc", 1)); + fx.svc.poll_once().await; + + *fx.mock.view_404.lock().unwrap() = true; + let (st, body, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}/preview/0"), None).await; + assert_eq!((st, body["code"].as_str()), (410, Some("expired"))); + let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; + assert_eq!(j["status"], "expired"); + + let (st, _, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{id}/approve"), + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 410); + assert!(fx.mock.approve_calls.lock().unwrap().is_empty()); +} + +#[actix_rt::test] +async fn approve_gone_on_comfy_side_expires_the_job() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let id = job["id"].as_str().unwrap().to_string(); + fx.mock.queue_pending.lock().unwrap().clear(); + fx.mock + .history + .lock() + .unwrap() + .insert("prompt-1".into(), finished_entry("abc", 1)); + fx.svc.poll_once().await; + *fx.mock.approve_gone.lock().unwrap() = true; + let (st, _, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{id}/approve"), + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 410); + let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; + assert_eq!(j["status"], "expired"); +} + +#[actix_rt::test] +async fn approve_ingests_links_and_is_idempotent() { + let fx = fx().await; + let original = fx.put(PHOTO, b"the original photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let id = job["id"].as_str().unwrap().to_string(); + fx.mock.queue_pending.lock().unwrap().clear(); + fx.mock + .history + .lock() + .unwrap() + .insert("prompt-1".into(), finished_entry("abc", 1)); + fx.svc.poll_once().await; + + // What ComfyUI's approve would have written: a PNG with the XMP link. + let saved_rel = "2009_200/aiedits/019_19.JPG.restored.png"; + let xmp = producer_xmp(&original, "restored", &[]); + write_png(&fx.root.join(saved_rel), Some(&xmp), None, 7); + *fx.mock.approve_path.lock().unwrap() = "Z:\\2009_200\\aiedits\\019_19.JPG.restored.png".into(); + + let (st, v, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{id}/approve"), + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 200, "{v}"); + assert_eq!(v["path"], saved_rel); + assert_eq!(v["library_id"], 1); + assert_eq!(v["kind"], "restored"); + assert_eq!(v["source"], "xmp"); + for k in ["hash", "created_at", "references", "width", "height"] { + assert!(v.get(k).is_some(), "missing {k}: {v}"); + } + assert_eq!(fx.mock.approve_calls.lock().unwrap().as_slice(), ["abc"]); + + // Visible in /image/versions right away. + let (_, vs, _) = call(&fx, "GET", &format!("/image/versions?path={PHOTO}"), None).await; + assert_eq!(vs["versions"].as_array().unwrap().len(), 1); + assert_eq!(vs["versions"][0]["path"], saved_rel); + + // Approving again: same version, ComfyUI not asked, no duplicate row. + let (st, v2, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{id}/approve"), + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 200); + assert_eq!(v2["hash"], v["hash"]); + assert_eq!(fx.mock.approve_calls.lock().unwrap().len(), 1); + let (_, vs, _) = call(&fx, "GET", &format!("/image/versions?path={PHOTO}"), None).await; + assert_eq!(vs["versions"].as_array().unwrap().len(), 1); + let rows: i64 = { + #[derive(QueryableByName)] + struct N { + #[diesel(sql_type = diesel::sql_types::BigInt)] + n: i64, + } + sql_query("SELECT COUNT(*) AS n FROM image_exif WHERE rel_path = ?") + .bind::(saved_rel) + .get_result::(&mut *fx.conn.lock().unwrap()) + .unwrap() + .n + }; + assert_eq!(rows, 1); + + let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; + assert_eq!(j["status"], "approved"); + assert_eq!(j["pending"][0]["approved"], true); + assert_eq!(j["saved"][0]["path"], saved_rel); +} + +#[actix_rt::test] +async fn approve_of_a_path_outside_every_library_is_reported_not_lost() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let id = job["id"].as_str().unwrap().to_string(); + fx.mock.queue_pending.lock().unwrap().clear(); + fx.mock + .history + .lock() + .unwrap() + .insert("prompt-1".into(), finished_entry("abc", 1)); + fx.svc.poll_once().await; + *fx.mock.approve_path.lock().unwrap() = "C:\\somewhere\\else.png".into(); + let (st, body, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{id}/approve"), + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 502); + assert!( + body["error"] + .as_str() + .unwrap() + .contains("C:\\somewhere\\else.png") + ); +} + +#[actix_rt::test] +async fn approve_before_ready_is_a_400() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let (st, _, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{}/approve", job["id"].as_str().unwrap()), + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 400); + let (st, _, _) = call( + &fx, + "POST", + "/enhance/jobs/nope/approve", + Some(json!({"n": 0})), + ) + .await; + assert_eq!(st, 404); +} + +#[actix_rt::test] +async fn discard_hides_the_job_from_the_photo_list_only() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + let (_, a, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let (_, b, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + let (st, _, _) = call( + &fx, + "POST", + &format!("/enhance/jobs/{}/discard", a["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(st, 200); + let (_, list, _) = call(&fx, "GET", &format!("/enhance/jobs?path={PHOTO}"), None).await; + let ids: Vec<&str> = list["jobs"] + .as_array() + .unwrap() + .iter() + .map(|j| j["id"].as_str().unwrap()) + .collect(); + assert_eq!(ids, [b["id"].as_str().unwrap()]); + // Still fetchable by id. + let (st, _, _) = call( + &fx, + "GET", + &format!("/enhance/jobs/{}", a["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(st, 200); + let (st, _, _) = call(&fx, "POST", "/enhance/jobs/nope/discard", None).await; + assert_eq!(st, 404); +} + +#[actix_rt::test] +async fn unreachable_active_jobs_become_unavailable_after_the_grace_period() { + // Submit against the mock, then point nothing at it: use a service whose + // client targets a dead port but whose DB holds an active job. + let fx = fx_with(Some("http://127.0.0.1:9".into()), |c| { + c.unreachable_grace_secs = 1 + }) + .await; + let now = chrono::Utc::now().timestamp(); + let dao_job = crate::database::enhance_dao::EnhanceJobRow { + id: "j1".into(), + original_hash: "h".into(), + library_id: 1, + rel_path: "a.jpg".into(), + template: "restore".into(), + prompt: None, + reference_hash: None, + reference_path: None, + comfy_prompt_id: "p".into(), + status: "running".into(), + error: None, + queue_position: None, + pending: "[]".into(), + saved: "[]".into(), + discarded: 0, + created_at: now, + updated_at: now, + }; + SqliteEnhanceJobDao::from_shared(fx.conn.clone()) + .insert(&dao_job) + .unwrap(); + fx.svc.poll_once().await; + assert_eq!( + fx.svc.get_job("j1").unwrap().status, + "running", + "within grace" + ); + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + fx.svc.poll_once().await; + let j = fx.svc.get_job("j1").unwrap(); + assert_eq!(j.status, "unavailable"); + assert_eq!(j.error.as_deref(), Some("ComfyUI isn't running")); +} + +#[actix_rt::test] +async fn idle_comfy_is_freed_once_and_activity_rearms_it() { + let fx = fx_with(None, |c| c.idle_free_secs = 1).await; + *fx.mock.newest.lock().unwrap() = Some(finished_entry("x", 100)); + fx.svc.idle_check().await; // baseline + assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 0); + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + fx.svc.idle_check().await; + fx.svc.idle_check().await; + assert_eq!( + fx.mock.free_calls.load(Ordering::SeqCst), + 1, + "freed exactly once" + ); + + // Someone runs a workflow in ComfyUI's own page: newest history changes. + *fx.mock.newest.lock().unwrap() = Some(finished_entry("y", 200)); + fx.svc.idle_check().await; + assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 1); + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + fx.svc.idle_check().await; + assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 2); +} + +#[actix_rt::test] +async fn busy_queue_blocks_the_free() { + let fx = fx_with(None, |c| c.idle_free_secs = 1).await; + fx.mock.queue_running.lock().unwrap().push("someone".into()); + fx.svc.idle_check().await; + tokio::time::sleep(std::time::Duration::from_millis(1100)).await; + fx.svc.idle_check().await; + assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 0); +} diff --git a/src/enhance/path_map.rs b/src/enhance/path_map.rs index 1071959..7ceb098 100644 --- a/src/enhance/path_map.rs +++ b/src/enhance/path_map.rs @@ -82,7 +82,8 @@ impl PathMap { .filter_map(|e| { let prefix = format!("{}/", e.to.replace('\\', "/")); let head = text.get(..prefix.len())?; - head.eq_ignore_ascii_case(&prefix).then_some((e, prefix.len())) + head.eq_ignore_ascii_case(&prefix) + .then_some((e, prefix.len())) }) .max_by_key(|(_, len)| *len)?; let rest = &text[best.1..]; diff --git a/src/enhance/service.rs b/src/enhance/service.rs index 5b4bd5f..8a3ba53 100644 --- a/src/enhance/service.rs +++ b/src/enhance/service.rs @@ -13,9 +13,7 @@ use std::time::{Duration, Instant}; use log::{info, warn}; use serde::{Deserialize, Serialize}; -use super::comfy::{ - ApproveError, ComfyClient, ComfyError, HistoryEntry, Staged, SubmitError, -}; +use super::comfy::{ApproveError, ComfyClient, ComfyError, HistoryEntry, Staged, SubmitError}; use super::config::{EnhanceConfig, Loaded, PromptMode}; use super::path_map::PathMap; use super::template::Template; @@ -35,8 +33,7 @@ const JOB_RETENTION_SECS: i64 = 14 * 24 * 3600; pub const JOBS_PER_PHOTO: i64 = 20; pub const MSG_UNAVAILABLE: &str = "ComfyUI isn't running"; -pub const MSG_EXPIRED: &str = - "This result expired (ComfyUI was restarted). Run it again."; +pub const MSG_EXPIRED: &str = "This result expired (ComfyUI was restarted). Run it again."; pub const MSG_DISABLED: &str = "Enhance isn't configured"; /// A library file a request names. @@ -247,10 +244,12 @@ impl EnhanceService { } pub fn default_template(&self) -> Option<&str> { - let wanted = self.config.default_name()?; - self.template(wanted) + self.config + .default_template + .as_deref() + .and_then(|w| self.template(w)) + .or_else(|| self.templates.first()) .map(|t| t.cfg.name.as_str()) - .or_else(|| self.templates.first().map(|t| t.cfg.name.as_str())) } /// Cached `GET /system_stats` probe. @@ -315,12 +314,9 @@ impl EnhanceService { None => None, }; let filled = template - .fill( - &source, - prompt, - reference_path.as_deref(), - &mut || rand::random::() >> 12, - ) + .fill(&source, prompt, reference_path.as_deref(), &mut || { + rand::random::() >> 12 + }) .map_err(|e| EnhanceError::BadRequest(e.0))?; if !self.probe_now().await { @@ -329,7 +325,10 @@ impl EnhanceService { let prompt_id = match client.submit(&filled.prompt).await { Ok(id) => id, Err(SubmitError::Comfy(e)) => return Err(comfy_err(e)), - Err(SubmitError::Rejected { message, node_errors }) => { + Err(SubmitError::Rejected { + message, + node_errors, + }) => { // A reference-less run of a template that wants one fails // validation on the dropped node's consumers. return Err(if filled.dropped_reference && node_errors { @@ -481,7 +480,11 @@ impl EnhanceService { for job in active { self.unreachable_since.lock().unwrap().remove(&job.id); if let Some((running, ahead)) = queue.position_of(&job.comfy_prompt_id) { - let status = if running { STATUS_RUNNING } else { STATUS_QUEUED }; + let status = if running { + STATUS_RUNNING + } else { + STATUS_QUEUED + }; let mut u = JobUpdate::status(status); u.queue_position = Some(ahead as i32); if job.status != status || job.queue_position != u.queue_position { @@ -501,7 +504,9 @@ impl EnhanceService { u.error = Some(MSG_EXPIRED.to_string()); self.update_job(&job.id, &u); } - Err(ComfyError::Unreachable(_)) => self.note_unreachable(std::slice::from_ref(&job)), + Err(ComfyError::Unreachable(_)) => { + self.note_unreachable(std::slice::from_ref(&job)) + } Err(e) => warn!("enhance: /history for job {} failed: {e}", job.id), } } @@ -527,7 +532,11 @@ impl EnhanceService { let mut u = JobUpdate::status(STATUS_READY); u.pending = serde_json::to_string(&entry.staged).ok(); self.update_job(&job.id, &u); - info!("enhance: job {} ready ({} result(s))", job.id, entry.staged.len()); + info!( + "enhance: job {} ready ({} result(s))", + job.id, + entry.staged.len() + ); } fn note_unreachable(&self, jobs: &[EnhanceJobRow]) { diff --git a/src/enhance/template.rs b/src/enhance/template.rs index cd8a381..55c0e0d 100644 --- a/src/enhance/template.rs +++ b/src/enhance/template.rs @@ -184,9 +184,7 @@ impl Template { let user_text = prompt.map(str::trim).filter(|t| !t.is_empty()); match (&self.prompt, user_text) { (None, Some(_)) => { - return Err(FillError( - "this template doesn't take a prompt".to_string(), - )); + return Err(FillError("this template doesn't take a prompt".to_string())); } (Some(_), None) if self.prompt_required() => { return Err(FillError( @@ -344,7 +342,9 @@ mod tests { #[test] fn missing_required_titles_are_rejected_with_a_clear_message() { let no_source = RESTORE.replace("imageapi:source", "something else"); - let err = load(&no_source, PromptMode::Append).unwrap_err().to_string(); + let err = load(&no_source, PromptMode::Append) + .unwrap_err() + .to_string(); assert!(err.contains("imageapi:source"), "{err}"); let no_save = RESTORE.replace("imageapi:save", "something else"); let err = load(&no_save, PromptMode::Append).unwrap_err().to_string(); @@ -372,7 +372,9 @@ mod tests { let mut v: Value = serde_json::from_str(RESTORE).unwrap(); v["13"]["inputs"] = json!({"prompt": "base"}); let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap(); - let f = t.fill("Z:\\a.jpg", Some("more"), None, &mut seeds()).unwrap(); + let f = t + .fill("Z:\\a.jpg", Some("more"), None, &mut seeds()) + .unwrap(); assert_eq!(input(&f.prompt, "13", "prompt"), "base more"); } @@ -393,7 +395,9 @@ mod tests { assert_eq!(input(&f.prompt, "36", "require_approval"), true); // KSampler.seed and TextGenerate's sampling_mode.seed, both fresh. let a = input(&f.prompt, "5", "seed").as_u64().unwrap(); - let b = input(&f.prompt, "14", "sampling_mode.seed").as_u64().unwrap(); + let b = input(&f.prompt, "14", "sampling_mode.seed") + .as_u64() + .unwrap(); assert!(a > 1000 && b > 1000 && a != b); assert_ne!(a, 774744764437710); // Two fills with a real RNG differ. @@ -433,7 +437,9 @@ mod tests { #[test] fn default_join_is_a_single_space_sentence() { let t = load(RESTORE, PromptMode::Append).unwrap(); - let f = t.fill("p", Some("Make it warmer."), None, &mut seeds()).unwrap(); + let f = t + .fill("p", Some("Make it warmer."), None, &mut seeds()) + .unwrap(); let text = input(&f.prompt, "13", "value").as_str().unwrap(); assert!(text.ends_with("Remove any phone interface overlays. Make it warmer.")); } @@ -444,7 +450,9 @@ mod tests { let err = t.fill("p", None, None, &mut seeds()).unwrap_err(); assert!(err.0.contains("instruction")); assert!(t.fill("p", Some(" "), None, &mut seeds()).is_err()); - let f = t.fill("p", Some("make it black and white"), None, &mut seeds()).unwrap(); + let f = t + .fill("p", Some("make it black and white"), None, &mut seeds()) + .unwrap(); assert_eq!(input(&f.prompt, "13", "value"), "make it black and white"); } @@ -482,7 +490,11 @@ mod tests { assert!(!wf.contains_key("38")); // No input anywhere still links to node 38, and the leftover export // path can't leak. - assert!(!serde_json::to_string(&f.prompt).unwrap().contains("[\"38\"")); + assert!( + !serde_json::to_string(&f.prompt) + .unwrap() + .contains("[\"38\"") + ); assert!(wf["36"]["inputs"].get("reference_1").is_none()); // Links to other nodes survive. assert_eq!(input(&f.prompt, "36", "source"), &json!(["37", 1]));