feat(enhance): discard, tests against a mock ComfyUI, docs (#192)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Cameron CordesandClaude Opus 5.5 authored and cameron committed 2026-10-30 19:41:05 +00:00
1 parent 3a21761cfe
commit 97d9301a71
10 files changed
+1165 -48

No files matched your search

+10 -3
View File
@@ -93,8 +93,11 @@ pub trait EnhanceJobDao: Sync + Send {
fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError>;
fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError>;
/// Newest first.
fn list_for_hash(&mut self, original_hash: &str, limit: i64)
-> Result<Vec<EnhanceJobRow>, DbError>;
fn list_for_hash(
&mut self,
original_hash: &str,
limit: i64,
) -> Result<Vec<EnhanceJobRow>, DbError>;
/// Jobs the poller still has to advance (`queued` / `running`).
fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, 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).
+7 -2
View File
@@ -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| {
+2 -2
View File
@@ -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);
}
}
+21 -13
View File
@@ -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<Mutex<Box<dyn ExifDao>>> = 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(),
+985
View File
@@ -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<Vec<Value>>,
queue_running: Mutex<Vec<String>>,
queue_pending: Mutex<Vec<String>>,
/// prompt_id -> history entry
history: Mutex<HashMap<String, Value>>,
newest: Mutex<Option<Value>>,
approve_calls: Mutex<Vec<String>>,
approve_path: Mutex<String>,
approve_gone: Mutex<bool>,
free_calls: AtomicUsize,
reject_prompts: Mutex<Option<Value>>,
view_404: Mutex<bool>,
}
async fn start_mock() -> (Arc<Mock>, 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<Mock>, body: web::Json<Value>| 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<Mock>| async move {
let item = |id: &String| json!([1, id, {}, {}, []]);
HttpResponse::Ok().json(json!({
"queue_running": m.queue_running.lock().unwrap().iter().map(item).collect::<Vec<_>>(),
"queue_pending": m.queue_pending.lock().unwrap().iter().map(item).collect::<Vec<_>>(),
}))
}),
)
.route(
"/history",
web::get().to(|m: Data<Mock>| 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<Mock>, id: web::Path<String>| 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<Mock>, q: web::Query<HashMap<String, String>>| 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<Mock>, body: web::Json<Value>| 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<Mock>, body: web::Json<Value>| 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<Mock>,
svc: Data<EnhanceService>,
state: Data<AppState>,
exif: Data<Mutex<Box<dyn ExifDao>>>,
conn: Arc<Mutex<SqliteConnection>>,
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<String>, 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::<Text, _>(&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<Mutex<Box<dyn EnhanceJobDao>>> = 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<dyn ExifDao>,
)),
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::<Text, _>(rel)
.bind::<Text, _>(&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<Value>) -> (u16, Value, Vec<u8>) {
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<Value>) -> (u16, Value, Vec<u8>) {
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::<Text, _>(saved_rel)
.get_result::<N>(&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);
}
+2 -1
View File
@@ -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..];
+27 -18
View File
@@ -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::<u64>() >> 12,
)
.fill(&source, prompt, reference_path.as_deref(), &mut || {
rand::random::<u64>() >> 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]) {
+21 -9
View File
@@ -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]));