diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 754d8a3..c0f018b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -8,6 +8,7 @@ on: branches: [main] push: branches: [main] + workflow_dispatch: # Superseded PR pushes cancel their older run. Every main push gets its own # group (run_id), so no main commit's run is ever cancelled or dropped as pending. @@ -146,6 +147,32 @@ jobs: python -m coverage combine python -m coverage report --fail-under=85 + perf-budget: + # The reference wall-clock measurement (TESTING.md "Performance budgets"). + # Shared-runner timing varies, so it never gates a pull request: it runs on + # each push to main and on demand. A red run on main means fix the code. + name: perf budget + if: github.event_name == 'push' || github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + + - uses: actions/setup-python@v7 + with: + python-version: "3.13" + cache: pip + cache-dependency-path: requirements-dev.txt + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements-dev.txt + + - name: Workspace filter budget + env: + PROMPTMANAGER_PERF_BUDGET: "1" + run: python -m unittest -v tests.test_api_workspace_images.TestWorkspaceImages.test_filter_budget + lint: name: black and flake8 runs-on: ubuntu-latest diff --git a/TESTING.md b/TESTING.md index 0e91696..7d83e54 100644 --- a/TESTING.md +++ b/TESTING.md @@ -125,6 +125,37 @@ not certify every existing feature, arbitrary host versions, GPU/model changes, cross-pack installation or visual design changes. Those still need targeted validation. +## Performance budgets + +A spec performance budget gets two tests: + +- A structural test that always runs, on every OS/Python job and under the + coverage tracer. It asserts counts, never time: SELECTs per request, the + `EXPLAIN QUERY PLAN` of the hot query, cache builds on cold and warm requests, + and calls of the sort key per cursor page. Example: + `test_filter_work_is_bounded` in `tests/test_api_workspace_images.py`. +- A wall-clock test that runs only with `PROMPTMANAGER_PERF_BUDGET=1` and + without a tracer (`sys.gettrace()` or a `sys.monitoring` tool). It skips + otherwise. It warms up once, then measures three rounds and asserts the best + round's p95; the failure message lists every timing. Example: + `test_filter_budget` (Images filters, spec D12: p95 < 150 ms at 50,000 files + and 20,000 links). + +The `perf budget` job in `.github/workflows/test.yml` is the reference +measurement: `ubuntu-latest`, one Python, no coverage, only the wall-clock tests. +It runs on every push to `main` and on `workflow_dispatch`, never on pull +requests, because shared-runner timing must not block a merge. Dispatch it on a +branch to measure before merging. +Shared Windows and macOS runners do not run them; the structural tests cover +portability. Run it locally with +`PROMPTMANAGER_PERF_BUDGET=1 python -m unittest tests.test_api_workspace_images.TestWorkspaceImages.test_filter_budget` +and quote the printed p95 in the PR body. + +A red `perf budget` run on `main` means the code got slower: fix the code (or +change the budget in the spec with a reason) in a follow-up PR. Never rerun it to +green. It is not a required check, and it cannot be one while it skips pull +requests (`.github/branch-protection.json`). + ## Windows portability preflight Before changing paths, SQLite, streams, restore or cleanup: @@ -392,6 +423,141 @@ npm run --prefix tools/vendor test:workflows -- shell-concept.spec.mjs \ --project chromium --project touch --project firefox --project webkit ``` +## Workspace data fields for the concept pages (#249) + +These fields are additive. The legacy `/prompt_manager/images/output` listing (no +`mode=workspace`) and the legacy pages do not change, and existing workspace pages +ignore the new keys. + +Workspace image items (`/prompt_manager/images/output?mode=workspace`, and Library +image galleries) gain three card facts. Library list previews carry none of them +and skip the parse. + +- `linked: {status, positive_id, title, title_source, category, rating}`. One + `get_image_link_summaries` SELECT per 200 paths, matched exactly as the + inspection route matches (exact or normalized `file_path`, no basename + fallback). `status` is `linked`, `ambiguous` (several owners), `none`, or + `unavailable` when the database read failed (logged, no error text). Treat + anything but `linked` as the filename fallback. In pair scope (the image has + exactly one distinct pair), `positive_id` is the pair's positive, and the title + comes from the pair title, then the positive's Library title, then the text's + first line. In positive scope it comes from the Library title, then the first + line. A first line longer than 80 code points is cut to 80 plus `…`, like + `library.mjs` `title()`. `title_source` is `stored`, `text` or null. `category` + and `rating` come from the linked positive (D13), not from the pair. The title + gate is shared with the inspector (`_verified_title`): a `stored` title always + shows. A `text` title shows only when the file's own positive, parsed from its + records by the shared parser, gives the same display title. Otherwise the title + is null on the card and in the inspector header. Videos and unreadable files + have no parsed positive, so their `text` titles are withheld. +- `seed`: a decimal string, or null. Always a string, because `JSON.parse` rounds + numbers past 2**53 and ComfyUI's randomized seeds are often that large. It is + `extract_generation_params` from `utils/generation_params.py` (the inspector's + legacy parser, #251), run on the `prompt`/`workflow` records read with the + dimensions. A PNG is read by the inspector's own reader + (`utils.prompt_inspector.read_inspection_file`): text after IDAT counts, the + shared byte budgets apply, and a file with conflicting duplicate `prompt` or + `workflow` chunks gets no card seed, as the inspector refuses it (the card + keeps its dimensions). Other formats use Pillow's header `info`, as the + inspector's header read does. The records go through the inspection + projection first (`utils/inspection_response.project_graph`), which carries + integers past 2**53 as their exact digits, so the card and the inspector read + the same digits. The parse runs after the file handle closes. A record over + 2 MiB (UTF-8), a parser error, a recursion error, a placeholder, or a + non-integral, boolean or non-finite seed gives null. A card never shows a seed + that differs from the inspector's + (`test_card_seed_equals_inspector_seed_when_both_present`, which `JSON.parse`s + the listing in Node and compares `String(seed)` with the inspector's seed, + including one past 2**53). The card shows no seed for A1111 `parameters` files + or for a BasicGuider custom sampler, where the inspector does. +- Header facts are cached. Dimensions, `seed` and the parsed positive's display + title are read once per `(path, mtime)` of the current gallery scan and kept + beside the scan's per-file values. A monitor invalidation or TTL rescan keeps + the facts of unchanged files and drops the rest, and a roots change drops them + all. A header that could not be read is not stored. Links and favorites are + never cached (`test_warm_page_reuses_card_facts_per_scan`: a warm page reads + and parses no header). Measured cold cost is about 3.4 ms per 50 KB record set + and 25 ms per 400 KB workflow. +- `favorite: bool`. One favorites SELECT per root per page. A read failure logs a + warning and shows false. + +Filters (`mode=workspace`; anything invalid is 400 "Invalid gallery filters"): + +- `q` (at most 500 characters). An entry passes when the case-folded relative path + contains it, or when its linked positive's text, Library title, a tag name, or + its pair title contains it (SQLite `LIKE` with escaping, ASCII case folding + only). Search does not apply the card title gate, and does not search pair tags. +- `category` (exact and case-sensitive, at most 200 characters, stripped; blank + means none) and `tags` (a comma list, blanks dropped, order-free, AND, at most 30 + tags of at most 200 characters) admit linked images only, through the linked + positive. +- Ambiguous files never match a linked filter, and they rank as unlinked. +- `sort` is `newest | oldest | name | name-desc | rating`. `rating` orders by + rating 5 to 1, then linked but unrated, then unlinked and ambiguous; ties break + newest first. A signed cursor covers every filter and the sort, so reusing it + with a change is 400. +- One link read (`get_image_link_matches`, streamed and bounded to the listed + paths) serves the filters, the rating sort and the count. + +`positive_count` is the number of distinct linked positives over the full matching +set, not the page: + +- Present on every first page, and on every page of a filtered (`q`, `category`, + `tags`) or `sort=rating` listing. +- Absent on a plain listing's cursor pages, which read no links. The frontend + keeps the first page's value. +- Null when a plain first page's link read failed. The listing still returns 200, + and consumers then hide "· K positive prompts". A `q` search alone also + survives the failure: it matches filenames only, as before linked search, and + every page's count is null (logged with the traceback). A failure under + `category`, `tags` or the rating sort stays a safe 500. + +Library parents (list, detail and command receipts) gain `run_count` (an int; the +D23 migration estimate for rows from before 3.2.4) and `last_used_at`, a UTC ISO +string. A naive value reads as UTC, and an unreadable or out-of-range value is +null. The edit token ignores both. `variant_total` has two meanings. On the list +response it counts the active variants over all filtered parents (the same +snapshot as `total`, the same on every cursor page). On the detail response it +counts one positive's variants in the chosen archive scope. The Library caption +reads it from the list state. + +Favorites store (D7 = A): `PUT` and `DELETE /prompt_manager/images/favorite` with +`{root_index, path}`. Both are idempotent and answer +`{success, favorite, root_index, path}`. The boundary returns 400 for an unknown +root or a malformed body, and 403 for traversal, an absolute or drive path, NUL, a +non-media name, or a key over the stored bound. PUT needs a listed file (404 +otherwise). DELETE also clears a path that is no longer listed. Send the DELETE +body as `application/json`. + +- The table is `image_favorites(root_key, relative_path, created_at)`. `root_key` + is the root's real path (`os.path.realpath`), case-normalized with `/` + separators. It is stored server-side only and never returned. Folder labels are + not unique across roots, so they are not used as keys. +- `relative_path` is `normcase(path)` with `/`. On Windows a case-only rename + keeps the star. On POSIX the key is the exact spelling, so `A.png` and `a.png` + stay separate (tests swap in `ntpath` through `py.api.images._path_flavor`). + Moving the install or changing a root's real path orphans its rows, which are + harmless and never shown. +- Schema validation (`image_favorites_schema_error`) runs at startup and on + restore. Old backups without the table restore, and backup plus restore keeps + WAL-resident favorites and closes every handle (`tests/test_database_backup.py`). + +Budgets: `test_filter_work_is_bounded` always runs. It asserts that a plain first +page makes 1 link SELECT and a plain cursor page none, that every filter shape +makes exactly 1, the covering-index plan, and that the per-file values are built +once per scan and only for new files after a monitor invalidation. The opt-in +`test_filter_budget` (`PROMPTMANAGER_PERF_BUDGET=1`, the `perf budget` job) times +the D12 p95 < 150 ms budget; see "Performance budgets". + +```sh +python -m unittest tests.test_api_workspace_images tests.test_card_seed \ + tests.test_image_favorites tests.test_library tests.test_api_library \ + tests.test_database_backup -q +``` + +Installing these fields needs a maintainer ComfyUI restart (Python, routes and +schema). Automated tests use disposable databases and output folders. + ## Prompt sending and cross-pack compatibility `tests/js/prompt-bridge-contracts.test.js` runs the Toolkit's shared protocol diff --git a/database/image_favorites.py b/database/image_favorites.py new file mode 100644 index 0000000..4334926 --- /dev/null +++ b/database/image_favorites.py @@ -0,0 +1,139 @@ +"""Persisted image favorites, keyed by gallery root + relative path. + +The root key is the root's real, case-normalized path and the relative key is +case-normalized the same way (see ``py/api/images.py``); both stay +server-side. Folder names alone are not unique across roots, so a favorite +follows the directory itself when the configured roots are reordered. +Optional in old backups, required once a database is current (see +``library_schema.py``). +""" + +import re +import sqlite3 +from typing import Iterable + +TABLE = "image_favorites" +MAX_KEY_LENGTH = 4096 +_LOOKUP_CHUNK = 500 +_DDL = ( + f"CREATE TABLE {TABLE} (root_key TEXT NOT NULL, relative_path TEXT NOT NULL, " + "created_at TEXT NOT NULL, PRIMARY KEY(root_key, relative_path))" +) +_AUTOINDEX = f"sqlite_autoindex_{TABLE}_1" +_COLUMNS = ( + ("root_key", "TEXT", 1, 1), + ("relative_path", "TEXT", 1, 2), + ("created_at", "TEXT", 1, 0), +) + + +def _compact(sql: str) -> str: + return re.sub(r"\s+", "", sql).lower() + + +def image_favorites_schema_error( + conn: sqlite3.Connection, *, allow_missing: bool = True +) -> str | None: + """Reject unsupported favorites schemas and invalid persisted values. + + Reads only. A missing table is acceptable for backups that predate it + unless ``allow_missing`` is false. + """ + objects = conn.execute( + "SELECT type,name,sql FROM sqlite_master WHERE name=? " + "OR (type IN ('index','trigger') AND tbl_name=?)", + (TABLE, TABLE), + ).fetchall() + if not objects: + return None if allow_missing else "Missing image favorites schema" + tables = [sql for kind, _name, sql in objects if kind == "table"] + others = {(kind, name, sql) for kind, name, sql in objects if kind != "table"} + if len(tables) != 1 or others != {("index", _AUTOINDEX, None)}: + return "Image favorites name is reserved by an unsupported schema object" + if _compact(tables[0] or "") != _compact(_DDL): + return "Image favorites has an invalid table definition" + columns = tuple( + (row[1], row[2].upper(), row[3], row[5]) + for row in conn.execute(f"PRAGMA table_info({TABLE})") + ) + if columns != _COLUMNS: + return "Image favorites has invalid columns" + # Bytes bound characters from above (UTF-8 is at most 4 bytes each). + limit = MAX_KEY_LENGTH * 4 + if conn.execute( + f"SELECT 1 FROM {TABLE} WHERE typeof(root_key) != 'text' " + "OR typeof(relative_path) != 'text' OR typeof(created_at) != 'text' " + f"OR length(CAST(root_key AS BLOB)) > {limit} " + f"OR length(CAST(relative_path AS BLOB)) > {limit} " + "OR length(CAST(created_at AS BLOB)) > 64 LIMIT 1" + ).fetchone(): + return "Image favorites contains invalid values" + for root, path in conn.execute( + f"SELECT CAST(root_key AS BLOB), CAST(relative_path AS BLOB) FROM {TABLE} " + f"WHERE length(CAST(root_key AS BLOB)) > {MAX_KEY_LENGTH} " + f"OR length(CAST(relative_path AS BLOB)) > {MAX_KEY_LENGTH}" + ): + for raw in (root, path): + if len(bytes(raw).decode("utf-8", "replace")) > MAX_KEY_LENGTH: + return "Image favorites contains invalid values" + return None + + +def migrate_image_favorites(conn: sqlite3.Connection) -> None: + """Install the table atomically; validate before and after creating it.""" + conn.execute("SAVEPOINT image_favorites_schema") + try: + error = image_favorites_schema_error(conn) + if error: + raise sqlite3.DatabaseError(error) + conn.execute(_DDL.replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS", 1)) + error = image_favorites_schema_error(conn, allow_missing=False) + if error: + raise sqlite3.DatabaseError(error) + conn.execute("RELEASE image_favorites_schema") + except BaseException: + conn.execute("ROLLBACK TO image_favorites_schema") + conn.execute("RELEASE image_favorites_schema") + raise + + +class ImageFavoriteOperations: + """PromptDatabase mixin: set, clear and look up image favorites.""" + + def set_image_favorite(self, root_key: str, relative_path: str) -> bool: + """Mark one file as a favorite; repeating it is a no-op. Returns True.""" + from .models import utc_now_iso # models imports this module's schema + + with self.model.get_connection() as conn: + conn.execute( + "INSERT OR IGNORE INTO image_favorites VALUES (?, ?, ?)", + (root_key, relative_path, utc_now_iso()), + ) + return True + + def clear_image_favorite(self, root_key: str, relative_path: str) -> bool: + """Remove one favorite; True when a row was removed.""" + with self.model.get_connection() as conn: + cursor = conn.execute( + "DELETE FROM image_favorites WHERE root_key=? AND relative_path=?", + (root_key, relative_path), + ) + return cursor.rowcount > 0 + + def get_image_favorites( + self, root_key: str, relative_paths: Iterable[str] + ) -> set[str]: + """The subset of *relative_paths* favorited under *root_key*.""" + wanted = list(dict.fromkeys(relative_paths)) + found: set[str] = set() + conn = self.model.get_connection() + for start in range(0, len(wanted), _LOOKUP_CHUNK): + chunk = wanted[start : start + _LOOKUP_CHUNK] + marks = ",".join("?" * len(chunk)) + rows = conn.execute( # only "?" placeholders are interpolated + "SELECT relative_path FROM image_favorites WHERE root_key=? " + f"AND relative_path IN ({marks})", # nosec B608 + (root_key, *chunk), + ) + found.update(path for (path,) in rows) + return found diff --git a/database/library.py b/database/library.py index 218f783..ce562df 100644 --- a/database/library.py +++ b/database/library.py @@ -1,5 +1,6 @@ """Read-only positive-parent views over existing prompt and pair tables.""" +import datetime import sqlite3 from collections.abc import Iterator from contextlib import contextmanager @@ -13,10 +14,79 @@ _ELIGIBLE = """(EXISTS(SELECT 1 FROM prompt_pairs ep WHERE ep.positive_prompt_id _FROM = "prompts p LEFT JOIN prompt_library_state ls ON ls.prompt_id=p.id" # The characters JavaScript's String.prototype.trim() removes, so a title the # display treats as blank also falls back to the prompt text for name sorting. -_BLANK = ( - "char(9,10,11,12,13,32,160,5760,8192,8193,8194,8195,8196,8197,8198,8199," - "8200,8201,8202,8232,8233,8239,8287,12288,65279)" -) +_TRIM_CODES = (9, 10, 11, 12, 13, 32, 160, 5760, *range(8192, 8203)) +_TRIM_CODES += (8232, 8233, 8239, 8287, 12288, 65279) +_BLANK = f"char({','.join(map(str, _TRIM_CODES))})" +_JS_TRIM = "".join(map(chr, _TRIM_CODES)) +# The Library's display title: a stored title, else the prompt's first line. +DISPLAY_TITLE_LIMIT = 80 + + +def _utf16_units(text: str) -> int: + """JavaScript ``length``: characters beyond the BMP count as two units.""" + return len(text) + sum(1 for ch in text if ord(ch) > 0xFFFF) + + +def _utf16_prefix(text: str, units: int) -> str: + """The longest prefix within ``units`` UTF-16 units, never splitting a pair.""" + used = 0 + for index, ch in enumerate(text): + used += 2 if ord(ch) > 0xFFFF else 1 + if used > units: + return text[:index] + return text + + +def titled(stored_titles: tuple, text: object) -> tuple[str | None, str | None]: + """Return ``(title, source)``: a stored title, else the text's first line. + + ``source`` is ``"stored"`` for the first non-blank stored title (blank as + in JavaScript's trim, like the Library view), ``"text"`` for the first line, + and None when both are empty. The first line is trimmed and, past + DISPLAY_TITLE_LIMIT UTF-16 units, cut with an ellipsis like the Library's + ``libraryTitle()`` (never inside a surrogate pair). + """ + for title in stored_titles: + if isinstance(title, str) and title.strip(_JS_TRIM): + return title, "stored" + if not isinstance(text, str): + return None, None + first = text.split("\n")[0].strip(_JS_TRIM) + if _utf16_units(first) > DISPLAY_TITLE_LIMIT: + first = _utf16_prefix(first, DISPLAY_TITLE_LIMIT) + "\u2026" + return (first, "text") if first else (None, None) + + +def display_title(stored_titles: tuple, text: object) -> str | None: + """The Library display title of ``titled``, without its source.""" + return titled(stored_titles, text)[0] + + +def utc_iso(value: object) -> str | None: + """A stored timestamp as UTC ISO text, or None when absent or unreadable. + + Usage writes are already UTC ISO; a naive value (SQLite's CURRENT_TIMESTAMP + form) is read as UTC, and an offset is converted to UTC. One bad row must + not fail a whole Library read, so a value past the datetime range after + conversion (OverflowError) is unreadable too. Python 3.11+ fromisoformat + accepts more ISO forms than 3.10 (e.g. a 2-digit fraction); every form the + app writes (utc_now_iso, the 3.2.4 backfill, CURRENT_TIMESTAMP, Z) parses + the same on both, so only hand-edited values can differ by version. + """ + if not isinstance(value, str): + return None + text = value.strip() + if text.endswith(("Z", "z")): + text = text[:-1] + "+00:00" + try: + stamp = datetime.datetime.fromisoformat(text) + if stamp.tzinfo is None: + stamp = stamp.replace(tzinfo=datetime.timezone.utc) + return stamp.astimezone(datetime.timezone.utc).isoformat() + except (ValueError, OverflowError): + return None + + _NAME = ( f"LOWER(CASE WHEN TRIM(COALESCE(ls.title,''),{_BLANK})!='' " "THEN ls.title ELSE p.text END)" @@ -145,6 +215,7 @@ class LibraryOperations: ) -> list[dict]: rows = conn.execute( "SELECT p.id,p.text,ls.title,p.category,p.rating,p.notes,p.created_at," + "p.run_count,p.last_used_at," + (f"{sort_expression} AS _sort," if sort_expression else "") + "(SELECT COUNT(*) FROM prompt_pairs v WHERE v.positive_prompt_id=p.id " "AND v.archived=0) AS variant_count," @@ -165,6 +236,9 @@ class LibraryOperations: ): by_id[row["prompt_id"]]["tags"].append(row["name"]) for row in result: + # Usage is shown, never compared: the token hashes owner fields only. + row["run_count"] = row["run_count"] if type(row["run_count"]) is int else 0 + row["last_used_at"] = utc_iso(row["last_used_at"]) row["edit_token"] = library_edit_token(epoch, row) return result @@ -196,9 +270,15 @@ class LibraryOperations: ceiling = after["ceiling"] where += " AND p.id<=?" params.append(ceiling) - total = conn.execute( - "SELECT COUNT(*) FROM " + _FROM + " WHERE " + where, params - ).fetchone()[0] + # One SELECT counts the filtered parents and their active variants. + total, variant_total = conn.execute( + "SELECT COUNT(*),COALESCE(SUM((SELECT COUNT(*) FROM prompt_pairs v " + "WHERE v.positive_prompt_id=p.id AND v.archived=0)),0) FROM " + + _FROM + + " WHERE " + + where, + params, + ).fetchone() expression, direction = _ORDERS[sort] if after is not None: operator = ">" if direction == "ASC" else "<" @@ -230,6 +310,7 @@ class LibraryOperations: "library_epoch": epoch, "items": rows, "total": total, + "variant_total": variant_total, "has_more": has_more, "next_key": next_key, } diff --git a/database/models.py b/database/models.py index 448fd3a..85f1f23 100644 --- a/database/models.py +++ b/database/models.py @@ -16,6 +16,7 @@ from pathlib import Path from typing import Optional, Set, Tuple from .library_schema import library_schema_error, migrate_library_state +from .image_favorites import image_favorites_schema_error, migrate_image_favorites from .execution_history_schema import ( history_schema_error, migrate_execution_history, @@ -298,6 +299,7 @@ class PromptModel: self._migrate_prompt_pairs(conn) migrate_execution_history(conn) migrate_library_state(conn) + migrate_image_favorites(conn) def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None: conn.execute("SAVEPOINT prompt_pair_schema") @@ -1016,6 +1018,9 @@ class PromptModel: Args: path: File to inspect + require_library_state: Also require the tables every current + database has (Library state and image favorites), as a + migrated restore stage must; old backups may lack them. Returns: (ok, reason): reason is "ok" when the file is acceptable, otherwise @@ -1070,6 +1075,11 @@ class PromptModel: ) if library_error: return False, library_error + favorites_error = image_favorites_schema_error( + conn, allow_missing=not require_library_state + ) + if favorites_error: + return False, favorites_error return True, "ok" except sqlite3.Error as e: return False, f"Invalid SQLite database: {e}" diff --git a/database/operations.py b/database/operations.py index e210199..3fb9f18 100644 --- a/database/operations.py +++ b/database/operations.py @@ -6,11 +6,12 @@ import sqlite3 import json import datetime import os -from typing import Optional, List, Dict, Any +from typing import Optional, List, Dict, Any, Set, Tuple from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations -from .library import LibraryOperations +from .library import LibraryOperations, titled +from .image_favorites import ImageFavoriteOperations from .library_commands import LibraryCommandOperations from .vocabulary import VocabularyOperations from .execution_history import ExecutionHistoryOperations @@ -179,6 +180,119 @@ AND p.id NOT IN (SELECT prompt_id FROM live_hint_prompts) """ +# One SELECT reads the link summaries of at most this many paths: three bound +# parameters per path stays under SQLite's historical 999-variable limit. +LINK_SUMMARY_BATCH = 200 +# A pair link counts only when an image has exactly one distinct pair, as in +# get_image_prompt_info; its positive, title and metadata then describe the card. +_LINK_SUMMARY_SQL = """SELECT gi.image_path, gi.file_path, gi.prompt_id, + p.text, p.category, p.rating, ls.title, + pp.positive_prompt_id, pp.title, pt.text, pt.category, pt.rating, pls.title +FROM generated_images gi +JOIN prompts p ON gi.prompt_id = p.id +LEFT JOIN prompt_library_state ls ON ls.prompt_id = p.id +LEFT JOIN prompt_pairs pp ON pp.id = ( + SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END + FROM image_prompt_pairs ip WHERE ip.image_id = gi.id) +LEFT JOIN prompts pt ON pt.id = pp.positive_prompt_id +LEFT JOIN prompt_library_state pls ON pls.prompt_id = pp.positive_prompt_id +WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}""" +# One image row describes a card only when no other row claims the same file +# (as in get_image_link_summaries, several rows are ambiguous). The pair scope +# and the linked positive follow _LINK_SUMMARY_SQL; category, rating, tags and +# the searchable text and titles come from that positive (spec D12/D13). Every +# writer since the 3.2.4 migration stores file_path, so it is the path key. +# The prompt-level conditions are uncorrelated IN sets SQLite builds once, and +# the CASE evaluates the per-row ambiguity probe only for admitted rows. +_LINK_FILTER_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id +FROM generated_images gi +LEFT JOIN prompt_pairs pp ON pp.id = ( + SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END + FROM image_prompt_pairs ip WHERE ip.image_id = gi.id) +JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id) +WHERE gi.file_path IS NOT NULL AND CASE WHEN {admit} THEN NOT EXISTS ( + SELECT 1 FROM generated_images other + WHERE other.file_path = gi.file_path AND other.id <> gi.id) ELSE 0 END""" +# The positive's text, its Library title or one of its tag names, or the pair +# title, contains :pattern (LIKE, so ASCII case is ignored as in Library search). +_LINK_FILTER_HIT = """(p.id IN ( + SELECT hp.id FROM prompts hp + LEFT JOIN prompt_library_state hls ON hls.prompt_id = hp.id + WHERE hp.text LIKE :pattern ESCAPE '\\' OR hls.title LIKE :pattern ESCAPE '\\' + OR hp.id IN (SELECT ht.prompt_id FROM prompt_tags ht JOIN tags hn + ON hn.id = ht.tag_id WHERE hn.name LIKE :pattern ESCAPE '\\')) + OR pp.title LIKE :pattern ESCAPE '\\')""" +_LINK_FILTER_TAG = ( + "p.id IN (SELECT ft.prompt_id FROM prompt_tags ft JOIN tags fn " + "ON fn.id = ft.tag_id WHERE fn.name = :tag{index})" +) + + +def link_filter_query( + q: str, category: Optional[str], tags: List[str], *, every_link: bool = False +) -> Tuple[str, Dict[str, Any]]: + """The workspace filter SELECT and its named parameters (all bound). + + Only code constants and enumerate indices are formatted into the SQL; see + PromptDatabase.get_image_link_matches for the semantics. + """ + params: Dict[str, Any] = {} + hit = "1" + if q: + hit = _LINK_FILTER_HIT + params["pattern"] = f"%{escape_like(q)}%" + admitted = [] + if category is not None: + admitted.append("p.category = :category") + params["category"] = category + for index, tag in enumerate(tags): + admitted.append(_LINK_FILTER_TAG.format(index=index)) + params[f"tag{index}"] = tag + if q and not admitted and not every_link: + admitted.append(hit) + sql = _LINK_FILTER_SQL.format(hit=hit, admit=" AND ".join(admitted) or "1") + return sql, params + + +_ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz") + + +def _like_suffix(value: object, suffix: str) -> bool: + """Python form of SQLite's ``value LIKE '%' || suffix``: ASCII-only case.""" + return isinstance(value, str) and value.translate(_ASCII_FOLD).endswith( + suffix.translate(_ASCII_FOLD) + ) + + +def _empty_link_summary(status: str) -> Dict[str, Any]: + return { + "status": status, + "positive_id": None, + "title": None, + "title_source": None, + "category": None, + "rating": None, + } + + +def _link_summary(row: tuple) -> Dict[str, Any]: + if row[7] is not None: + positive, text, category, rating = row[7], row[9], row[10], row[11] + titles = (row[8], row[12]) + else: + positive, text, category, rating = row[2], row[3], row[4], row[5] + titles = (row[6],) + title, source = titled(titles, text) + return { + "status": "linked", + "positive_id": positive, + "title": title, + "title_source": source, + "category": category, + "rating": rating, + } + + class PromptDatabase( PairOperations, LibraryOperations, @@ -188,6 +302,7 @@ class PromptDatabase( ExecutionContextOperations, ExecutionOutputOperations, HistoryReadOperations, + ImageFavoriteOperations, ): """Database operations class for managing prompts.""" @@ -2160,6 +2275,100 @@ class PromptDatabase( "parameters": row[8] if len(row) > 8 else None, } + def get_image_link_summaries( + self, paths: List[str], allow_filename_fallback: bool = True + ) -> Dict[str, Dict[str, Any]]: + """Summarize each path's linked saved positive for an image card. + + Matching equals get_image_prompt_info: exact image_path or normalized + file_path rows win; otherwise, with the fallback, rows whose image_path + ends with the basename. One match is linked, several are ambiguous. + One SELECT reads up to LINK_SUMMARY_BATCH paths. + """ + unique = list(dict.fromkeys(paths)) + summaries: Dict[str, Dict[str, Any]] = {} + for start in range(0, len(unique), LINK_SUMMARY_BATCH): + batch = unique[start : start + LINK_SUMMARY_BATCH] + summaries.update(self._link_summary_batch(batch, allow_filename_fallback)) + return summaries + + def _link_summary_batch( + self, paths: List[str], allow_filename_fallback: bool + ) -> Dict[str, Dict[str, Any]]: + normalized = {path: normalize_image_path(path) for path in paths} + basenames = {path: os.path.basename(path) for path in paths} + marks = ",".join("?" for _ in paths) + params: List[Any] = [*paths, *normalized.values()] + fallback = "" + if allow_filename_fallback: + fallback = "".join(" OR gi.image_path LIKE ? ESCAPE '\\'" for _ in paths) + params.extend(f"%{escape_like(basenames[path])}" for path in paths) + sql = _LINK_SUMMARY_SQL.format(marks=marks, fallback=fallback) + with self.model.get_connection() as conn: + rows = conn.execute(sql, params).fetchall() + summaries = {} + for path in paths: + exact = [ + row for row in rows if row[0] == path or row[1] == normalized[path] + ] + matches = exact or ( + [row for row in rows if _like_suffix(row[0], basenames[path])] + if allow_filename_fallback + else [] + ) + if len(matches) == 1: + summaries[path] = _link_summary(matches[0]) + else: + summaries[path] = _empty_link_summary( + "ambiguous" if matches else "none" + ) + return summaries + + def get_image_link_matches( + self, + q: str, + category: Optional[str], + tags: List[str], + *, + every_link: bool = False, + wanted: Optional[Set[str]] = None, + ) -> Dict[str, Tuple[bool, Optional[int], int]]: + """Map normalized paths of unambiguously linked images to + (hit, rating, positive id). + + Rows are limited to images whose linked positive has *category* and + every tag in *tags* (exact names). ``hit`` says whether *q* occurs in + the positive's text, Library title or a tag name, or in the pair title. + Without category or tags only hits are returned, unless *every_link* + asks for every linked image (the rating sort and the positive count). + The positive id is the card's linked positive. One SELECT whose rows + are streamed; with *wanted* (the listing's normalized paths) only those + paths are kept, so memory is bounded by the listing, not the database. + """ + sql, params = link_filter_query(q, category, tags, every_link=every_link) + matches: Dict[str, Tuple[bool, Optional[int], int]] = {} + with self.model.get_connection() as conn: + for path, rating, matched, positive in conn.execute(sql, params): + if wanted is None or path in wanted: + matches[path] = (bool(matched), rating, positive) + return matches + + def get_image_paths_matching( + self, q: str, category: Optional[str], tags: List[str] + ) -> Set[str]: + """Normalized paths whose linked positive matches *q*, *category* and *tags*. + + See get_image_link_matches; an empty *q* matches every linked image + the category and tags admit. + """ + return { + path + for path, (hit, _rating, _positive) in self.get_image_link_matches( + q, category, tags + ).items() + if hit or not q + } + def get_prompt_id_for_image(self, image_path: str) -> Optional[int]: """Return the prompt_id linked to an image, or None.""" with self.model.get_connection() as conn: diff --git a/py/api/__init__.py b/py/api/__init__.py index ed24041..2e15b90 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -283,6 +283,16 @@ class PromptManagerAPI( self._html_cache = {} # Cached HTML file contents keyed by path self._gallery_cache = {} # dict: path_str -> (files, timestamp) self._gallery_cache_ttl = 30 # Cache TTL in seconds + # Workspace listing values and sort orders derived from a cached scan + # (py/api/images.py _workspace_listed/_workspace_orders). + self._workspace_listed_cache = {} + self._workspace_order_cache = {} + # Card facts (dimensions, seed, positive title) per root and + # (path, mtime), kept and pruned with the per-file values + # (py/api/images.py _card_facts). + self._workspace_card_cache = {} + # Favorite keys of each root's cached scan (images._favorite_listing). + self._favorite_listing_cache = {} # Single-flight locks for long jobs, keyed by job family; see # _acquire_long_job(). Created lazily so no lock binds to an event # loop before the server's loop exists. @@ -348,9 +358,22 @@ class PromptManagerAPI( def invalidate_gallery_cache(self): """Invalidate the cached gallery file listing. - Called by the image monitor when new files are detected. + Called by the image monitor when new files are detected. The workspace + keeps its per-file values and card facts, so the next request rebuilds + only new or changed files (see ImageRoutesMixin._workspace_listed and + _card_facts); its sort orders + are dropped, as the rescan replaces them anyway, and so is the + favorites lookup, which holds the old scan list. """ self._gallery_cache = {} + self._workspace_order_cache = {} + self._favorite_listing_cache = {} + + def reset_gallery_roots_cache(self): + """Drop every cached listing and derived workspace value (roots changed).""" + self.invalidate_gallery_cache() + self._workspace_listed_cache = {} + self._workspace_card_cache = {} def add_routes(self, routes): """Register all API routes with the ComfyUI server. diff --git a/py/api/admin.py b/py/api/admin.py index f5f3b0c..4360e00 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -897,7 +897,7 @@ class AdminRoutesMixin: return False GalleryConfig.MONITORING_DIRECTORIES = roots self._cached_output_dir = None - self._gallery_cache = {} + self.reset_gallery_roots_cache() return True def _persist_settings(self): diff --git a/py/api/images.py b/py/api/images.py index acf8dfb..e46e49f 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -11,8 +11,10 @@ import re import threading import time as _time import urllib.parse +from bisect import bisect_left, bisect_right from collections.abc import Mapping from pathlib import Path +from typing import NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError @@ -29,14 +31,123 @@ from .inspection_service import ( ) try: + from ...database.image_favorites import MAX_KEY_LENGTH + from ...database.library import display_title + from ...database.models import normalize_image_path + from ...utils.card_seed import card_seed, header_params, parsed_positive from ...utils.parallel import map_parallel from ...utils.inspection_response import encode_inspection_response + from ...utils.prompt_inspector import MetadataReadError, read_inspection_file except ImportError: + from database.image_favorites import MAX_KEY_LENGTH + from database.library import display_title + from database.models import normalize_image_path + from utils.card_seed import card_seed, header_params, parsed_positive from utils.parallel import map_parallel from utils.inspection_response import encode_inspection_response + from utils.prompt_inspector import MetadataReadError, read_inspection_file IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS +# The header records the card parses (see utils.card_seed.header_params). +_CARD_RECORDS = ("prompt", "workflow") +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + + +def _card_header(path: Path) -> tuple: + """``(width, height, records)`` of one image card, every handle closed. + + A PNG's records come from the inspector's own reader + (``read_inspection_file``: text after IDAT, the shared byte budgets, and + the refusal of conflicting duplicate prompt or workflow chunks), so the + card reads exactly what the inspector reads. When that reader refuses the + file, the card keeps its dimensions and gets no records. Other formats + keep Pillow's header ``info``, as the inspector's header read does. + Raises OSError or ValueError for a file that cannot be read at all. + """ + with path.open("rb") as stream: + if stream.read(len(_PNG_SIGNATURE)) != _PNG_SIGNATURE: + stream.seek(0) + with Image.open(stream) as source: + width, height = source.size + return width, height, {k: source.info.get(k) for k in _CARD_RECORDS} + try: + inspection = read_inspection_file(path) + except MetadataReadError: + with path.open("rb") as stream, Image.open(stream) as source: + width, height = source.size + return width, height, None + records = {k: inspection.metadata.get(k) for k in _CARD_RECORDS} + return inspection.width, inspection.height, records + + +class _CardFacts(NamedTuple): + """One image file's card facts, derived from its header alone. + + Cached per gallery scan by ``(path, mtime)`` (see ``_card_facts``); links + and favorites are read fresh on every page. + """ + + width: int | None + height: int | None + seed: str | None + positive_title: str | None # _positive_title of the parsed positive + + +_NO_CARD_FACTS = _CardFacts(None, None, None, None) + + +def _unavailable_link() -> dict: + """A card link summary when the database could not be read.""" + return { + "status": "unavailable", + "positive_id": None, + "title": None, + "title_source": None, + "category": None, + "rating": None, + } + + +def _positive_title(positive: str | None) -> str | None: + """The display title of a file's parsed positive (None without one).""" + return None if positive is None else display_title((), positive) + + +def _verified_title(summary: dict | None, positive_title: str | None) -> str | None: + """A linked file's Library title, withholding unconfirmed saved prompt text. + + The one rule for the card and the inspector. A stored title is user data and + always shows. A title taken from the saved positive's first line shows only + when ``positive_title`` (``_positive_title`` of the file's own positive as + the shared parser reads it from the file's records) is the same title: a + saved association never replaces file evidence. + """ + if not isinstance(summary, dict) or summary.get("status") != "linked": + return None + title = summary.get("title") + if summary.get("title_source") != "text": + return title + confirmed = positive_title is not None and positive_title == title + return title if confirmed else None + + +def _card_link(summary: dict, positive_title: str | None) -> dict: + """The card's link summary with an unconfirmed first-line title withheld.""" + if summary.get("title_source") != "text": + return summary + return {**summary, "title": _verified_title(summary, positive_title)} + + +def _inspected_positive(observed: dict) -> str | None: + """The positive the shared parser reads from the inspected file's records. + + Videos give None, as on the card, which reads no video metadata. + """ + inspection = observed.get("inspection") + if inspection is None or observed["file"]["media_type"] == "video": + return None + return parsed_positive(header_params(inspection.metadata)) def _public_path(path): @@ -238,14 +349,66 @@ def _workspace_root(query: dict, roots: list[Path]) -> int | None: return index +WORKSPACE_ORDERS = ("newest", "oldest", "name", "name-desc", "rating") +# Library-compatible bounds on the linked-positive filters (database/library.py). +MAX_FILTER_CATEGORY = 200 +MAX_FILTER_TAGS = 30 +MAX_FILTER_TAG = 200 + + +class _Listed(NamedTuple): + """One listed file plus the derived values every filter and sort reads. + + Built once per gallery scan (see ``_workspace_orders``); the first four + fields are the ``(path, mtime, root, index)`` entry the page formatter takes. + """ + + path: Path + mtime: float + root: Path + index: int + folded: str # relative POSIX path, case-folded (filename search) + folder: str # str(relative.parent), native separators (subfolder filter) + name: str # case-folded basename (name sorts) + tie: str # sha256 of the relative POSIX path (bounded tie-break) + video: bool + normalized: str # normalize_image_path, compared with linked rows + + +def _listed(path: Path, mtime: float, root: Path, index: int) -> _Listed: + text, prefix = str(path), os.path.join(str(root), "") + relative = ( + text[len(prefix) :] if text.startswith(prefix) else str(path.relative_to(root)) + ) + posix = relative if os.sep == "/" else relative.replace(os.sep, "/") + name = relative.rpartition(os.sep)[2] + return _Listed( + path, + mtime, + root, + index, + posix.casefold(), + os.path.dirname(relative) or ".", + name.casefold(), + hashlib.sha256(posix.encode("utf-8")).hexdigest(), + os.path.splitext(name)[1].lower() in VIDEO_EXTENSIONS, + normalize_image_path(text), + ) + + def _workspace_key(entry: tuple, order: str) -> tuple: - path, mtime, root, index = entry - relative = path.relative_to(root).as_posix() - tie = hashlib.sha256(relative.encode("utf-8")).hexdigest() - name = path.name.casefold() + """The keyset sort key of *entry* (a _Listed or a 4-tuple entry).""" + if not isinstance(entry, _Listed): + entry = _listed(*entry) if order in ("name", "name-desc"): - return (name, tie, index, mtime) - return ((-mtime if order == "newest" else mtime), name, tie, index) + return (entry.name, entry.tie, entry.index, entry.mtime) + mtime = -entry.mtime if order in ("newest", "rating") else entry.mtime + return (mtime, entry.name, entry.tie, entry.index) + + +def _rated_key(entry: _Listed, rank: int) -> tuple: + """Highest rated: rating 5..1, linked unrated (0), unlinked (-1); then newest.""" + return (-rank, *_workspace_key(entry, "rating")) def _workspace_cursor(key: tuple, scope: str) -> str: @@ -276,6 +439,8 @@ def _workspace_cursor_key(token: str, scope: str, order: str) -> tuple: if order in ("name", "name-desc") else ((float, int), str, str, int) ) + if order == "rating": + types = (int, *types) if len(key) != len(types) or not all( type(value) in (kind if isinstance(kind, tuple) else (kind,)) for value, kind in zip(key, types) @@ -328,6 +493,57 @@ async def _json_body_or_error(request): return None, _json_error("Invalid JSON body", 400) +MAX_FAVORITE_PATH = 4096 + + +def _favorite_target(data: Mapping, roots: list[Path]): + """``((index, relative), None)`` for a favorite body, else ``(None, error)``. + + Rejects unknown roots, traversal, absolute or drive paths and non-media + names before any filesystem or database access. + """ + index, relative = data.get("root_index"), data.get("path") + if type(index) is not int or not 0 <= index < len(roots): + return None, bad_request("Invalid gallery root") + if ( + not isinstance(relative, str) + or not relative + or len(relative) > MAX_FAVORITE_PATH + ): + return None, bad_request("path must be a relative media path") + if ( + _has_traversal(relative) + or "\0" in relative + or re.match(r"^[a-zA-Z]:", relative) + ): + return None, _forbidden() + if Path(relative).suffix.lower() not in MEDIA_EXTENSIONS: + return None, _forbidden("Only media files can be favorites") + return (index, relative), None + + +# The path module whose case rules key favorites (tests swap in ntpath). +_path_flavor = os.path + + +def _favorite_relative_key(relative: str) -> str: + """Case-normalized POSIX key of a root-relative path. + + Lowercased with "/" separators where the platform folds case (Windows), + identity on POSIX, so a case-only rename keeps its favorite on Windows + and distinct Linux files stay distinct. + """ + return _path_flavor.normcase(relative).replace(_path_flavor.sep, "/") + + +def _favorite_root_key(root) -> str: + """Server-side identity of a gallery root: its real, case-normalized path. + + Never sent in a response; folder names alone are not unique across roots. + """ + return _favorite_relative_key(os.path.realpath(str(root))) + + def _iter_media_files(output_path, max_files): """Yield media files under *output_path* with bounded effort. @@ -407,6 +623,87 @@ def _output_image_entry( return entry +def _workspace_filters(query: Mapping) -> dict: + """Validate the workspace listing's filter and sort parameters. + + ``q`` (at most 500 characters) matches the relative path, case-folded, and + the linked positive's text, title and tag names. ``category`` is exact, + ``tags`` is a comma list that must all be present, ``sort=rating`` ranks by + the linked positive. Raises ValueError for any unknown or oversized value. + """ + text = query.get("q", "").strip() + folded = text.casefold() + media = query.get("media", "all") + order = query.get("sort", "newest") + folder = query.get("subfolder", "").strip().replace("\\", "/") + category = query.get("category", "").strip() or None + tags = sorted( + {tag.strip() for tag in query.get("tags", "").split(",") if tag.strip()} + ) + if ( + len(text) > 500 + or len(folded) > 500 + or media not in ("all", "image", "video") + or order not in WORKSPACE_ORDERS + or (category is not None and len(category) > MAX_FILTER_CATEGORY) + or len(tags) > MAX_FILTER_TAGS + or any(len(tag) > MAX_FILTER_TAG for tag in tags) + ): + raise ValueError("Invalid gallery filters") + if folder and (_has_traversal(folder) or len(folder) > 1024): + raise ValueError("Invalid gallery folder") + return { + "text": text, + "query": folded, + "media": media, + "order": order, + "folder": folder, + "category": category, + "tags": tags, + } + + +# A plain listing's cursor page leaves positive_count out (see +# _select_workspace_page); None instead means the count could not be read. +_COUNT_OMITTED = object() + + +def _filtered_page( + entries: list, start: int, size: int, filters: dict, links: dict | None +) -> tuple[int, list, int | None]: + """Count the *entries* the filters admit and take *size* from *start* on. + + *links* maps normalized paths to (hit, rating, positive id) for every + unambiguously linked image (see PromptDatabase.get_image_link_matches); a + category or tag filter admits only those, and ``q`` admits a filename match + or a linked-positive hit. Returns the total, the page and the number of + distinct linked positives among all admitted entries (None without links). + """ + query, media, folder = filters["query"], filters["media"], filters["folder"] + members_only = filters["category"] is not None or bool(filters["tags"]) + native_folder = folder.replace("/", os.sep) + nested = native_folder + os.sep + total, page, positives = 0, [], set() + for position, entry in enumerate(entries): + if folder and not ( + entry.folder == native_folder or entry.folder.startswith(nested) + ): + continue + if media != "all" and (media == "video") != entry.video: + continue + match = None if links is None else links.get(entry.normalized) + if members_only and match is None: + continue + if query and query not in entry.folded and not (match and match[0]): + continue + total += 1 + if match is not None: + positives.add(match[2]) + if position >= start and len(page) < size: + page.append(entry) + return total, page, None if links is None else len(positives) + + def _in_subfolder(rel_dir, subfolder): return rel_dir == subfolder or rel_dir.startswith(subfolder + os.sep) @@ -523,6 +820,14 @@ class ImageRoutesMixin: async def get_gallery_subfolders_route(request): return await self.get_gallery_subfolders(request) + @routes.put("/prompt_manager/images/favorite") + async def set_image_favorite_route(request): + return await self.update_image_favorite(request, favorite=True) + + @routes.delete("/prompt_manager/images/favorite") + async def clear_image_favorite_route(request): + return await self.update_image_favorite(request, favorite=False) + # POST: drops the server-side listing cache and rewalks the output roots. @routes.post("/prompt_manager/gallery/rescan") async def rescan_gallery_route(request): @@ -756,32 +1061,26 @@ class ImageRoutesMixin: return bad_request("limit and offset must be integers") try: root_index = _workspace_root(request.query, output_dirs) - query = request.query.get("q", "").strip().casefold() - media = request.query.get("media", "all") - order = request.query.get("sort", "newest") - folder = request.query.get("subfolder", "").strip().replace("\\", "/") - if ( - len(query) > 500 - or media not in ("all", "image", "video") - or order not in ("newest", "oldest", "name", "name-desc") - ): - raise ValueError("Invalid gallery filters") - if folder and (_has_traversal(folder) or len(folder) > 1024): - raise ValueError("Invalid gallery folder") + filters = _workspace_filters(request.query) scope = hashlib.sha256( json.dumps( [ - query, - media, - order, - folder, + filters["query"], + filters["text"], + filters["media"], + filters["order"], + filters["folder"], root_index, [str(root) for root in output_dirs], + filters["category"], + filters["tags"], ] ).encode() ).hexdigest() token = request.query.get("cursor") - after = _workspace_cursor_key(token, scope, order) if token else None + after = ( + _workspace_cursor_key(token, scope, filters["order"]) if token else None + ) except ValueError as exc: return bad_request(str(exc)) @@ -790,75 +1089,193 @@ class ImageRoutesMixin: for index, root in enumerate(output_dirs) if root_index is None or index == root_index ] - page, total, has_more, next_cursor = await self._run_in_executor( - self._select_workspace_page, - listings, - folder, - query, - media, - order, - after, - limit, - scope, + page, total, has_more, next_cursor, positive_count = ( + await self._run_in_executor( + self._select_workspace_page, listings, filters, after, limit, scope + ) ) images = await self._run_in_executor(self._format_workspace_page, page) - return web.json_response( - { - "success": True, - "images": images, - "total": total, - "limit": limit, - "has_more": has_more, - "next_cursor": next_cursor, - } + result = { + "success": True, + "images": images, + "total": total, + "positive_count": positive_count, + "limit": limit, + "has_more": has_more, + "next_cursor": next_cursor, + } + if positive_count is _COUNT_OMITTED: + del result["positive_count"] + return web.json_response(result) + + def _workspace_listed(self, index: int, root: Path, files: list) -> list: + """The _Listed values of one root's scan, rebuilt only where it changed. + + The gallery cache hands out the same list object until it rescans. + A rescan reuses the value of every unchanged ``(path, mtime)`` and, + when nothing changed, the previous list itself, so the cached sort + orders stay valid. Only a gallery-roots change drops these values. + """ + key = (index, str(root)) + cached = self._workspace_listed_cache.get(key) + if cached is not None and cached[0] is files: + return cached[1] + previous = cached[2] if cached is not None else {} + listed = [] + for path, mtime in files: + entry = previous.get((path, mtime)) + listed.append(entry or _listed(path, mtime, root, index)) + if ( + cached is not None + and len(listed) == len(cached[1]) + and all(new is old for new, old in zip(listed, cached[1])) + ): + listed, lookup = cached[1], previous + else: + lookup = {(entry.path, entry.mtime): entry for entry in listed} + # Card facts follow the per-file values: keep unchanged files only. + cards = self._workspace_card_cache.get(key) + if cards: + self._workspace_card_cache[key] = { + file: facts for file, facts in cards.items() if file in lookup + } + self._workspace_listed_cache[key] = (files, listed, lookup) + return listed + + def _workspace_orders(self, listings: list[tuple]) -> dict: + """The selected roots' files in newest order, plus lazily added orders.""" + parts = tuple( + self._workspace_listed(index, root, files) + for index, root, files in listings ) + key = tuple((index, str(root)) for index, root, _files in listings) + cached = self._workspace_order_cache.get(key) + if cached is not None and all( + old is new for old, new in zip(cached[0], parts, strict=True) + ): + return cached[1] + newest = sorted( + (entry for part in parts for entry in part), + key=lambda entry: _workspace_key(entry, "newest"), + ) + orders = {"newest": newest} + self._workspace_order_cache[key] = (parts, orders) + return orders + + def _workspace_wanted(self, listings: list[tuple]) -> frozenset: + """The selected listing's normalized paths (cached with its orders).""" + orders = self._workspace_orders(listings) + if "wanted" not in orders: + orders["wanted"] = frozenset(entry.normalized for entry in orders["newest"]) + return orders["wanted"] + + def _workspace_order(self, listings: list[tuple], order: str) -> list: + """The listing sorted by a filesystem order (built lazily, cached).""" + orders = self._workspace_orders(listings) + if order not in orders: + if order == "name-desc": + orders[order] = self._workspace_order(listings, "name")[::-1] + else: + orders[order] = sorted( + orders["newest"], key=lambda entry: _workspace_key(entry, order) + ) + return orders[order] - @staticmethod def _select_workspace_page( + self, listings: list[tuple], - folder: str, - query: str, - media: str, - order: str, + filters: dict, after: tuple | None, limit: int, scope: str, ) -> tuple: - entries = [] - for index, root, files in listings: - for path, mtime in files: - relative = path.relative_to(root) - if folder and not _in_subfolder( - str(relative.parent), folder.replace("/", os.sep) - ): - continue - if query and query not in relative.as_posix().casefold(): - continue - kind = "video" if path.suffix.lower() in VIDEO_EXTENSIONS else "image" - if media != "all" and media != kind: - continue - entries.append((path, mtime, root, index)) - total = len(entries) - reverse = order == "name-desc" - entries.sort(key=lambda entry: _workspace_key(entry, order), reverse=reverse) + """One keyset page of the filtered listing (blocking; run in executor). + + One SELECT (get_image_link_matches) reads every linked positive of the + listing: the filters, the rating sort and the positive count share it. + A plain listing (no ``q``, category, tags or rating sort) reads it only + for its first page; its cursor pages read no links and the returned + count is _COUNT_OMITTED, so the response leaves ``positive_count`` out + and the frontend keeps the first page's value. On that plain first + page a failed read leaves the listing usable and the count unknown + (None); so does a failed read under ``q`` alone, whose search then + matches filenames only (as before the linked search existed). + """ + order, category, tags = filters["order"], filters["category"], filters["tags"] + needed = ( + bool(filters["query"]) + or category is not None + or bool(tags) + or order == "rating" + ) + links = None + if needed or after is None: + try: + links = self.db.get_image_link_matches( + filters["text"], + category, + tags, + every_link=True, + wanted=self._workspace_wanted(listings), + ) + except Exception: + # Category, tags and the rating sort cannot be served without + # the links. A plain listing and a filename search can: the + # search then matches filenames only, and the count is None. + if category is not None or tags or order == "rating": + raise + self.logger.warning( + "Workspace prompt links could not be read; %s", + ( + "searching filenames only" + if filters["query"] + else "positive count unknown" + ), + exc_info=True, + ) + if order == "rating": + + def rank(entry): + match = links.get(entry.normalized) + return -1 if match is None else (match[1] or 0) + + # Stable: within one rank the newest order (the tie-break) stays. + entries = sorted( + self._workspace_order(listings, "newest"), + key=lambda entry: -rank(entry), + ) + + def key(entry): + return _rated_key(entry, rank(entry)) + + else: + entries = self._workspace_order(listings, order) + + def key(entry): + return _workspace_key(entry, order) + + start = 0 if after is not None: - entries = [ - entry - for entry in entries - if ( - _workspace_key(entry, order) < after - if reverse - else _workspace_key(entry, order) > after - ) - ] - page = entries[:limit] - has_more = len(entries) > limit - next_cursor = ( - _workspace_cursor(_workspace_key(page[-1], order), scope) - if has_more - else None + if order == "name-desc": + ascending = self._workspace_order(listings, "name") + start = len(entries) - bisect_left(ascending, after, key=key) + else: + start = bisect_right(entries, after, key=key) + total, page, positive_count = _filtered_page( + entries, start, limit + 1, filters, links + ) + if not needed and after is not None: + positive_count = _COUNT_OMITTED + has_more = len(page) > limit + page = page[:limit] + next_cursor = _workspace_cursor(key(page[-1]), scope) if has_more else None + return ( + [(entry.path, entry.mtime, entry.root, entry.index) for entry in page], + total, + has_more, + next_cursor, + positive_count, ) - return page, total, has_more, next_cursor async def serve_image(self, request): """Serve the actual image file using streamed FileResponse.""" @@ -1679,6 +2096,91 @@ class ImageRoutesMixin: {"success": False, "error": _safe_error(e)}, status=500 ) + async def update_image_favorite(self, request, *, favorite: bool): + """PUT/DELETE one favorite, named by ``{root_index, path}``. + + A favorite is set only on a file in the root's listing; clearing also + accepts a well-formed path that is no longer listed, so a deleted + file's favorite can still be removed. Both are idempotent. + """ + data, error = await _json_body_or_error(request) + if error is not None: + return error + roots = list(self._get_all_output_dirs()) + target, error = _favorite_target(data, roots) + if error is not None: + return error + index, relative = target + root = roots[index] + relative = Path(relative).as_posix() + try: + files = await self._get_gallery_files(root) + listing = await self._run_in_executor(self._favorite_listing, root, files) + key = _favorite_relative_key(relative) + listed = listing.get(key) + if listed is None and favorite: + return web.json_response( + {"success": False, "error": "Image file not found"}, status=404 + ) + root_key = _favorite_root_key(root) + if favorite and max(len(root_key), len(key)) > MAX_KEY_LENGTH: + # A stored key past the schema bound would fail the next + # startup or restore validation. + return _forbidden("Path is too long to favorite") + update = ( + self.db.set_image_favorite if favorite else self.db.clear_image_favorite + ) + await self._run_in_executor(update, root_key, key) + except Exception: + self.logger.exception("Favorite update failed") + return web.json_response( + {"success": False, "error": "Favorite update failed"}, status=500 + ) + return web.json_response( + { + "success": True, + "favorite": favorite, + "root_index": index, + "path": listed or relative, + } + ) + + def _favorite_listing(self, root: Path, files: list) -> dict[str, str]: + """``{relative key: listing POSIX spelling}`` for one root's scan. + + Built once per scan (the gallery cache hands out the same list until + it rescans), so a toggle is a dictionary lookup. + """ + cache_key = str(root) + cached = self._favorite_listing_cache.get(cache_key) + if cached is not None and cached[0] is files: + return cached[1] + listing = {} + for path, _mtime in files: + posix = path.relative_to(root).as_posix() + listing.setdefault(_favorite_relative_key(posix), posix) + self._favorite_listing_cache[cache_key] = (files, listing) + return listing + + def _workspace_favorites(self, page: list[tuple]) -> set[tuple[int, str]]: + """``(root index, listing path)`` of the favorites among *page* entries.""" + by_root: dict[tuple[int, Path], dict[str, str]] = {} + for path, _mtime, root, index in page: + posix = path.relative_to(root).as_posix() + by_root.setdefault((index, root), {})[_favorite_relative_key(posix)] = posix + try: + found = set() + for (index, root), keys in by_root.items(): + for key in self.db.get_image_favorites( + _favorite_root_key(root), list(keys) + ): + found.add((index, keys[key])) + return found + except Exception: + # The listing stays usable; stars show as not favorited. + self.logger.warning("Image favorites could not be read") + return set() + async def inspect_image(self, request): """Read prompt evidence for one explicitly rooted physical media file.""" roots = list(self._get_all_output_dirs()) @@ -1712,7 +2214,7 @@ class ImageRoutesMixin: try: def inspect() -> bytes | None: - saved, unavailable = None, False + saved, unavailable, summary = None, False, None try: saved = self.db.get_image_prompt_info( str(path), False, report_ambiguity=True @@ -1720,6 +2222,13 @@ class ImageRoutesMixin: except Exception: unavailable = True self.logger.warning("Saved inspection display could not be read") + if saved and saved.get("association_status") != "ambiguous": + try: + summary = self.db.get_image_link_summaries([str(path)], False) + summary = summary[str(path)] + except Exception: + # The link stays known; only its title is withheld. + self.logger.warning("Inspection link title could not be read") try: observed = inspect_current_file( path, @@ -1731,6 +2240,9 @@ class ImageRoutesMixin: ) except CurrentFileChanged: return None + observed["library"]["title"] = _verified_title( + summary, _positive_title(_inspected_positive(observed)) + ) return encode_inspection_response( image={ "root_index": index, @@ -1764,6 +2276,7 @@ class ImageRoutesMixin: ) def _inspection_library(self, saved: dict | None, *, unavailable: bool) -> dict: + """The inspector's Library facts; the route sets the title afterwards.""" result = { "status": "unavailable" if unavailable else "none", "scope": "saved-library", @@ -1785,7 +2298,6 @@ class ImageRoutesMixin: status="linked", positive_id=pair.get("positive_prompt_id"), pair_id=pair.get("id"), - title=pair.get("title"), category=pair.get("category"), tags=pair.get("tags"), metadata_scope="pair", @@ -1897,24 +2409,106 @@ class ImageRoutesMixin: if _in_subfolder(str(entry[0].relative_to(entry[2]).parent), subfolder) ] - def _format_workspace_page(self, page: list[tuple]) -> list[dict]: + def _workspace_link_summaries(self, paths: list[str]) -> dict[str, dict]: + """Card link summaries, matched exactly as the inspection route does.""" + try: + return self.db.get_image_link_summaries(paths, False) + except Exception: + # The listing stays usable; cards fall back to their filenames. + self.logger.warning("Workspace link summaries could not be read") + return {} + + def _card_parse(self, path: Path, records: dict | None) -> tuple: + """One card's ``(seed, positive title)``; ``(None, None)`` on any error. + + One file's unexpected parse failure must not fail the whole page: its + card shows no seed and withholds a first-line title. + """ + try: + params = header_params(records) + return card_seed(params), _positive_title(parsed_positive(params)) + except Exception: + self.logger.warning( + "Workspace card facts could not be parsed for %s", + path.name, + exc_info=True, + ) + return None, None + + def _read_card_facts(self, path: Path, *, parse: bool) -> _CardFacts | None: + """Read one image's header; ``parse`` adds the seed and positive title. + + None when the file cannot be read at all (corrupt, locked, vanishing). + """ + try: + width, height, records = _card_header(path) + except (OSError, ValueError, Image.DecompressionBombError): + # A corrupt or disappearing preview must not hide the listing. + return None + if not parse: + return _CardFacts(width, height, None, None) + # Parsed after the handle closes, so a concurrent rename or delete + # (which Windows refuses on open files) is not held up by it. + return _CardFacts(width, height, *self._card_parse(path, records)) + + def _card_facts(self, entry: tuple, *, parse: bool) -> _CardFacts: + """One listed image's card facts, read once per ``(path, mtime)``. + + Kept beside the scan's per-file values (``_workspace_listed``) and + invalidated the same way: a rescan keeps the facts of unchanged files + and drops the rest, and a roots change drops them all. Only files of + the current scan are stored, so the cache is bounded by the listing. + """ + path, mtime, root, index = entry[:4] + root_key, key = (index, str(root)), (path, mtime) + cached = self._workspace_card_cache.get(root_key, {}).get(key) + if cached is not None: + return cached + facts = self._read_card_facts(path, parse=parse) + if facts is None: + # Not stored: a locked or half-written file may read next time. + return _NO_CARD_FACTS + listed = self._workspace_listed_cache.get(root_key) + if parse and listed is not None and key in listed[2]: + self._workspace_card_cache.setdefault(root_key, {})[key] = facts + return facts + + def _format_workspace_page( + self, page: list[tuple], *, linked: bool = True + ) -> list[dict]: + """Workspace descriptors; ``linked`` adds card facts (link, seed, favorite).""" images = self._format_output_page(page, scoped=True) - paths = { - (entry[3], entry[0].relative_to(entry[2]).as_posix()): entry[0] + entries = { + (entry[3], entry[0].relative_to(entry[2]).as_posix()): entry for entry in page } + summaries = ( + self._workspace_link_summaries( + [str(entry[0]) for entry in entries.values()] + ) + if linked + else {} + ) for image in images: - image["width"] = None - image["height"] = None - if image["is_video"]: - continue - try: - with paths[(image["root_index"], image["path"])].open("rb") as stream: - with Image.open(stream) as source: - image["width"], image["height"] = source.size - except (OSError, ValueError, Image.DecompressionBombError): - # A corrupt or disappearing preview must not hide the listing. + entry = entries[(image["root_index"], image["path"])] + # Videos: the card reads no video metadata. Library previews show + # no card facts, so they skip the parse (``parse=linked``). + facts = ( + _NO_CARD_FACTS + if image["is_video"] + else self._card_facts(entry, parse=linked) + ) + image["width"], image["height"] = facts.width, facts.height + if not linked: continue + image["seed"] = facts.seed + summary = summaries.get(str(entry[0])) or _unavailable_link() + image["linked"] = _card_link(summary, facts.positive_title) + if linked: + # Image cards carry the star; Library previews show none. + favorites = self._workspace_favorites(page) + for image in images: + image["favorite"] = (image["root_index"], image["path"]) in favorites return images def _format_output_page(self, page, *, scoped=False): diff --git a/py/api/library.py b/py/api/library.py index 3532e49..ba9e458 100644 --- a/py/api/library.py +++ b/py/api/library.py @@ -233,7 +233,9 @@ class LibraryRoutesMixin: self.add_library_negative ) - def _library_media(self, records: list[dict], roots: list[Path]) -> list[dict]: + def _library_media( + self, records: list[dict], roots: list[Path], *, linked: bool = True + ) -> list[dict]: page = [] for record in records: try: @@ -251,7 +253,7 @@ class LibraryRoutesMixin: break except (OSError, ValueError, RuntimeError): continue - return self._format_workspace_page(page) + return self._format_workspace_page(page, linked=linked) @pair_response async def get_library_options(self, request: web.Request) -> web.Response: @@ -296,7 +298,10 @@ class LibraryRoutesMixin: def previews() -> None: for parent in result["items"]: - parent["previews"] = self._library_media(records[parent["id"]], roots) + # Previews belong to their parent; a per-parent link read is N+1. + parent["previews"] = self._library_media( + records[parent["id"]], roots, linked=False + ) await self._run_in_executor(previews) return web.json_response({"success": True, **result}) diff --git a/tests/js/support/inspection-seed.mjs b/tests/js/support/inspection-seed.mjs new file mode 100644 index 0000000..9aaa338 --- /dev/null +++ b/tests/js/support/inspection-seed.mjs @@ -0,0 +1,39 @@ +// The seed the workspace inspector shows for encoded inspection responses, through +// its own path: projectInspectionPayload with the server facts and the legacy +// parser fill (comfy-metadata.js, passed explicitly; Node has no window global). +// Used by tests/test_api_workspace_images.py to compare card and inspector seeds. +// stdin: {listing, inspections: [{name, body}]}; each body is the exact response +// JSON text, and listing (optional) is the exact workspace listing response text. +// stdout: {inspector: {name: {status, value}}, card: {path: String(seed) | null}}. +// Inspector status is 'error' where the page cannot project the record. The card +// values come from JSON.parse of the listing, as the browser reads them, so a +// seed sent as a JSON number past 2**53 shows up rounded here. +import { createRequire } from 'node:module'; +import { fileURLToPath, pathToFileURL } from 'node:url'; +import path from 'node:path'; + +const here = path.dirname(fileURLToPath(import.meta.url)); +const web = path.join(here, '..', '..', '..', 'web', 'js'); +const require = createRequire(import.meta.url); +const parser = require(path.join(web, 'comfy-metadata.js')); +const { projectInspectionPayload } = await import(pathToFileURL(path.join(web, 'workspace', 'inspection-state.mjs')).href); + +const chunks = []; +for await (const chunk of process.stdin) chunks.push(chunk); +const { listing = null, inspections } = JSON.parse(Buffer.concat(chunks).toString('utf8')); +const inspector = {}; +for (const { name, body } of inspections) { + const identity = { kind: 'rooted-file', key: `0:${name}`, root_index: 0, path: name, root_scope: null }; + try { + const record = projectInspectionPayload(JSON.parse(body), identity, null, parser.extractGenerationFacts, false, parser); + inspector[name] = { status: record.generation.seed.status, value: record.generation.seed.value }; + } catch (error) { + // The page shows no record here, so no seed either. + inspector[name] = { status: 'error', value: String(error && error.name) }; + } +} +const card = {}; +for (const item of listing === null ? [] : JSON.parse(listing).images) { + card[item.path] = item.seed === null || item.seed === undefined ? null : String(item.seed); +} +process.stdout.write(JSON.stringify({ inspector, card })); diff --git a/tests/js/workspace-library-title.test.js b/tests/js/workspace-library-title.test.js new file mode 100644 index 0000000..e8cadf2 --- /dev/null +++ b/tests/js/workspace-library-title.test.js @@ -0,0 +1,15 @@ +import test from 'node:test'; +import assert from 'node:assert/strict'; + +const { libraryTitle } = await import('../../web/js/workspace/library.mjs').catch(() => ({})); + +test('library title counts UTF-16 units and never splits a surrogate pair', () => { + assert.equal(typeof libraryTitle, 'function', 'libraryTitle is exported'); + assert.equal(libraryTitle(' short line \nsecond'), 'short line'); + assert.equal(libraryTitle('a'.repeat(80)), 'a'.repeat(80)); + assert.equal(libraryTitle('a'.repeat(81)), `${'a'.repeat(80)}…`); + // 80 emoji are 160 UTF-16 units: 40 fit. + assert.equal(libraryTitle('😀'.repeat(80)), `${'😀'.repeat(40)}…`); + // 79 units then an emoji: the emoji would straddle the limit, so it is dropped whole. + assert.equal(libraryTitle(`${'a'.repeat(79)}😀b`), `${'a'.repeat(79)}…`); +}); diff --git a/tests/test_api_library.py b/tests/test_api_library.py index cd30b55..6cc6329 100644 --- a/tests/test_api_library.py +++ b/tests/test_api_library.py @@ -188,6 +188,10 @@ class TestLibraryAPI(ImageAPITestCase): async def test_maximum_parent_page_uses_bounded_sql_queries_for_previews(self): for index in range(200): positive = self._save_prompt(f"Parent {index}", tags=[f"tag {index}"]) + if index % 3 == 0: + self.api.db.save_prompt_pair(f"Parent {index}", f"Negative {index}") + if index % 4 == 0: + self.api.db.record_prompt_use(positive, times=index + 1) make_png(self.output_dir / f"{index}.png") self.api.db.link_image_to_prompt( positive, str(self.output_dir / f"{index}.png") @@ -232,6 +236,32 @@ class TestLibraryAPI(ImageAPITestCase): for item in result["items"] ) ) + # Usage fields and the variant total ride on the same five SELECTs. + self.assertEqual(result["variant_total"], 67) + self.assertEqual( + result["variant_total"], + sum(item["variant_count"] for item in result["items"]), + ) + used = { + int(item["text"].split()[-1]): item + for item in result["items"] + if item["last_used_at"] is not None + } + self.assertEqual(sorted(used), list(range(0, 200, 4))) + self.assertTrue( + all( + item["run_count"] == index + 1 + and item["last_used_at"].endswith("+00:00") + for index, item in used.items() + ) + ) + self.assertTrue( + all( + type(item["run_count"]) is int + for item in result["items"] + if item["last_used_at"] is None + ) + ) async def test_long_name_cursor_stays_small_and_preserves_full_sorting(self): prefix = "猫" * 10_000 @@ -438,6 +468,38 @@ class TestLibraryAPI(ImageAPITestCase): ) return response, await response.json() + async def test_unconvertible_last_use_does_not_fail_list_detail_or_receipt(self): + positive = self._save_prompt("Cat") + other = self._save_prompt("Dog") + with self.api.db.model.get_connection() as connection: + connection.execute( + "UPDATE prompts SET last_used_at=? WHERE id=?", + ("9999-12-31T23:59:59-01:00", positive), + ) + connection.execute( + "UPDATE prompts SET last_used_at=? WHERE id=?", + ("0001-01-01T00:00:00+01:00", other), + ) + response, listed = await self.get() + self.assertEqual(response.status, 200, listed) + self.assertEqual( + [item["last_used_at"] for item in listed["items"]], [None, None] + ) + response, detail = await self.get(f"/{positive}") + self.assertEqual(response.status, 200, detail) + self.assertIsNone(detail["positive"]["last_used_at"]) + response, receipt = await self.patch_parent( + positive, + { + "expected_epoch": detail["library_epoch"], + "if_match": detail["positive"]["edit_token"], + "metadata": {"title": "Portrait"}, + }, + ) + self.assertEqual(response.status, 200, receipt) + self.assertEqual(receipt["positive"]["title"], "Portrait") + self.assertIsNone(receipt["positive"]["last_used_at"]) + async def test_native_parent_patch_strict_validation(self): from database.execution_history_schema import rotate_history_epoch diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 425c9df..83aa576 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -1,16 +1,135 @@ """Opt-in workspace gallery contracts against disposable SQLite and media.""" +import asyncio +import json import os +import shutil +import subprocess # nosec B404 +import sys import unittest from pathlib import Path from unittest.mock import AsyncMock, patch +from PIL import Image, PngImagePlugin + +from utils.generation_params import extract_from_records + try: from test_api_images import ImageAPITestCase, make_png from open_handles import assert_closed + from generation_params_cases import CASES, encoder, ksampler except ImportError: from tests.test_api_images import ImageAPITestCase, make_png from tests.open_handles import assert_closed + from tests.generation_params_cases import CASES, encoder, ksampler + +# Wall-clock budgets run only on request and untraced (TESTING.md +# "Performance budgets"); their structure is asserted everywhere. +PERF_BUDGET = os.environ.get("PROMPTMANAGER_PERF_BUDGET") == "1" +TRACED = sys.gettrace() is not None or ( + sys.version_info >= (3, 12) + and any(sys.monitoring.get_tool(tool) for tool in range(6)) +) +SEED_HARNESS = ( + Path(__file__).resolve().parent / "js" / "support" / "inspection-seed.mjs" +) +NODE_TIMEOUT_SECONDS = 120 + + +def embedded_png(path, positive, negative=""): + """A PNG whose API prompt graph resolves to *positive* and *negative*. + + The sampler has a model link, as every ComfyUI KSampler does, so the shared + parser (which needs one) and the server resolver both read the prompts. + """ + graph = { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}}, + "3": { + "class_type": "KSampler", + "inputs": {"positive": ["1", 0], "negative": ["2", 0], "model": ["4", 0]}, + }, + "4": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "fixture.safetensors"}, + }, + } + info = PngImagePlugin.PngInfo() + info.add_text("prompt", json.dumps(graph)) + with Image.new("RGB", (4, 4)) as image: + image.save(path, pnginfo=info) + os.utime(path, (100, 100)) + return path + + +def records_png(path, records, mtime=100): + """A PNG carrying the given text records (None values are left out).""" + info = PngImagePlugin.PngInfo() + for key, text in records.items(): + if text is not None: + info.add_text(key, text) + with Image.new("RGB", (4, 4)) as image: + image.save(path, pnginfo=info) + os.utime(path, (mtime, mtime)) + return path + + +def png_chunk(kind, data): + """One PNG chunk: length, type, data and CRC.""" + import struct + import zlib + + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + +def png_with_text_chunks(path, before=(), after=(), mtime=100): + """A PNG with ``(keyword, text)`` tEXt chunks before and after its IDAT.""" + import io + + buffer = io.BytesIO() + with Image.new("RGB", (4, 4)) as image: + image.save(buffer, format="PNG") + data = buffer.getvalue() + idat = data.index(b"IDAT") - 4 + iend = data.index(b"IEND") - 4 + + def chunks(records): + return b"".join( + png_chunk(b"tEXt", key.encode("latin-1") + b"\0" + text.encode("latin-1")) + for key, text in records + ) + + path.write_bytes( + data[:idat] + chunks(before) + data[idat:iend] + chunks(after) + data[iend:] + ) + os.utime(path, (mtime, mtime)) + return path + + +def seeded_records(positive, seed): + """API prompt records for a one-sampler graph with *positive* and *seed*.""" + graph = { + "4": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "concept.safetensors"}, + }, + "6": encoder(positive), + "7": encoder("blurry"), + "3": ksampler(("6", 0), ("7", 0), seed=seed), + } + return {"prompt": json.dumps(graph)} + + +def header_params_of_images(): + """The card parser the images module calls (unpatched).""" + from py.api import images + + return images.header_params class TestWorkspaceImages(ImageAPITestCase): @@ -300,6 +419,1246 @@ class TestWorkspaceImages(ImageAPITestCase): }, ) + def _link_fixtures(self): + """Linked (titled, first-line, pair, untitled pair), ambiguous and none.""" + db = self.api.db + long_line = "lantern " * 15 + titled = self._save_prompt("Titled positive", category="Night", rating=5) + with db.model.get_connection() as connection: + connection.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (titled, "Stored library title"), + ) + plain = self._save_prompt(long_line) + pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair title") + bare = db.save_prompt_pair("Bare pair positive\nmore", "Blur") + first, second = self._save_prompt("Owner one"), self._save_prompt("Owner two") + files = { + name: self.media(f"{name}.png") + for name in ("titled", "pair", "ambiguous", "none") + } + # First-line titles come from files whose own prompt evidence agrees. + files["plain"] = embedded_png(self.output_dir / "plain.png", long_line) + files["bare"] = embedded_png( + self.output_dir / "bare.png", "Bare pair positive\nmore", "Blur" + ) + db.link_image_to_prompt(titled, str(files["titled"])) + db.link_image_to_prompt(plain, str(files["plain"])) + for name, pair_id in (("pair", pair), ("bare", bare)): + positive = db.get_prompt_pair(pair_id)["positive_prompt_id"] + image = db.link_image_to_prompt(positive, str(files[name])) + db.link_image_to_pair(pair_id, image, {"source": "fixture"}) + db.link_image_to_prompt(first, str(files["ambiguous"])) + db.link_image_to_prompt(second, str(files["ambiguous"])) + return { + "titled": ("linked", titled, "Stored library title"), + "plain": ("linked", plain, long_line.strip()[:80] + "\u2026"), + "pair": ( + "linked", + db.get_prompt_pair(pair)["positive_prompt_id"], + "Pair title", + ), + "bare": ( + "linked", + db.get_prompt_pair(bare)["positive_prompt_id"], + "Bare pair positive", + ), + "ambiguous": ("ambiguous", None, None), + "none": ("none", None, None), + } + + async def test_workspace_items_carry_link_summary_matching_inspection(self): + expected = self._link_fixtures() + _, data = await self.get(limit=50) + items = {Path(item["filename"]).stem: item for item in data["images"]} + self.assertEqual(set(items), set(expected)) + for name, (status, positive_id, title) in expected.items(): + with self.subTest(file=name): + linked = items[name]["linked"] + self.assertEqual( + set(linked), + { + "status", + "positive_id", + "title", + "title_source", + "category", + "rating", + }, + ) + self.assertEqual( + (linked["status"], linked["positive_id"], linked["title"]), + (status, positive_id, title), + ) + response = await self.client.get( + f"/prompt_manager/images/inspect/{name}.png", + params={"root": 0, "root_scope": items[name]["root_scope"]}, + ) + self.assertEqual(response.status, 200) + library = (await response.json())["library"] + self.assertEqual( + (library["status"], library["positive_id"], library["title"]), + (linked["status"], linked["positive_id"], linked["title"]), + ) + self.assertEqual( + ( + items["titled"]["linked"]["category"], + items["titled"]["linked"]["rating"], + ), + ("Night", 5), + ) + + async def test_stored_title_shows_on_card_and_inspector_without_evidence(self): + # A stored title is user data: shown on both, even when it equals the + # first line and the file carries no prompt metadata. + prompt = self._save_prompt("Same words\nrest") + with self.api.db.model.get_connection() as connection: + connection.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (prompt, "Same words"), + ) + path = self.media("stored.png") + self.api.db.link_image_to_prompt(prompt, str(path)) + _, data = await self.get() + linked = data["images"][0]["linked"] + self.assertEqual( + (linked["title"], linked["title_source"]), ("Same words", "stored") + ) + library = await self.inspected("stored.png", data["images"][0]) + self.assertEqual(library["title"], linked["title"]) + self.assertEqual(library["positive_id"], prompt) + + async def test_unconfirmed_text_title_is_withheld_on_card_and_inspector(self): + prompt = self._save_prompt("Secret first line\nrest") + path = self.media("plain.png") + self.api.db.link_image_to_prompt(prompt, str(path)) + _, data = await self.get() + item = data["images"][0] + response = await self.client.get( + "/prompt_manager/images/inspect/plain.png", + params={"root": 0, "root_scope": item["root_scope"]}, + ) + body = await response.json() + self.assertEqual(body["library"]["positive_id"], prompt) + self.assertIsNone(body["library"]["title"]) + self.assertNotIn("Secret first line", json.dumps(body)) + self.assertIsNone(item["linked"]["title"]) + self.assertNotIn("Secret first line", json.dumps(data)) + + async def test_summary_failure_keeps_inspector_link(self): + prompt = self._save_prompt("Readable link") + path = self.media("a.png") + self.api.db.link_image_to_prompt(prompt, str(path)) + _, data = await self.get() + with patch.object( + self.api.db, + "get_image_link_summaries", + side_effect=RuntimeError("summary gone"), + ): + library = await self.inspected("a.png", data["images"][0]) + self.assertEqual(library["status"], "linked") + self.assertEqual(library["positive_id"], prompt) + self.assertIsNone(library["title"]) + + async def inspected(self, name, item): + response = await self.client.get( + f"/prompt_manager/images/inspect/{name}", + params={"root": 0, "root_scope": item["root_scope"]}, + ) + self.assertEqual(response.status, 200) + return (await response.json())["library"] + + async def test_link_summary_uses_one_query_per_page(self): + prompt = self._save_prompt("Paged prompt") + for index in range(12): + path = self.media(f"{index:02}.png", mtime=100 + index) + if index % 2: + self.api.db.link_image_to_prompt(prompt, str(path)) + statements = [] + original = self.api.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + with patch.object( + self.api.db.model, "get_connection", side_effect=trace_connection + ): + response, data = await self.get(limit=12) + self.assertEqual(response.status, 200) + self.assertEqual(len(data["images"]), 12) + self.assertEqual( + sorted(item["linked"]["status"] for item in data["images"]), + ["linked"] * 6 + ["none"] * 6, + ) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + # The page's link summary, the positive count's one link read and + # the page's one favorites read. + self.assertEqual(len(selects), 3, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) + self.assertEqual(data["positive_count"], 1) # six files, one prompt + + async def test_link_summary_runs_in_executor_and_survives_database_errors(self): + import threading + + self.media("a.png") + event_thread = threading.get_ident() + threads = [] + + def broken(paths, allow_filename_fallback=True): + threads.append(threading.get_ident()) + raise RuntimeError("database is gone") + + with patch.object(self.api.db, "get_image_link_summaries", side_effect=broken): + response, data = await self.get() + self.assertEqual(response.status, 200) + self.assertTrue(threads) + self.assertNotEqual(threads[0], event_thread) + self.assertEqual(data["images"][0]["linked"]["status"], "unavailable") + self.assertIsNone(data["images"][0]["linked"]["title"]) + self.assertIsNone(data["images"][0]["linked"]["title_source"]) + self.assertNotIn("database is gone", str(data)) + + async def test_items_carry_seed_from_header_records(self): + records_png(self.output_dir / "seeded.png", seeded_records("a fox", 42001)) + self.media("plain.png") + self.media("clip.mp4") + self.media("broken.png").write_bytes(b"invalid") + _, data = await self.get() + seeds = {item["filename"]: item["seed"] for item in data["images"]} + self.assertEqual( + seeds, + { + "seeded.png": "42001", + "plain.png": None, + "clip.mp4": None, + "broken.png": None, + }, + ) + + async def test_card_reads_each_header_once(self): + from py.api import images + + for index in range(3): + records_png( + self.output_dir / f"{index}.png", seeded_records("a fox", index) + ) + self.media("clip.mp4") + with ( + patch.object( + images, "read_inspection_file", wraps=images.read_inspection_file + ) as read, + patch("PIL.Image.open", wraps=Image.open) as opened, + ): + _, data = await self.get() + self.assertEqual(len(data["images"]), 4) + # One inspector read per PNG (dimensions and records); no decoder open. + self.assertEqual(read.call_count, 3) + self.assertEqual(opened.call_count, 0) + + async def test_card_reads_png_records_as_the_inspector_does(self): + one, two = seeded_records("a fox", 1), seeded_records("a fox", 2) + # Conflicting duplicate prompt chunks: the inspector refuses the file, + # so the card shows no seed (Pillow alone would keep the last one). + png_with_text_chunks( + self.output_dir / "conflict.png", + before=[("prompt", one["prompt"]), ("prompt", two["prompt"])], + mtime=300, + ) + # Text after IDAT: Pillow's header open misses it; the inspector reads it. + png_with_text_chunks( + self.output_dir / "late.png", after=[("prompt", one["prompt"])], mtime=200 + ) + png_with_text_chunks( + self.output_dir / "repeat.png", + before=[("prompt", one["prompt"]), ("prompt", one["prompt"])], + ) + response, data = await self.get() + self.assertEqual(response.status, 200, data) + items = {item["filename"]: item for item in data["images"]} + self.assertIsNone(items["conflict.png"]["seed"]) + self.assertEqual( + (items["conflict.png"]["width"], items["conflict.png"]["height"]), (4, 4) + ) + self.assertEqual(items["late.png"]["seed"], "1") + self.assertEqual(items["repeat.png"]["seed"], "1") + shown = await self._inspector_seeds(data["images"]) + self.assertNotEqual(shown["conflict.png"]["status"], "known") + self.assertEqual(shown["late.png"], {"status": "known", "value": "1"}) + + async def test_deeply_nested_record_still_lists_and_inspects(self): + nested = "[" * 900 + "1" + "]" * 900 + workflow = ( + '{"nodes":[{"id":%s,"type":"KSampler","widgets_values":[1]}],' + '"links":[[%s,1,0,2,0]]}' % (nested, nested) + ) + records_png(self.output_dir / "ok.png", seeded_records("a fox", 5), mtime=200) + records_png(self.output_dir / "evil.png", {"workflow": workflow}) + prompt = self._save_prompt("Evil first line") + self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "evil.png")) + response, data = await self.get() + self.assertEqual(response.status, 200, data) + items = {item["filename"]: item for item in data["images"]} + self.assertEqual(items["ok.png"]["seed"], "5") + response = await self.client.get( + "/prompt_manager/images/inspect/evil.png", + params={"root": 0, "root_scope": items["evil.png"]["root_scope"]}, + ) + self.assertEqual(response.status, 200, await response.text()) + library = (await response.json())["library"] + self.assertEqual(library["title"], items["evil.png"]["linked"]["title"]) + shown = (await self._inspector_seeds([items["evil.png"]]))["evil.png"] + card = items["evil.png"]["seed"] + if card is None: + self.assertNotEqual(shown["status"], "known") + else: + self.assertEqual(shown, {"status": "known", "value": card}) + + async def test_warm_page_reuses_card_facts_per_scan(self): + from py.api import images + + a = records_png(self.output_dir / "a.png", seeded_records("a fox", 1), 300) + b = records_png(self.output_dir / "b.png", seeded_records("a cat", 2), 200) + gone = records_png(self.output_dir / "gone.png", seeded_records("x", 3), 100) + real_read, real_parse = images.read_inspection_file, images.header_params + read, parsed = [], [] + + def reading(path): + read.append(path.name) + return real_read(path) + + def parsing(info): + parsed.append(info) + return real_parse(info) + + async def page(): + read.clear() + parsed.clear() + response, data = await self.get() + self.assertEqual(response.status, 200, data) + return {item["filename"]: item for item in data["images"]} + + with ( + patch.object(images, "read_inspection_file", side_effect=reading), + patch.object(images, "header_params", side_effect=parsing), + ): + cold = await page() + self.assertEqual(sorted(read), ["a.png", "b.png", "gone.png"]) + self.assertEqual(len(parsed), 3) + # A warm page reads and parses no header. + prompt = self._save_prompt("a fox") + self.api.db.link_image_to_prompt(prompt, str(a)) + warm = await page() + self.assertEqual((read, parsed), ([], [])) + self.assertEqual(warm["a.png"]["seed"], "1") + self.assertEqual((warm["a.png"]["width"], warm["a.png"]["height"]), (4, 4)) + # Links are never cached: the new link's confirmed title shows. + self.assertIsNone(cold["a.png"]["linked"]["title"]) + self.assertEqual(warm["a.png"]["linked"]["title"], "a fox") + # A monitor invalidation keeps the values of unchanged files. + records_png(b, seeded_records("a cat", 22), mtime=250) + gone.unlink() + self.api.invalidate_gallery_cache() + changed = await page() + self.assertEqual(read, ["b.png"]) + self.assertEqual(changed["b.png"]["seed"], "22") + cards = self.api._workspace_card_cache[(0, str(self.output_dir))] + self.assertEqual( + sorted(path.name for path, _mtime in cards), ["a.png", "b.png"] + ) + # A header that cannot be read (locked, half-written) is not stored. + records_png(self.output_dir / "c.png", seeded_records("a cow", 4), 400) + self.api.invalidate_gallery_cache() + with patch.object(images, "_card_header", side_effect=OSError("busy")): + locked = await page() + self.assertIsNone(locked["c.png"]["seed"]) + self.assertEqual(locked["a.png"]["seed"], "1") + self.assertEqual((await page())["c.png"]["seed"], "4") + self.assertEqual(read, ["c.png"]) + # A roots change drops them, as it drops the per-file values. + self.api.reset_gallery_roots_cache() + self.assertEqual(self.api._workspace_card_cache, {}) + await page() + self.assertEqual(sorted(read), ["a.png", "b.png", "c.png"]) + + async def test_unbounded_digit_string_seed_still_lists(self): + records_png( + self.output_dir / "huge.png", seeded_records("a fox", "9" * 5000), 200 + ) + records_png(self.output_dir / "ok.png", seeded_records("a fox", 5)) + response, data = await self.get() + self.assertEqual(response.status, 200, data) + items = {item["filename"]: item for item in data["images"]} + self.assertIsNone(items["huge.png"]["seed"]) + self.assertEqual(items["ok.png"]["seed"], "5") + + async def test_one_cards_parse_error_spares_the_page(self): + prompt = self._save_prompt("a fox") + for name, mtime in (("bad.png", 200), ("ok.png", 100)): + path = records_png( + self.output_dir / name, seeded_records("a fox", 5), mtime + ) + self.api.db.link_image_to_prompt(prompt, str(path)) + real_parse = header_params_of_images() + + def parsing(info): + if parsing.calls == 0: + parsing.calls += 1 + raise RuntimeError("unexpected parser bug") + return real_parse(info) + + parsing.calls = 0 + with ( + patch("py.api.images.header_params", side_effect=parsing), + self.assertLogs(self.api.logger, "WARNING") as logged, + ): + response, data = await self.get() + self.assertEqual(response.status, 200, data) + items = {item["filename"]: item for item in data["images"]} + self.assertIsNone(items["bad.png"]["seed"]) + self.assertIsNone(items["bad.png"]["linked"]["title"]) + self.assertEqual(items["bad.png"]["linked"]["status"], "linked") + self.assertEqual(items["ok.png"]["seed"], "5") + self.assertEqual(items["ok.png"]["linked"]["title"], "a fox") + self.assertNotIn("unexpected parser bug", json.dumps(data)) + self.assertIn("unexpected parser bug", "\n".join(logged.output)) + + async def test_library_previews_skip_card_facts(self): + records_png(self.output_dir / "a.png", seeded_records("a fox", 5)) + page = await self.api._collect_output_images( + self.api._get_all_output_dirs(), "" + ) + with patch("py.api.images.header_params") as parsed: + items = self.api._format_workspace_page(page, linked=False) + parsed.assert_not_called() + self.assertNotIn("seed", items[0]) + self.assertEqual((items[0]["width"], items[0]["height"]), (4, 4)) + + async def test_card_facts_parse_after_the_file_is_closed(self): + records_png(self.output_dir / "a.png", seeded_records("a fox", 5)) + streams, closed = [], [] + real_open, real_parse = Path.open, header_params_of_images() + + def opening(path, *args, **kwargs): + stream = real_open(path, *args, **kwargs) + if path.name == "a.png": + streams.append(stream) + return stream + + def parsing(info): + closed.append(bool(streams) and all(s.closed for s in streams)) + return real_parse(info) + + with ( + patch.object(Path, "open", autospec=True, side_effect=opening), + patch("py.api.images.header_params", side_effect=parsing), + ): + _, data = await self.get() + self.assertEqual(data["images"][0]["seed"], "5") + self.assertEqual(closed, [True]) + + def _seed_files(self): + """Concept-style one-sampler files plus every shared parser case.""" + expected = {} + for i in range(2): + for j in range(3): + name = f"concept-{i}-{j}.png" + seed = 42001 + i * 100 + j + records_png( + self.output_dir / name, seeded_records(f"concept {i} {j}", seed) + ) + expected[name] = str(seed) + # ComfyUI's randomized seeds mostly exceed 2**53, where JSON.parse rounds. + big = 2**63 + 12345 + records_png(self.output_dir / "big-seed.png", seeded_records("big", big)) + expected["big-seed.png"] = str(big) + for case, records in CASES.items(): + records_png(self.output_dir / f"case-{case}.png", records) + records_png( + self.output_dir / "a1111.png", + {"parameters": "a fox\nNegative prompt: blur\nSteps: 20, Seed: 9"}, + ) + self.media("bare.png") + return expected + + async def _inspector_seeds(self, items): + """The seed the inspector shows for each listed item (Node harness).""" + return (await self._node_seeds(items))["inspector"] + + async def _node_seeds(self, items, listing=None): + """The harness output: inspector seeds and, from *listing*, card seeds. + + *listing* is the exact listing response text; the harness JSON.parses + it as the browser does, so the card seeds are what a page would read. + """ + node = shutil.which("node") + if node is None: # pragma: no cover - CI runners and dev hosts ship Node + if os.environ.get("CI"): + self.fail("Node.js is required in CI for the seed parity check") + self.skipTest("Node.js is not installed") + bodies = [] + for item in items: + response = await self.client.get( + f"/prompt_manager/images/inspect/{item['path']}", + params={"root": 0, "root_scope": item["root_scope"]}, + ) + self.assertEqual(response.status, 200, item["path"]) + bodies.append({"name": item["path"], "body": await response.text()}) + # Fixed local Node harness, no shell. + completed = subprocess.run( # nosec B603 + [node, str(SEED_HARNESS)], + input=json.dumps({"listing": listing, "inspections": bodies}).encode( + "utf-8" + ), + capture_output=True, + timeout=NODE_TIMEOUT_SECONDS, + check=False, + ) + self.assertEqual(completed.returncode, 0, completed.stderr.decode("utf-8")) + return json.loads(completed.stdout) + + async def test_card_seed_equals_inspector_seed_when_both_present(self): + expected = self._seed_files() + response = await self.client.get( + "/prompt_manager/images/output", + params={"mode": "workspace", "limit": 200}, + ) + listing = await response.text() + items = json.loads(listing)["images"] + self.assertEqual(len(items), len(CASES) + len(expected) + 2) + seeds = await self._node_seeds(items, listing) + inspector, browser = seeds["inspector"], seeds["card"] + both = 0 + for item in items: + with self.subTest(file=item["path"]): + card, shown = item["seed"], inspector[item["path"]] + # Sent as a decimal string, so JSON.parse keeps every digit. + self.assertTrue(card is None or isinstance(card, str), card) + self.assertEqual(browser[item["path"]], card) + if item["path"] in expected: + self.assertEqual(card, expected[item["path"]]) + if card is None: + continue + # The card never shows a seed the inspector does not. + self.assertEqual(shown, {"status": "known", "value": card}) + both += 1 + self.assertGreaterEqual(both, len(expected) + 10) + self.assertGreater(int(browser["big-seed.png"]), 2**53) + + async def test_text_title_needs_the_files_parsed_positive(self): + prompt = self._save_prompt("Saved words\nrest") + other = records_png( + self.output_dir / "other.png", seeded_records("Different words", 1) + ) + shared = records_png( + self.output_dir / "shared.png", CASES["shared-prompt-pipeline"] + ) + clip = self.media("clip.mp4") + shared_prompt = self._save_prompt( + extract_from_records(CASES["shared-prompt-pipeline"]["prompt"], None)[ + "positivePrompt" + ] + ) + for path, owner in ((other, prompt), (shared, shared_prompt), (clip, prompt)): + self.api.db.link_image_to_prompt(owner, str(path)) + _, data = await self.get() + items = {item["filename"]: item for item in data["images"]} + for name in ("other.png", "clip.mp4"): + with self.subTest(file=name): + self.assertEqual(items[name]["linked"]["status"], "linked") + self.assertEqual(items[name]["linked"]["title_source"], "text") + self.assertIsNone(items[name]["linked"]["title"]) + library = await self.inspected("other.png", items["other.png"]) + self.assertIsNone(library["title"]) + self.assertNotIn("Saved words", json.dumps(data)) + # The resolver cannot decide this graph; the shared parser (and so the + # inspector's prompt) confirms the saved first line on both. + title = items["shared.png"]["linked"]["title"] + self.assertIsNotNone(title) + self.assertEqual( + (await self.inspected("shared.png", items["shared.png"]))["title"], title + ) + + # B3: search, category/tag filters and rating sort over linked positives. + + def _names(self, data): + return [item["filename"] for item in data["images"]] + + def _filter_fixtures(self): + """Linked, pair-linked, ambiguous and unlinked files for the B3 filters.""" + db = self.api.db + night_ab = self._save_prompt( + "A ZEBRA under neon", category="Night", tags=["a", "b"], rating=5 + ) + night_a = self._save_prompt("Quiet harbor", category="Night", tags=["a"]) + day_ab = self._save_prompt( + "Morning field", category="Day", tags=["a", "b", "zebra-tag"], rating=3 + ) + titled = self._save_prompt("Plain words", category="Day", rating=1) + with db.model.get_connection() as connection: + connection.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (titled, "Zebra stored title"), + ) + percent = self._save_prompt("100% wild_card", category="Day") + pair_positive = self._save_prompt("Pair positive", category="Night", rating=4) + pair = db.save_prompt_pair("Pair positive", "Noise", title="Zebra pair") + owners = ( + self._save_prompt("Zebra one", category="Night"), + self._save_prompt("Zebra two", category="Night"), + ) + files = { + name: self.media(f"{name}.png", mtime=100 + index) + for index, name in enumerate( + ( + "night_ab", + "night_a", + "day_ab", + "titled", + "percent", + "paired", + "ambiguous", + "zebra_unlinked", + "plain_unlinked", + ) + ) + } + db.link_image_to_prompt(night_ab, str(files["night_ab"])) + # A non-normalized spelling of the same file still compares equal. + db.link_image_to_prompt( + night_a, str(self.output_dir / "sub" / ".." / "night_a.png") + ) + db.link_image_to_prompt(day_ab, str(files["day_ab"])) + db.link_image_to_prompt(titled, str(files["titled"])) + db.link_image_to_prompt(percent, str(files["percent"])) + self.assertEqual(db.get_prompt_pair(pair)["positive_prompt_id"], pair_positive) + image = db.link_image_to_prompt(pair_positive, str(files["paired"])) + db.link_image_to_pair(pair, image, {"source": "fixture"}) + for owner in owners: + db.link_image_to_prompt(owner, str(files["ambiguous"])) + return files + + async def test_positive_count_counts_distinct_linked_positives(self): + files = self._filter_fixtures() + db = self.api.db + night_ab = db.get_prompt_id_for_image(str(files["night_ab"])) + # A second file of one positive still counts that positive once. + again = self.media("again.png", mtime=50) + db.link_image_to_prompt(night_ab, str(again)) + cases = ( + # Six unambiguous positives; the ambiguous file and unlinked count none. + ({}, 10, 6), + # Text, title, tag, pair title (its positive) and the filename match. + ({"q": "ZEBRA"}, 6, 4), + ({"category": "Night"}, 4, 3), + ({"tags": "a,b"}, 3, 2), + ({"q": "plain_unlinked"}, 1, 0), + ({"media": "video"}, 0, 0), + ({"sort": "rating"}, 10, 6), + ({"sort": "name", "category": "Day"}, 3, 3), + ) + for params, total, positives in cases: + with self.subTest(**params): + response, data = await self.get(**params) + self.assertEqual(response.status, 200, data) + self.assertEqual(data["total"], total) + self.assertEqual(data["positive_count"], positives) + # The count covers the whole matching set. A filter or the rating + # sort reads the links on every page, so its cursor pages carry it too. + for params in ({"q": "ZEBRA"}, {"sort": "rating"}): + with self.subTest(cursor=params): + response, first = await self.get(limit=2, **params) + response, second = await self.get( + limit=2, cursor=first["next_cursor"], **params + ) + self.assertEqual(response.status, 200, second) + self.assertEqual(second["positive_count"], first["positive_count"]) + # A plain cursor page skips the read and omits the key (absent, not + # null): the frontend keeps the first page's value. + for params in ({}, {"sort": "oldest"}, {"sort": "name"}): + with self.subTest(plain_cursor=params): + response, first = await self.get(limit=2, **params) + self.assertEqual(first["positive_count"], 6) + response, second = await self.get( + limit=2, cursor=first["next_cursor"], **params + ) + self.assertEqual(response.status, 200, second) + self.assertNotIn("positive_count", second) + self.assertEqual(second["total"], 10) + # One SELECT counts the positives; the page's link summary is the other. + statements = [] + original = db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + with patch.object(db.model, "get_connection", side_effect=trace_connection): + response, data = await self.get() + self.assertEqual(response.status, 200) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + self.assertEqual(len(selects), 3, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) + _, first = await self.get(limit=2) + statements.clear() + with patch.object(db.model, "get_connection", side_effect=trace_connection): + response, data = await self.get(limit=2, cursor=first["next_cursor"]) + self.assertEqual(response.status, 200) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + # The page's link summary and its favorites read only. + self.assertEqual(len(selects), 2, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) + self.assertNotIn("gi.file_path, p.rating", selects[0]) + # Without filters the listing outlives a database failure; the count + # is then unknown rather than wrong. + with patch.object( + db, "get_image_link_matches", side_effect=RuntimeError("database gone") + ): + response, data = await self.get() + self.assertEqual(response.status, 200, data) + self.assertEqual(data["total"], 10) + self.assertIsNone(data["positive_count"]) + self.assertNotIn("database gone", json.dumps(data)) + # A linked filter or the rating sort needs the read, so its failure is + # an error (filename-only search survives it, see below). + for params in ( + {"category": "Night"}, + {"tags": "a"}, + {"q": "zebra", "category": "Night"}, + {"sort": "rating"}, + ): + with ( + self.subTest(failing=params), + patch.object( + db, "get_image_link_matches", side_effect=RuntimeError("gone") + ), + ): + response, data = await self.get(**params) + self.assertEqual(response.status, 500, data) + self.assertNotIn("images", data) + + async def test_filename_search_survives_a_database_failure(self): + for name, mtime in (("zebra-1.png", 300), ("zebra-2.png", 200)): + self.media(name, mtime) + self.media("other.png", 100) + with ( + patch.object( + self.api.db, + "get_image_link_matches", + side_effect=RuntimeError("database gone"), + ), + self.assertLogs(self.api.logger, "WARNING") as logged, + ): + response, first = await self.get(q="ZEBRA", limit=1) + self.assertEqual(response.status, 200, first) + response, second = await self.get( + q="ZEBRA", limit=1, cursor=first["next_cursor"] + ) + self.assertEqual(response.status, 200, second) + # Filename matches only; the count of linked positives is unknown. + self.assertEqual(self._names(first), ["zebra-1.png"]) + self.assertEqual(self._names(second), ["zebra-2.png"]) + for page in (first, second): + self.assertEqual(page["total"], 2) + self.assertIsNone(page["positive_count"]) + self.assertNotIn("database gone", json.dumps(page)) + self.assertTrue( + all(record.exc_info for record in logged.records), logged.output + ) + + async def test_search_matches_filename_prompt_text_and_tags(self): + from database.models import normalize_image_path + + files = self._filter_fixtures() + _, data = await self.get(q="ZEBRA", sort="name") + # Text, stored title, tag, pair title and filename; never an ambiguous link. + self.assertEqual( + self._names(data), + [ + "day_ab.png", + "night_ab.png", + "paired.png", + "titled.png", + "zebra_unlinked.png", + ], + ) + self.assertEqual(data["total"], 5) + _, data = await self.get(q="0% w") + self.assertEqual(self._names(data), ["percent.png"]) + _, data = await self.get(q="%") + self.assertEqual(self._names(data), ["percent.png"]) + _, data = await self.get(q="harbor") + self.assertEqual(self._names(data), ["night_a.png"]) + self.assertEqual( + self.api.db.get_image_paths_matching("zebra", None, []), + { + normalize_image_path(str(files[name])) + for name in ("night_ab", "day_ab", "titled", "paired") + }, + ) + + async def test_category_and_tags_filter_linked_images_only(self): + self._filter_fixtures() + cases = ( + ({"category": "Night"}, ["night_ab.png", "night_a.png", "paired.png"]), + ({"category": "night"}, []), + ({"tags": "a,b"}, ["night_ab.png", "day_ab.png"]), + ({"tags": " a , b ,"}, ["night_ab.png", "day_ab.png"]), + ({"category": "Night", "tags": "a,b"}, ["night_ab.png"]), + ({"category": "Night", "q": "harbor"}, ["night_a.png"]), + ({"category": "Night", "q": "paired"}, ["paired.png"]), + ({"category": "Day", "q": "zebra"}, ["day_ab.png", "titled.png"]), + ({"category": ""}, None), + ({"category": " Night "}, ["night_ab.png", "night_a.png", "paired.png"]), + ({"category": " "}, None), + ) + _, everything = await self.get(sort="name") + for params, expected in cases: + with self.subTest(**params): + response, data = await self.get(sort="name", **params) + self.assertEqual(response.status, 200, data) + if expected is None: + self.assertEqual(data["images"], everything["images"]) + else: + self.assertEqual(self._names(data), sorted(expected)) + self.assertEqual(data["total"], len(expected)) + for params in ( + {"category": "x" * 201}, + {"tags": ",".join(f"t{i}" for i in range(31))}, + {"tags": "x" * 201}, + ): + response, data = await self.get(**params) + self.assertEqual(response.status, 400, params) + # The filters read one SELECT besides the page's link summary. + statements = [] + original = self.api.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + with patch.object( + self.api.db.model, "get_connection", side_effect=trace_connection + ): + response, _ = await self.get( + category="Night", tags="a", q="zebra", sort="rating" + ) + self.assertEqual(response.status, 200) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + self.assertEqual(len(selects), 3, selects) + self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) + + async def test_rating_sort_orders_linked_then_unlinked(self): + self._filter_fixtures() + response, data = await self.get(sort="rating") + self.assertEqual(response.status, 200, data) + self.assertEqual( + self._names(data), + [ + "night_ab.png", # 5 + "paired.png", # 4, the pair's positive + "day_ab.png", # 3 + "titled.png", # 1 + # linked without a rating, newest first + "percent.png", + "night_a.png", + # unlinked (and ambiguous) last, newest first + "plain_unlinked.png", + "zebra_unlinked.png", + "ambiguous.png", + ], + ) + _, first = await self.get(sort="rating", limit=4) + _, rest = await self.get(sort="rating", limit=10, cursor=first["next_cursor"]) + self.assertEqual(first["images"] + rest["images"], data["images"]) + _, night = await self.get(sort="rating", category="Night") + self.assertEqual( + self._names(night), ["night_ab.png", "paired.png", "night_a.png"] + ) + + async def test_cursor_rejects_changed_filters(self): + self._filter_fixtures() + _, first = await self.get(category="Night", limit=1) + token = first["next_cursor"] + self.assertIsInstance(token, str) + response, _ = await self.get(category="Night", limit=1, cursor=token) + self.assertEqual(response.status, 200) + for params in ( + {"category": "Day"}, + {}, + {"category": "Night", "tags": "a"}, + {"category": "Night", "q": "zebra"}, + {"category": "Night", "sort": "rating"}, + ): + with self.subTest(**params): + response, _ = await self.get(cursor=token, **params) + self.assertEqual(response.status, 400) + _, rated = await self.get(sort="rating", limit=1) + response, _ = await self.get(sort="rating", cursor=rated["next_cursor"]) + self.assertEqual(response.status, 200) + response, _ = await self.get(sort="newest", cursor=rated["next_cursor"]) + self.assertEqual(response.status, 400) + _, tagged = await self.get(tags="a,b", limit=1) + response, _ = await self.get(tags="b,a", cursor=tagged["next_cursor"]) + self.assertEqual(response.status, 200) + response, _ = await self.get(tags="a", cursor=tagged["next_cursor"]) + self.assertEqual(response.status, 400) + for params in ({"sort": "rated"}, {"sort": "rating-desc"}, {"media": "x"}): + response, _ = await self.get(**params) + self.assertEqual(response.status, 400, params) + + async def test_windows_style_link_paths_compare_normalized(self): + import ntpath + + def windows_normalize(path): + return ntpath.normcase(ntpath.normpath(path)) + + prompt = self._save_prompt("Case words", category="Night", rating=2) + listed = self.media("Sub/Night.PNG") + self.media("Sub/other.png") + stored = str(listed).upper().replace("/", "\\") + with ( + patch("database.operations.normalize_image_path", windows_normalize), + patch("py.api.images.normalize_image_path", windows_normalize), + ): + self.api.db.link_image_to_prompt(prompt, stored) + _, data = await self.get(category="Night") + self.assertEqual(self._names(data), ["Night.PNG"]) + _, data = await self.get(q="case words") + self.assertEqual(self._names(data), ["Night.PNG"]) + _, data = await self.get(sort="rating") + self.assertEqual(self._names(data), ["Night.PNG", "other.png"]) + + def _unlisted_link_rows(self, count): + """*count* linked rows for files outside the listing, written in bulk.""" + prompt = self._save_prompt("Elsewhere", category="Night", rating=5) + with self.api.db.model.get_connection() as conn: + conn.executemany( + "INSERT INTO generated_images(prompt_id,image_path,filename,file_path)" + " VALUES(?,?,?,?)", + ( + (prompt, f"/elsewhere/{i}.png", f"{i}.png", f"/elsewhere/{i}.png") + for i in range(count) + ), + ) + + async def test_rating_sort_reads_only_listed_paths_past_the_old_cap(self): + from database.models import normalize_image_path + + files = self._filter_fixtures() + # More rows than the former 500,000-row cap, none of them listed. + await asyncio.get_running_loop().run_in_executor( + None, self._unlisted_link_rows, 500_001 + ) + response, data = await self.get(sort="rating", limit=3) + self.assertEqual(response.status, 200, data) + self.assertEqual( + self._names(data), ["night_ab.png", "paired.png", "day_ab.png"] + ) + response, data = await self.get(category="Night", sort="name") + self.assertEqual(response.status, 200, data) + self.assertEqual( + self._names(data), ["night_a.png", "night_ab.png", "paired.png"] + ) + wanted = {normalize_image_path(str(files["night_ab"])), "/elsewhere/7.png"} + elsewhere = self.api.db.get_prompt_id_for_image("/elsewhere/7.png") + self.assertEqual( + self.api.db.get_image_link_matches( + "", "Night", [], every_link=True, wanted=wanted + ), + { + normalize_image_path(str(files["night_ab"])): ( + True, + 5, + self.api.db.get_prompt_id_for_image(str(files["night_ab"])), + ), + "/elsewhere/7.png": (True, 5, elsewhere), + }, + ) + + def _budget_fixture(self, files, links): + """*files* listed paths and *links* linked rows, written in bulk.""" + from database.models import normalize_image_path + + listed = [ + (self.output_dir / f"day-{i % 50:02}" / f"img_{i:05}.png", 1_000.0 + i) + for i in range(files) + ] + prompts = min(2_000, links) + with self.api.db.model.get_connection() as conn: + conn.executemany( + "INSERT INTO tags(id,name) VALUES(?,?)", + [(i + 1, f"tag-{i}") for i in range(20)], + ) + conn.executemany( + "INSERT INTO prompts(id,text,category,rating,hash) VALUES(?,?,?,?,?)", + [ + ( + i + 1, + f"prompt {i} " + "lantern " * (i % 7) + "harbor " * 20, + f"cat-{i % 10}", + (i % 6) or None, + f"hash-{i}", + ) + for i in range(prompts) + ], + ) + conn.executemany( + "INSERT OR IGNORE INTO prompt_tags(prompt_id,tag_id) VALUES(?,?)", + [(i + 1, t) for i in range(prompts) for t in (i % 20 + 1, 1)], + ) + conn.executemany( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + [(i + 1, f"Title {i}") for i in range(0, prompts, 3)], + ) + conn.executemany( + "INSERT INTO generated_images(prompt_id,image_path,filename,file_path)" + " VALUES(?,?,?,?)", + [ + ( + i % prompts + 1, + str(listed[i * 2][0]), + listed[i * 2][0].name, + normalize_image_path(str(listed[i * 2][0])), + ) + for i in range(links) + ], + ) + return listed + + _BUDGET_REQUESTS = ( + # First, so the timed loop also measures a cursor page of each. + {}, + {"sort": "oldest"}, + {"q": "lantern"}, + {"q": "img_0123"}, + {"q": "Title 1"}, + {"category": "cat-3"}, + {"tags": "tag-0,tag-2"}, + {"sort": "rating"}, + {"sort": "rating", "category": "cat-4"}, + {"q": "harbor", "category": "cat-5", "tags": "tag-5"}, + {"sort": "name", "q": "tag-7"}, + {"subfolder": "day-07", "q": "lantern"}, + ) + + @staticmethod + def _stub_page(page): + return [{"filename": entry[0].name} for entry in page] + + async def test_filter_work_is_bounded(self): + """The budget's structure, counted rather than timed (runs everywhere).""" + import math + + from py.api import images as images_module + from database.operations import link_filter_query + + listed = await asyncio.get_running_loop().run_in_executor( + None, self._budget_fixture, 5_000, 2_000 + ) + scans = [listed] + built, keyed = [], [] + real_listed, real_key = images_module._listed, images_module._workspace_key + + def counting_listed(*args): + built.append(args[0]) + return real_listed(*args) + + def counting_key(entry, order): + keyed.append(order) + return real_key(entry, order) + + statements = [] + original = self.api.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + async def scan(_root): + return scans[-1] + + def selects(): + found = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + statements.clear() + return found + + with ( + patch.object(self.api, "_get_gallery_files", side_effect=scan), + patch.object(self.api, "_format_workspace_page", self._stub_page), + patch.object(images_module, "_listed", counting_listed), + patch.object(images_module, "_workspace_key", counting_key), + patch.object( + self.api.db.model, "get_connection", side_effect=trace_connection + ), + ): + response, _ = await self.get(limit=10) + self.assertEqual(response.status, 200) + self.assertEqual(len(built), len(listed)) + # No filter still reads the linked positives once, for the count. + self.assertEqual(len(selects()), 1) + newest = self.api._workspace_order_cache[((0, str(self.output_dir)),)] + newest = newest[1]["newest"] + for params in self._BUDGET_REQUESTS: + with self.subTest(**params): + built.clear() + response, data = await self.get(limit=10, **params) + self.assertEqual(response.status, 200, data) + self.assertEqual(built, []) + self.assertEqual(len(selects()), 1) + self.assertIs( + self.api._workspace_order_cache[((0, str(self.output_dir)),)][1][ + "newest" + ], + newest, + ) + bound = math.ceil(math.log2(len(listed))) + 3 + for order in ("newest", "name", "name-desc", "oldest", "rating"): + with self.subTest(cursor=order): + await self.get(sort=order, limit=10) # builds a lazy order + _, first = await self.get(sort=order, limit=10) + keyed.clear() + selects() + response, page = await self.get( + sort=order, limit=10, cursor=first["next_cursor"] + ) + self.assertEqual(response.status, 200) + self.assertEqual(len(page["images"]), 10) + self.assertLessEqual(len(keyed), bound) + # A plain cursor page reads no links; rating needs them. + plain = order != "rating" + self.assertEqual(len(selects()), 0 if plain else 1) + self.assertEqual("positive_count" in page, not plain) + key = ((0, str(self.output_dir)),) + # A TTL rescan that found the same files rebuilds nothing. + built.clear() + scans.append(list(listed)) + response, data = await self.get(limit=10) + self.assertEqual(response.status, 200) + self.assertEqual(built, []) + self.assertIs(self.api._workspace_order_cache[key][1]["newest"], newest) + # One changed file is the only one rebuilt. + added = (self.output_dir / "day-00" / "new.png", 9_999_999.0) + scans.append([added, *listed[1:]]) + response, data = await self.get(limit=10) + self.assertEqual(built, [added[0]]) + self.assertEqual(data["total"], len(listed)) + self.assertEqual(data["images"][0]["filename"], "new.png") + # The monitor invalidates the gallery when it sees a new file: the + # next request rebuilds only that file's values, and search finds it. + self.api.invalidate_gallery_cache() + self.assertNotEqual(self.api._workspace_listed_cache, {}) + monitored = (self.output_dir / "day-01" / "monitored.png", 9_999_999.5) + scans.append([monitored, *scans[-1]]) + built.clear() + response, data = await self.get(limit=10, q="monitored") + self.assertEqual(response.status, 200) + self.assertEqual(built, [monitored[0]]) + self.assertEqual(self._names(data), ["monitored.png"]) + # A roots change drops every derived value. + with patch("py.config.GalleryConfig.MONITORING_DIRECTORIES", []): + self.assertTrue(self.api._apply_gallery_roots([str(self.output_dir)])) + self.assertEqual(self.api._workspace_listed_cache, {}) + self.assertEqual(self.api._workspace_order_cache, {}) + # The per-row ambiguity probe is an index search, never a scan. + sql, params = link_filter_query("lantern", "cat-3", ["tag-1"], every_link=True) + with self.api.db.model.get_connection() as conn: + plan = [row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)] + self.assertIn( + "SEARCH other USING COVERING INDEX idx_image_file_path (file_path=?)", plan + ) + self.assertFalse([step for step in plan if step.startswith("SCAN other")]) + + @unittest.skipUnless( + PERF_BUDGET and not TRACED, + "wall-clock budget: set PROMPTMANAGER_PERF_BUDGET=1 and run untraced", + ) + async def test_filter_budget(self): + """p95 < 150 ms at 50,000 files / 20,000 links (TESTING.md budgets).""" + import math + import time + + # Built off the event loop: the bulk inserts take a few hundred ms. + listed = await asyncio.get_running_loop().run_in_executor( + None, self._budget_fixture, 50_000, 20_000 + ) + rounds, plain = [], {"first": [], "cursor": []} + with ( + patch.object( + self.api, "_get_gallery_files", AsyncMock(return_value=listed) + ), + patch.object(self.api, "_format_workspace_page", self._stub_page), + ): + # The first request after a scan builds the cached sort orders. + started = time.perf_counter() + response, _ = await self.get(sort="rating") + cold = time.perf_counter() - started + self.assertEqual(response.status, 200) + for _round in range(3): + timings = [] + for index in range(20): + requests = self._BUDGET_REQUESTS + shape = requests[index % len(requests)] + params = dict(shape, limit=100) + if index >= len(requests): + _, first = await self.get(**params) + if first["next_cursor"]: + params["cursor"] = first["next_cursor"] + started = time.perf_counter() + response, data = await self.get(**params) + timings.append(time.perf_counter() - started) + if shape.get("sort", "newest") in ("newest", "oldest") and not ( + set(shape) - {"sort"} + ): + # Plain first pages read the links for the count; + # plain cursor pages read none. + plain["cursor" if "cursor" in params else "first"].append( + timings[-1] + ) + self.assertEqual(response.status, 200, (params, data)) + timings.sort() + rounds.append(timings) + p95s = [timings[math.ceil(0.95 * len(timings)) - 1] for timings in rounds] + best = min(p95s) + plain_ms = { + kind: sorted(round(t * 1000, 1) for t in found) + for kind, found in plain.items() + } + print( + "\nworkspace filter budget: best p95=" + f"{best * 1000:.1f} ms, rounds=" + f"{[round(p * 1000, 1) for p in p95s]} ms, cold={cold * 1000:.1f} ms, " + f"plain listing ms={plain_ms}", + file=sys.stderr, + ) + self.assertLess( + best, 0.150, [[round(t * 1000, 1) for t in timings] for timings in rounds] + ) + class TestWorkspaceCursorEncoding(unittest.TestCase): def test_ntfs_unicode_basename_cursor_is_bounded_and_round_trips(self): diff --git a/tests/test_card_seed.py b/tests/test_card_seed.py new file mode 100644 index 0000000..466c92c --- /dev/null +++ b/tests/test_card_seed.py @@ -0,0 +1,236 @@ +"""Card facts from header records: the shared legacy parser, bounded and exact.""" + +import json +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils.card_seed import ( # noqa: E402 + MAX_RECORD_BYTES, + card_seed, + header_params, + parsed_positive, +) +from utils.generation_params import ( # noqa: E402 + GenerationParamsError, + extract_from_records, + extract_generation_params, +) +from utils.inspection_response import ( # noqa: E402 + MAX_PROJECTION_DEPTH, + project_graph, +) + +try: + from tests.generation_params_cases import CASES, encoder, ksampler # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from generation_params_cases import CASES, encoder, ksampler # noqa: E402 + + +def graph(*samplers, extra=None): + nodes = {"6": encoder("a lantern at dusk"), "7": encoder("blurry")} + for index, sampler in enumerate(samplers): + nodes[str(10 + index)] = sampler + nodes.update(extra or {}) + return {"prompt": json.dumps(nodes)} + + +def seeded(seed, **settings): + return ksampler(("6", 0), ("7", 0), seed=seed, **settings) + + +def nested_workflow(depth=900): + """A small workflow whose node and link ids nest ``depth`` arrays deep.""" + nested = "[" * depth + "1" + "]" * depth + return ( + '{"nodes":[{"id":%s,"type":"KSampler","widgets_values":[1]}],' + '"links":[[%s,1,0,2,0]]}' % (nested, nested) + ) + + +def info_seed(info): + return card_seed(header_params(info)) + + +class CardSeed(unittest.TestCase): + def test_card_seed_single_sampler(self): + info = graph(seeded(42001)) + self.assertEqual(info_seed(info), "42001") + self.assertEqual(parsed_positive(header_params(info)), "a lantern at dusk") + + def test_card_seed_two_different_seeds_follows_shared_parser(self): + # Ruling: no separate heuristic; the card shows what the legacy parser + # (and so the inspector) shows for several samplers. + info = graph(seeded(111), seeded(222)) + expected = extract_from_records(info["prompt"], None)["seed"] + self.assertIn(expected, (111, 222)) + self.assertEqual(info_seed(info), str(expected)) + + def test_card_seed_linked_input_follows_shared_parser(self): + primitive = {"class_type": "PrimitiveNode", "inputs": {"value": 555}} + info = graph(seeded(["20", 0]), extra={"20": primitive}) + self.assertEqual(info_seed(info), "555") + dangling = graph(seeded(["99", 0])) + self.assertEqual( + extract_from_records(dangling["prompt"], None)["seed"], "Unknown" + ) + self.assertIsNone(info_seed(dangling)) + + def test_card_seed_advanced_sampler_noise_seed(self): + advanced = seeded(None) + advanced["class_type"] = "KSamplerAdvanced" + del advanced["inputs"]["seed"] + advanced["inputs"]["noise_seed"] = 77 + self.assertEqual(info_seed(graph(advanced)), "77") + + def test_card_seed_oversized_record_is_none(self): + padding = "x" * (MAX_RECORD_BYTES + 1) + big_prompt = graph(seeded(1), extra={"30": encoder(padding)}) + self.assertGreater(len(big_prompt["prompt"]), MAX_RECORD_BYTES) + self.assertIsNone(header_params(big_prompt)) + big_workflow = {**graph(seeded(1)), "workflow": json.dumps({"x": padding})} + self.assertIsNone(header_params(big_workflow)) + # Bytes, not characters: two-byte text over the budget is oversized. + wide = "é" * (MAX_RECORD_BYTES // 2 + 1) + self.assertLessEqual(len(wide), MAX_RECORD_BYTES) + self.assertIsNone(header_params({"prompt": json.dumps({"a": wide})})) + self.assertEqual(info_seed(graph(seeded(1))), "1") + + def test_records_past_the_inspection_projection_bound_give_nothing(self): + # The inspector hands the browser parser only graphs its projection + # parsed; a deeper graph is "limited" there, so the card skips it too. + self.assertEqual( + project_graph(nested_workflow(), "workflow")["status"], "limited" + ) + self.assertIsNone(header_params({"workflow": nested_workflow()})) + # A limited workflow beside a parsed prompt leaves the prompt's facts. + both = header_params({**graph(seeded(1)), "workflow": nested_workflow()}) + self.assertEqual(card_seed(both), "1") + + def test_depth_bound_is_the_inspectors(self): + def deep(levels): + nested = "[" * levels + "1" + "]" * levels + extra = '{"class_type":"Note","inputs":{"v":%s}}' % nested + text = graph(seeded(3))["prompt"] + return {"prompt": text[:-1] + ',"30":' + extra + "}"} + + # root object, node object, inputs object: three levels before "v". + inside = MAX_PROJECTION_DEPTH - 3 + self.assertEqual(info_seed(deep(inside)), "3") + self.assertIsNone(header_params(deep(inside + 1))) + + def test_strict_projection_rules_apply(self): + duplicate = '{"10": {"class_type": "KSampler", "inputs": {}}, "10": {}}' + self.assertEqual(project_graph(duplicate, "prompt")["status"], "malformed") + self.assertIsNone(header_params({"prompt": duplicate})) + + def test_record_size_is_utf8_bytes_at_both_bounds(self): + frame = '{"a":"%s"}' + ascii_at = frame % ("x" * (MAX_RECORD_BYTES - 8)) + self.assertEqual(len(ascii_at), MAX_RECORD_BYTES) + self.assertIsNotNone(header_params({"prompt": ascii_at})) + four_byte = frame % ("\U0001f600" * ((MAX_RECORD_BYTES - 8) // 4)) + self.assertLessEqual(len(four_byte.encode()), MAX_RECORD_BYTES) + self.assertIsNotNone(header_params({"prompt": four_byte})) + padded = four_byte + " " * (MAX_RECORD_BYTES - len(four_byte.encode()) + 1) + self.assertIsNone(header_params({"prompt": padded})) + + def test_placeholder_and_non_integer_seeds_are_none(self): + for value in ("randomize", " 12", True, 1.5, 7.0, 1e21, None): + with self.subTest(seed=value): + self.assertIsNone(info_seed(graph(seeded(value)))) + self.assertIsNone(card_seed({"seed": "Unknown"})) + self.assertIsNone(card_seed({"seed": 7.0})) + self.assertEqual(info_seed(graph(seeded(-3))), "-3") + + def test_digit_string_seed_is_what_the_inspector_reads(self): + # The projection carries big integers as text and keeps text as text; + # the inspector shows a digit string as the seed, and so does the card. + self.assertEqual(info_seed(graph(seeded("123"))), "123") + + def test_seed_beyond_double_precision_is_the_recorded_integer(self): + # JSON.parse would round it; the projection carries the record's digits. + case = CASES["seed-beyond-double-precision"] + rounded = extract_from_records(case["prompt"], None)["seed"] + self.assertEqual(rounded, float(4974052037430000001)) + self.assertEqual(info_seed({"prompt": case["prompt"]}), "4974052037430000001") + self.assertEqual(info_seed(graph(seeded(-(2**60) - 1))), str(-(2**60) - 1)) + + def test_seed_is_always_a_decimal_string(self): + # JSON.parse rounds numbers past 2**53, so every seed travels as text. + for value in (0, 7, 2**53, 2**53 + 1, 2**64, -(2**63)): + with self.subTest(seed=value): + self.assertEqual(card_seed({"seed": value}), str(value)) + self.assertEqual(card_seed({"seed": "18446744073709551616"}), str(2**64)) + + def test_digit_string_seed_is_bounded(self): + # 2**64 has 20 digits; a longer string is no seed, and an unbounded one + # must not reach int() (Python's int max-str-digits raises ValueError). + self.assertEqual(card_seed({"seed": "9" * 20}), "9" * 20) + self.assertEqual(card_seed({"seed": "-" + "9" * 20}), "-" + "9" * 20) + self.assertIsNone(card_seed({"seed": "1" * 21})) + self.assertIsNone(card_seed({"seed": "9" * 5000})) + self.assertIsNone(info_seed(graph(seeded("9" * 5000)))) + + def test_seed_conversion_error_gives_none(self): + class Huge(int): + def __str__(self): + raise ValueError("Exceeds the limit for integer string conversion") + + self.assertIsNone(card_seed({"seed": Huge(5)})) + + def test_workflow_only_record(self): + case = next( + case + for case in CASES.values() + if case["prompt"] is None + and case["workflow"] is not None + and isinstance(extract_from_records(None, case["workflow"])["seed"], int) + ) + expected = extract_from_records(None, case["workflow"])["seed"] + self.assertEqual(info_seed({"workflow": case["workflow"]}), str(expected)) + + def test_records_the_parser_cannot_use_give_nothing(self): + self.assertIsNone(header_params({})) + self.assertIsNone(header_params(None)) + self.assertIsNone(header_params({"parameters": "a cat\nSteps: 20, Seed: 9"})) + self.assertIsNone(header_params({"prompt": b"{}"})) + self.assertIsNone(card_seed(None)) + self.assertIsNone(parsed_positive(None)) + # The browser parser throws here; the card shows nothing. + thrown = CASES["null-workflow-node"] + self.assertIsNone(header_params(thrown)) + unparsable = header_params(CASES["unparsable-records"]) + self.assertIsNone(card_seed(unparsable)) + self.assertIsNone(parsed_positive(unparsable)) + + def test_every_shared_case_reads_the_projected_graphs(self): + for name, case in CASES.items(): + with self.subTest(case=name): + info = {k: v for k, v in case.items() if v is not None} + params = header_params(info) + graphs = {key: project_graph(text, key) for key, text in info.items()} + values = { + key: graph["value"] + for key, graph in graphs.items() + if graph["status"] == "parsed" + } + if not values: + self.assertIsNone(params) + continue + try: + expected = extract_generation_params(values) + except GenerationParamsError: + self.assertIsNone(params) + continue + self.assertEqual(params, expected) + # Prompt text is the same as from a plain parse of the records. + if len(values) == len(info): + raw = extract_from_records(case["prompt"], case["workflow"]) + self.assertEqual(params["positivePrompt"], raw["positivePrompt"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_database.py b/tests/test_database.py index e7b57a0..36bad16 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1432,3 +1432,167 @@ class TestImagePaging(DatabaseTestCase): self.assertEqual(len(second), 2) self.assertNotEqual({img["id"] for img in first}, {img["id"] for img in second}) self.assertEqual(len(self.db.search_images_by_prompt("paged")), 5) + + +class TestImageLinkSummaries(DatabaseTestCase): + """Batched card summaries follow get_image_prompt_info's matching rules.""" + + def _statements(self, call): + statements = [] + original = self.db.model.get_connection + + def traced(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + try: + with patch.object(self.db.model, "get_connection", side_effect=traced): + result = call() + finally: + original().set_trace_callback(None) + selects = [s for s in statements if s.lstrip().upper().startswith("SELECT")] + return result, selects + + def test_summary_status_matches_image_prompt_info_for_every_rule(self): + owner = self._save("Owner prompt\nsecond line", category="Cat", rating=4) + rival = self._save("Rival prompt") + exact = os.path.join("out", "one", "same.png") + other = os.path.join("out", "two", "same.png") + unique = os.path.join("out", "three", "unique.png") + shared = os.path.join("out", "shared.png") + self.db.link_image_to_prompt(owner, exact) + self.db.link_image_to_prompt(rival, other) + self.db.link_image_to_prompt(owner, unique) + self.db.link_image_to_prompt(owner, shared) + self.db.link_image_to_prompt(rival, shared) + probes = [ + exact, + os.path.join("out", "x", "..", "one", "same.png"), + os.path.join("elsewhere", "same.png"), + os.path.join("elsewhere", "unique.png"), + shared, + "missing.png", + ] + for fallback in (True, False): + summaries = self.db.get_image_link_summaries(probes, fallback) + self.assertEqual(set(summaries), set(probes)) + for probe in probes: + with self.subTest(probe=probe, fallback=fallback): + info = self.db.get_image_prompt_info( + probe, fallback, report_ambiguity=True + ) + summary = summaries[probe] + if info is None: + expected = "none" + elif info.get("association_status") == "ambiguous": + expected = "ambiguous" + else: + expected = "linked" + self.assertEqual(summary["status"], expected) + if expected == "linked": + self.assertEqual(summary["positive_id"], info["prompt_id"]) + else: + self.assertEqual( + summary, + { + "status": expected, + "positive_id": None, + "title": None, + "title_source": None, + "category": None, + "rating": None, + }, + ) + linked = self.db.get_image_link_summaries([unique])[unique] + self.assertEqual( + linked, + { + "status": "linked", + "positive_id": owner, + "title": "Owner prompt", + "title_source": "text", + "category": "Cat", + "rating": 4, + }, + ) + + def test_summary_title_prefers_stored_title_then_first_line(self): + long_line = "word " * 30 + titled = self._save("Untitled text") + plain = self._save(f" {long_line}\nsecond") + blank = self._save("Blank title text") + with self.db.model.get_connection() as conn: + conn.executemany( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + [(titled, "Stored title"), (blank, " \u3000 ")], + ) + paths = {} + for name, prompt in (("titled", titled), ("plain", plain), ("blank", blank)): + paths[name] = os.path.join("out", f"{name}.png") + self.db.link_image_to_prompt(prompt, paths[name]) + summaries = self.db.get_image_link_summaries(list(paths.values())) + sources = { + name: (summaries[path]["title"], summaries[path]["title_source"]) + for name, path in paths.items() + } + self.assertEqual(sources["titled"], ("Stored title", "stored")) + # Cut like the Library's title(): 80 characters, then an ellipsis. + self.assertEqual(sources["plain"], (long_line.strip()[:80] + "\u2026", "text")) + self.assertEqual(sources["blank"], ("Blank title text", "text")) + + def test_display_title_matches_library_title_cut(self): + from database.library import display_title + + exact = "x" * 80 + self.assertEqual(display_title((), exact), exact) + self.assertEqual(display_title((), exact + "y"), exact + "\u2026") + self.assertEqual(display_title((), " short \nnext"), "short") + self.assertIsNone(display_title((" ",), " \n")) + self.assertEqual(display_title(("Kept ",), "text"), "Kept ") + + def test_summary_uses_one_select_per_200_paths(self): + prompt = self._save("Batch") + paths = [os.path.join("out", f"{index}.png") for index in range(201)] + for path in paths[:5]: + self.db.link_image_to_prompt(prompt, path) + result, selects = self._statements( + lambda: self.db.get_image_link_summaries(paths[:200]) + ) + self.assertEqual(len(selects), 1) + self.assertEqual(len(result), 200) + self.assertEqual( + [result[path]["status"] for path in paths[:6]], ["linked"] * 5 + ["none"] + ) + _, selects = self._statements(lambda: self.db.get_image_link_summaries(paths)) + self.assertEqual(len(selects), 2) + _, selects = self._statements(lambda: self.db.get_image_link_summaries([])) + self.assertEqual(selects, []) + + def test_windows_style_paths_normalize(self): + import ntpath + + def windows_normalize(path): + return ntpath.normcase(ntpath.normpath(path)) + + prompt = self._save("Windows prompt") + stored = r"C:\ComfyUI\Output\Sub\Image.PNG" + with patch("database.operations.normalize_image_path", windows_normalize): + self.db.link_image_to_prompt(prompt, stored) + probes = [ + stored, + r"c:\comfyui\output\SUB\image.png", + r"C:\ComfyUI\Output\x\..\Sub\Image.PNG", + "C:/ComfyUI/Output/Sub/Image.PNG", + r"D:\ComfyUI\Output\Sub\Other.png", + ] + summaries = self.db.get_image_link_summaries(probes, False) + for probe in probes: + with self.subTest(probe=probe): + info = self.db.get_image_prompt_info(probe, False) + expected = "linked" if info else "none" + self.assertEqual(summaries[probe]["status"], expected) + self.assertEqual( + [summaries[probe]["status"] for probe in probes], + ["linked"] * 4 + ["none"], + ) diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 68c68f5..dc7436a 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -428,6 +428,72 @@ class TestRestoreFromFile(BackupTestCase): PromptModel.verify_database_file(dst, require_library_state=True)[0] ) + def test_restore_old_backup_without_favorites(self): + self.db.set_image_favorite("output", "live.png") + src = self._make_source_db(["predates favorites"]) + with closing(sqlite3.connect(src)) as conn: + conn.execute("DROP TABLE image_favorites") + conn.commit() + self.assertTrue(PromptModel.verify_database_file(src)[0]) + + self.db.model.restore_from_file(src) + + self.assertEqual(_texts(self.path), ["predates favorites"]) + self.assertEqual(self.db.get_image_favorites("output", ["live.png"]), set()) + self.assertTrue(self.db.set_image_favorite("output", "new.png")) + ok, reason = PromptModel.verify_database_file( + self.path, require_library_state=True + ) + self.assertTrue(ok, reason) + with closing(sqlite3.connect(src)) as conn: # the source is untouched + self.assertIsNone( + conn.execute( + "SELECT 1 FROM sqlite_master WHERE name='image_favorites'" + ).fetchone() + ) + + def test_restore_preserves_favorites_and_closes_handles(self): + self.db.set_image_favorite("output", "live.png") + with tempfile.TemporaryDirectory() as directory: + src_path = os.path.join(directory, "source.db") + src = PromptDatabase(src_path) + try: + self._save(src, "with favorites") + src.set_image_favorite("output", "sub/kept.png") + src.set_image_favorite("extra/root", "kept.png") + finally: + src.close_all() + # The favorites still sit in the WAL; a backup must carry them. + backup = os.path.join(directory, "backup.db") + reopened = PromptDatabase(src_path) + try: + self.assertTrue(reopened.model.backup_database(backup)) + finally: + reopened.close_all() + + self.db.model.restore_from_file(backup) + + self.assertEqual( + self.db.get_image_favorites("output", ["sub/kept.png", "live.png"]), + {"sub/kept.png"}, + ) + self.assertEqual( + self.db.get_image_favorites("extra/root", ["kept.png"]), {"kept.png"} + ) + PromptModel.close_all_instances() + # Windows refuses to delete a file a restore left open. + assert_closed(self, directory) + backups = list(Path(self.tmpdir).glob("prompts.db.backup_*")) + self.assertEqual(len(backups), 1) + with closing(sqlite3.connect(str(backups[0]))) as conn: + self.assertEqual( + conn.execute( + "SELECT root_key, relative_path FROM image_favorites" + ).fetchall(), + [("output", "live.png")], + ) + assert_closed(self, str(backups[0])) + def test_restore_without_a_live_database_returns_empty_backup_path(self): src = self._make_source_db(["fresh"]) self.db.close_all() diff --git a/tests/test_image_favorites.py b/tests/test_image_favorites.py new file mode 100644 index 0000000..b2d4a2c --- /dev/null +++ b/tests/test_image_favorites.py @@ -0,0 +1,383 @@ +"""Persisted image favorites (spec D7): store, schema and rooted API boundary.""" + +import json +import ntpath +import os +import sqlite3 +import tempfile +import unittest +from contextlib import closing +from pathlib import Path +from unittest.mock import patch + +try: + from test_api_images import ImageAPITestCase, make_png + from open_handles import assert_closed +except ImportError: + from tests.test_api_images import ImageAPITestCase, make_png + from tests.open_handles import assert_closed + +from database.image_favorites import ( + MAX_KEY_LENGTH, + TABLE, + image_favorites_schema_error, + migrate_image_favorites, +) +from database.models import PromptModel +from database.operations import PromptDatabase + +FAVORITE = "/prompt_manager/images/favorite" + + +class TestImageFavoriteApi(ImageAPITestCase): + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + async def send(self, method, body): + response = await self.client.request(method, FAVORITE, json=body) + return response, await response.json() + + async def workspace(self, **params): + response = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace", **params} + ) + self.assertEqual(response.status, 200) + data = await response.json() + return {image["path"]: image for image in data["images"]} + + async def test_favorite_round_trip_and_idempotent(self): + make_png(self.output_dir / "sub" / "a.png") + make_png(self.output_dir / "b.png") + items = await self.workspace() + self.assertIs(items["sub/a.png"]["favorite"], False) + self.assertIs(items["b.png"]["favorite"], False) + + body = {"root_index": 0, "path": "sub/a.png"} + for _ in range(2): # a repeated PUT is a no-op, not an error + response, data = await self.send("PUT", body) + self.assertEqual(response.status, 200) + self.assertEqual( + data, + { + "success": True, + "favorite": True, + "root_index": 0, + "path": "sub/a.png", + }, + ) + items = await self.workspace() + self.assertIs(items["sub/a.png"]["favorite"], True) + self.assertIs(items["b.png"]["favorite"], False) + count = self.api.db.model.get_connection().execute( + f"SELECT COUNT(*) FROM {TABLE}" + ) + self.assertEqual(count.fetchone()[0], 1) + + for _ in range(2): + response, data = await self.send("DELETE", body) + self.assertEqual(response.status, 200) + self.assertTrue(data["success"]) + self.assertIs(data["favorite"], False) + items = await self.workspace() + self.assertIs(items["sub/a.png"]["favorite"], False) + + def rows(self): + return [ + tuple(row) + for row in self.api.db.model.get_connection().execute( + f"SELECT root_key, relative_path FROM {TABLE} ORDER BY 1, 2" + ) + ] + + async def test_favorite_is_keyed_by_root_not_index(self): + second = self.comfy_dir / "second" + make_png(second / "a.png") + make_png(self.output_dir / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + response, data = await self.send("PUT", {"root_index": 1, "path": "a.png"}) + self.assertEqual(response.status, 200) + self.assertNotIn(str(self.comfy_dir), json.dumps(data)) + root_key = Path(os.path.normcase(os.path.realpath(second))).as_posix() + self.assertEqual(self.rows(), [(root_key, "a.png")]) + # The same root at another index keeps its favorite; the other root + # with the same relative path does not inherit it. + self.api._get_all_output_dirs = lambda: [second, self.output_dir] + self.api.invalidate_gallery_cache() + self.assertIs((await self.workspace(root=0))["a.png"]["favorite"], True) + self.assertIs((await self.workspace(root=1))["a.png"]["favorite"], False) + + async def test_roots_sharing_a_folder_name_keep_separate_favorites(self): + # Two roots under different parents, both shown as "output". + with tempfile.TemporaryDirectory() as other_parent: + other = Path(os.path.realpath(other_parent)) / "output" + make_png(other / "a.png") + make_png(self.output_dir / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, other] + response = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace"} + ) + labels = { + item["root_index"]: item["root_dir"] + for item in (await response.json())["images"] + } + self.assertEqual(labels, {0: "output", 1: "output"}) + + await self.send("PUT", {"root_index": 0, "path": "a.png"}) + self.assertIs((await self.workspace(root=0))["a.png"]["favorite"], True) + self.assertIs((await self.workspace(root=1))["a.png"]["favorite"], False) + + await self.send("PUT", {"root_index": 1, "path": "a.png"}) + response, data = await self.send( + "DELETE", {"root_index": 0, "path": "a.png"} + ) + self.assertEqual(response.status, 200) + self.assertNotIn(other_parent, json.dumps(data)) + self.assertIs((await self.workspace(root=0))["a.png"]["favorite"], False) + self.assertIs((await self.workspace(root=1))["a.png"]["favorite"], True) + listing = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace"} + ) + self.assertNotIn(other_parent, await listing.text()) + + async def test_windows_case_only_rename_keeps_the_star(self): + make_png(self.output_dir / "Sub" / "A.png") + with patch("py.api.images._path_flavor", ntpath): + response, data = await self.send( + "PUT", {"root_index": 0, "path": "Sub/A.png"} + ) + self.assertEqual(response.status, 200) + self.assertEqual(data["path"], "Sub/A.png") # the listing's spelling + (self.output_dir / "Sub" / "A.png").rename( + self.output_dir / "Sub" / "a.png" + ) + (self.output_dir / "Sub").rename(self.output_dir / "sub") + self.api.invalidate_gallery_cache() + self.assertIs((await self.workspace())["sub/a.png"]["favorite"], True) + response, data = await self.send( + "DELETE", {"root_index": 0, "path": "SUB/A.PNG"} + ) + self.assertEqual(response.status, 200) + self.assertEqual(data["path"], "sub/a.png") + self.assertIs((await self.workspace())["sub/a.png"]["favorite"], False) + self.assertEqual(self.rows(), []) + + async def test_posix_case_distinct_files_keep_independent_stars(self): + make_png(self.output_dir / "A.png") + make_png(self.output_dir / "a.png") + if len(list(self.output_dir.iterdir())) != 2: + self.skipTest("case-insensitive filesystem") + await self.send("PUT", {"root_index": 0, "path": "A.png"}) + items = await self.workspace() + self.assertIs(items["A.png"]["favorite"], True) + self.assertIs(items["a.png"]["favorite"], False) + + async def test_favorite_rejects_traversal_and_unknown_root(self): + make_png(self.output_dir / "a.png") + make_png(self.comfy_dir / "outside.png") + (self.output_dir / "notes.txt").write_text("not media") + make_png(self.output_dir / "thumbnails" / "a_thumb.png") + cases = [ + ({"root_index": 0, "path": "../outside.png"}, 403), + ({"root_index": 0, "path": "sub/../../outside.png"}, 403), + ({"root_index": 0, "path": str(self.comfy_dir / "outside.png")}, 403), + ({"root_index": 0, "path": "/a.png"}, 403), + ({"root_index": 0, "path": "C:/a.png"}, 403), + ({"root_index": 0, "path": "a.png\0.png"}, 403), + ({"root_index": 0, "path": "notes.txt"}, 403), + ({"root_index": 1, "path": "a.png"}, 400), + ({"root_index": -1, "path": "a.png"}, 400), + ({"root_index": "0", "path": "a.png"}, 400), + ({"root_index": True, "path": "a.png"}, 400), + ({"root_index": 0}, 400), + ({"root_index": 0, "path": ""}, 400), + ({"root_index": 0, "path": 7}, 400), + ({"root_index": 0, "path": "x" * 5000 + ".png"}, 400), + ({"root_index": 0, "path": "missing.png"}, 404), + ({"root_index": 0, "path": "thumbnails/a_thumb.png"}, 404), + ] + for method in ("PUT", "DELETE"): + for body, status in cases: + if method == "DELETE" and status == 404: + continue # unlisting a stale favorite is allowed (below) + with self.subTest(method=method, body=str(body)[:60]): + response, data = await self.send(method, body) + self.assertEqual(response.status, status) + self.assertIs(data["success"], False) + self.assertNotIn(str(self.comfy_dir), data["error"]) + response = await self.client.put(FAVORITE, data=b"[1]") + self.assertEqual(response.status, 400) + response = await self.client.put(FAVORITE, data=b"{not json") + self.assertEqual(response.status, 400) + count = self.api.db.model.get_connection().execute( + f"SELECT COUNT(*) FROM {TABLE}" + ) + self.assertEqual(count.fetchone()[0], 0) + + async def test_delete_clears_a_favorite_whose_file_is_gone(self): + target = make_png(self.output_dir / "gone.png") + await self.send("PUT", {"root_index": 0, "path": "gone.png"}) + target.unlink() + self.api.invalidate_gallery_cache() + response, data = await self.send( + "DELETE", {"root_index": 0, "path": "gone.png"} + ) + self.assertEqual(response.status, 200) + self.assertIs(data["favorite"], False) + count = self.api.db.model.get_connection().execute( + f"SELECT COUNT(*) FROM {TABLE}" + ) + self.assertEqual(count.fetchone()[0], 0) + + async def test_legacy_listing_is_unchanged_by_favorites(self): + make_png(self.output_dir / "a.png") + await self.send("PUT", {"root_index": 0, "path": "a.png"}) + response = await self.client.get("/prompt_manager/images/output") + data = await response.json() + self.assertEqual(len(data["images"]), 1) + self.assertNotIn("favorite", data["images"][0]) + + async def test_toggle_lookup_is_built_once_per_scan(self): + make_png(self.output_dir / "a.png") + files = await self.api._get_gallery_files(self.output_dir) + first = self.api._favorite_listing(self.output_dir, files) + self.assertEqual(first, {"a.png": "a.png"}) + self.assertIs(self.api._favorite_listing(self.output_dir, files), first) + make_png(self.output_dir / "b.png") + self.api.invalidate_gallery_cache() + rescanned = await self.api._get_gallery_files(self.output_dir) + self.assertEqual( + self.api._favorite_listing(self.output_dir, rescanned), + {"a.png": "a.png", "b.png": "b.png"}, + ) + self.api.reset_gallery_roots_cache() + self.assertEqual(self.api._favorite_listing_cache, {}) + + async def test_monitor_invalidation_drops_the_favorites_lookup(self): + make_png(self.output_dir / "a.png") + await self.send("PUT", {"root_index": 0, "path": "a.png"}) + self.assertNotEqual(self.api._favorite_listing_cache, {}) + # The image monitor's invalidation must not keep the old scan alive. + self.api.invalidate_gallery_cache() + self.assertEqual(self.api._favorite_listing_cache, {}) + self.assertIs((await self.workspace())["a.png"]["favorite"], True) + + async def test_put_rejects_a_key_beyond_the_stored_bound(self): + make_png(self.output_dir / "a.png") + too_long = "/" + "r" * MAX_KEY_LENGTH + with patch("py.api.images._favorite_root_key", return_value=too_long): + response, data = await self.send("PUT", {"root_index": 0, "path": "a.png"}) + self.assertEqual(response.status, 403) + self.assertFalse(data["success"]) + self.assertNotIn(too_long, json.dumps(data)) + self.assertEqual(self.rows(), []) + # The live database still validates (a stored over-long key would + # fail startup and restore). + conn = self.api.db.model.get_connection() + self.assertIsNone(image_favorites_schema_error(conn)) + + async def test_favorite_store_failure_keeps_the_listing(self): + make_png(self.output_dir / "a.png") + + def broken(*_args, **_kwargs): + raise sqlite3.OperationalError("database is locked") + + self.api.db.get_image_favorites = broken + items = await self.workspace() + self.assertIs(items["a.png"]["favorite"], False) + self.api.db.set_image_favorite = broken + response, data = await self.send("PUT", {"root_index": 0, "path": "a.png"}) + self.assertEqual(response.status, 500) + self.assertEqual(data, {"success": False, "error": "Favorite update failed"}) + + +class TestImageFavoriteSchema(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.mkdtemp() + self.path = os.path.join(self.tmpdir, "favorites.db") + + def tearDown(self): + PromptModel.close_all_instances() + assert_closed(self, self.tmpdir) + for name in os.listdir(self.tmpdir): + os.unlink(os.path.join(self.tmpdir, name)) + os.rmdir(self.tmpdir) + + def test_store_round_trip_and_bounded_lookup(self): + db = PromptDatabase(self.path) + try: + self.assertTrue(db.set_image_favorite("output", "a.png")) + self.assertTrue(db.set_image_favorite("output", "a.png")) + db.set_image_favorite("other", "a.png") + self.assertEqual( + db.get_image_favorites("output", ["a.png", "b.png"]), {"a.png"} + ) + self.assertEqual(db.get_image_favorites("output", []), set()) + many = [f"{n}.png" for n in range(2000)] + ["a.png"] + self.assertEqual(db.get_image_favorites("output", many), {"a.png"}) + self.assertTrue(db.clear_image_favorite("output", "a.png")) + self.assertFalse(db.clear_image_favorite("output", "a.png")) + self.assertEqual(db.get_image_favorites("output", ["a.png"]), set()) + created = db.model.get_connection().execute( + f"SELECT created_at FROM {TABLE}" + ) + self.assertRegex(created.fetchone()[0], r"^\d{4}-\d\d-\d\dT") + finally: + db.close_all() + + def test_schema_rejects_unsupported_definitions_and_values(self): + variants = { + "missing": None, + "wrong table": f"CREATE TABLE {TABLE} (x TEXT)", + "view": f"CREATE VIEW {TABLE} AS SELECT 1", + "nullable key": ( + f"CREATE TABLE {TABLE} (root_key TEXT, relative_path TEXT, " + "created_at TEXT, PRIMARY KEY(root_key, relative_path))" + ), + } + for label, ddl in variants.items(): + with self.subTest(label): + with closing(sqlite3.connect(":memory:")) as conn: + if ddl: + conn.execute(ddl) + error = image_favorites_schema_error(conn) + if ddl is None: + self.assertIsNone(error) + self.assertIsNotNone( + image_favorites_schema_error(conn, allow_missing=False) + ) + else: + self.assertIsNotNone(error) + with self.assertRaises(sqlite3.DatabaseError): + migrate_image_favorites(conn) + with closing(sqlite3.connect(":memory:")) as conn: + migrate_image_favorites(conn) + migrate_image_favorites(conn) # idempotent + self.assertIsNone(image_favorites_schema_error(conn, allow_missing=False)) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', 'a.png', 'now')") + self.assertIsNone(image_favorites_schema_error(conn)) + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', x'00ff', 'now')") + self.assertIsNotNone(image_favorites_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE} WHERE typeof(relative_path)='blob'") + conn.execute(f"INSERT INTO {TABLE} VALUES ('r', ?, 'now')", ("x" * 5000,)) + self.assertIsNotNone(image_favorites_schema_error(conn)) + conn.execute(f"DELETE FROM {TABLE}") + conn.execute(f"CREATE INDEX extra ON {TABLE}(created_at)") + self.assertIsNotNone(image_favorites_schema_error(conn)) + + def test_verify_rejects_a_backup_with_an_invalid_favorites_table(self): + db = PromptDatabase(self.path) + db.close_all() + PromptModel.close_all_instances() + with closing(sqlite3.connect(self.path)) as conn: + conn.execute(f"DROP TABLE {TABLE}") + conn.execute(f"CREATE TABLE {TABLE} (anything TEXT)") + conn.commit() + ok, reason = PromptModel.verify_database_file(self.path) + self.assertFalse(ok) + self.assertIn("favorites", reason.lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_library.py b/tests/test_library.py index 5ddfd00..088706c 100644 --- a/tests/test_library.py +++ b/tests/test_library.py @@ -487,6 +487,115 @@ class TestLibrary(unittest.TestCase): rotate_history_epoch(self.db.model.get_connection()) self.assertNotEqual(token(), before) + def test_library_items_expose_run_count_and_last_used(self): + import datetime + + used = self.save("Used") + never = self.save("Never") + self.assertTrue(self.db.record_prompt_use(used)) + self.assertTrue(self.db.record_prompt_use(used, times=2)) + legacy = self.save("Legacy naive") + offset = self.save("Offset") + broken = self.save("Broken") + zulu = self.save("Zulu") + # Parse but cannot convert to UTC: past the datetime range either way. + floor = self.save("Floor") + ceiling = self.save("Ceiling") + with self.db.model.get_connection() as connection: + for prompt_id, value in ( + (legacy, "2026-10-07 18:39:06"), + (zulu, "2026-10-07T18:39:06Z"), + (floor, "0001-01-01T00:00:00+01:00"), + (ceiling, "9999-12-31T23:59:59-01:00"), + (offset, "2026-10-07T20:39:06.250+02:00"), + (broken, "yesterday"), + ): + connection.execute( + "UPDATE prompts SET last_used_at=?,run_count=1 WHERE id=?", + (value, prompt_id), + ) + items = {row["id"]: row for row in self.db.list_library()["items"]} + self.assertEqual(items[used]["run_count"], 3) + stamp = datetime.datetime.fromisoformat(items[used]["last_used_at"]) + self.assertEqual(stamp.utcoffset(), datetime.timedelta(0)) + self.assertLess( + abs(datetime.datetime.now(datetime.timezone.utc) - stamp), + datetime.timedelta(minutes=5), + ) + self.assertEqual(items[never]["run_count"], 0) + self.assertIsNone(items[never]["last_used_at"]) + self.assertEqual(items[legacy]["last_used_at"], "2026-10-07T18:39:06+00:00") + self.assertEqual( + items[offset]["last_used_at"], "2026-10-07T18:39:06.250000+00:00" + ) + self.assertIsNone(items[broken]["last_used_at"]) + self.assertEqual(items[zulu]["last_used_at"], "2026-10-07T18:39:06+00:00") + self.assertIsNone(items[floor]["last_used_at"]) + self.assertIsNone(items[ceiling]["last_used_at"]) + self.assertIsNone( + self.db.get_library_positive(floor)["positive"]["last_used_at"] + ) + self.assertEqual(type(items[broken]["run_count"]), int) + detail = self.db.get_library_positive(used)["positive"] + self.assertEqual(detail["run_count"], 3) + self.assertEqual(detail["last_used_at"], items[used]["last_used_at"]) + + def test_variant_total_matches_filtered_cards(self): + archived = self.db.save_prompt_pair("Cat", "Blur", category="A") + self.db.save_prompt_pair("Cat", "Noise", category="A") + self.db.save_prompt_pair("Cat", "Grain", category="B") + self.db.archive_prompt_pair(archived) + self.db.save_prompt_pair("Dog", "Blur", category="B") + self.save("Legacy") + everything = self.db.list_library() + self.assertEqual(everything["variant_total"], 3) + self.assertEqual( + everything["variant_total"], + sum(row["variant_count"] for row in everything["items"]), + ) + cats = self.db.list_library(text="Cat") + self.assertEqual( + (cats["total"], cats["variant_total"]), + (1, 2), + ) + self.assertEqual(self.db.list_library(category="B")["variant_total"], 3) + self.assertEqual(self.db.list_library(text="nothing")["variant_total"], 0) + first = self.db.list_library(limit=1) + self.assertEqual(first["variant_total"], 3) + # A later parent is past the cursor's ceiling, like the total. + self.db.save_prompt_pair("Bird", "Blur") + second = self.db.list_library(limit=1, after=first["next_key"]) + self.assertEqual( + (second["total"], second["variant_total"]), + (first["total"], 3), + ) + + def test_edit_token_unchanged_by_usage_fields(self): + positive = self.save("Cat", category="A", tags=["a"], rating=3) + + def parent(): + return self.db.get_library_positive(positive)["positive"] + + before = parent() + listed = self.db.list_library()["items"][0] + self.assertEqual(listed["edit_token"], before["edit_token"]) + self.db.record_prompt_use(positive, times=4) + with self.db.model.get_connection() as connection: + connection.execute( + "UPDATE prompts SET last_used_at='2001-01-01T00:00:00+00:00' " + "WHERE id=?", + (positive,), + ) + after = parent() + self.assertNotEqual( + (after["run_count"], after["last_used_at"]), + (before["run_count"], before["last_used_at"]), + ) + self.assertEqual(after["edit_token"], before["edit_token"]) + self.assertEqual( + self.db.list_library()["items"][0]["edit_token"], before["edit_token"] + ) + def test_library_snapshot_reads_epoch_and_ceiling_in_one_select_and_ends(self): self.save("Cat") connection = self.db.model.get_connection() @@ -1584,3 +1693,16 @@ class TestLibrary(unittest.TestCase): ) with self.assertRaisesRegex(ValueError, "does not belong to this positive"): self.db.list_library_images(positive, pair_id=foreign, archive="all") + + +class DisplayTitleUnits(unittest.TestCase): + """The first-line title is cut like the Library's JavaScript title().""" + + def test_title_counts_utf16_units_and_never_splits_a_surrogate_pair(self): + from database.library import display_title + + self.assertEqual(display_title((), " short line \nsecond"), "short line") + self.assertEqual(display_title((), "a" * 80), "a" * 80) + self.assertEqual(display_title((), "a" * 81), "a" * 80 + "…") + self.assertEqual(display_title((), "\U0001f600" * 80), "\U0001f600" * 40 + "…") + self.assertEqual(display_title((), "a" * 79 + "\U0001f600b"), "a" * 79 + "…") diff --git a/utils/card_seed.py b/utils/card_seed.py new file mode 100644 index 0000000..5fd0ebb --- /dev/null +++ b/utils/card_seed.py @@ -0,0 +1,115 @@ +"""Workspace card facts from the metadata records a header open already loaded. + +The card shows what the inspector shows, so it reads exactly what the +inspector's browser parser reads: each ``prompt``/``workflow`` record through +the inspection projection (``project_graph``: strict JSON, the depth and node +bounds, big integers and decimals as text), and only graphs it parsed, through +the shared parser (``extract_generation_params``, the Python twin of the +browser parser the inspector fills from since #251). No separate seed +heuristic, no graph resolution and no extra file reads. One parse feeds both +the seed and the positive that confirms a first-line card title. +""" + +import re +from collections.abc import Mapping + +from .generation_params import ( + PLACEHOLDERS, + GenerationParamsError, + extract_generation_params, +) +from .inspection_response import project_graph + +__all__ = [ + "MAX_RECORD_BYTES", + "MAX_SEED_DIGITS", + "card_seed", + "header_params", + "parsed_positive", +] + +# Each record is parsed only up to this many UTF-8 bytes; past it, nothing is read. +MAX_RECORD_BYTES = 2 * 1024 * 1024 +_RECORD_KEYS = ("prompt", "workflow") +# A digit-string seed has at most this many digits (2**64 has 20). Longer text +# is no seed, and never reaches int(), whose max-str-digits guard raises. +MAX_SEED_DIGITS = 20 +_INTEGER = re.compile(r"-?[0-9]{1,%d}\Z" % MAX_SEED_DIGITS) + + +def _utf8_size_over(text: str, limit: int) -> bool: + """True when ``text`` is over ``limit`` UTF-8 bytes, encoding only if unclear. + + A code point takes one to four bytes, so most records are decided from + their length alone without copying them. + """ + if len(text) > limit: + return True + if len(text) * 4 <= limit: + return False + return len(text.encode("utf-8", "surrogatepass")) > limit + + +def header_params(info) -> dict | None: + """The shared parser's result for the inspector's view of ``info``'s records. + + ``info`` is a Pillow ``Image.info`` (or the inspection read's metadata). + A record the projection does not parse is left out, as the inspector + leaves it out. None when no record parses, a record is not text or is + over ``MAX_RECORD_BYTES``, or the parser cannot use the graphs. + """ + if not isinstance(info, Mapping): + return None + graphs = {} + for key in _RECORD_KEYS: + text = info.get(key) + if text is None: + continue + if not isinstance(text, str) or _utf8_size_over(text, MAX_RECORD_BYTES): + return None + graph = project_graph(text, key) + if graph["status"] == "parsed": + graphs[key] = graph["value"] + if not graphs: + return None + try: + return extract_generation_params(graphs) + except (GenerationParamsError, RecursionError): + # Defence in depth: a parser failure gives no card facts, never a 500. + return None + + +def card_seed(params: dict | None) -> str | None: + """The parsed seed as a decimal string, else None. + + Always text: ``JSON.parse`` in the browser rounds a number past 2**53, and + ComfyUI's randomized seeds are often that large. The projection carries + integers beyond 2**53 as their exact digits, so a digit string is the + recorded seed (the inspector shows it too), up to ``MAX_SEED_DIGITS`` + digits. A placeholder, other or longer text, a boolean or a decimal gives + None. + """ + if not isinstance(params, Mapping): + return None + value = params.get("seed") + if isinstance(value, bool): + return None + try: + if isinstance(value, int): + return str(value) + if isinstance(value, str) and _INTEGER.match(value): + return str(int(value)) + except ValueError: + # Python's integer string conversion limit: no seed, never a 500. + return None + return None + + +def parsed_positive(params: dict | None) -> str | None: + """The parsed positive prompt text, or None for the placeholder or blank.""" + if not isinstance(params, Mapping): + return None + text = params.get("positivePrompt") + if not isinstance(text, str) or text == PLACEHOLDERS["positivePrompt"]: + return None + return text if text.strip() else None diff --git a/utils/inspection_response.py b/utils/inspection_response.py index 02c37d3..8dffb4e 100644 --- a/utils/inspection_response.py +++ b/utils/inspection_response.py @@ -161,6 +161,21 @@ def _graph(value: object, *, workflow: bool = False) -> dict: return {"status": "parsed", "value": value} +def project_graph(text: str, kind: Literal["prompt", "workflow"]) -> dict: + """One record's graph as the inspector projects it: ``{status, value}``. + + ``status`` is ``parsed`` (``value`` is the projected graph the browser + parser receives), ``limited`` (past the depth or node bound) or + ``malformed``. Shared with the workspace card so both read one bound. + """ + try: + return _graph(_strict_projection(text, kind=kind), workflow=kind == "workflow") + except _ProjectionLimit: + return {"status": "limited", "value": None} + except (ValueError, TypeError, KeyError, OverflowError): + return {"status": "malformed", "value": None} + + def _metadata(inspection: "InspectionRead | None") -> dict: formats = ( {} @@ -213,16 +228,15 @@ def _metadata(inspection: "InspectionRead | None") -> dict: graph["status"] = failed or "limited" elif formats.get(key) == "invalid": graph["status"] = "malformed" + elif carrier is None and formats.get(key) == "reserialized-json": + graph["status"] = "limited" + elif carrier is None: + graph = project_graph(raw.get(key), key) else: try: - if carrier is not None: - value = carrier.get(key) - if isinstance(value, str): - value = _strict_projection(value, kind=key) - elif formats.get(key) == "reserialized-json": - raise _ProjectionLimit() - else: - value = _strict_projection(raw[key], kind=key) + value = carrier.get(key) + if isinstance(value, str): + value = _strict_projection(value, kind=key) graph = _graph(value, workflow=key == "workflow") except _ProjectionLimit: graph["status"] = "limited" diff --git a/web/js/workspace/library.mjs b/web/js/workspace/library.mjs index 389df2c..fc62309 100644 --- a/web/js/workspace/library.mjs +++ b/web/js/workspace/library.mjs @@ -28,7 +28,14 @@ function el(tag,text,className) { const node=document.createElement(tag);if(text!==undefined) node.textContent=text;if(className) node.className=className;return node; } function link(text,href,className) {const node=el('a',text,className);node.href=href;return node;} -function title(text) {const first=text.split('\n')[0].trim();return first.length>80?`${first.slice(0,80)}…`:first;} +// Cut at 80 UTF-16 units, never inside a surrogate pair (database/library.py titled() matches). +export function libraryTitle(text) { + const first=text.split('\n')[0].trim(); + if (first.length<=80) return first; + const cut=first.slice(0,80); + return `${/[\uD800-\uDBFF]$/.test(cut)?cut.slice(0,-1):cut}…`; +} +const title=libraryTitle; /** The full stored title; the first-line fallback only when it is blank. */ function displayTitle(positive) {return typeof positive.title==='string' && positive.title.trim()?positive.title:title(positive.text);} function focusHeading(main) {const heading=main.querySelector('h1');if(!heading)return;heading.tabIndex=-1;heading.focus();}