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/enhance.example.toml b/enhance.example.toml new file mode 100644 index 0000000..9cc80a9 --- /dev/null +++ b/enhance.example.toml @@ -0,0 +1,65 @@ +# ImageApi Enhance configuration (#192). Copy to `enhance.toml` next to +# `.env` (or point ENHANCE_CONFIG_PATH at it). Every key is optional: delete +# a line to get its default (trailing `default:` comment). Without `url`, +# Enhance is off and GET /enhance/templates reports `available: false`. +# +# ComfyUI has no auth and stays LAN-only; the apps only ever talk to +# ImageApi, which relays previews. Nothing in this file is a secret. +# +# Templates are API-format workflows (ComfyUI -> Workflow -> Export (API)). +# ImageApi finds the nodes it fills in by `_meta.title`, so re-exporting +# after editing the graph keeps working: +# imageapi:source required Load Image From Path; gets the photo's path +# imageapi:save required Save Image Version; forced to require_approval +# imageapi:prompt optional text node; its `text`, `prompt` or `value` +# input (first one present) gets the user's text +# imageapi:reference optional second loader for a context photo. With no +# reference in the request the node and every +# link to it are removed from the prompt. +# A template missing a required title (or with a prompt node that has none of +# the three inputs) is rejected at startup with a log line; the rest still load. +# Never edit the template files from here: ImageApi only reads them. + +url = "http://192.168.10.218:8188" # default: unset (Enhance disabled) +timeout_secs = 15 # default: 15; per ComfyUI request +probe_ttl_secs = 10 # default: 10; how long `available` is cached +poll_secs = 2 # default: 2; job poll interval +idle_free_secs = 300 # default: 300; ComfyUI idle (empty queue, newest + # history entry this old) -> POST /free once. 0 = off +unreachable_grace_secs = 120 # default: 120; active job -> `unavailable` after this +default_template = "restore" # default: the first template +warn_above_megapixels = 1.5 # default: 1.5; clients warn above this original size + +# Library path prefix <-> ComfyUI path prefix; longest prefix wins and +# separators are converted, in both directions. With no entries, paths pass +# through unchanged. A mapped drive (Z:\) only exists in the logged-in +# session; use a UNC prefix (\\nas\) if ComfyUI ever runs as a service. +[[path_map]] +from = "/media/nas/" +to = 'Z:\' + +[[template]] +name = "restore" # stable id clients send +label = "Restore" +description = "Sharpen and clean up the photo (about 3 s)." +file = "/media/nas/image_fibo2_restore_turbo_imageapi_API.json" # relative = beside this file +# append: the template's prompt is a full base instruction; the user's text +# is appended after prompt_join. Empty text leaves the prompt as is. +prompt_mode = "append" # default: "append" | "replace" +prompt_join = " " # default: " " (single space); e.g. "\n\nAdditional instructions: " +output_megapixels = 1.0 # default: 1.0; ~1 MP output, for the client's size warning + +[[template]] +name = "edit" +label = "Edit" +description = "Change the photo as you describe." +file = "/media/nas/image_fibo2_edit_turbo_imageapi_API.json" +# replace: the user's text IS the prompt, so it is required (empty = 400). +prompt_mode = "replace" + +# The Dev variant (50 steps, ~50 s) is no better on the archive test set, so +# it is not the default; configure it if you want it: +# [[template]] +# name = "restore-dev" +# label = "Restore (Dev, slow)" +# file = "/media/nas/image_fibo2_restore_imageapi_API.json" diff --git a/migrations/2026-10-30-100000_add_enhance_jobs/down.sql b/migrations/2026-10-30-100000_add_enhance_jobs/down.sql new file mode 100644 index 0000000..d432e23 --- /dev/null +++ b/migrations/2026-10-30-100000_add_enhance_jobs/down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS enhance_jobs; diff --git a/migrations/2026-10-30-100000_add_enhance_jobs/up.sql b/migrations/2026-10-30-100000_add_enhance_jobs/up.sql new file mode 100644 index 0000000..5485e2d --- /dev/null +++ b/migrations/2026-10-30-100000_add_enhance_jobs/up.sql @@ -0,0 +1,31 @@ +-- Enhance on demand (#192): one row per ComfyUI run submitted for a photo. +-- +-- Keyed to the photo by `original_hash` (the blake3 content_hash of the file +-- that was submitted), like every other derived table, so a job survives a +-- rename or move; `library_id` + `rel_path` say which file it was at submit +-- time (and are what the API echoes back as `path`). +CREATE TABLE enhance_jobs ( + id TEXT PRIMARY KEY NOT NULL, -- uuid, ImageApi's own + original_hash TEXT NOT NULL, + library_id INTEGER NOT NULL, + rel_path TEXT NOT NULL, + template TEXT NOT NULL, + prompt TEXT, + reference_hash TEXT, + reference_path TEXT, -- library-relative, same library + comfy_prompt_id TEXT NOT NULL, + -- queued | running | ready | approved | failed | expired | unavailable + status TEXT NOT NULL, + error TEXT, + queue_position INTEGER, + -- json: [{id, view: {filename, subfolder, type}}], ComfyUI's staged results + pending TEXT NOT NULL DEFAULT '[]', + -- json: [{n, comfy_path, library_id, rel_path}], approved results + saved TEXT NOT NULL DEFAULT '[]', + -- 1 = the user threw the result away; hidden from the per-photo list + discarded INTEGER NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + updated_at BIGINT NOT NULL +); +CREATE INDEX idx_enhance_jobs_hash ON enhance_jobs(original_hash, created_at); +CREATE INDEX idx_enhance_jobs_status ON enhance_jobs(status); diff --git a/src/database/enhance_dao.rs b/src/database/enhance_dao.rs new file mode 100644 index 0000000..67633c3 --- /dev/null +++ b/src/database/enhance_dao.rs @@ -0,0 +1,343 @@ +//! Enhance jobs (#192): one row per ComfyUI run submitted for a photo. +//! +//! The row is the only state: the poller advances it from ComfyUI's +//! `/queue` and `/history`, the endpoints read it, approve appends to +//! `saved`. `pending` / `saved` hold small JSON documents owned by +//! `crate::enhance`; this layer treats them as opaque text. + +use std::sync::{Arc, Mutex}; + +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{BigInt, Integer, Nullable, Text}; +use diesel::sqlite::SqliteConnection; + +use crate::database::{DbError, DbErrorKind, connect}; +use crate::otel::trace_db_call; + +pub const STATUS_QUEUED: &str = "queued"; +pub const STATUS_RUNNING: &str = "running"; +pub const STATUS_READY: &str = "ready"; +pub const STATUS_APPROVED: &str = "approved"; +pub const STATUS_FAILED: &str = "failed"; +pub const STATUS_EXPIRED: &str = "expired"; +pub const STATUS_UNAVAILABLE: &str = "unavailable"; + +#[derive(Debug, Clone, PartialEq, Eq, QueryableByName)] +pub struct EnhanceJobRow { + #[diesel(sql_type = Text)] + pub id: String, + #[diesel(sql_type = Text)] + pub original_hash: String, + #[diesel(sql_type = Integer)] + pub library_id: i32, + #[diesel(sql_type = Text)] + pub rel_path: String, + #[diesel(sql_type = Text)] + pub template: String, + #[diesel(sql_type = Nullable)] + pub prompt: Option, + #[diesel(sql_type = Nullable)] + pub reference_hash: Option, + #[diesel(sql_type = Nullable)] + pub reference_path: Option, + #[diesel(sql_type = Text)] + pub comfy_prompt_id: String, + #[diesel(sql_type = Text)] + pub status: String, + #[diesel(sql_type = Nullable)] + pub error: Option, + #[diesel(sql_type = Nullable)] + pub queue_position: Option, + #[diesel(sql_type = Text)] + pub pending: String, + #[diesel(sql_type = Text)] + pub saved: String, + #[diesel(sql_type = Integer)] + pub discarded: i32, + #[diesel(sql_type = BigInt)] + pub created_at: i64, + #[diesel(sql_type = BigInt)] + pub updated_at: i64, +} + +const COLUMNS: &str = "id, original_hash, library_id, rel_path, template, prompt, reference_hash, \ + reference_path, comfy_prompt_id, status, error, queue_position, pending, saved, \ + discarded, created_at, updated_at"; + +/// What the poller / approve changes on an existing job. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct JobUpdate { + pub status: String, + pub error: Option, + pub queue_position: Option, + /// `Some` replaces the `pending` document. + pub pending: Option, + /// `Some` replaces the `saved` document. + pub saved: Option, +} + +impl JobUpdate { + pub fn status(status: &str) -> Self { + JobUpdate { + status: status.to_string(), + error: None, + queue_position: None, + pending: None, + saved: None, + } + } +} + +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>; + /// 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>; + /// Hide a job from `list_for_hash`. + fn discard(&mut self, id: &str) -> Result<(), DbError>; + /// Delete jobs last touched before `cutoff` (unix seconds) that are not + /// active. Returns the number removed. + fn prune(&mut self, cutoff: i64) -> Result; +} + +pub struct SqliteEnhanceJobDao { + connection: Arc>, +} + +impl Default for SqliteEnhanceJobDao { + fn default() -> Self { + Self::new() + } +} + +impl SqliteEnhanceJobDao { + pub fn new() -> Self { + SqliteEnhanceJobDao { + connection: Arc::new(Mutex::new(connect())), + } + } + + /// Test-only: share a connection with other DAOs. + #[cfg(test)] + pub fn from_shared(connection: Arc>) -> Self { + SqliteEnhanceJobDao { connection } + } +} + +impl EnhanceJobDao for SqliteEnhanceJobDao { + fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError> { + trace_db_call("insert", "enhance_job_insert", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query(format!( + "INSERT INTO enhance_jobs ({COLUMNS}) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)" + )) + .bind::(&job.id) + .bind::(&job.original_hash) + .bind::(job.library_id) + .bind::(&job.rel_path) + .bind::(&job.template) + .bind::, _>(&job.prompt) + .bind::, _>(&job.reference_hash) + .bind::, _>(&job.reference_path) + .bind::(&job.comfy_prompt_id) + .bind::(&job.status) + .bind::, _>(&job.error) + .bind::, _>(job.queue_position) + .bind::(&job.pending) + .bind::(&job.saved) + .bind::(job.discarded) + .bind::(job.created_at) + .bind::(job.updated_at) + .execute(&mut *conn) + .map(|_| ()) + .map_err(|e| anyhow::anyhow!("Insert error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::InsertError, e)) + } + + fn get(&mut self, id: &str) -> Result, DbError> { + trace_db_call("query", "enhance_job_get", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query(format!("SELECT {COLUMNS} FROM enhance_jobs WHERE id = ?")) + .bind::(id) + .load::(&mut *conn) + .map(|rows| rows.into_iter().next()) + .map_err(|e| anyhow::anyhow!("Query error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::QueryError, e)) + } + + fn list_for_hash( + &mut self, + original_hash: &str, + limit: i64, + ) -> Result, DbError> { + trace_db_call("query", "enhance_job_list_for_hash", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query(format!( + "SELECT {COLUMNS} FROM enhance_jobs WHERE original_hash = ? AND discarded = 0 \ + ORDER BY created_at DESC, rowid DESC LIMIT ?" + )) + .bind::(original_hash) + .bind::(limit) + .load::(&mut *conn) + .map_err(|e| anyhow::anyhow!("Query error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::QueryError, e)) + } + + fn list_active(&mut self) -> Result, DbError> { + trace_db_call("query", "enhance_job_list_active", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query(format!( + "SELECT {COLUMNS} FROM enhance_jobs WHERE status IN ('queued', 'running') \ + ORDER BY created_at, rowid" + )) + .load::(&mut *conn) + .map_err(|e| anyhow::anyhow!("Query error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::QueryError, e)) + } + + fn update(&mut self, id: &str, u: &JobUpdate) -> Result<(), DbError> { + trace_db_call("update", "enhance_job_update", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query( + "UPDATE enhance_jobs SET status = ?, error = ?, queue_position = ?, \ + pending = COALESCE(?, pending), saved = COALESCE(?, saved), updated_at = ? \ + WHERE id = ?", + ) + .bind::(&u.status) + .bind::, _>(&u.error) + .bind::, _>(u.queue_position) + .bind::, _>(&u.pending) + .bind::, _>(&u.saved) + .bind::(chrono::Utc::now().timestamp()) + .bind::(id) + .execute(&mut *conn) + .map(|_| ()) + .map_err(|e| anyhow::anyhow!("Update error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::UpdateError, e)) + } + + fn discard(&mut self, id: &str) -> Result<(), DbError> { + trace_db_call("update", "enhance_job_discard", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query("UPDATE enhance_jobs SET discarded = 1, updated_at = ? WHERE id = ?") + .bind::(chrono::Utc::now().timestamp()) + .bind::(id) + .execute(&mut *conn) + .map(|_| ()) + .map_err(|e| anyhow::anyhow!("Update error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::UpdateError, e)) + } + + fn prune(&mut self, cutoff: i64) -> Result { + trace_db_call("delete", "enhance_job_prune", |_span| { + let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); + sql_query( + "DELETE FROM enhance_jobs WHERE updated_at < ? \ + AND status NOT IN ('queued', 'running')", + ) + .bind::(cutoff) + .execute(&mut *conn) + .map_err(|e| anyhow::anyhow!("Delete error: {}", e)) + }) + .map_err(|e| DbError::log(DbErrorKind::UpdateError, e)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::database::test::in_memory_db_connection; + + fn dao() -> SqliteEnhanceJobDao { + SqliteEnhanceJobDao::from_shared(Arc::new(Mutex::new(in_memory_db_connection()))) + } + + fn job(id: &str, hash: &str, status: &str, created: i64) -> EnhanceJobRow { + EnhanceJobRow { + id: id.into(), + original_hash: hash.into(), + library_id: 1, + rel_path: "a.jpg".into(), + template: "restore".into(), + prompt: Some("warmer".into()), + reference_hash: None, + reference_path: None, + comfy_prompt_id: format!("p-{id}"), + status: status.into(), + error: None, + queue_position: None, + pending: "[]".into(), + saved: "[]".into(), + discarded: 0, + created_at: created, + updated_at: created, + } + } + + #[test] + fn insert_get_list_update_prune() { + let mut d = dao(); + d.insert(&job("a", "h1", "queued", 100)).unwrap(); + d.insert(&job("b", "h1", "ready", 200)).unwrap(); + d.insert(&job("c", "h2", "running", 300)).unwrap(); + + assert_eq!(d.get("a").unwrap().unwrap().comfy_prompt_id, "p-a"); + assert!(d.get("zzz").unwrap().is_none()); + + let ids: Vec = d + .list_for_hash("h1", 10) + .unwrap() + .into_iter() + .map(|j| j.id) + .collect(); + assert_eq!(ids, ["b", "a"], "newest first"); + assert_eq!(d.list_for_hash("h1", 1).unwrap().len(), 1); + + let active: Vec = d.list_active().unwrap().into_iter().map(|j| j.id).collect(); + assert_eq!(active, ["a", "c"]); + + d.update( + "a", + &JobUpdate { + status: "ready".into(), + error: None, + queue_position: None, + pending: Some("[{\"id\":\"x\"}]".into()), + saved: None, + }, + ) + .unwrap(); + let a = d.get("a").unwrap().unwrap(); + assert_eq!(a.status, "ready"); + assert_eq!(a.pending, "[{\"id\":\"x\"}]"); + assert_eq!(a.saved, "[]", "unset columns are left alone"); + assert!(a.updated_at >= 100); + + d.discard("b").unwrap(); + 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). + assert_eq!(d.prune(i64::MAX).unwrap(), 2); + assert!(d.get("c").unwrap().is_some()); + assert!(d.get("b").unwrap().is_none()); + } +} diff --git a/src/database/mod.rs b/src/database/mod.rs index bdbfc3f..a2d1b7a 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -43,6 +43,7 @@ pub struct UtcBackfillCandidate { pub mod calendar_dao; pub mod daily_summary_dao; +pub mod enhance_dao; pub mod insight_generation_job_dao; pub mod insights_dao; pub mod knowledge_dao; @@ -59,6 +60,7 @@ pub mod version_dao; pub use calendar_dao::{CalendarEventDao, SqliteCalendarEventDao}; pub use daily_summary_dao::{DailySummaryDao, InsertDailySummary, SqliteDailySummaryDao}; +pub use enhance_dao::{EnhanceJobDao, SqliteEnhanceJobDao}; pub use insight_generation_job_dao::{InsightGenerationJobDao, SqliteInsightGenerationJobDao}; pub use insights_dao::{InsightDao, SqliteInsightDao}; pub use knowledge_dao::{ diff --git a/src/database/schema.rs b/src/database/schema.rs index ed2cb02..e59accf 100644 --- a/src/database/schema.rs +++ b/src/database/schema.rs @@ -41,6 +41,28 @@ diesel::table! { } } +diesel::table! { + enhance_jobs (id) { + id -> Text, + original_hash -> Text, + library_id -> Integer, + rel_path -> Text, + template -> Text, + prompt -> Nullable, + reference_hash -> Nullable, + reference_path -> Nullable, + comfy_prompt_id -> Text, + status -> Text, + error -> Nullable, + queue_position -> Nullable, + pending -> Text, + saved -> Text, + discarded -> Integer, + created_at -> BigInt, + updated_at -> BigInt, + } +} + diesel::table! { entities (id) { id -> Integer, @@ -406,6 +428,7 @@ diesel::allow_tables_to_appear_in_same_query!( calendar_events, clip_failures, daily_conversation_summaries, + enhance_jobs, entities, entity_facts, entity_photo_links, diff --git a/src/enhance/comfy.rs b/src/enhance/comfy.rs new file mode 100644 index 0000000..3189f65 --- /dev/null +++ b/src/enhance/comfy.rs @@ -0,0 +1,428 @@ +//! Thin client for the parts of ComfyUI's HTTP API Enhance uses (#192). +//! +//! | Step | Call | +//! |-----------|----------------------------------------| +//! | reachable | `GET /system_stats` | +//! | submit | `POST /prompt {prompt, client_id}` | +//! | status | `GET /queue`, `GET /history/{id}` | +//! | preview | `GET /view?filename=&subfolder=&type=` | +//! | approve | `POST /imageapi/approve {id}` | +//! | free VRAM | `POST /free` | + +use std::time::Duration; + +use bytes::Bytes; +use serde::Deserialize; +use serde_json::{Value, json}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ComfyError { + /// Connect/timeout failure: ComfyUI isn't running (a normal state). + Unreachable(String), + /// ComfyUI answered with a status we don't handle. + Http { status: u16, body: String }, + /// ComfyUI answered with something we can't read. + BadResponse(String), +} + +impl std::fmt::Display for ComfyError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ComfyError::Unreachable(e) => write!(f, "ComfyUI is unreachable: {e}"), + ComfyError::Http { status, body } => write!(f, "ComfyUI returned {status}: {body}"), + ComfyError::BadResponse(e) => write!(f, "unexpected ComfyUI response: {e}"), + } + } +} + +impl std::error::Error for ComfyError {} + +impl From for ComfyError { + fn from(e: reqwest::Error) -> Self { + if e.is_connect() || e.is_timeout() || e.is_request() { + ComfyError::Unreachable(e.without_url().to_string()) + } else { + ComfyError::BadResponse(e.without_url().to_string()) + } + } +} + +/// `POST /prompt` outcomes other than success. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SubmitError { + Comfy(ComfyError), + /// ComfyUI refused the prompt (HTTP 400): validation failed. + Rejected { + message: String, + /// `node_errors` was non-empty (a node is missing an input, a model + /// or a link target). + node_errors: bool, + }, +} + +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}") + } + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ApproveError { + Comfy(ComfyError), + /// The staged result is gone (ComfyUI restarted: its temp folder is cleared). + Gone, + /// ComfyUI refused the approval (bad id, write failure). + Refused(String), +} + +/// Where `GET /view` finds a staged result. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, Deserialize)] +pub struct ViewRef { + pub filename: String, + #[serde(default)] + pub subfolder: String, + #[serde(default = "temp", rename = "type")] + pub kind: String, +} + +fn temp() -> String { + "temp".into() +} + +/// A result the save node staged: ComfyUI's pending id, where to preview it, +/// and the library path when it was already saved (never, with approval on). +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, Deserialize)] +pub struct Staged { + pub id: String, + pub view: ViewRef, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct HistoryEntry { + pub completed: bool, + /// `success` / `error`. + pub status_str: String, + /// ComfyUI's exception message when the run failed. + pub error: Option, + pub staged: Vec, + /// Last event time, unix millis. + pub timestamp_ms: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct QueueState { + pub running: Vec, + pub pending: Vec, +} + +impl QueueState { + pub fn is_empty(&self) -> bool { + self.running.is_empty() && self.pending.is_empty() + } + + /// Jobs ahead of `prompt_id`; `None` when it isn't queued. + pub fn position_of(&self, prompt_id: &str) -> Option<(bool, usize)> { + if self.running.iter().any(|p| p == prompt_id) { + return Some((true, 0)); + } + self.pending + .iter() + .position(|p| p == prompt_id) + .map(|i| (false, self.running.len() + i)) + } +} + +#[derive(Clone)] +pub struct ComfyClient { + http: reqwest::Client, + base: String, + /// Client id sent with submissions. + client_id: String, +} + +impl ComfyClient { + pub fn new(base_url: &str, timeout: Duration) -> Self { + let http = reqwest::Client::builder() + .timeout(timeout) + .connect_timeout(Duration::from_secs(3).min(timeout)) + .build() + .expect("reqwest client"); + ComfyClient { + http, + base: base_url.trim_end_matches('/').to_string(), + client_id: "imageapi".to_string(), + } + } + + fn url(&self, path: &str) -> String { + format!("{}{}", self.base, path) + } + + /// `GET /system_stats` with its own (short) timeout. + pub async fn probe(&self, timeout: Duration) -> bool { + self.http + .get(self.url("/system_stats")) + .timeout(timeout) + .send() + .await + .is_ok_and(|r| r.status().is_success()) + } + + pub async fn submit(&self, prompt: &Value) -> Result { + let body = json!({ "prompt": prompt, "client_id": self.client_id }); + let resp = self + .http + .post(self.url("/prompt")) + .json(&body) + .send() + .await + .map_err(|e| SubmitError::Comfy(e.into()))?; + let status = resp.status(); + let text = resp + .text() + .await + .map_err(|e| SubmitError::Comfy(e.into()))?; + let json: Value = serde_json::from_str(&text).unwrap_or(Value::Null); + if status.is_success() { + // A 200 can still carry node_errors for outputs that were pruned; + // the prompt itself queued. + return json["prompt_id"] + .as_str() + .map(str::to_string) + .ok_or_else(|| { + SubmitError::Comfy(ComfyError::BadResponse(format!("no prompt_id in {text}"))) + }); + } + if status.as_u16() == 400 { + return Err(SubmitError::Rejected { + message: rejection_message(&json, &text), + node_errors: json["node_errors"] + .as_object() + .is_some_and(|m| !m.is_empty()), + }); + } + Err(SubmitError::Comfy(ComfyError::Http { + status: status.as_u16(), + body: truncate(&text), + })) + } + + pub async fn queue(&self) -> Result { + let json: Value = self.get_json("/queue").await?; + let ids = |key: &str| -> Vec { + json[key] + .as_array() + .map(|a| { + a.iter() + .filter_map(|item| item.get(1).and_then(Value::as_str)) + .map(str::to_string) + .collect() + }) + .unwrap_or_default() + }; + Ok(QueueState { + running: ids("queue_running"), + pending: ids("queue_pending"), + }) + } + + /// `GET /history/{id}`; `None` when ComfyUI doesn't know the id. + pub async fn history(&self, prompt_id: &str) -> Result, ComfyError> { + let json: Value = self + .get_json(&format!("/history/{}", urlencoding::encode(prompt_id))) + .await?; + Ok(json.get(prompt_id).map(parse_history_entry)) + } + + /// The newest finished run, for the idle check. + pub async fn newest_history(&self) -> Result, ComfyError> { + let json: Value = self.get_json("/history?max_items=1").await?; + Ok(json + .as_object() + .and_then(|m| m.values().next()) + .map(parse_history_entry)) + } + + /// `GET /view`: the staged PNG. `Ok(None)` for 404. + pub async fn view(&self, r: &ViewRef) -> Result, ComfyError> { + let resp = self + .http + .get(self.url("/view")) + .query(&[ + ("filename", r.filename.as_str()), + ("subfolder", r.subfolder.as_str()), + ("type", r.kind.as_str()), + ]) + .send() + .await?; + if resp.status().as_u16() == 404 { + return Ok(None); + } + if !resp.status().is_success() { + return Err(ComfyError::Http { + status: resp.status().as_u16(), + body: truncate(&resp.text().await.unwrap_or_default()), + }); + } + let content_type = resp + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .unwrap_or("image/png") + .to_string(); + Ok(Some((content_type, resp.bytes().await?))) + } + + /// `POST /imageapi/approve`: ComfyUI writes the file; returns its path + /// as ComfyUI sees it. + pub async fn approve(&self, pending_id: &str) -> Result { + let resp = self + .http + .post(self.url("/imageapi/approve")) + .json(&json!({ "id": pending_id })) + .send() + .await + .map_err(|e| ApproveError::Comfy(e.into()))?; + let status = resp.status(); + let text = resp + .text() + .await + .map_err(|e| ApproveError::Comfy(e.into()))?; + let json: Value = serde_json::from_str(&text).unwrap_or(Value::Null); + if status.is_success() { + return json["path"].as_str().map(str::to_string).ok_or_else(|| { + ApproveError::Comfy(ComfyError::BadResponse(format!("no path in {text}"))) + }); + } + let message = json["error"].as_str().unwrap_or(&text).to_string(); + match status.as_u16() { + 404 => Err(ApproveError::Gone), + 400 | 500 => Err(ApproveError::Refused(truncate(&message))), + s => Err(ApproveError::Comfy(ComfyError::Http { + status: s, + body: truncate(&text), + })), + } + } + + /// `POST /free`: unload models and release cached memory. + pub async fn free(&self) -> Result<(), ComfyError> { + let resp = self + .http + .post(self.url("/free")) + .json(&json!({ "unload_models": true, "free_memory": true })) + .send() + .await?; + if resp.status().is_success() { + Ok(()) + } else { + Err(ComfyError::Http { + status: resp.status().as_u16(), + body: truncate(&resp.text().await.unwrap_or_default()), + }) + } + } + + async fn get_json(&self, path: &str) -> Result { + let resp = self.http.get(self.url(path)).send().await?; + if !resp.status().is_success() { + return Err(ComfyError::Http { + status: resp.status().as_u16(), + body: truncate(&resp.text().await.unwrap_or_default()), + }); + } + Ok(resp.json().await?) + } +} + +fn truncate(s: &str) -> String { + s.chars().take(300).collect() +} + +fn rejection_message(json: &Value, raw: &str) -> String { + let err = &json["error"]; + let head = err["message"].as_str().or_else(|| err.as_str()); + let detail = err["details"].as_str().filter(|d| !d.is_empty()); + // Per-node errors say which input is wrong. + let nodes: Vec = json["node_errors"] + .as_object() + .map(|m| { + m.values() + .filter_map(|n| { + let class = n["class_type"].as_str().unwrap_or("node"); + let first = n["errors"].as_array()?.first()?; + Some(format!( + "{class}: {}", + first["message"].as_str().unwrap_or("invalid") + )) + }) + .collect() + }) + .unwrap_or_default(); + let mut parts: Vec = Vec::new(); + parts.extend(head.map(str::to_string)); + parts.extend(detail.map(str::to_string)); + parts.extend(nodes); + if parts.is_empty() { + truncate(raw) + } else { + truncate(&parts.join("; ")) + } +} + +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 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| { + let (kind, data) = (m.get(0)?.as_str()?, m.get(1)?); + (kind == "execution_error").then(|| { + data["exception_message"] + .as_str() + .unwrap_or("execution error") + .trim() + .to_string() + }) + }); + let timestamp_ms = messages + .iter() + .filter_map(|m| m.get(1)?.get("timestamp")?.as_i64()) + .max(); + + let mut staged = Vec::new(); + if let Some(outputs) = entry["outputs"].as_object() { + for out in outputs.values() { + let Some(results) = out["imageapi_versions"].as_array() else { + continue; + }; + let images = out["images"].as_array().cloned().unwrap_or_default(); + for (i, r) in results.iter().enumerate() { + let (Some(id), Some(image)) = (r["id"].as_str(), images.get(i)) else { + continue; + }; + if let Ok(view) = serde_json::from_value::(image.clone()) { + staged.push(Staged { + id: id.to_string(), + view, + }); + } + } + } + } + HistoryEntry { + completed, + status_str, + error, + staged, + timestamp_ms, + } +} diff --git a/src/enhance/config.rs b/src/enhance/config.rs new file mode 100644 index 0000000..1cac49f --- /dev/null +++ b/src/enhance/config.rs @@ -0,0 +1,263 @@ +//! Enhance configuration (#192): one `enhance.toml`, next to `llm.toml`. +//! +//! Path from `ENHANCE_CONFIG_PATH`, else `./enhance.toml`. A missing file at +//! the default path means Enhance is simply off (no `url`); a missing +//! explicit path is an error. Nothing here is a secret: ComfyUI has no auth +//! and stays LAN-only. `enhance.example.toml` is the annotated schema (a +//! test parses it). + +use std::path::{Path, PathBuf}; + +use anyhow::{Context as _, Result, bail}; +use serde::Deserialize; + +use super::path_map::PathMap; + +const DEFAULT_PATH: &str = "enhance.toml"; + +/// How the app's text combines with the template's own prompt. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, serde::Serialize, Default)] +#[serde(rename_all = "lowercase")] +pub enum PromptMode { + /// The template's prompt is a full base instruction; the app's text is + /// appended after `prompt_join`. Empty text leaves the prompt as is. + #[default] + Append, + /// The app's text *is* the prompt (and is required). + Replace, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct EnhanceConfig { + /// ComfyUI base URL, e.g. `http://192.168.10.218:8188`. Unset = Enhance + /// is disabled (`/enhance/templates` reports `available: false`). + pub url: Option, + /// Per-request timeout for ComfyUI calls. + pub timeout_secs: u64, + /// How long the `available` probe result is cached. + pub probe_ttl_secs: u64, + /// Job poll interval. + pub poll_secs: u64, + /// Idle time (empty queue, newest history entry this old) before + /// ImageApi calls `POST /free` once. 0 disables it. + pub idle_free_secs: u64, + /// An active job whose ComfyUI stays unreachable this long becomes + /// `unavailable`. + pub unreachable_grace_secs: u64, + /// Template preselected by clients; default = the first one. + pub default_template: Option, + /// Clients warn above this many megapixels in the original. + pub warn_above_megapixels: f64, + /// Library path prefix <-> ComfyUI path prefix. Longest prefix wins. + pub path_map: Vec, + #[serde(rename = "template")] + pub templates: Vec, +} + +impl Default for EnhanceConfig { + fn default() -> Self { + Self { + url: None, + timeout_secs: 15, + probe_ttl_secs: 10, + poll_secs: 2, + idle_free_secs: 300, + unreachable_grace_secs: 120, + default_template: None, + warn_above_megapixels: 1.5, + path_map: Vec::new(), + templates: Vec::new(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct PathMapping { + /// Prefix as ImageApi sees it, e.g. `/media/nas/`. + pub from: String, + /// Prefix as ComfyUI sees it, e.g. `Z:\` or `\\nas\`. + pub to: String, +} + +#[derive(Debug, Clone, PartialEq, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TemplateConfig { + /// Stable id clients send (`POST /enhance { template }`). + pub name: String, + pub label: String, + #[serde(default)] + pub description: String, + /// API-format workflow JSON (ComfyUI -> Workflow -> Export (API)). + /// Relative paths resolve against the config file's directory. + pub file: PathBuf, + #[serde(default)] + pub prompt_mode: PromptMode, + /// Text between the template's prompt and the app's text (append mode). + #[serde(default = "default_join")] + pub prompt_join: String, + /// Approximate output size, for the client's large-original warning. + #[serde(default = "default_output_mp")] + pub output_megapixels: f64, +} + +fn default_join() -> String { + " ".into() +} + +fn default_output_mp() -> f64 { + 1.0 +} + +impl EnhanceConfig { + pub fn validate(&self) -> Result<()> { + if let Some(url) = &self.url + && !(url.starts_with("http://") || url.starts_with("https://")) + { + bail!("url must start with http:// or https://, got {url:?}"); + } + for (key, v) in [ + ("timeout_secs", self.timeout_secs), + ("probe_ttl_secs", self.probe_ttl_secs), + ("poll_secs", self.poll_secs), + ("unreachable_grace_secs", self.unreachable_grace_secs), + ] { + if v == 0 { + bail!("{key} must be greater than 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(); + for t in &self.templates { + let ok = !t.name.is_empty() + && t.name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'); + if !ok { + bail!( + "template name {:?} must be non-empty and use only letters, digits, '_' and '-'", + t.name + ); + } + if !seen.insert(t.name.as_str()) { + bail!("duplicate template name {:?}", t.name); + } + if t.label.trim().is_empty() { + bail!("template {:?} needs a label", t.name); + } + if t.output_megapixels.is_nan() || t.output_megapixels <= 0.0 { + bail!("template {:?}: output_megapixels must be > 0", t.name); + } + } + if let Some(d) = &self.default_template + && !self.templates.iter().any(|t| &t.name == d) + { + bail!("default_template {d:?} is not a configured template"); + } + // Parses every mapping (rejects empty prefixes). + PathMap::new(&self.path_map)?; + Ok(()) + } + + /// The preselected template's name. + pub fn default_name(&self) -> Option<&str> { + self.default_template + .as_deref() + .or_else(|| self.templates.first().map(|t| t.name.as_str())) + } +} + +/// Where the config was read from, and what it said. +#[derive(Debug)] +pub struct Loaded { + pub config: EnhanceConfig, + /// Directory relative template `file`s resolve against. + pub base_dir: PathBuf, + /// The file, when one was read. + pub path: Option, +} + +/// Read `enhance.toml` (see module docs). A bad file is an error so a typo +/// fails startup instead of silently disabling the feature. +pub fn load() -> Result { + let (path, explicit) = match std::env::var("ENHANCE_CONFIG_PATH") { + Ok(p) if !p.trim().is_empty() => (PathBuf::from(p.trim()), true), + _ => (PathBuf::from(DEFAULT_PATH), false), + }; + match std::fs::read_to_string(&path) { + Ok(text) => { + let config = parse(&text).with_context(|| format!("in {}", path.display()))?; + let base_dir = path + .parent() + .filter(|p| !p.as_os_str().is_empty()) + .map(Path::to_path_buf) + .unwrap_or_else(|| PathBuf::from(".")); + Ok(Loaded { + config, + base_dir, + path: Some(path), + }) + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound && !explicit => Ok(Loaded { + config: EnhanceConfig::default(), + base_dir: PathBuf::from("."), + path: None, + }), + Err(e) => Err(e).with_context(|| format!("reading enhance config {}", path.display())), + } +} + +pub fn parse(text: &str) -> Result { + let config: EnhanceConfig = toml::from_str(text)?; + config.validate()?; + Ok(config) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn example_file_parses() { + let cfg = parse(include_str!("../../enhance.example.toml")).unwrap(); + assert!(cfg.url.is_some()); + assert_eq!(cfg.templates[0].name, "restore"); + assert_eq!(cfg.default_name(), Some("restore")); + assert!( + cfg.templates + .iter() + .any(|t| t.prompt_mode == PromptMode::Replace) + ); + } + + #[test] + fn empty_file_is_disabled_defaults() { + let cfg = parse("").unwrap(); + assert_eq!(cfg, EnhanceConfig::default()); + assert!(cfg.url.is_none()); + } + + #[test] + fn rejects_unknown_keys_and_bad_values() { + assert!(parse("nope = 1").is_err()); + assert!(parse("url = \"ftp://x\"").is_err()); + assert!(parse("poll_secs = 0").is_err()); + assert!(parse("default_template = \"x\"").is_err()); + let dup = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n[[template]]\nname='a'\nlabel='B'\nfile='b.json'"; + assert!(parse(dup).is_err()); + assert!(parse("[[template]]\nname='a b'\nlabel='A'\nfile='a.json'").is_err()); + assert!(parse("[[path_map]]\nfrom=''\nto='Z:\\\\'").is_err()); + } + + #[test] + fn template_defaults() { + let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap(); + let t = &cfg.templates[0]; + assert_eq!(t.prompt_mode, PromptMode::Append); + assert_eq!(t.prompt_join, " "); + assert_eq!(t.output_megapixels, 1.0); + } +} diff --git a/src/enhance/handlers.rs b/src/enhance/handlers.rs new file mode 100644 index 0000000..1757b68 --- /dev/null +++ b/src/enhance/handlers.rs @@ -0,0 +1,523 @@ +//! HTTP surface of Enhance (#192). ImageApi relays everything, previews +//! included: the apps never talk to ComfyUI. +//! +//! - `GET /enhance/templates` +//! - `POST /enhance` +//! - `GET /enhance/jobs?path=` +//! - `GET /enhance/jobs/{id}` +//! - `GET /enhance/jobs/{id}/preview/{n}` +//! - `POST /enhance/jobs/{id}/approve` +//! - `POST /enhance/jobs/{id}/discard` + +use std::path::{Component, Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use actix_web::{HttpResponse, Responder, get, post, web}; +use log::warn; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +use super::config::PromptMode; +use super::service::{EnhanceError, EnhanceService, SourceFile}; +use crate::data::Claims; +use crate::database::ExifDao; +use crate::database::enhance_dao::{EnhanceJobRow, STATUS_APPROVED, STATUS_READY}; +use crate::libraries::{self, Library}; +use crate::state::AppState; +use crate::utils::normalize_path; +use crate::versions::handlers::build_response; + +type ExifData = web::Data>>; + +fn error_response(e: &EnhanceError) -> HttpResponse { + let (mut builder, code) = match e { + EnhanceError::Disabled => (HttpResponse::ServiceUnavailable(), "disabled"), + EnhanceError::Unavailable => (HttpResponse::ServiceUnavailable(), "unavailable"), + EnhanceError::UnknownTemplate(_) => (HttpResponse::NotFound(), "unknown_template"), + EnhanceError::BadRequest(_) => (HttpResponse::BadRequest(), "bad_request"), + EnhanceError::Rejected(_) => (HttpResponse::UnprocessableEntity(), "rejected"), + EnhanceError::Expired => (HttpResponse::Gone(), "expired"), + EnhanceError::NotFound => (HttpResponse::NotFound(), "not_found"), + EnhanceError::Unplaceable(_) => (HttpResponse::BadGateway(), "unplaceable"), + EnhanceError::Upstream(_) => (HttpResponse::BadGateway(), "upstream"), + EnhanceError::Internal(_) => (HttpResponse::InternalServerError(), "internal"), + }; + builder.json(json!({ "error": e.message(), "code": code })) +} + +fn bad(msg: &str) -> HttpResponse { + error_response(&EnhanceError::BadRequest(msg.to_string())) +} + +// ---- views -------------------------------------------------------------- + +#[derive(Debug, Serialize)] +struct PendingView { + /// Index for `preview/{n}` and `approve { n }`. + n: usize, + /// Path of the preview on this API. + preview: String, + approved: bool, +} + +#[derive(Debug, Serialize)] +struct SavedView { + n: usize, + path: Option, + library_id: Option, +} + +#[derive(Debug, Serialize)] +pub struct JobView { + id: String, + status: String, + #[serde(skip_serializing_if = "Option::is_none")] + queue_position: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, + path: String, + library_id: i32, + template: String, + #[serde(skip_serializing_if = "Option::is_none")] + prompt: Option, + #[serde(skip_serializing_if = "Option::is_none")] + reference_path: Option, + created_at: i64, + updated_at: i64, + pending: Vec, + saved: Vec, +} + +pub fn job_view(job: &EnhanceJobRow) -> JobView { + let saved = EnhanceService::saved_of(job); + // An expired job's staged files are gone; don't offer them. + let pending = if job.status == STATUS_READY || job.status == STATUS_APPROVED { + EnhanceService::pending_of(job) + .iter() + .enumerate() + .map(|(n, _)| PendingView { + n, + preview: format!("/enhance/jobs/{}/preview/{n}", job.id), + approved: saved.iter().any(|s| s.n == n), + }) + .collect() + } else { + Vec::new() + }; + let queued = matches!(job.status.as_str(), "queued" | "running"); + JobView { + id: job.id.clone(), + status: job.status.clone(), + queue_position: job.queue_position.filter(|_| queued), + error: job.error.clone(), + path: job.rel_path.clone(), + library_id: job.library_id, + template: job.template.clone(), + prompt: job.prompt.clone(), + reference_path: job.reference_path.clone(), + created_at: job.created_at, + updated_at: job.updated_at, + pending, + saved: saved + .into_iter() + .map(|s| SavedView { + n: s.n, + path: s.rel_path, + library_id: s.library_id, + }) + .collect(), + } +} + +// ---- file resolution ---------------------------------------------------- + +fn safe_rel(rel: &str) -> bool { + !rel.is_empty() + && Path::new(rel) + .components() + .all(|c| matches!(c, Component::Normal(_))) +} + +/// An indexed, hashed file in a library. +fn resolve_file( + app_state: &AppState, + exif_dao: &Mutex>, + path: &str, + library: Option<&str>, +) -> Result { + let library_id = match libraries::resolve_library_param_state(app_state, library) { + Ok(l) => l.map(|l| l.id), + Err(msg) => return Err(bad(&msg)), + }; + let rel = normalize_path(path); + if !safe_rel(&rel) { + return Err(bad("invalid path")); + } + let mut rows = exif_dao + .lock() + .expect("Unable to lock ExifDao") + .get_exif_batch(library_id, std::slice::from_ref(&rel)) + .map_err(|_| error_response(&EnhanceError::Internal("lookup failed".into())))?; + rows.sort_by_key(|r| r.id); + let Some(row) = rows.iter().find(|r| r.content_hash.is_some()) else { + return Err(error_response(&EnhanceError::NotFound)); + }; + let Some(lib) = app_state.library_by_id(row.library_id) else { + return Err(error_response(&EnhanceError::NotFound)); + }; + let abs = Path::new(&lib.root_path).join(&rel); + if !abs.is_file() { + return Err(error_response(&EnhanceError::NotFound)); + } + Ok(SourceFile { + library_id: row.library_id, + rel_path: rel, + abs, + hash: row.content_hash.clone().unwrap_or_default(), + }) +} + +/// 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)> { + let mut ordered: Vec<&Library> = libs.iter().collect(); + ordered.sort_by_key(|l| l.id != prefer); + for lib in ordered { + if let Ok(rest) = abs.strip_prefix(&lib.root_path) { + let rel = rest.to_str()?.replace('\\', "/"); + if safe_rel(&rel) { + return Some((lib.id, rel)); + } + } + } + None +} + +// ---- templates ---------------------------------------------------------- + +#[derive(Serialize)] +struct TemplateView { + name: String, + label: String, + description: String, + accepts_prompt: bool, + prompt_mode: PromptMode, + /// Replace-mode templates need non-empty text (empty is a 400). + prompt_required: bool, + accepts_reference: bool, + default: bool, + /// Approximate output size; clients warn when the original is larger. + output_megapixels: f64, + warn_above_megapixels: f64, +} + +#[get("/enhance/templates")] +pub async fn get_templates(_: Claims, svc: web::Data) -> impl Responder { + let available = svc.available().await; + let reason = if !svc.is_configured() { + Some(super::service::MSG_DISABLED) + } else if !available { + Some(super::service::MSG_UNAVAILABLE) + } else { + None + }; + let default = svc.default_template().map(str::to_string); + let templates: Vec = svc + .templates() + .iter() + .map(|t| TemplateView { + name: t.cfg.name.clone(), + label: t.cfg.label.clone(), + description: t.cfg.description.clone(), + accepts_prompt: t.accepts_prompt(), + prompt_mode: t.cfg.prompt_mode, + prompt_required: t.prompt_required(), + accepts_reference: t.accepts_reference(), + default: default.as_deref() == Some(t.cfg.name.as_str()), + output_megapixels: t.cfg.output_megapixels, + warn_above_megapixels: svc.config.warn_above_megapixels, + }) + .collect(); + let mut body = json!({ + "available": available, + "default": default, + "warn_above_megapixels": svc.config.warn_above_megapixels, + "templates": templates, + }); + if let Some(r) = reason { + body["reason"] = json!(r); + } + HttpResponse::Ok().json(body) +} + +// ---- jobs --------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +pub struct LibraryQuery { + pub library: Option, +} + +#[derive(Debug, Deserialize)] +pub struct SubmitBody { + pub path: String, + pub template: String, + #[serde(default)] + pub prompt: Option, + #[serde(default)] + pub reference_path: Option, + #[serde(default)] + pub library: Option, +} + +#[post("/enhance")] +#[tracing::instrument(name = "enhance_submit", skip_all)] +pub async fn submit( + _: Claims, + query: web::Query, + body: web::Json, + svc: web::Data, + app_state: web::Data, + exif_dao: ExifData, +) -> impl Responder { + if !svc.is_configured() { + return error_response(&EnhanceError::Disabled); + } + let library = body.library.as_deref().or(query.library.as_deref()); + let photo = match resolve_file(&app_state, &exif_dao, &body.path, library) { + Ok(f) => f, + Err(r) => return r, + }; + 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()); + match resolve_file(&app_state, &exif_dao, p, lib.as_deref()) { + Ok(f) => Some(f), + Err(r) => return r, + } + } + None => None, + }; + match svc + .submit( + &photo, + &body.template, + body.prompt.as_deref(), + reference.as_ref(), + ) + .await + { + Ok(job) => HttpResponse::Ok().json(job_view(&job)), + Err(e) => error_response(&e), + } +} + +#[get("/enhance/jobs/{id}")] +pub async fn get_job( + _: Claims, + id: web::Path, + svc: web::Data, +) -> impl Responder { + match svc.get_job(&id) { + Ok(job) => HttpResponse::Ok().json(job_view(&job)), + Err(e) => error_response(&e), + } +} + +#[derive(Debug, Deserialize)] +pub struct JobsQuery { + pub path: String, + pub library: Option, +} + +/// A photo's recent jobs, newest first, so the viewer can pick up an enhance +/// in progress after reopening. Discarded jobs are not listed. +#[get("/enhance/jobs")] +pub async fn list_jobs( + _: Claims, + q: web::Query, + svc: web::Data, + app_state: web::Data, + exif_dao: ExifData, +) -> impl Responder { + let file = match resolve_file(&app_state, &exif_dao, &q.path, q.library.as_deref()) { + Ok(f) => f, + Err(r) => return r, + }; + match svc.jobs_for_hash(&file.hash) { + Ok(jobs) => { + let views: Vec = jobs.iter().map(job_view).collect(); + HttpResponse::Ok().json(json!({ "jobs": views })) + } + Err(e) => error_response(&e), + } +} + +#[get("/enhance/jobs/{id}/preview/{n}")] +pub async fn preview( + _: Claims, + path: web::Path<(String, usize)>, + svc: web::Data, +) -> impl Responder { + let (id, n) = path.into_inner(); + let job = match svc.get_job(&id) { + Ok(j) => j, + Err(e) => return error_response(&e), + }; + match svc.preview(&job, n).await { + Ok((content_type, bytes)) => HttpResponse::Ok() + .content_type(content_type) + .insert_header(("Cache-Control", "private, max-age=3600")) + .body(bytes), + Err(e) => error_response(&e), + } +} + +#[post("/enhance/jobs/{id}/discard")] +pub async fn discard( + _: Claims, + id: web::Path, + svc: web::Data, +) -> impl Responder { + match svc.discard(&id) { + Ok(job) => HttpResponse::Ok().json(job_view(&job)), + Err(e) => error_response(&e), + } +} + +#[derive(Debug, Deserialize)] +pub struct ApproveBody { + #[serde(default)] + pub n: usize, +} + +/// Ask ComfyUI to write result `n` into the library, ingest that file now, +/// and answer with the version (an `/image/versions` `versions[]` entry). +#[post("/enhance/jobs/{id}/approve")] +#[tracing::instrument(name = "enhance_approve", skip_all)] +pub async fn approve( + _: Claims, + id: web::Path, + body: Option>, + svc: web::Data, + app_state: web::Data, + exif_dao: ExifData, +) -> impl Responder { + let n = body.map(|b| b.n).unwrap_or(0); + let job = match svc.get_job(&id) { + Ok(j) => j, + Err(e) => return error_response(&e), + }; + let libs = app_state.libraries.clone(); + let prefer = job.library_id; + let place = move |abs: &Path| place_in_libraries(&libs, prefer, abs); + let (job, saved) = match svc.approve(&id, n, &place).await { + Ok(v) => v, + Err(e) => return error_response(&e), + }; + let (Some(library_id), Some(rel)) = (saved.library_id, saved.rel_path.clone()) else { + return error_response(&EnhanceError::Unplaceable( + "the saved file isn't in a library".into(), + )); + }; + let Some(library) = app_state.library_by_id(library_id).cloned() else { + return error_response(&EnhanceError::Unplaceable("unknown library".into())); + }; + let abs = PathBuf::from(&library.root_path).join(&rel); + + // The file was written by another machine over a share; give a lagging + // mount a moment to show it. + for _ in 0..25 { + if abs.is_file() { + break; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + if !abs.is_file() { + return error_response(&EnhanceError::Unplaceable(format!( + "ComfyUI saved {rel} but it isn't visible to ImageApi yet; it will appear after the next scan" + ))); + } + + 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 hash = match ingested { + Ok(Some(h)) => h, + _ => { + warn!("enhance: ingest of {rel} failed"); + return HttpResponse::Accepted().json(json!({ + "linked": false, + "path": rel, + "library_id": library_id, + "job": job_view(&job), + })); + } + }; + + // Which version row did that become? + let entry = { + 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(), + }; + build_response(dao.as_mut(), &original, job.library_id) + .ok() + .and_then(|r| r.versions.into_iter().find(|v| v.hash == hash)) + }; + match entry { + Some(v) => HttpResponse::Ok().json(v), + None => HttpResponse::Accepted().json(json!({ + "linked": false, + "path": rel, + "library_id": library_id, + "job": job_view(&job), + })), + } +} + +/// Register the saved file (if the watcher hasn't yet) and link it as a +/// version. Idempotent. Returns the file's content hash. +pub fn ingest_saved( + library: &Library, + rel: &str, + abs: &Path, + exif_dao: &Arc>>, + version_dao: &Arc>>, +) -> Option { + let rel_s = rel.to_string(); + let existing = exif_dao + .lock() + .expect("Unable to lock ExifDao") + .get_exif_batch(Some(library.id), std::slice::from_ref(&rel_s)) + .ok()?; + let file = vec![(abs.to_path_buf(), rel_s.clone())]; + if existing.is_empty() { + crate::watcher::register_new_files(library, file.clone(), exif_dao); + } + let linked = crate::versions::link_new_files(library, &file, exif_dao, version_dao) + + crate::versions::sweep_name_links(library, exif_dao, version_dao); + if linked > 0 { + crate::duplicates::invalidate_caches(); + } + exif_dao + .lock() + .expect("Unable to lock ExifDao") + .get_exif_batch(Some(library.id), std::slice::from_ref(&rel_s)) + .ok()? + .into_iter() + .find_map(|r| r.content_hash) +} diff --git a/src/enhance/integration_tests.rs b/src/enhance/integration_tests.rs new file mode 100644 index 0000000..6abe693 --- /dev/null +++ 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/mod.rs b/src/enhance/mod.rs new file mode 100644 index 0000000..2e31846 --- /dev/null +++ b/src/enhance/mod.rs @@ -0,0 +1,25 @@ +//! Enhance on demand (#192): run a ComfyUI restore/edit workflow on a photo +//! from the apps, approve the result in-app, and have it land in the library +//! as a version (#191). +//! +//! - [`config`]: `enhance.toml` (templates, path map, timings). +//! - [`template`]: workflow loading/validation and per-photo filling. +//! - [`comfy`]: the ComfyUI HTTP calls used. +//! - [`service`]: submit, poll, approve, idle `/free`. +//! - [`handlers`]: the HTTP endpoints. +//! +//! ImageApi never writes under a library root: ComfyUI writes the approved +//! file, ImageApi only ingests it. + +pub mod comfy; +pub mod config; +pub mod handlers; +pub mod path_map; +pub mod service; +pub mod template; + +#[cfg(test)] +mod integration_tests; + +pub use handlers::{approve, discard, get_job, get_templates, list_jobs, preview, submit}; +pub use service::EnhanceService; diff --git a/src/enhance/path_map.rs b/src/enhance/path_map.rs new file mode 100644 index 0000000..7ceb098 --- /dev/null +++ b/src/enhance/path_map.rs @@ -0,0 +1,175 @@ +//! Library path <-> ComfyUI path mapping (#192). +//! +//! ImageApi sees `/media/nas/Media/Pictures/a.jpg`; ComfyUI (on Windows) +//! sees `Z:\Media\Pictures\a.jpg` or `\\nas\Media\Pictures\a.jpg`. Each +//! mapping is a prefix pair; the longest matching prefix wins and the +//! separators are converted. Used in both directions: library -> ComfyUI on +//! submit, and ComfyUI -> library to recognise the path `approve` returns. +//! With no mappings configured, paths pass through unchanged. + +use std::path::{Path, PathBuf}; + +use anyhow::{Result, bail}; + +use super::config::PathMapping; + +#[derive(Debug, Clone)] +struct Entry { + /// Library side, `/`-separated, always ending in `/`. + from: String, + /// ComfyUI side without a trailing separator. + to: String, + /// Separator ComfyUI's side uses. + to_sep: char, +} + +#[derive(Debug, Clone, Default)] +pub struct PathMap { + entries: Vec, +} + +impl PathMap { + pub fn new(mappings: &[PathMapping]) -> Result { + let mut entries = Vec::new(); + for m in mappings { + let from = m.from.replace('\\', "/"); + let to = m.to.trim(); + if from.trim_matches('/').is_empty() || to.is_empty() { + bail!("path_map entries need non-empty from and to ({m:?})"); + } + let to_sep = if to.contains('\\') { '\\' } else { '/' }; + entries.push(Entry { + from: format!("{}/", from.trim_end_matches('/')), + to: to.trim_end_matches(['/', '\\']).to_string(), + to_sep, + }); + } + Ok(PathMap { entries }) + } + + pub fn is_identity(&self) -> bool { + self.entries.is_empty() + } + + /// The path ComfyUI should open for a library file. `None` when mappings + /// exist but none covers the file. + pub fn to_comfy(&self, abs: &Path) -> Option { + let text = abs.to_str()?.replace('\\', "/"); + if self.entries.is_empty() { + return Some(abs.to_string_lossy().into_owned()); + } + let best = self + .entries + .iter() + .filter(|e| text.starts_with(e.from.as_str())) + .max_by_key(|e| e.from.len())?; + let rest = &text[best.from.len()..]; + let sep = best.to_sep.to_string(); + Some(format!("{}{}{}", best.to, sep, rest.replace('/', &sep))) + } + + /// The library-side absolute path of a path ComfyUI reported. Matching + /// is case-insensitive (Windows). `None` when no mapping covers it or it + /// tries to climb out with `..`. + pub fn from_comfy(&self, comfy: &str) -> Option { + let text = comfy.trim().replace('\\', "/"); + if self.entries.is_empty() { + return Some(PathBuf::from(text)); + } + let best = self + .entries + .iter() + .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())) + }) + .max_by_key(|(_, len)| *len)?; + let rest = &text[best.1..]; + if rest.split('/').any(|c| c == "..") { + return None; + } + Some(PathBuf::from(format!("{}{}", best.0.from, rest))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn map(pairs: &[(&str, &str)]) -> PathMap { + let m: Vec = pairs + .iter() + .map(|(f, t)| PathMapping { + from: f.to_string(), + to: t.to_string(), + }) + .collect(); + PathMap::new(&m).unwrap() + } + + #[test] + fn drive_letter_round_trip() { + let m = map(&[("/media/nas/", "Z:\\")]); + let c = m + .to_comfy(Path::new("/media/nas/Media/Pictures/a b.jpg")) + .unwrap(); + assert_eq!(c, "Z:\\Media\\Pictures\\a b.jpg"); + assert_eq!( + m.from_comfy(&c).unwrap(), + PathBuf::from("/media/nas/Media/Pictures/a b.jpg") + ); + // Drive letters are case-insensitive. + assert_eq!( + m.from_comfy("z:\\x\\y.png").unwrap(), + PathBuf::from("/media/nas/x/y.png") + ); + } + + #[test] + fn unc_target_and_missing_trailing_slash() { + let m = map(&[("/media/nas", "\\\\nas\\")]); + assert_eq!( + m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(), + "\\\\nas\\Media\\a.jpg" + ); + assert_eq!( + m.from_comfy("\\\\nas\\Media\\a.jpg").unwrap(), + PathBuf::from("/media/nas/Media/a.jpg") + ); + } + + #[test] + fn longest_prefix_wins_both_ways() { + let m = map(&[("/media/nas/", "Z:\\"), ("/media/nas/Media/", "Y:\\")]); + assert_eq!( + m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(), + "Y:\\a.jpg" + ); + assert_eq!( + m.to_comfy(Path::new("/media/nas/Other/a.jpg")).unwrap(), + "Z:\\Other\\a.jpg" + ); + assert_eq!( + m.from_comfy("Y:\\a.jpg").unwrap(), + PathBuf::from("/media/nas/Media/a.jpg") + ); + } + + #[test] + fn uncovered_and_escaping_paths_are_none() { + let m = map(&[("/media/nas/", "Z:\\")]); + assert!(m.to_comfy(Path::new("/elsewhere/a.jpg")).is_none()); + assert!(m.to_comfy(Path::new("/media/nasty/a.jpg")).is_none()); + assert!(m.from_comfy("C:\\a.png").is_none()); + assert!(m.from_comfy("Z:\\..\\etc\\passwd").is_none()); + } + + #[test] + fn no_mappings_is_identity() { + let m = PathMap::default(); + assert_eq!(m.to_comfy(Path::new("/a/b.jpg")).unwrap(), "/a/b.jpg"); + assert_eq!(m.from_comfy("/a/b.png").unwrap(), PathBuf::from("/a/b.png")); + } +} diff --git a/src/enhance/service.rs b/src/enhance/service.rs new file mode 100644 index 0000000..8a3ba53 --- /dev/null +++ b/src/enhance/service.rs @@ -0,0 +1,742 @@ +//! Enhance orchestration (#192): submit, poll, approve, idle `/free`. +//! +//! The job row (`enhance_jobs`) is the only state that matters. ComfyUI is +//! allowed to be down (a normal state) and to forget everything when it +//! restarts (unapproved results are cleared), so every transition tolerates +//! both. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use log::{info, warn}; +use serde::{Deserialize, Serialize}; + +use super::comfy::{ApproveError, ComfyClient, ComfyError, HistoryEntry, Staged, SubmitError}; +use super::config::{EnhanceConfig, Loaded, PromptMode}; +use super::path_map::PathMap; +use super::template::Template; +use crate::database::enhance_dao::{ + EnhanceJobDao, EnhanceJobRow, JobUpdate, STATUS_APPROVED, STATUS_EXPIRED, STATUS_FAILED, + STATUS_QUEUED, STATUS_READY, STATUS_RUNNING, STATUS_UNAVAILABLE, +}; + +/// Short timeout for the `available` probe, so the templates call stays snappy +/// when ComfyUI's host is off. +const PROBE_TIMEOUT: Duration = Duration::from_secs(2); +/// How often the idle check looks at ComfyUI. +const IDLE_CHECK_EVERY: Duration = Duration::from_secs(30); +/// Finished jobs are kept this long. +const JOB_RETENTION_SECS: i64 = 14 * 24 * 3600; +/// Jobs shown per photo. +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_DISABLED: &str = "Enhance isn't configured"; + +/// A library file a request names. +#[derive(Debug, Clone)] +pub struct SourceFile { + pub library_id: i32, + pub rel_path: String, + pub abs: PathBuf, + pub hash: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum EnhanceError { + /// No `url` configured. + Disabled, + /// ComfyUI isn't reachable. + Unavailable, + UnknownTemplate(String), + BadRequest(String), + /// ComfyUI refused the workflow. + Rejected(String), + /// The result is gone (ComfyUI restarted). + Expired, + NotFound, + /// Approved, but ImageApi can't place the file in a library. + Unplaceable(String), + Upstream(String), + Internal(String), +} + +impl EnhanceError { + pub fn message(&self) -> String { + match self { + EnhanceError::Disabled => MSG_DISABLED.into(), + EnhanceError::Unavailable => MSG_UNAVAILABLE.into(), + EnhanceError::UnknownTemplate(t) => format!("unknown template {t:?}"), + EnhanceError::BadRequest(m) + | EnhanceError::Rejected(m) + | EnhanceError::Unplaceable(m) + | EnhanceError::Upstream(m) + | EnhanceError::Internal(m) => m.clone(), + EnhanceError::Expired => MSG_EXPIRED.into(), + EnhanceError::NotFound => "not found".into(), + } + } +} + +fn comfy_err(e: ComfyError) -> EnhanceError { + match e { + ComfyError::Unreachable(_) => EnhanceError::Unavailable, + other => EnhanceError::Upstream(other.to_string()), + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SavedEntry { + pub n: usize, + pub comfy_path: String, + pub library_id: Option, + pub rel_path: Option, +} + +/// What the idle check decided. +#[derive(Debug, Default)] +pub struct IdleState { + initialised: bool, + marker: Option, + seen_at_ms: i64, + freed: bool, +} + +impl IdleState { + /// One observation of a reachable ComfyUI. `marker` identifies the + /// newest history entry (its last event time): when it changes, or the + /// queue isn't empty, or one of our jobs is active, ComfyUI is in use. + /// Returns true once per quiet spell when it has been idle for + /// `idle_ms` and `/free` should be called. + /// + /// The marker's *value* is ComfyUI's clock; only changes in it matter, + /// and "how long ago" is measured on ImageApi's clock, so clock skew + /// between the hosts can't cause an early unload. After ImageApi + /// (re)starts the first observation counts as activity. + pub fn observe( + &mut self, + now_ms: i64, + queue_empty: bool, + active_jobs: usize, + marker: Option, + idle_ms: i64, + ) -> bool { + if !self.initialised || marker != self.marker || !queue_empty || active_jobs > 0 { + self.initialised = true; + self.marker = marker; + self.seen_at_ms = now_ms; + self.freed = false; + return false; + } + if !self.freed && now_ms - self.seen_at_ms >= idle_ms { + self.freed = true; + return true; + } + false + } + + /// ComfyUI went away: the next observation starts a fresh baseline. + pub fn reset(&mut self) { + *self = IdleState::default(); + } + + /// `/free` failed; try again on the next check. + pub fn free_failed(&mut self) { + self.freed = false; + } +} + +pub struct EnhanceService { + pub config: EnhanceConfig, + templates: Vec