Compare commits

..
Author SHA1 Message Date
Vito Sansevero 387bca2c51 test: exercise real page workflows and enforce frontend coverage 2026-10-16 06:30:56 -07:00
Vito e9fe586442 fix: repair prompt edit hashes, auto-tag review, and stale UI responses (#180)
* fix: keep edited prompt hashes consistent

Fixes #176

* fix: align gallery auto-tag review with scoped image APIs

Fixes #177

* fix: keep Tags prompt results on the current filter

Fixes #178

* fix: ignore stale metadata in admin image viewers

Fixes #179

* test: compare canonical gallery paths on macOS
2026-10-16 06:14:06 -07:00
Vito 6d9f3fc8fb fix: preserve gallery image identity and ignore stale metadata (#175)
* fix: preserve gallery root identity across media and metadata

Fixes #173

* fix: ignore stale gallery metadata responses

Fixes #174

* test: use platform-valid gallery filenames on Windows
2026-10-16 05:35:03 -07:00
dependabot[bot] 911b801883 deps(pip)(deps-dev): update aiohttp requirement (#172)
---
updated-dependencies:
- dependency-name: aiohttp
  dependency-version: 3.14.5
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-16 05:23:48 -07:00
dependabot[bot] a576cff69f deps(pip)(deps-dev): update black requirement (#171)
Updates the requirements on [black](https://github.com/psf/black) to permit the latest version.
- [Release notes](https://github.com/psf/black/releases)
- [Changelog](https://github.com/psf/black/blob/main/CHANGES.md)
- [Commits](https://github.com/psf/black/compare/26.10.0...26.10.1)

---
updated-dependencies:
- dependency-name: black
  dependency-version: 26.10.1
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-16 05:23:28 -07:00
Vito cf0cef3f25 feat: vendor Viewer.js reproducibly and automate release updates (#170)
Refs #169. Pin and verify pristine distribution assets, upgrade to 1.15.2, and add tested maintainer tooling plus desktop and touch browser checks.
2026-10-16 04:56:25 -07:00
Vito 7b6430e20a fix: stop misleading thumbnail prompts on gallery startup (#168)
* fix: count pending thumbnails accurately across gallery roots

* fix: clear thumbnails across all configured gallery roots

* fix: log skipped thumbnail directory cleanup
2026-10-16 04:56:07 -07:00
Vito 611c19d2f4 fix: keep infinite-scroll gallery rendering and viewer work bounded (#166) 2026-10-16 04:17:20 -07:00
ghosttigerllc-bitandkai 2763e9f9fa Read A1111/Forge generation parameters: fix the SELECT and parse multi-line prompts (#165)
* get_image_prompt_info: include gi.parameters in SELECT and return value

The function selects workflow_data, prompt_metadata and generation_time from
generated_images but omits the parameters column, so the A1111/Forge
generation-parameters string (steps, sampler, CFG, seed, LoRA list) never
reaches the API even when present in the database. Add it to the SELECT and
include it in the returned dict.

* parse A1111/Forge parameters: multi-line prompts, structured settings

_parse_comfyui_prompt took lines[0] of the parameters chunk as the whole
positive prompt, truncating every multi-line prompt to its first line, and
matched 'negative prompt:' only as a full line start. Split on the trailing
'Steps: ...' settings line instead, keep the full multi-line positive prompt,
and extract Steps / Sampler / Schedule type / CFG scale / Seed as structured
fields alongside the raw string.

Together with the get_image_prompt_info SELECT fix (previous commit), the
full A1111/Forge generation recipe now flows from the database to the API.

* style: apply black to the A1111 parameters parser

CI's black --check flagged py/api/__init__.py (slice spacing and two long
lines in the settings extraction). Formatting only; the AST is unchanged.

---------

Co-authored-by: kai <kai@hermes.local>
2026-10-16 03:50:26 -07:00
Vito Sansevero 3bb21c621c chore: release 3.3.1
- Fix Prompt Manager UI routing behind a ComfyUI proxy mount such as
  SwarmUI's /ComfyBackendDirect/ (#164, thanks @ghosttigerllc-bit)
- Restore runtime minimums Pillow>=8.0.0 and watchdog>=2.1.0: nothing
  requires the newer floors, and raising them forces upgrades in users'
  shared ComfyUI environment
2026-10-14 03:26:05 -07:00
ghosttigerllc-bitandZayne 374ce6e7a8 Fix Prompt Manager UI routing behind a ComfyUI proxy (#164)
* Fix gallery folder filters behind SwarmUI proxy

* Route Prompt Manager UI requests through proxy mount

* Use Comfy frontend base for Prompt Manager node requests

---------

Thanks @ghosttigerllc-bit

Co-authored-by: Zayne <noreply@gtb.ai>
2026-10-14 03:18:29 -07:00
dependabot[bot] e2ec6df3d4 deps(pip)(deps-dev): update numpy requirement from >=1.24 to >=2.2.6 (#162)
Updates the requirements on [numpy](https://github.com/numpy/numpy) to permit the latest version.
- [Release notes](https://github.com/numpy/numpy/releases)
- [Changelog](https://github.com/numpy/numpy/blob/main/doc/RELEASE_WALKTHROUGH.rst)
- [Commits](https://github.com/numpy/numpy/compare/v1.24.0...v2.2.6)

---
updated-dependencies:
- dependency-name: numpy
  dependency-version: 2.2.6
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:46:02 -07:00
dependabot[bot] 38daedf131 deps(pip)(deps-dev): update flake8 requirement from >=7.0.0 to >=7.4.1 (#160)
Updates the requirements on [flake8](https://github.com/pycqa/flake8) to permit the latest version.
- [Commits](https://github.com/pycqa/flake8/compare/7.0.0...7.4.1)

---
updated-dependencies:
- dependency-name: flake8
  dependency-version: 7.4.1
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:45:59 -07:00
dependabot[bot] 1e1b92a1ba deps(pip)(deps-dev): update aiohttp requirement from >=3.8.0 to >=3.14.4 (#159)
---
updated-dependencies:
- dependency-name: aiohttp
  dependency-version: 3.14.4
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:45:55 -07:00
dependabot[bot] ef66d7bd3e deps(pip)(deps): update watchdog requirement from >=2.1.0 to >=6.0.0 (#156)
Updates the requirements on [watchdog](https://github.com/gorakhargosh/watchdog) to permit the latest version.
- [Release notes](https://github.com/gorakhargosh/watchdog/releases)
- [Changelog](https://github.com/gorakhargosh/watchdog/blob/master/changelog.rst)
- [Commits](https://github.com/gorakhargosh/watchdog/compare/v2.1.0...v6.0.0)

---
updated-dependencies:
- dependency-name: watchdog
  dependency-version: 6.0.0
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:45:52 -07:00
dependabot[bot] f00d6408d8 deps(pip)(deps-dev): update pre-commit requirement (#163)
Updates the requirements on [pre-commit](https://github.com/pre-commit/pre-commit) to permit the latest version.
- [Release notes](https://github.com/pre-commit/pre-commit/releases)
- [Changelog](https://github.com/pre-commit/pre-commit/blob/main/CHANGELOG.md)
- [Commits](https://github.com/pre-commit/pre-commit/compare/v4.0.0...v4.6.2)

---
updated-dependencies:
- dependency-name: pre-commit
  dependency-version: 4.6.2
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:32:43 -07:00
dependabot[bot] 6962fecf86 deps(pip)(deps-dev): update coverage requirement from >=7.0 to >=7.16.2 (#161)
Updates the requirements on [coverage](https://github.com/coveragepy/coveragepy) to permit the latest version.
- [Release notes](https://github.com/coveragepy/coveragepy/releases)
- [Changelog](https://github.com/coveragepy/coveragepy/blob/main/CHANGES.rst)
- [Commits](https://github.com/coveragepy/coveragepy/compare/7.0.0...7.16.2)

---
updated-dependencies:
- dependency-name: coverage
  dependency-version: 7.16.2
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:32:38 -07:00
dependabot[bot] 90d706c943 deps(pip)(deps-dev): update bandit requirement from >=1.7.0 to >=1.9.4 (#158)
Updates the requirements on [bandit](https://github.com/PyCQA/bandit) to permit the latest version.
- [Release notes](https://github.com/PyCQA/bandit/releases)
- [Commits](https://github.com/PyCQA/bandit/compare/1.7.0...1.9.4)

---
updated-dependencies:
- dependency-name: bandit
  dependency-version: 1.9.4
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:32:31 -07:00
dependabot[bot] d17be56e1b deps(pip)(deps-dev): update black requirement from >=24.0.0 to >=26.10.0 (#157)
Updates the requirements on [black](https://github.com/psf/black) to permit the latest version.
- [Release notes](https://github.com/psf/black/releases)
- [Changelog](https://github.com/psf/black/blob/main/CHANGES.md)
- [Commits](https://github.com/psf/black/compare/24.1a1...26.10.0)

---
updated-dependencies:
- dependency-name: black
  dependency-version: 26.10.0
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:32:28 -07:00
dependabot[bot] 94366938d1 deps(pip)(deps-dev): update setuptools requirement from >=61 to >=84.0.0 (#155)
Updates the requirements on [setuptools](https://github.com/pypa/setuptools) to permit the latest version.
- [Release notes](https://github.com/pypa/setuptools/releases)
- [Changelog](https://github.com/pypa/setuptools/blob/main/NEWS.rst)
- [Commits](https://github.com/pypa/setuptools/compare/v61.0.0...v84.0.0)

---
updated-dependencies:
- dependency-name: setuptools
  dependency-version: 84.0.0
  dependency-type: direct:development
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:32:23 -07:00
dependabot[bot] e13ea3c92f deps(pip)(deps): update pillow requirement from >=8.0.0 to >=12.3.0 (#154)
Updates the requirements on [pillow](https://github.com/python-pillow/Pillow) to permit the latest version.
- [Release notes](https://github.com/python-pillow/Pillow/releases)
- [Changelog](https://github.com/python-pillow/Pillow/blob/main/CHANGES.rst)
- [Commits](https://github.com/python-pillow/Pillow/compare/8.0.0...12.3.0)

---
updated-dependencies:
- dependency-name: pillow
  dependency-version: 12.3.0
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-11 07:32:20 -07:00
Vito da322f8c94 Hardening: security fixes, queue-time run counting, background jobs, 85%+ test coverage (#153)
* chore: valid packaging metadata, flake8 config file, ignore local config with secrets

pyproject.toml carried requirements.txt comment lines as dependencies, which
fails PEP 508 parsing. flake8 never read [tool.flake8] from pyproject, so the
settings move to .flake8 and pre-commit runs the full rule set. config.json
(holds the CivitAI API key and machine paths) and other local artifacts are
now ignored. coverage joins the dev requirements.

* ci: test on Linux, Windows and macOS with a coverage gate and lint job

* refactor: drop validate_workflow_name, the column it validated no longer exists

* chore: keep agent worktrees out of coverage measurement

* fix: image routes serve only media files inside allowed directories

serve_image and serve_output_image now reject anything whose extension
is not in IMAGE_EXTENSIONS, fail closed when no output directory is
configured, and answer 403 (not 404) for dot-dot, absolute and
symlink-escaping paths. Containment is checked after realpath+normcase
so it behaves the same on Windows, macOS and Linux. 500 responses use
safe_error_message so OSError text no longer leaks absolute paths.

* fix: bound page sizes on list endpoints

/search, /recent and /images/recent share parse_page_params: limit is
clamped to [1, MAX_PAGE_LIMIT=500], offset to >= 0, and non-integer
values answer 400 instead of 500. /search now honours offset and all
three return a pagination block with the effective limit and offset.

* perf: thumbnail generation no longer blocks the event loop

generate_thumbnails_with_progress is split into _thumbnail_targets
(scan, executor), _generate_one (one PIL/ffmpeg job, executor) and an
async driver that awaits each file via _run_in_executor, emits one
progress frame per file and yields between files. Per-file failures
are reported as file_error events (error stays fatal for the UI) and
the loop continues. The wildcard Access-Control-Allow-Origin header is
gone. The blocking POST variant reuses the same helpers. The scan never
follows symlinks and is capped by MAX_SCAN_DEPTH / MAX_THUMBNAIL_FILES.

* fix: cap every request-driven resource in the image routes

- /images/recent, /images/all and /images/output parse limit/offset
  through parse_page_params on every path (absent, 0, negative, float
  strings), and the clamped values are what reach the DB. /images/all
  gets its own MAX_BULK_LIMIT ceiling instead of an unbounded query;
  /images/recent caps its fetch window at MAX_RECENT_IMAGES_WINDOW.
- parse_page_params also clamps offset to MAX_PAGE_OFFSET so values past
  int64 no longer surface as a 500 from SQLite.
- Gallery and thumbnail scans share _iter_media_files: no symlink
  following, MAX_SCAN_DEPTH, and MAX_GALLERY_FILES / MAX_THUMBNAIL_FILES
  caps. Thumbnail size comes only from named presets.
- JSON bodies on /images/link and /images/generate-thumbnails are read
  in bounded chunks; over MAX_JSON_BODY_BYTES answers 413, malformed or
  non-object JSON answers 400.
- /images/link only accepts media paths inside the allowed directories
  (403 otherwise), closing a path-existence oracle.

* test: cover image and prompt routes

Every route in py/api/images.py and py/api/prompts.py now has a success
path and an error path test: save/update/delete/rate/tag/category/bulk
endpoints, image list/link/unlink/delete/metadata/scan endpoints, the
thumbnail SSE stream, clear-thumbnails, and the video thumbnail chain
(fake cv2 success/failure, ffmpeg absent or failing, PIL placeholder).
Combined statement+branch coverage: images.py 97.7%, prompts.py 97.4%.

* fix: autotag endpoints no longer read arbitrary paths or run on GET

* fix: CivitAI key is never returned or sent off-host

* fix: autotag/single checks real paths and DB rows against output dirs

* fix: CivitAI downloads never follow redirects

* refactor: autotag model loading is injectable

* test: cover autotag and lora integration

* fix: saving LoRA settings without a key keeps the stored CivitAI key

* refactor: share the internal error message in autotag routes

* fix: never treat negative-only PromptManager nodes as the run prompt

run_prompt_nodes fell back to every PromptManager node whenever no positive
node was traced, including graphs that do have a sampler and whose only
PromptManager node feeds its negative input. Fall back only when the graph
has no recognisable sampler (new has_sampler helper). Samplers named like
one are recognised even without a model link. Adversarial graphs (cycles,
10k-node fan-out, non-dict inputs) are covered.

* fix: do not follow negative text inputs when tracing the positive prompt

_TEXT_KEY matched input names such as text_negative, so a styler node feeding
a sampler's positive input pulled its negative PromptManager node into the
positive trace. Input names matching (^|_)neg(ative)?($|_) are now skipped
(new follows_input helper).

* fix: count prompt runs when ComfyUI dequeues the prompt, not when it is submitted

The on-prompt hook runs before ComfyUI validates the request, so counting
there bumped run_count for prompts that were rejected or deleted from the
queue. The hook now only resolves the positive PromptManager texts, mints a
UUID4 prompt_id when the client sent none, and stashes the hashes under that
id (QueuedRuns, bounded to 1000 entries). PromptQueue.get is wrapped
idempotently; when a prompt is dequeued for execution its stashed runs are
counted (or remembered as pending first use). Servers without a prompt_queue
fall back to counting at submit time with one warning.

__init__ defines _global_db before the monitoring block and passes a factory
that raises RuntimeError while the database is unavailable; the package is
now imported under test with server and folder_paths stubbed.

* fix: keep batch image queue aligned when unrelated workflows save images

The metadata strategy returned None for images whose embedded workflow has
no positive PromptManager node, so the handler popped the batch FIFO for
them and every later batch image linked to the wrong prompt. It now returns
a tri-state: linked (prompt found), foreign (no PromptManager prompt in the
workflow, so the image is skipped and no fallback applies) or unknown (only
then the queue, snapshot, tracker and recent-prompt fallbacks run).

* fix: image monitor validates directories before creating the observer

start_monitoring assigned self.observer before checking that any directory
existed, so with no valid directory the monitor was stuck "already running"
and stop_monitoring raised on join of a never-started observer. Directories
are validated first and only a started observer is kept; stop_monitoring is
safe to call at any time and always joins. New is_monitoring property.

The per-image threading.Timer is replaced by one daemon worker: paths are
deduplicated and bounded (MAX_PENDING_IMAGES), processed in creation order
(batch images pop the prompt queue in order) after the processing delay and
once the file size is unchanged across two reads.

* refactor: one prompt hash implementation

utils.hashing.generate_prompt_hash is now the only prompt hash:
PromptManagerBase._generate_hash delegates to it and the inline
hashlib.sha256 copies in the prompt tracker and image monitor are gone.

Remove PromptManagerBase._open_web_interface (it targeted a web_interface
folder that never existed) and the unused restart_gallery.py script.

IS_CHANGED on PromptManager and PromptManagerText now includes category and
tags, so editing either re-runs the node and updates the saved metadata.

* fix: stop ComfyUI from loading the admin bundle into the canvas page

WEB_DIRECTORY = "web" made ComfyUI inject every .js under web/ (admin.js,
gallery.js, vendored Tailwind and ViewerJS) into the canvas page. The node
extension moves to web/comfy/prompt_manager.js (content unchanged; its
../../scripts/app.js import resolves against the served URL, which is the
same) and WEB_DIRECTORY points at that folder only. Admin and gallery assets
are still served from web/ by the extension's own routes.

tests/test_sort_contract.py still reads web/prompt_manager.js and must be
updated to the new path by its owner.

* fix: image monitor recent-prompt fallback reads the paginated result

get_recent_prompts returns {"prompts": [...]} but _get_fallback_prompt
indexed it like a list, so the last-resort link to the most recent prompt
always failed silently. Accept both shapes.

* test: cover prompt_manager_base, prompt_manager and prompt_search_list

Node base: init wiring, metadata update on re-save, run counting, database
error paths, tag parsing, search and recent API wrappers, gallery start
failure, LoRA trigger-word injection (disabled, missing modules, enabled).
PromptManager: missing CLIP error, encoding of the combined text with a fake
CLIP, tracking failure tolerance, empty text. PromptManagerText: tracking
failure tolerance. PromptSearchList: filters, multipart and LoRA-only
skipping, newline collapsing, preview, empty and error results.

* fix: diagnostics dependency check survives a watchdog without __version__

check_dependencies read watchdog.__version__ for its log line; recent
watchdog releases have no such attribute, so the AttributeError escaped and
took the whole diagnostics run down. Read it with a default.

Also covers utils/diagnostics.py: database and images-table checks against
a real temporary database, file system check, ComfyUI output detection
(folder_paths, configured directories, relative fallback, nothing found),
dependencies, the full run and the test image link.

* test: cover utils/logging_config

Isolated logger manager on a temporary log directory: memory buffer order,
level filter and 1000-entry cap, rotation and backup listing, tail and whole
file reads, path traversal rejection, truncation (including removal
errors), level changes across prompt_manager loggers, handler rebuilds.

* test: cover utils/comfyui_integration

Singleton re-initialisation keeps registrations, latest prompt per thread,
cross-thread fallback with the five-minute recency limit, global lookup for
the current thread, and cleanup of old registrations.

* test: cover __init__

Version parsing from pyproject (and the unknown fallback), and a server
whose add_on_prompt_handler raises: nodes and monitoring still come up.

* docs: point at the moved canvas extension under web/comfy

* fix: give every thread its own sqlite connection

PromptModel.get_connection() handed every thread the same
check_same_thread=False connection, so a commit or rollback from the
watchdog, executor or queue-hook thread could land on another thread's
half-done write (sqlite3 raised 'bad parameter or other API misuse'
under concurrent saves). Connections now live in threading.local, each
with WAL, Row factory, foreign_keys and busy_timeout, and are tracked so
they can be closed. PromptModel.close() / PromptDatabase.close() close
the calling thread's connection so tests and restores can delete the
file on Windows.

* fix: saving an existing prompt concurrently returns its id

Node execution, the queue hook and the image monitor all do
get_prompt_by_hash() then save_prompt(). When two of them race on the
same text the second INSERT failed on UNIQUE(hash) and surfaced an
IntegrityError. save_prompt now catches it, re-fetches by hash and
returns the existing id, so eight threads saving one prompt end with a
single row and every record_prompt_use counted.

* fix: backups include WAL data and restores are integrity-checked

The database runs in WAL mode, so a shutil.copy2 of prompts.db silently
dropped every committed row still in prompts.db-wal; the backup
download and the pre-restore safety copy both did that. Restore also
overwrote the live file with connections open, read the whole upload
into memory before checking the 100 MB cap and never ran
PRAGMA integrity_check.

PromptModel.backup_database now uses the SQLite online backup API,
verify_database_file opens a candidate read-only and checks header,
integrity and the prompts table, and restore_from_file backs up the
live database, closes every thread's connection, copies the file,
removes stale -wal/-shm and re-runs migrations. The admin handlers
stream the upload to a temp file with an early size cutoff, reject
anything verify_database_file refuses with a 400 and keep the
{success: ...} JSON shape. The schema-init connection is now closed
instead of leaking.

* fix: migrations are idempotent and schema init runs once per process

Every PromptDatabase() (nodes, image monitor, API, diagnostics) re-ran
table creation and all migrations. Schema init is now cached per
normalised path (class-level set under a lock) and re-run only when the
file disappeared, after a restore, or after reset_schema_cache().

The copy-and-swap migrations drop a stale *_new table left by an
interrupted run, execute inside an explicit transaction with
PRAGMA foreign_keys = OFF (DROP TABLE prompts used to cascade-delete
every generated_images row), roll back cleanly and log and re-raise
sqlite3.Error instead of swallowing it.

* fix: image uniqueness by full path, escaped LIKE search, no usage stamp on save

generated_images was UNIQUE(prompt_id, filename) on the basename, so
ComfyUI's same-named files in different date folders were silently
ignored by INSERT OR IGNORE, and the lastrowid == 0 branch after it
never ran. The table now carries file_path (normcase/normpath/abspath
of the image path) with UNIQUE(prompt_id, file_path); the migration
collapses rows that normalise to one file and supersedes the old
basename migration, which would otherwise undo it on every start.

search_prompts and the tag search escaped nothing, so 100% matched
'100 percent' and a_b matched aXb; user text now goes through
escape_like with ESCAPE. save_prompt no longer stamps last_used_at, the
usage backfill runs only on the upgrade that adds the columns, and the
Recently Used / Most Used orders put never-run prompts (NULL) after run
ones instead of first.

* fix: escape LIKE wildcards in image prompt search

search_images_by_prompt interpolated the user's text into LIKE
unescaped, so '50%' also matched '50 percent'. It now goes through the
same escape_like helper as search_prompts. Also drops an unused import
and an unused cursor assignment flagged by pyflakes.

* test: cover database operations

Behaviour tests for export (json/csv/errors), duplicate detection and
merging, image metadata storage and lookup, missing-file cleanup,
statistics, bulk operations, orphan pruning, consistency checks,
subfolder listing, tag queries and management edge cases, legacy tag
column parsing, config-driven path resolution, model housekeeping and
the TEXT -> INTEGER prompt_id migration. database/operations.py and
database/models.py are both above 85% statement+branch coverage.

* test: isolate the clock and log buffer from background threads of other tests

* fix: gallery roots are restricted to ComfyUI directories

A gallery root could be set to '/' (or 'C:\') through POST
/prompt_manager/settings because the deny list only matched a few
system prefixes. The root was persisted to config.json and fed into the
image-serving containment check, which then allowed any file on disk to
be read through /prompt_manager/images/serve/.

- GalleryConfig.validate_gallery_root() accepts a directory only when its
  realpath lies inside a ComfyUI directory (folder_paths output/input/
  temp/user) or a directory listed in PROMPT_MANAGER_EXTRA_GALLERY_ROOTS;
  filesystem roots, the home directory, files and missing paths are
  rejected.
- save_settings returns 400 and persists nothing when any root fails.
- _get_all_output_dirs() filters configured roots through the validator
  and never falls back to a rejected root.
- PROMPT_MANAGER_CONFIG_PATH relocates config.json; load/save default
  to it.

* fix: log config accepts only known keys

POST /prompt_manager/logs/config forwarded the whole request body into
PromptManagerLogger.update_config(), so any key could be injected into
the logging configuration. The handler now accepts only LOG_CONFIG_KEYS
(level, max_file_size, backup_count, console_logging, file_logging,
buffer_size), type-checks each value (level name, non-negative int,
bool) and answers 400 listing unknown keys. Malformed or non-object
bodies are 400 instead of 500.

* fix: no absolute server paths in API responses

Settings, diagnostics, duplicate/output scans and the logging endpoints
returned absolute filesystem paths, and several error envelopes echoed
str(e) which for OSError carries the full path.

- _public_path() renders a path relative to the ComfyUI base directory
  (folder_paths.base_path, else the output dir's parent) or to the
  parent of an opted-in extra gallery root, falling back to the
  basename; _public_error() reduces OSError to strerror + basename.
- GET /settings returns gallery roots and monitored directories in that
  form; POST /settings resolves relative roots against the same anchors
  (GalleryConfig.resolve_gallery_root) so the GET output round-trips.
- Diagnostics output_dirs, monitored_directories and the missing-DB
  message no longer carry absolute paths; the image monitor is reached
  through the package-or-top-level import fallback instead of a
  three-dot relative import that failed outside the package.
- Duplicate scan 'path' is now relative to the output root and the
  delete endpoint accepts relative paths.
- /logs/files 'path' and /logs/stats 'log_directory' use the public
  form; unexpected errors answer with a generic message.

* perf: run duplicate and output scans off the event loop

The duplicate scan walked the output tree and hashed every file on the
aiohttp event loop, and the duplicate delete removed files inline; the
output scan ran its per-file database work one executor hop at a time.

- _find_duplicate_images_sync / _describe_media_file do the walk and
  hashing; find_duplicate_images() runs them via _run_in_executor.
- _delete_duplicate_images_sync / _delete_one_duplicate do the removal
  (absolute or output-relative paths) off the loop.
- The output scan is split into _collect_output_media_sync,
  _extract_batch_metadata_sync and _ingest_scan_batch_sync (one
  executor call per batch for all DB work) with a small SSE generator.
- Media collection is shared and de-duplicates by normcase'd path.

* refactor: settings reads the image monitor through its module

get_settings located the ImageMonitor singleton by scanning sys.modules
for any module with 'image_monitor' in its name, a workaround for the
package being importable under two identities. It now uses the same
try-package / except-top-level import as the rest of the API
(_current_image_monitor) and falls back to the configured roots when no
monitor is running.

* fix: clamp the logs limit and reject non-integers

GET /prompt_manager/logs accepted any limit up to 1000 and answered 500
when the value was not an integer. The limit is now clamped to [1, 500]
and a non-integer value is a 400.

* test: cover core API module and gzip middleware

* refactor: split diagnostics, maintenance and HTML serving into helpers

- run_diagnostics composes four module-level probes (_diagnose_database,
  _diagnose_dependencies, _diagnose_output_dirs, _diagnose_image_monitor).
- run_maintenance is table-driven: _maintenance_runners maps each
  operation name to a callable and _run_maintenance_operations reports
  per-operation success or a public error message.
- _read_json answers 400 for malformed or non-object JSON bodies in
  settings, duplicate delete, test-link and maintenance instead of 500.
- test_image_link delegates to _link_test_image.
- The four HTML page routes share _serve_html_page.

* test: cover admin and core API routes

Route tests for stats, cleanup, duplicate scan/delete, settings,
diagnostics, test-link, maintenance (every operation, success and
failure), the output scan internals, and contract-level tests for
backup/restore (status codes and envelope only).

* test: make admin, logging and diagnostics tests independent of test order

The diagnostics message no longer carries the absolute path, the restore test
closes its raw sqlite connections before the live file is replaced, the logs
tests declare their own folder_paths base directory, and the isolated logger
filters every handler to the test thread.

* fix: truncating logs no longer rolls the main log over or pads it with zeros

truncate_logs emptied prompt_manager.log through a second open(), leaving the
rotating handler's stream at its old offset: the next record was written past
the new end (zero padding on POSIX) and tell() still reported the old size, so
the file rolled over and resurrected the backup that had just been deleted.
It also logged once per file from inside the deletion loop, which could do the
same with a small size limit. The main log is now truncated through the
handler's own stream and a single summary line is logged afterwards.

* fix: HTML escaping covers attribute context

escapeHtml now escapes & < > " ' (shared web/js/html-escape.js, loaded by
admin, gallery and metadata pages) instead of the textContent/innerHTML
round-trip that left quotes intact, so a tag, category, checkpoint or file
name containing a quote can no longer break out of a data-*, value or title
attribute.

Inline onclick handlers that interpolated server data now carry data-action
and data-* attributes and are dispatched through one delegated listener per
page (web/js/data-actions.js); the img onerror fallback uses
data-fallback-src with a capturing error listener. Every remaining template
value written into an attribute or text node (prompts, checkpoints, file
paths, log lines, directory names, categories, suggestions) is escaped.

metadata.html no longer builds an "explorer /select" shell command from the
file name (Electron-only branch); the path is copied to the clipboard.

* refactor: remove shadowed and dead frontend code

admin.js defined parsePNGMetadata, extractComfyUIData, showFullPrompt,
showWorkflowData, downloadWorkflowJSON and copyAllMetadata twice; the first
definitions were unreachable because the later ones shadow them.
updateMetadataPanel, showMetadataError, copyPrompt and tryFallbackMetadata
only ever wrote to #metadataContent, an element no page creates, so they are
gone too. gallery.js defined escapeHtml and showNotification twice. The admin
keydown/resize handlers referenced closeImageViewer, previousImage, nextImage
and applyImageSizing, none of which exist, and the imageViewMode state they
used; only the Escape-closes-gallery branch remains.

Debug console.log calls are removed; the ones that were the only statement of
a catch block are console.warn so parse failures stay visible.

* fix: stale responses, double submits and leaked listeners

web/js/api-client.js wraps fetch with one failure path (ApiError for non-2xx,
success:false and non-JSON bodies) and latestOnly(), which aborts the previous
in-flight request and resolves superseded calls to STALE. The admin search and
recent-page loads share one guard, the gallery page load has its own, so fast
typing or clicking can no longer paint an older result over a newer one.

In-place prompt editing used a blur listener with {once:true} plus a keydown
listener added on every edit, which sent the PUT twice and stacked handlers.
The list now has one keydown/focusout pair bound at construction and
finishEdit() clears the edit state before touching the element, so a save is
sent exactly once. Mutations re-fetch the page currently shown (refreshList)
instead of calling search(), and delete, bulk delete, edit, restore and scan no
longer reload the whole page.

Pagination state is committed only after a successful fetch via the pure
ListState.nextPageState(current, response) in both admin and gallery, so a
failed or aborted page request leaves the current page intact. The gallery
lightbox removed its keydown handlers on a non-existent 'remove' event; a
MutationObserver now drops them when the modal leaves the document.

web/js/sse-stream.js reads server-sent events over fetch so POST-only progress
endpoints can stream.

* fix: follow the POST-only autotag streams and the new LoRA and list contracts

autotag/download and autotag/start are POST-only now; EventSource can only
GET, so both pages stream their progress through SseStream.connect(url,
{method: 'POST'}), which keeps the onmessage/onerror/close surface the
handlers already use. autotag/single is called with image_id and the re-tag
dialog shows the server's filename instead of splitting an absolute path.

lora/status reports has_civitai_api_key instead of the key: the field shows a
'key is stored' placeholder, an empty field keeps the key on save and a
'Clear the stored key on save' checkbox sends clear_civitai_api_key.

Thumbnail generation reports per-file failures as file_error events; both
progress views count them (and show error_count from progress) without
closing the stream. The auto-tag review follows pagination.has_more on the
capped /images/all endpoint, and /search is paged like /recent so paging in
search mode no longer switches the list back to recent prompts.

* refactor: one PNG metadata parser

admin.js, gallery.js and metadata.html each carried a copy of the PNG chunk
walker that reads ComfyUI's tEXt/iTXt entries (admin.js had two until the
shadowed one was removed). web/js/png-metadata.js is now the only parser:
it checks the signature, bounds-checks every chunk so a truncated file or a
bogus length cannot throw, decodes tEXt as Latin-1 and iTXt as UTF-8, skips
compressed zTXt/iTXt payloads instead of decoding them as garbage, and stops
at IEND. The three page-level parsePNGMetadata functions delegate to it.

* fix: tags deep link renders and admin modals close

admin.html loaded tags-page.js after admin.js, so opening the page at #/tags
found no TagsPageManager and rendered nothing; the script order is now
theme-bridge, html-escape, api-client, png-metadata, (list-state, sse-stream,
data-actions), viewer, comfy-metadata, prompt-list-sort, view-router,
tags-page, admin. Route resolution lives in the pure ViewRouter.resolveView().

The auto-tag modal's close buttons called a bare hideModal(), which only
gallery.js defines; they call window.admin.hideModal(). PromptAdmin is
constructed from startAdmin() on DOMContentLoaded (or immediately when the
document is already parsed) instead of at script evaluation time.

* test: cover frontend modules

browser-globals.test.js loads every page module through vm with a fake
window, the way the pages do, and asserts the global each one exposes
(HtmlEscape/escapeHtml, DataActions, ApiClient, ListState, SseStream,
PngMetadata, ViewRouter, PromptListSort, ComfyMetadata). sse-stream gains
cases for a POST body, missing handlers and close() before or during
delivery. Every pure module is above 85% line coverage under
node --test --experimental-test-coverage.

* fix: page image listings in SQL and drop three unused helpers

get_recent_images and search_images_by_prompt take limit/offset so the recent
and search image routes no longer over-fetch a 5000-row window and slice it.
generate_content_hash, PromptModel.migrate_database and find_duplicates had no
callers outside their own tests. The isolated logger tests re-apply their
thread filter after update_config rebuilds the handlers. Long lines in the
touched database and image modules are wrapped so the full flake8 hook passes.

* fix: restore through SQLite's backup API, close_all for shared files, safer output scans

restore_from_file copied the uploaded file over the live database and deleted
its WAL while other PromptModel instances (nodes, image monitor, queue hook)
still held connections, which corrupts on POSIX and fails on Windows. It now
copies pages through the online backup API into the live connection so every
open handle sees the new content. PromptDatabase.close_all() closes the
connections of every thread, which tests need before deleting a database.
Output scans skip symlinked files, filesystem entries get a stable digest id
instead of a per-process hash(), and thumbnail streaming stops when the
client disconnects.

* style: wrap long lines in node modules

Split tooltip, description and docstring strings that exceed 88 columns
using implicit concatenation; string values are unchanged. Drop the unused
typing imports from prompt_search_list and mark the deliberately late node
imports in __init__ with noqa: E402.

* style: wrap long lines in the API layer, autotag and config

Split long f-string log messages, response messages, URL builders and
docstrings; DEFAULT_PROMPT and SYSTEM_PROMPT are rewritten as implicit
concatenations with identical values. Drop the unused shutil import in
admin and the unused HFCacheInfo import in autotag; keep the
hf_hub_download/snapshot_download availability probe with noqa: F401 and
mark the deliberately late imports in config with noqa: E402.

* style: wrap long lines and drop unused imports in utils

Re-wrap module docstrings, comments and f-string log messages; remove the
unused typing.List and PngInfo imports; mark the PIL availability probe in
diagnostics with noqa: F401; turn two placeholder-less f-strings into plain
strings; strip trailing whitespace inside the diagnostics SQL.

* style: wrap long lines in tests

Split long docstrings and SQL literals (values unchanged) and drop the
unused os import from test_path_security.

* chore: flake8 ignores the gitignored standalone tagger script

* fix: escape image dimensions and workflow JSON before rendering

openGalleryModal interpolated width, height and file_size (stored from a
POST body, so possibly strings) straight into data-caption and the card
text. showWorkflowData wrote JSON.stringify(workflow) raw into a blank
same-origin window, so a workflow containing </pre><script> ran. Both now
go through ImageHelpers (new pure module, node:test covered); a string
file_size no longer reaches formatFileSize's toFixed either, and a blocked
popup shows a notice instead of throwing.

* fix: escape remaining server values rendered through innerHTML

Maintenance and diagnostics results, review-tag errors, video info and
metadata error messages, consistency issues, recent image filenames,
output directory names, the prompt rating attribute, gallery media types
and duplicate group headers were interpolated unescaped. The copy
fallback modal also inlined the prompt text inside <textarea>, which a
closing tag could break out of; it now sets textarea.value.

* fix: address gallery auto-tag review images by path, not digest id

The gallery review flow builds its list from /prompt_manager/images/output,
whose id is a 16-hex path digest, but posted it as image_id to
autotag/single, so the server looked up a database row and every request
404ed. autotagSingleBody (pure, tested) now sends path for output-folder
entries and keeps image_id for database rows, which the admin review flow
(/images/all) still uses.

* fix: read thumbnail generation progress over POST with SseStream

images/generate-thumbnails/progress is becoming POST-only, which the
browser EventSource cannot request. Both thumbnail progress consumers in
the gallery now use SseStream.connect(url, {method: 'POST'}) exactly as
the autotag streams do. SseStream gains addEventListener/removeEventListener
for named event: frames (tested) so the status, start, progress,
file_error, complete and error listeners stay as they were, and an onclose
guard rejects when the stream ends without a terminal frame.

* fix: drop the unused copy-text data-action handler

Nothing in the admin markup emits data-action="copy-text"; the prompt
copy buttons call copyToClipboard directly.

* fix(lora): confine LoraManager paths to custom_nodes and hide server paths

- is_civitai_url also refuses URLs with userinfo or a non-443 port, so
  https://evil.com\@civitai.com/x and https://civitai.com:8443/x are rejected.
- resolve_lora_manager_path / custom_nodes_directories: lora/enable and
  lora/scan only accept a LoraManager directory that is a direct child of
  ComfyUI's custom_nodes (folder_paths base_path or registered custom_nodes
  folders, else this package's parent) and whose name contains "lora",
  compared via realpath+normcase; relative public paths resolve against the
  API path anchors. Anything else is 400, so find_lora_directories can no
  longer widen the image allow-list with arbitrary trees.
- lora/enable persists through PromptManagerConfig.save_to_file(), which
  honours PROMPT_MANAGER_CONFIG_PATH, instead of a hard-coded config.json.
- detect/status/enable return path and detected_path through _public_path;
  every lora handler reports exceptions via _public_error, and lora/scan
  errors raised after the SSE stream opened become an error event.
- Tests close the database before unlinking it.

* fix(autotag): contain batch rows, require media extensions, bound the scan

- The batch tagger runs every generated_images row through
  _resolve_allowed_image_path; rows that exist but lie outside the served
  output directories are skipped and counted in skipped_outside_roots on
  progress and complete events, and the engine receives the resolved path.
- _resolve_allowed_image_path requires a media extension (.png .jpg .jpeg
  .webp .gif .bmp .tiff .tif), so autotag/single returns 403 for other files.
- autotag/start holds a module-level lock; a concurrent start gets 409
  {"success": false, "error": "Auto-tagging is already running"}.
- scan_output_dir walks with a 50,000-file and 12-level cap, skips symlinked
  files and thumbnails, returns path through _public_path, relative_path in
  POSIX form, and a truncated flag.
- Autotag handlers report exceptions via the shared _public_error; the
  local copy is removed. Tests close the database before unlinking it.

* test: close every database before unlinking and anchor public paths in image tests

Each owned test that opens a PromptDatabase now calls close_all() before
removing the file, and temp directories holding a database use
ignore_cleanup_errors, so the suite tears down cleanly on Windows. The
image API fixture lays out a fake ComfyUI tree (<tmp>/output) and pins
GalleryConfig.path_anchors to it so responses can be checked for
root-relative paths.

* fix: validate image file_info before it reaches the generated_images columns

link_image_to_prompt stored size, dimensions and format straight from the
request body, so a string such as '"><img src=x>' ended up in INTEGER
columns and was echoed back to the gallery. Size and dimensions are now
coerced with int() and bounded (size 0..10^12, width/height 0..65535),
format is matched against a whitelist of media formats, and anything
else is stored as NULL while the image still links. Non-dict metadata or
file_info no longer aborts the link.

* fix: clamp paging on tag listings and the page-derived recent offset

tags/stats, tags/{tag}/prompts and tags/filter parsed limit/offset with a
bare int(), so limit=999999 was passed straight to SQLite and a huge
offset raised OverflowError (500). They now share parse_page_params
(limit 1..500, offset 0..MAX_PAGE_OFFSET, 400 on non-integers), and
/recent clamps the offset it derives from ?page= the same way.

* fix: make the thumbnail progress stream POST-only

GET /prompt_manager/images/generate-thumbnails/progress wrote thumbnails
to disk, so any prefetch, link preview or <img src> could start a full
output-directory scan. The route now accepts POST only (GET answers
405); quality still comes from the query string so the SSE client can
keep the same URL shape.

* fix: describe images relative to the ComfyUI tree in API responses

Gallery and prompt responses carried absolute server paths: /images/output
entries (path, root_dir), thumbnails_path, the images/prompt lookup, every
image list (image_path plus its normalised file_path duplicate) and the
preview images attached to prompt listings. They now go through
_public_path: /images/output 'path' is relative to its root (what the
delete and autotag endpoints resolve against), 'root_dir' and every
other image_path are relative to the ComfyUI base, and file_path is
dropped. recent/all/search image lists also gain url/thumbnail_url via
the same enrichment the prompt-images route already had. Remaining
str(e) error bodies in both modules use the basename-only error text.

* fix: check thumbnail directory containment with a real path test

clear_thumbnails compared resolved paths with str.startswith, so a
thumbnails directory that resolved to a sibling with a shared prefix
(out vs out2) passed the safety check and was emptied. It now uses the
module's _is_within (realpath + normcase + Path.is_relative_to).

* fix: decode tag names from the URL only once

The tag lookup, rename and delete routes ran urllib.parse.unquote on
match_info values that aiohttp had already percent-decoded, so a tag
literally named '%41' (requested as %2541) was handled as 'A'. The
second decode is dropped.

* fix: validate tags on every tag mutation route

add_prompt_tag, add_tags_to_prompt, bulk_add_tags, merge_tags and
rename_tag accepted whatever the body carried; a non-string tag raised
AttributeError (500) and over-long or control-character tags were
stored. They now run utils.validators.validate_tags through a shared
tags_error helper and answer 400 with the validator's message. Blank
entries in add_tags_to_prompt are still skipped rather than rejected.

* fix: export every prompt by paging instead of truncating at 10000

/prompt_manager/export fetched a single page of 10000 prompts and
silently dropped the rest. It now walks the table in EXPORT_PAGE_SIZE
pages (1000) inside one executor job until a short page is returned.

* fix: gate the LoraManager serving roots and accept relative link paths

A LoraManager path hand-edited into config.json went straight into
find_lora_directories and so into the image-serving allow-list. When a
path is configured it must now pass lora_utils.resolve_lora_manager_path
(a custom_nodes child whose name contains 'lora'); otherwise, or when the
resolver is unavailable, no LoRA directories are served. An empty path
keeps the custom_nodes-only auto-detection.

POST /images/link now also accepts the relative forms the API hands out
(root-relative 'sub/a.png' and ComfyUI-relative 'output/sub/a.png'),
resolving them against the allowed roots and anchors; '..' segments in
a relative path are refused and unknown files are 404.

* fix: review flow uses the served image url and the lora-dirs test pins its config

Image paths in API responses are now relative to the ComfyUI tree, so the
auto-tag review page can no longer derive a serve URL from image_path; the
server always provides url. The lora-dirs serving test assumed an empty
LoraManager path, which only held in a checkout without a local config.json.

* test: stop leaking a mocked utils.comfyui_integration and mock the IS_CHANGED clock

test_image_monitor_singleton replaced utils.comfyui_integration in
sys.modules with a MagicMock and never restored it, so running it before
test_comfyui_integration failed eight tests. The mock was unused; drop it
and install the server stub with setdefault in every owned test module.

The PromptSearchList IS_CHANGED test slept 10 ms between two time.time()
reads; the Windows clock ticks every 15.6 ms, so it was flaky there. Mock
the clock with an incrementing counter instead.

* fix: resolve_text follows a node shared by two branches

The cycle guard used one visited set for the whole walk, so a string node
read twice (a diamond, or StringConcatenate with the same source on both
inputs) made the second branch None and the prompt text unknowable. Cycle
detection is now scoped to the current path; a separate visit counter
keeps MAX_VISITS as the bound on total work.

* fix: bound resolve_text on diamond chains with memoisation, depth and length caps

A per-path visited set alone makes a chain of diamonds (node i reads node
i+1 on both inputs) exponential in depth. Each node's text is now
memoised so it is expanded once, every expansion is charged to the shared
MAX_VISITS budget, chains deeper than MAX_DEPTH (64) are unknowable so
recursion stays bounded, and text longer than MAX_TEXT_LENGTH (100k) is
unknowable so StringConcatenate doubling cannot exhaust memory.

* fix: static routes resolve under web/ with an extension allow-list

/prompt_manager/lib/{filepath} and /js/{filepath} only rejected '..' and a
leading '/'. On Windows os.path.join(root, 'C:/x') or a leading backslash
re-anchors the path, so any file on the host could be read. Requests are
now refused when the path is absolute on any platform (slash, backslash,
drive letter, UNC prefix), carries an empty, '.' or '..' segment, has an
extension outside the allow-list, or resolves (symlinks included) outside
the static directory; the resolved real path is served with its MIME type.

* feat: capped JSON body reader and single-flight lock for long admin jobs

Add _read_json_body(request, max_bytes=1_000_000) to py/api: reads the
body in chunks, answers 413 past the cap (declared or streamed) and 400
for invalid or non-object JSON. Every JSON POST in the admin and logging
handlers goes through it; it is exposed as PromptManagerAPI._read_json_body
for the other route mixins.

Add _acquire_long_job(name) / _long_job_busy_response(label) on the API
object: one asyncio.Lock per job family, handed out as an async context
manager or None when busy, so a handler answers 409 instead of starting a
second scan. The duplicate scan and delete share the 'duplicates' family.

* fix: harden database restore and keep server paths out of admin responses

Restore: verify_database_file and restore_from_file run through
_run_in_executor; the upload is also opened read-only with
PRAGMA trusted_schema=OFF and refused when sqlite_master lists a trigger
or a view (a prompts database has neither, a crafted one could use them
to run SQL against the live data); at most MAX_SAFETY_BACKUPS (10)
'<db>.backup_*' files are kept next to the live database, oldest first.

Responses: backup_created and the maintenance statistics database_path
go through _public_path, and backup/restore failure messages through
_public_error, so no absolute server path appears in an admin body.

* fix: validate hand-edited gallery roots everywhere and bound output scans

A root written straight into config.json skipped the containment check
the settings endpoint applies. GalleryConfig.update_config (and so
load_from_file) now drops entries that fail validate_gallery_root with a
warning, _find_comfyui_output_dir validates the configured directory
before using it, and ImageMonitor.start_monitoring validates every
directory before scheduling the observer.

The duplicate cleaner only deletes files with a media extension and
never anything under a thumbnails directory. _collect_media_files walks
with os.walk(followlinks=False), skips symlinked files and directories,
prunes thumbnails, stops MAX_SCAN_DEPTH (12) levels down and after
MAX_SCAN_FILES (50000) files.

* fix: bound log file size and backup count in the logging config endpoint

max_file_size must be between 1 MB and 1024 MB (bytes) and backup_count
between 0 and 50; anything else is a 400. Also close every database
connection in the logging test teardown.

* fix: private config.json, merged settings save, synthetic test-link path

PromptManagerConfig.save_to_file and the settings endpoint write
config.json through write_private_json: created with mode 0600 (chmod on
POSIX, best effort on Windows) because the file can hold an API key.
_persist_settings merges web_ui and gallery.monitoring.directories into
the existing file instead of overwriting it, so the integrations section
and any other key survive a settings save.

The test-link diagnostic no longer stores a client-supplied path; it
records a marker file name under the output directory.

* fix: image worker thread exits on shutdown and restarts cleanly

ImageGenerationHandler.shutdown() clears the pending set, sets a stop
event the worker's delay and settle waits listen to, and puts a None
sentinel on the queue; the worker exits within the current wait. Each
worker owns its queue and stop event so a draining worker never swallows
an image scheduled after shutdown. ImageMonitor.stop_monitoring calls it.

* test: close all database connections before deleting temp files

Windows refuses to unlink an open SQLite file and executor threads keep
their own connections, so every owned test that creates a PromptDatabase
calls close_all() before cleanup; the diagnostics temp dir ignores
cleanup errors.

* fix: output scan runs as a background job that survives the browser going away

Closing the dashboard tab or clicking outside the scan modal raised
ClientConnectionResetError out of scan_images and abandoned the scan
mid-batch, leaving a half-scanned database. Two overlapping scans could
also run at once.

- py/api/scan_job.py: ScanJob fans progress out to any number of SSE
  subscribers (latest snapshot first, then live events)
- POST /prompt_manager/scan starts the scan as an asyncio task, or attaches
  to the one already running; a write to a dead socket ends that stream
  quietly while the scan continues
- GET /prompt_manager/scan/status reports the running flag and last event
- dashboard: dismissing the modal while a scan runs minimizes it to a
  progress pill, the page re-attaches to a running scan after a reload,
  a lost stream reconnects, and leaving the page mid-scan asks first

* feat: worker-thread setting speeds up the output scan and thumbnail generation

Both jobs processed one file at a time on a single core. Pillow decodes the
whole image to reach trailing PNG text chunks and releases the GIL while it
does, so a bounded thread pool scales almost linearly: on a 1 GB sample of
400 PNGs the metadata pass dropped from 18 s to 2.5 s at 8 threads, and 120
thumbnails from 4.6 s to 0.6 s.

- utils/parallel.py: map_parallel(func, items, workers) keeps order, runs on
  the caller's thread when workers is 1, and caps the pool at the item count
- PromptManagerConfig.WORKER_THREADS defaults to half the detected cores,
  is clamped to max_worker_threads() when loaded from config.json, and is
  persisted under "performance"
- settings API returns worker_threads and cpu_count and validates saves;
  the settings modal gets a "Worker threads" slider capped at the core count
- the scan reads each metadata batch on the configured threads; thumbnail
  generation processes batches of files per executor job, still emitting
  one progress event per file

Threads rather than processes: this runs inside ComfyUI's server, where
forking a CUDA-initialised process is unsafe and the spawn start method
would re-import ComfyUI's main module in every child.

* fix: scan modal ignores backdrop clicks so a focus click cannot dismiss it

Only its own buttons dismiss it: Cancel before a scan starts, Hide (minimize
to the progress pill) while one runs.

* feat: scan reads prompts embedded in videos; unreadable files no longer log errors

The output scan handed every .mp4 to Pillow and logged one ERROR per file.
ComfyUI's video savers store the queued graph as JSON in the container's
comment tag, so those files carry prompts worth linking.

- utils/video_metadata.py reads the comment/description tag through
  ffprobe with a timeout, normalised to the PNG text-chunk shape so the
  existing parser handles both; without ffprobe videos are skipped and a
  single INFO notice is logged
- _extract_comfyui_metadata routes videos there and logs unreadable images
  as warnings, since the scan carries on past them

Test hygiene: every aiohttp test base class overrode tearDownAsync without
calling super(), so no test client was ever closed. That produced thousands
of ResourceWarnings, "Unclosed client session" errors from asyncio that
broke log assertions, and doubled the suite's run time.

* fix: thumbnail runs skip fresh thumbnails before the cap and report what remains

The per-run cap of 10,000 applied to the directory walk itself, so a
library with more files than that never had thumbnails generated for the
rest, and a run that found 9,987 existing thumbnails spent its whole
budget on skipping them.

- _plan_thumbnails walks the whole tree (bounded like the gallery listing),
  counts files whose thumbnail is already newer as skipped without queuing
  them, and caps only the pending files; the leftover count is returned
- start and complete payloads carry skipped and remaining; total_images
  counts every file considered and the processing rate uses the files
  actually worked on
- the gallery page shows a warning naming the pending count instead of a
  success toast when a run stopped at the cap

* feat(gallery): warn with the pending count when a thumbnail run stops at the cap

* fix: thumbnails are written atomically and served under mtime-versioned urls

Cards showed "Failed to load" for thumbnails that were valid on disk. The
browser had fetched them while the thumbnail run was still writing the
file and, because image responses are cached for an hour, kept the broken
bytes. The card also had no fallback to the original.

- image and video thumbnails are written to a hidden temp sibling and
  renamed into place, so a partial file is never visible at the final path
- every thumbnail_url (gallery listing, prompt images, duplicates, autotag
  listing) carries ?v=<mtime>, so a regenerated thumbnail gets a new URL
  instead of the cached old one
- a 404 from the image serve route is marked no-store; ComfyUI's cache
  middleware would otherwise let browsers cache the miss for an hour
- image-helpers gains fallbackImageSource(): the original is tried once
  before a card gives up

Test hygiene: the autotag tests now stub ComfyUI's server module like the
other API test modules, so the module passes when run on its own.

* fix(gallery): display settings apply at load, columns work, media info toggles

Default View Mode, Items per Page and Grid Columns were only applied when
Save was pressed, never when the page loaded; the column setting swapped a
Tailwind class that was never compiled (only 3, 6 and 8 exist) and a wider
breakpoint class overrode it anyway; Show Media Info was saved and never
read.

- gallery-settings.js validates what localStorage holds (view mode, limit
  option, column range, booleans) and falls back to defaults per value
- settings are applied in the constructor before the first load; changing
  the page size reloads, toggling media info re-renders
- the column count sets an explicit grid template through a CSS variable
  that wins on desktop widths
- media info lines honour the toggle in grid and list views
- image load failures are handled by one delegated listener that tries the
  original before showing the failed-to-load card

* feat(gallery): duplicate scan runs as a background job with progress; infinite scrolling

The duplicate scan was one blocking request with a spinner: no progress,
a modal that could not be put away, and a result lost if the tab closed.

- POST /prompt_manager/scan_duplicates/stream starts the scan as a
  background job (same ScanJob fan-out as the output scan) or attaches to
  the running one; GET /prompt_manager/scan_duplicates/status reports it.
  The job holds the duplicates lock, so deletes are refused while it runs.
- hashing runs on the configured worker threads and reports progress
  about every percent
- the gallery modal shows a progress bar, closing it minimizes to a pill
  that reopens it, results wait in the pill if the modal was closed, and a
  reload re-attaches to a running scan

Infinite scrolling is a new Display Setting (off by default): the next page
is appended when the viewport nears the bottom, pagination controls hide,
and the load-more rule lives in list-state.js where it is unit tested.
Thumb-size buttons and the column setting no longer override each other.

* fix(gallery): Rescan Output Folder drops the server listing cache and rewalks

The button only re-requested the listing, which the server answers from a
30-second in-memory cache, so a rescan showed exactly what was already on
screen.

- POST /prompt_manager/gallery/rescan invalidates the gallery cache,
  rewalks every output root and returns the file and root counts
- the button calls it, closes the settings modal, reloads the folder tree
  and the first page, and reports what it found

* feat: tag search box autocompletes, narrowed by the tags already typed

- PromptDatabase.suggest_tags(prefix, with_tags, limit): tags starting
  with the prefix, counted over the prompts that carry every context tag,
  context tags excluded, ordered by count then name; LIKE wildcards in the
  prefix are escaped
- GET /prompt_manager/tags/suggest?q=&with=&limit= exposes it (limit
  clamped to 1..50, context capped at 10 tags, registered before the
  {tag_name} routes)
- tag-autocomplete.js splits the comma-separated input into committed tags
  and the one being typed, builds the query, applies a pick and handles
  wrap-around keyboard selection
- the dashboard Tags box shows suggestions with counts on focus and while
  typing, arrow keys and Enter/Tab pick, Escape closes, and each pick
  immediately offers the next tag narrowed by the previous ones

* feat(dashboard): infinite scrolling for the prompt list, saved default plus session toggle

- PromptManagerConfig.INFINITE_SCROLL (web_ui.infinite_scroll) is reported
  and saved through the settings API; only booleans are accepted
- Settings modal gains an "Infinite Scrolling" default; saving applies it
  at once
- the results header gains an "Infinite scroll" checkbox that overrides the
  default for the current session without saving
- when on, the next page is appended as the list nears its bottom (same
  load-more rule as the gallery), page controls hide, a short page keeps
  loading until the list fills, and a refresh after an edit restarts from
  the top

* fix: first CI run on Windows and macOS; tests now enforce closed database files

The first push of the branch ran the Tests matrix on all three operating
systems. Every Python job failed; the causes were:

- numpy missing in CI: the WD14 engine tests import it (requirements-dev)
- macOS: temp directories resolve through /private; the admin API tests
  now use the real path like the gallery roots they compare against
- Windows: the startup banner printed an emoji, which raises
  UnicodeEncodeError on a legacy-code-page console and aborted loading
  the whole node; the banner is ASCII now and a test loads the package
  under a cp1252 stdout
- Windows: gallery subfolders came back with backslashes, which neither
  the folder tree nor the search filter understand; they are "/" joined
  on every platform. A relative DEFAULT_DB_PATH kept its forward slashes
  after being joined to the extension root; it is normalised now
- Windows: sqlite connections opened with `with sqlite3.connect()` are
  not closed by the context manager and linger until garbage collection,
  which holds a file lock there. Production sites (vacuum, diagnostics,
  restore inspection) close explicitly; a restore that swaps the API's
  database object closes the previous one; PromptModel keeps a registry
  so a shutdown or a test can close every instance
- Windows: served files were still open when their temp directory was
  removed (tests now read the body), a log download test compared text
  mode newlines, two tests compared paths with a hard-coded separator

tests/open_handles.py makes the Linux run fail the way Windows would:
every database test base checks, right before deleting its files, that
nothing in the process still holds them open.

* fix: Windows case-folded gallery roots in tests, "/" subfolders everywhere, longer busy timeout

Second CI run on Windows:
- gallery roots are stored resolved and case-folded there; the admin API
  tests now compare against that form instead of the mixed-case temp path
- the subfolder fallback for images outside every root also uses "/", and
  the test expects it
- a relative diagnostics db path is normalised after joining
- "database is locked" under many writers on a slow disk: connections now
  wait up to 30 seconds (connect timeout and busy_timeout), long enough for
  a scan or thumbnail run while ComfyUI saves images

* fix: serve every gallery video format, validate settings before applying, dedupe output dirs

Review findings on #153 (Codex):
- the image serve route allowed only three video types while the gallery
  lists seven, so .avi/.mkv/.m4v/.wmv cards opened to a 403; the serve
  allow-list now covers every gallery video format
- a settings payload with one invalid field still applied the valid ones
  before returning 400; the whole payload is validated first and nothing
  changes on failure. webui_display_mode is validated too

Third Windows CI run:
- the same output directory spelled twice (different case or a link) was
  listed twice in diagnostics and served roots; both dedupe on the
  resolved, case-folded path
- the root-less subfolder test expects the "/" form

Code-quality bot: index definitions are (name, target) pairs instead of
implicitly concatenated strings; redundant local asyncio imports and
two trivial lambdas in test stubs removed.

* docs: version 3.3.0, README changelog and refreshed sections

Bump to 3.3.0: new features (worker threads, background scans, tag
autocomplete, infinite scrolling, video prompt scanning) and API contract
changes (POST-only side-effect routes, relative image paths, paginated
search) go beyond a patch release.

README: changelog entries for 3.2.4 and 3.3.0, refreshed bulk scanning,
gallery settings, search, configuration and testing sections, and the
roadmap's recently-completed list.
2026-10-08 18:53:11 -07:00
Vito 6c7e4d5b35 feat: release 3.2.4 — recently used sorting, run counts, positive-prompt image linking (#152)
* chore: start 3.2.4

* feat: track prompt usage (run_count, last_used_at) with server-side sorting

- Migration adds last_used_at and run_count; existing prompts are estimated
  from linked images (count, newest image time, else created_at). The
  backfill only touches NULL rows, so it self-heals rows written after a
  downgrade or an interrupted upgrade.
- Every node run records a use, including re-runs of an existing prompt;
  a tracking failure is logged and never blocks generation.
- search_prompts/get_recent_prompts accept a whitelisted sort key
  (last_used, created, run_count, rating, text); default stays newest first.

Refs #90

* feat: sort param on /recent and /search; microsecond last_used_at

- /prompt_manager/recent and /search accept ?sort= (whitelisted in the DB
  layer; unknown values fall back to newest first).
- last_used_at is written with microseconds from Python so a re-run in the
  same millisecond as another save still sorts first; ordering stays correct
  against millisecond values from the backfill.

Refs #90

* feat: server-side sort in admin UI with Recently Used / Most Used and run badges

- New web/js/prompt-list-sort.js owns the sort options, validation and
  request URLs; admin.js fills the dropdown from it and re-fetches from the
  server on change, so every sort spans all pages (the old client-side sort
  only reordered the current page).
- Default sort is Recently Used; prompt cards show a run-count badge with
  the last-used time on hover.
- The node's Recent button now lists recently used prompts.
- Contract test keeps the UI sort keys and the node's key in sync with the
  server whitelist.

Refs #90

* fix: link images and count runs by the positive prompt, not node execution

With a negative-prompt PromptManager node in the workflow, every executed
PromptManager node pushed onto the image-linking queue, so images could be
linked to the negative prompt or to a previous run. ComfyUI also skips
unchanged nodes, so re-runs of the same prompt were never counted.

- utils/prompt_graph.py finds the PromptManager nodes feeding a sampler's
  positive input (KSampler, CFG/Basic guiders, ControlNet and combine
  pass-throughs, PromptManagerText into encoders), falling back to all
  PromptManager nodes when none can be traced.
- Runs are counted by a ComfyUI on-prompt hook when a workflow is queued;
  prompts not yet saved are counted when the node first saves them.
- Nodes receive the queued graph (hidden PROMPT/UNIQUE_ID). Only the run's
  positive prompt becomes the current prompt; only batch items with linked
  text are pushed onto the FIFO queue.
- The image monitor links from the image's own metadata first, using the
  positive node, and falls back to the queue for batch items.

Refs #90

* fix: address review — count cached runs of linked text, keep repeated pending runs

- Text from known pure string nodes (core PrimitiveString/StringConcatenate,
  WAS Text Concatenate/Multiline, pysssss ShowText) is resolved at queue
  time, so cached re-runs of such prompts are counted and their images link
  from metadata. Unknown nodes are never guessed; the node compares its real
  text with the queue-time resolution and falls back to counting and
  queueing itself on any mismatch.
- A new prompt queued several times before its first run keeps every
  pending run; the first save records all of them.
- Registration flag is a threading.Event; consistent unittest imports;
  logging instead of print in the init fallback.

* refactor: unquote PendingFirstUse annotation so Tuple usage is visible
2026-10-07 13:52:09 -07:00
Vito aa768f57dd fix: release 3.2.3 — metadata parsing, null-safe saves, auto-hide, tag search (#151)
* fix: release 3.2.3 — metadata parsing, null-safe saves, auto-hide, tag search

- fix(#75): stop SaveImage patch from writing the last-run prompt into every
  PromptManager node (negative prompts were saved as copies of the positive)
- fix(#75): replace four duplicated metadata parsers with a shared
  web/js/comfy-metadata.js that traces sampler links (KSampler, CFGGuider /
  SamplerCustomAdvanced, LoRA chains, ControlNet pass-through, any checkpoint
  loader) and recovers per-node text for images saved by older versions
- fix(#143): treat null JSON body fields as empty in prompt API handlers
- fix(#149): keep auto-hide timeout 0 (disabled) instead of falling back to
  the default; validate result_timeout on save
- feat(#150): support tag:name filters in the node search box
- docs(#142): scan dialog mentions A1111/Forge metadata support
- docs(#144): fix ComfyUI-Lora-Manager links
- ci: run JS tests with node --test

* fix: bound metadata graph traversal to prevent tab freeze on crafted PNGs

Link resolution tried every linked input recursively with only a depth cap,
so a PNG whose prompt graph has cycles and fan-out could force ~inputs^25
paths and hang the gallery/metadata page. Each lookup now shares a visited
set (each node expanded once) with a hard cap of 500 visits.

* fix: read the matching output of multi-output settings nodes

KSampler settings linked to a node like SamplerCombo all resolved to its
first value, so steps, cfg and sampler showed the sampler name. Scalar
resolution now prefers the linked node's input with the same field name.

* fix: address review — trust API text, restore workflow-only parsing

- Workflow widget text now only replaces API text for metadata showing the
  pre-3.2.3 corruption signature (2+ PromptManager nodes with identical API
  text but different widgets). Otherwise the API graph wins, so expanded
  dynamic prompts and linked text display what actually ran.
- Images with only a UI workflow are converted to the API graph shape and
  traced the same way: prompts, KSamplerAdvanced, UNET/checkpoint loaders,
  KJNodes Set/Get, Reroute, bypassed and muted nodes.
2026-10-07 11:27:40 -07:00
152 changed files with 28941 additions and 5095 deletions
+4
View File
@@ -14,6 +14,10 @@ omit =
*/node_modules/*
*/static/*
*/media/*
*/htmlcov/*
*/.claude/*
restart_gallery.py
standalone_tagger.py
branch = true
parallel = true
+6
View File
@@ -0,0 +1,6 @@
[flake8]
max-line-length = 88
extend-ignore = E203, W503
exclude = reference/, .git, __pycache__, web/, venv/, htmlcov/, build/, dist/, .claude/, docs/, standalone_tagger.py
per-file-ignores =
tests/*:E402
+25
View File
@@ -0,0 +1,25 @@
{
"required_status_checks": {
"strict": true,
"contexts": [
"py3.10 on ubuntu-latest",
"py3.11 on ubuntu-latest",
"py3.12 on ubuntu-latest",
"py3.10 on windows-latest",
"py3.11 on windows-latest",
"py3.12 on windows-latest",
"py3.10 on macos-latest",
"py3.11 on macos-latest",
"py3.12 on macos-latest",
"coverage gate (85%)",
"black and flake8",
"frontend modules",
"viewerjs"
]
},
"enforce_admins": true,
"required_pull_request_reviews": null,
"restrictions": null,
"allow_force_pushes": false,
"allow_deletions": false
}
+12
View File
@@ -46,3 +46,15 @@ updates:
update-types:
- "minor"
- "patch"
# Maintainer browser test dependencies; Viewer.js itself has an asset-aware updater.
- package-ecosystem: "npm"
directory: "/tools/vendor"
schedule:
interval: "weekly"
allow:
- dependency-name: "@playwright/test"
labels:
- "dependencies"
commit-message:
prefix: "deps(vendor-tools)"
+43 -10
View File
@@ -11,9 +11,12 @@ on:
jobs:
test:
runs-on: ubuntu-latest
name: py${{ matrix.python-version }} on ${{ matrix.os }}
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
python-version: ["3.10", "3.11", "3.12"]
steps:
@@ -23,14 +26,8 @@ jobs:
uses: actions/setup-python@v7
with:
python-version: ${{ matrix.python-version }}
- name: Cache pip dependencies
uses: actions/cache@v6
with:
path: ~/.cache/pip
key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('requirements-dev.txt') }}
restore-keys: |
${{ runner.os }}-pip-${{ matrix.python-version }}-
cache: pip
cache-dependency-path: requirements-dev.txt
- name: Install dependencies
run: |
@@ -38,9 +35,45 @@ jobs:
pip install -r requirements-dev.txt
- name: Run tests
run: python -m unittest discover tests/ -v
run: python -m unittest discover -s tests -v
coverage:
name: coverage gate (85%)
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v7
- uses: actions/setup-python@v7
with:
python-version: "3.12"
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: Measure coverage
run: |
python -m coverage run -m unittest discover -s tests
python -m coverage combine
python -m coverage report --fail-under=85
lint:
name: black and flake8
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v7
- uses: actions/setup-python@v7
with:
python-version: "3.12"
- run: pip install black flake8
- run: black --check .
- run: flake8 .
js-test:
name: frontend modules
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v7
+47
View File
@@ -0,0 +1,47 @@
name: Viewer.js vendor and compatibility
on:
pull_request:
branches: [main]
push:
branches: [main]
workflow_dispatch:
permissions:
contents: read
jobs:
viewerjs:
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@v7
- uses: actions/setup-python@v7
with:
python-version: '3.12'
cache: pip
cache-dependency-path: requirements-dev.txt
- run: python -m pip install -r requirements-dev.txt
- uses: actions/setup-node@v7
with:
node-version: '24'
cache: npm
cache-dependency-path: tools/vendor/package-lock.json
- run: npm ci --prefix tools/vendor --ignore-scripts
- run: npm run --prefix tools/vendor check
- run: npm run --prefix tools/vendor test
- run: npx --prefix tools/vendor playwright install --with-deps chromium firefox webkit
- name: Frontend unit, viewer, workflow, and coverage checks
run: npm run --prefix tools/vendor test:coverage
- name: Firefox and WebKit page workflows
run: npm run --prefix tools/vendor test:workflows -- --project firefox --project webkit
- name: Upload browser traces and frontend coverage
if: always()
uses: actions/upload-artifact@v7
with:
name: frontend-test-reports
path: |
tools/vendor/test-results/
coverage/js/
!coverage/js/raw/**
retention-days: 14
+83
View File
@@ -0,0 +1,83 @@
name: Update vendored Viewer.js
on:
schedule:
- cron: '30 7 * * *'
workflow_dispatch:
inputs:
version:
description: 'Stable version to vendor; leave blank for the latest npm release'
required: false
type: string
permissions:
contents: read
concurrency:
group: viewerjs-update
cancel-in-progress: false
jobs:
update:
runs-on: ubuntu-latest
permissions:
contents: write
pull-requests: write
actions: write
env:
GH_TOKEN: ${{ github.token }}
steps:
- uses: actions/checkout@v7
with:
ref: ${{ github.event.repository.default_branch }}
- uses: actions/setup-node@v7
with:
node-version: '24'
cache: npm
cache-dependency-path: tools/vendor/package-lock.json
- run: npm ci --prefix tools/vendor --ignore-scripts
- name: Update the dependency and regenerate assets
env:
VIEWER_VERSION: ${{ inputs.version }}
run: |
if [ -n "$VIEWER_VERSION" ]; then
npm run --prefix tools/vendor update -- --version "$VIEWER_VERSION"
else
npm run --prefix tools/vendor update
fi
- name: Detect changes
id: changes
run: |
if git diff --quiet -- tools/vendor/package.json tools/vendor/package-lock.json web/lib/viewerjs; then
echo 'changed=false' >> "$GITHUB_OUTPUT"
else
echo 'changed=true' >> "$GITHUB_OUTPUT"
fi
- if: steps.changes.outputs.changed == 'true'
run: |
npm run --prefix tools/vendor check
npm run --prefix tools/vendor test
node --test tests/js/*.test.js
npx --prefix tools/vendor playwright install --with-deps chromium
npm run --prefix tools/vendor test:browser
- name: Open the update PR
if: steps.changes.outputs.changed == 'true'
env:
BASE_BRANCH: ${{ github.event.repository.default_branch }}
run: |
version=$(node -p "require('./web/lib/viewerjs/provenance.json').version")
branch="deps/viewerjs-$version"
existing=$(gh pr list --repo "$GITHUB_REPOSITORY" --state all --head "$branch" --json number --jq 'length')
if [ "$existing" != '0' ]; then
echo "A PR already exists for Viewer.js $version; leaving it for review."
exit 0
fi
remote=$(git ls-remote origin "refs/heads/$branch")
if [ -z "$remote" ]; then
git switch -c "$branch"
git add tools/vendor/package.json tools/vendor/package-lock.json web/lib/viewerjs
git -c user.name='github-actions[bot]' -c user.email='41898282+github-actions[bot]@users.noreply.github.com' commit -m "deps: update vendored Viewer.js to $version"
git push origin "HEAD:refs/heads/$branch"
fi
gh pr create --repo "$GITHUB_REPOSITORY" --base "$BASE_BRANCH" --head "$branch" --title "deps: update vendored Viewer.js to $version" --body-file tools/vendor/update-pr.md
gh workflow run viewerjs-check.yml --repo "$GITHUB_REPOSITORY" --ref "$branch"
+14
View File
@@ -9,6 +9,12 @@ AGENTS.md
CLAUDE.md
# Local files
# Local PromptManager state (config.json holds API keys and machine paths)
config.json
resume.sh
*.backup_*
promptmanager_logs/
custom_data/
data/
docs/
logs/
@@ -194,3 +200,11 @@ cython_debug/
# PyPI configuration file
.pypirc
# Maintainer-only JavaScript tooling and browser test artifacts
node_modules/
test-results/
playwright-report/
# Combined frontend coverage reports
coverage/
+1 -1
View File
@@ -9,7 +9,7 @@ repos:
- id: flake8
name: flake8
entry: flake8 --select=E9,F63,F7,F82 --show-source --exclude=reference/,web/
entry: flake8 --show-source
language: system
types: [python]
+58 -8
View File
@@ -139,6 +139,9 @@ The comprehensive gallery system provides complete management of all ComfyUI out
_Powered by [ViewerJS](https://github.com/fengyuanchen/viewerjs) for professional image viewing capabilities._
Viewer.js is vendored locally. See [Viewer.js maintenance](tools/vendor/README.md)
for reproducible updates, verification, and automated update PRs.
### 🔬 ComfyUI Workflow Analysis
- **PNG Metadata Parsing**: Advanced PNG chunk analysis to extract embedded ComfyUI data
@@ -164,7 +167,7 @@ _Powered by [ViewerJS](https://github.com/fengyuanchen/viewerjs) for professiona
- **Full-text search** across all stored prompts
- **Category filtering** for organized prompt collections
- **Tag-based search** with support for multiple tags
- **Tag-based search** with support for multiple tags, with autocomplete that narrows to tags seen together with the ones already typed
- **Rating filters** to find your best prompts
- **Date range filtering** for temporal searches
- **Recent prompts** quick access
@@ -306,9 +309,12 @@ The comprehensive gallery system provides complete access to all your ComfyUI ou
Configure your viewing experience in the settings panel:
- **Performance Settings**: Grid columns, image quality, lazy loading
- **Display Settings**: Default view, items per page, grid columns, media info, infinite scrolling
- **Performance Settings**: Image quality, lazy loading
- **Video Settings**: Autoplay, mute by default, loop videos
- **Thumbnail Management**: Generate optimized thumbnails for faster loading
- **Thumbnail Management**: Generate optimized thumbnails for faster loading (runs on the configured worker threads; a run that stops at the per-run cap tells you how many remain)
- **Find Duplicate Images**: Hashes every file on the worker threads with a progress bar; close the window and it keeps running, results wait in a pill
- **Rescan Output Folder**: Drops the server-side listing cache and rewalks the folders
#### **🖼️ Image Viewing Experience**
@@ -358,14 +364,16 @@ Use the standalone metadata viewer to analyze any ComfyUI-generated PNG:
Import existing ComfyUI images into your database:
1. **Open admin dashboard** and click **"📸 Scan Images"**
2. **Start scan** to analyze your entire output directory
3. **Monitor progress** with real-time status updates
2. **Start scan** to analyze every configured output directory
3. **Monitor progress** with real-time status updates. The scan runs on the server: closing the window only minimizes it to a progress pill, and reloading the page re-attaches to a running scan
4. **Review results** showing:
- Images found and processed
- Prompts extracted and linked
- Any errors or issues encountered
5. **Access imported data** through the normal gallery interface
The scan reads PNG text chunks and A1111-style parameters, and, when `ffprobe` is installed, the prompt ComfyUI's video savers embed in MP4/MOV/WebM files. Files are read on the number of **worker threads** set in Settings → Performance (default: half your cores), which makes a large library several times faster.
### 🏷️ AI AutoTag
Automatically tag your entire image collection using AI vision models:
@@ -664,7 +672,8 @@ CREATE TABLE generated_images (
- **`utils/diagnostics.py`** - System diagnostics and health checks
- **`web/admin.html`** - Advanced admin dashboard with metadata panel
- **`web/index.html`** - Simple web interface
- **`web/js/prompt_manager.js`** - Dashboard JavaScript
- **`web/comfy/prompt_manager.js`** - ComfyUI canvas extension (the only file ComfyUI loads)
- **`web/js/admin.js`**, **`web/js/gallery.js`** - Dashboard and gallery JavaScript
- **`web/js/tags-page.js`** - Tag management JavaScript
- **`web/metadata.html`** - Standalone PNG metadata viewer
@@ -699,8 +708,11 @@ ComfyUI_PromptManager/
│ ├── gallery.html # Image gallery
│ ├── index.html # Simple web interface
│ ├── metadata.html # Standalone metadata viewer
│ ├── comfy/
│ │ └── prompt_manager.js # ComfyUI canvas extension (WEB_DIRECTORY)
│ └── js/
│ ├── prompt_manager.js # Dashboard JavaScript
│ ├── admin.js # Dashboard JavaScript
│ ├── gallery.js # Gallery JavaScript
│ └── tags-page.js # Tag management JavaScript
├── tests/
│ ├── __init__.py
@@ -763,11 +775,15 @@ Configure the web interface behavior:
# Web UI settings (PromptManagerConfig class)
RESULT_TIMEOUT = 5 # Seconds to auto-hide results in ComfyUI node
WEBUI_DISPLAY_MODE = 'newtab' # 'popup' or 'newtab' for Web UI button
INFINITE_SCROLL = False # Dashboard list loads the next page on scroll by default
SHOW_TEST_BUTTON = False # Show API test button in node UI
```
All of these are editable in the dashboard Settings modal and saved to `config.json` (which is private to the installation and holds the CivitAI key, so it is never committed).
### Performance Tuning
- **Worker threads** (Settings → Performance): threads used by the image rescan, thumbnail generation and duplicate detection. Default is half the detected cores, maximum is the core count. Pillow releases the GIL while decoding, so the scan and thumbnails scale almost linearly; keep a few cores free while ComfyUI is generating.
- The database automatically creates indexes for optimal search performance
- Regular `VACUUM` operations keep the database optimized
- Consider backing up the database periodically
@@ -811,9 +827,14 @@ db.model.backup_database("backup_prompts.db")
```bash
cd ComfyUI_PromptManager
python -m pytest tests/ -v
pip install -r requirements-dev.txt
python -m unittest discover -s tests # Python (what CI runs)
node --test tests/js/*.test.js # frontend pure modules
python -m coverage run -m unittest discover -s tests && python -m coverage combine && python -m coverage report
```
CI runs the suite on Linux, Windows and macOS for Python 3.10 to 3.12 with an 85% coverage gate. Tests that delete a database file first check that nothing in the process still holds it open (`tests/open_handles.py`), so a connection leak fails on Linux the way it would on Windows.
### Code Style
The project follows PEP 8 guidelines with:
@@ -874,6 +895,10 @@ MIT License - see LICENSE file for details.
### Recently Completed
- **✅ Security & Reliability Hardening (3.3.0)**: validated gallery roots, contained file serving, safe backups and restores, background scans, 85%+ test coverage on all three operating systems
- **✅ Worker Threads**: scan, thumbnails and duplicate detection run on a configurable thread pool
- **✅ Tag Autocomplete & Infinite Scrolling**: faster browsing in the dashboard and gallery
- **✅ Queue-Time Run Counting (3.2.4)**: run counts and image linking follow the ComfyUI queue
- **✅ LoRA Manager Integration**: Import LoRA metadata, trigger words, and preview images
- **✅ Folder Filter**: Browse and filter prompts by output subdirectory
- **✅ Multi-Directory Gallery**: Scan multiple output directories simultaneously
@@ -893,6 +918,31 @@ MIT License - see LICENSE file for details.
## Changelog
### v3.3.1 (Proxy Routing Fix)
- **🌐 Proxy support**: the dashboard, gallery, tag page and metadata viewer now work when ComfyUI is reached through a proxy mount such as SwarmUI's `/ComfyBackendDirect/`. Scripts, images, API calls, progress streams and gallery folder filters stay under the mount, and the canvas node's buttons use ComfyUI's `fileURL()`. Direct ComfyUI URLs are unchanged. Thanks to @ghosttigerllc-bit ([#164](https://github.com/ComfyAssets/ComfyUI_PromptManager/pull/164))
- **📦 Dependencies**: the runtime minimums stay at `Pillow>=8.0.0` and `watchdog>=2.1.0`, so updating Prompt Manager never forces upgrades in the shared ComfyUI environment
### v3.3.0 (Hardening, Background Jobs & Worker Threads)
- **🔒 Security**: gallery roots restricted to ComfyUI directories and validated wherever they are read; media-only, fail-closed image serving; Windows path escape in static routes closed; quote-safe HTML escaping with inline handlers moved to data attributes; the CivitAI key is never returned and only sent to HTTPS CivitAI hosts; auto-tag paths contained; side-effect routes are POST-only; backups include the WAL and restores are size-capped, integrity-checked and applied through SQLite's online backup API
- **🧵 Worker Threads**: a Settings slider (default half the cores) used by the image rescan, thumbnail generation and duplicate detection. Measured: a 1 GB sample scanned in 2.5 s instead of 18 s at 8 threads
- **🔁 Background Jobs**: the image rescan and the duplicate scan run on the server, stream progress, survive a closed tab, minimize to a progress pill and re-attach after a reload
- **🎬 Video Prompts**: the scan reads the prompt ComfyUI's video savers embed in MP4/MOV/WebM containers (needs `ffprobe`)
- **🖼️ Thumbnails**: written atomically and served under versioned URLs so a regenerated thumbnail never shows a stale cached copy; the per-run cap applies only to thumbnails that still need generating and the gallery says how many remain
- **🏷️ Tag Autocomplete**: the dashboard tag box suggests tags, narrowed to those that co-occur with the tags already typed
- **♾️ Infinite Scrolling**: optional in the gallery and the dashboard list (saved default plus a session toggle)
- **⚙️ Gallery Settings** now apply at load; grid columns, media info and the Rescan button work as labelled
- **🐛 Correctness**: run counts increment when ComfyUI dequeues a prompt; negative prompts are never traced; per-thread SQLite connections and idempotent migrations (an old rebuild could delete every image row); `truncate_logs` no longer rolls the log over; the ComfyUI canvas loads only the canvas extension; Windows console banners are ASCII-safe; subfolders use `/` on every platform
- **🧪 Tests & CI**: 1,460 Python and 142 JavaScript tests (from 510 and 26), coverage 96% with an 85% gate, CI on Linux, Windows and macOS for Python 3.10–3.12, flake8 clean
- **⚠️ API changes**: image paths in responses are relative to the ComfyUI tree (use `url`); `autotag/download`, `autotag/start` and the thumbnail progress stream are POST-only; `lora/status` returns `has_civitai_api_key` instead of the key; `/search` is paginated
### v3.2.4 (Recently Used Sorting & Queue-Time Run Counts)
- **🕒 Recently Used Sort**: sort the dashboard by last use, most used or rating, server-side across pages
- **🔢 Run Counts**: `run_count` and `last_used_at` are recorded through the ComfyUI queue hook, so cached nodes still count
- **🔗 Positive-Prompt Image Linking**: generated images link to the positive prompt traced through the queued graph
### v3.2.1 (LoRA Manager Integration)
- **🔗 LoRA Manager Integration**: Import LoRA metadata, trigger words, and CivitAI example images from [ComfyUI-Lora-Manager](https://github.com/willmiao/ComfyUI-Lora-Manager) into your prompt database
+83
View File
@@ -0,0 +1,83 @@
# Testing and merge confidence
Routine changes should merge after the required checks pass and the diff is
reviewed. Add a regression test for each bug; prove it fails before the fix.
## What runs
| Layer | Real components | Replaced boundaries |
| --- | --- | --- |
| Python tests | Domain logic, SQLite, API handlers, filesystem checks | ComfyUI host and ML runtimes where needed |
| JavaScript tests | Frontend modules and controller methods | DOM/network fixtures for focused edge cases |
| Viewer checks | Production Viewer.js and frontend methods in Chromium desktop/touch | Controlled images and page fixtures |
| Page workflows | Production HTML, startup, click handlers, HTTP routes, SQLite, PNG files, thumbnails | ComfyUI host services and the tagging model |
Page workflows cover save/edit/reload, duplicate-edit rejection, gallery image
identity, Tags selection and viewer cleanup, delayed responses, auto-tag review
and persistence, and infinite scrolling without replacing existing cards.
They run in Chromium, Firefox, WebKit, and Chromium touch emulation. No GPU,
model downloads, running ComfyUI instance, or personal configuration is required.
The server binds to localhost and creates disposable directories and a database;
its reset endpoint exists only in the test harness.
## Run locally
Use a Python environment containing `requirements-dev.txt` and Node.js 24:
```sh
python -m unittest discover -s tests -q
npm ci --prefix tools/vendor --ignore-scripts
npx --prefix tools/vendor playwright install --with-deps chromium firefox webkit
npm run --prefix tools/vendor test
npm run --prefix tools/vendor test:coverage
npm run --prefix tools/vendor test:workflows -- --project firefox --project webkit
```
`PM_TEST_PYTHON` selects the Python executable; `PM_TEST_PORT` selects the local
port (default 8765). Set these environment variables when Python is outside PATH
or the port is occupied. Test servers are never reused. On unsupported Linux
distributions, browser system dependencies may need manual installation; CI uses
Ubuntu. For a quick browser run, use `test:workflows -- --project chromium`.
## Coverage and failures
Python CI enforces 85% combined statement/branch coverage. Frontend coverage merges
Node V8 unit coverage with Chromium page-workflow coverage through c8. Every
`web/js/*.js` file appears, including untouched files at zero coverage. The
controller tests provide V8 filenames so their executed code is counted too.
`tools/vendor/frontend-coverage-baseline.json` holds a floor for every file and
metric. CI rejects regressions, missing measurements, and new files without a
reviewed baseline. These floors reflect current coverage, with two percentage
points of allowance for runtime/timing differences; they do not claim 80% coverage
of the large controllers. New files need at least 80% line coverage when added via
the baseline updater. Raise floors with `npm run --prefix tools/vendor
coverage:update`; it reruns the suites and cannot lower an existing floor.
Review baseline changes like code changes.
The `viewerjs` check runs on every PR, including backend-only changes. GitHub
artifacts named `frontend-test-reports` contain HTML coverage, a JSON summary,
and failure screenshots/traces. Open traces with Playwright's `show-trace` tool.
Live checks remain useful for ComfyUI host compatibility, real inference/GPU or
model dependency changes, and visual design changes. The automated fixtures do
not certify those environments or every existing product feature.
## Required merge checks
`.github/branch-protection.json` records the required checks: all nine Python/OS
combinations, Python coverage, formatting, frontend modules, and `viewerjs`.
The branch must be current with main, and administrators follow the same checks.
It does not require an additional human approval.
A repository administrator applies the configuration after the new workflow is
passing, then verifies it:
```sh
gh api --method PUT repos/ComfyAssets/ComfyUI_PromptManager/branches/main/protection \
--input .github/branch-protection.json
gh api repos/ComfyAssets/ComfyUI_PromptManager/branches/main/protection/required_status_checks
```
Update the manifest whenever required job names change. Keep the frontend check
free of path filters so a required check is never silently skipped.
+43 -7
View File
@@ -5,7 +5,8 @@ gallery system.
This module provides two main node types:
- PromptManager: CLIP encoding node that outputs CONDITIONING
- PromptManagerText: Text-only node that outputs STRING with prepend/append functionality
- PromptManagerText: Text-only node that outputs STRING with prepend/append
functionality
Both nodes share the same database backend for persistent prompt storage and include
an automatic image gallery system that monitors ComfyUI output directories and links
@@ -37,9 +38,9 @@ def get_version():
return "unknown"
from .prompt_manager import PromptManager
from .prompt_manager_text import PromptManagerText
from .prompt_search_list import PromptSearchList
from .prompt_manager import PromptManager # noqa: E402
from .prompt_manager_text import PromptManagerText # noqa: E402
from .prompt_search_list import PromptSearchList # noqa: E402
NODE_CLASS_MAPPINGS = {
"PromptManager": PromptManager,
@@ -53,8 +54,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"PromptSearchList": "Prompt Search List",
}
# Define path to web directory for UI components
WEB_DIRECTORY = "web"
# Canvas extension only. ComfyUI loads every .js under this directory into the
# canvas page, so the admin/gallery bundles and vendored libraries under web/
# are served by our own routes instead (see py/api).
WEB_DIRECTORY = "web/comfy"
# Add API routes (same pattern as ComfyUI_Assets)
try:
@@ -82,6 +85,10 @@ except Exception as e:
except Exception:
pass
# Shared database, set up with the monitoring system below. Stays None when that
# fails so the queue hook can report the database as unavailable instead of crashing.
_global_db = None
# Start image monitoring globally at module import time
# This ensures images are linked regardless of which node is used in the workflow
try:
@@ -111,12 +118,41 @@ except Exception as e:
except Exception:
print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}")
def _db_factory():
"""The shared database for the queue hook; called lazily per request."""
if _global_db is None:
raise RuntimeError("PromptManager database unavailable")
return _global_db
# Count prompt runs when ComfyUI dequeues workflows (it skips unchanged nodes, so
# node execution alone misses re-runs). Separate from monitoring so either can fail.
try:
from .py.config import server_instance
from .utils.usage_tracking import register_queue_hook
register_queue_hook(server_instance, _db_factory)
except Exception as e:
try:
from .utils.logging_config import get_logger
get_logger("prompt_manager.init").error(f"Failed to register usage hook: {e}")
except Exception:
import logging
logging.getLogger("prompt_manager.init").error(
f"Failed to register usage hook: {e}"
)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
# Print startup message with loaded tools
print()
print(f"\033[94m[ComfyUI-PromptManager] Version:\033[0m {get_version()}")
for node_key, display_name in NODE_DISPLAY_NAME_MAPPINGS.items():
print(f"🫶 \033[94mLoaded:\033[0m {display_name}")
# ASCII only: a Windows console using a legacy code page cannot print emoji,
# and a UnicodeEncodeError here would abort loading the whole node.
print(f" \033[94mLoaded:\033[0m {display_name}")
print(f"\033[94mTotal: {len(NODE_CLASS_MAPPINGS)} tools loaded\033[0m")
print()
+523 -196
View File
@@ -2,10 +2,16 @@
Database schema and models for KikoTextEncode prompt storage.
"""
import contextlib
import datetime
import shutil
import sqlite3
import os
import threading
from typing import Optional
import weakref
from contextlib import closing
from pathlib import Path
from typing import Set, Tuple
# Import logging system
try:
@@ -17,10 +23,53 @@ except ImportError:
sys.path.insert(0, current_dir)
from utils.logging_config import get_logger
# last_used_at is UTC ISO text so ORDER BY can compare it as text. The backfill
# writes milliseconds (SQLite's %f); live writes use microseconds so runs in the
# same millisecond still order correctly. Both forms sort correctly together.
ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
def utc_now_iso() -> str:
"""Current UTC time as ISO text with microseconds.
Example: 2026-10-07T18:39:06.758123+00:00
"""
return datetime.datetime.now(datetime.timezone.utc).isoformat(
timespec="microseconds"
)
def normalize_image_path(path: str) -> str:
"""Canonical form of an image path for generated_images.file_path.
Absolute, with '..' and '.' collapsed and the platform's case folding
applied, so two spellings of one file (or the same basename in two date
folders) compare the way the filesystem would.
"""
return os.path.normcase(os.path.normpath(os.path.abspath(path)))
# How long a connection waits for a locked database before giving up. Thirty
# seconds covers a scan or a thumbnail run writing while ComfyUI saves images,
# and the slow disks of CI runners.
BUSY_TIMEOUT_SECONDS = 30
SQLITE_HEADER = b"SQLite format 3\x00"
REQUIRED_PROMPT_COLUMNS = ("id", "text", "created_at")
class PromptModel:
"""Database model for prompt storage and schema management."""
# Paths (normalised, see _schema_key) whose schema has been created and
# migrated in this process. Many PromptDatabase instances point at the
# same file; only the first one pays for the schema check.
_initialized_paths: Set[str] = set()
_init_lock = threading.Lock()
# Every live model, so a shutdown or a test teardown can close them all
# before the database file is deleted (Windows refuses otherwise).
_instances: "weakref.WeakSet[PromptModel]" = weakref.WeakSet()
_instances_lock = threading.Lock()
def __init__(self, db_path: str):
"""
Initialize the database model.
@@ -31,8 +80,16 @@ class PromptModel:
self.logger = get_logger("prompt_manager.database.models")
self.logger.debug(f"Initializing database model with path: {db_path}")
self.db_path = db_path
self._conn: Optional[sqlite3.Connection] = None
# One connection per thread (see get_connection); every open connection
# is also tracked here so restore_from_file can close them all.
self._local = threading.local()
self._connections: Set[sqlite3.Connection] = set()
self._conn_lock = threading.Lock()
# Bumped by _close_all_connections so threads holding a closed
# connection reopen on their next get_connection().
self._generation = 0
with PromptModel._instances_lock:
PromptModel._instances.add(self)
self._ensure_database_exists()
def _ensure_database_exists(self) -> None:
@@ -45,17 +102,59 @@ class PromptModel:
Raises:
Exception: If database creation fails
"""
key = self._schema_key()
with PromptModel._init_lock:
if key in PromptModel._initialized_paths and os.path.exists(self.db_path):
return
try:
conn = sqlite3.connect(self.db_path, timeout=BUSY_TIMEOUT_SECONDS)
try:
conn.execute("PRAGMA journal_mode = WAL")
conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
conn.execute("PRAGMA foreign_keys = ON")
self._create_tables(conn)
self._create_indexes(conn)
conn.commit()
finally:
conn.close()
except Exception as e:
self.logger.error(f"Error creating database: {e}")
raise
PromptModel._initialized_paths.add(key)
def _schema_key(self) -> str:
"""Normalised identity of db_path for the once-per-process schema cache."""
return os.path.normcase(os.path.realpath(self.db_path))
@classmethod
def reset_schema_cache(cls) -> None:
"""Forget which databases were initialised (tests and restores)."""
with cls._init_lock:
cls._initialized_paths.clear()
@contextlib.contextmanager
def _table_rebuild(self, conn: sqlite3.Connection):
"""Run a copy-and-swap table rebuild with foreign keys disabled.
DROP TABLE on a parent table with foreign_keys ON performs an implicit
DELETE that cascades into child rows (generated_images, prompt_tags),
so the rebuild runs with the pragma off. The pragma is a no-op inside
a transaction, hence the commits around it. Errors roll back and
propagate after the pragma is restored.
"""
conn.commit()
conn.execute("PRAGMA foreign_keys = OFF")
# Explicit BEGIN: in sqlite3's legacy transaction mode DDL autocommits,
# so without it a failed copy would leave the *_new table behind.
conn.execute("BEGIN")
try:
with sqlite3.connect(self.db_path) as conn:
conn.execute("PRAGMA journal_mode = WAL")
conn.execute("PRAGMA busy_timeout = 5000")
conn.execute("PRAGMA foreign_keys = ON")
self._create_tables(conn)
self._create_indexes(conn)
conn.commit()
except Exception as e:
self.logger.error(f"Error creating database: {e}")
yield
conn.commit()
except BaseException:
conn.rollback()
raise
finally:
conn.execute("PRAGMA foreign_keys = ON")
def _create_tables(self, conn: sqlite3.Connection) -> None:
"""
@@ -78,7 +177,9 @@ class PromptModel:
tags TEXT,
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
notes TEXT,
hash TEXT UNIQUE
hash TEXT UNIQUE,
last_used_at TIMESTAMP,
run_count INTEGER NOT NULL DEFAULT 0
)
""")
@@ -97,8 +198,9 @@ class PromptModel:
workflow_data TEXT,
prompt_metadata TEXT,
parameters TEXT,
file_path TEXT,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
UNIQUE(prompt_id, filename)
UNIQUE(prompt_id, file_path)
)
""")
@@ -119,9 +221,6 @@ class PromptModel:
)
""")
# Add unique constraint to existing databases (migration)
self._migrate_add_unique_constraint(conn)
# Check if we need to migrate from old schema with workflow_name
self._migrate_workflow_name_removal(conn)
@@ -131,6 +230,12 @@ class PromptModel:
# Migrate JSON tags to normalized junction tables
self._migrate_json_tags_to_junction(conn)
# Usage tracking columns (after any migration that rebuilds prompts)
self._migrate_add_usage_columns(conn)
# Image uniqueness by full path instead of basename (3.2.4)
self._migrate_image_uniqueness_by_path(conn)
def _create_indexes(self, conn: sqlite3.Connection) -> None:
"""
Create indexes for better query performance.
@@ -145,41 +250,100 @@ class PromptModel:
- Hash values for duplicate detection
"""
indexes = [
"CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)",
"CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)",
"CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)",
"CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)",
"CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)",
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)",
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)",
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)",
"CREATE INDEX IF NOT EXISTS idx_prompt_tags_tag ON prompt_tags(tag_id)",
"CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name)",
("idx_prompts_text", "prompts(text)"),
("idx_prompts_category", "prompts(category)"),
("idx_prompts_created_at", "prompts(created_at)"),
("idx_prompts_hash", "prompts(hash)"),
("idx_prompts_rating", "prompts(rating)"),
("idx_prompts_last_used", "prompts(last_used_at)"),
("idx_prompts_run_count", "prompts(run_count)"),
("idx_prompt_images", "generated_images(prompt_id)"),
("idx_image_path", "generated_images(image_path)"),
("idx_image_file_path", "generated_images(file_path)"),
("idx_generation_time", "generated_images(generation_time)"),
("idx_prompt_tags_tag", "prompt_tags(tag_id)"),
("idx_tags_name", "tags(name)"),
]
for index_sql in indexes:
conn.execute(index_sql)
for name, target in indexes:
conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}")
def get_connection(self) -> sqlite3.Connection:
"""
Get the persistent database connection.
Get the calling thread's database connection.
Returns a thread-safe, reusable connection protected by a lock.
The connection is created on first call and reused thereafter.
Callers use it as a context manager; the lock is held for the
duration of the ``with`` block.
Each thread gets its own ``sqlite3.Connection`` (kept in
``threading.local``), created on first use and reused by that thread
thereafter. Every connection is configured with WAL journaling,
``row_factory = sqlite3.Row``, ``PRAGMA foreign_keys = ON`` and
``PRAGMA busy_timeout`` of BUSY_TIMEOUT_SECONDS so concurrent writers
wait instead of failing. Callers use it as a context manager
(``with conn:``), which commits or rolls back that thread's transaction
and leaves the connection open. No lock is held while a connection is in use.
Returns:
sqlite3.Connection: Configured database connection
sqlite3.Connection: The current thread's configured connection
"""
conn = getattr(self._local, "conn", None)
if conn is not None and self._local.generation == self._generation:
return conn
# check_same_thread=False only so restore_from_file can close other
# threads' connections; each connection is still used by one thread.
conn = sqlite3.connect(
self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS
)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode = WAL")
conn.execute("PRAGMA foreign_keys = ON")
conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}")
self._local.conn = conn
self._local.generation = self._generation
with self._conn_lock:
self._connections.add(conn)
return conn
@classmethod
def close_all_instances(cls) -> None:
"""Close every connection of every live model in this process."""
with cls._instances_lock:
models = list(cls._instances)
for model in models:
model.close_all()
def close_all(self) -> None:
"""Close every thread's connection; they reopen lazily on next use.
Needed before deleting the database file (Windows keeps open files
locked) and by tests that ran queries through executor threads.
"""
with self._conn_lock:
if self._conn is None:
self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
self._conn.row_factory = sqlite3.Row
self._conn.execute("PRAGMA journal_mode = WAL")
self._conn.execute("PRAGMA foreign_keys = ON")
self._conn.execute("PRAGMA busy_timeout = 5000")
return self._conn
connections = list(self._connections)
self._connections.clear()
self._generation += 1
self._local.conn = None
for conn in connections:
try:
conn.close()
except sqlite3.Error as e:
self.logger.warning(f"Error closing database connection: {e}")
def close(self) -> None:
"""Close the calling thread's connection, if it has one.
Other threads keep their connections. Call this before deleting or
replacing the database file: Windows refuses to delete an open file.
"""
conn = getattr(self._local, "conn", None)
if conn is None:
return
self._local.conn = None
with self._conn_lock:
self._connections.discard(conn)
try:
conn.close()
except sqlite3.Error as e:
self.logger.warning(f"Error closing database connection: {e}")
def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
"""
@@ -190,16 +354,20 @@ class PromptModel:
This migration handles legacy schema updates by removing the deprecated
workflow_name column while preserving all other data.
Raises:
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
"""
try:
# Check if workflow_name column exists
cursor = conn.execute("PRAGMA table_info(prompts)")
columns = [column[1] for column in cursor.fetchall()]
if "workflow_name" not in columns:
return
if "workflow_name" in columns:
self.logger.info("Migrating database: removing workflow_name column")
# Create new table without workflow_name
self.logger.info("Migrating database: removing workflow_name column")
with self._table_rebuild(conn):
# A previous attempt may have died between CREATE and RENAME
conn.execute("DROP TABLE IF EXISTS prompts_new")
conn.execute("""
CREATE TABLE prompts_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -213,23 +381,20 @@ class PromptModel:
hash TEXT UNIQUE
)
""")
# Copy data from old table to new table
conn.execute("""
INSERT INTO prompts_new (id, text, created_at, updated_at, category, tags, rating, notes, hash)
SELECT id, text, created_at, updated_at, category, tags, rating, notes, hash
INSERT INTO prompts_new (id, text, created_at, updated_at,
category, tags, rating, notes, hash)
SELECT id, text, created_at, updated_at, category, tags, rating,
notes, hash
FROM prompts
""")
# Drop old table and rename new one
conn.execute("DROP TABLE prompts")
conn.execute("ALTER TABLE prompts_new RENAME TO prompts")
self.logger.info("Database migration completed")
self.logger.info("Database migration completed")
except Exception as e:
self.logger.error(f"Migration error: {e}")
# If migration fails, the table creation will handle it
except sqlite3.Error as e:
self.logger.error(f"Migration error (workflow_name removal): {e}")
raise
def _migrate_foreign_key_types(self, conn: sqlite3.Connection) -> None:
"""
@@ -240,18 +405,19 @@ class PromptModel:
Converts prompt_id from TEXT to INTEGER type to match the prompts table's
primary key type, ensuring referential integrity.
Raises:
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
"""
try:
# Check if generated_images table exists and has TEXT prompt_id
cursor = conn.execute("PRAGMA table_info(generated_images)")
columns = {column[1]: column[2] for column in cursor.fetchall()}
if columns.get("prompt_id") != "TEXT":
return
if "prompt_id" in columns and columns["prompt_id"] == "TEXT":
self.logger.info(
"Migrating foreign key types: prompt_id TEXT -> INTEGER"
)
# Create new table with correct types
self.logger.info("Migrating foreign key types: prompt_id TEXT -> INTEGER")
with self._table_rebuild(conn):
conn.execute("DROP TABLE IF EXISTS generated_images_new")
conn.execute("""
CREATE TABLE generated_images_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -269,130 +435,27 @@ class PromptModel:
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE
)
""")
# Copy data, converting prompt_id from TEXT to INTEGER
conn.execute("""
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters)
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename, generation_time,
file_size, width, height, format, workflow_data, prompt_metadata, parameters
SELECT id, CAST(prompt_id AS INTEGER), image_path, filename,
generation_time,
file_size, width, height, format, workflow_data,
prompt_metadata, parameters
FROM generated_images
WHERE prompt_id != '' AND prompt_id IS NOT NULL
AND CAST(prompt_id AS INTEGER) IN (SELECT id FROM prompts)
""")
# Drop old table and rename new one
conn.execute("DROP TABLE generated_images")
conn.execute(
"ALTER TABLE generated_images_new RENAME TO generated_images"
)
self.logger.info("Foreign key migration completed")
self.logger.info("Foreign key migration completed")
except Exception as e:
self.logger.error(f"Foreign key migration error: {e}")
# If migration fails, continue with existing schema
def _migrate_add_unique_constraint(self, conn: sqlite3.Connection) -> None:
"""
Add UNIQUE constraint on (prompt_id, filename) to prevent duplicate image entries.
This migration:
1. Checks if the constraint already exists
2. Removes duplicate entries (keeping the most recent)
3. Recreates the table with the UNIQUE constraint
Args:
conn: Active database connection
"""
try:
# Check if the unique constraint already exists by looking at table info
cursor = conn.execute("PRAGMA index_list(generated_images)")
indexes = cursor.fetchall()
# Check if we have a unique index on prompt_id, filename
has_unique_constraint = False
for idx in indexes:
if idx[2] == 1: # unique flag
cursor = conn.execute(f"PRAGMA index_info({idx[1]})")
columns = [col[2] for col in cursor.fetchall()]
if "prompt_id" in columns and "filename" in columns:
has_unique_constraint = True
break
if has_unique_constraint:
return # Already migrated
self.logger.info(
"Migrating database: adding UNIQUE constraint on (prompt_id, filename)"
)
# First, remove duplicates keeping only the most recent (highest id)
conn.execute("""
DELETE FROM generated_images
WHERE id NOT IN (
SELECT MAX(id) FROM generated_images
GROUP BY prompt_id, filename
)
""")
duplicates_removed = conn.total_changes
if duplicates_removed > 0:
self.logger.info(
f"Removed {duplicates_removed} duplicate image entries"
)
# Create new table with UNIQUE constraint
conn.execute("""
CREATE TABLE IF NOT EXISTS generated_images_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
prompt_id INTEGER NOT NULL,
image_path TEXT NOT NULL,
filename TEXT NOT NULL,
generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
file_size INTEGER,
width INTEGER,
height INTEGER,
format TEXT,
workflow_data TEXT,
prompt_metadata TEXT,
parameters TEXT,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE,
UNIQUE(prompt_id, filename)
)
""")
# Copy data
conn.execute("""
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters)
SELECT id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters
FROM generated_images
""")
# Drop old table and rename new one
conn.execute("DROP TABLE generated_images")
conn.execute("ALTER TABLE generated_images_new RENAME TO generated_images")
# Recreate indexes
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)"
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)"
)
self.logger.info("UNIQUE constraint migration completed successfully")
except Exception as e:
self.logger.error(f"UNIQUE constraint migration error: {e}")
# Continue with existing schema if migration fails
except sqlite3.Error as e:
self.logger.error(f"Migration error (foreign key types): {e}")
raise
def _migrate_json_tags_to_junction(self, conn: sqlite3.Connection) -> None:
"""
@@ -400,6 +463,9 @@ class PromptModel:
Runs once: skips if the tags table already has data. Uses json_each()
to extract tag names from the JSON arrays stored in prompts.tags.
Raises:
sqlite3.Error: If the migration fails (logged)
"""
try:
cursor = conn.execute("SELECT COUNT(*) FROM tags")
@@ -419,7 +485,8 @@ class PromptModel:
conn.execute(
"INSERT OR IGNORE INTO tags (name) "
"SELECT DISTINCT je.value FROM prompts, json_each(prompts.tags) AS je "
"WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND prompts.tags != '[]'"
"WHERE prompts.tags IS NOT NULL AND prompts.tags != '' AND "
"prompts.tags != '[]'"
)
# Populate junction table
@@ -434,21 +501,142 @@ class PromptModel:
tag_count = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0]
link_count = conn.execute("SELECT COUNT(*) FROM prompt_tags").fetchone()[0]
self.logger.info(
f"Tag migration complete: {tag_count} unique tags, {link_count} prompt-tag links"
f"Tag migration complete: {tag_count} unique tags, {link_count} "
"prompt-tag links"
)
except Exception as e:
self.logger.error(f"Tag junction migration error: {e}")
except sqlite3.Error as e:
self.logger.error(f"Migration error (tag junction): {e}")
raise
def migrate_database(self) -> None:
def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None:
"""
Apply any pending database migrations.
Add last_used_at and run_count (3.2.4) and estimate usage for existing rows.
This method serves as an entry point for future schema migrations.
Add new migration logic here as the database evolves.
The backfill runs only when the columns are added, i.e. on the upgrade
from a version without usage tracking. Afterwards a NULL last_used_at
means "never run": saving a prompt does not stamp it, only
record_prompt_use does, and "Recently Used" relies on that.
The estimate comes from linked images: run_count is the image count (at
least 1) and last_used_at the newest image time, falling back to created_at.
Timestamps are compared with julianday() because prompts use ISO 'T'
timestamps while generated_images uses SQLite's CURRENT_TIMESTAMP format.
"""
# Future migrations can be added here
pass
try:
columns = {row[1] for row in conn.execute("PRAGMA table_info(prompts)")}
added = False
if "last_used_at" not in columns:
self.logger.info("Adding prompt usage column last_used_at")
conn.execute("ALTER TABLE prompts ADD COLUMN last_used_at TIMESTAMP")
added = True
if "run_count" not in columns:
self.logger.info("Adding prompt usage column run_count")
conn.execute(
"ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL "
"DEFAULT 0"
)
added = True
if not added:
return
cursor = conn.execute(f"""
UPDATE prompts SET
run_count = MAX(run_count, 1, (
SELECT COUNT(*) FROM generated_images gi
WHERE gi.prompt_id = prompts.id
)),
last_used_at = strftime('{ISO_TIMESTAMP_FORMAT}', MAX(
COALESCE(julianday(created_at), julianday('now')),
COALESCE((
SELECT MAX(julianday(gi.generation_time))
FROM generated_images gi WHERE gi.prompt_id = prompts.id
), 0)
))
WHERE last_used_at IS NULL
""")
if cursor.rowcount:
self.logger.info(f"Estimated usage for {cursor.rowcount} prompts")
except sqlite3.Error as e:
self.logger.error(f"Migration error (usage columns): {e}")
raise
def _migrate_image_uniqueness_by_path(self, conn: sqlite3.Connection) -> None:
"""
Rebuild generated_images so an image is unique per (prompt_id, file_path).
Older schemas were unique per (prompt_id, filename), i.e. the basename,
so ComfyUI's same-named files in different date folders were silently
dropped. file_path holds normalize_image_path(image_path); rows whose
paths normalise to the same file collapse to the newest one.
Raises:
sqlite3.Error: If the rebuild fails (logged, nothing half-applied)
"""
try:
columns = {
row[1] for row in conn.execute("PRAGMA table_info(generated_images)")
}
if "file_path" in columns:
return
self.logger.info("Migrating database: image uniqueness by full path")
with self._table_rebuild(conn):
rows = conn.execute(
"SELECT id, prompt_id, image_path FROM generated_images ORDER BY id"
).fetchall()
keep = {}
for row_id, prompt_id, image_path in rows:
keep[(prompt_id, normalize_image_path(image_path))] = row_id
conn.execute("DROP TABLE IF EXISTS generated_images_new")
conn.execute("""
CREATE TABLE generated_images_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
prompt_id INTEGER NOT NULL,
image_path TEXT NOT NULL,
filename TEXT NOT NULL,
generation_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
file_size INTEGER,
width INTEGER,
height INTEGER,
format TEXT,
workflow_data TEXT,
prompt_metadata TEXT,
parameters TEXT,
file_path TEXT,
FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE
CASCADE,
UNIQUE(prompt_id, file_path)
)
""")
conn.executemany(
"""
INSERT INTO generated_images_new
(id, prompt_id, image_path, filename, generation_time, file_size,
width, height, format, workflow_data, prompt_metadata, parameters,
file_path)
SELECT id, prompt_id, image_path, filename, generation_time,
file_size,
width, height, format, workflow_data, prompt_metadata,
parameters,
?
FROM generated_images WHERE id = ?
""",
[(file_path, row_id) for (_, file_path), row_id in keep.items()],
)
conn.execute("DROP TABLE generated_images")
conn.execute(
"ALTER TABLE generated_images_new RENAME TO generated_images"
)
dropped = len(rows) - len(keep)
if dropped:
self.logger.info(f"Collapsed {dropped} duplicate image rows")
self.logger.info("Image path uniqueness migration completed")
except sqlite3.Error as e:
self.logger.error(f"Migration error (image path uniqueness): {e}")
raise
def vacuum_database(self) -> None:
"""
@@ -458,7 +646,7 @@ class PromptModel:
improving query performance and reducing file size.
"""
try:
with sqlite3.connect(self.db_path) as conn:
with closing(sqlite3.connect(self.db_path)) as conn:
conn.execute("VACUUM")
conn.commit()
except Exception as e:
@@ -477,12 +665,14 @@ class PromptModel:
total_prompts = cursor.fetchone()["total_prompts"]
cursor = conn.execute(
"SELECT COUNT(DISTINCT category) as unique_categories FROM prompts WHERE category IS NOT NULL"
"SELECT COUNT(DISTINCT category) as unique_categories FROM "
"prompts WHERE category IS NOT NULL"
)
unique_categories = cursor.fetchone()["unique_categories"]
cursor = conn.execute(
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS NOT NULL"
"SELECT AVG(rating) as avg_rating FROM prompts WHERE rating IS "
"NOT NULL"
)
avg_rating = cursor.fetchone()["avg_rating"]
@@ -504,19 +694,156 @@ class PromptModel:
def backup_database(self, backup_path: str) -> bool:
"""
Create a backup of the database.
Write a consistent copy of the database to backup_path.
Uses the SQLite online backup API through this thread's connection,
so rows that are committed but still sit in the -wal file are
included and other threads may keep reading and writing meanwhile.
Args:
backup_path: Path where the backup should be saved
backup_path: Path where the backup should be saved (overwritten)
Returns:
bool: True if backup was successful, False otherwise
"""
try:
import shutil
shutil.copy2(self.db_path, backup_path)
if os.path.exists(backup_path) and os.path.samefile(
backup_path, self.db_path
):
self.logger.error("Refusing to back up the database onto itself")
return False
if os.path.exists(backup_path):
os.remove(backup_path)
source = self.get_connection()
destination = sqlite3.connect(backup_path)
try:
source.backup(destination)
finally:
destination.close()
return True
except Exception as e:
except (OSError, sqlite3.Error) as e:
self.logger.error(f"Error creating database backup: {e}")
return False
@staticmethod
def verify_database_file(path: str) -> Tuple[bool, str]:
"""
Check that a file is a healthy prompts database before restoring it.
Opens the file read-only, runs PRAGMA integrity_check and confirms the
prompts table has the columns every version of this extension needs.
Args:
path: File to inspect
Returns:
(ok, reason): reason is "ok" when the file is acceptable, otherwise
a human-readable explanation. Never raises for bad input.
"""
if not os.path.isfile(path):
return False, "File not found"
try:
with open(path, "rb") as fh:
header = fh.read(len(SQLITE_HEADER))
except OSError as e:
return False, f"Cannot read file: {e}"
if header != SQLITE_HEADER:
return False, "Not a SQLite database file"
uri = Path(os.path.abspath(path)).as_uri() + "?mode=ro"
try:
conn = sqlite3.connect(uri, uri=True)
except sqlite3.Error as e:
return False, f"Invalid SQLite database: {e}"
try:
row = conn.execute("PRAGMA integrity_check").fetchone()
verdict = row[0] if row else "no result"
if verdict != "ok":
return False, f"Integrity check failed: {verdict}"
table = conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' AND name='prompts'"
).fetchone()
if not table:
return False, "Database does not contain a 'prompts' table"
columns = {r[1] for r in conn.execute("PRAGMA table_info(prompts)")}
missing = [c for c in REQUIRED_PROMPT_COLUMNS if c not in columns]
if missing:
return (
False,
"prompts table is missing required columns: " + ", ".join(missing),
)
return True, "ok"
except sqlite3.Error as e:
return False, f"Invalid SQLite database: {e}"
finally:
conn.close()
def _replace_contents_from(self, src_path: str) -> None:
"""Copy every page of *src_path* into the live database.
Goes through SQLite's online backup API into this thread's normal
connection, so SQLite takes the locks itself and the other
PromptModel instances that share the file (nodes, image monitor,
queue hook, API) simply see the new pages on their next statement.
Nothing is swapped underneath open handles and no -wal/-shm file has
to be deleted, which also keeps Windows happy. When SQLite refuses
(page sizes differ while the target is in WAL mode) fall back to a
file copy with every connection of this instance closed first.
"""
source = sqlite3.connect(
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
)
try:
try:
source.backup(self.get_connection())
return
except sqlite3.OperationalError as e:
self.logger.warning(
f"Online restore failed ({e}); falling back to a file copy"
)
finally:
source.close()
self.close_all()
shutil.copy2(src_path, self.db_path)
for suffix in ("-wal", "-shm"):
stale = self.db_path + suffix
if os.path.exists(stale):
os.remove(stale)
def restore_from_file(self, src_path: str) -> str:
"""
Replace the live database with the file at src_path.
The current database is first backed up next to itself as
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), the
pages of ``src_path`` are copied into the live database through
SQLite's backup API (visible to every open connection) and the schema
migrations run on the restored content.
Callers should run verify_database_file(src_path) first.
Args:
src_path: SQLite file to restore
Returns:
Path of the backup taken of the previous database, or "" when
there was no database to back up.
Raises:
RuntimeError: If the previous database could not be backed up
OSError: If the restored file could not be copied into place
"""
backup_path = ""
if os.path.exists(self.db_path):
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
backup_path = f"{self.db_path}.backup_{timestamp}"
if not self.backup_database(backup_path):
raise RuntimeError("Could not back up the current database")
self.logger.info(f"Current database backed up to: {backup_path}")
self._replace_contents_from(src_path)
with PromptModel._init_lock:
PromptModel._initialized_paths.discard(self._schema_key())
self._ensure_database_exists()
self.logger.info(f"Database restored from {src_path}")
return backup_path
+308 -160
View File
@@ -6,19 +6,21 @@ import sqlite3
import json
import datetime
import os
from typing import Optional, List, Dict, Any, Union
from typing import Optional, List, Dict, Any
from .models import PromptModel
from .models import PromptModel, normalize_image_path, utc_now_iso
# Import logging system
try:
from ..utils.logging_config import get_logger
from ..utils.hashing import generate_prompt_hash
except ImportError:
import sys
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, current_dir)
from utils.logging_config import get_logger
from utils.hashing import generate_prompt_hash
# Subquery to fetch tags from junction table, embedded in SELECT statements
@@ -28,6 +30,41 @@ TAG_SUBQUERY = (
"WHERE pt.prompt_id = prompts.id) AS _tag_list"
)
# Server-side sort orders, keyed by the value the UI sends. Only these strings ever
# reach ORDER BY; unknown keys fall back to DEFAULT_SORT.
# last_used_at is NULL until the prompt has run, so "IS NULL" first keeps
# never-run prompts after every run one (SQLite sorts NULL first in DESC).
SORT_ORDERS = {
"last_used_desc": "last_used_at IS NULL, last_used_at DESC, id DESC",
"created_desc": "created_at DESC, id DESC",
"created_asc": "created_at ASC, id ASC",
"run_count_desc": (
"run_count DESC, last_used_at IS NULL, last_used_at DESC, id DESC"
),
"rating_desc": "rating IS NULL, rating DESC, id DESC",
"rating_asc": "rating IS NULL, rating ASC, id DESC",
"text_asc": "text COLLATE NOCASE ASC, id ASC",
"text_desc": "text COLLATE NOCASE DESC, id DESC",
}
DEFAULT_SORT = "created_desc"
def order_by_clause(sort: Optional[str]) -> str:
"""Return a whitelisted ORDER BY expression for a UI sort key."""
return SORT_ORDERS.get(sort or DEFAULT_SORT, SORT_ORDERS[DEFAULT_SORT])
LIKE_ESCAPE = "\\"
def escape_like(term: str) -> str:
"""Escape LIKE wildcards in user text; use with ``LIKE ? ESCAPE '\\'``."""
return (
term.replace(LIKE_ESCAPE, LIKE_ESCAPE + LIKE_ESCAPE)
.replace("%", LIKE_ESCAPE + "%")
.replace("_", LIKE_ESCAPE + "_")
)
def _resolve_db_path(db_path: Optional[str] = None) -> str:
"""Resolve the database path from config, falling back to defaults.
@@ -54,7 +91,7 @@ def _resolve_db_path(db_path: Optional[str] = None) -> str:
db_path = "prompts.db"
if not os.path.isabs(db_path):
db_path = os.path.join(extension_root, db_path)
db_path = os.path.normpath(os.path.join(extension_root, db_path))
# Ensure parent directory exists for custom paths
parent = os.path.dirname(db_path)
@@ -64,6 +101,51 @@ def _resolve_db_path(db_path: Optional[str] = None) -> str:
return db_path
# Bounds for the file_info a link request may attach to an image. Anything
# outside them (or of the wrong type) is stored as NULL instead of being
# trusted into an INTEGER or TEXT column and echoed back to the gallery.
MAX_IMAGE_FILE_SIZE = 10**12
MAX_IMAGE_DIMENSION = 65535
IMAGE_FORMATS = frozenset(
{"PNG", "JPEG", "JPG", "WEBP", "GIF", "BMP", "TIFF", "MP4", "WEBM", "MOV"}
)
def _bounded_int(value: Any, upper: int) -> Optional[int]:
"""``value`` as an int within ``[0, upper]``; None for anything else."""
if isinstance(value, bool):
return None
try:
number = int(value)
except (TypeError, ValueError, OverflowError):
return None
return number if 0 <= number <= upper else None
def _known_image_format(value: Any) -> Optional[str]:
"""Upper-cased *value* when it names a supported media format, else None."""
if not isinstance(value, str):
return None
name = value.strip().upper()
return name if name in IMAGE_FORMATS else None
def _coerce_file_info(metadata: Any) -> tuple:
"""``(size, width, height, format)`` from request metadata, each validated."""
file_info = metadata.get("file_info") if isinstance(metadata, dict) else None
if not isinstance(file_info, dict):
return None, None, None, None
dimensions = file_info.get("dimensions")
if not isinstance(dimensions, (list, tuple)) or len(dimensions) < 2:
dimensions = (None, None)
return (
_bounded_int(file_info.get("size"), MAX_IMAGE_FILE_SIZE),
_bounded_int(dimensions[0], MAX_IMAGE_DIMENSION),
_bounded_int(dimensions[1], MAX_IMAGE_DIMENSION),
_known_image_format(file_info.get("format")),
)
class PromptDatabase:
"""Database operations class for managing prompts."""
@@ -81,6 +163,19 @@ class PromptDatabase:
self.model = PromptModel(db_path)
self.logger.debug("Database operations initialized successfully")
def close(self) -> None:
"""Close the calling thread's database connection (see PromptModel.close)."""
self.model.close()
@staticmethod
def close_all_instances() -> None:
"""Close the connections of every PromptDatabase/PromptModel in the process."""
PromptModel.close_all_instances()
def close_all(self) -> None:
"""Close every thread's connection (see PromptModel.close_all)."""
self.model.close_all()
def save_prompt(
self,
text: str,
@@ -115,32 +210,66 @@ class PromptDatabase:
raise ValueError("Rating must be between 1 and 5")
self.logger.debug(
f"Saving prompt: text_length={len(text)}, category={category}, tags={tags}, rating={rating}"
f"Saving prompt: text_length={len(text)}, category={category}, "
f"tags={tags}, rating={rating}"
)
try:
with self.model.get_connection() as conn:
cursor = conn.execute(
"""
INSERT INTO prompts (
text, category, tags, rating, notes, hash, created_at,
updated_at
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
""",
(
text.strip(),
category,
rating,
notes,
prompt_hash,
datetime.datetime.now(datetime.timezone.utc).isoformat(),
datetime.datetime.now(datetime.timezone.utc).isoformat(),
),
)
prompt_id = cursor.lastrowid
if tags:
self._sync_prompt_tags(conn, prompt_id, tags)
conn.commit()
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
return prompt_id
except sqlite3.IntegrityError:
# Another caller inserted the same hash between the caller's
# duplicate check and this INSERT (the with-block rolled back).
existing = self.get_prompt_by_hash(prompt_hash) if prompt_hash else None
if existing is None:
raise
self.logger.debug(
f"Prompt with hash {prompt_hash} was saved concurrently; "
f"returning existing ID {existing['id']}"
)
return existing["id"]
def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
"""
Count runs of a prompt and mark it as most recently used.
Args:
prompt_id: The prompt ID
times: Number of runs to add (queued runs counted on first save)
Returns:
True if the prompt exists and was updated
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
"""
INSERT INTO prompts (
text, category, tags, rating, notes, hash, created_at, updated_at
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
""",
(
text.strip(),
category,
rating,
notes,
prompt_hash,
datetime.datetime.now(datetime.timezone.utc).isoformat(),
datetime.datetime.now(datetime.timezone.utc).isoformat(),
),
"UPDATE prompts SET run_count = run_count + ?, last_used_at = ?"
" WHERE id = ?",
(max(1, int(times)), utc_now_iso(), prompt_id),
)
prompt_id = cursor.lastrowid
if tags:
self._sync_prompt_tags(conn, prompt_id, tags)
conn.commit()
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
return prompt_id
return cursor.rowcount > 0
def get_prompt_by_id(self, prompt_id: int) -> Optional[Dict[str, Any]]:
"""
@@ -191,6 +320,7 @@ class PromptDatabase:
limit: int = 100,
offset: int = 0,
tag_partial: bool = False,
sort: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""
Search prompts with various filters.
@@ -207,6 +337,7 @@ class PromptDatabase:
limit: Maximum number of results
offset: Number of results to skip
tag_partial: Use LIKE matching for tags instead of exact match
sort: Sort key from SORT_ORDERS (default: newest first)
Returns:
List of dictionaries containing prompt data
@@ -215,8 +346,8 @@ class PromptDatabase:
params = []
if text:
query_parts.append("AND text LIKE ?")
params.append(f"%{text}%")
query_parts.append("AND text LIKE ? ESCAPE '\\'")
params.append(f"%{escape_like(text)}%")
if category:
query_parts.append("AND category = ?")
@@ -228,9 +359,10 @@ class PromptDatabase:
query_parts.append(
"AND prompts.id IN ("
" SELECT pt.prompt_id FROM prompt_tags pt"
" JOIN tags t ON pt.tag_id = t.id WHERE t.name LIKE ?)"
" JOIN tags t ON pt.tag_id = t.id"
" WHERE t.name LIKE ? ESCAPE '\\')"
)
params.append(f"%{tag}%")
params.append(f"%{escape_like(tag)}%")
else:
query_parts.append(
"AND prompts.id IN ("
@@ -275,7 +407,7 @@ class PromptDatabase:
query_parts.append("AND created_at <= ?")
params.append(date_to)
query_parts.append("ORDER BY created_at DESC LIMIT ? OFFSET ?")
query_parts.append(f"ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?")
params.extend([limit, offset])
query = " ".join(query_parts)
@@ -291,13 +423,16 @@ class PromptDatabase:
return prompts
def get_recent_prompts(self, limit: int = 10, offset: int = 0) -> Dict[str, Any]:
def get_recent_prompts(
self, limit: int = 10, offset: int = 0, sort: Optional[str] = None
) -> Dict[str, Any]:
"""
Get the most recent prompts with pagination support.
Args:
limit: Maximum number of prompts to return
offset: Number of prompts to skip (for pagination)
sort: Sort key from SORT_ORDERS (default: newest first)
Returns:
Dictionary containing prompt data and pagination info
@@ -310,7 +445,8 @@ class PromptDatabase:
# Get paginated results
cursor = conn.execute(
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts ORDER BY created_at DESC LIMIT ? OFFSET ?",
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts"
f" ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?",
(limit, offset),
)
rows = cursor.fetchall()
@@ -345,7 +481,8 @@ class PromptDatabase:
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? ORDER BY created_at DESC LIMIT ?",
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? "
"ORDER BY created_at DESC LIMIT ?",
(category, limit),
)
rows = cursor.fetchall()
@@ -421,7 +558,7 @@ class PromptDatabase:
params.append(datetime.datetime.now(datetime.timezone.utc).isoformat())
params.append(prompt_id)
query = f"UPDATE prompts SET {', '.join(updates)} WHERE id = ?"
cursor = conn.execute(query, params)
conn.execute(query, params)
if tags is not None:
self._sync_prompt_tags(conn, prompt_id, tags)
conn.execute(
@@ -495,7 +632,8 @@ class PromptDatabase:
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
"SELECT DISTINCT TRIM(category) as category FROM prompts WHERE category IS NOT NULL AND TRIM(category) != '' ORDER BY category"
"SELECT DISTINCT TRIM(category) as category FROM prompts WHERE "
"category IS NOT NULL AND TRIM(category) != '' ORDER BY category"
)
return [row["category"] for row in cursor.fetchall()]
@@ -536,7 +674,8 @@ class PromptDatabase:
rel = os.path.relpath(parent, root)
if not rel.startswith(".."):
if rel != ".":
folders.add(rel)
# Clients split and filter on "/" whatever the OS
folders.add(rel.replace(os.sep, "/"))
made_relative = True
break
except ValueError:
@@ -546,7 +685,7 @@ class PromptDatabase:
# Preserve full relative-style path instead of collapsing
# to just the basename, which loses hierarchy and creates
# ambiguity (e.g. foo/bar and baz/bar both become "bar").
folders.add(parent)
folders.add(parent.replace(os.sep, "/"))
if include_ancestors:
ancestors = set()
@@ -596,8 +735,8 @@ class PromptDatabase:
search_clause = ""
params: list = []
if search:
search_clause = "HAVING t.name LIKE ?"
params.append(f"%{search}%")
search_clause = "HAVING t.name LIKE ? ESCAPE '\\'"
params.append(f"%{escape_like(search)}%")
sort_map = {
"alpha_desc": "tag COLLATE NOCASE DESC",
@@ -638,6 +777,53 @@ class PromptDatabase:
"has_more": (offset + limit) < total,
}
def suggest_tags(
self, prefix: str, with_tags: List[str], limit: int = 15
) -> List[Dict[str, Any]]:
"""Tag autocomplete narrowed by co-occurrence.
Returns tags whose name starts with ``prefix`` (case-insensitive),
counted over the prompts that carry every tag in ``with_tags``; the
context tags themselves are left out. Ordered by count, then name.
"""
seen = set()
context = []
for raw in with_tags or []:
name = str(raw).strip()
if name and name.lower() not in seen:
seen.add(name.lower())
context.append(name)
where = ["t.name LIKE ? ESCAPE '\\'"]
params: list = [f"{escape_like(prefix or '')}%"]
if context:
placeholders = ",".join("?" * len(context))
where.append(
"pt.prompt_id IN ("
" SELECT pt2.prompt_id FROM prompt_tags pt2"
" JOIN tags t2 ON t2.id = pt2.tag_id"
f" WHERE t2.name COLLATE NOCASE IN ({placeholders})"
" GROUP BY pt2.prompt_id"
" HAVING COUNT(DISTINCT lower(t2.name)) = ?)"
)
params.extend(context)
params.append(len(context))
where.append(f"t.name COLLATE NOCASE NOT IN ({placeholders})")
params.extend(context)
params.append(max(1, int(limit)))
sql = (
"SELECT t.name AS name, COUNT(*) AS count"
" FROM prompt_tags pt JOIN tags t ON t.id = pt.tag_id"
f" WHERE {' AND '.join(where)}"
" GROUP BY t.id"
" ORDER BY count DESC, t.name COLLATE NOCASE ASC"
" LIMIT ?"
)
with self.model.get_connection() as conn:
rows = conn.execute(sql, params).fetchall()
return [{"name": row["name"], "count": row["count"]} for row in rows]
def get_prompts_by_tags(
self, tags: List[str], mode: str = "and", limit: int = 20, offset: int = 0
) -> Dict[str, Any]:
@@ -870,7 +1056,8 @@ class PromptDatabase:
conn.commit()
self.logger.info(
f"Merged {tags_merged} tags into '{target_tag}', affected {affected} prompts"
f"Merged {tags_merged} tags into '{target_tag}', affected {affected} "
"prompts"
)
return {
"success": True,
@@ -889,7 +1076,8 @@ class PromptDatabase:
with self.model.get_connection() as conn:
cursor = conn.execute(
"SELECT COUNT(*) as total FROM prompts "
"WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = prompts.id)"
"WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = "
"prompts.id)"
)
return cursor.fetchone()["total"]
@@ -1002,7 +1190,8 @@ class PromptDatabase:
tag_id = tag_map.get(tag_name)
if tag_id:
conn.execute(
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)",
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
"VALUES (?, ?)",
(prompt_id, tag_id),
)
@@ -1091,79 +1280,6 @@ class PromptDatabase:
self.logger.error(f"Error exporting prompts: {e}")
return False
def find_duplicates(self) -> List[Dict[str, Any]]:
"""
Find duplicate prompts based on text content without removing them.
Returns:
List of duplicate groups, each containing:
- text: The duplicate text content
- prompts: List of prompt records with same text
"""
self.logger.info("Scanning for duplicate prompts")
try:
with self.model.get_connection() as conn:
# Find duplicates by text content (case-insensitive)
# Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility
# We'll sort the IDs manually after fetching
cursor = conn.execute("""
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
GROUP_CONCAT(id) as ids,
GROUP_CONCAT(created_at) as created_dates
FROM prompts
GROUP BY LOWER(TRIM(text))
HAVING COUNT(*) > 1
""")
duplicate_groups = cursor.fetchall()
self.logger.debug(
f"Found {len(duplicate_groups)} groups of duplicate prompts"
)
result = []
for group in duplicate_groups:
ids = group["ids"].split(",")
created_dates = group["created_dates"].split(",")
# Sort IDs by created_at date
id_date_pairs = list(zip(ids, created_dates))
id_date_pairs.sort(key=lambda x: x[1]) # Sort by date
ids = [pair[0] for pair in id_date_pairs]
# Get full details for all prompts in this duplicate group
prompts = []
for prompt_id in ids:
cursor = conn.execute(
"SELECT id, text, category, rating, created_at, updated_at "
"FROM prompts WHERE id = ?",
(int(prompt_id),),
)
prompt_data = cursor.fetchone()
if prompt_data:
prompt_dict = dict(prompt_data)
prompt_dict["tags"] = self._get_prompt_tags(
conn, int(prompt_id)
)
prompts.append(prompt_dict)
if prompts:
result.append(
{
"text": prompts[0][
"text"
], # Use the actual text (not normalized)
"prompts": prompts,
}
)
self.logger.info(f"Found {len(result)} groups with duplicates")
return result
except Exception as e:
self.logger.error(f"Error finding duplicates: {e}")
return []
def cleanup_duplicates(self) -> int:
"""
Remove duplicate prompts based on text content, preserving all image links.
@@ -1208,7 +1324,8 @@ class PromptDatabase:
duplicate_ids = sorted_ids[1:]
self.logger.debug(
f"Merging duplicates: keeping {primary_id}, removing {duplicate_ids}"
f"Merging duplicates: keeping {primary_id}, removing "
f"{duplicate_ids}"
)
# Get primary prompt details
@@ -1248,7 +1365,8 @@ class PromptDatabase:
if total_removed > 0:
self.logger.info(
f"Removed {total_removed} duplicate prompts, transferred {total_images_transferred} images"
f"Removed {total_removed} duplicate prompts, transferred "
f"{total_images_transferred} images"
)
else:
self.logger.info("No duplicate prompts found")
@@ -1301,48 +1419,46 @@ class PromptDatabase:
)
if not cursor.fetchone():
self.logger.warning(
f"Prompt ID {prompt_id_int} not found in database, skipping image linking"
f"Prompt ID {prompt_id_int} not found in database, skipping "
"image linking"
)
return 0
# Proceed with linking
if not isinstance(metadata, dict):
metadata = {}
filename = os.path.basename(image_path)
file_info = metadata.get("file_info", {}) if metadata else {}
file_path = normalize_image_path(image_path)
file_size, width, height, image_format = _coerce_file_info(metadata)
# Use INSERT OR IGNORE to skip duplicates (same prompt_id + filename)
# INSERT OR IGNORE skips a file already linked to this prompt
# (same prompt_id + normalised full path)
cursor = conn.execute(
"""
INSERT OR IGNORE INTO generated_images
(prompt_id, image_path, filename, file_size, width, height, format,
workflow_data, prompt_metadata, parameters)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
(prompt_id, image_path, filename, file_path, file_size, width,
height, format, workflow_data, prompt_metadata, parameters)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
prompt_id_int,
image_path,
filename,
file_info.get("size"),
(
file_info.get("dimensions", [None, None])[0]
if file_info.get("dimensions")
else None
),
(
file_info.get("dimensions", [None, None])[1]
if file_info.get("dimensions")
else None
),
file_info.get("format"),
json.dumps(metadata.get("workflow", {}) if metadata else {}),
json.dumps(metadata.get("prompt", {}) if metadata else {}),
json.dumps(metadata.get("parameters", {}) if metadata else {}),
file_path,
file_size,
width,
height,
image_format,
json.dumps(metadata.get("workflow", {})),
json.dumps(metadata.get("prompt", {})),
json.dumps(metadata.get("parameters", {})),
),
)
conn.commit()
if cursor.lastrowid == 0:
if cursor.rowcount == 0:
self.logger.debug(
f"Image {filename} already linked to prompt {prompt_id_int}"
f"Image {file_path} already linked to prompt {prompt_id_int}"
)
return 0
@@ -1373,12 +1489,15 @@ class PromptDatabase:
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
def get_recent_images(self, limit: int = 50) -> List[Dict[str, Any]]:
def get_recent_images(
self, limit: int = 50, offset: int = 0
) -> List[Dict[str, Any]]:
"""
Get recently generated images across all prompts.
Get recently generated images across all prompts, newest first.
Args:
limit: Maximum number of images to return
offset: Number of newest images to skip (paging)
Returns:
List of image records with prompt text
@@ -1390,9 +1509,9 @@ class PromptDatabase:
FROM generated_images gi
LEFT JOIN prompts p ON gi.prompt_id = p.id
ORDER BY gi.generation_time DESC
LIMIT ?
LIMIT ? OFFSET ?
""",
(limit,),
(int(limit), max(0, int(offset))),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
@@ -1414,7 +1533,8 @@ class PromptDatabase:
sql = (
"SELECT gi.*, p.text as prompt_text, "
"(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt "
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS _prompt_tags_list "
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS "
"_prompt_tags_list "
"FROM generated_images gi "
"INNER JOIN prompts p ON gi.prompt_id = p.id "
"WHERE gi.image_path IS NOT NULL AND gi.image_path != '' "
@@ -1438,12 +1558,16 @@ class PromptDatabase:
result.append(data)
return result
def search_images_by_prompt(self, search_term: str) -> List[Dict[str, Any]]:
def search_images_by_prompt(
self, search_term: str, limit: Optional[int] = None, offset: int = 0
) -> List[Dict[str, Any]]:
"""
Search images by prompt text.
Search images by prompt text, newest first.
Args:
search_term: Text to search for in prompt content
limit: Maximum number of images to return (None for no limit)
offset: Number of newest matches to skip (paging)
Returns:
List of image records with prompt text
@@ -1454,10 +1578,15 @@ class PromptDatabase:
SELECT gi.*, p.text as prompt_text
FROM generated_images gi
JOIN prompts p ON gi.prompt_id = p.id
WHERE p.text LIKE ?
WHERE p.text LIKE ? ESCAPE '\\'
ORDER BY gi.generation_time DESC
LIMIT ? OFFSET ?
""",
(f"%{search_term}%",),
(
f"%{escape_like(search_term)}%",
-1 if limit is None else int(limit),
max(0, int(offset)),
),
)
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
@@ -1628,7 +1757,8 @@ class PromptDatabase:
if cursor.rowcount > 0:
self.logger.debug(
f"Transferred {cursor.rowcount} images from prompt {dup_id} to {primary_id}"
f"Transferred {cursor.rowcount} images from prompt {dup_id} "
f"to {primary_id}"
)
return transferred_count
@@ -1726,16 +1856,20 @@ class PromptDatabase:
}
def update_prompt_text(self, prompt_id: int, new_text: str) -> bool:
"""Update only the text of a prompt.
"""Update prompt text and its deduplication hash atomically.
Returns:
True if a row was updated.
Raises:
sqlite3.IntegrityError: Another prompt already has the new hash.
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
"UPDATE prompts SET text = ?, updated_at = ? WHERE id = ?",
"UPDATE prompts SET text = ?, hash = ?, updated_at = ? WHERE id = ?",
(
new_text,
generate_prompt_hash(new_text),
datetime.datetime.now(datetime.timezone.utc).isoformat(),
prompt_id,
),
@@ -1806,7 +1940,8 @@ class PromptDatabase:
tag_id = tag_map.get(tag_name)
if tag_id:
cursor = conn.execute(
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)",
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
"VALUES (?, ?)",
(pid, tag_id),
)
if cursor.rowcount > 0:
@@ -1838,20 +1973,29 @@ class PromptDatabase:
conn.commit()
return count
def get_image_prompt_info(self, image_path: str) -> Optional[Dict[str, Any]]:
def get_image_prompt_info(
self, image_path: str, allow_filename_fallback: bool = True
) -> Optional[Dict[str, Any]]:
"""Look up prompt + image metadata for a given image path.
Returns a dict with prompt_id, text, category, tags, rating, notes,
workflow_data, prompt_metadata, generation_time — or None.
workflow_data, prompt_metadata, generation_time, parameters — or None.
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
"""SELECT gi.prompt_id, p.text, p.category, p.rating, p.notes,
gi.workflow_data, gi.prompt_metadata, gi.generation_time
gi.workflow_data, gi.prompt_metadata, gi.generation_time,
gi.parameters
FROM generated_images gi
JOIN prompts p ON gi.prompt_id = p.id
WHERE gi.image_path = ? OR gi.image_path LIKE ?""",
(image_path, f"%{os.path.basename(image_path)}"),
WHERE gi.image_path = ? OR gi.file_path = ?
OR (? AND gi.image_path LIKE ? ESCAPE '\\')""",
(
image_path,
normalize_image_path(image_path),
allow_filename_fallback,
f"%{escape_like(os.path.basename(image_path))}",
),
)
row = cursor.fetchone()
if not row:
@@ -1875,14 +2019,16 @@ class PromptDatabase:
"workflow_data": workflow,
"prompt_metadata": prompt_meta,
"generation_time": row[7],
"parameters": row[8] if len(row) > 8 else None,
}
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:
cursor = conn.execute(
"SELECT prompt_id FROM generated_images WHERE image_path = ?",
(image_path,),
"SELECT prompt_id FROM generated_images"
" WHERE image_path = ? OR file_path = ?",
(image_path, normalize_image_path(image_path)),
)
row = cursor.fetchone()
return row[0] if row else None
@@ -1937,7 +2083,8 @@ class PromptDatabase:
""")
for ref in cursor.fetchall():
issues.append(
f"prompt_tags entry references non-existent prompt {ref['prompt_id']}"
"prompt_tags entry references non-existent prompt "
f"{ref['prompt_id']}"
)
# Check for orphaned image entries
@@ -1948,7 +2095,8 @@ class PromptDatabase:
""")
for ref in cursor.fetchall():
issues.append(
f"Image {ref['id']} references non-existent prompt {ref['prompt_id']}"
f"Image {ref['id']} references non-existent prompt "
f"{ref['prompt_id']}"
)
return issues
+61 -36
View File
@@ -4,7 +4,7 @@ with persistent prompt storage and search capabilities.
"""
import time
from typing import Any, Tuple
from typing import Any, Dict, Optional, Tuple
try:
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
@@ -51,7 +51,9 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
{
"multiline": True,
"dynamicPrompts": True,
"tooltip": "The text prompt to be encoded and saved to database.",
"tooltip": (
"The text prompt to be encoded and saved to database."
),
},
),
"clip": (
@@ -64,49 +66,73 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
IO.STRING,
{
"default": "",
"tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')",
"tooltip": (
"Optional category for organizing prompts "
"(e.g., 'landscapes', 'portraits')"
),
},
),
"tags": (
IO.STRING,
{
"default": "",
"tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')",
"tooltip": (
"Comma-separated tags for the prompt "
"(e.g., 'anime, detailed, sunset')"
),
},
),
"search_text": (
IO.STRING,
{
"default": "",
"tooltip": 'Search past prompts by text; add tag:name (or tag:"two words") to filter by tag',
"tooltip": (
"Search past prompts by text; add tag:name "
'(or tag:"two words") to filter by tag'
),
},
),
"prepend_text": (
IO.STRING,
{
"tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)"
"tooltip": (
"Text to prepend to the main prompt (connected STRING "
"nodes will be added before the main text)"
)
},
),
"append_text": (
IO.STRING,
{
"tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)"
"tooltip": (
"Text to append to the main prompt (connected STRING "
"nodes will be added after the main text)"
)
},
),
},
"hidden": {
# The queued graph and this node's id: used to tell the positive
# prompt of a run from negative-prompt PromptManager nodes
"prompt": "PROMPT",
"unique_id": "UNIQUE_ID",
},
}
RETURN_TYPES = (IO.CONDITIONING, IO.STRING)
OUTPUT_TOOLTIPS = (
"A conditioning containing the embedded text used to guide the diffusion model.",
"The final combined text string (with prepend/append applied) that was encoded.",
"A conditioning containing the embedded text used to guide the diffusion "
"model.",
"The final combined text string (with prepend/append applied) that was "
"encoded.",
)
FUNCTION = "encode_prompt"
OUTPUT_NODE = True
CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
DESCRIPTION = (
"Encodes a text prompt using a CLIP model into an embedding that can be used to guide "
"the diffusion model towards generating specific images. Additionally saves all prompts "
"Encodes a text prompt using a CLIP model into an embedding that can be used "
"to guide the diffusion model towards generating specific images. "
"Additionally saves all prompts "
"to a local SQLite database with optional metadata for search and retrieval."
)
@@ -119,6 +145,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
search_text: str = "",
prepend_text: str = "",
append_text: str = "",
prompt: Optional[Dict[str, Any]] = None,
unique_id: Optional[str] = None,
) -> Tuple[Any]:
"""
Encode the text prompt and save it to the database.
@@ -133,7 +161,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
append_text: Text to append to the main prompt
Returns:
Tuple containing the conditioning for the diffusion model and the final text string
Tuple containing the conditioning for the diffusion model and the final
text string
Raises:
RuntimeError: If clip input is invalid
@@ -154,7 +183,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
# Use the combined text for encoding
encoding_text = final_text
# For database storage, save the original main text with metadata about prepend/append
# For database storage, save the original main text with metadata
# about prepend/append
storage_text = text
# Validate CLIP model
@@ -181,30 +211,24 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
extended_tags.append(f"append:{append_text.strip()[:50]}")
try:
prompt_id = self._save_prompt_to_database(
prompt_id = self._track_prompt_execution(
text=storage_text.strip(),
encoding_text=encoding_text.strip(),
category=category.strip() if category else None,
tags=extended_tags if extended_tags else None,
additional_data={
"category": category.strip() if category else None,
"tags": extended_tags,
"prepend_text": (
prepend_text.strip() if prepend_text else None
),
"append_text": append_text.strip() if append_text else None,
"final_text": encoding_text.strip(),
},
prompt_graph=prompt,
unique_id=unique_id,
)
# Set current prompt for image tracking
if prompt_id:
execution_id = self.prompt_tracker.set_current_prompt(
prompt_text=encoding_text.strip(),
additional_data={
"category": category.strip() if category else None,
"tags": extended_tags,
"prompt_id": prompt_id,
"prepend_text": (
prepend_text.strip() if prepend_text else None
),
"append_text": append_text.strip() if append_text else None,
"final_text": encoding_text.strip(),
},
)
self.logger.debug(
f"Set execution context: {execution_id} for prompt ID: {prompt_id}"
)
self.logger.debug(f"Tracked prompt ID: {prompt_id}")
except Exception as e:
# Log error but don't fail the encoding
@@ -249,10 +273,11 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
"""
ComfyUI method to determine if node needs re-execution.
Returns a hash of input values that affect the conditioning output.
This enables proper branch execution - only re-execute when inputs change.
Returns a hash of input values that affect the conditioning output or
the saved prompt's metadata. This enables proper branch execution -
only re-execute when inputs change.
"""
import hashlib
combined = f"{text}|{prepend_text}|{append_text}"
combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}"
return hashlib.sha256(combined.encode()).hexdigest()
+86 -39
View File
@@ -6,9 +6,7 @@ search, gallery system management, and cleanup — extracted from the
duplicate code in prompt_manager.py and prompt_manager_text.py.
"""
import hashlib
import os
import webbrowser
from typing import Any, Dict, List, Optional
try:
@@ -24,6 +22,9 @@ try:
from .utils.comfyui_integration import get_comfyui_integration
from .utils.image_monitor import get_image_monitor
from .utils.prompt_tracker import get_prompt_tracker
from .utils.hashing import generate_prompt_hash
from .utils.prompt_graph import resolve_text, run_prompt_nodes
from .utils.usage_tracking import PENDING_FIRST_USE
except ImportError:
import sys
@@ -32,6 +33,9 @@ except ImportError:
from utils.comfyui_integration import get_comfyui_integration
from utils.image_monitor import get_image_monitor
from utils.prompt_tracker import get_prompt_tracker
from utils.hashing import generate_prompt_hash
from utils.prompt_graph import resolve_text, run_prompt_nodes
from utils.usage_tracking import PENDING_FIRST_USE
class PromptManagerBase:
@@ -56,14 +60,23 @@ class PromptManagerBase:
self.logger.debug(f"{self.__class__.__name__} node initialization completed")
def _save_prompt_to_database(
self, text: str, category: Optional[str] = None, tags: Optional[list] = None
self,
text: str,
category: Optional[str] = None,
tags: Optional[list] = None,
count_run: bool = False,
) -> Optional[int]:
"""Save the prompt to the SQLite database.
Runs are normally counted by the queue hook (utils/usage_tracking.py), so a
save only counts one when asked to (batch items whose text the hook could not
see) or when the hook queued this prompt before it existed.
Args:
text: The prompt text
category: Optional category
tags: List of tags
count_run: Count this execution as a run of the prompt
Returns:
The prompt ID if saved successfully, None otherwise
@@ -82,6 +95,8 @@ class PromptManagerBase:
prompt_id=existing["id"], category=category, tags=tags
)
self.logger.debug("Updated metadata for existing prompt")
if count_run:
self._record_use(existing["id"])
return existing["id"]
self.logger.debug(
@@ -93,6 +108,10 @@ class PromptManagerBase:
if prompt_id:
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
runs = pending.consume(prompt_hash) + (1 if count_run else 0)
if runs:
self._record_use(prompt_id, times=runs)
else:
self.logger.warning("Failed to save prompt - no ID returned")
@@ -102,6 +121,68 @@ class PromptManagerBase:
self.logger.error(f"Error saving prompt to database: {e}")
return None
def _prompt_role(
self, prompt_graph: Any, unique_id: Any, text: str
) -> Dict[str, bool]:
"""How this node's execution relates to the run, from the queued graph.
Returns flags:
is_run_prompt: feeds a sampler's positive input (images belong to it)
seen_at_queue: the queue hook resolved exactly this text, so it already
counted the run and the image's metadata can identify the prompt
known: the graph and node id were available
"""
if not isinstance(prompt_graph, dict) or unique_id is None:
return {"is_run_prompt": True, "seen_at_queue": False, "known": False}
node_id = str(unique_id)
resolved = resolve_text(prompt_graph, node_id)
return {
"is_run_prompt": node_id in run_prompt_nodes(prompt_graph),
"seen_at_queue": resolved is not None
and generate_prompt_hash(resolved) == generate_prompt_hash(text),
"known": True,
}
def _track_prompt_execution(
self,
text: str,
encoding_text: str,
category: Optional[str],
tags: Optional[list],
additional_data: Dict[str, Any],
prompt_graph: Any = None,
unique_id: Any = None,
) -> Optional[int]:
"""Save the prompt and, for the run's positive prompt, register it for linking.
Negative-prompt nodes are saved but never become the prompt their images are
linked to. Text the queue hook could resolve is counted there and linked from
the image's own metadata; anything else (PromptSearchList batches, unknown
text nodes, or a resolution that doesn't match) is counted here and pushed
onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept.
"""
role = self._prompt_role(prompt_graph, unique_id, text)
batch_item = (
role["known"] and role["is_run_prompt"] and not role["seen_at_queue"]
)
prompt_id = self._save_prompt_to_database(
text=text, category=category, tags=tags, count_run=batch_item
)
if prompt_id and role["is_run_prompt"]:
self.prompt_tracker.set_current_prompt(
prompt_text=encoding_text,
additional_data={**additional_data, "prompt_id": prompt_id},
push_to_queue=batch_item or not role["known"],
)
return prompt_id
def _record_use(self, prompt_id: int, times: int = 1) -> None:
"""Count runs; a tracking failure must never block generation."""
try:
self.db.record_prompt_use(prompt_id, times=times)
except Exception as e:
self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}")
def _inject_lora_trigger_words(self, text: str) -> str:
"""Append LoRA trigger words if integration is enabled.
@@ -142,16 +223,8 @@ class PromptManagerBase:
return modified
def _generate_hash(self, text: str) -> str:
"""Generate SHA256 hash for the prompt text.
Args:
text: The prompt text to hash
Returns:
Hexadecimal string representation of the SHA256 hash
"""
normalized_text = text.strip().lower()
return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest()
"""The prompt's deduplication hash (see utils.hashing.generate_prompt_hash)."""
return generate_prompt_hash(text)
def _parse_tags(self, tags_string: str) -> Optional[list]:
"""Parse comma-separated tags string into a list.
@@ -202,32 +275,6 @@ class PromptManagerBase:
self.logger.error(f"Error searching prompts: {e}")
return []
def _open_web_interface(self):
"""Open the web interface in the default browser."""
try:
current_dir = os.path.dirname(os.path.abspath(__file__))
web_dir = os.path.join(current_dir, "web_interface")
if os.path.exists(web_dir):
index_path = os.path.join(web_dir, "index.html")
if os.path.exists(index_path):
webbrowser.open(f"file://{index_path}")
self.logger.info("Web interface opened in browser")
else:
self.logger.warning(
f"Web interface directory found but no index.html. "
f"Please check {web_dir} for setup instructions"
)
else:
self.logger.info(
"Web interface not yet implemented. This feature will open a "
"web-based prompt management interface when the web_interface "
"directory is created."
)
except Exception as e:
self.logger.error(f"Error opening web interface: {e}")
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
"""API method for JavaScript UI to search prompts."""
return self._search_prompts(search_text=search_text)
+57 -34
View File
@@ -4,7 +4,7 @@ without CLIP encoding, while maintaining all database and search features.
"""
import time
from typing import Tuple
from typing import Any, Dict, Optional, Tuple
try:
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
@@ -53,7 +53,9 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
{
"multiline": True,
"dynamicPrompts": True,
"tooltip": "The text prompt to be processed and saved to database.",
"tooltip": (
"The text prompt to be processed and saved to database."
),
},
)
},
@@ -62,51 +64,74 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
IO.STRING,
{
"default": "",
"tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')",
"tooltip": (
"Optional category for organizing prompts "
"(e.g., 'landscapes', 'portraits')"
),
},
),
"tags": (
IO.STRING,
{
"default": "",
"tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')",
"tooltip": (
"Comma-separated tags for the prompt "
"(e.g., 'anime, detailed, sunset')"
),
},
),
"search_text": (
IO.STRING,
{
"default": "",
"tooltip": 'Search past prompts by text; add tag:name (or tag:"two words") to filter by tag',
"tooltip": (
"Search past prompts by text; add tag:name "
'(or tag:"two words") to filter by tag'
),
},
),
"prepend_text": (
IO.STRING,
{
"default": "",
"tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)",
"tooltip": (
"Text to prepend to the main prompt (connected STRING "
"nodes will be added before the main text)"
),
},
),
"append_text": (
IO.STRING,
{
"default": "",
"tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)",
"tooltip": (
"Text to append to the main prompt (connected STRING "
"nodes will be added after the main text)"
),
},
),
},
"hidden": {
# The queued graph and this node's id: used to tell the positive
# prompt of a run from negative-prompt PromptManager nodes
"prompt": "PROMPT",
"unique_id": "UNIQUE_ID",
},
}
RETURN_TYPES = (IO.STRING,)
OUTPUT_TOOLTIPS = (
"The final combined text string (with prepend/append applied) ready for use in other nodes.",
"The final combined text string (with prepend/append applied) ready for use "
"in other nodes.",
)
FUNCTION = "process_text"
OUTPUT_NODE = True
CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
DESCRIPTION = (
"Processes and manages text prompts with database storage and search capabilities. "
"Outputs a plain STRING that can be used with any node that accepts text input. "
"Includes all PromptManager features: categorization, tagging, search, and prepend/append functionality."
"Processes and manages text prompts with database storage and search "
"capabilities. Outputs a plain STRING that can be used with any node that "
"accepts text input. Includes all PromptManager features: categorization, "
"tagging, search, and prepend/append functionality."
)
def process_text(
@@ -117,6 +142,8 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
search_text: str = "",
prepend_text: str = "",
append_text: str = "",
prompt: Optional[Dict[str, Any]] = None,
unique_id: Optional[str] = None,
) -> Tuple[str]:
"""
Process the text prompt and save it to the database.
@@ -145,7 +172,8 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
# Inject LoRA trigger words if integration is enabled
final_text = self._inject_lora_trigger_words(final_text)
# For database storage, save the original main text with metadata about prepend/append
# For database storage, save the original main text with metadata
# about prepend/append
storage_text = text
# Save prompt to database and set execution context for gallery tracking
@@ -162,29 +190,23 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
extended_tags.append(f"append:{append_text.strip()[:50]}")
try:
prompt_id = self._save_prompt_to_database(
prompt_id = self._track_prompt_execution(
text=storage_text.strip(),
encoding_text=final_text.strip(),
category=category.strip() if category else None,
tags=extended_tags if extended_tags else None,
additional_data={
"category": category.strip() if category else None,
"tags": extended_tags,
"prepend_text": (
prepend_text.strip() if prepend_text else None
),
"append_text": append_text.strip() if append_text else None,
},
prompt_graph=prompt,
unique_id=unique_id,
)
# Set current prompt for image tracking
if prompt_id:
execution_id = self.prompt_tracker.set_current_prompt(
prompt_text=final_text.strip(),
additional_data={
"category": category.strip() if category else None,
"tags": extended_tags,
"prompt_id": prompt_id,
"prepend_text": (
prepend_text.strip() if prepend_text else None
),
"append_text": append_text.strip() if append_text else None,
},
)
self.logger.debug(
f"Set execution context: {execution_id} for prompt ID: {prompt_id}"
)
self.logger.debug(f"Tracked prompt ID: {prompt_id}")
except Exception as e:
# Log error but don't fail the processing
@@ -221,10 +243,11 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
"""
ComfyUI method to determine if node needs re-execution.
Returns a hash of input values that affect the text output.
This enables proper branch execution - only re-execute when inputs change.
Returns a hash of input values that affect the text output or the
saved prompt's metadata. This enables proper branch execution - only
re-execute when inputs change.
"""
import hashlib
combined = f"{text}|{prepend_text}|{append_text}"
combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}"
return hashlib.sha256(combined.encode()).hexdigest()
+7 -3
View File
@@ -6,7 +6,6 @@ OUTPUT_IS_LIST=True, allowing direct connection to nodes that accept list inputs
"""
import re
from typing import Any, Dict, List, Tuple
try:
from .utils.logging_config import get_logger
@@ -80,7 +79,9 @@ class PromptSearchList(ComfyNodeABC):
IO.STRING,
{
"default": "",
"tooltip": "Comma-separated list of tags to filter by (partial match)",
"tooltip": (
"Comma-separated list of tags to filter by (partial match)"
),
},
),
"min_rating": (
@@ -105,7 +106,10 @@ class PromptSearchList(ComfyNodeABC):
"BOOLEAN",
{
"default": True,
"tooltip": "Skip prompts containing Clip_1/Clip_2/etc. multi-part markers",
"tooltip": (
"Skip prompts containing Clip_1/Clip_2/etc. "
"multi-part markers"
),
},
),
},
+339 -209
View File
@@ -13,7 +13,9 @@ import datetime
import functools
import gzip as gzip_module
import json
import ntpath
import os
import re
from pathlib import Path
from aiohttp import web
@@ -29,6 +31,7 @@ from .lora_integration import LoraIntegrationMixin
try:
from ...database.operations import PromptDatabase
from ...utils.logging_config import get_logger
from ...utils.video_metadata import is_video_path, read_video_metadata
except ImportError:
import sys
@@ -37,6 +40,7 @@ except ImportError:
)
from database.operations import PromptDatabase
from utils.logging_config import get_logger
from utils.video_metadata import is_video_path, read_video_metadata
def _get_project_root():
@@ -44,6 +48,147 @@ def _get_project_root():
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
def _public_path(path) -> str:
"""Render a server path for an API response without revealing the layout.
Returns the path relative to the ComfyUI base directory (or to the parent
of an opted-in extra gallery root), as a POSIX-style string, or just the
basename when the path is not under any of those anchors.
"""
if not path:
return ""
from ..config import GalleryConfig
try:
canonical = Path(os.path.normcase(os.path.realpath(str(path))))
except (OSError, ValueError):
return os.path.basename(str(path))
for anchor in GalleryConfig.path_anchors():
if canonical.is_relative_to(Path(anchor)):
return canonical.relative_to(Path(anchor)).as_posix() or "."
return canonical.name
def _public_error(exc: BaseException) -> str:
"""Error text safe for API responses: OSError paths reduced to basenames."""
if isinstance(exc, OSError):
parts = [exc.strerror or type(exc).__name__]
if exc.filename:
parts.append(os.path.basename(str(exc.filename)))
return ": ".join(parts)
return str(exc) or type(exc).__name__
# ── Static file serving ────────────────────────────────────────────────
# Only these extensions are ever served from web/lib and web/js; the map
# doubles as the allow-list.
STATIC_MIME_TYPES = {
".js": "application/javascript",
".css": "text/css",
".map": "application/json",
".woff": "font/woff",
".woff2": "font/woff2",
".ttf": "font/ttf",
".png": "image/png",
".svg": "image/svg+xml",
".ico": "image/x-icon",
".html": "text/html",
}
_STATIC_SEPARATORS = re.compile(r"[\\/]+")
def _resolve_static_file(root, filepath):
"""Real path of ``filepath`` under ``root``, or None when it must not be served.
The request path is untrusted. It is rejected when it is empty, carries a
``..`` segment, is absolute on any platform (a leading slash or backslash,
a drive letter such as ``C:`` or a UNC ``\\\\server\\share`` prefix all
re-anchor ``os.path.join`` on Windows), has an extension outside
STATIC_MIME_TYPES, or resolves (symlinks included) outside ``root``.
"""
if not isinstance(filepath, str) or not filepath:
return None
if (
os.path.isabs(filepath)
or ntpath.isabs(filepath)
or ntpath.splitdrive(filepath)[0]
):
return None
segments = _STATIC_SEPARATORS.split(filepath)
if any(segment in ("", ".", "..") for segment in segments):
return None
if os.path.splitext(filepath)[1].lower() not in STATIC_MIME_TYPES:
return None
try:
root_real = Path(os.path.normcase(os.path.realpath(root)))
candidate = Path(
os.path.normcase(os.path.realpath(os.path.join(root, *segments)))
)
except (OSError, ValueError):
return None
if candidate == root_real or not candidate.is_relative_to(root_real):
return None
return str(candidate)
# ── JSON request bodies ────────────────────────────────────────────────
JSON_BODY_MAX_BYTES = 1_000_000
_JSON_READ_CHUNK = 64 * 1024
def _json_error(message, status):
return web.json_response({"success": False, "error": message}, status=status)
async def _read_json_body(request, max_bytes=JSON_BODY_MAX_BYTES):
"""Read and parse a JSON object body of at most ``max_bytes``.
The body is read in chunks so an oversized or chunked upload is refused
as soon as it crosses the cap rather than buffered whole.
Returns:
(data, None) on success, or (None, response) where ``response`` is a
413 (body too large) or 400 (invalid JSON, or not a JSON object).
"""
too_large = _json_error(f"Request body too large (max {max_bytes} bytes)", 413)
declared = request.content_length
if declared is not None and declared > max_bytes:
return None, too_large
chunks = []
size = 0
async for chunk in request.content.iter_chunked(_JSON_READ_CHUNK):
size += len(chunk)
if size > max_bytes:
return None, too_large
chunks.append(chunk)
try:
data = json.loads(b"".join(chunks).decode("utf-8"))
except (UnicodeDecodeError, ValueError):
return None, _json_error("Request body must be valid JSON", 400)
if not isinstance(data, dict):
return None, _json_error("Request body must be a JSON object", 400)
return data, None
class _LongJobGuard:
"""Async context manager handed out by ``_acquire_long_job``."""
def __init__(self, lock):
self._lock = lock
async def __aenter__(self):
# The caller checked ``locked()`` on this same loop iteration, so this
# acquire never waits.
await self._lock.acquire()
return self
async def __aexit__(self, exc_type, exc, tb):
self._lock.release()
return False
# ── Gzip compression middleware ────────────────────────────────────────
_GZIP_MIN_SIZE = 1024 # Only compress bodies larger than 1 KB
_GZIP_TYPES = frozenset(
@@ -127,6 +272,14 @@ 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
# 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.
self._long_jobs: dict = {}
# The current output scan (py/api/scan_job.ScanJob), if one was started.
self._scan_job = None
# The current duplicate scan (also a ScanJob), if one was started.
self._duplicate_job = None
# Run cleanup on initialization to remove any existing duplicates
try:
@@ -140,6 +293,35 @@ class PromptManagerAPI(
self.logger.info("PromptManager API initialization completed")
_public_path = staticmethod(_public_path)
_public_error = staticmethod(_public_error)
_read_json_body = staticmethod(_read_json_body)
def _acquire_long_job(self, name):
"""Claim the single-flight slot for a long job family, or None when busy.
Usage in a handler::
job = self._acquire_long_job("duplicates")
if job is None:
return self._long_job_busy_response("Duplicate scan")
async with job:
...
Families are free-form strings ("duplicates", "thumbnails",
"autotag", ...); the lock is released when the ``async with`` block
exits, including on exceptions. Must be called from the event loop.
"""
lock = self._long_jobs.setdefault(name, asyncio.Lock())
if lock.locked():
return None
return _LongJobGuard(lock)
@staticmethod
def _long_job_busy_response(label):
"""409 envelope for a long job that is already running."""
return _json_error(f"{label} already running", 409)
async def _run_in_executor(self, func, *args, **kwargs):
"""Run a blocking function in the default thread pool executor.
@@ -181,191 +363,31 @@ class PromptManagerAPI(
@routes.get("/prompt_manager/web")
async def serve_web_ui(request):
try:
html_path = os.path.join(_get_project_root(), "web", "index.html")
if os.path.exists(html_path):
with open(html_path, "r", encoding="utf-8") as f:
html_content = f.read()
return web.Response(
text=html_content, content_type="text/html", charset="utf-8"
)
else:
return web.Response(
text="<h1>Web UI not found</h1><p>HTML file not located at expected path.</p>",
content_type="text/html",
status=404,
)
except Exception as e:
self.logger.exception("Failed to load web UI")
return web.Response(
text="<h1>Error</h1><p>Failed to load web UI. Check server logs for details.</p>",
content_type="text/html",
status=500,
)
return self._serve_html_page("index.html", "Web UI", cache=False)
@routes.get("/prompt_manager/gallery.html")
async def serve_gallery_ui(request):
try:
html_path = os.path.join(
_get_project_root(),
"web",
"metadata.html",
)
if html_path not in self._html_cache:
if os.path.exists(html_path):
with open(html_path, "r", encoding="utf-8") as f:
self._html_cache[html_path] = f.read()
else:
return web.Response(
text="<h1>Gallery not found</h1><p>gallery.html file not located at expected path.</p>",
content_type="text/html",
status=404,
)
return web.Response(
text=self._html_cache[html_path],
content_type="text/html",
charset="utf-8",
)
except Exception as e:
self.logger.exception("Failed to load gallery")
return web.Response(
text="<h1>Error</h1><p>Failed to load gallery. Check server logs for details.</p>",
content_type="text/html",
status=500,
)
return self._serve_html_page("metadata.html", "Gallery")
@routes.get("/prompt_manager/admin")
async def serve_admin_ui(request):
try:
html_path = os.path.join(
_get_project_root(),
"web",
"admin.html",
)
if html_path not in self._html_cache:
if os.path.exists(html_path):
with open(html_path, "r", encoding="utf-8") as f:
self._html_cache[html_path] = f.read()
else:
return web.Response(
text="<h1>Admin UI not found</h1>",
content_type="text/html",
status=404,
)
return web.Response(
text=self._html_cache[html_path],
content_type="text/html",
charset="utf-8",
)
except Exception as e:
self.logger.exception("Failed to load admin UI")
return web.Response(
text="<h1>Error</h1><p>Failed to load admin UI. Check server logs for details.</p>",
content_type="text/html",
status=500,
)
return self._serve_html_page("admin.html", "Admin UI")
@routes.get("/prompt_manager/gallery")
async def serve_gallery_admin_ui(request):
try:
html_path = os.path.join(
_get_project_root(),
"web",
"gallery.html",
)
if html_path not in self._html_cache:
if os.path.exists(html_path):
with open(html_path, "r", encoding="utf-8") as f:
self._html_cache[html_path] = f.read()
else:
return web.Response(
text="<h1>Gallery not found</h1><p>gallery.html file not located at expected path.</p>",
content_type="text/html",
status=404,
)
return web.Response(
text=self._html_cache[html_path],
content_type="text/html",
charset="utf-8",
)
except Exception as e:
self.logger.exception("Failed to load gallery")
return web.Response(
text="<h1>Error</h1><p>Failed to load gallery. Check server logs for details.</p>",
content_type="text/html",
status=500,
)
return self._serve_html_page("gallery.html", "Gallery")
# ── Static file serving ───────────────────────────────────────
@routes.get("/prompt_manager/lib/{filepath:.*}")
async def serve_lib_static(request):
"""Serve static library files (JS, CSS) from web/lib directory."""
MIME_TYPES = {
".js": "application/javascript",
".css": "text/css",
".json": "application/json",
".map": "application/json",
}
filepath = request.match_info.get("filepath", "")
# Security: prevent directory traversal
if ".." in filepath or filepath.startswith("/"):
return web.Response(text="Forbidden", status=403)
file_path = os.path.join(_get_project_root(), "web", "lib", filepath)
if not os.path.exists(file_path) or not os.path.isfile(file_path):
return web.Response(text=f"Not Found: {filepath}", status=404)
ext = os.path.splitext(file_path)[1].lower()
content_type = MIME_TYPES.get(ext, "application/octet-stream")
with open(file_path, "rb") as f:
content = f.read()
return web.Response(body=content, content_type=content_type)
return self._serve_static_file("lib", request.match_info.get("filepath"))
@routes.get("/prompt_manager/js/{filepath:.*}")
async def serve_js_static(request):
"""Serve static JavaScript files from web/js directory."""
MIME_TYPES = {
".js": "application/javascript",
".css": "text/css",
".json": "application/json",
".map": "application/json",
}
filepath = request.match_info.get("filepath", "")
if ".." in filepath or filepath.startswith("/"):
return web.Response(text="Forbidden", status=403)
file_path = os.path.join(_get_project_root(), "web", "js", filepath)
if not os.path.exists(file_path) or not os.path.isfile(file_path):
return web.Response(text=f"Not Found: {filepath}", status=404)
ext = os.path.splitext(file_path)[1].lower()
content_type = MIME_TYPES.get(ext, "application/octet-stream")
with open(file_path, "rb") as f:
content = f.read()
return web.Response(body=content, content_type=content_type)
return self._serve_static_file("js", request.match_info.get("filepath"))
# ── Register domain-specific routes from mixins ───────────────
@@ -390,6 +412,47 @@ class PromptManagerAPI(
self.logger.info("All routes registered with decorator pattern")
def _serve_static_file(self, subdir, filepath):
"""Serve ``web/<subdir>/<filepath>``; 403 outside it, 404 when missing."""
root = os.path.join(_get_project_root(), "web", subdir)
file_path = _resolve_static_file(root, filepath)
if file_path is None:
return web.Response(text="Forbidden", status=403)
if not os.path.isfile(file_path):
return web.Response(text="Not Found", status=404)
ext = os.path.splitext(file_path)[1].lower()
with open(file_path, "rb") as f:
content = f.read()
return web.Response(body=content, content_type=STATIC_MIME_TYPES[ext])
def _serve_html_page(self, filename, title, cache=True):
"""Serve ``web/<filename>``; 404 when missing, 500 when unreadable."""
html_path = os.path.join(_get_project_root(), "web", filename)
try:
if cache and html_path in self._html_cache:
content = self._html_cache[html_path]
elif os.path.exists(html_path):
with open(html_path, "r", encoding="utf-8") as f:
content = f.read()
if cache:
self._html_cache[html_path] = content
else:
return web.Response(
text=f"<h1>{title} not found</h1>"
"<p>HTML file not located at expected path.</p>",
content_type="text/html",
status=404,
)
return web.Response(text=content, content_type="text/html", charset="utf-8")
except Exception:
self.logger.exception(f"Failed to load {title}")
return web.Response(
text=f"<h1>Error</h1><p>Failed to load {title}. "
"Check server logs for details.</p>",
content_type="text/html",
status=500,
)
# ── Shared utilities used by multiple mixins ──────────────────────
def _enrich_images(self, images):
@@ -415,7 +478,8 @@ class PromptManagerAPI(
rel_path = img_path.resolve().relative_to(output_path.resolve())
image["relative_path"] = str(rel_path)
image["url"] = (
f"/prompt_manager/images/serve/{url_quote(rel_path.as_posix(), safe='/')}"
"/prompt_manager/images/serve/"
f"{url_quote(rel_path.as_posix(), safe='/')}"
)
# Check for thumbnail
@@ -423,11 +487,13 @@ class PromptManagerAPI(
thumb_rel = (
f"thumbnails/{rel_no_ext.as_posix()}_thumb{rel_path.suffix}"
)
thumb_abs = output_path / thumb_rel
if thumb_abs.exists():
image["thumbnail_url"] = (
f"/prompt_manager/images/serve/{url_quote(thumb_rel, safe='/')}"
)
from .prompts import thumbnail_url_for
thumbnail_url = thumbnail_url_for(
thumb_rel, output_path / thumb_rel
)
if thumbnail_url:
image["thumbnail_url"] = thumbnail_url
break # Found matching root, stop searching
except (ValueError, RuntimeError):
continue # Try next root
@@ -476,17 +542,18 @@ class PromptManagerAPI(
from ..config import GalleryConfig
if GalleryConfig.MONITORING_DIRECTORIES:
configured_dir = Path(GalleryConfig.MONITORING_DIRECTORIES[0]).resolve()
if configured_dir.exists() and configured_dir.is_dir():
configured = GalleryConfig.MONITORING_DIRECTORIES[0]
# Same check as the settings endpoint: a root hand-edited
# into config.json must not open a directory outside ComfyUI
ok, reason = GalleryConfig.validate_gallery_root(configured)
if ok:
configured_dir = Path(configured).resolve()
self.logger.info(
f"Using configured monitoring directory: {configured_dir}"
)
self._cached_output_dir = str(configured_dir)
return self._cached_output_dir
else:
self.logger.warning(
f"Configured directory does not exist: {GalleryConfig.MONITORING_DIRECTORIES[0]}"
)
self.logger.warning(f"Ignoring configured directory: {reason}")
except ImportError:
self.logger.debug(
"GalleryConfig not available, skipping configured directory check"
@@ -506,7 +573,8 @@ class PromptManagerAPI(
output_dir = current_dir / "output"
if output_dir.exists() and output_dir.is_dir():
self.logger.debug(
f"Found ComfyUI output directory via upward search: {output_dir}"
"Found ComfyUI output directory via upward search: "
f"{output_dir}"
)
self._cached_output_dir = str(output_dir)
return self._cached_output_dir
@@ -557,39 +625,72 @@ class PromptManagerAPI(
return None
def _get_all_output_dirs(self):
"""Get all configured output directories, falling back to auto-detect.
"""Get the directories images may be served from.
Configured gallery roots are filtered through
``GalleryConfig.validate_gallery_root`` so a root that was hand-edited
into config.json (or persisted before validation existed) can never
expose files outside ComfyUI's directories.
Returns:
List[Path]: Valid output directory paths, possibly empty.
List[Path]: Validated directory paths, possibly empty.
"""
from ..config import GalleryConfig
configured = list(GalleryConfig.MONITORING_DIRECTORIES)
output_dirs = []
if GalleryConfig.MONITORING_DIRECTORIES:
for d in GalleryConfig.MONITORING_DIRECTORIES:
p = Path(d).resolve()
if p.is_dir():
output_dirs.append(p)
seen = set()
for d in configured:
ok, reason = GalleryConfig.validate_gallery_root(d)
if not ok:
self.logger.warning(f"Ignoring configured gallery root: {reason}")
continue
real = os.path.realpath(d)
key = os.path.normcase(real) # one entry per directory, however spelled
if key in seen:
continue
seen.add(key)
output_dirs.append(Path(real))
# Fallback to auto-detect if no configured dirs are valid
if not output_dirs:
fallback = self._find_comfyui_output_dir()
if fallback:
output_dirs.append(Path(fallback))
if output_dirs:
return output_dirs
return output_dirs
if configured:
# Every configured root was rejected: never let auto-detection
# resurrect it, serve ComfyUI's own output directory instead.
return self._comfyui_output_dir_list()
fallback = self._find_comfyui_output_dir()
return [Path(fallback)] if fallback else []
def _comfyui_output_dir_list(self):
"""ComfyUI's output directory as a one-element list, or [] if unknown."""
try:
import folder_paths
output_dir = folder_paths.get_output_directory()
except (ImportError, AttributeError, OSError):
return []
if isinstance(output_dir, str) and os.path.isdir(output_dir):
return [Path(os.path.realpath(output_dir))]
return []
def _extract_comfyui_metadata(self, image_path):
"""Extract ComfyUI workflow metadata from PNG image files."""
"""Embedded ComfyUI metadata of an image (PNG text chunks) or a video.
Videos go through ffprobe (see utils/video_metadata). Unreadable files
are a warning, not an error: a scan over a large output tree meets
truncated and foreign files, and it carries on past them.
"""
if is_video_path(image_path):
return read_video_metadata(image_path)
try:
with Image.open(image_path) as img:
metadata = {}
if hasattr(img, "text"):
for key, value in img.text.items():
metadata[key] = value
return metadata
return dict(getattr(img, "text", None) or {})
except Exception as e:
self.logger.error(f"Error reading {image_path}: {e}")
self.logger.warning(
f"Unreadable image {os.path.basename(str(image_path))}: {e}"
)
return {}
def _parse_comfyui_prompt(self, metadata):
@@ -602,18 +703,43 @@ class PromptManagerAPI(
"negative_prompt": None,
}
# Check for A1111 style parameters first (like parse-metadata.py)
# Check for A1111/Forge style parameters first.
# The chunk is one block: a possibly multi-line positive prompt, then a
# "Negative prompt:" section, then a trailing "Steps: ..." settings line.
# Taking lines[0] as the positive prompt truncates any multi-line prompt
# to one line, so split on the settings line instead.
if "parameters" in metadata:
import re as _re
params = metadata["parameters"]
lines = params.splitlines()
if lines:
result["positive_prompt"] = lines[0].strip()
for line in lines:
if line.lower().startswith("negative prompt:"):
result["negative_prompt"] = line.split(":", 1)[1].strip()
break
body = params
settings = ""
_sm = _re.search(r"\nSteps:\s", body)
if _sm:
settings = body[_sm.start() + 1 :].strip()
body = body[: _sm.start()]
negative = None
_nm = _re.search(r"(?im)^\s*negative prompt:\s*", body)
if _nm:
negative = body[_nm.end() :].strip()
body = body[: _nm.start()]
elif body.lower().lstrip().startswith("negative prompt:"):
negative = body.lower().lstrip()[len("negative prompt:") :].strip()
body = ""
result["positive_prompt"] = body.strip() or None
if negative:
result["negative_prompt"] = negative
# Store raw parameters too
result["parameters"]["parameters"] = params
if settings:
for _f in ("Steps", "Sampler", "Schedule type", "CFG scale", "Seed"):
_fm = _re.search(
r"(?i)(?:^|,\s*)" + _re.escape(_f) + r":\s*([^,]+)", settings
)
if _fm:
result["parameters"][_f.lower().replace(" ", "_")] = _fm.group(
1
).strip()
# If no A1111 format found, proceed with ComfyUI parsing
if result["positive_prompt"] is None:
@@ -710,7 +836,8 @@ class PromptManagerAPI(
New format: inputs is a list of connection objects
inputs = [
{"name": "text", "type": "STRING", "link": null, "widget": {"name": "text"}},
{"name": "text", "type": "STRING", "link": null,
"widget": {"name": "text"}},
{"name": "clip", "type": "CLIP", "link": 11}
]
"""
@@ -773,7 +900,10 @@ class PromptManagerAPI(
return None
def _extract_positive_prompt_from_comfyui_data(self, data):
"""Extract positive prompt from ComfyUI data, handling both old and new formats."""
"""Extract positive prompt from ComfyUI data.
Handles both old and new formats.
"""
if not isinstance(data, dict):
return None
+1185 -917
View File
File diff suppressed because it is too large Load Diff
+533 -334
View File
@@ -2,10 +2,53 @@
import asyncio
import json
import os
import threading
from pathlib import Path
from typing import Optional
from urllib.parse import quote
from .prompts import thumbnail_url_for
from aiohttp import web
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
_BUSY_ERROR = "Auto-tagging is already running"
# Only these may be handed to the tagging engine or listed by the scanner.
_MEDIA_EXTENSIONS = frozenset(
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}
)
# Bounds for scan_output_dir: a runaway output tree must not pin the server.
_SCAN_MAX_FILES = 50_000
_SCAN_MAX_DEPTH = 12
# One batch run at a time; a second autotag/start gets 409 while this is held.
_autotag_batch_lock = threading.Lock()
def _file_exists(path) -> bool:
try:
return Path(str(path)).exists()
except (OSError, ValueError):
return False
def path_is_within(candidate, root) -> bool:
"""True when ``candidate`` is strictly inside ``root`` on the real filesystem.
Both sides go through ``os.path.realpath`` (symlinks, ``..``) and
``os.path.normcase`` (Windows case/separator folding) before the
``Path.is_relative_to`` check, so ``/out2`` is never treated as inside
``/out`` and a symlink inside ``root`` pointing elsewhere is rejected.
"""
try:
child = Path(os.path.normcase(os.path.realpath(str(candidate))))
parent = Path(os.path.normcase(os.path.realpath(str(root))))
except (OSError, RuntimeError, ValueError):
return False
return child != parent and child.is_relative_to(parent)
class AutotagRoutesMixin:
"""Mixin providing auto-tagging API endpoints."""
@@ -15,11 +58,11 @@ class AutotagRoutesMixin:
async def get_autotag_models_route(request):
return await self.get_autotag_models(request)
@routes.get("/prompt_manager/autotag/download/{model_type}")
@routes.post("/prompt_manager/autotag/download/{model_type}")
async def download_autotag_model_route(request):
return await self.download_autotag_model(request)
@routes.get("/prompt_manager/autotag/start")
@routes.post("/prompt_manager/autotag/start")
async def start_autotag_route(request):
return await self.start_autotag(request)
@@ -61,253 +104,17 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"Get autotag models error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
async def download_autotag_model(self, request):
"""Download an AutoTag model with streaming progress."""
model_type = request.match_info.get("model_type")
async def stream_response():
try:
from ..autotag import get_autotag_service
service = get_autotag_service()
if model_type not in service.models_config:
yield f"data: {json.dumps({'type': 'error', 'message': f'Invalid model type: {model_type}'})}\n\n"
return
yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Starting download...'})}\n\n"
progress_data = {"last_progress": 0}
def progress_callback(status: str, progress: float):
progress_data["last_progress"] = progress
loop = asyncio.get_event_loop()
success = await loop.run_in_executor(
None, lambda: service.download_model(model_type, progress_callback)
)
if success:
yield f"data: {json.dumps({'type': 'complete', 'progress': 100, 'status': 'Download complete'})}\n\n"
else:
yield f"data: {json.dumps({'type': 'error', 'message': 'Download failed'})}\n\n"
except Exception as e:
self.logger.exception("Download model error")
yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n"
response = web.StreamResponse(
status=200,
reason="OK",
headers={
"Content-Type": "text/event-stream",
"Cache-Control": "no-cache",
"Connection": "keep-alive",
},
)
await response.prepare(request)
async for chunk in stream_response():
await response.write(chunk.encode("utf-8"))
await response.write_eof()
return response
async def start_autotag(self, request):
"""Start batch auto-tagging with streaming progress."""
model_type = request.query.get("model_type", "gguf")
custom_prompt = request.query.get("prompt", "")
skip_tagged = request.query.get("skip_tagged", "true").lower() == "true"
keep_in_memory = request.query.get("keep_in_memory", "true").lower() == "true"
use_gpu = True
# WD14 threshold params
general_threshold = request.query.get("general_threshold")
character_threshold = request.query.get("character_threshold")
try:
if general_threshold is not None:
general_threshold = float(general_threshold)
if character_threshold is not None:
character_threshold = float(character_threshold)
except (ValueError, TypeError):
return web.json_response(
{
"success": False,
"error": "general_threshold and character_threshold must be numeric",
},
status=400,
{"success": False, "error": self._public_error(e)}, status=500
)
async def stream_response():
try:
from ..autotag import get_autotag_service
import time as _time
service = get_autotag_service()
status = service.get_models_status()
if not status.get(model_type, {}).get("downloaded"):
yield f"data: {json.dumps({'type': 'error', 'message': f'Model {model_type} not downloaded'})}\n\n"
return
yield f"data: {json.dumps({'type': 'progress', 'progress': 0, 'status': 'Loading model...'})}\n\n"
loop = asyncio.get_event_loop()
try:
await loop.run_in_executor(
None, lambda: service.load_model(model_type, use_gpu)
)
except Exception as e:
self.logger.exception("Failed to load model")
yield f"data: {json.dumps({'type': 'error', 'message': 'Failed to load model. Check server logs for details.'})}\n\n"
return
if custom_prompt:
service.custom_prompt = custom_prompt
yield f"data: {json.dumps({'type': 'progress', 'progress': 5, 'status': 'Model loaded. Fetching all images from database...'})}\n\n"
images = await self._run_in_executor(self.db.get_all_images)
total_files = len(images)
if total_files == 0:
yield f"data: {json.dumps({'type': 'complete', 'processed': 0, 'tagged': 0, 'skipped': 0, 'status': 'No images with linked prompts found'})}\n\n"
service.unload_model()
return
yield f"data: {json.dumps({'type': 'progress', 'progress': 10, 'status': f'Found {total_files} images. Processing...'})}\n\n"
processed = 0
tagged = 0
skipped = 0
errors = 0
tagged_prompt_ids = set()
last_update_time = _time.monotonic()
for i, image_data in enumerate(images):
image_path = image_data.get("image_path")
prompt_id = image_data.get("prompt_id")
if not image_path or not prompt_id:
skipped += 1
continue
if prompt_id in tagged_prompt_ids:
skipped += 1
now = _time.monotonic()
if (now - last_update_time) >= 0.5 or i == total_files - 1:
progress = 10 + int((i + 1) / total_files * 85)
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (prompt already processed)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n"
await asyncio.sleep(0.01)
last_update_time = now
continue
if not Path(image_path).exists():
skipped += 1
now = _time.monotonic()
if (now - last_update_time) >= 0.5 or i == total_files - 1:
progress = 10 + int((i + 1) / total_files * 85)
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (file missing)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n"
await asyncio.sleep(0.01)
last_update_time = now
continue
if skip_tagged:
prompt_tags = image_data.get("prompt_tags", [])
if isinstance(prompt_tags, str):
prompt_tags = [
t.strip() for t in prompt_tags.split(",") if t.strip()
]
real_tags = [
t
for t in prompt_tags
if t != "auto-scanned"
and not t.startswith("prepend:")
and not t.startswith("append:")
]
if real_tags:
tagged_prompt_ids.add(prompt_id)
skipped += 1
now = _time.monotonic()
if (now - last_update_time) >= 0.5 or i == total_files - 1:
progress = 10 + int((i + 1) / total_files * 85)
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Skipping {i+1}/{total_files} (already tagged)...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n"
await asyncio.sleep(0.01)
last_update_time = now
continue
try:
tags = await loop.run_in_executor(
None,
lambda p=str(image_path): service.generate_tags(
p,
general_threshold=general_threshold,
character_threshold=character_threshold,
),
)
processed += 1
if tags:
existing_prompt = await self._run_in_executor(
self.db.get_prompt_by_id, prompt_id
)
if existing_prompt:
existing_tags = existing_prompt.get("tags", [])
if isinstance(existing_tags, str):
existing_tags = [
t.strip()
for t in existing_tags.split(",")
if t.strip()
]
new_tags = [t for t in tags if t not in existing_tags]
if new_tags:
all_tags = existing_tags + new_tags
await self._run_in_executor(
self.db.update_prompt_metadata,
prompt_id,
tags=all_tags,
)
tagged += 1
else:
skipped += 1
else:
skipped += 1
else:
skipped += 1
tagged_prompt_ids.add(prompt_id)
progress = 10 + int((i + 1) / total_files * 85)
yield f"data: {json.dumps({'type': 'progress', 'progress': progress, 'status': f'Processing {i+1}/{total_files}...', 'processed': processed, 'tagged': tagged, 'skipped': skipped})}\n\n"
await asyncio.sleep(0.01)
last_update_time = _time.monotonic()
except Exception as img_err:
self.logger.error(f"Error processing {image_path}: {img_err}")
errors += 1
processed += 1
if not keep_in_memory:
service.unload_model()
model_status = "Model unloaded"
else:
model_status = "Model kept in memory"
yield f"data: {json.dumps({'type': 'complete', 'progress': 100, 'processed': processed, 'tagged': tagged, 'skipped': skipped, 'errors': errors, 'status': 'Complete', 'model_status': model_status, 'model_loaded': keep_in_memory})}\n\n"
except Exception as e:
self.logger.error(f"AutoTag error: {e}")
import traceback
self.logger.exception("AutoTag error")
yield f"data: {json.dumps({'type': 'error', 'message': 'An internal error occurred. Check server logs for details.'})}\n\n"
async def _stream_sse(self, request, events):
"""Write an async iterator of dict events as server-sent events.
A client disconnect raises ConnectionResetError from ``write``; the
event generator is then closed so its cleanup runs before re-raising
(aiohttp logs premature disconnects at debug level).
"""
response = web.StreamResponse(
status=200,
reason="OK",
@@ -317,20 +124,409 @@ class AutotagRoutesMixin:
"Connection": "keep-alive",
},
)
await response.prepare(request)
async for chunk in stream_response():
await response.write(chunk.encode("utf-8"))
try:
async for event in events:
await response.write(f"data: {json.dumps(event)}\n\n".encode("utf-8"))
except (ConnectionResetError, asyncio.CancelledError):
self.logger.info("Client disconnected; stopping autotag stream")
await events.aclose()
raise
await response.write_eof()
return response
async def download_autotag_model(self, request):
"""Download an AutoTag model with streaming progress (POST only)."""
model_type = request.match_info.get("model_type")
async def events():
try:
from ..autotag import get_autotag_service
service = get_autotag_service()
if model_type not in service.models_config:
yield {
"type": "error",
"message": f"Invalid model type: {model_type}",
}
return
yield {
"type": "progress",
"progress": 0,
"status": "Starting download...",
}
def progress_callback(status: str, progress: float):
pass
success = await self._run_in_executor(
service.download_model, model_type, progress_callback
)
if success:
yield {
"type": "complete",
"progress": 100,
"status": "Download complete",
}
else:
yield {"type": "error", "message": "Download failed"}
except Exception:
self.logger.exception("Download model error")
yield {
"type": "error",
"message": _INTERNAL_ERROR,
}
return await self._stream_sse(request, events())
async def _read_start_params(self, request) -> Optional[dict]:
"""Merge query-string and optional JSON-body parameters for autotag/start.
Returns None when a body is present but is not a JSON object.
"""
params = dict(request.query)
if request.can_read_body:
try:
body = await request.json()
except (json.JSONDecodeError, UnicodeDecodeError):
return None
if not isinstance(body, dict):
return None
params.update(body)
return params
@staticmethod
def _real_tags(tags) -> list:
"""Normalise a tag field to a list without bookkeeping markers."""
if isinstance(tags, str):
tags = [t.strip() for t in tags.split(",") if t.strip()]
return [
t
for t in (tags or [])
if t != "auto-scanned"
and not t.startswith("prepend:")
and not t.startswith("append:")
]
async def _apply_tags_to_prompt(self, prompt_id, tags) -> bool:
"""Merge ``tags`` into the prompt; True when something new was added."""
existing_prompt = await self._run_in_executor(
self.db.get_prompt_by_id, prompt_id
)
if not existing_prompt:
return False
existing_tags = existing_prompt.get("tags", [])
if isinstance(existing_tags, str):
existing_tags = [t.strip() for t in existing_tags.split(",") if t.strip()]
new_tags = [t for t in tags if t not in existing_tags]
if not new_tags:
return False
await self._run_in_executor(
self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags
)
return True
async def _autotag_batch_events(self, service, opts: dict):
"""Yield SSE progress dicts while tagging every image linked to a prompt."""
import time as _time
model_type = opts["model_type"]
keep_in_memory = opts["keep_in_memory"]
status = service.get_models_status()
if not status.get(model_type, {}).get("downloaded"):
yield {"type": "error", "message": f"Model {model_type} not downloaded"}
return
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
try:
await self._run_in_executor(service.load_model, model_type, True)
except Exception:
self.logger.exception("Failed to load model")
yield {
"type": "error",
"message": "Failed to load model. Check server logs for details.",
}
return
if opts["custom_prompt"]:
service.custom_prompt = opts["custom_prompt"]
yield {
"type": "progress",
"progress": 5,
"status": "Model loaded. Fetching all images from database...",
}
images = await self._run_in_executor(self.db.get_all_images)
total = len(images)
if total == 0:
yield {
"type": "complete",
"processed": 0,
"tagged": 0,
"skipped": 0,
"errors": 0,
"skipped_outside_roots": 0,
"status": "No images with linked prompts found",
}
service.unload_model()
return
yield {
"type": "progress",
"progress": 10,
"status": f"Found {total} images. Processing...",
}
counts = {
"processed": 0,
"tagged": 0,
"skipped": 0,
"errors": 0,
"skipped_outside_roots": 0,
}
tagged_prompt_ids = set()
last_update = _time.monotonic()
finished = False
def progress_event(i, message):
return {
"type": "progress",
"progress": 10 + int((i + 1) / total * 85),
"status": message,
"processed": counts["processed"],
"tagged": counts["tagged"],
"skipped": counts["skipped"],
"skipped_outside_roots": counts["skipped_outside_roots"],
}
try:
for i, image_data in enumerate(images):
image_path = image_data.get("image_path")
prompt_id = image_data.get("prompt_id")
skip_reason = None
allowed_path = None
if not image_path or not prompt_id:
skip_reason = "unlinked"
elif prompt_id in tagged_prompt_ids:
skip_reason = "prompt already processed"
else:
# Rows can point anywhere (restored DB, hand edits): only
# media files inside the served output roots are tagged.
allowed_path = self._resolve_allowed_image_path(str(image_path))
if allowed_path is None:
if _file_exists(image_path):
skip_reason = "outside allowed directories"
counts["skipped_outside_roots"] += 1
else:
skip_reason = "file missing"
elif opts["skip_tagged"] and self._real_tags(
image_data.get("prompt_tags", [])
):
tagged_prompt_ids.add(prompt_id)
skip_reason = "already tagged"
if skip_reason is not None:
counts["skipped"] += 1
now = _time.monotonic()
if (now - last_update) >= 0.5 or i == total - 1:
yield progress_event(
i, f"Skipping {i + 1}/{total} ({skip_reason})..."
)
await asyncio.sleep(0.01)
last_update = now
continue
try:
tags = await self._run_in_executor(
service.generate_tags,
allowed_path,
general_threshold=opts["general_threshold"],
character_threshold=opts["character_threshold"],
)
counts["processed"] += 1
if tags and await self._apply_tags_to_prompt(prompt_id, tags):
counts["tagged"] += 1
else:
counts["skipped"] += 1
tagged_prompt_ids.add(prompt_id)
except Exception as img_err:
self.logger.error(
f"Error processing {os.path.basename(str(image_path))}: "
f"{img_err}"
)
counts["errors"] += 1
counts["processed"] += 1
yield progress_event(i, f"Processing {i + 1}/{total}...")
await asyncio.sleep(0.01)
last_update = _time.monotonic()
finished = True
if not keep_in_memory:
service.unload_model()
yield {
"type": "complete",
"progress": 100,
**counts,
"status": "Complete",
"model_status": (
"Model kept in memory" if keep_in_memory else "Model unloaded"
),
"model_loaded": keep_in_memory,
}
finally:
if not finished and not keep_in_memory:
service.unload_model()
async def start_autotag(self, request):
"""Start batch auto-tagging with streaming progress (POST only)."""
params = await self._read_start_params(request)
if params is None:
return self._bad_request("Request body must be a JSON object")
opts = {
"model_type": params.get("model_type", "gguf"),
"custom_prompt": params.get("prompt", ""),
"skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true",
"keep_in_memory": str(params.get("keep_in_memory", "true")).lower()
== "true",
"general_threshold": params.get("general_threshold"),
"character_threshold": params.get("character_threshold"),
}
try:
for key in ("general_threshold", "character_threshold"):
if opts[key] is not None:
opts[key] = float(opts[key])
except (ValueError, TypeError):
return self._bad_request(
"general_threshold and character_threshold must be numeric"
)
async def events():
try:
from ..autotag import get_autotag_service
async for event in self._autotag_batch_events(
get_autotag_service(), opts
):
yield event
except Exception:
self.logger.exception("AutoTag error")
yield {
"type": "error",
"message": _INTERNAL_ERROR,
}
if not _autotag_batch_lock.acquire(blocking=False):
return web.json_response(
{"success": False, "error": _BUSY_ERROR}, status=409
)
try:
return await self._stream_sse(request, events())
finally:
_autotag_batch_lock.release()
def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
"""Return the real path when it is a media file inside an output directory."""
if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
return None
for output_dir in self._get_all_output_dirs():
if path_is_within(raw_path, output_dir):
real = os.path.realpath(raw_path)
return real if os.path.isfile(real) else None
return None
@staticmethod
def _forbidden(raw_path: str):
return web.json_response(
{
"success": False,
"error": (
f"{os.path.basename(str(raw_path).rstrip('/' + os.sep))} is not "
"an image inside the configured output directories"
),
},
status=403,
)
async def _select_autotag_target(self, data: dict):
"""Pick the image to tag from an autotag/single body.
Returns ``(image_path, prompt_id, None)`` on success or
``(None, None, error_response)`` when the request must be rejected.
"""
image_id = data.get("image_id")
if image_id is not None:
if isinstance(image_id, bool) or not isinstance(image_id, (int, str)):
return None, None, self._bad_request("image_id must be an integer")
try:
image_id = int(image_id)
except (TypeError, ValueError):
return None, None, self._bad_request("image_id must be an integer")
record = await self._run_in_executor(self.db.get_image_by_id, image_id)
if not record or not record.get("image_path"):
return (
None,
None,
web.json_response(
{"success": False, "error": f"Image {image_id} not found"},
status=404,
),
)
stored_path = record["image_path"]
resolved = self._resolve_allowed_image_path(stored_path)
if resolved is None:
return None, None, self._forbidden(stored_path)
return resolved, record.get("prompt_id"), None
raw_path = data.get("path") or data.get("image_path")
if not raw_path or not isinstance(raw_path, str):
return None, None, self._bad_request("image_id or path is required")
if "root_index" in data:
root_index = data["root_index"]
if isinstance(root_index, bool) or not isinstance(root_index, int):
return None, None, self._bad_request("root_index must be an integer")
roots = self._get_all_output_dirs()
if not 0 <= root_index < len(roots):
return None, None, self._forbidden(raw_path)
root = roots[root_index]
if Path(raw_path).is_absolute() or ".." in raw_path.replace(
"\\", "/"
).split("/"):
return None, None, self._forbidden(raw_path)
raw_path = str(Path(root) / raw_path)
if not path_is_within(raw_path, root):
return None, None, self._forbidden(raw_path)
resolved = self._resolve_allowed_image_path(raw_path)
if resolved is None:
return None, None, self._forbidden(raw_path)
prompt_id = None
try:
prompt_id = await self._run_in_executor(
self.db.get_prompt_id_for_image, raw_path
)
if prompt_id is None and resolved != raw_path:
prompt_id = await self._run_in_executor(
self.db.get_prompt_id_for_image, resolved
)
except Exception as e:
self.logger.warning(f"Could not find linked prompt: {e}")
return resolved, prompt_id, None
@staticmethod
def _bad_request(message: str):
return web.json_response({"success": False, "error": message}, status=400)
async def autotag_single(self, request):
"""Generate tags for a single image."""
"""Generate tags for a single image selected by image_id or gallery path."""
try:
data = await request.json()
image_path = data.get("image_path")
if not isinstance(data, dict):
return self._bad_request("Request body must be a JSON object")
model_type = data.get("model_type", "gguf")
custom_prompt = data.get("prompt")
use_gpu = data.get("use_gpu", True)
@@ -342,18 +538,13 @@ class AutotagRoutesMixin:
if character_threshold is not None:
character_threshold = float(character_threshold)
except (ValueError, TypeError):
return web.json_response(
{
"success": False,
"error": "general_threshold and character_threshold must be numeric",
},
status=400,
return self._bad_request(
"general_threshold and character_threshold must be numeric"
)
if not image_path:
return web.json_response(
{"success": False, "error": "image_path is required"}, status=400
)
image_path, prompt_id, error = await self._select_autotag_target(data)
if error is not None:
return error
from ..autotag import get_autotag_service
@@ -363,44 +554,34 @@ class AutotagRoutesMixin:
not service.is_model_loaded()
or service.get_loaded_model_type() != model_type
):
loop = asyncio.get_event_loop()
await loop.run_in_executor(
None, lambda: service.load_model(model_type, use_gpu)
)
await self._run_in_executor(service.load_model, model_type, use_gpu)
if custom_prompt:
service.custom_prompt = custom_prompt
loop = asyncio.get_event_loop()
tags = await loop.run_in_executor(
None,
lambda: service.generate_tags(
image_path,
general_threshold=general_threshold,
character_threshold=character_threshold,
),
tags = await self._run_in_executor(
service.generate_tags,
image_path,
general_threshold=general_threshold,
character_threshold=character_threshold,
)
prompt_id = None
try:
prompt_id = await self._run_in_executor(
self.db.get_prompt_id_for_image, image_path
)
except Exception as e:
self.logger.warning(f"Could not find linked prompt: {e}")
return web.json_response(
{
"success": True,
"tags": tags,
"prompt_id": prompt_id,
"image_path": image_path,
"filename": os.path.basename(image_path),
}
)
except json.JSONDecodeError:
return self._bad_request("Request body must be valid JSON")
except Exception as e:
self.logger.error(f"AutoTag single error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def apply_autotag(self, request):
"""Apply selected tags to a prompt."""
@@ -445,7 +626,9 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"Apply autotag error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def unload_autotag_model(self, request):
"""Manually unload the AutoTag model from memory."""
@@ -472,7 +655,37 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"Unload autotag model error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
@staticmethod
def _collect_scan_files(output_path: Path):
"""Media files under ``output_path`` from a bounded, symlink-free walk.
Descends at most ``_SCAN_MAX_DEPTH`` levels, skips ``thumbnails``
trees and symlinked files, and stops after ``_SCAN_MAX_FILES`` files.
Returns ``(paths, truncated)``.
"""
found = []
base_depth = len(output_path.parts)
for dirpath, dirnames, filenames in os.walk(output_path):
depth = len(Path(dirpath).parts) - base_depth
dirnames[:] = sorted(
name
for name in dirnames
if name != "thumbnails" and depth + 1 <= _SCAN_MAX_DEPTH
)
for name in sorted(filenames):
if Path(name).suffix.lower() not in _MEDIA_EXTENSIONS:
continue
candidate = Path(dirpath) / name
if candidate.is_symlink():
continue
if len(found) >= _SCAN_MAX_FILES:
return found, True
found.append(candidate)
return found, False
async def scan_output_dir(self, request):
"""Scan ComfyUI output directory for images."""
@@ -485,63 +698,49 @@ class AutotagRoutesMixin:
)
output_path = Path(output_dir)
image_extensions = [
".png",
".jpg",
".jpeg",
".gif",
".webp",
".bmp",
".tiff",
]
files, truncated = await self._run_in_executor(
self._collect_scan_files, output_path
)
has_thumbnails = (output_path / "thumbnails").is_dir()
images = []
seen_paths = set()
for ext in image_extensions:
for pattern in [f"*{ext.lower()}", f"*{ext.upper()}"]:
for image_path in output_path.rglob(pattern):
if "thumbnails" not in image_path.parts:
normalized_path = str(image_path).lower()
if normalized_path not in seen_paths:
seen_paths.add(normalized_path)
rel_path = image_path.relative_to(output_path)
thumbnail_url = None
thumbnails_dir = output_path / "thumbnails"
if thumbnails_dir.exists():
rel_path_no_ext = rel_path.with_suffix("")
thumbnail_rel_path = f"thumbnails/{rel_path_no_ext.as_posix()}_thumb{image_path.suffix}"
thumbnail_abs_path = (
thumbnails_dir
/ f"{rel_path_no_ext.as_posix()}_thumb{image_path.suffix}"
)
if thumbnail_abs_path.exists():
from urllib.parse import quote
thumbnail_url = f'/prompt_manager/images/serve/{quote(thumbnail_rel_path, safe="/")}'
from urllib.parse import quote as url_quote
images.append(
{
"filename": image_path.name,
"path": str(image_path),
"relative_path": str(rel_path),
"url": f'/prompt_manager/images/serve/{url_quote(rel_path.as_posix(), safe="/")}',
"thumbnail_url": thumbnail_url,
}
)
for image_path in files:
rel_path = image_path.relative_to(output_path)
thumbnail_url = None
if has_thumbnails:
thumb_rel = (
f"thumbnails/{rel_path.with_suffix('').as_posix()}"
f"_thumb{image_path.suffix}"
)
thumbnail_url = thumbnail_url_for(
thumb_rel, output_path / thumb_rel
)
images.append(
{
"filename": image_path.name,
"path": self._public_path(image_path),
"relative_path": rel_path.as_posix(),
"url": "/prompt_manager/images/serve/"
+ quote(rel_path.as_posix(), safe="/"),
"thumbnail_url": thumbnail_url,
}
)
images.sort(key=lambda x: x["filename"])
self.logger.info(f"Found {len(images)} images in output directory")
return web.json_response(
{"success": True, "images": images, "count": len(images)}
{
"success": True,
"images": images,
"count": len(images),
"truncated": truncated,
}
)
except Exception as e:
self.logger.error(f"Scan output dir error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
+1060 -773
View File
File diff suppressed because it is too large Load Diff
+123 -31
View File
@@ -4,6 +4,86 @@ import os
from aiohttp import web
# Keys PromptManagerLogger.update_config() understands. Anything else is
# rejected so a client cannot inject arbitrary entries into the logging config.
LOG_CONFIG_KEYS = frozenset(
{
"level",
"max_file_size",
"backup_count",
"console_logging",
"file_logging",
"buffer_size",
}
)
_LOG_INT_KEYS = frozenset({"max_file_size", "backup_count", "buffer_size"})
_LOG_BOOL_KEYS = frozenset({"console_logging", "file_logging"})
_LOG_LEVEL_NAMES = frozenset({"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"})
# Inclusive bounds for the integer keys that size the log files on disk
_MEGABYTE = 1024 * 1024
_LOG_INT_BOUNDS = {
"max_file_size": (1 * _MEGABYTE, 1024 * _MEGABYTE),
"backup_count": (0, 50),
}
_LOG_INT_BOUND_LABELS = {
"max_file_size": "between 1 MB and 1024 MB (in bytes)",
"backup_count": "between 0 and 50",
}
DEFAULT_LOG_LIMIT = 100
MAX_LOG_LIMIT = 500
def _parse_limit(raw):
"""Parse a ``limit`` query value, clamped to [1, MAX_LOG_LIMIT].
Raises:
ValueError: when ``raw`` is not an integer.
"""
if raw is None or raw == "":
return DEFAULT_LOG_LIMIT
limit = int(raw)
return max(1, min(limit, MAX_LOG_LIMIT))
def _validate_log_config(data):
"""Validate a log-config payload against LOG_CONFIG_KEYS.
Returns:
(clean, error): ``clean`` is the normalised dict to apply, ``error``
a message describing the first problem (``clean`` is then ``None``).
"""
if not isinstance(data, dict):
return None, "Request body must be a JSON object"
unknown = sorted(key for key in data if key not in LOG_CONFIG_KEYS)
if unknown:
return None, (
f"Unknown logging config keys: {', '.join(unknown)}. "
f"Allowed: {', '.join(sorted(LOG_CONFIG_KEYS))}"
)
clean = {}
for key, value in data.items():
if key == "level":
if not isinstance(value, str) or value.upper() not in _LOG_LEVEL_NAMES:
return None, (
f"Invalid log level. Must be one of: {sorted(_LOG_LEVEL_NAMES)}"
)
clean[key] = value.upper()
elif key in _LOG_INT_KEYS:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
return None, f"{key} must be a non-negative integer"
bounds = _LOG_INT_BOUNDS.get(key)
if bounds and not bounds[0] <= value <= bounds[1]:
return None, f"{key} must be {_LOG_INT_BOUND_LABELS[key]}"
clean[key] = value
elif key in _LOG_BOOL_KEYS:
if not isinstance(value, bool):
return None, f"{key} must be a boolean"
clean[key] = value
return clean, None
class LoggingRoutesMixin:
"""Mixin providing logging-related API endpoints."""
@@ -56,14 +136,15 @@ class LoggingRoutesMixin:
try:
logger_manager = self._get_logger_manager()
limit = int(request.query.get("limit", 100))
try:
limit = _parse_limit(request.query.get("limit"))
except ValueError:
return web.json_response(
{"success": False, "error": "limit must be an integer", "logs": []},
status=400,
)
level = request.query.get("level", None)
if limit > 1000:
limit = 1000
elif limit < 1:
limit = 1
logs = logger_manager.get_recent_logs(limit=limit, level=level)
return web.json_response(
@@ -79,14 +160,18 @@ class LoggingRoutesMixin:
except Exception as e:
self.logger.error(f"Get logs error: {e}")
return web.json_response(
{"success": False, "error": str(e), "logs": []}, status=500
{"success": False, "error": self._public_error(e), "logs": []},
status=500,
)
async def get_log_files(self, request):
"""Get information about available log files."""
try:
logger_manager = self._get_logger_manager()
log_files = logger_manager.get_log_files()
log_files = [
{**entry, "path": self._public_path(entry.get("path"))}
for entry in logger_manager.get_log_files()
]
return web.json_response(
{"success": True, "files": log_files, "count": len(log_files)}
@@ -95,7 +180,8 @@ class LoggingRoutesMixin:
except Exception as e:
self.logger.error(f"Get log files error: {e}")
return web.json_response(
{"success": False, "error": str(e), "files": []}, status=500
{"success": False, "error": self._public_error(e), "files": []},
status=500,
)
async def download_log_file(self, request):
@@ -130,7 +216,9 @@ class LoggingRoutesMixin:
except Exception as e:
self.logger.error(f"Download log file error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def truncate_logs(self, request):
"""Truncate all log files."""
@@ -148,7 +236,9 @@ class LoggingRoutesMixin:
except Exception as e:
self.logger.error(f"Truncate logs error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def get_log_config(self, request):
"""Get current logging configuration."""
@@ -160,27 +250,23 @@ class LoggingRoutesMixin:
except Exception as e:
self.logger.error(f"Get log config error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def update_log_config(self, request):
"""Update logging configuration."""
"""Update logging configuration (whitelisted keys only)."""
try:
data = await request.json()
data, error_response = await self._read_json_body(request)
if error_response is not None:
return error_response
clean, error = _validate_log_config(data)
if error:
return web.json_response({"success": False, "error": error}, status=400)
logger_manager = self._get_logger_manager()
if "level" in data:
valid_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
if data["level"].upper() not in valid_levels:
return web.json_response(
{
"success": False,
"error": f"Invalid log level. Must be one of: {valid_levels}",
},
status=400,
)
data["level"] = data["level"].upper()
logger_manager.update_config(data)
logger_manager.update_config(clean)
return web.json_response(
{
@@ -192,16 +278,22 @@ class LoggingRoutesMixin:
except Exception as e:
self.logger.error(f"Update log config error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": "Failed to update logging configuration"},
status=500,
)
async def get_log_stats(self, request):
"""Get logging statistics."""
try:
logger_manager = self._get_logger_manager()
stats = logger_manager.get_log_stats()
stats = dict(logger_manager.get_log_stats())
stats["log_directory"] = self._public_path(stats.get("log_directory"))
return web.json_response({"success": True, "stats": stats})
except Exception as e:
self.logger.error(f"Get log stats error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
+259 -222
View File
@@ -1,11 +1,15 @@
"""LoraManager integration API routes for PromptManager."""
import json
import os
from pathlib import Path
from aiohttp import web
_LORA_PATH_ERROR = (
"LoraManager path must be a directory under ComfyUI's custom_nodes "
"whose name contains 'lora'"
)
class LoraIntegrationMixin:
"""Mixin providing LoraManager detection, scanning, and trigger word endpoints."""
@@ -35,6 +39,15 @@ class LoraIntegrationMixin:
async def lora_refresh_cache_route(request):
return await self.lora_refresh_cache(request)
def _error_response(self, exc: Exception, status: int = 500):
return web.json_response(
{"success": False, "error": self._public_error(exc)}, status=status
)
@staticmethod
def _bad_lora_request(message: str):
return web.json_response({"success": False, "error": message}, status=400)
# ── Detection ────────────────────────────────────────────────────
async def lora_detect(self, request):
@@ -47,12 +60,12 @@ class LoraIntegrationMixin:
{
"success": True,
"detected": path is not None,
"path": path or "",
"path": self._public_path(path) if path else "",
}
)
except Exception as e:
self.logger.error(f"LoraManager detection failed: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return self._error_response(e)
# ── Status ───────────────────────────────────────────────────────
@@ -74,20 +87,49 @@ class LoraIntegrationMixin:
{
"success": True,
"enabled": config["enabled"],
"path": config["path"],
"path": self._public_path(config["path"]),
"trigger_words_enabled": config["trigger_words_enabled"],
"civitai_api_key": config.get("civitai_api_key", ""),
"has_civitai_api_key": bool(config.get("civitai_api_key")),
"detected": detected_path is not None,
"detected_path": detected_path or "",
"detected_path": (
self._public_path(detected_path) if detected_path else ""
),
"trigger_cache_loaded": cache.is_loaded,
}
)
except Exception as e:
self.logger.error(f"LoraManager status check failed: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return self._error_response(e)
# ── Enable / Disable ─────────────────────────────────────────────
async def _validated_lora_manager_path(self, path):
"""Canonical LoraManager path for ``path`` or ``(None, error_response)``.
The user's path (or, when empty, the auto-detected install) must be a
LoraManager-looking directory that is a direct child of ComfyUI's
``custom_nodes`` with ``lora`` in its name; the trees under it are
later served as images, so nothing outside custom_nodes may be used.
"""
from ..lora_utils import detect_lora_manager, resolve_lora_manager_path
candidate = ""
if path:
candidate = await self._run_in_executor(resolve_lora_manager_path, path)
if candidate is None:
return None, self._bad_lora_request(_LORA_PATH_ERROR)
resolved = await self._run_in_executor(detect_lora_manager, candidate)
if not resolved:
return None, self._bad_lora_request(
"LoraManager not found at the specified path"
)
allowed = await self._run_in_executor(resolve_lora_manager_path, resolved)
if allowed is None:
return None, self._bad_lora_request(_LORA_PATH_ERROR)
return allowed, None
async def lora_enable(self, request):
"""Enable or disable LoraManager integration and save to config.json."""
try:
@@ -95,23 +137,29 @@ class LoraIntegrationMixin:
enabled = data.get("enabled", False)
path = data.get("path", "")
trigger_words = data.get("trigger_words_enabled", False)
civitai_key = data.get("civitai_api_key", "")
from ..config import IntegrationConfig, PromptManagerConfig
from ..lora_utils import detect_lora_manager, get_trigger_cache
from ..lora_utils import get_trigger_cache
# The key is never echoed back to the UI, so a blank or absent
# field means "unchanged"; clearing is an explicit request.
civitai_key = IntegrationConfig.CIVITAI_API_KEY
submitted_key = data.get("civitai_api_key")
if submitted_key is not None and not isinstance(submitted_key, str):
return self._bad_lora_request("civitai_api_key must be a string")
if data.get("clear_civitai_api_key"):
civitai_key = ""
elif submitted_key and submitted_key.strip():
civitai_key = submitted_key.strip()
if not isinstance(path, str):
return self._bad_lora_request("path must be a string")
# If enabling, validate the path
if enabled:
resolved = await self._run_in_executor(detect_lora_manager, path)
if not resolved:
return web.json_response(
{
"success": False,
"error": "LoraManager not found at the specified path",
},
status=400,
)
path = resolved
path, error = await self._validated_lora_manager_path(path)
if error is not None:
return error
# Update in-memory config
IntegrationConfig.LORA_MANAGER_ENABLED = enabled
@@ -119,12 +167,8 @@ class LoraIntegrationMixin:
IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED = trigger_words
IntegrationConfig.CIVITAI_API_KEY = civitai_key
# Persist to config.json
config_dir = os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)
config_file = os.path.join(config_dir, "config.json")
PromptManagerConfig.save_to_file(config_file)
# Persist to the configured config.json
PromptManagerConfig.save_to_file()
# Load trigger word cache if enabling
cache = get_trigger_cache()
@@ -138,49 +182,36 @@ class LoraIntegrationMixin:
{
"success": True,
"enabled": enabled,
"path": path,
"path": self._public_path(path),
"trigger_words_enabled": trigger_words,
}
)
except Exception as e:
self.logger.error(f"LoraManager enable/disable failed: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return self._error_response(e)
# ── Scan LoRA example images ─────────────────────────────────────
async def lora_scan(self, request):
"""Scan LoraManager metadata and import LoRA info + preview images.
Streams progress as SSE, matching the existing scan pattern.
Streams progress as SSE, matching the existing scan pattern. Errors
raised once the stream is open are reported as an ``error`` event.
"""
try:
from ..config import IntegrationConfig
from ..lora_utils import (
download_civitai_images,
find_lora_directories,
get_example_prompt_from_metadata,
get_lora_image_cache_dir,
get_preview_images_from_metadata,
get_trigger_words_from_metadata,
get_model_name_from_metadata,
read_lora_metadata,
)
from ..lora_utils import resolve_lora_manager_path
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return web.json_response(
{
"success": False,
"error": "LoraManager integration is not enabled",
},
status=400,
)
return self._bad_lora_request("LoraManager integration is not enabled")
lm_path = IntegrationConfig.LORA_MANAGER_PATH
if not lm_path:
return web.json_response(
{"success": False, "error": "LoraManager path not configured"},
status=400,
)
return self._bad_lora_request("LoraManager path not configured")
allowed = await self._run_in_executor(resolve_lora_manager_path, lm_path)
if allowed is None:
return self._bad_lora_request(_LORA_PATH_ERROR)
response = web.StreamResponse(
status=200,
@@ -193,167 +224,185 @@ class LoraIntegrationMixin:
line = f"data: {json.dumps(data)}\n\n"
await response.write(line.encode("utf-8"))
await send_progress(
{
"type": "progress",
"status": "Clearing previous lora-manager imports...",
"progress": 0,
}
)
# Clear previous imports so reimport is always clean
await self._run_in_executor(
self.db.delete_prompts_by_category, "lora-manager"
)
await send_progress(
{
"type": "progress",
"status": "Finding LoRA directories...",
"progress": 2,
}
)
lora_dirs = await self._run_in_executor(find_lora_directories, lm_path)
# Collect all metadata files
meta_files = []
for d in lora_dirs:
dir_path = Path(d)
meta_files.extend(dir_path.rglob("*.metadata.json"))
total = len(meta_files)
imported = 0
skipped = 0
await send_progress(
{
"type": "progress",
"status": f"Found {total} LoRA metadata files",
"progress": 5,
"total": total,
}
)
cache_dir = get_lora_image_cache_dir()
for i, meta_file in enumerate(meta_files):
metadata = await self._run_in_executor(read_lora_metadata, meta_file)
if not metadata:
skipped += 1
continue
model_name = get_model_name_from_metadata(metadata)
trigger_words = get_trigger_words_from_metadata(metadata)
# Collect all images: local previews + downloaded civitai examples
preview_paths = await self._run_in_executor(
get_preview_images_from_metadata, metadata, meta_file
)
civitai_paths = await self._run_in_executor(
download_civitai_images,
metadata,
meta_file,
cache_dir,
IntegrationConfig.CIVITAI_API_KEY,
)
# Merge, local first, dedup
seen = set(preview_paths)
all_images = list(preview_paths)
for cp in civitai_paths:
if cp not in seen:
all_images.append(cp)
seen.add(cp)
# Build prompt text: prefer example prompt, then model name
example_prompt = get_example_prompt_from_metadata(metadata)
prompt_text = example_prompt or model_name
# Build tags
tags = ["lora-manager", f"lora:{model_name}"]
tags.extend(trigger_words)
# Save to database via existing mechanism
try:
import hashlib
prompt_hash = hashlib.sha256(
prompt_text.strip().lower().encode("utf-8")
).hexdigest()
existing = await self._run_in_executor(
self.db.get_prompt_by_hash, prompt_hash
)
if existing:
# Link all images
for pp in all_images:
await self._run_in_executor(
self.db.link_image_to_prompt,
existing["id"],
pp,
)
skipped += 1
else:
prompt_id = await self._run_in_executor(
self.db.save_prompt,
prompt_text,
"lora-manager", # category
tags,
None, # rating
None, # notes
prompt_hash,
)
if prompt_id:
for pp in all_images:
await self._run_in_executor(
self.db.link_image_to_prompt,
prompt_id,
pp,
)
imported += 1
else:
skipped += 1
except Exception as e:
self.logger.warning(f"Failed to import LoRA {model_name}: {e}")
skipped += 1
# Progress update for every LoRA
progress = int(5 + (90 * (i + 1) / max(total, 1)))
img_count = len(all_images)
status = f"{model_name}"
if img_count:
status += f" ({img_count} images)"
await send_progress(
{
"type": "progress",
"status": status,
"progress": progress,
"processed": i + 1,
"imported": imported,
"skipped": skipped,
}
)
await send_progress(
{
"type": "complete",
"progress": 100,
"total": total,
"imported": imported,
"skipped": skipped,
}
)
try:
await self._lora_scan_stream(allowed, send_progress)
except Exception as e:
self.logger.exception("LoRA scan failed")
await send_progress({"type": "error", "message": self._public_error(e)})
await response.write_eof()
return response
except Exception as e:
self.logger.error(f"LoRA scan failed: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return self._error_response(e)
async def _lora_scan_stream(self, lm_path, send_progress):
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``."""
from ..config import IntegrationConfig
from ..lora_utils import (
download_civitai_images,
find_lora_directories,
get_example_prompt_from_metadata,
get_lora_image_cache_dir,
get_preview_images_from_metadata,
get_trigger_words_from_metadata,
get_model_name_from_metadata,
read_lora_metadata,
)
await send_progress(
{
"type": "progress",
"status": "Clearing previous lora-manager imports...",
"progress": 0,
}
)
# Clear previous imports so reimport is always clean
await self._run_in_executor(self.db.delete_prompts_by_category, "lora-manager")
await send_progress(
{
"type": "progress",
"status": "Finding LoRA directories...",
"progress": 2,
}
)
lora_dirs = await self._run_in_executor(find_lora_directories, lm_path)
# Collect all metadata files
meta_files = []
for d in lora_dirs:
dir_path = Path(d)
meta_files.extend(dir_path.rglob("*.metadata.json"))
total = len(meta_files)
imported = 0
skipped = 0
await send_progress(
{
"type": "progress",
"status": f"Found {total} LoRA metadata files",
"progress": 5,
"total": total,
}
)
cache_dir = get_lora_image_cache_dir()
for i, meta_file in enumerate(meta_files):
metadata = await self._run_in_executor(read_lora_metadata, meta_file)
if not metadata:
skipped += 1
continue
model_name = get_model_name_from_metadata(metadata)
trigger_words = get_trigger_words_from_metadata(metadata)
# Collect all images: local previews + downloaded civitai examples
preview_paths = await self._run_in_executor(
get_preview_images_from_metadata, metadata, meta_file
)
civitai_paths = await self._run_in_executor(
download_civitai_images,
metadata,
meta_file,
cache_dir,
IntegrationConfig.CIVITAI_API_KEY,
)
# Merge, local first, dedup
seen = set(preview_paths)
all_images = list(preview_paths)
for cp in civitai_paths:
if cp not in seen:
all_images.append(cp)
seen.add(cp)
# Build prompt text: prefer example prompt, then model name
example_prompt = get_example_prompt_from_metadata(metadata)
prompt_text = example_prompt or model_name
# Build tags
tags = ["lora-manager", f"lora:{model_name}"]
tags.extend(trigger_words)
# Save to database via existing mechanism
try:
import hashlib
prompt_hash = hashlib.sha256(
prompt_text.strip().lower().encode("utf-8")
).hexdigest()
existing = await self._run_in_executor(
self.db.get_prompt_by_hash, prompt_hash
)
if existing:
# Link all images
for pp in all_images:
await self._run_in_executor(
self.db.link_image_to_prompt,
existing["id"],
pp,
)
skipped += 1
else:
prompt_id = await self._run_in_executor(
self.db.save_prompt,
prompt_text,
"lora-manager", # category
tags,
None, # rating
None, # notes
prompt_hash,
)
if prompt_id:
for pp in all_images:
await self._run_in_executor(
self.db.link_image_to_prompt,
prompt_id,
pp,
)
imported += 1
else:
skipped += 1
except Exception as e:
self.logger.warning(f"Failed to import LoRA {model_name}: {e}")
skipped += 1
# Progress update for every LoRA
progress = int(5 + (90 * (i + 1) / max(total, 1)))
img_count = len(all_images)
status = f"{model_name}"
if img_count:
status += f" ({img_count} images)"
await send_progress(
{
"type": "progress",
"status": status,
"progress": progress,
"processed": i + 1,
"imported": imported,
"skipped": skipped,
}
)
await send_progress(
{
"type": "complete",
"progress": 100,
"total": total,
"imported": imported,
"skipped": skipped,
}
)
# ── Trigger word endpoints ───────────────────────────────────────
@@ -364,17 +413,11 @@ class LoraIntegrationMixin:
from ..lora_utils import get_trigger_cache
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return web.json_response(
{"success": False, "error": "LoraManager integration not enabled"},
status=400,
)
return self._bad_lora_request("LoraManager integration not enabled")
lora_name = request.query.get("name", "")
if not lora_name:
return web.json_response(
{"success": False, "error": "Missing 'name' query parameter"},
status=400,
)
return self._bad_lora_request("Missing 'name' query parameter")
cache = get_trigger_cache()
if not cache.is_loaded:
@@ -388,7 +431,7 @@ class LoraIntegrationMixin:
)
except Exception as e:
self.logger.error(f"Trigger word lookup failed: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return self._error_response(e)
async def lora_refresh_cache(self, request):
"""Force-refresh the trigger word cache from disk."""
@@ -397,17 +440,11 @@ class LoraIntegrationMixin:
from ..lora_utils import get_trigger_cache
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return web.json_response(
{"success": False, "error": "LoraManager integration not enabled"},
status=400,
)
return self._bad_lora_request("LoraManager integration not enabled")
lm_path = IntegrationConfig.LORA_MANAGER_PATH
if not lm_path:
return web.json_response(
{"success": False, "error": "LoraManager path not configured"},
status=400,
)
return self._bad_lora_request("LoraManager path not configured")
cache = get_trigger_cache()
count = await self._run_in_executor(cache.load, lm_path)
@@ -416,4 +453,4 @@ class LoraIntegrationMixin:
)
except Exception as e:
self.logger.error(f"Trigger cache refresh failed: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return self._error_response(e)
+323 -91
View File
@@ -2,8 +2,93 @@
import datetime
import json
import os
import sqlite3
from aiohttp import web
from urllib.parse import quote
# Upper bounds on page sizes and offsets for every list endpoint. The offset
# ceiling keeps values inside SQLite's 64-bit range and bounds skip work.
MAX_PAGE_LIMIT = 500
MAX_PAGE_OFFSET = 10_000_000
# The export walks the whole table in pages of this size instead of asking
# for one giant result set (and silently truncating large libraries).
EXPORT_PAGE_SIZE = 1000
def parse_page_params(query, default_limit=50, max_limit=MAX_PAGE_LIMIT):
"""Return ``(limit, offset)`` from a query mapping, clamped to safe bounds.
``limit`` is clamped to ``[1, max_limit]`` and ``offset`` to
``[0, MAX_PAGE_OFFSET]``. Non-integer values raise ``ValueError`` so the
caller can answer 400.
"""
limit = int(query.get("limit", default_limit))
offset = int(query.get("offset", 0))
return max(1, min(limit, max_limit)), max(0, min(offset, MAX_PAGE_OFFSET))
def bad_request(message):
return web.json_response({"success": False, "error": message}, status=400)
def tags_error(tags):
"""A 400 response when *tags* fail ``validate_tags``, else None."""
try:
validate_tags(tags)
except ValueError as exc:
return bad_request(str(exc))
return None
def safe_error_message(exc):
"""Describe *exc* for a client without leaking absolute server paths.
OSError carries the offending path in ``filename``; only its basename
is echoed back together with ``strerror``.
"""
if isinstance(exc, OSError):
reason = exc.strerror or exc.__class__.__name__
if exc.filename:
return f"{reason}: {os.path.basename(str(exc.filename))}"
return reason
return str(exc) or exc.__class__.__name__
def thumbnail_url_for(thumb_rel, thumb_abs):
"""Serve URL for an existing thumbnail, or None when the file is missing.
The URL carries the file's mtime as ``?v=``: thumbnails are served with a
long cache lifetime, so a regenerated file must get a new URL or browsers
keep showing whatever they cached for the old one.
"""
try:
version = int(thumb_abs.stat().st_mtime)
except OSError:
return None
return f"/prompt_manager/images/serve/{quote(thumb_rel, safe='/')}?v={version}"
def publish_image_paths(images, public_path):
"""Image dicts for a response: ``image_path`` in its public (ComfyUI
relative) form and the absolute ``file_path`` duplicate removed.
``url``/``thumbnail_url`` added by enrichment are left untouched, so the
frontend keeps working without ever seeing the server's directory layout.
"""
published = []
for image in images:
if not isinstance(image, dict):
published.append(image)
continue
image = {k: v for k, v in image.items() if k != "file_path"}
if "image_path" in image:
image["image_path"] = public_path(image["image_path"])
published.append(image)
return published
try:
from ...utils.validators import (
@@ -50,6 +135,10 @@ class PromptRoutesMixin:
async def get_tags_filter_route(request):
return await self.get_tags_filter(request)
@routes.get("/prompt_manager/tags/suggest")
async def suggest_tags_route(request):
return await self.suggest_tags(request)
# Bulk tag operations (register BEFORE {tag_name} to avoid path param match)
@routes.post("/prompt_manager/tags/merge")
async def merge_tags_route(request):
@@ -122,6 +211,15 @@ class PromptRoutesMixin:
async def export_prompts_route(request):
return await self.export_prompts(request)
def _present_prompts(self, prompts):
"""Prompts ready for a response: image urls added, server paths hidden."""
self._enrich_prompt_images(prompts)
for prompt in prompts:
prompt["images"] = publish_image_paths(
prompt.get("images", []), self._public_path
)
return prompts
async def search_prompts(self, request):
"""Search for prompts using multiple filter criteria."""
try:
@@ -129,9 +227,14 @@ class PromptRoutesMixin:
category = request.query.get("category", "").strip()
tags_str = request.query.get("tags", "").strip()
min_rating = request.query.get("min_rating", 0)
limit = int(request.query.get("limit", 50))
try:
limit, offset = parse_page_params(request.query)
except ValueError:
return bad_request("limit and offset must be integers")
folder = request.query.get("folder", "").strip() or None
# Validated against a whitelist in the database layer
sort = request.query.get("sort") or None
tags = None
if tags_str:
@@ -149,18 +252,33 @@ class PromptRoutesMixin:
tags=tags,
rating_min=min_rating,
limit=limit,
offset=offset,
folder=folder,
sort=sort,
)
self._enrich_prompt_images(results)
self._present_prompts(results)
return web.json_response(
{"success": True, "results": results, "count": len(results)}
{
"success": True,
"results": results,
"count": len(results),
"pagination": {
"limit": limit,
"offset": offset,
"count": len(results),
},
}
)
except Exception as e:
self.logger.error(f"Search error: {e}", exc_info=True)
return web.json_response(
{"success": False, "error": f"Search failed: {str(e)}", "results": []},
{
"success": False,
"error": f"Search failed: {safe_error_message(e)}",
"results": [],
},
status=500,
)
@@ -179,27 +297,29 @@ class PromptRoutesMixin:
return web.json_response({"success": True, "subfolders": subfolders})
except Exception as e:
self.logger.error(f"Subfolders error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def get_recent_prompts(self, request):
"""Retrieve recently created prompts with pagination support."""
"""Retrieve prompts with pagination and an optional sort (default newest
first)."""
try:
limit = int(request.query.get("limit", 50))
page = int(request.query.get("page", 1))
offset = int(request.query.get("offset", 0))
try:
limit, offset = parse_page_params(request.query)
page = int(request.query.get("page", 1))
except ValueError:
return bad_request("limit, offset and page must be integers")
if page > 1 and offset == 0:
offset = (page - 1) * limit
if limit > 1000:
limit = 1000
elif limit < 1:
limit = 1
offset = min((page - 1) * limit, MAX_PAGE_OFFSET)
# Validated against a whitelist in the database layer
sort = request.query.get("sort") or None
results = await self._run_in_executor(
self.db.get_recent_prompts, limit=limit, offset=offset
self.db.get_recent_prompts, limit=limit, offset=offset, sort=sort
)
self._enrich_prompt_images(results["prompts"])
self._present_prompts(results["prompts"])
return web.json_response(
{
@@ -222,7 +342,7 @@ class PromptRoutesMixin:
return web.json_response(
{
"success": False,
"error": f"Failed to get recent prompts: {str(e)}",
"error": f"Failed to get recent prompts: {safe_error_message(e)}",
"results": [],
"pagination": {"total": 0, "page": 1, "total_pages": 0},
},
@@ -239,7 +359,7 @@ class PromptRoutesMixin:
return web.json_response(
{
"success": False,
"error": f"Failed to get categories: {str(e)}",
"error": f"Failed to get categories: {safe_error_message(e)}",
"categories": [],
},
status=500,
@@ -255,7 +375,7 @@ class PromptRoutesMixin:
return web.json_response(
{
"success": False,
"error": f"Failed to get tags: {str(e)}",
"error": f"Failed to get tags: {safe_error_message(e)}",
"tags": [],
},
status=500,
@@ -265,13 +385,9 @@ class PromptRoutesMixin:
"""Get tags with usage counts, search, sort, and pagination."""
try:
try:
limit = int(request.query.get("limit", 50))
offset = int(request.query.get("offset", 0))
except (ValueError, TypeError):
return web.json_response(
{"success": False, "error": "Invalid limit or offset parameter"},
status=400,
)
limit, offset = parse_page_params(request.query)
except ValueError:
return bad_request("limit and offset must be integers")
search = request.query.get("search", "").strip() or None
sort = request.query.get("sort", "alpha_asc")
@@ -297,32 +413,30 @@ class PromptRoutesMixin:
)
except Exception as e:
self.logger.error(f"Tags stats error: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def get_tag_prompts(self, request):
"""Get prompts for a single tag."""
try:
from urllib.parse import unquote
tag_name = unquote(request.match_info.get("tag_name", ""))
# aiohttp has already percent-decoded match_info; decoding again
# would turn a tag literally named "%41" into "A".
tag_name = request.match_info.get("tag_name", "")
if not tag_name:
return web.json_response(
{"success": False, "error": "Tag name required"}, status=400
)
try:
limit = int(request.query.get("limit", 20))
offset = int(request.query.get("offset", 0))
except (ValueError, TypeError):
return web.json_response(
{"success": False, "error": "Invalid limit or offset parameter"},
status=400,
)
limit, offset = parse_page_params(request.query, default_limit=20)
except ValueError:
return bad_request("limit and offset must be integers")
result = await self._run_in_executor(
self.db.get_prompts_by_tags, [tag_name], "and", limit, offset
)
self._enrich_prompt_images(result["prompts"])
self._present_prompts(result["prompts"])
return web.json_response(
{
@@ -339,7 +453,46 @@ class PromptRoutesMixin:
)
except Exception as e:
self.logger.error(f"Tag prompts error: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def suggest_tags(self, request):
"""Autocomplete for the tag search box.
``q`` is the prefix being typed, ``with`` the comma-separated tags
already entered: only tags that co-occur with all of them are offered.
"""
try:
prefix = request.query.get("q", "").strip()[:100]
context = []
seen = set()
for raw in request.query.get("with", "").split(","):
name = raw.strip()
if name and name.lower() not in seen:
seen.add(name.lower())
context.append(name)
context = context[:10]
try:
limit = int(request.query.get("limit", 15))
except ValueError:
limit = 15
limit = max(1, min(limit, 50))
suggestions = await self._run_in_executor(
self.db.suggest_tags, prefix, context, limit
)
return web.json_response({"success": True, "suggestions": suggestions})
except Exception as e:
self.logger.error(f"Tag suggest error: {e}")
return web.json_response(
{
"success": False,
"error": "Failed to suggest tags",
"suggestions": [],
},
status=500,
)
async def get_tags_filter(self, request):
"""Get prompts matching multiple tags with AND/OR mode, or untagged prompts."""
@@ -348,20 +501,13 @@ class PromptRoutesMixin:
if untagged:
try:
limit = int(request.query.get("limit", 20))
offset = int(request.query.get("offset", 0))
except (ValueError, TypeError):
return web.json_response(
{
"success": False,
"error": "Invalid limit or offset parameter",
},
status=400,
)
limit, offset = parse_page_params(request.query, default_limit=20)
except ValueError:
return bad_request("limit and offset must be integers")
result = await self._run_in_executor(
self.db.get_untagged_prompts, limit, offset
)
self._enrich_prompt_images(result["prompts"])
self._present_prompts(result["prompts"])
return web.json_response(
{
"success": True,
@@ -389,18 +535,14 @@ class PromptRoutesMixin:
mode = "and"
try:
limit = int(request.query.get("limit", 20))
offset = int(request.query.get("offset", 0))
except (ValueError, TypeError):
return web.json_response(
{"success": False, "error": "Invalid limit or offset parameter"},
status=400,
)
limit, offset = parse_page_params(request.query, default_limit=20)
except ValueError:
return bad_request("limit and offset must be integers")
result = await self._run_in_executor(
self.db.get_prompts_by_tags, tags_list, mode, limit, offset
)
self._enrich_prompt_images(result["prompts"])
self._present_prompts(result["prompts"])
return web.json_response(
{
@@ -418,14 +560,16 @@ class PromptRoutesMixin:
)
except Exception as e:
self.logger.error(f"Tags filter error: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def rename_tag_endpoint(self, request):
"""Rename a tag across all prompts."""
try:
from urllib.parse import unquote
tag_name = unquote(request.match_info.get("tag_name", ""))
# aiohttp has already percent-decoded match_info; decoding again
# would turn a tag literally named "%41" into "A".
tag_name = request.match_info.get("tag_name", "")
if not tag_name:
return web.json_response(
{"success": False, "error": "Tag name required"}, status=400
@@ -437,11 +581,15 @@ class PromptRoutesMixin:
return web.json_response(
{"success": False, "error": "Invalid JSON body"}, status=400
)
new_name = (body.get("new_name") or "").strip()
if not new_name:
raw_name = body.get("new_name")
if not raw_name:
return web.json_response(
{"success": False, "error": "New tag name required"}, status=400
)
error = tags_error([raw_name])
if error is not None:
return error
new_name = raw_name.strip()
result = await self._run_in_executor(
self.db.rename_tag_all_prompts, tag_name, new_name
@@ -455,19 +603,22 @@ class PromptRoutesMixin:
if result.get("skipped_count", 0) > 0:
resp["skipped_count"] = result["skipped_count"]
resp["warning"] = (
f"{result['skipped_count']} prompt(s) had corrupted tag data and were skipped"
f"{result['skipped_count']} prompt(s) had corrupted tag data "
"and were skipped"
)
return web.json_response(resp)
except Exception as e:
self.logger.error(f"Rename tag error: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def delete_tag_endpoint(self, request):
"""Delete a tag from all prompts."""
try:
from urllib.parse import unquote
tag_name = unquote(request.match_info.get("tag_name", ""))
# aiohttp has already percent-decoded match_info; decoding again
# would turn a tag literally named "%41" into "A".
tag_name = request.match_info.get("tag_name", "")
if not tag_name:
return web.json_response(
{"success": False, "error": "Tag name required"}, status=400
@@ -484,12 +635,15 @@ class PromptRoutesMixin:
if result.get("skipped_count", 0) > 0:
resp["skipped_count"] = result["skipped_count"]
resp["warning"] = (
f"{result['skipped_count']} prompt(s) had corrupted tag data and were skipped"
f"{result['skipped_count']} prompt(s) had corrupted tag data "
"and were skipped"
)
return web.json_response(resp)
except Exception as e:
self.logger.error(f"Delete tag error: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def merge_tags_endpoint(self, request):
"""Merge source tags into a target tag."""
@@ -501,16 +655,22 @@ class PromptRoutesMixin:
{"success": False, "error": "Invalid JSON body"}, status=400
)
source_tags = body.get("source_tags", [])
target_tag = (body.get("target_tag") or "").strip()
raw_target = body.get("target_tag")
if not source_tags:
return web.json_response(
{"success": False, "error": "Source tags required"}, status=400
)
if not target_tag:
if not raw_target:
return web.json_response(
{"success": False, "error": "Target tag required"}, status=400
)
if not isinstance(source_tags, list):
return bad_request("Source tags must be a list")
error = tags_error([*source_tags, raw_target])
if error is not None:
return error
target_tag = raw_target.strip()
result = await self._run_in_executor(
self.db.merge_tags, source_tags, target_tag
@@ -524,12 +684,15 @@ class PromptRoutesMixin:
if result.get("skipped_count", 0) > 0:
resp["skipped_count"] = result["skipped_count"]
resp["warning"] = (
f"{result['skipped_count']} prompt(s) had corrupted tag data and were skipped"
f"{result['skipped_count']} prompt(s) had corrupted tag data "
"and were skipped"
)
return web.json_response(resp)
except Exception as e:
self.logger.error(f"Merge tags error: {e}", exc_info=True)
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": safe_error_message(e)}, status=500
)
async def save_prompt(self, request):
"""Save a new prompt with metadata and duplicate detection."""
@@ -604,7 +767,10 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Save error: {e}", exc_info=True)
return web.json_response(
{"success": False, "error": f"Failed to save prompt: {str(e)}"},
{
"success": False,
"error": f"Failed to save prompt: {safe_error_message(e)}",
},
status=500,
)
@@ -634,7 +800,10 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Delete error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to delete prompt: {str(e)}"},
{
"success": False,
"error": f"Failed to delete prompt: {safe_error_message(e)}",
},
status=500,
)
@@ -671,6 +840,11 @@ class PromptRoutesMixin:
{"success": False, "error": "Prompt not found"}, status=404
)
except sqlite3.IntegrityError:
return web.json_response(
{"success": False, "error": "Another prompt already has this text"},
status=409,
)
except ValueError:
return web.json_response(
{"success": False, "error": "Invalid prompt ID"}, status=400
@@ -678,7 +852,10 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Update prompt error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to update prompt: {str(e)}"},
{
"success": False,
"error": f"Failed to update prompt: {safe_error_message(e)}",
},
status=500,
)
@@ -715,7 +892,10 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Update rating error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to update rating: {str(e)}"},
{
"success": False,
"error": f"Failed to update rating: {safe_error_message(e)}",
},
status=500,
)
@@ -724,12 +904,16 @@ class PromptRoutesMixin:
try:
prompt_id = int(request.match_info["prompt_id"])
data = await request.json()
new_tag = (data.get("tag") or "").strip()
raw_tag = data.get("tag")
if not new_tag:
if not raw_tag:
return web.json_response(
{"success": False, "error": "Tag cannot be empty"}, status=400
)
error = tags_error([raw_tag])
if error is not None:
return error
new_tag = raw_tag.strip()
prompt = await self._run_in_executor(self.db.get_prompt_by_id, prompt_id)
if not prompt:
@@ -758,7 +942,11 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Add tag error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to add tag: {str(e)}"}, status=500
{
"success": False,
"error": f"Failed to add tag: {safe_error_message(e)}",
},
status=500,
)
async def add_tags_to_prompt(self, request):
@@ -778,6 +966,12 @@ class PromptRoutesMixin:
{"success": False, "error": "Tags must be a non-empty list"},
status=400,
)
# Blank entries are skipped below, so only the rest is validated.
error = tags_error(
[t for t in new_tags if not (isinstance(t, str) and not t.strip())]
)
if error is not None:
return error
prompt = await self._run_in_executor(self.db.get_prompt_by_id, prompt_id)
if not prompt:
@@ -816,7 +1010,11 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Add tags error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to add tags: {str(e)}"}, status=500
{
"success": False,
"error": f"Failed to add tags: {safe_error_message(e)}",
},
status=500,
)
async def remove_prompt_tag(self, request):
@@ -853,7 +1051,10 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Remove tag error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to remove tag: {str(e)}"},
{
"success": False,
"error": f"Failed to remove tag: {safe_error_message(e)}",
},
status=500,
)
@@ -883,7 +1084,10 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Bulk delete error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to delete prompts: {str(e)}"},
{
"success": False,
"error": f"Failed to delete prompts: {safe_error_message(e)}",
},
status=500,
)
@@ -899,6 +1103,11 @@ class PromptRoutesMixin:
{"success": False, "error": "No prompt IDs or tags provided"},
status=400,
)
if not isinstance(new_tags, list):
return bad_request("Tags must be a list")
error = tags_error(new_tags)
if error is not None:
return error
updated_count = await self._run_in_executor(
self.db.bulk_add_tags, prompt_ids, new_tags
@@ -915,7 +1124,11 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Bulk add tags error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to add tags: {str(e)}"}, status=500
{
"success": False,
"error": f"Failed to add tags: {safe_error_message(e)}",
},
status=500,
)
async def bulk_set_category(self, request):
@@ -945,14 +1158,27 @@ class PromptRoutesMixin:
except Exception as e:
self.logger.error(f"Bulk set category error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to set category: {str(e)}"},
{
"success": False,
"error": f"Failed to set category: {safe_error_message(e)}",
},
status=500,
)
def _all_prompts_for_export(self):
"""Every prompt, fetched in EXPORT_PAGE_SIZE pages (blocking)."""
prompts, offset = [], 0
while True:
page = self.db.search_prompts(limit=EXPORT_PAGE_SIZE, offset=offset)
prompts.extend(page)
if len(page) < EXPORT_PAGE_SIZE:
return prompts
offset += EXPORT_PAGE_SIZE
async def export_prompts(self, request):
"""Export all prompts to JSON."""
try:
prompts = await self._run_in_executor(self.db.search_prompts, limit=10000)
prompts = await self._run_in_executor(self._all_prompts_for_export)
export_data = {
"export_date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
@@ -966,13 +1192,19 @@ class PromptRoutesMixin:
text=json_data,
content_type="application/json",
headers={
"Content-Disposition": f'attachment; filename="prompt_manager_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}.json"'
"Content-Disposition": (
'attachment; filename="prompt_manager_'
f'{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}.json"'
)
},
)
except Exception as e:
self.logger.error(f"Export error: {e}")
return web.json_response(
{"success": False, "error": f"Failed to export prompts: {str(e)}"},
{
"success": False,
"error": f"Failed to export prompts: {safe_error_message(e)}",
},
status=500,
)
+61
View File
@@ -0,0 +1,61 @@
"""Progress fan-out for the background output scan.
The scan runs as an asyncio task that outlives the HTTP request that started
it, so an accidental tab close does not abandon a half-finished scan. Every
SSE subscriber (the request that started the scan, a reconnecting page, a
second tab) receives the latest event as a snapshot and then every event
published after it, until the job finishes.
"""
import asyncio
class ScanJob:
"""One scan run: its latest event, completion flag and live subscribers."""
def __init__(self):
self.last_event = None
self.finished = False
self.task = None
self._subscribers = set()
@property
def running(self):
return not self.finished
@property
def subscriber_count(self):
return len(self._subscribers)
def publish(self, payload):
"""Record ``payload`` as the latest event and hand it to every subscriber."""
self.last_event = payload
for queue in list(self._subscribers):
queue.put_nowait(payload)
def finish(self):
"""Mark the job done and release every subscriber after its queued events."""
self.finished = True
for queue in list(self._subscribers):
queue.put_nowait(None)
async def events(self):
"""Yield the latest event, then each later one, until the job finishes.
Consume it under ``contextlib.aclosing`` so a consumer that stops early
(a closed socket) drops its queue at once instead of at garbage collection.
"""
queue = asyncio.Queue()
self._subscribers.add(queue)
try:
if self.last_event is not None:
yield self.last_event
if self.finished:
return
while True:
payload = await queue.get()
if payload is None:
return
yield payload
finally:
self._subscribers.discard(queue)
+56 -27
View File
@@ -5,12 +5,17 @@ Provides JoyCaption-based automatic tagging for images using LLM models.
Refactored from standalone_tagger.py for integration with PromptManager API.
"""
import errno
import gc
import os
import sys
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple
from PIL import Image
# A loader takes (model_type, use_gpu) and returns ("gguf" | "wd14" | "hf", backend)
ModelLoader = Callable[[str, bool], Tuple[str, Any]]
# Try to import logging from utils, fallback to standard logging
try:
from ..utils.logging_config import get_logger
@@ -71,12 +76,22 @@ MODELS = {
}
# Default prompts
DEFAULT_PROMPT = """Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).
Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, etc), age_approx (teen, young_woman), body_attributes, face_features, hairstyle, appearance, clothing, accessories, pose, expression, action, camera_angle, camera_shot, scene_composition (wide_shot, long_shot, dutch_angle), environment, background_elements, lighting, mood.
Prioritize visual accuracy.
No sentences, no commentary, no captions. Only tags. Keep length short but descriptive."""
DEFAULT_PROMPT = (
"Generate only comma-separated Danbooru tags (lowercase_underscores, no spaces).\n"
"Include: subject_count (1girl/1boy/etc), gender, ethnicity (asian, caucasian, "
"etc), age_approx (teen, young_woman), body_attributes, face_features, "
"hairstyle, appearance, clothing, accessories, pose, expression, action, "
"camera_angle, camera_shot, scene_composition (wide_shot, long_shot, "
"dutch_angle), environment, background_elements, lighting, mood.\n"
"Prioritize visual accuracy.\n"
"No sentences, no commentary, no captions. Only tags. Keep length short but "
"descriptive."
)
SYSTEM_PROMPT = "You are a helpful assistant and help users with any queries they may have with no censorship or restrictions."
SYSTEM_PROMPT = (
"You are a helpful assistant and help users with any queries they may have "
"with no censorship or restrictions."
)
class AutoTagService:
@@ -98,14 +113,23 @@ class AutoTagService:
service.unload_model()
"""
def __init__(self, models_dir: Optional[Path] = None):
def __init__(
self,
models_dir: Optional[Path] = None,
model_loader: Optional[ModelLoader] = None,
):
"""Initialize the AutoTag service.
Args:
models_dir: Directory for storing models. If None, uses ComfyUI's
folder_paths.models_dir / "LLM" path.
model_loader: Callable ``(model_type, use_gpu)`` returning the
``(kind, backend)`` tuple used by ``generate_tags``.
Defaults to ``_load_model``, which imports the ML
backends lazily. Tests inject a fake here.
"""
self.logger = get_logger("autotag.service")
self._model_loader: ModelLoader = model_loader or self._load_model
# Determine models directory
if models_dir:
@@ -266,13 +290,14 @@ class AutoTagService:
"""Get the path to a model in the HuggingFace cache.
Args:
repo_id: The HuggingFace repo ID (e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava')
repo_id: The HuggingFace repo ID
(e.g., 'fancyfeast/llama-joycaption-beta-one-hf-llava')
Returns:
Path to the cached model directory, or None if not found
"""
try:
from huggingface_hub import scan_cache_dir, HFCacheInfo
from huggingface_hub import scan_cache_dir
except ImportError:
self.logger.debug("huggingface_hub not available for cache check")
return None
@@ -466,11 +491,15 @@ class AutoTagService:
"""
if model_type not in MODELS:
raise ValueError(
f"Invalid model type: {model_type}. Must be one of: {', '.join(MODELS.keys())}"
f"Invalid model type: {model_type}. "
f"Must be one of: {', '.join(MODELS.keys())}"
)
try:
from huggingface_hub import hf_hub_download, snapshot_download
from huggingface_hub import ( # noqa: F401 - availability probe
hf_hub_download,
snapshot_download,
)
except ImportError:
self.logger.error("huggingface_hub not installed")
if progress_callback:
@@ -591,13 +620,7 @@ class AutoTagService:
raise RuntimeError(f"Model {model_type} not downloaded")
try:
if model_type == "gguf":
self._tagger = self._load_gguf_tagger(use_gpu)
elif model_type.startswith("wd14"):
self._tagger = self._load_wd14_tagger(model_type, use_gpu)
else:
self._tagger = self._load_hf_tagger()
self._tagger = self._model_loader(model_type, use_gpu)
self._current_model_type = model_type
self.logger.info(f"Model {model_type} loaded successfully")
return True
@@ -608,6 +631,14 @@ class AutoTagService:
self._current_model_type = None
raise RuntimeError(f"Failed to load model: {e}")
def _load_model(self, model_type: str, use_gpu: bool = True) -> Tuple[str, Any]:
"""Default loader: dispatch to a backend, importing it only now."""
if model_type == "gguf":
return self._load_gguf_tagger(use_gpu)
if model_type.startswith("wd14"):
return self._load_wd14_tagger(model_type, use_gpu)
return self._load_hf_tagger()
def _load_gguf_tagger(self, use_gpu: bool = True):
"""Load GGUF-based tagger."""
from llama_cpp import Llama
@@ -690,14 +721,11 @@ class AutoTagService:
self._current_model_type = None
gc.collect()
# Try to clear CUDA cache if available
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except ImportError:
pass
# Clear the CUDA cache only if torch is already loaded; importing
# it here would pull in a heavy dependency the backend never used.
torch = sys.modules.get("torch")
if torch is not None and torch.cuda.is_available():
torch.cuda.empty_cache()
self.logger.info("Model unloaded, memory freed")
@@ -720,7 +748,8 @@ class AutoTagService:
Args:
image_path: Path to the image file
prompt: Custom prompt for tag generation (LLM models only). Uses default if None.
prompt: Custom prompt for tag generation (LLM models only).
Uses default if None.
general_threshold: Confidence threshold for general tags (WD14 only).
character_threshold: Confidence threshold for character tags (WD14 only).
@@ -735,7 +764,7 @@ class AutoTagService:
raise RuntimeError("No model loaded. Call load_model() first.")
if not os.path.exists(image_path):
raise FileNotFoundError(f"Image not found: {image_path}")
raise FileNotFoundError(errno.ENOENT, "Image not found", str(image_path))
use_prompt = prompt or self._custom_prompt
+293 -20
View File
@@ -22,7 +22,7 @@ Example:
extension_name = "PromptManager"
# Get server instance and routes (same pattern as ComfyUI_Assets)
from server import PromptServer
from server import PromptServer # noqa: E402
server_instance = PromptServer.instance
routes = server_instance.routes
@@ -30,8 +30,9 @@ routes = server_instance.routes
# Extension info
extension_uri = None # Will be set in __init__.py
import os
from typing import Dict, Any, List
import os # noqa: E402
from pathlib import Path # noqa: E402
from typing import Any, Dict, List, Optional, Tuple # noqa: E402
# Import logging system
try:
@@ -46,6 +47,108 @@ except ImportError:
# Initialize logger for config operations
config_logger = get_logger("prompt_manager.config")
# Environment variable listing extra directories (os.pathsep-separated) that
# may be used as gallery roots in addition to ComfyUI's own directories.
EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
# Environment variable overriding where config.json is read from and saved to.
CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH"
# ComfyUI folder_paths getters whose directories are valid gallery parents.
_COMFYUI_DIRECTORY_GETTERS = (
"get_output_directory",
"get_input_directory",
"get_temp_directory",
"get_user_directory",
)
def _canonical_path(path: str) -> str:
"""Return a path normalised for comparison (realpath + normcase)."""
return os.path.normcase(os.path.realpath(path))
def write_private_json(path: str, data: Any) -> None:
"""Write ``data`` as JSON to ``path`` readable by the current user only.
config.json can hold an API key (integrations.lora_manager), so it is
created with mode 0600 and, on POSIX, an existing more open file is
tightened too. Windows has no POSIX modes; there the chmod is a best
effort that only affects the read-only bit. Parent directories are
created as needed. Raises OSError when the file cannot be written.
"""
import json
parent_dir = os.path.dirname(path)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
if os.name == "posix":
os.chmod(path, 0o600)
def _is_filesystem_root(path: str) -> bool:
"""True for '/' on POSIX and drive roots such as 'C:\\' on Windows."""
return os.path.dirname(path) == path
def _import_folder_paths():
"""ComfyUI's folder_paths module, or None outside ComfyUI."""
try:
import folder_paths
except ImportError:
return None
return folder_paths
def _comfyui_base_directory() -> Optional[str]:
"""ComfyUI's base directory: folder_paths.base_path, else the output dir's
parent."""
folder_paths = _import_folder_paths()
if folder_paths is None:
return None
base = getattr(folder_paths, "base_path", None)
if isinstance(base, str) and base:
return base
getter = getattr(folder_paths, "get_output_directory", None)
if getter is None:
return None
try:
output_dir = getter()
except Exception:
return None
if isinstance(output_dir, str) and output_dir:
return os.path.dirname(os.path.realpath(output_dir))
return None
def _comfyui_directories() -> List[str]:
"""Directories reported by ComfyUI's folder_paths module, if importable."""
folder_paths = _import_folder_paths()
if folder_paths is None:
return []
found = []
for getter_name in _COMFYUI_DIRECTORY_GETTERS:
getter = getattr(folder_paths, getter_name, None)
if getter is None:
continue
try:
directory = getter()
except Exception:
continue
if isinstance(directory, str) and directory:
found.append(directory)
return found
def _extra_gallery_roots() -> List[str]:
"""Directories the user opted in through PROMPT_MANAGER_EXTRA_GALLERY_ROOTS."""
raw = os.environ.get(EXTRA_GALLERY_ROOTS_ENV, "")
return [entry.strip() for entry in raw.split(os.pathsep) if entry.strip()]
class GalleryConfig:
"""Configuration class for the gallery monitoring and image processing system.
@@ -106,8 +209,9 @@ class GalleryConfig:
Returns:
Dict[str, Any]: A nested dictionary containing all gallery configuration
sections: monitoring, tracking, database, web_interface, and performance.
Each section contains the relevant configuration parameters as key-value pairs.
sections: monitoring, tracking, database, web_interface, and
performance. Each section contains the relevant configuration
parameters as key-value pairs.
Example:
config = GalleryConfig.get_config()
@@ -166,7 +270,9 @@ class GalleryConfig:
if "enabled" in monitoring:
cls.MONITORING_ENABLED = monitoring["enabled"]
if "directories" in monitoring:
cls.MONITORING_DIRECTORIES = monitoring["directories"]
directories = cls._validated_directories(monitoring["directories"])
if directories is not None:
cls.MONITORING_DIRECTORIES = directories
if "extensions" in monitoring:
cls.SUPPORTED_EXTENSIONS = monitoring["extensions"]
if "processing_delay" in monitoring:
@@ -200,6 +306,127 @@ class GalleryConfig:
if "metadata_extraction_timeout" in performance:
cls.METADATA_EXTRACTION_TIMEOUT = performance["metadata_extraction_timeout"]
@classmethod
def _validated_directories(cls, directories: Any) -> Optional[List[str]]:
"""Gallery roots from a config payload that pass validate_gallery_root.
A hand-edited config.json is as untrusted as a settings request:
entries that fail validation are dropped with a warning (naming only
the basename), and a value that is not a list is ignored entirely
(``None`` is returned so the current roots stay as they are).
"""
if not isinstance(directories, list):
config_logger.warning(
"Ignoring gallery monitoring directories: expected a list"
)
return None
kept = []
for entry in directories:
ok, reason = cls.validate_gallery_root(entry)
if ok:
kept.append(entry)
continue
label = os.path.basename(entry) if isinstance(entry, str) else entry
config_logger.warning(f"Ignoring gallery root {label!r}: {reason}")
return kept
@classmethod
def allowed_gallery_parents(cls) -> List[str]:
"""Canonical directories under which gallery roots may live.
Filesystem roots are never allowed as parents, even when listed in
the environment, because that would re-open the whole disk.
"""
parents = []
for candidate in _comfyui_directories() + _extra_gallery_roots():
try:
canonical = _canonical_path(candidate)
except (OSError, ValueError):
continue
if _is_filesystem_root(canonical) or not os.path.isdir(canonical):
continue
if canonical not in parents:
parents.append(canonical)
return parents
@classmethod
def path_anchors(cls) -> List[str]:
"""Canonical directories that relative gallery paths are resolved against.
The ComfyUI base directory comes first, followed by the parent of each
directory listed in PROMPT_MANAGER_EXTRA_GALLERY_ROOTS, so an extra
root is addressed by its own name (``gallery/sub``) rather than by an
absolute path. The same anchors drive the public (relative) form of
paths in API responses.
"""
anchors = []
base = _comfyui_base_directory()
candidates = [base] if base else []
candidates.extend(os.path.dirname(root) for root in _extra_gallery_roots())
for candidate in candidates:
try:
canonical = _canonical_path(candidate)
except (OSError, ValueError):
continue
if canonical not in anchors:
anchors.append(canonical)
return anchors
@classmethod
def resolve_gallery_root(cls, path: str) -> str:
"""Canonical absolute path for a gallery root given in any accepted form.
Absolute paths are canonicalised as-is. Relative paths are tried
against each :meth:`path_anchors` entry and the first existing match
wins; otherwise the first anchor (or the current directory when there
is none) is used, so the caller's existence check reports it.
"""
path = path.strip()
if os.path.isabs(path):
return _canonical_path(path)
anchors = cls.path_anchors()
for anchor in anchors:
candidate = os.path.join(anchor, path)
if os.path.exists(candidate):
return _canonical_path(candidate)
return _canonical_path(os.path.join(anchors[0], path) if anchors else path)
@classmethod
def validate_gallery_root(cls, path: Any) -> Tuple[bool, str]:
"""Check whether ``path`` may be used as a gallery root.
Returns:
(True, "") when the directory lies inside one of
:meth:`allowed_gallery_parents`; otherwise (False, reason).
"""
if not isinstance(path, str) or not path.strip():
return False, "Gallery root must be a non-empty path"
try:
canonical = cls.resolve_gallery_root(path)
except (OSError, ValueError):
return False, "Gallery root could not be resolved"
if _is_filesystem_root(canonical):
return False, "Gallery root cannot be a filesystem root"
if canonical == _canonical_path(os.path.expanduser("~")):
return False, "Gallery root cannot be the home directory"
if not os.path.exists(canonical):
return False, "Gallery root does not exist"
if not os.path.isdir(canonical):
return False, "Gallery root is not a directory"
candidate = Path(canonical)
for parent in cls.allowed_gallery_parents():
if candidate.is_relative_to(Path(parent)):
return True, ""
return (
False,
"Gallery root must be inside a ComfyUI directory "
f"(output, input, temp, user) or one listed in {EXTRA_GALLERY_ROOTS_ENV}",
)
class IntegrationConfig:
"""Configuration for third-party extension integrations.
@@ -272,11 +499,36 @@ class PromptManagerConfig:
RESULT_TIMEOUT = 5 # Seconds to auto-hide results in ComfyUI node
SHOW_TEST_BUTTON = False # Show API test button in node UI
WEBUI_DISPLAY_MODE = "newtab" # 'popup' or 'newtab'
INFINITE_SCROLL = False # Dashboard list: load the next page on scroll by default
# Performance settings
MAX_SEARCH_RESULTS = 100
ENABLE_FUZZY_SEARCH = False # Requires fuzzywuzzy
AUTO_BACKUP_INTERVAL = 24 # Hours
# Threads for the output scan and thumbnail generation; half the cores by
# default so ComfyUI's own generation keeps CPU headroom. 1..max_worker_threads().
DEFAULT_WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2)
WORKER_THREADS = DEFAULT_WORKER_THREADS
@classmethod
def max_worker_threads(cls) -> int:
"""Largest accepted worker-thread count: the detected core count."""
return os.cpu_count() or 1
@classmethod
def clamp_worker_threads(cls, value, fallback: int) -> int:
"""Coerce a worker count from config.json into ``[1, max_worker_threads()]``.
Non-numeric values (and booleans) return ``fallback`` so a hand-edited
file cannot disable the setting.
"""
if isinstance(value, bool):
return fallback
if isinstance(value, str) and value.strip().isdigit():
value = int(value.strip())
if not isinstance(value, int):
return fallback
return max(1, min(value, cls.max_worker_threads()))
@classmethod
def get_config(cls) -> Dict[str, Any]:
@@ -304,18 +556,33 @@ class PromptManagerConfig:
"result_timeout": cls.RESULT_TIMEOUT,
"show_test_button": cls.SHOW_TEST_BUTTON,
"webui_display_mode": cls.WEBUI_DISPLAY_MODE,
"infinite_scroll": cls.INFINITE_SCROLL,
},
"performance": {
"max_search_results": cls.MAX_SEARCH_RESULTS,
"enable_fuzzy_search": cls.ENABLE_FUZZY_SEARCH,
"auto_backup_interval": cls.AUTO_BACKUP_INTERVAL,
"worker_threads": cls.WORKER_THREADS,
},
"gallery": GalleryConfig.get_config(),
"integrations": IntegrationConfig.get_config(),
}
@classmethod
def load_from_file(cls, config_path: str):
def get_config_path(cls) -> str:
"""Path of the persisted config.json.
Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file
lives at the repository root next to ``pyproject.toml``.
"""
override = os.environ.get(CONFIG_PATH_ENV, "").strip()
if override:
return override
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
return os.path.join(repo_root, "config.json")
@classmethod
def load_from_file(cls, config_path: Optional[str] = None):
"""Load configuration settings from a JSON file.
Reads configuration from the specified JSON file and updates the current
@@ -324,7 +591,8 @@ class PromptManagerConfig:
Args:
config_path (str): Path to the JSON configuration file to load.
Can be relative or absolute path.
Can be relative or absolute path. Defaults to
:meth:`get_config_path`.
Raises:
The method handles all exceptions internally and logs errors rather
@@ -336,6 +604,9 @@ class PromptManagerConfig:
"""
import json
if config_path is None:
config_path = cls.get_config_path()
if os.path.exists(config_path):
try:
with open(config_path, "r") as f:
@@ -348,7 +619,7 @@ class PromptManagerConfig:
config_logger.info(f"Config file not found: {config_path}, using defaults")
@classmethod
def save_to_file(cls, config_path: str):
def save_to_file(cls, config_path: Optional[str] = None):
"""Save the current configuration to a JSON file.
Serializes the complete configuration (including gallery settings) to
@@ -357,6 +628,7 @@ class PromptManagerConfig:
Args:
config_path (str): Path where the JSON configuration file should be saved.
Parent directories will be created if they don't exist.
Defaults to :meth:`get_config_path`.
Raises:
The method handles all exceptions internally and logs errors rather
@@ -366,15 +638,11 @@ class PromptManagerConfig:
PromptManagerConfig.save_to_file('backup_config.json')
PromptManagerConfig.save_to_file('/etc/comfyui/prompt_manager.json')
"""
import json
if config_path is None:
config_path = cls.get_config_path()
try:
config = cls.get_config()
os.makedirs(os.path.dirname(config_path), exist_ok=True)
with open(config_path, "w") as f:
json.dump(config, f, indent=2)
write_private_json(config_path, cls.get_config())
config_logger.info(f"Saved configuration to {config_path}")
except Exception as e:
config_logger.error(f"Error saving config to {config_path}: {e}")
@@ -390,7 +658,8 @@ class PromptManagerConfig:
Args:
new_config (Dict[str, Any]): Nested dictionary containing configuration
updates. Should follow the same structure as returned by get_config().
Valid top-level keys are: 'database', 'web_ui', 'performance', 'gallery'.
Valid top-level keys are: 'database', 'web_ui', 'performance',
'gallery'.
Example:
new_settings = {
@@ -414,6 +683,8 @@ class PromptManagerConfig:
cls.SHOW_TEST_BUTTON = web_ui["show_test_button"]
if "webui_display_mode" in web_ui:
cls.WEBUI_DISPLAY_MODE = web_ui["webui_display_mode"]
if isinstance(web_ui.get("infinite_scroll"), bool):
cls.INFINITE_SCROLL = web_ui["infinite_scroll"]
performance = new_config.get("performance", {})
if "max_search_results" in performance:
@@ -422,6 +693,10 @@ class PromptManagerConfig:
cls.ENABLE_FUZZY_SEARCH = performance["enable_fuzzy_search"]
if "auto_backup_interval" in performance:
cls.AUTO_BACKUP_INTERVAL = performance["auto_backup_interval"]
if "worker_threads" in performance:
cls.WORKER_THREADS = cls.clamp_worker_threads(
performance["worker_threads"], cls.WORKER_THREADS
)
# Update gallery config
if "gallery" in new_config:
@@ -434,8 +709,6 @@ class PromptManagerConfig:
# Load configuration on import
try:
config_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
config_file = os.path.join(config_dir, "config.json")
PromptManagerConfig.load_from_file(config_file)
PromptManagerConfig.load_from_file()
except Exception as e:
config_logger.error(f"Error during config initialization: {e}")
+181 -4
View File
@@ -15,6 +15,7 @@ import threading
import urllib.request
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from urllib.parse import urlsplit
try:
from ..utils.logging_config import get_logger
@@ -108,6 +109,103 @@ def _looks_like_lora_manager(path: Path) -> bool:
return (path / "py").is_dir() or (path / "lora_manager").is_dir()
def _public_path_anchors() -> List[str]:
"""Directories the API's public (relative) paths are rendered against."""
try:
from .config import GalleryConfig
except ImportError:
try:
from config import GalleryConfig
except ImportError:
return []
try:
return list(GalleryConfig.path_anchors())
except Exception:
return []
def custom_nodes_directories() -> List[Path]:
"""ComfyUI ``custom_nodes`` directories a LoraManager install may live in.
Inside ComfyUI these come from ``folder_paths`` (``base_path`` plus any
registered ``custom_nodes`` folders); outside it, the parent of this
package, which is ``custom_nodes`` for a normal install.
"""
found: List[Path] = []
def add(candidate) -> None:
try:
resolved = Path(candidate).resolve()
except (OSError, RuntimeError, ValueError):
return
if resolved not in found:
found.append(resolved)
try:
import folder_paths
except ImportError:
folder_paths = None
if folder_paths is not None:
base = getattr(folder_paths, "base_path", None)
if isinstance(base, str) and base:
add(Path(base) / "custom_nodes")
getter = getattr(folder_paths, "get_folder_paths", None)
if callable(getter):
try:
registered = getter("custom_nodes") or []
except Exception:
registered = []
for entry in registered:
add(entry)
if not found:
add(Path(__file__).resolve().parent.parent.parent)
return found
def resolve_lora_manager_path(path) -> Optional[str]:
"""Canonical form of ``path`` when it may be used as the LoraManager install.
Accepted: an existing directory that is a direct child of one of
:func:`custom_nodes_directories` and whose name contains ``lora``
(case-insensitive). A relative path is tried against the ComfyUI root
and each custom_nodes directory. Anything else, including a symlink
that escapes custom_nodes, yields None.
"""
if not isinstance(path, str) or not path.strip():
return None
path = path.strip()
roots = []
for root in custom_nodes_directories():
try:
roots.append(Path(os.path.normcase(os.path.realpath(str(root)))))
except (OSError, ValueError):
continue
candidates = [path]
if not os.path.isabs(path):
# Public (API) paths are relative to these anchors; the ComfyUI root
# and the custom_nodes directories cover hand-typed relative paths.
bases = list(_public_path_anchors())
comfy_root = find_comfyui_root()
if comfy_root:
bases.append(comfy_root)
bases.extend(root.parent for root in roots)
bases.extend(roots)
candidates = [os.path.join(str(base), path) for base in bases]
for candidate in candidates:
try:
canonical = Path(os.path.normcase(os.path.realpath(candidate)))
is_dir = canonical.is_dir()
except (OSError, RuntimeError, ValueError):
continue
if not is_dir or "lora" not in canonical.name.lower():
continue
if canonical.parent in roots:
return str(canonical)
return None
# ── Metadata reading ─────────────────────────────────────────────────
@@ -298,16 +396,95 @@ def get_preview_image_from_metadata(
_THUMB_MAX_SIZE = 512
# Only these hosts may ever receive the CivitAI API key, and only over HTTPS.
CIVITAI_HOSTS = frozenset(
{"civitai.com", "www.civitai.com", "api.civitai.com", "image.civitai.com"}
)
MAX_CIVITAI_DOWNLOAD_BYTES = 50 * 1024 * 1024
_DOWNLOAD_CHUNK_BYTES = 64 * 1024
class _RefuseRedirects(urllib.request.HTTPRedirectHandler):
"""Never follow a redirect: urllib would re-send the Bearer header to it."""
def redirect_request(self, req, fp, code, msg, headers, newurl):
logger.warning(f"Refusing redirect from {req.full_url} to {newurl}")
return None
def _build_opener(*handlers) -> urllib.request.OpenerDirector:
"""Build the outbound opener; extra handlers let tests serve canned responses."""
return urllib.request.build_opener(_RefuseRedirects, *handlers)
_opener = _build_opener()
def _open_url(req: urllib.request.Request, timeout: float = 10):
"""Open ``req`` through the redirect-refusing opener."""
return _opener.open(req, timeout=timeout)
def is_civitai_url(url: str) -> bool:
"""True when ``url`` is an HTTPS URL on an allow-listed CivitAI host."""
try:
parts = urlsplit(url)
port = parts.port
except ValueError:
return False
hostname = (parts.hostname or "").lower()
return (
parts.scheme == "https"
and hostname in CIVITAI_HOSTS
and parts.username is None
and parts.password is None
and port in (None, 443)
)
def _read_capped(resp, cap: int) -> Optional[bytes]:
"""Read a response body, refusing anything larger than ``cap`` bytes."""
declared = resp.headers.get("Content-Length")
if declared is not None:
try:
if int(declared) > cap:
return None
except ValueError:
pass
chunks = []
total = 0
while True:
chunk = resp.read(_DOWNLOAD_CHUNK_BYTES)
if not chunk:
break
total += len(chunk)
if total > cap:
return None
chunks.append(chunk)
return b"".join(chunks)
def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
"""Download a single image, resize to thumbnail, save as JPEG."""
"""Download a single CivitAI image, resize to thumbnail, save as JPEG.
Refuses any URL that is not HTTPS on an allow-listed CivitAI host so the
API key is never sent elsewhere, and refuses bodies larger than
MAX_CIVITAI_DOWNLOAD_BYTES.
"""
if not is_civitai_url(url):
logger.warning(f"Refusing non-CivitAI download URL: {url}")
return None
try:
headers = {"User-Agent": "ComfyUI-PromptManager/1.0"}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
req = urllib.request.Request(url, headers=headers)
with urllib.request.urlopen(req, timeout=10) as resp:
raw = resp.read()
with _open_url(req, timeout=10) as resp:
raw = _read_capped(resp, MAX_CIVITAI_DOWNLOAD_BYTES)
if raw is None:
logger.warning(f"Refusing oversized download: {url}")
return None
# Resize to thumbnail to save disk space
from io import BytesIO
@@ -361,7 +538,7 @@ def download_civitai_images(
if not isinstance(img, dict):
continue
url = img.get("url", "")
if not isinstance(url, str) or not url.startswith("http"):
if not isinstance(url, str) or not is_civitai_url(url):
continue
url_hash = hashlib.md5(url.encode()).hexdigest()[:12]
+7 -7
View File
@@ -1,9 +1,14 @@
[project]
name = "promptmanager"
description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite."
version = "3.2.3"
version = "3.3.1"
license = {file = "LICENSE"}
dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"]
requires-python = ">=3.10"
dependencies = ["watchdog>=2.1.0", "Pillow>=8.0.0"]
[build-system]
requires = ["setuptools>=84.0.0"]
build-backend = "setuptools.build_meta"
[project.urls]
Repository = "https://github.com/ComfyAssets/ComfyUI_PromptManager"
@@ -18,10 +23,5 @@ Icon = "https://avatars.githubusercontent.com/u/213204677?s=200"
line-length = 88
target-version = ["py310"]
[tool.flake8]
max-line-length = 88
extend-ignore = ["E203", "W503"]
exclude = ["reference/", ".git", "__pycache__", "web/", "venv/"]
[tool.bandit]
exclude_dirs = ["tests", "reference"]
+11 -7
View File
@@ -1,14 +1,18 @@
# Runtime dependencies
watchdog>=2.1.0
Pillow>=8.0.0
aiohttp>=3.8.0
watchdog>=6.0.0
Pillow>=12.3.0
aiohttp>=3.14.5
# Tests and coverage
coverage>=7.16.2
numpy>=2.2.6 # WD14 engine tests
# Code quality
black>=24.0.0
flake8>=7.0.0
black>=26.10.1
flake8>=7.4.1
# Security scanning
bandit[toml]>=1.7.0
bandit[toml]>=1.9.4
# Pre-commit hooks
pre-commit>=4.0.0
pre-commit>=4.6.2
+4 -26
View File
@@ -1,26 +1,4 @@
# Core dependencies for PromptManager
# Note: Most dependencies are already included with ComfyUI
# Already included with Python standard library:
# - sqlite3
# - hashlib
# - json
# - datetime
# - os
# - typing
# - threading
# - uuid
# Required for gallery functionality:
watchdog>=2.1.0 # For file system monitoring
Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)
# Optional dependencies for enhanced search functionality:
# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)
# sqlalchemy>=1.4.0 # For advanced ORM features (optional)
# Development dependencies (optional):
# pytest>=6.0.0 # For running tests
# black>=22.0.0 # For code formatting
# flake8>=4.0.0 # For linting
# mypy>=0.910 # For type checking
# Runtime dependencies for PromptManager.
# aiohttp, numpy and torch come from ComfyUI itself and are not listed here.
watchdog>=2.1.0 # Output directory monitoring
Pillow>=8.0.0 # Image metadata extraction and thumbnails
-64
View File
@@ -1,64 +0,0 @@
#!/usr/bin/env python3
"""
ComfyUI_PromptManager Gallery System Restart Script.
This script reinitializes and restarts the automatic image gallery monitoring system
for ComfyUI_PromptManager. Use this script after code modifications or when the
gallery system needs to be reloaded.
The script performs the following operations:
1. Initializes the database connection
2. Creates a PromptTracker instance for linking prompts to images
3. Sets up the ImageMonitor for real-time file system monitoring
4. Starts monitoring ComfyUI output directories
5. Reports system status and readiness
Usage:
python restart_gallery.py
This is particularly useful during development when you need to restart the
gallery monitoring system without restarting the entire ComfyUI application.
"""
import sys
import os
# Add current directory to Python path
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)
try:
# Import and initialize the gallery system
from database.operations import PromptDatabase
from utils.prompt_tracker import get_prompt_tracker
from utils.image_monitor import get_image_monitor
print("[RESTART] Restarting PromptManager Gallery System...")
# Initialize components using singleton getters
db = PromptDatabase()
tracker = get_prompt_tracker(db)
monitor = get_image_monitor(db, tracker)
print("[SUCCESS] Components initialized successfully")
# Test image monitoring directories
status = monitor.get_status()
print(f"[INFO] Monitoring status: {status}")
# Start monitoring
monitor.start_monitoring()
print("[SUCCESS] Image monitoring restarted")
print("\n[READY] Gallery system restart completed!")
print("Generate some images now to test the automatic linking.")
except ImportError as e:
print(f"[ERROR] Import error: {e}")
print("Make sure you're running this from the PromptManager directory")
except Exception as e:
print(f"[ERROR] Error: {e}")
import traceback
traceback.print_exc()
+1
View File
@@ -0,0 +1 @@
"""Browser workflow fixtures; never imported by production code."""
+154
View File
@@ -0,0 +1,154 @@
"""Serve production pages and API routes with disposable state for browser tests."""
import argparse
import json
import os
import shutil
from pathlib import Path
import sys
import tempfile
import time
import types
from aiohttp import web
from PIL import Image
from PIL.PngImagePlugin import PngInfo
class TestModel:
default_prompt = "Describe the image"
wd14_general_threshold = 0.35
wd14_character_threshold = 0.85
custom_prompt = ""
loaded = None
def get_models_status(self) -> dict:
return {
name: {"downloaded": True, "loaded": self.loaded == name}
for name in ("gguf", "wd14-vit", "wd14-convnext", "wd14-swinv2")
}
def is_model_loaded(self) -> bool:
return self.loaded is not None
def get_loaded_model_type(self) -> str | None:
return self.loaded
def load_model(self, model_type: str, use_gpu: bool = True) -> bool:
self.loaded = model_type
return True
def unload_model(self) -> None:
self.loaded = None
def generate_tags(self, image_path: str, **kwargs: object) -> list[str]:
return ["generated", Path(image_path).parent.parent.name]
def create_app(root: Path) -> web.Application:
os.environ["PROMPT_MANAGER_CONFIG_PATH"] = str(root / "config.json")
os.chdir(root)
app = web.Application()
routes = web.RouteTableDef()
server = types.ModuleType("server")
server.PromptServer = types.SimpleNamespace(
instance=types.SimpleNamespace(app=app, routes=routes)
)
sys.modules["server"] = server
folder_paths = types.ModuleType("folder_paths")
folder_paths.base_path = str(root)
folder_paths.models_dir = str(root / "models")
for name in ("output", "input", "temp", "user"):
directory = root / name
directory.mkdir()
setattr(folder_paths, f"get_{name}_directory", lambda path=directory: str(path))
sys.modules["folder_paths"] = folder_paths
from py.config import GalleryConfig, PromptManagerConfig
PromptManagerConfig.DEFAULT_DB_PATH = str(root / "browser.db")
GalleryConfig.MONITORING_DIRECTORIES = [str(root / "output"), str(root / "input")]
from py.api import PromptManagerAPI
from py import autotag
from utils.hashing import generate_prompt_hash
model = TestModel()
autotag.get_autotag_service = lambda: model
api = PromptManagerAPI()
api.add_routes(routes)
def seed(extra_images: int = 0) -> None:
with api.db.model.get_connection() as conn:
for table in ("generated_images", "prompt_tags", "tags", "prompts"):
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
for index, (directory, text, tag) in enumerate(
[("output", "Cat portrait", "cats"), ("input", "Dog landscape", "dogs")]
):
image_path = root / directory / "shared" / "same.png"
image_path.parent.mkdir(exist_ok=True)
info = PngInfo()
info.add_text(
"prompt",
json.dumps(
{"1": {"class_type": "CLIPTextEncode", "inputs": {"text": text}}}
),
)
with Image.new("RGB", (64, 64), (index * 255, 64, 128)) as image:
image.save(image_path, pnginfo=info)
stamp = time.time() - index * 10
os.utime(image_path, (stamp, stamp))
pid = api.db.save_prompt(
text, tags=[tag], prompt_hash=generate_prompt_hash(text)
)
api.db.link_image_to_prompt(pid, str(image_path))
if index == 0:
for number in range(extra_images):
extra = root / directory / "bulk" / f"{number:04}.png"
extra.parent.mkdir(exist_ok=True)
shutil.copyfile(image_path, extra)
api.db.link_image_to_prompt(pid, str(extra))
api._generate_gallery_thumbnails_sync(api._get_all_output_dirs(), (150, 150))
api.invalidate_gallery_cache()
@routes.post("/__test__/reset")
async def reset(request: web.Request) -> web.Response:
data = await request.json() if request.can_read_body else {}
extra = data.get("extra_images", 0)
if (
isinstance(extra, bool)
or not isinstance(extra, int)
or not 0 <= extra <= 200
):
raise web.HTTPBadRequest()
await api._run_in_executor(seed, extra)
return web.json_response({"success": True})
@routes.get("/api/settings/Comfy.ColorPalette")
async def palette(request: web.Request) -> web.Response:
return web.json_response("dark")
async def cleanup(application: web.Application) -> None:
api.db.close_all()
seed()
app.on_cleanup.append(cleanup)
app.add_routes(routes)
return app
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8765)
args = parser.parse_args()
with tempfile.TemporaryDirectory(prefix="pm-browser-") as directory:
web.run_app(
create_app(Path(directory).resolve()), host="127.0.0.1", port=args.port
)
if __name__ == "__main__":
main()
+141
View File
@@ -0,0 +1,141 @@
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
test.beforeEach(async ({ request, page, browserName }) => {
await startCoverage(page, browserName);
await expect(await request.post('/__test__/reset')).toBeOK();
page.workflowErrors = [];
page.on('pageerror', error => page.workflowErrors.push(error.message));
page.on('response', response => {
if (response.status() >= 500 && response.url().includes('/prompt_manager/')) {
page.workflowErrors.push(`${response.status()} ${response.url()}`);
}
});
await page.addInitScript(() => { localStorage.setItem('pm_update_notice_dismissed', '3.2.0'); });
});
test.afterEach(async ({ page, browserName }) => {
await saveCoverage(page, browserName);
expect(page.workflowErrors).toEqual([]);
});
test('actual gallery form saves a prompt that the admin page can edit and reload', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
await page.locator('#addPromptBtn').click();
await page.locator('#addPromptText').fill('Browser saved landscape');
await page.locator('#addPromptCategory').fill('Browser tests');
await page.locator('#addPromptTagInput').fill('browser-tag');
await page.locator('#addPromptTagInput').press('Enter');
await page.locator('#saveNewPromptBtn').click();
await expect(page.locator('#addPromptModal')).toBeHidden();
await page.goto('/prompt_manager/admin');
const text = page.locator('.prompt-text').filter({ hasText: 'Browser saved landscape' });
await expect(text).toBeVisible();
const id = await text.getAttribute('data-id');
await page.locator(`[data-action="edit"][data-prompt-id="${id}"]`).click();
const editor = page.locator(`.prompt-text[data-id="${id}"]`);
await editor.fill('Browser edited portrait');
await editor.press('Enter');
await expect(page.locator('.prompt-text').filter({ hasText: 'Browser edited portrait' })).toBeVisible();
await page.reload();
await expect(page.locator('.prompt-text').filter({ hasText: 'Browser edited portrait' })).toBeVisible();
const saved = await request.post('/prompt_manager/save', { data: { text: 'Browser saved landscape' } });
expect((await saved.json()).prompt_id).not.toBe(Number(id));
});
test('actual gallery review generates and persists tags for the selected root', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
await page.locator('#autoTagBtn').click();
const generated = page.waitForResponse(response => response.url().endsWith('/prompt_manager/autotag/single'));
await page.locator('#startReviewBtn').click();
const generation = await generated;
expect(generation.request().postDataJSON().root_index).toBe(0);
await expect(page.locator('#reviewTagsContainer')).toContainText('generated');
const result = page.waitForResponse(response => response.url().endsWith('/prompt_manager/autotag/apply'));
await page.locator('#applyReviewBtn').click();
expect((await (await result).json()).success).toBe(true);
const data = await (await request.get('/prompt_manager/recent?limit=20')).json();
expect(data.results.find(prompt => prompt.text === 'Cat portrait').tags).toContain('generated');
expect(data.results.find(prompt => prompt.text === 'Dog landscape').tags).not.toContain('generated');
});
test('actual Tags page reloads the selected filter and opens its image viewer', async ({ page }) => {
await page.goto('/prompt_manager/admin#/tags');
await page.locator('[data-tag-name="cats"]').click();
await expect(page.locator('#tagPromptsGrid')).toContainText('Cat portrait');
await page.locator('[data-tag-name="dogs"]').click();
await expect(page.locator('#tagPromptsGrid')).toContainText('Dog landscape');
await expect(page.locator('#tagPromptsGrid')).not.toContainText('Cat portrait');
await page.locator('#tagPromptsGrid img').first().click();
await expect(page.locator('.viewer-canvas img')).toBeVisible();
await page.locator('.viewer-container').getByRole('button', { name: 'Close', exact: true }).click();
await expect(page.locator('.viewer-container')).toHaveCount(0);
});
test('actual gallery serves colliding filenames and shows the matching prompt', async ({ page }) => {
await page.goto('/prompt_manager/gallery');
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(2);
const images = page.locator('#galleryGrid img[data-original]');
const first = await images.nth(0).getAttribute('data-original');
const second = await images.nth(1).getAttribute('data-original');
expect(first).not.toBe(second);
await page.locator('#galleryGrid .image-item[data-image-index="0"]').click();
await expect(page.locator('#metadata-content')).toContainText('Cat portrait');
await page.keyboard.press('ArrowRight');
await expect(page.locator('#metadata-content')).toContainText('Dog landscape');
await page.keyboard.press('Escape');
await expect(page.locator('#metadata-sidebar')).toHaveCount(0);
});
test('an edit collision is rejected by the backend and the page restores the original text', async ({ page }) => {
await page.goto('/prompt_manager/admin');
const original = page.locator('.prompt-text').filter({ hasText: 'Cat portrait' });
const id = await original.getAttribute('data-id');
await page.locator(`[data-action="edit"][data-prompt-id="${id}"]`).click();
const editor = page.locator(`.prompt-text[data-id="${id}"]`);
await editor.fill('Dog landscape');
const rejected = page.waitForResponse(response => response.request().method() === 'PUT' && response.url().endsWith(`/prompts/${id}`));
await editor.press('Enter');
expect((await rejected).status()).toBe(409);
await expect(editor).toHaveText('Cat portrait');
await page.reload();
await expect(page.locator(`.prompt-text[data-id="${id}"]`)).toHaveText('Cat portrait');
});
test('a delayed real Tags response cannot replace a newer selection', async ({ page }) => {
let release;
const delayed = new Promise(resolve => { release = resolve; });
let received;
const waiting = new Promise(resolve => { received = resolve; });
await page.route('**/prompt_manager/tags/cats/prompts?*', async route => {
const response = await route.fetch();
received();
await delayed;
await route.fulfill({ response });
});
await page.goto('/prompt_manager/admin#/tags');
await page.locator('[data-tag-name="cats"]').click();
await waiting;
await page.locator('[data-tag-name="dogs"]').click();
await expect(page.locator('#tagPromptsGrid')).toContainText('Dog landscape');
const completed = page.waitForResponse(response => response.url().includes('/tags/cats/prompts?'));
release();
await completed;
await expect(page.locator('#tagPromptsGrid')).not.toContainText('Cat portrait');
});
test('gallery infinite scroll appends a real backend page while retaining existing cards', async ({ page, request }) => {
await expect(await request.post('/__test__/reset', { data: { extra_images: 108 } })).toBeOK();
await page.goto('/prompt_manager/gallery');
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(100);
await page.locator('#settingsBtn').click();
await page.locator('#infiniteScrollToggle').check();
await page.locator('#saveSettingsBtn').click();
await expect(page.locator('#settingsModal')).toBeHidden();
const first = page.locator('#galleryGrid .image-item').first();
await first.evaluate(element => { element.dataset.preserved = 'yes'; });
await page.locator('#galleryScroll').evaluate(element => { element.scrollTop = element.scrollHeight; });
await expect(page.locator('#galleryGrid .image-item')).toHaveCount(110);
await expect(first).toHaveAttribute('data-preserved', 'yes');
});
+50
View File
@@ -0,0 +1,50 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const vm = require('node:vm');
function fixture() {
const panel = { innerHTML: '', isConnected: true };
const requests = new Map();
const context = vm.createContext({ document: { readyState: 'loading', addEventListener() {},
getElementById: id => id === 'metadata-content' ? panel : null },
window: {}, console: { error() {} },
});
const Admin = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/admin.js'), 'utf8') + '\nPromptAdmin;', context, { filename: path.join(__dirname, '../../web/js/admin.js') });
const manager = Object.create(Admin.prototype);
manager.extractImageMetadata = url => new Promise((resolve, reject) => requests.set(url, { resolve, reject }));
manager.updateMetadataSidebar = metadata => { panel.innerHTML = metadata?.imagePath || 'No metadata'; };
manager.hideModal = () => {};
const load = url => manager.loadMetadataForImage({ getAttribute: () => url });
return { manager, panel, requests, load };
}
test('admin metadata and copy actions stay on the selected image', async () => {
const { manager, panel, requests, load } = fixture();
const old = load('old.png');
const current = load('current.png');
requests.get('current.png').resolve({ imagePath: 'current.png' }); await current;
requests.get('old.png').resolve({ imagePath: 'old.png' }); await old;
assert.equal(panel.innerHTML, 'current.png');
assert.equal(manager.currentMetadata.imagePath, 'current.png');
});
test('navigation clears old action data and late errors cannot replace current metadata', async () => {
const { manager, panel, requests, load } = fixture();
manager.currentMetadata = { imagePath: 'previous.png' };
const old = load('old.png');
assert.equal(manager.currentMetadata, null);
const current = load('current.png');
requests.get('current.png').resolve({ imagePath: 'current.png' }); await current;
requests.get('old.png').reject(new Error('old error')); await old;
assert.equal(panel.innerHTML, 'current.png');
});
for (const close of ['close gallery', 'detach sidebar']) {
test(`${close} prevents pending admin metadata from restoring copy state`, async () => {
const { manager, panel, requests, load } = fixture();
const pending = load('old.png');
if (close === 'close gallery') manager.closeGallery();
else panel.isConnected = false;
requests.get('old.png').resolve({ imagePath: 'old.png' }); await pending;
assert.equal(manager.currentMetadata, null);
assert.notEqual(panel.innerHTML, 'old.png');
});
}
+199
View File
@@ -0,0 +1,199 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { createApiClient, resolveUrl, fetchUrl, ApiError, latestOnly, STALE } = require("../../web/js/api-client.js");
test("Prompt Manager URLs follow the page's proxy mount without losing query strings", () => {
const url = "/prompt_manager/images/output?limit=100&subfolder=2025%2F04%2F02";
assert.equal(resolveUrl(url, "/ComfyBackendDirect/prompt_manager/gallery"),
"/ComfyBackendDirect" + url);
assert.equal(resolveUrl(url, "/prompt_manager/gallery"), url);
assert.equal(resolveUrl(url, "/nested/proxy/prompt_manager/admin"), "/nested/proxy" + url);
assert.equal(resolveUrl("/other/api?x=1", "/ComfyBackendDirect/prompt_manager/gallery"), "/other/api?x=1");
assert.equal(resolveUrl("/ComfyBackendDirect/prompt_manager/images/serve/a.png", "/ComfyBackendDirect/prompt_manager/gallery"),
"/ComfyBackendDirect/prompt_manager/images/serve/a.png");
assert.equal(resolveUrl("/prompt_manager/images/serve/a.png?root=1", "/ComfyBackendDirect/prompt_manager/gallery"),
"/ComfyBackendDirect/prompt_manager/images/serve/a.png?root=1");
});
test("fetchUrl sends Prompt Manager requests under the current proxy mount", async () => {
const previousFetch = globalThis.fetch;
const previousLocation = globalThis.location;
const calls = [];
globalThis.location = { pathname: "/ComfyBackendDirect/prompt_manager/admin" };
globalThis.fetch = async (url, init) => {
calls.push([url, init]);
return { ok: true };
};
try {
await fetchUrl("/prompt_manager/settings?show=all", { method: "GET" });
assert.deepEqual(calls, [["/ComfyBackendDirect/prompt_manager/settings?show=all", { method: "GET" }]]);
} finally {
globalThis.fetch = previousFetch;
globalThis.location = previousLocation;
}
});
/** A fetch stand-in whose responses the test resolves by hand, in any order. */
function fakeFetch() {
const calls = [];
const impl = (url, init = {}) =>
new Promise((resolve, reject) => {
const call = { url, init, resolve, reject, aborted: false };
if (init.signal) {
init.signal.addEventListener("abort", () => {
call.aborted = true;
reject(new DOMException("The operation was aborted.", "AbortError"));
});
}
calls.push(call);
});
impl.calls = calls;
return impl;
}
function jsonResponse(status, body, { invalid = false } = {}) {
return {
ok: status >= 200 && status < 300,
status,
statusText: status === 200 ? "OK" : `Status ${status}`,
json: async () => {
if (invalid) throw new SyntaxError("Unexpected token");
return body;
},
};
}
test("get returns the parsed body of a successful response", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const pending = api.get("/prompt_manager/recent?page=1");
fetch.calls[0].resolve(jsonResponse(200, { success: true, results: [1, 2] }));
assert.deepEqual(await pending, { success: true, results: [1, 2] });
assert.equal(fetch.calls[0].init.method, "GET");
});
test("get forwards an AbortSignal to fetch", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const controller = new AbortController();
const pending = api.get("/x", { signal: controller.signal });
assert.equal(fetch.calls[0].init.signal, controller.signal);
controller.abort();
await assert.rejects(pending, (err) => err.name === "AbortError");
});
test("non-2xx responses throw ApiError carrying the status and server message", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const pending = api.get("/x");
fetch.calls[0].resolve(jsonResponse(404, { success: false, error: "Prompt not found" }));
await assert.rejects(pending, (err) => {
assert.ok(err instanceof ApiError);
assert.ok(err instanceof Error);
assert.equal(err.status, 404);
assert.equal(err.message, "Prompt not found");
return true;
});
});
test("non-2xx responses without a JSON body fall back to the status text", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const pending = api.get("/x");
fetch.calls[0].resolve(jsonResponse(500, null, { invalid: true }));
await assert.rejects(pending, (err) => err.status === 500 && err.message === "Status 500");
});
test("a 2xx body with success:false throws ApiError", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const pending = api.post("/x", {});
fetch.calls[0].resolve(jsonResponse(200, { success: false, error: "Nope" }));
await assert.rejects(pending, (err) => err instanceof ApiError && err.status === 200 && err.message === "Nope");
});
test("a 2xx body that is not JSON throws ApiError", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const pending = api.get("/x");
fetch.calls[0].resolve(jsonResponse(200, null, { invalid: true }));
await assert.rejects(pending, (err) => err instanceof ApiError && err.status === 200);
});
test("post, put and del send JSON bodies with the right method", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const p1 = api.post("/a", { text: "x" });
const p2 = api.put("/b", { rating: 3 });
const p3 = api.del("/c");
assert.equal(fetch.calls[0].init.method, "POST");
assert.equal(fetch.calls[0].init.body, JSON.stringify({ text: "x" }));
assert.equal(fetch.calls[0].init.headers["Content-Type"], "application/json");
assert.equal(fetch.calls[1].init.method, "PUT");
assert.equal(fetch.calls[1].init.body, JSON.stringify({ rating: 3 }));
assert.equal(fetch.calls[2].init.method, "DELETE");
assert.equal(fetch.calls[2].init.body, undefined);
for (const c of fetch.calls) c.resolve(jsonResponse(200, { success: true }));
await Promise.all([p1, p2, p3]);
});
test("network failures propagate as ordinary errors", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const pending = api.get("/x");
fetch.calls[0].reject(new TypeError("Failed to fetch"));
await assert.rejects(pending, (err) => err instanceof TypeError && !(err instanceof ApiError));
});
test("createApiClient without an argument uses the global fetch", async () => {
const original = globalThis.fetch;
globalThis.fetch = async () => jsonResponse(200, { success: true, via: "global" });
try {
assert.deepEqual(await createApiClient().get("/x"), { success: true, via: "global" });
} finally {
globalThis.fetch = original;
}
});
test("latestOnly: a newer call aborts the older one and the older result is STALE", async () => {
const fetch = fakeFetch();
const api = createApiClient(fetch);
const latest = latestOnly();
const first = latest((signal) => api.get("/search?q=a", { signal }));
const second = latest((signal) => api.get("/search?q=ab", { signal }));
assert.equal(fetch.calls[0].aborted, true, "first request aborted");
assert.equal(fetch.calls[1].aborted, false);
fetch.calls[1].resolve(jsonResponse(200, { success: true, q: "ab" }));
assert.deepEqual(await second, { success: true, q: "ab" });
assert.equal(await first, STALE);
});
test("latestOnly: an older response that arrives after a newer request is STALE even without abort", async () => {
const latest = latestOnly();
let resolveFirst;
const first = latest(() => new Promise((r) => { resolveFirst = r; }));
const second = latest(async () => "new");
assert.equal(await second, "new");
resolveFirst("old");
assert.equal(await first, STALE);
});
test("latestOnly: sequential calls each return their own result", async () => {
const latest = latestOnly();
assert.equal(await latest(async () => 1), 1);
assert.equal(await latest(async () => 2), 2);
});
test("latestOnly: errors from the current call propagate, errors from superseded calls do not", async () => {
const latest = latestOnly();
let rejectFirst;
const first = latest(() => new Promise((_, rej) => { rejectFirst = rej; }));
const second = latest(async () => { throw new Error("boom"); });
await assert.rejects(second, /boom/);
rejectFirst(new Error("old failure"));
assert.equal(await first, STALE);
});
+74
View File
@@ -0,0 +1,74 @@
// Run with: node --test tests/js/
// Each page module is a UMD-style IIFE: CommonJS export under Node, a window
// global in the browser. This loads every module the way a browser would and
// checks the global name the pages rely on.
const test = require("node:test");
const assert = require("node:assert/strict");
const fs = require("node:fs");
const path = require("node:path");
const vm = require("node:vm");
const MODULES = {
"html-escape.js": ["HtmlEscape", "escapeHtml", "escapeAttr"],
"data-actions.js": ["DataActions"],
"api-client.js": ["ApiClient"],
"list-state.js": ["ListState"],
"sse-stream.js": ["SseStream"],
"png-metadata.js": ["PngMetadata"],
"view-router.js": ["ViewRouter"],
"prompt-list-sort.js": ["PromptListSort"],
"comfy-metadata.js": ["ComfyMetadata"],
};
function loadInFakeBrowser(file) {
const source = fs.readFileSync(path.join(__dirname, "../../web/js", file), "utf8");
const window = {};
window.window = window;
const context = vm.createContext({
window,
globalThis: window,
console,
TextDecoder,
TextEncoder,
AbortController,
URLSearchParams,
Symbol,
Object,
Array,
Number,
String,
Math,
Set,
Map,
JSON,
Error,
Promise,
DataView,
Uint8Array,
ArrayBuffer,
});
vm.runInContext(source, context, { filename: file });
return window;
}
for (const [file, globals] of Object.entries(MODULES)) {
test(`${file} exposes ${globals.join(", ")} on window when module is undefined`, () => {
const window = loadInFakeBrowser(file);
for (const name of globals) {
assert.ok(window[name], `${name} missing`);
assert.ok(["object", "function"].includes(typeof window[name]), `${name} is a ${typeof window[name]}`);
}
});
}
test("the browser escapeHtml global is the same function as HtmlEscape.escapeHtml", () => {
const window = loadInFakeBrowser("html-escape.js");
assert.equal(window.escapeHtml, window.HtmlEscape.escapeHtml);
assert.equal(window.escapeAttr, window.escapeHtml);
assert.equal(window.escapeHtml(`<a href="x">`), "&lt;a href=&quot;x&quot;&gt;");
});
test("the browser ViewRouter global resolves the tags deep link", () => {
const window = loadInFakeBrowser("view-router.js");
assert.equal(window.ViewRouter.resolveView("#/tags/a"), "tags");
});
+49
View File
@@ -0,0 +1,49 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { readAction, dispatch } = require("../../web/js/data-actions.js");
test("readAction parses promptId and index as integers", () => {
const parsed = readAction({ action: "open-film", promptId: "42", index: "3" });
assert.equal(parsed.action, "open-film");
assert.equal(parsed.promptId, 42);
assert.equal(parsed.index, 3);
});
test("readAction turns missing or non-numeric ids into null", () => {
assert.equal(readAction({ action: "edit" }).promptId, null);
assert.equal(readAction({ action: "edit", promptId: "abc" }).promptId, null);
assert.equal(readAction({ action: "edit", index: "" }).index, null);
});
test("readAction passes string fields through untouched", () => {
const parsed = readAction({ action: "remove-tag", tag: `a "quoted" tag`, url: "/x?y=1&z=2", copyText: "p" });
assert.equal(parsed.tag, `a "quoted" tag`);
assert.equal(parsed.url, "/x?y=1&z=2");
assert.equal(parsed.copyText, "p");
});
test("readAction tolerates a missing dataset", () => {
assert.equal(readAction(undefined).action, "");
assert.equal(readAction(null).promptId, null);
});
test("dispatch calls the matching handler with the parsed action and extras", () => {
const calls = [];
const handled = dispatch(
{ action: "edit", promptId: "7" },
{ edit: (a, el) => calls.push([a.promptId, el]) },
"element",
);
assert.equal(handled, true);
assert.deepEqual(calls, [[7, "element"]]);
});
test("dispatch returns false and calls nothing for an unknown or empty action", () => {
let called = 0;
const handlers = { edit: () => called++ };
assert.equal(dispatch({ action: "other" }, handlers), false);
assert.equal(dispatch({}, handlers), false);
assert.equal(dispatch({ action: "constructor" }, handlers), false, "no prototype lookups");
assert.equal(called, 0);
});
+76
View File
@@ -0,0 +1,76 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const vm = require('node:vm');
const path = require('node:path');
function deferred() {
let resolve, reject;
const promise = new Promise((yes, no) => { resolve = yes; reject = no; });
return { promise, resolve, reject };
}
function fixture() {
const context = vm.createContext({ document: { addEventListener() {}, getElementById() {} },
window: {}, console: { error() {} }, escapeHtml: value => value });
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
const manager = Object.create(Manager.prototype);
const requests = new Map();
manager.extractImageMetadata = url => {
const pending = deferred();
requests.set(url, pending);
return pending.promise;
};
manager.displayMetadata = (metadata, panel) => {
panel.innerHTML = metadata?.positivePrompt || 'No metadata';
manager.currentMetadata = metadata;
};
manager.stopViewerSlideshow = () => {};
return { manager, requests, panel: { innerHTML: '', isConnected: true } };
}
test('a late response cannot overwrite the selected image or copy state', async () => {
const { manager, requests, panel } = fixture();
const first = manager.loadImageMetadata('old.png', panel);
const second = manager.loadImageMetadata('new.png', panel);
requests.get('new.png').resolve({ positivePrompt: 'new prompt' });
await second;
requests.get('old.png').resolve({ positivePrompt: 'old prompt' });
await first;
assert.equal(panel.innerHTML, 'new prompt');
assert.equal(manager.currentMetadata.positivePrompt, 'new prompt');
});
test('navigation clears old action data while the new request is pending', async () => {
const { manager, requests, panel } = fixture();
manager.currentMetadata = { positivePrompt: 'old prompt' };
const pending = manager.loadImageMetadata('new.png', panel);
assert.equal(manager.currentMetadata, null);
requests.get('new.png').resolve(null);
await pending;
assert.equal(manager.currentMetadata, null);
});
test('a stale error cannot replace newer metadata', async () => {
const { manager, requests, panel } = fixture();
const first = manager.loadImageMetadata('old.png', panel);
const second = manager.loadImageMetadata('new.png', panel);
requests.get('new.png').resolve({ positivePrompt: 'new prompt' });
await second;
requests.get('old.png').reject(new Error('old request failed'));
await first;
assert.equal(panel.innerHTML, 'new prompt');
});
for (const close of ['destroy viewer', 'remove modal']) {
test(`closing via ${close} prevents pending metadata from restoring action data`, async () => {
const { manager, requests, panel } = fixture();
const pending = manager.loadImageMetadata('old.png', panel);
if (close === 'destroy viewer') manager.destroyViewer();
else panel.isConnected = false;
requests.get('old.png').resolve({ positivePrompt: 'closed prompt' });
await pending;
assert.equal(manager.currentMetadata, null);
assert.notEqual(panel.innerHTML, 'closed prompt');
});
}
+209
View File
@@ -0,0 +1,209 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const vm = require('node:vm');
const path = require('node:path');
const { escapeHtml } = require('../../web/js/html-escape.js');
function fixture(count = 2000) {
const elements = new Map();
function element() {
return {
dataset: {}, children: [], listeners: {}, classList: { add() {}, remove() {} },
set innerHTML(html) { this.html = html; this.children = [{ html }]; },
get innerHTML() { return this.html || ''; },
insertAdjacentHTML(position, html) { this.children.push({ html }); },
addEventListener(type, handler) { (this.listeners[type] ||= []).push(handler); },
querySelectorAll() { return []; },
remove() {},
};
}
const document = {
addEventListener() {},
getElementById(id) { if (!elements.has(id)) elements.set(id, element()); return elements.get(id); },
createElement: element,
};
const instances = [];
const timers = new Map();
let timerId = 0;
class Viewer {
constructor(holder, options) {
this.element = holder; this.options = { fullscreen: true, ...options };
this.index = options.initialViewIndex; this.isShown = true;
this.length = (holder.innerHTML.match(/<img /g) || []).length;
instances.push(this);
}
requestFullscreen() { this.fullscreenRequested = true; }
exitFullscreen() { this.fullscreenExited = true; }
show() { this.shown = true; }
view(index) { this.index = index; }
clearEnforceFocus() { this.focusCleared = true; }
destroy() { this.destroyed = true; }
}
const context = vm.createContext({ document, window: {}, Viewer, escapeHtml,
console, setTimeout(callback) { timers.set(++timerId, callback); return timerId; },
clearTimeout(id) { timers.delete(id); },
ApiClient: { resolveUrl: value => value, STALE: Symbol() },
ListState: { nextPageState: (state, next) => ({ ...state, ...next }) },
});
const source = fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8');
const Manager = vm.runInContext(source + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
const manager = Object.create(Manager.prototype);
Object.assign(manager, { images: Array.from({ length: count }, (_, i) => ({
filename: `${i}.png`, url: `/original/${i}.png`, thumbnail_url: `/thumb/${i}.jpg`,
relative_path: `${i}.png`, size: 100, modified_time: 1,
})), viewMode: 'grid', viewer: null, currentPage: 1, limit: 100, total: count,
getSettings: () => ({ showImageInfo: true }), });
return { manager, document, instances, timers };
}
for (const mode of ['grid', 'list']) {
test(`${mode} append retains existing cards and adds only the new page`, () => {
const { manager, document, timers } = fixture(100);
manager.viewMode = mode;
manager.renderGallery();
const container = document.getElementById(mode === 'grid' ? 'galleryGrid' : 'galleryList');
assert.equal(timers.size, 0);
const first = container.children[0];
manager.images.push({ ...manager.images[0], filename: 'new.png' });
manager.renderGallery({ from: 100 });
assert.equal(container.children[0], first);
assert.equal(container.children.length, 2);
assert.match(container.children[1].html, /new.png/);
assert.doesNotMatch(container.children[1].html, /alt="0.png"/);
assert.match(container.children[1].html, /data-image-index="100"/);
});
}
test('loadImages passes the previous length for append, and zero for replacement', async () => {
const { manager } = fixture(200);
const offsets = [];
Object.assign(manager, { imagesRequest: async () => ({ total: 500, images: [manager.images[0]] }),
renderGallery: options => offsets.push(options.from), updateStats() {},
updatePagination() {}, maybeLoadMore() {}, showLoading() {}, });
await manager.loadImages(3, { append: true });
await manager.loadImages(1);
assert.deepEqual(offsets, [200, 0]);
assert.equal(manager.images.length, 1);
});
test('viewer opening is bounded and selects the clicked original', () => {
const { manager, instances } = fixture();
manager.openViewer(1500);
const viewer = instances[0];
assert.ok(viewer.element.innerHTML.match(/<img /g).length <= 100);
assert.match(viewer.element.innerHTML, /data-original="\/original\/1500.png"/);
const urls = [...viewer.element.innerHTML.matchAll(/data-original="([^"]+)"/g)].map(m => m[1]);
assert.equal(urls[viewer.options.initialViewIndex], '/original/1500.png');
assert.equal(viewer.shown, true);
});
test('previous/next cross window edges and wrap only at gallery edges', () => {
const { manager, instances } = fixture();
manager.openViewer(1500);
const first = manager.viewer;
first.index = 99;
const lastUrl = [...first.element.innerHTML.matchAll(/data-original="\/original\/(\d+).png"/g)].at(-1);
first.next(true);
const next = manager.viewer;
const urls = [...next.element.innerHTML.matchAll(/data-original="\/original\/(\d+).png"/g)];
assert.equal(Number(urls[next.index][1]), Number(lastUrl[1]) + 1);
manager.openViewer(0);
manager.viewer.prev(true);
const wrapped = manager.viewer;
const wrappedUrls = [...wrapped.element.innerHTML.matchAll(/data-original="\/original\/(\d+).png"/g)];
assert.equal(Number(wrappedUrls[wrapped.index][1]), 1999);
assert.ok(instances.length >= 3);
});
test('delegated card clicks bind once, open appended indices, and preserve video dispatch', () => {
const { manager, document } = fixture();
const opened = [];
manager.openViewer = index => opened.push(index);
manager.showVideoModal = (...args) => opened.push(args);
manager.addFallbackClickHandlers();
manager.addFallbackClickHandlers();
const grid = document.getElementById('galleryGrid');
assert.equal(grid.listeners.click.length, 1);
const item = { dataset: { imageIndex: '1500' }, querySelector: () => ({ dataset: {
original: '/original/1500.png', caption: 'caption', isVideo: 'false',
} }) };
const click = () => grid.listeners.click[0]({ target: { closest: () => item }, preventDefault() {}, stopPropagation() {} });
click();
item.querySelector = () => ({ dataset: { original: '/video.mp4', caption: 'video', isVideo: 'true' } });
click();
assert.deepEqual(opened, [1500, ['/video.mp4', 'video', 1500]]);
});
test('navigation inside the window reuses the viewer and honors non-looping edges', () => {
const { manager, instances } = fixture();
manager.openViewer(1500);
const viewer = manager.viewer;
viewer.next(true);
assert.equal(manager.viewer, viewer);
assert.equal(viewer.index, 51);
viewer.prev(true);
assert.equal(viewer.index, 50);
manager.openViewer(0);
const edge = manager.viewer;
edge.prev(false);
assert.equal(manager.viewer, edge);
assert.equal(edge.index, 0);
assert.equal(instances.length, 2);
});
test('append leaves an open viewer intact; a full render destroys it', () => {
const { manager } = fixture();
manager.openViewer(1500);
const viewer = manager.viewer;
manager.renderGallery({ from: 1900 });
assert.equal(manager.viewer, viewer);
manager.renderGallery();
assert.equal(viewer.destroyed, true);
assert.equal(manager.viewer, null);
});
test('each viewed image updates sidebar metadata without waiting for DOM mutations', () => {
const { manager, document } = fixture();
const urls = [];
document.querySelector = () => ({});
const sidebar = document.getElementById('metadata-sidebar');
sidebar.querySelector = () => ({});
manager.loadImageMetadata = url => urls.push(url);
manager.openViewer(1500);
manager.viewer.options.viewed({ detail: { originalImage: { dataset: { original: '/original/1500.png' } } } });
manager.viewer.options.viewed({ detail: { originalImage: { dataset: { original: '/original/1501.png' } } } });
assert.deepEqual(urls, ['/original/1500.png', '/original/1501.png']);
});
test('slideshow advances across windows without preloading every original and stops on hide', () => {
const { manager, timers } = fixture();
manager.openViewer(1500);
let prevented = false;
manager.viewer.options.play({ preventDefault() { prevented = true; } });
assert.equal(prevented, true);
assert.equal(manager.viewer.fullscreenRequested, true);
assert.equal(timers.size, 1);
manager.viewer.index = 99;
const tick = timers.values().next().value;
timers.clear();
tick();
const viewer = manager.viewer;
const urls = [...viewer.element.innerHTML.matchAll(/data-original="([^"]+)"/g)];
assert.equal(urls[viewer.index][1], '/original/1550.png');
assert.equal(timers.size, 1);
viewer.options.hide();
assert.equal(viewer.fullscreenExited, true);
assert.equal(timers.size, 0);
});
test('destroying a viewer removes its focus trap and slideshow timer', () => {
const { manager, timers } = fixture();
manager.openViewer(1500);
const viewer = manager.viewer;
viewer.options.play({ preventDefault() {} });
manager.destroyViewer();
assert.equal(viewer.focusCleared, true);
assert.equal(viewer.destroyed, true);
assert.equal(timers.size, 0);
});
+88
View File
@@ -0,0 +1,88 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const vm = require('node:vm');
const ImageHelpers = require('../../web/js/image-helpers.js');
function deferred() {
let resolve;
const promise = new Promise(yes => { resolve = yes; });
return { promise, resolve };
}
function fixture() {
const elements = new Map();
const requests = [];
const document = { addEventListener() {}, getElementById(id) {
if (!elements.has(id)) elements.set(id, {});
return elements.get(id);
} };
const context = vm.createContext({ document, window: {}, ImageHelpers,
escapeHtml: value => value, console: { error() {} },
ApiClient: { resolveUrl: value => value, fetchUrl: (url, options) => {
const request = deferred();
requests.push({ url, body: JSON.parse(options.body), ...request });
return request.promise;
} },
});
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
const manager = Object.create(Manager.prototype);
const notifications = [];
Object.assign(manager, { autoTagState: { reviewImages: [
{ id: 'digest', path: 'same.png', root_index: 1, url: '/same.png' },
{ id: 'other', path: 'next.png', root_index: 0, url: '/next.png' },
], reviewIndex: 0, currentTags: [], modelType: 'gguf' },
renderReviewTags() {}, showNotification: (...args) => notifications.push(args), hideModal() {}, loadImages() {},
});
const reply = (index, data) => requests[index].resolve({ json: async () => data });
return { manager, requests, reply, notifications };
}
test('gallery review retains root identity and applies tags with the returned prompt ID', async () => {
const { manager, requests, reply } = fixture();
const load = manager.loadNextReviewImage();
assert.equal(requests[0].body.root_index, 1);
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
await load;
manager.loadNextReviewImage = () => {};
const apply = manager.applyReviewTags();
assert.deepEqual(requests[1].body, { prompt_id: 42, tags: ['smile'] });
reply(1, { success: true });
await apply;
});
test('unlinked images do not submit an invalid apply request or advance', async () => {
const { manager, requests, reply } = fixture();
const load = manager.loadNextReviewImage();
reply(0, { success: true, prompt_id: null, tags: ['smile'] });
await load;
const apply = manager.applyReviewTags();
if (requests[1]) reply(1, { success: false, error: "Unexpected apply" });
await apply;
assert.equal(requests.length, 1);
assert.equal(manager.autoTagState.reviewIndex, 0);
});
test('a failed apply keeps the current image available for retry', async () => {
const { manager, requests, reply } = fixture();
const load = manager.loadNextReviewImage();
reply(0, { success: true, prompt_id: 42, tags: ['smile'] });
await load;
const apply = manager.applyReviewTags();
reply(1, { success: false, error: 'Try again' });
await apply;
assert.equal(manager.autoTagState.reviewIndex, 0);
assert.equal(requests.length, 2);
});
test('late generation responses cannot restore tags from a skipped or closed image', async () => {
const { manager, reply } = fixture();
const first = manager.loadNextReviewImage();
manager.autoTagState.reviewIndex = 1;
const second = manager.loadNextReviewImage();
reply(1, { success: true, prompt_id: 2, tags: ['current'] });
await second;
reply(0, { success: true, prompt_id: 1, tags: ['stale'] });
await first;
assert.deepEqual(Array.from(manager.autoTagState.currentTags), ['current']);
const third = manager.loadNextReviewImage();
manager.cancelReview();
reply(2, { success: true, prompt_id: 2, tags: ['closed'] });
await third;
assert.equal(manager.autoTagState.currentTags.length, 0);
});
+30
View File
@@ -0,0 +1,30 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const vm = require('node:vm');
const path = require('node:path');
const ImageHelpers = require('../../web/js/image-helpers.js');
test('gallery metadata requests retain root identity without including cache parameters in filenames', async () => {
const requests = [];
const context = vm.createContext({ document: { addEventListener() {} }, window: {}, console,
ImageHelpers, ApiClient: { fetchUrl: async url => {
requests.push(url);
return { ok: true, json: async () => ({ success: true, prompt: { text: url } }) };
} },
});
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
const manager = Object.create(Manager.prototype);
manager.extractImageMetadata = async () => null;
manager.displayMetadata = (metadata, panel) => { panel.prompt = metadata.positivePrompt; };
for (const root of [0, 1]) {
const panel = {};
const expected = `/prompt_manager/images/prompt/sub%2Fsame.png?root=${root}`;
await manager.loadImageMetadata(`/prompt_manager/images/serve/sub/same.png?v=123&root=${root}`, panel);
assert.equal(panel.prompt, expected);
}
assert.deepEqual(requests, [
'/prompt_manager/images/prompt/sub%2Fsame.png?root=0',
'/prompt_manager/images/prompt/sub%2Fsame.png?root=1',
]);
});
+63
View File
@@ -0,0 +1,63 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const {
DEFAULT_SETTINGS,
normalizeGallerySettings,
gridColumnsStyle,
LIMIT_OPTIONS,
} = require("../../web/js/gallery-settings.js");
test("defaults are a complete, sane settings object", () => {
const s = normalizeGallerySettings(null);
assert.deepEqual(s, DEFAULT_SETTINGS);
assert.equal(s.defaultViewMode, "grid");
assert.equal(s.defaultLimit, 100);
assert.equal(s.gridColumns, 8);
assert.equal(s.showImageInfo, true);
assert.equal(s.sidebarCollapsedByDefault, false);
});
test("stored values override defaults and unknown keys are kept", () => {
const s = normalizeGallerySettings({ defaultViewMode: "list", defaultLimit: 500, gridColumns: 10, showImageInfo: false, videoLoop: false, extra: 1 });
assert.equal(s.defaultViewMode, "list");
assert.equal(s.defaultLimit, 500);
assert.equal(s.gridColumns, 10);
assert.equal(s.showImageInfo, false);
assert.equal(s.videoLoop, false);
assert.equal(s.extra, 1);
});
test("invalid stored values fall back to defaults instead of breaking the page", () => {
const s = normalizeGallerySettings({ defaultViewMode: "mosaic", defaultLimit: 7, gridColumns: "lots", showImageInfo: "no" });
assert.equal(s.defaultViewMode, "grid");
assert.equal(s.defaultLimit, 100);
assert.equal(s.gridColumns, 8);
assert.equal(s.showImageInfo, true);
});
test("numeric strings from <select> values are accepted and columns are clamped", () => {
assert.equal(normalizeGallerySettings({ defaultLimit: "200" }).defaultLimit, 200);
assert.equal(normalizeGallerySettings({ gridColumns: "10" }).gridColumns, 10);
assert.equal(normalizeGallerySettings({ gridColumns: 40 }).gridColumns, 12);
assert.equal(normalizeGallerySettings({ gridColumns: 1 }).gridColumns, 2);
assert.deepEqual(LIMIT_OPTIONS, [50, 100, 200, 500]);
});
test("normalize never mutates its input", () => {
const raw = { gridColumns: 4 };
normalizeGallerySettings(raw);
assert.deepEqual(raw, { gridColumns: 4 });
});
test("gridColumnsStyle produces an explicit grid template", () => {
assert.equal(gridColumnsStyle(10), "repeat(10, minmax(0, 1fr))");
assert.equal(gridColumnsStyle("6"), "repeat(6, minmax(0, 1fr))");
assert.equal(gridColumnsStyle("x"), "repeat(8, minmax(0, 1fr))");
});
test("infinite scrolling is off by default and must be a boolean", () => {
assert.equal(DEFAULT_SETTINGS.infiniteScroll, false);
assert.equal(normalizeGallerySettings({ infiniteScroll: true }).infiniteScroll, true);
assert.equal(normalizeGallerySettings({ infiniteScroll: "yes" }).infiniteScroll, false);
});
+90
View File
@@ -0,0 +1,90 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const vm = require('node:vm');
function fixture(status, settings = {}) {
const timers = [];
const prompts = [];
const requests = [];
const buttons = {};
const storage = new Map();
const document = {
addEventListener() {},
createElement: () => ({ addEventListener() {} }),
body: { appendChild() {}, removeChild() {} },
getElementById: id => buttons[id] ||= { addEventListener(type, handler) { this[type] = handler; } },
};
const context = vm.createContext({ document, window: {}, console: { debug() {}, warn() {} },
setTimeout: callback => timers.push(callback),
localStorage: { setItem: (key, value) => storage.set(key, value) },
ApiClient: { fetchUrl: async () => ({ json: async () => ({ success: true,
total: 10000, images: Array.from({ length: 50 }, () => ({ thumbnail_url: null })),
}) }) },
});
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/gallery.js'), 'utf8') + '\nGalleryManager;', context, { filename: path.join(__dirname, '../../web/js/gallery.js') });
const manager = Object.create(Manager.prototype);
Object.assign(manager, { getSettings: () => ({ checkThumbnailsAtStartup: true, ...settings }),
api: { get: async url => { requests.push(url); return status; } },
});
const originalPrompt = manager.showThumbnailGenerationPrompt.bind(manager);
manager.showThumbnailGenerationPrompt = (...args) => prompts.push(args);
async function check() {
await manager.checkThumbnailsAtStartup();
await timers.shift()();
}
return { manager, check, prompts, requests, settings, buttons, storage, originalPrompt };
}
test('startup uses exact collection counts instead of extrapolating the newest 50 images', async () => {
const f = fixture({ success: true, total_images: 10000, pending: 50, ready: 9950, collection_id: 'a' });
await f.check();
assert.deepEqual(f.requests, ['/prompt_manager/images/thumbnails/status']);
assert.equal(f.prompts.length, 0);
});
test('startup prompt reports the actual number needing generation', async () => {
const f = fixture({ success: true, total_images: 10000, pending: 1200, collection_id: 'a' });
await f.check();
assert.equal(f.prompts.length, 1);
assert.deepEqual(f.prompts[0].slice(0, 2), [10000, 1200]);
});
test('a complete collection never prompts again', async () => {
const f = fixture({ success: true, total_images: 10000, pending: 0, collection_id: 'a' });
await f.check();
await f.check();
assert.equal(f.prompts.length, 0);
});
test('Not Now persists the collection state and suppresses the unchanged prompt', async () => {
const status = { success: true, total_images: 100, pending: 30, collection_id: 'a' };
const f = fixture(status);
await f.check();
f.originalPrompt(...f.prompts[0]);
f.buttons.thumbnailPromptNo.click();
Object.assign(f.settings, JSON.parse(f.storage.get('gallerySettings')));
f.prompts.length = 0;
await f.check();
assert.equal(f.prompts.length, 0);
status.pending = 40;
await f.check();
assert.equal(f.prompts.length, 1);
});
test('disabled checks do not request collection status', async () => {
const f = fixture({}, { checkThumbnailsAtStartup: false });
await f.check();
assert.equal(f.requests.length, 0);
assert.equal(f.prompts.length, 0);
});
test('saving gallery settings retains the remembered thumbnail dismissal', () => {
const f = fixture({}, { thumbnailPromptDismissed: 'remembered-state' });
f.manager.showNotification = () => {};
f.manager.hideSettings = () => {};
f.manager.applySettings = () => {};
f.manager.saveSettings();
assert.equal(JSON.parse(f.storage.get('gallerySettings')).thumbnailPromptDismissed, 'remembered-state');
});
+56
View File
@@ -0,0 +1,56 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { escapeHtml, escapeAttr } = require("../../web/js/html-escape.js");
test("escapes the five HTML-significant characters", () => {
assert.equal(escapeHtml(`&<>"'`), "&amp;&lt;&gt;&quot;&#39;");
});
test("a double quote cannot break out of an attribute", () => {
const payload = `" onmouseover="alert(1)`;
const escaped = escapeHtml(payload);
assert.ok(!escaped.includes(`"`), escaped);
assert.equal(escaped, `&quot; onmouseover=&quot;alert(1)`);
assert.equal(`<a data-tag="${escaped}">`, `<a data-tag="&quot; onmouseover=&quot;alert(1)">`);
});
test("a single quote cannot break out of a single-quoted attribute", () => {
const escaped = escapeHtml(`' onclick='alert(1)`);
assert.ok(!escaped.includes(`'`), escaped);
assert.equal(escaped, `&#39; onclick=&#39;alert(1)`);
});
test("a script tag is neutralised in text context", () => {
assert.equal(escapeHtml("<script>x</script>"), "&lt;script&gt;x&lt;/script&gt;");
});
test("unicode and ordinary text pass through unchanged", () => {
assert.equal(escapeHtml("héllo wörld 日本語 🙂"), "héllo wörld 日本語 🙂");
assert.equal(escapeHtml("plain text, no specials"), "plain text, no specials");
});
test("already escaped text is escaped again, never double-decoded", () => {
assert.equal(escapeHtml("&amp;"), "&amp;amp;");
});
test("null and undefined become the empty string", () => {
assert.equal(escapeHtml(null), "");
assert.equal(escapeHtml(undefined), "");
});
test("numbers and booleans are coerced to strings", () => {
assert.equal(escapeHtml(42), "42");
assert.equal(escapeHtml(0), "0");
assert.equal(escapeHtml(1.5), "1.5");
assert.equal(escapeHtml(false), "false");
});
test("objects are coerced through String()", () => {
assert.equal(escapeHtml({ toString: () => "<b>" }), "&lt;b&gt;");
});
test("escapeAttr is the same function as escapeHtml", () => {
assert.equal(escapeAttr, escapeHtml);
assert.equal(escapeAttr(`a"b`), "a&quot;b");
});
+123
View File
@@ -0,0 +1,123 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const {
formatImageCaption,
workflowDocumentHtml,
autotagSingleBody,
fallbackImageSource,
} = require("../../web/js/image-helpers.js");
const WD14 = { modelType: "wd14-vit", generalThreshold: 0.35, characterThreshold: 0.85, prompt: "ignored" };
const CAPTION = { modelType: "blip", generalThreshold: 0.35, characterThreshold: 0.85, prompt: "describe it" };
test("autotagSingleBody sends image_id for a database row", () => {
const body = autotagSingleBody({ id: 42, prompt_id: 7, image_path: "/out/a.png" }, WD14);
assert.equal(body.image_id, 42);
assert.equal("path" in body, false);
});
test("autotagSingleBody sends path for an output-folder entry whose id is a digest", () => {
const body = autotagSingleBody({ id: "3f2a9c1b0d4e5f67", path: "/out/sub/a.png", filename: "a.png" }, WD14);
assert.equal(body.path, "/out/sub/a.png");
assert.equal("image_id" in body, false);
});
test("autotagSingleBody prefers path when a numeric-looking string id has no prompt link", () => {
const body = autotagSingleBody({ id: "1234567890abcdef", path: "/out/b.png" }, WD14);
assert.deepEqual(Object.keys(body).filter((k) => k === "path" || k === "image_id"), ["path"]);
});
test("autotagSingleBody adds wd14 thresholds and no prompt", () => {
const body = autotagSingleBody({ id: 1 }, WD14);
assert.equal(body.model_type, "wd14-vit");
assert.equal(body.general_threshold, 0.35);
assert.equal(body.character_threshold, 0.85);
assert.equal("prompt" in body, false);
});
test("autotagSingleBody adds the prompt for caption models and no thresholds", () => {
const body = autotagSingleBody({ id: 1 }, CAPTION);
assert.equal(body.model_type, "blip");
assert.equal(body.prompt, "describe it");
assert.equal("general_threshold" in body, false);
});
test("autotagSingleBody does not mutate its inputs", () => {
const image = Object.freeze({ id: 1 });
const settings = Object.freeze({ ...WD14 });
assert.doesNotThrow(() => autotagSingleBody(image, settings));
});
const XSS = '"><img src=x onerror=1>';
const formatFileSize = (bytes) => `${bytes} B`;
test("formatImageCaption joins dimensions and size with the separator", () => {
const caption = formatImageCaption({ width: 640, height: 480, file_size: 2048 }, { formatFileSize });
assert.equal(caption, "640×480 | 2048 B");
});
test("formatImageCaption honours a custom separator and omits a missing size", () => {
assert.equal(formatImageCaption({ width: 1, height: 2 }, { formatFileSize, separator: " • " }), "1×2");
});
test("formatImageCaption falls back to Unknown size when a dimension is missing", () => {
assert.equal(formatImageCaption({ width: 640 }, { formatFileSize }), "Unknown size");
assert.equal(formatImageCaption({}, { formatFileSize }), "Unknown size");
assert.equal(formatImageCaption(null, { formatFileSize }), "Unknown size");
});
test("formatImageCaption escapes string dimensions so they cannot break out of an attribute", () => {
const caption = formatImageCaption({ width: XSS, height: 10 }, { formatFileSize });
assert.ok(!caption.includes("<"), caption);
assert.ok(!caption.includes('"'), caption);
assert.ok(caption.includes("&quot;&gt;&lt;img"), caption);
});
test("formatImageCaption ignores a non-numeric file_size instead of formatting it", () => {
const caption = formatImageCaption({ width: 1, height: 1, file_size: XSS }, { formatFileSize });
assert.equal(caption, "1×1");
});
test("formatImageCaption escapes whatever formatFileSize returns", () => {
const caption = formatImageCaption({ width: 1, height: 1, file_size: 5 }, { formatFileSize: () => "<b>" });
assert.equal(caption, "1×1 | &lt;b&gt;");
});
test("workflowDocumentHtml escapes the JSON so it cannot close the pre block", () => {
const html = workflowDocumentHtml({ note: "</pre><script>alert(1)</script>" });
assert.ok(html.includes("&lt;/pre&gt;&lt;script&gt;"), html);
assert.ok(!html.includes("</pre><script>"), html);
assert.ok(html.startsWith("<!DOCTYPE html>"));
assert.ok(html.includes("<title>ComfyUI Workflow Data</title>"));
});
test("workflowDocumentHtml pretty-prints the workflow with two-space indentation", () => {
const html = workflowDocumentHtml({ a: 1 });
assert.ok(html.includes('{\n &quot;a&quot;: 1\n}'), html);
});
test("fallbackImageSource offers the original once when a thumbnail fails", () => {
const dataset = { original: "/serve/a.png", thumbnail: "/serve/thumbnails/a_thumb.png?v=1" };
assert.equal(fallbackImageSource("/serve/thumbnails/a_thumb.png?v=1", dataset), "/serve/a.png");
});
test("fallbackImageSource gives up when the original itself failed or was already tried", () => {
assert.equal(fallbackImageSource("/serve/a.png", { original: "/serve/a.png" }), null);
assert.equal(fallbackImageSource("/serve/t.png", { original: "/serve/a.png", fellBack: "1" }), null);
assert.equal(fallbackImageSource("/serve/t.png", {}), null);
assert.equal(fallbackImageSource("/serve/t.png", undefined), null);
});
const { metadataLookupUrl } = require('../../web/js/image-helpers.js');
test('metadata lookup keeps root identity and strips thumbnail cache query parameters', () => {
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/sub/a.png?root=1&v=123'), '/prompt_manager/images/prompt/sub%2Fa.png?root=1');
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/a.png?root=0'), '/prompt_manager/images/prompt/a.png?root=0');
});
test('metadata lookup preserves encoded filename characters and works behind a proxy', () => {
assert.equal(metadataLookupUrl('https://example.test/ComfyBackendDirect/prompt_manager/images/serve/space%20%23%3F%2520%20caf%C3%A9.png?root=1'), '/prompt_manager/images/prompt/space%20%23%3F%2520%20caf%C3%A9.png?root=1');
assert.equal(metadataLookupUrl('/prompt_manager/images/serve/a.png?v=123'), '/prompt_manager/images/prompt/a.png');
assert.equal(metadataLookupUrl('/elsewhere/a.png'), null);
});
+65
View File
@@ -0,0 +1,65 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { nextPageState, shouldSaveEdit, hasMorePages, shouldLoadMore } = require("../../web/js/list-state.js");
test("hasMorePages follows pagination.has_more only while pages keep arriving", () => {
assert.equal(hasMorePages({ pagination: { has_more: true } }, 5000), true);
assert.equal(hasMorePages({ pagination: { has_more: false } }, 5000), false);
assert.equal(hasMorePages({ pagination: { has_more: true } }, 0), false, "an empty page ends the loop");
assert.equal(hasMorePages({}, 10), false, "no pagination block means a single page");
assert.equal(hasMorePages(null, 10), false);
});
const current = { page: 2, limit: 50, total: 120, totalPages: 3 };
test("nextPageState moves to the requested page once the fetch succeeded", () => {
assert.deepEqual(nextPageState(current, { page: 3, total: 120 }), { page: 3, limit: 50, total: 120, totalPages: 3 });
});
test("nextPageState clamps a page beyond the end when the total shrank", () => {
assert.deepEqual(nextPageState(current, { page: 3, total: 60 }), { page: 2, limit: 50, total: 60, totalPages: 2 });
});
test("nextPageState never goes below page 1 and has at least one page", () => {
assert.deepEqual(nextPageState(current, { page: 0, total: 0 }), { page: 1, limit: 50, total: 0, totalPages: 1 });
assert.equal(nextPageState(current, { page: -5, total: 10 }).page, 1);
});
test("nextPageState keeps the current page and total when the response omits them", () => {
assert.deepEqual(nextPageState(current, {}), current);
assert.deepEqual(nextPageState(current, { total: "many", page: "x" }), current);
});
test("nextPageState accepts a new limit and recomputes totalPages", () => {
assert.deepEqual(nextPageState(current, { page: 1, limit: 25, total: 120 }), { page: 1, limit: 25, total: 120, totalPages: 5 });
assert.equal(nextPageState(current, { limit: 0 }).limit, 50, "a zero limit is ignored");
});
test("nextPageState does not mutate its input", () => {
const before = { ...current };
nextPageState(current, { page: 1, total: 5 });
assert.deepEqual(current, before);
});
test("shouldSaveEdit is true only for a non-empty text that changed", () => {
assert.equal(shouldSaveEdit("old", "new"), true);
assert.equal(shouldSaveEdit("same", "same"), false);
assert.equal(shouldSaveEdit("same", " same \n"), false, "whitespace-only changes are not saved");
assert.equal(shouldSaveEdit("old", " "), false, "an emptied prompt is not saved");
assert.equal(shouldSaveEdit("old", null), false);
});
test("shouldLoadMore asks for the next page only near the bottom with pages left", () => {
const base = {
infiniteScroll: true, loading: false, page: 1, limit: 100, total: 1000,
scrollTop: 0, clientHeight: 800, scrollHeight: 4000, threshold: 600,
};
assert.equal(shouldLoadMore(base), false, "far from the bottom");
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700 }), true, "within the threshold");
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700, loading: true }), false, "a load is in flight");
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700, infiniteScroll: false }), false, "setting off");
assert.equal(shouldLoadMore({ ...base, scrollTop: 2700, page: 10 }), false, "last page reached");
assert.equal(shouldLoadMore({ ...base, scrollHeight: 700 }), true, "content shorter than the viewport");
assert.equal(shouldLoadMore({ ...base, total: 0, scrollHeight: 700 }), false, "nothing to load");
});
+102
View File
@@ -0,0 +1,102 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { parsePngTextChunks, PNG_SIGNATURE } = require("../../web/js/png-metadata.js");
const SIGNATURE = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
function u32(n) {
return [(n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff];
}
/** A chunk with a dummy CRC (the parser does not verify it). */
function chunk(type, dataBytes) {
return [...u32(dataBytes.length), ...Buffer.from(type, "latin1"), ...dataBytes, 0, 0, 0, 0];
}
const IHDR = chunk("IHDR", [...u32(1), ...u32(1), 8, 6, 0, 0, 0]);
const IEND = chunk("IEND", []);
function tEXt(keyword, text) {
return chunk("tEXt", [...Buffer.from(keyword, "latin1"), 0, ...Buffer.from(text, "latin1")]);
}
function iTXt(keyword, text, { compressed = false } = {}) {
return chunk("iTXt", [
...Buffer.from(keyword, "latin1"), 0,
compressed ? 1 : 0, 0, // compression flag, compression method
0, // empty language tag
0, // empty translated keyword
...Buffer.from(text, "utf8"),
]);
}
function png(...chunks) {
return Uint8Array.from([...SIGNATURE, ...IHDR, ...chunks.flat(), ...IEND]);
}
test("exports the PNG signature", () => {
assert.deepEqual([...PNG_SIGNATURE], SIGNATURE);
});
test("reads a tEXt chunk written by ComfyUI", () => {
const bytes = png(tEXt("prompt", '{"a":1}'));
assert.deepEqual(parsePngTextChunks(bytes), { prompt: '{"a":1}' });
});
test("reads several text chunks and an uncompressed iTXt chunk with UTF-8 text", () => {
const bytes = png(tEXt("prompt", '{"a":1}'), iTXt("workflow", '{"b":"日本語 🙂"}'), tEXt("Software", "ComfyUI"));
assert.deepEqual(parsePngTextChunks(bytes), {
prompt: '{"a":1}',
workflow: '{"b":"日本語 🙂"}',
Software: "ComfyUI",
});
});
test("tEXt is Latin-1, so byte 0xE9 is an e-acute", () => {
const bytes = png(chunk("tEXt", [...Buffer.from("Comment", "latin1"), 0, 0xe9]));
assert.deepEqual(parsePngTextChunks(bytes), { Comment: "é" });
});
test("skips compressed iTXt and zTXt payloads instead of returning garbage", () => {
const zTXt = chunk("zTXt", [...Buffer.from("Comment", "latin1"), 0, 0, 0x78, 0x9c, 1, 2, 3]);
const bytes = png(iTXt("workflow", "zzz", { compressed: true }), zTXt, tEXt("prompt", "ok"));
assert.deepEqual(parsePngTextChunks(bytes), { prompt: "ok" });
});
test("ignores non-text chunks and a text chunk without a NUL separator", () => {
const bytes = png(chunk("IDAT", [1, 2, 3]), chunk("tEXt", [...Buffer.from("nonul", "latin1")]), tEXt("k", "v"));
assert.deepEqual(parsePngTextChunks(bytes), { k: "v" });
});
test("accepts an ArrayBuffer as well as a Uint8Array", () => {
const bytes = png(tEXt("prompt", "x"));
assert.deepEqual(parsePngTextChunks(bytes.buffer.slice(bytes.byteOffset, bytes.byteOffset + bytes.byteLength)), { prompt: "x" });
});
test("returns {} for garbage, empty and non-PNG input", () => {
assert.deepEqual(parsePngTextChunks(new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])), {});
assert.deepEqual(parsePngTextChunks(new Uint8Array(0)), {});
assert.deepEqual(parsePngTextChunks(Uint8Array.from(Buffer.from("GIF89a....................", "latin1"))), {});
assert.deepEqual(parsePngTextChunks(null), {});
assert.deepEqual(parsePngTextChunks("not bytes"), {});
});
test("a file truncated inside a text chunk yields only the chunks that were complete", () => {
const full = png(tEXt("prompt", "first"), tEXt("workflow", "second-second-second"));
const cutInsideSecond = full.subarray(0, full.length - IEND.length - 6);
assert.deepEqual(parsePngTextChunks(cutInsideSecond), { prompt: "first" });
const cutInsideFirst = full.subarray(0, SIGNATURE.length + IHDR.length + 10);
assert.deepEqual(parsePngTextChunks(cutInsideFirst), {});
});
test("a chunk whose declared length overruns the file is ignored", () => {
const bogus = [...u32(0x7fffffff), ...Buffer.from("tEXt", "latin1"), ...Buffer.from("k\0v", "latin1")];
const bytes = Uint8Array.from([...SIGNATURE, ...IHDR, ...tEXt("ok", "1"), ...bogus]);
assert.deepEqual(parsePngTextChunks(bytes), { ok: "1" });
});
test("stops at IEND even if bytes follow it", () => {
const bytes = Uint8Array.from([...png(tEXt("a", "1")), ...tEXt("b", "2")]);
assert.deepEqual(parsePngTextChunks(bytes), { a: "1" });
});
+66
View File
@@ -0,0 +1,66 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const sort = require("../../web/js/prompt-list-sort.js");
test("default sort is recently used and is a listed option", () => {
assert.equal(sort.DEFAULT_SORT, "last_used_desc");
assert.ok(sort.SORT_OPTIONS.some((o) => o.value === sort.DEFAULT_SORT));
});
test("options include Recently used and Most used, each with a label", () => {
const values = sort.SORT_OPTIONS.map((o) => o.value);
assert.ok(values.includes("last_used_desc"));
assert.ok(values.includes("run_count_desc"));
for (const option of sort.SORT_OPTIONS) assert.ok(option.label.trim(), option.value);
assert.equal(new Set(values).size, values.length, "values are unique");
});
test("options cannot be mutated by callers", () => {
assert.throws(() => sort.SORT_OPTIONS.push({ value: "x", label: "x" }));
// Non-strict code fails silently on frozen objects, so check the value survived
const first = sort.SORT_OPTIONS[0].value;
sort.SORT_OPTIONS[0].value = "hacked";
assert.equal(sort.SORT_OPTIONS[0].value, first);
assert.ok(sort.SORT_OPTIONS.every(Object.isFrozen));
});
test("normalizeSort keeps known values and falls back for anything else", () => {
assert.equal(sort.normalizeSort("run_count_desc"), "run_count_desc");
for (const bad of [undefined, null, "", "created_at; DROP TABLE prompts", 42, "LAST_USED_DESC"]) {
assert.equal(sort.normalizeSort(bad), sort.DEFAULT_SORT, String(bad));
}
});
test("buildRecentUrl computes offset from page and always sends a valid sort", () => {
assert.equal(
sort.buildRecentUrl({ page: 3, limit: 25, sort: "rating_desc" }),
"/prompt_manager/recent?limit=25&offset=50&page=3&sort=rating_desc",
);
assert.equal(
sort.buildRecentUrl({ sort: "bogus" }),
"/prompt_manager/recent?limit=50&offset=0&page=1&sort=last_used_desc",
);
});
test("withSort returns a copy with the sort set and leaves the input untouched", () => {
const params = new URLSearchParams({ text: "castle", sort: "text_asc" });
const result = sort.withSort(params, "run_count_desc");
assert.equal(result.get("sort"), "run_count_desc");
assert.equal(result.get("text"), "castle");
assert.equal(params.get("sort"), "text_asc");
assert.equal(sort.withSort(new URLSearchParams(), "nope").get("sort"), sort.DEFAULT_SORT);
});
test("formatRunCount labels zero, one and many runs", () => {
assert.equal(sort.formatRunCount(0), "Never run");
assert.equal(sort.formatRunCount(1), "1 run");
assert.equal(sort.formatRunCount(642), "642 runs");
assert.equal(sort.formatRunCount(1234), "1,234 runs");
});
test("formatRunCount hides the badge for missing or invalid counts", () => {
for (const bad of [undefined, null, -1, 1.5, "3", NaN]) {
assert.equal(sort.formatRunCount(bad), "", String(bad));
}
});
+78
View File
@@ -0,0 +1,78 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const {
initialState,
reduceScanEvent,
fromStatus,
dismissAction,
progressLabel,
} = require("../../web/js/scan-progress.js");
test("initialState is idle with zero counts", () => {
assert.deepEqual(initialState(), {
status: "idle", progress: 0, statusText: "", processed: 0, found: 0, added: 0, linked: 0, message: "",
});
});
test("a progress event moves the scan to running and copies the counts", () => {
const before = initialState();
const after = reduceScanEvent(before, {
type: "progress", progress: 42, status: "Processing file 84/200...", processed: 84, found: 12,
});
assert.equal(after.status, "running");
assert.equal(after.progress, 42);
assert.equal(after.statusText, "Processing file 84/200...");
assert.equal(after.processed, 84);
assert.equal(after.found, 12);
assert.equal(before.status, "idle", "reducer must not mutate its input");
});
test("a complete event finishes the scan with the final totals", () => {
const running = reduceScanEvent(initialState(), { type: "progress", progress: 90, status: "x", processed: 9, found: 1 });
const done = reduceScanEvent(running, { type: "complete", processed: 10, found: 3, added: 2, linked: 1 });
assert.equal(done.status, "done");
assert.equal(done.progress, 100);
assert.deepEqual([done.processed, done.found, done.added, done.linked], [10, 3, 2, 1]);
});
test("an error event fails the scan and keeps the server message", () => {
const failed = reduceScanEvent(initialState(), { type: "error", message: "No output directories found." });
assert.equal(failed.status, "failed");
assert.equal(failed.message, "No output directories found.");
});
test("unknown events leave the state untouched", () => {
const state = reduceScanEvent(initialState(), { type: "progress", progress: 5, status: "s", processed: 1, found: 0 });
assert.equal(reduceScanEvent(state, { type: "mystery" }), state);
assert.equal(reduceScanEvent(state, null), state);
});
test("fromStatus rebuilds a running scan from the status endpoint", () => {
const state = fromStatus({
running: true,
last_event: { type: "progress", progress: 30, status: "Processing file 3/10...", processed: 3, found: 1 },
});
assert.equal(state.status, "running");
assert.equal(state.progress, 30);
assert.equal(state.processed, 3);
});
test("fromStatus is idle when nothing is running, even if a last event exists", () => {
assert.equal(fromStatus({ running: false, last_event: { type: "complete", processed: 10 } }).status, "idle");
assert.equal(fromStatus({ running: false, last_event: null }).status, "idle");
assert.equal(fromStatus(null).status, "idle");
});
test("dismissing the modal minimizes a running scan and closes otherwise", () => {
const running = reduceScanEvent(initialState(), { type: "progress", progress: 1, status: "s", processed: 0, found: 0 });
assert.equal(dismissAction(running), "minimize");
assert.equal(dismissAction(initialState()), "close");
assert.equal(dismissAction(reduceScanEvent(running, { type: "complete" })), "close");
});
test("progressLabel summarises a running scan for the minimized pill", () => {
const running = reduceScanEvent(initialState(), { type: "progress", progress: 42, status: "s", processed: 84, found: 12 });
assert.equal(progressLabel(running), "Scanning 42% · 84 files · 12 prompts");
assert.equal(progressLabel(initialState()), "");
});
+234
View File
@@ -0,0 +1,234 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { parseSseChunk, connect } = require("../../web/js/sse-stream.js");
test("parseSseChunk splits complete events and keeps the unfinished tail", () => {
const { events, rest } = parseSseChunk('data: {"a":1}\n\ndata: {"b":2}\n\ndata: {"c"');
assert.deepEqual(events.map((e) => e.data), ['{"a":1}', '{"b":2}']);
assert.equal(rest, 'data: {"c"');
});
test("parseSseChunk joins multi-line data, reads event names and ignores comments", () => {
const { events, rest } = parseSseChunk(": keep-alive\nevent: progress\ndata: line1\ndata: line2\nid: 7\n\n");
assert.deepEqual(events, [{ event: "progress", data: "line1\nline2" }]);
assert.equal(rest, "");
});
test("parseSseChunk accepts CRLF line endings and a missing space after the colon", () => {
const { events } = parseSseChunk("data:x\r\n\r\ndata: y\r\n\r\n");
assert.deepEqual(events.map((e) => e.data), ["x", "y"]);
});
test("parseSseChunk returns nothing for an empty or whitespace buffer", () => {
assert.deepEqual(parseSseChunk(""), { events: [], rest: "" });
assert.deepEqual(parseSseChunk("\n\n"), { events: [], rest: "" });
});
function streamOf(chunks) {
const encoder = new TextEncoder();
return new ReadableStream({
start(controller) {
for (const c of chunks) controller.enqueue(encoder.encode(c));
controller.close();
},
});
}
function until(check, timeoutMs = 500) {
return new Promise((resolve, reject) => {
const started = Date.now();
(function poll() {
if (check()) return resolve();
if (Date.now() - started > timeoutMs) return reject(new Error("timed out"));
setTimeout(poll, 2);
})();
});
}
test("connect POSTs and delivers each event to onmessage, then onclose", async () => {
const seen = [];
let closed = false;
const calls = [];
const fetchImpl = async (url, init) => {
calls.push({ url, init });
return new Response(streamOf(['data: {"n":1}\n\ndata: {"n', '":2}\n\n']), { status: 200 });
};
const source = connect("/prompt_manager/autotag/start?x=1", { method: "POST", fetchImpl });
source.onmessage = (e) => seen.push(JSON.parse(e.data).n);
source.onclose = () => { closed = true; };
await until(() => closed);
assert.deepEqual(seen, [1, 2]);
assert.equal(calls[0].init.method, "POST");
assert.equal(calls[0].url, "/prompt_manager/autotag/start?x=1");
assert.equal(source.readyState, source.CLOSED);
});
test("connect keeps streaming requests under the current proxy mount", async () => {
const previousApiClient = globalThis.ApiClient;
const previousLocation = globalThis.location;
globalThis.ApiClient = require("../../web/js/api-client.js");
globalThis.location = { pathname: "/ComfyBackendDirect/prompt_manager/gallery" };
let requestedUrl;
try {
const source = connect("/prompt_manager/images/generate-thumbnails/progress?quality=80", {
method: "POST",
fetchImpl: async (url) => {
requestedUrl = url;
return new Response(streamOf([]), { status: 200 });
},
});
await until(() => source.readyState === source.CLOSED);
assert.equal(requestedUrl, "/ComfyBackendDirect/prompt_manager/images/generate-thumbnails/progress?quality=80");
} finally {
globalThis.ApiClient = previousApiClient;
globalThis.location = previousLocation;
}
});
test("connect reports a non-2xx response through onerror and never onmessage", async () => {
let error = null;
const fetchImpl = async () => new Response("nope", { status: 405 });
const source = connect("/x", { fetchImpl });
source.onmessage = () => assert.fail("no message expected");
source.onerror = (err) => { error = err; };
await until(() => error !== null);
assert.match(error.message, /405/);
assert.equal(source.readyState, source.CLOSED);
});
test("connect reports a network failure through onerror", async () => {
let error = null;
const source = connect("/x", { fetchImpl: async () => { throw new TypeError("Failed to fetch"); } });
source.onerror = (err) => { error = err; };
await until(() => error !== null);
assert.equal(error.message, "Failed to fetch");
});
test("close() aborts the request and suppresses later callbacks", async () => {
let aborted = false;
let errors = 0;
const fetchImpl = (url, init) =>
new Promise((resolve, reject) => {
init.signal.addEventListener("abort", () => {
aborted = true;
reject(new DOMException("aborted", "AbortError"));
});
resolve(new Response(new ReadableStream({ start() {} }), { status: 200 }));
});
const source = connect("/x", { fetchImpl });
source.onerror = () => errors++;
await until(() => source.readyState === source.OPEN);
source.close();
assert.equal(aborted, true);
assert.equal(source.readyState, source.CLOSED);
await new Promise((r) => setTimeout(r, 10));
assert.equal(errors, 0);
});
test("connect forwards a request body and an Accept header", async () => {
let init = null;
const fetchImpl = async (_url, i) => { init = i; return new Response(streamOf([]), { status: 200 }); };
const source = connect("/x", { method: "POST", body: '{"a":1}', fetchImpl });
await until(() => source.readyState === source.CLOSED);
assert.equal(init.body, '{"a":1}');
assert.equal(init.headers.Accept, "text/event-stream");
});
test("events arriving with no onmessage handler are dropped without error", async () => {
const fetchImpl = async () => new Response(streamOf(["data: 1\n\n", "data: 2\n\n"]), { status: 200 });
const source = connect("/x", { fetchImpl });
let errors = 0;
source.onerror = () => errors++;
await until(() => source.readyState === source.CLOSED);
assert.equal(errors, 0);
});
test("close() during delivery stops further messages and skips onclose", async () => {
const seen = [];
let closed = false;
const fetchImpl = async () => new Response(streamOf(["data: 1\n\ndata: 2\n\ndata: 3\n\n"]), { status: 200 });
const source = connect("/x", { fetchImpl });
source.onmessage = (e) => { seen.push(e.data); if (e.data === "1") source.close(); };
source.onclose = () => { closed = true; };
await until(() => source.readyState === source.CLOSED);
await new Promise((r) => setTimeout(r, 10));
assert.deepEqual(seen, ["1"]);
assert.equal(closed, false);
});
test("close() before the response arrives discards it silently", async () => {
let resolveFetch;
let errors = 0;
const fetchImpl = () => new Promise((resolve) => { resolveFetch = resolve; });
const source = connect("/x", { fetchImpl });
source.onerror = () => errors++;
source.close();
resolveFetch(new Response(streamOf(["data: late\n\n"]), { status: 200 }));
await new Promise((r) => setTimeout(r, 10));
assert.equal(errors, 0);
assert.equal(source.readyState, source.CLOSED);
});
test("connect defaults to GET and the global fetch", async () => {
const original = globalThis.fetch;
let init = null;
globalThis.fetch = async (_url, i) => { init = i; return new Response(streamOf([]), { status: 200 }); };
try {
const source = connect("/x");
await until(() => source.readyState === source.CLOSED);
assert.equal(init.method, "GET");
} finally {
globalThis.fetch = original;
}
});
test("addEventListener receives named events and unnamed ones go to 'message'", async () => {
const named = [];
const plain = [];
let closed = false;
const fetchImpl = async () =>
new Response(
streamOf(['event: start\ndata: {"total":3}\n\n', 'data: {"n":1}\n\n', 'event: complete\ndata: {"ok":true}\n\n']),
{ status: 200 },
);
const source = connect("/prompt_manager/images/generate-thumbnails/progress", { method: "POST", fetchImpl });
source.addEventListener("start", (e) => named.push(["start", JSON.parse(e.data).total]));
source.addEventListener("complete", (e) => named.push(["complete", JSON.parse(e.data).ok]));
source.addEventListener("message", (e) => plain.push(JSON.parse(e.data).n));
source.onclose = () => { closed = true; };
await until(() => closed);
assert.deepEqual(named, [["start", 3], ["complete", true]]);
assert.deepEqual(plain, [1]);
});
test("removeEventListener stops delivery and onmessage still sees every event", async () => {
const all = [];
const progress = [];
let closed = false;
const fetchImpl = async () =>
new Response(streamOf(["event: progress\ndata: 1\n\n", "event: progress\ndata: 2\n\n"]), { status: 200 });
const source = connect("/x", { fetchImpl });
const listener = (e) => {
progress.push(e.data);
source.removeEventListener("progress", listener);
};
source.addEventListener("progress", listener);
source.onmessage = (e) => all.push(`${e.event}:${e.data}`);
source.onclose = () => { closed = true; };
await until(() => closed);
assert.deepEqual(progress, ["1"]);
assert.deepEqual(all, ["progress:1", "progress:2"]);
});
test("a listener that closes the source stops later events", async () => {
const seen = [];
const fetchImpl = async () =>
new Response(streamOf(["event: error\ndata: boom\n\n", "event: progress\ndata: 2\n\n"]), { status: 200 });
const source = connect("/x", { fetchImpl });
source.addEventListener("error", (e) => { seen.push(e.data); source.close(); });
source.addEventListener("progress", (e) => seen.push(e.data));
await until(() => source.readyState === source.CLOSED);
await new Promise((r) => setTimeout(r, 10));
assert.deepEqual(seen, ["boom"]);
});
+39
View File
@@ -0,0 +1,39 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const {
splitTagInput,
applySuggestion,
suggestUrl,
moveSelection,
} = require("../../web/js/tag-autocomplete.js");
test("splitTagInput separates committed tags from the tag being typed", () => {
assert.deepEqual(splitTagInput("asian, portrait, sm"), { committed: ["asian", "portrait"], current: "sm" });
assert.deepEqual(splitTagInput("asian,"), { committed: ["asian"], current: "" });
assert.deepEqual(splitTagInput("asian, "), { committed: ["asian"], current: "" });
assert.deepEqual(splitTagInput("asi"), { committed: [], current: "asi" });
assert.deepEqual(splitTagInput(""), { committed: [], current: "" });
assert.deepEqual(splitTagInput(" asian ,, Asian , dog, "), { committed: ["asian", "dog"], current: "" });
});
test("applySuggestion replaces the partial tag and leaves the cursor ready for the next one", () => {
assert.equal(applySuggestion("asian, sm", "smile"), "asian, smile, ");
assert.equal(applySuggestion("", "asian"), "asian, ");
assert.equal(applySuggestion("asian, ", "portrait"), "asian, portrait, ");
assert.equal(applySuggestion("asian, asi", "asian"), "asian, ");
});
test("suggestUrl encodes the prefix and the committed tags", () => {
assert.equal(suggestUrl("asian, sm"), "/prompt_manager/tags/suggest?q=sm&with=asian&limit=15");
assert.equal(suggestUrl("a&b, c d, x", 5), "/prompt_manager/tags/suggest?q=x&with=a%26b%2Cc+d&limit=5");
assert.equal(suggestUrl(""), "/prompt_manager/tags/suggest?q=&with=&limit=15");
});
test("moveSelection wraps around the list and starts from nothing selected", () => {
assert.equal(moveSelection(-1, 1, 3), 0);
assert.equal(moveSelection(-1, -1, 3), 2);
assert.equal(moveSelection(2, 1, 3), 0);
assert.equal(moveSelection(0, -1, 3), 2);
assert.equal(moveSelection(1, 1, 0), -1);
});
+71
View File
@@ -0,0 +1,71 @@
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const vm = require('node:vm');
function fixture() {
const requests = [];
const grid = { replaceChildren() {}, appendChild() {} };
const context = vm.createContext({ window: {}, URLSearchParams, console: { error() {} },
requestAnimationFrame() {}, document: { getElementById: id => id === 'tagPromptsGrid' ? grid : null,
querySelector: () => null, createElement: () => ({ appendChild() {} }) },
});
const Manager = vm.runInContext(fs.readFileSync(path.join(__dirname, '../../web/js/tags-page.js'), 'utf8') + '\nTagsPageManager;', context, { filename: path.join(__dirname, '../../web/js/tags-page.js') });
const manager = Object.create(Manager.prototype);
Object.assign(manager, { selectedTags: ['cats'], filterMode: 'and', promptLimit: 20,
promptOffset: 0, isLoadingPrompts: false, prompts: [], hasMorePrompts: false,
fetchJson(url) { return new Promise((resolve, reject) => requests.push({ url, resolve, reject })); },
renderTagPrompts() {}, appendTagPrompts() {}, updatePromptCountBadge() {}, updateLoadMoreButton() {},
});
const reply = (index, text, more = false) => requests[index].resolve({ success: true,
prompts: [{ text }], pagination: { has_more: more, total: more ? 40 : 1 } });
return { manager, requests, reply };
}
test('changing a filter starts a new request and late results cannot overwrite it', async () => {
const { manager, requests, reply } = fixture();
const old = manager.loadTagPrompts();
manager.selectedTags = ['dogs'];
const current = manager.loadTagPrompts();
assert.equal(requests.length, 2);
assert.match(requests[1].url, /dogs/);
reply(1, 'dog prompt'); await current;
reply(0, 'cat prompt', true); await old;
assert.equal(manager.prompts[0].text, 'dog prompt');
assert.equal(manager.promptTotal, 1);
assert.equal(manager.hasMorePrompts, false);
});
test('clearing selection invalidates pending results', async () => {
const { manager, reply } = fixture();
const old = manager.loadTagPrompts();
manager.selectedTags = [];
manager.clearPrompts();
reply(0, 'cat prompt'); await old;
assert.equal(manager.prompts.length, 0);
assert.equal(manager.promptTotal, 0);
assert.equal(manager.isLoadingPrompts, false);
});
test('a stale pagination response cannot append rows from the previous filter', async () => {
const { manager, requests, reply } = fixture();
Object.assign(manager, { prompts: [{ text: 'cat first page' }], hasMorePrompts: true, promptOffset: 20 });
const old = manager.loadMoreTagPrompts();
manager.selectedTags = ['dogs'];
const current = manager.loadTagPrompts();
assert.equal(requests.length, 2);
reply(1, 'dog prompt'); await current;
reply(0, 'cat second page', true); await old;
assert.equal(manager.prompts.length, 1);
assert.equal(manager.prompts[0].text, 'dog prompt');
assert.equal(manager.promptOffset, 20);
});
test('a stale failure cannot finish the new request or clear its loading state', async () => {
const { manager, requests, reply } = fixture();
const old = manager.loadTagPrompts();
manager.selectedTags = ['dogs'];
const current = manager.loadTagPrompts();
assert.equal(requests.length, 2);
requests[0].reject(new Error('old failure')); await old;
assert.equal(manager.isLoadingPrompts, true);
reply(1, 'dog prompt'); await current;
assert.equal(manager.isLoadingPrompts, false);
assert.equal(manager.prompts[0].text, 'dog prompt');
});
+30
View File
@@ -0,0 +1,30 @@
// Run with: node --test tests/js/
const test = require("node:test");
const assert = require("node:assert/strict");
const { resolveView, VIEWS } = require("../../web/js/view-router.js");
test("the tags deep link and its sub-routes resolve to the tags view", () => {
assert.equal(resolveView("#/tags"), "tags");
assert.equal(resolveView("#/tags/"), "tags");
assert.equal(resolveView("#/tags/landscape,portrait?mode=or"), "tags");
});
test("an empty, bare or unknown hash resolves to the prompts dashboard", () => {
assert.equal(resolveView(""), "prompts");
assert.equal(resolveView("#"), "prompts");
assert.equal(resolveView("#/"), "prompts");
assert.equal(resolveView("#/other"), "prompts");
assert.equal(resolveView("#/tagsX"), "prompts", "prefix must end at a path boundary");
assert.equal(resolveView("#tags"), "prompts");
});
test("non-string input resolves to the prompts dashboard", () => {
assert.equal(resolveView(undefined), "prompts");
assert.equal(resolveView(null), "prompts");
assert.equal(resolveView(42), "prompts");
});
test("every resolved value is a listed view", () => {
assert.ok(Object.isFrozen(VIEWS));
for (const hash of ["#/tags", "", "#/x"]) assert.ok(VIEWS.includes(resolveView(hash)));
});
+43
View File
@@ -0,0 +1,43 @@
"""Teardown guard: fail when a file about to be deleted is still open.
Windows refuses to delete an open file, so a sqlite connection that was not
closed explicitly turns into a PermissionError in CI there. Linux would let
the deletion through silently; this check makes the Linux run catch the same
mistake. It inspects /proc/self/fd and is a no-op where that does not exist.
"""
import os
def open_handles_under(path):
"""Basenames of this process's open files at *path* or below it."""
fd_dir = "/proc/self/fd"
if not os.path.isdir(fd_dir):
return []
real = os.path.realpath(str(path))
hits = []
for fd in os.listdir(fd_dir):
try:
target = os.readlink(os.path.join(fd_dir, fd))
except OSError:
continue
if not target.startswith(os.sep):
continue # sockets, pipes, anon inodes
target = target.replace(" (deleted)", "")
if (
target == real
or target.startswith(real + os.sep)
or target.startswith(real + "-")
):
hits.append(os.path.basename(target))
return sorted(hits)
def assert_closed(testcase, *paths):
"""Fail *testcase* if any of *paths* (files or directories) is still open."""
hits = [hit for path in paths for hit in open_handles_under(path)]
if hits:
testcase.fail(
"Files still open at teardown (Windows could not delete them): "
+ ", ".join(hits)
)
+63
View File
@@ -38,6 +38,10 @@ class APITestCase(AioHTTPTestCase):
return app
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
# Close every thread's connection first: Windows refuses to delete
# an open database file, and executor threads hold connections too.
self.api.db.close_all()
for path in (
self._temp_db.name,
self._temp_db.name + "-wal",
@@ -238,6 +242,65 @@ class TestUpdatePrompt(APITestCase):
self.assertTrue(data["success"])
class TestUsageSorting(APITestCase):
"""Server-side sort for /recent and /search (3.2.4, #90.1)."""
def _use(self, prompt_id, times=1):
for _ in range(times):
self.api.db.record_prompt_use(prompt_id)
async def _ids(self, path):
resp = await self.client.request("GET", path)
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertTrue(data["success"])
return [p["id"] for p in data["results"]]
async def test_recent_sorted_by_last_used_bubbles_rerun_prompt(self):
old = self._save_prompt("Old prompt")
self._save_prompt("New prompt")
self._use(old)
ids = await self._ids("/prompt_manager/recent?limit=1&sort=last_used_desc")
self.assertEqual(ids, [old])
async def test_recent_sorted_by_run_count(self):
# Most-run prompt is the oldest, so newest-first would give the opposite order
most = self._save_prompt("Run thrice")
least = self._save_prompt("Run once")
self._use(most, times=3)
self._use(least)
ids = await self._ids("/prompt_manager/recent?sort=run_count_desc")
self.assertEqual(ids, [most, least])
async def test_recent_default_is_newest_first(self):
first = self._save_prompt("First")
second = self._save_prompt("Second")
self._use(first)
self.assertEqual(await self._ids("/prompt_manager/recent"), [second, first])
async def test_unknown_sort_falls_back_instead_of_erroring(self):
self._save_prompt("Only")
ids = await self._ids(
"/prompt_manager/recent?sort=created_at;DROP%20TABLE%20prompts"
)
self.assertEqual(len(ids), 1)
async def test_search_respects_sort(self):
a = self._save_prompt("castle at dawn")
self._save_prompt("castle at dusk")
self._use(a)
ids = await self._ids("/prompt_manager/search?text=castle&sort=last_used_desc")
self.assertEqual(ids[0], a)
async def test_results_include_usage_fields(self):
pid = self._save_prompt("Tracked")
self._use(pid, times=2)
resp = await self.client.request("GET", "/prompt_manager/recent")
prompt = (await resp.json())["results"][0]
self.assertEqual(prompt["run_count"], 2)
self.assertTrue(prompt["last_used_at"])
class TestTags(APITestCase):
async def test_get_all_tags(self):
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
+252
View File
@@ -0,0 +1,252 @@
"""Route tests for GET /prompt_manager/backup and POST /prompt_manager/restore.
The backup must include rows still sitting in the WAL, the restore must
stream the upload to disk, reject anything that is not a healthy prompts
database and leave the live database untouched when it does.
"""
import glob
import os
import sqlite3
import sys
import tempfile
import unittest
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.models import PromptModel # noqa: E402
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
import aiohttp
from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase
from database.operations import PromptDatabase
from py.api import PromptManagerAPI, _public_path
from utils.hashing import generate_prompt_hash
def _texts(path):
conn = sqlite3.connect(path)
try:
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
finally:
conn.close()
class BackupRestoreTestCase(AioHTTPTestCase):
async def get_application(self):
self.tmpdir = tempfile.mkdtemp()
self.db_path = os.path.join(self.tmpdir, "prompts.db")
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(self.db_path)
self.api.add_routes(routes)
app.router.add_routes(routes)
return app
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
PromptModel.close_all_instances() # restore may have swapped api.db
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 _save(self, text):
return self.api.db.save_prompt(
text=text, prompt_hash=generate_prompt_hash(text)
)
def _write(self, name, data):
path = os.path.join(self.tmpdir, name)
with open(path, "wb") as fh:
fh.write(data)
return path
def _make_db_file(self, texts, name="upload.db"):
path = os.path.join(self.tmpdir, name)
db = PromptDatabase(path)
for text in texts:
db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
db.close_all()
return path
async def _restore(self, path, field="database_file"):
with open(path, "rb") as fh:
form = aiohttp.FormData()
form.add_field(
field,
fh.read(),
filename=os.path.basename(path),
content_type="application/octet-stream",
)
return await self.client.post("/prompt_manager/restore", data=form)
class TestBackupRoute(BackupRestoreTestCase):
async def test_backup_download_contains_rows_still_in_the_wal(self):
for n in range(3):
self._save(f"backup me {n}")
self.assertGreater(os.path.getsize(self.db_path + "-wal"), 0)
created = []
real_mkstemp = tempfile.mkstemp
def spy_mkstemp(*args, **kwargs):
result = real_mkstemp(*args, **kwargs)
created.append(result[1])
return result
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 200)
self.assertEqual(resp.content_type, "application/octet-stream")
disposition = resp.headers["Content-Disposition"]
self.assertIn("attachment", disposition)
self.assertIn("prompts_backup_", disposition)
self.assertTrue(disposition.rstrip('"').endswith(".db"))
body = await resp.read()
downloaded = self._write("downloaded.db", body)
self.assertEqual(
_texts(downloaded), ["backup me 0", "backup me 1", "backup me 2"]
)
# The server-side temp file is gone once the response is complete
self.assertEqual(len(created), 1)
self.assertFalse(os.path.exists(created[0]))
async def test_backup_of_missing_database_is_404(self):
self.api.db.model.db_path = os.path.join(self.tmpdir, "nope.db")
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 404)
data = await resp.json()
self.assertFalse(data["success"])
async def test_backup_failure_is_500(self):
self._save("x")
with patch.object(self.api.db.model, "backup_database", return_value=False):
resp = await self.client.get("/prompt_manager/backup")
self.assertEqual(resp.status, 500)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("backup", data["error"].lower())
class TestRestoreRoute(BackupRestoreTestCase):
async def test_restore_replaces_database_and_reports_backup(self):
self._save("old prompt")
old_db = self.api.db
upload = self._make_db_file(["restored a", "restored b"])
resp = await self._restore(upload)
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertTrue(data["success"])
self.assertEqual(data["prompt_count"], 2)
# The response names the safety backup without the server's layout
backups = glob.glob(self.db_path + ".backup_*")
self.assertEqual(len(backups), 1)
self.assertEqual(data["backup_created"], _public_path(backups[0]))
self.assertNotIn(self.tmpdir, data["backup_created"])
self.assertEqual(_texts(backups[0]), ["old prompt"])
self.assertIsNot(self.api.db, old_db)
self.assertEqual(self.api.db.model.db_path, self.db_path)
texts = sorted(p["text"] for p in self.api.db.search_prompts())
self.assertEqual(texts, ["restored a", "restored b"])
async def test_restore_leaves_no_temp_files_behind(self):
upload = self._make_db_file(["restored"])
created = []
real_mkstemp = tempfile.mkstemp
def spy_mkstemp(*args, **kwargs):
result = real_mkstemp(*args, **kwargs)
created.append(result[1])
return result
with patch("py.api.admin.tempfile.mkstemp", spy_mkstemp):
resp = await self._restore(upload)
self.assertEqual(resp.status, 200)
self.assertEqual(len(created), 1)
for suffix in ("", "-wal", "-shm", "-journal"):
self.assertFalse(os.path.exists(created[0] + suffix))
async def test_text_file_is_rejected_and_database_untouched(self):
self._save("keep me")
upload = self._write("notes.txt", b"this is not a database")
resp = await self._restore(upload)
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("SQLite", data["error"])
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
self.assertFalse(
any(n.startswith("prompts.db.backup_") for n in os.listdir(self.tmpdir))
)
async def test_database_without_prompts_table_is_rejected(self):
self._save("keep me")
path = os.path.join(self.tmpdir, "other.db")
conn = sqlite3.connect(path)
conn.execute("CREATE TABLE things (id INTEGER)")
conn.commit()
conn.close()
resp = await self._restore(path)
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("prompts", data["error"])
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
async def test_wrong_field_name_is_rejected(self):
upload = self._make_db_file(["restored"])
resp = await self._restore(upload, field="file")
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("database_file", data["error"])
async def test_empty_upload_is_rejected(self):
upload = self._write("empty.db", b"")
resp = await self._restore(upload)
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("empty", data["error"].lower())
async def test_upload_over_the_size_limit_is_rejected_while_streaming(self):
self._save("keep me")
self.api.restore_max_bytes = 1024
upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
resp = await self._restore(upload)
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("too large", data["error"].lower())
self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
async def test_non_multipart_request_is_rejected(self):
resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
self.assertIn(resp.status, (400, 500))
data = await resp.json()
self.assertFalse(data["success"])
if __name__ == "__main__":
unittest.main()
+809
View File
@@ -0,0 +1,809 @@
"""Tests for the core API module (py/api/__init__.py): helpers and base routes."""
import json
import os
import shutil
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
from database.models import PromptModel # noqa: E402
from aiohttp import web # noqa: E402
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
import py.api as api_module # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.api import PromptManagerAPI, _public_error, _public_path # noqa: E402
from py.config import GalleryConfig # noqa: E402
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
class FolderPathsFixture(unittest.TestCase):
"""Temp ComfyUI tree exposed through a stub folder_paths module."""
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.addCleanup(assert_closed, self, self.tmpdir) # just before the rmtree
self.addCleanup(PromptModel.close_all_instances)
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
self.output_dir = self.comfy_dir / "output"
self.output_dir.mkdir(parents=True)
self._orig_folder_paths = sys.modules.get("folder_paths")
self.addCleanup(self._restore_folder_paths)
self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None)
self.addCleanup(self._restore_extra_env)
self.install_folder_paths(base_path=str(self.comfy_dir))
def install_folder_paths(self, **attrs):
attrs.setdefault("get_output_directory", lambda: str(self.output_dir))
sys.modules["folder_paths"] = types.SimpleNamespace(**attrs)
def _restore_folder_paths(self):
if self._orig_folder_paths is None:
sys.modules.pop("folder_paths", None)
else:
sys.modules["folder_paths"] = self._orig_folder_paths
def _restore_extra_env(self):
if self._orig_extra is None:
os.environ.pop(EXTRA_ROOTS_ENV, None)
else:
os.environ[EXTRA_ROOTS_ENV] = self._orig_extra
class TestPublicPath(FolderPathsFixture):
def test_path_under_base_is_relative_posix(self):
target = self.output_dir / "renders" / "a.png"
self.assertEqual(_public_path(str(target)), "output/renders/a.png")
def test_base_itself_is_dot(self):
self.assertEqual(_public_path(str(self.comfy_dir)), ".")
def test_path_outside_base_is_basename_only(self):
secret = Path(self.tmpdir) / "secret" / "passwords.txt"
secret.parent.mkdir()
secret.write_text("x")
self.assertEqual(_public_path(str(secret)), "passwords.txt")
self.assertNotIn(self.tmpdir, _public_path(str(secret)))
def test_path_object_accepted(self):
self.assertEqual(_public_path(self.output_dir), "output")
def test_falls_back_to_output_parent_without_base_path(self):
self.install_folder_paths() # no base_path attribute
self.assertEqual(_public_path(str(self.output_dir / "x.png")), "output/x.png")
def test_env_extra_root_keeps_its_own_name(self):
extra = Path(self.tmpdir) / "nas" / "gallery"
(extra / "sub").mkdir(parents=True)
os.environ[EXTRA_ROOTS_ENV] = str(extra)
self.assertEqual(_public_path(str(extra / "sub")), "gallery/sub")
def test_without_folder_paths_returns_basename(self):
sys.modules.pop("folder_paths", None)
self.assertEqual(_public_path(str(self.output_dir / "img.png")), "img.png")
def test_empty_or_none_gives_empty_string(self):
self.assertEqual(_public_path(""), "")
self.assertEqual(_public_path(None), "")
class TestPublicError(FolderPathsFixture):
def test_oserror_uses_strerror_and_basename(self):
err = FileNotFoundError(
2, "No such file or directory", str(self.tmpdir) + "/a/b.png"
)
text = _public_error(err)
self.assertIn("No such file or directory", text)
self.assertIn("b.png", text)
self.assertNotIn(self.tmpdir, text)
def test_oserror_without_filename(self):
err = PermissionError(13, "Permission denied")
self.assertEqual(_public_error(err), "Permission denied")
def test_oserror_without_strerror_uses_type_name(self):
err = OSError()
self.assertEqual(_public_error(err), "OSError")
def test_non_oserror_is_str(self):
self.assertEqual(_public_error(ValueError("bad value")), "bad value")
def test_non_oserror_with_empty_message_uses_type_name(self):
self.assertEqual(_public_error(RuntimeError()), "RuntimeError")
# ── gzip middleware ───────────────────────────────────────────────────
class TestGzipMiddleware(AioHTTPTestCase):
async def get_application(self):
big = json.dumps({"data": "x" * 4000})
async def big_json(request):
return web.json_response(text=big)
async def small_json(request):
return web.json_response({"ok": True})
async def already_encoded(request):
return web.Response(
text=big, content_type="text/plain", headers={"Content-Encoding": "br"}
)
async def binary(request):
return web.Response(body=b"\x89PNG" * 1000, content_type="image/png")
async def incompressible(request):
return web.Response(body=os.urandom(4096), content_type="text/plain")
async def streamed(request):
resp = web.StreamResponse()
await resp.prepare(request)
await resp.write(b"x" * 4000)
await resp.write_eof()
return resp
app = web.Application(middlewares=[api_module._gzip_middleware])
app.router.add_get("/prompt_manager/big", big_json)
app.router.add_get("/prompt_manager/small", small_json)
app.router.add_get("/prompt_manager/encoded", already_encoded)
app.router.add_get("/prompt_manager/binary", binary)
app.router.add_get("/prompt_manager/random", incompressible)
app.router.add_get("/prompt_manager/stream", streamed)
app.router.add_get("/other/big", big_json)
return app
async def _get(self, path, accept="gzip"):
headers = {"Accept-Encoding": accept} if accept else {}
return await self.client.request(
"GET",
path,
headers=headers,
auto_decompress=False,
skip_auto_headers=["Accept-Encoding"],
)
async def test_large_json_is_gzipped(self):
resp = await self._get("/prompt_manager/big")
self.assertEqual(resp.headers.get("Content-Encoding"), "gzip")
self.assertEqual(resp.headers.get("Vary"), "Accept-Encoding")
import gzip
self.assertIn(b'"data"', gzip.decompress(await resp.read()))
async def test_small_body_not_compressed(self):
resp = await self._get("/prompt_manager/small")
self.assertIsNone(resp.headers.get("Content-Encoding"))
async def test_other_paths_untouched(self):
resp = await self._get("/other/big")
self.assertIsNone(resp.headers.get("Content-Encoding"))
async def test_without_accept_encoding(self):
resp = await self._get("/prompt_manager/big", accept=None)
self.assertIsNone(resp.headers.get("Content-Encoding"))
async def test_already_encoded_untouched(self):
resp = await self._get("/prompt_manager/encoded")
self.assertEqual(resp.headers.get("Content-Encoding"), "br")
async def test_binary_type_untouched(self):
resp = await self._get("/prompt_manager/binary")
self.assertIsNone(resp.headers.get("Content-Encoding"))
async def test_incompressible_body_untouched(self):
resp = await self._get("/prompt_manager/random")
self.assertIsNone(resp.headers.get("Content-Encoding"))
async def test_stream_response_untouched(self):
resp = await self._get("/prompt_manager/stream")
self.assertIsNone(resp.headers.get("Content-Encoding"))
self.assertEqual(len(await resp.read()), 4000)
# ── JSON body reader ──────────────────────────────────────────────────
class TestReadJsonBody(AioHTTPTestCase):
"""_read_json_body: chunked read with a byte cap, object-only payloads."""
async def get_application(self):
async def echo(request):
data, error = await api_module._read_json_body(request)
if error is not None:
return error
return web.json_response({"success": True, "keys": sorted(data)})
async def small(request):
data, error = await api_module._read_json_body(request, max_bytes=64)
if error is not None:
return error
return web.json_response({"success": True, "data": data})
app = web.Application(client_max_size=4 * 1024 * 1024)
app.router.add_post("/echo", echo)
app.router.add_post("/small", small)
return app
async def test_valid_object_is_parsed(self):
resp = await self.client.request("POST", "/echo", json={"b": 1, "a": 2})
self.assertEqual(resp.status, 200)
self.assertEqual((await resp.json())["keys"], ["a", "b"])
async def test_default_cap_is_one_megabyte(self):
self.assertEqual(api_module.JSON_BODY_MAX_BYTES, 1_000_000)
body = json.dumps({"pad": "x" * 1_000_100}).encode()
resp = await self.client.request(
"POST", "/echo", data=body, headers={"Content-Type": "application/json"}
)
self.assertEqual(resp.status, 413)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("too large", data["error"].lower())
async def test_body_just_under_cap_is_accepted(self):
payload = {"pad": "x" * (1_000_000 - 20)}
body = json.dumps(payload).encode()
self.assertLessEqual(len(body), 1_000_000)
resp = await self.client.request(
"POST", "/echo", data=body, headers={"Content-Type": "application/json"}
)
self.assertEqual(resp.status, 200)
async def test_chunked_upload_without_content_length_is_capped(self):
async def gen():
for _ in range(10):
yield b'{"pad": "' + b"x" * 20 + b'"}'
resp = await self.client.request(
"POST",
"/small",
data=gen(),
headers={"Content-Type": "application/json"},
)
self.assertEqual(resp.status, 413)
async def test_declared_oversize_content_length_is_413(self):
body = b'{"pad": "' + b"x" * 200 + b'"}'
resp = await self.client.request(
"POST", "/small", data=body, headers={"Content-Type": "application/json"}
)
self.assertEqual(resp.status, 413)
async def test_invalid_json_is_400(self):
for body in (b"{oops", b"\xff\xfe", b""):
resp = await self.client.request(
"POST",
"/echo",
data=body,
headers={"Content-Type": "application/json"},
)
self.assertEqual(resp.status, 400, body)
self.assertFalse((await resp.json())["success"])
async def test_non_object_json_is_400(self):
for body in (b"[1, 2]", b'"text"', b"3", b"null", b"true"):
resp = await self.client.request(
"POST",
"/echo",
data=body,
headers={"Content-Type": "application/json"},
)
self.assertEqual(resp.status, 400, body)
self.assertIn("object", (await resp.json())["error"])
def test_helper_is_exposed_on_the_api_object(self):
self.assertIs(PromptManagerAPI._read_json_body, api_module._read_json_body)
class TestLongJobLock(unittest.TestCase):
"""_acquire_long_job: one job per family at a time, released on exit."""
def setUp(self):
self.api = PromptManagerAPI()
def _run(self, coro):
import asyncio
return asyncio.run(coro)
def test_second_acquire_is_refused_until_first_releases(self):
async def scenario():
first = self.api._acquire_long_job("duplicates")
self.assertIsNotNone(first)
async with first:
self.assertIsNone(self.api._acquire_long_job("duplicates"))
# Other families are independent
other = self.api._acquire_long_job("thumbnails")
self.assertIsNotNone(other)
async with other:
pass
again = self.api._acquire_long_job("duplicates")
self.assertIsNotNone(again)
async with again:
pass
self._run(scenario())
def test_lock_is_released_when_the_job_raises(self):
async def scenario():
with self.assertRaises(RuntimeError):
async with self.api._acquire_long_job("duplicates"):
raise RuntimeError("boom")
self.assertIsNotNone(self.api._acquire_long_job("duplicates"))
self._run(scenario())
def test_busy_response_is_409_with_envelope(self):
resp = self.api._long_job_busy_response("Duplicate scan")
self.assertEqual(resp.status, 409)
data = json.loads(resp.body)
self.assertFalse(data["success"])
self.assertEqual(data["error"], "Duplicate scan already running")
# ── UI and static routes ──────────────────────────────────────────────
class CoreRoutesTestCase(AioHTTPTestCase):
async def get_application(self):
self.tmpdir = tempfile.mkdtemp()
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
self.api.add_routes(routes)
app.router.add_routes(routes)
return app
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
self.api.db.close_all()
shutil.rmtree(self.tmpdir, ignore_errors=True)
def _fake_root(self):
"""Point _get_project_root at the temp dir and return its web/ dir."""
web_dir = Path(self.tmpdir) / "web"
(web_dir / "js").mkdir(parents=True)
(web_dir / "lib").mkdir()
orig_root = api_module._get_project_root
api_module._get_project_root = lambda: self.tmpdir
self.addCleanup(setattr, api_module, "_get_project_root", orig_root)
return web_dir
UI_PATHS = (
"/prompt_manager/web",
"/prompt_manager/gallery.html",
"/prompt_manager/admin",
"/prompt_manager/gallery",
)
class TestUiRoutes(CoreRoutesTestCase):
async def test_pages_are_served(self):
for path in UI_PATHS:
resp = await self.client.request("GET", path)
self.assertEqual(resp.status, 200, path)
self.assertIn("text/html", resp.headers["Content-Type"])
self.assertIn("<", await resp.text())
async def test_pages_are_cached_after_first_read(self):
await self.client.request("GET", "/prompt_manager/admin")
self.assertTrue(any(p.endswith("admin.html") for p in self.api._html_cache))
resp = await self.client.request("GET", "/prompt_manager/admin")
self.assertEqual(resp.status, 200)
async def test_missing_pages_are_404(self):
self._fake_root()
for path in UI_PATHS:
resp = await self.client.request("GET", path)
self.assertEqual(resp.status, 404, path)
async def test_unreadable_pages_are_500(self):
web_dir = self._fake_root()
for name in ("index.html", "metadata.html", "admin.html", "gallery.html"):
(web_dir / name).mkdir() # a directory cannot be read as a file
for path in UI_PATHS:
resp = await self.client.request("GET", path)
self.assertEqual(resp.status, 500, path)
self.assertNotIn(self.tmpdir, await resp.text())
async def test_test_route(self):
resp = await self.client.request("GET", "/prompt_manager/test")
self.assertTrue((await resp.json())["success"])
class TestStaticRoutes(CoreRoutesTestCase):
async def test_js_served_with_mime(self):
resp = await self.client.request("GET", "/prompt_manager/js/admin.js")
self.assertEqual(resp.status, 200)
self.assertIn("application/javascript", resp.headers["Content-Type"])
async def test_lib_served(self):
web_dir = self._fake_root()
(web_dir / "lib" / "x.css").write_text("body{}")
resp = await self.client.request("GET", "/prompt_manager/lib/x.css")
self.assertEqual(resp.status, 200)
self.assertIn("text/css", resp.headers["Content-Type"])
self.assertEqual(await resp.text(), "body{}")
async def test_real_lib_file_served_with_content_type(self):
resp = await self.client.request(
"GET", "/prompt_manager/lib/viewerjs/viewer.min.css"
)
self.assertEqual(resp.status, 200)
self.assertIn("text/css", resp.headers["Content-Type"])
async def test_allowed_asset_types_get_their_mime(self):
web_dir = self._fake_root()
expected = {
"a.map": "application/json",
"f.woff": "font/woff",
"f.woff2": "font/woff2",
"f.ttf": "font/ttf",
"i.png": "image/png",
"i.svg": "image/svg+xml",
"i.ico": "image/x-icon",
"p.html": "text/html",
}
for name in expected:
(web_dir / "lib" / name).write_bytes(b"\x00\x01")
for name, mime in expected.items():
resp = await self.client.request("GET", f"/prompt_manager/lib/{name}")
self.assertEqual(resp.status, 200, name)
self.assertIn(mime, resp.headers["Content-Type"], name)
async def test_disallowed_extensions_are_403_even_when_present(self):
web_dir = self._fake_root()
for name in ("x.bin", "x.py", "x.json", "noext", "x.JS.bak"):
(web_dir / "lib" / name).write_bytes(b"\x00")
(web_dir / "js" / name).write_bytes(b"\x00")
for prefix in ("lib", "js"):
for name in ("x.bin", "x.py", "x.json", "noext", "x.JS.bak"):
resp = await self.client.request(
"GET", f"/prompt_manager/{prefix}/{name}"
)
self.assertEqual(resp.status, 403, f"{prefix}/{name}")
async def test_traversal_is_403(self):
for prefix in ("lib", "js"):
for encoded in (
"..%2F..%2Fpyproject.toml",
"..%5C..%5Cpyproject.toml",
"sub%2F..%2F..%2F..%2Fpyproject.toml",
"%2E%2E%2Fpyproject.toml",
):
resp = await self.client.request(
"GET", f"/prompt_manager/{prefix}/{encoded}"
)
self.assertEqual(resp.status, 403, f"{prefix}/{encoded}")
async def test_absolute_drive_and_unc_paths_are_rejected_on_every_os(self):
# Built literally so the Windows-only escapes are exercised everywhere:
# os.path.join(root, "C:/x") drops root on Windows, "\\x" and UNC
# shares re-anchor the path as well.
candidates = (
"C:/Windows/win.ini",
"C:%5CWindows%5Cwin.ini",
"c:win.ini",
"%5C%5Cserver%5Cshare%5Cx.js",
"%2F%2Fserver%2Fshare%2Fx.js",
"%5CWindows%5Cwin.ini",
"%2Fetc%2Fpasswd",
)
for prefix in ("lib", "js"):
for encoded in candidates:
resp = await self.client.request(
"GET", f"/prompt_manager/{prefix}/{encoded}"
)
self.assertIn(resp.status, (403, 404), f"{prefix}/{encoded}")
self.assertNotIn("[extensions]", await resp.text())
async def test_symlink_escaping_web_dir_is_403(self):
web_dir = self._fake_root()
secret = Path(self.tmpdir) / "secret.js"
secret.write_text("secret")
link = web_dir / "lib" / "link.js"
try:
os.symlink(secret, link)
except (OSError, NotImplementedError):
self.skipTest("symlinks not available")
resp = await self.client.request("GET", "/prompt_manager/lib/link.js")
self.assertEqual(resp.status, 403)
async def test_empty_path_is_not_served(self):
for prefix in ("lib", "js"):
resp = await self.client.request("GET", f"/prompt_manager/{prefix}/")
self.assertIn(resp.status, (403, 404), prefix)
async def test_missing_is_404(self):
for prefix in ("lib", "js"):
resp = await self.client.request("GET", f"/prompt_manager/{prefix}/nope.js")
self.assertEqual(resp.status, 404, prefix)
async def test_directory_is_404(self):
web_dir = self._fake_root()
(web_dir / "lib" / "dir.js").mkdir()
resp = await self.client.request("GET", "/prompt_manager/lib/dir.js")
self.assertEqual(resp.status, 404)
async def test_root_of_static_dir_is_not_listed(self):
resp = await self.client.request("GET", "/prompt_manager/lib/tailwind")
self.assertEqual(resp.status, 403)
# ── instance helpers ──────────────────────────────────────────────────
class TestInstanceHelpers(FolderPathsFixture):
def setUp(self):
super().setUp()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
self.addCleanup(self.api.db.close_all) # runs before the rmtree cleanup
orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
self.addCleanup(setattr, GalleryConfig, "MONITORING_DIRECTORIES", orig_dirs)
def _run(self, coro):
import asyncio
return asyncio.run(coro)
def test_run_in_executor_with_kwargs(self):
def add(a, b=0):
return a + b
self.assertEqual(self._run(self.api._run_in_executor(add, 1, b=2)), 3)
self.assertEqual(self._run(self.api._run_in_executor(add, 1, 2)), 3)
def test_invalidate_gallery_cache(self):
self.api._gallery_cache = {"x": 1}
self.api.invalidate_gallery_cache()
self.assertEqual(self.api._gallery_cache, {})
def test_enrich_images_adds_urls_and_thumbnails(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
img = self.output_dir / "sub" / "a.png"
img.parent.mkdir()
img.write_bytes(b"x")
thumb = self.output_dir / "thumbnails" / "sub" / "a_thumb.png"
thumb.parent.mkdir(parents=True)
thumb.write_bytes(b"t")
outside = Path(self.tmpdir) / "elsewhere.png"
outside.write_bytes(b"o")
images = [
{"id": 7, "image_path": str(img)},
{"id": 8, "image_path": str(outside)},
{"id": 9, "image_path": ""},
{"image_path": str(self.output_dir / "missing.png")},
]
result = self.api._enrich_images(images)
self.assertEqual(result[0]["url"], "/prompt_manager/images/serve/sub/a.png")
self.assertRegex(
result[0]["thumbnail_url"],
r"^/prompt_manager/images/serve/thumbnails/sub/a_thumb\.png\?v=\d+$",
)
self.assertEqual(result[0]["relative_path"], os.path.join("sub", "a.png"))
self.assertEqual(result[1]["url"], "/prompt_manager/images/8/file")
self.assertNotIn("relative_path", result[1])
self.assertNotIn("url", result[2])
self.assertEqual(result[3]["url"], "/prompt_manager/images/serve/missing.png")
self.assertNotIn("thumbnail_url", result[3])
def test_enrich_prompt_images_walks_prompts(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
prompts = [
{"images": [{"id": 1, "image_path": str(self.output_dir / "p.png")}]},
{},
]
result = self.api._enrich_prompt_images(prompts)
self.assertEqual(
result[0]["images"][0]["url"], "/prompt_manager/images/serve/p.png"
)
def test_clean_nan_recursive(self):
nan = float("nan")
cleaned = self.api._clean_nan_recursive(
{"a": nan, "b": [nan, 1.5, {"c": nan}], "d": "x"}
)
self.assertEqual(cleaned, {"a": None, "b": [None, 1.5, {"c": None}], "d": "x"})
def test_find_output_dir_uses_configured_and_caches(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
self.api._cached_output_dir = None
first = self.api._find_comfyui_output_dir()
self.assertEqual(first, str(self.output_dir.resolve()))
GalleryConfig.MONITORING_DIRECTORIES = []
self.assertEqual(self.api._find_comfyui_output_dir(), first)
def test_find_output_dir_skips_configured_root_outside_comfyui(self):
outside = Path(self.tmpdir) / "outside"
outside.mkdir()
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
self.api._cached_output_dir = None
result = self.api._find_comfyui_output_dir()
self.assertNotEqual(
os.path.normcase(str(result or "")), os.path.normcase(str(outside))
)
def test_find_output_dir_skips_missing_configured(self):
GalleryConfig.MONITORING_DIRECTORIES = [os.path.join(self.tmpdir, "nope")]
self.api._cached_output_dir = None
result = self.api._find_comfyui_output_dir()
self.assertNotEqual(result, os.path.join(self.tmpdir, "nope"))
# ── metadata parsing ──────────────────────────────────────────────────
class TestMetadataParsing(unittest.TestCase):
def setUp(self):
self.api = PromptManagerAPI()
self.tmpdir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tmpdir, True)
def test_extract_metadata_from_png_and_garbage(self):
from PIL import Image, PngImagePlugin
path = os.path.join(self.tmpdir, "m.png")
info = PngImagePlugin.PngInfo()
info.add_text("prompt", "{}")
Image.new("RGB", (1, 1)).save(path, pnginfo=info)
self.assertEqual(self.api._extract_comfyui_metadata(path), {"prompt": "{}"})
garbage = os.path.join(self.tmpdir, "g.png")
with open(garbage, "wb") as f:
f.write(b"nope")
self.assertEqual(self.api._extract_comfyui_metadata(garbage), {})
def test_parse_a1111_parameters(self):
meta = {"parameters": "a cat\nNegative prompt: dog\nSteps: 20"}
parsed = self.api._parse_comfyui_prompt(meta)
self.assertEqual(parsed["positive_prompt"], "a cat")
self.assertEqual(parsed["negative_prompt"], "dog")
self.assertEqual(parsed["parameters"]["parameters"], meta["parameters"])
self.assertEqual(self.api._extract_readable_prompt(parsed), "a cat")
def test_parse_comfyui_fields(self):
meta = {
"prompt": json.dumps({"1": {"class_type": "X", "inputs": {}}}),
"workflow": "not json",
"steps": "20",
"cfg": "not-json-either",
}
parsed = self.api._parse_comfyui_prompt(meta)
self.assertIsInstance(parsed["prompt"], dict)
self.assertEqual(parsed["workflow"], "not json")
self.assertEqual(parsed["parameters"]["steps"], 20)
self.assertEqual(parsed["parameters"]["cfg"], "not-json-either")
parsed = self.api._parse_comfyui_prompt(
{"prompt": "plain text", "workflow": json.dumps({"nodes": []})}
)
self.assertEqual(parsed["prompt"], "plain text")
self.assertEqual(parsed["workflow"], {"nodes": []})
def test_extract_readable_prompt_variants(self):
extract = self.api._extract_readable_prompt
self.assertEqual(extract({"prompt": "plain"}), "plain")
self.assertEqual(extract({"prompt": ["a", "", "b"]}), "a b")
self.assertEqual(extract({"prompt": 42}), "42")
self.assertEqual(
extract({"prompt": None, "parameters": {"positive": ["x", "y"]}}), "x y"
)
self.assertIsNone(extract({"prompt": {}, "parameters": {}}))
graph = {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "neg"}},
"3": {
"class_type": "KSampler",
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
},
}
self.assertEqual(extract({"prompt": graph}), "pos")
self.assertEqual(extract({"prompt": {"x": 1}, "workflow": graph}), "pos")
def test_node_inputs_and_text_lookup(self):
self.assertEqual(self.api._get_node_inputs("nope"), {})
self.assertEqual(self.api._get_node_inputs({"inputs": 5}), {})
self.assertEqual(self.api._get_node_inputs({"inputs": {"a": 1}}), {"a": 1})
listed = self.api._get_node_inputs(
{"inputs": [{"name": "text", "link": None}, {"nope": 1}, "junk"]}
)
self.assertEqual(list(listed), ["text"])
self.assertIsNone(self.api._find_text_in_node(None))
self.assertEqual(self.api._find_text_in_node({"inputs": {"text": "hi"}}), "hi")
self.assertEqual(
self.api._find_text_in_node(
{"type": "CLIPTextEncode", "widgets_values": ["w"]}
),
"w",
)
self.assertIsNone(
self.api._find_text_in_node(
{"type": "CLIPTextEncode", "widgets_values": [" "]}
)
)
self.assertIsNone(self.api._find_text_in_node({"type": "KSampler"}))
def test_positive_prompt_from_nodes_array_and_fallbacks(self):
extract = self.api._extract_positive_prompt_from_comfyui_data
self.assertIsNone(extract("x"))
self.assertIsNone(extract({}))
self.assertIsNone(extract({"nodes": []}))
workflow = {
"nodes": [
{"id": 1, "type": "CLIPTextEncode", "widgets_values": ["pos"]},
{
"id": 2,
"type": "CLIPTextEncode",
"widgets_values": ["neg"],
"title": "Negative",
},
{"id": 3, "type": "KSampler", "inputs": [{"name": "seed", "link": 4}]},
"junk",
]
}
self.assertEqual(extract(workflow), "pos")
only_negative = {
"nodes": [
{
"id": 2,
"type": "CLIPTextEncode",
"widgets_values": ["neg"],
"title": "neg",
},
]
}
self.assertEqual(extract(only_negative), "neg")
bad_link = {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}},
"3": {
"class_type": "KSampler",
"inputs": {"positive": "oops", "negative": []},
},
"4": {
"class_type": "KSampler",
"inputs": {"positive": ["zz", 0], "negative": []},
},
}
self.assertEqual(extract(bad_link), "pos")
if __name__ == "__main__":
unittest.main()
File diff suppressed because it is too large Load Diff
+447
View File
@@ -0,0 +1,447 @@
"""Route tests for the logging API (py/api/logging_routes.py)."""
import json
import os
import shutil
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from aiohttp import web # noqa: E402
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.api import PromptManagerAPI # noqa: E402
from py.api.logging_routes import LOG_CONFIG_KEYS # noqa: E402
class FakeLoggerManager:
"""Stand-in for PromptManagerLogger with a scratch log directory."""
def __init__(self, log_dir):
self.log_dir = Path(log_dir)
self.config = {
"level": "INFO",
"max_file_size": 10 * 1024 * 1024,
"backup_count": 5,
"console_logging": True,
"file_logging": True,
"buffer_size": 1000,
}
self.update_calls = []
self.logs = []
def update_config(self, new_config):
self.update_calls.append(dict(new_config))
self.config.update(new_config)
def get_config(self):
return dict(self.config)
def get_recent_logs(self, limit=100, level=None):
logs = self.logs
if level:
logs = [entry for entry in logs if entry["level"] == level.upper()]
return list(reversed(logs[-limit:]))
def get_log_files(self):
files = []
for path in sorted(self.log_dir.glob("*.log*")):
files.append(
{
"filename": path.name,
"path": str(path),
"size": path.stat().st_size,
"modified": "2026-10-07T00:00:00",
"is_main": path.name == "prompt_manager.log",
}
)
return files
def truncate_logs(self):
truncated = []
for path in self.log_dir.glob("prompt_manager.log*"):
path.write_text("")
truncated.append(path.name)
return {"truncated": truncated, "errors": []}
def get_log_stats(self):
return {
"buffer_count": len(self.logs),
"level_counts": {},
"log_files_count": len(self.get_log_files()),
"total_log_size": 0,
"log_directory": str(self.log_dir),
"current_level": self.config["level"],
}
class LoggingAPITestCase(AioHTTPTestCase):
async def get_application(self):
self.tmpdir = tempfile.mkdtemp()
self.log_dir = Path(self.tmpdir) / "logs"
self.log_dir.mkdir()
self.manager = FakeLoggerManager(self.log_dir)
# Public paths are rendered relative to ComfyUI's base directory. Other
# test modules leave their own ``folder_paths`` stub behind, so declare
# ours explicitly instead of depending on import order.
self._saved_folder_paths = sys.modules.get("folder_paths")
fake_folder_paths = MagicMock()
fake_folder_paths.base_path = self.tmpdir
fake_folder_paths.get_output_directory.return_value = os.path.join(
self.tmpdir, "output"
)
sys.modules["folder_paths"] = fake_folder_paths
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
self.api._get_logger_manager = lambda: self.manager
self.api.add_routes(routes)
app.router.add_routes(routes)
return app
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
if self._saved_folder_paths is None:
sys.modules.pop("folder_paths", None)
else:
sys.modules["folder_paths"] = self._saved_folder_paths
self.api.db.close_all()
shutil.rmtree(self.tmpdir, ignore_errors=True)
async def _post_config(self, payload):
return await self.client.request(
"POST", "/prompt_manager/logs/config", json=payload
)
class TestUpdateLogConfigWhitelist(LoggingAPITestCase):
def test_whitelist_matches_logger_manager_keys(self):
self.assertEqual(LOG_CONFIG_KEYS, frozenset(self.manager.config.keys()))
async def test_unknown_keys_rejected_and_listed(self):
resp = await self._post_config({"level": "INFO", "evil": 1, "log_dir": "/x"})
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("evil", data["error"])
self.assertIn("log_dir", data["error"])
self.assertEqual(self.manager.update_calls, [])
async def test_invalid_level_rejected(self):
resp = await self._post_config({"level": "VERBOSE"})
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_level_must_be_a_string(self):
resp = await self._post_config({"level": 10})
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_level_is_normalised_to_upper_case(self):
resp = await self._post_config({"level": "debug"})
self.assertEqual(resp.status, 200)
self.assertEqual(self.manager.update_calls, [{"level": "DEBUG"}])
async def test_negative_int_rejected(self):
resp = await self._post_config({"backup_count": -1})
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_numeric_string_rejected(self):
resp = await self._post_config({"max_file_size": "1024"})
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_bool_is_not_an_int(self):
resp = await self._post_config({"buffer_size": True})
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_flags_must_be_bool(self):
resp = await self._post_config({"console_logging": "yes"})
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_valid_payload_is_applied(self):
payload = {
"level": "WARNING",
"max_file_size": 2 * 1024 * 1024,
"backup_count": 0,
"console_logging": False,
"file_logging": True,
"buffer_size": 50,
}
resp = await self._post_config(payload)
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertTrue(data["success"])
self.assertEqual(self.manager.update_calls, [payload])
self.assertEqual(data["config"]["level"], "WARNING")
async def test_max_file_size_is_bounded_to_1_to_1024_mb(self):
mb = 1024 * 1024
for bad in (0, mb - 1, 1024 * mb + 1, 10**12):
resp = await self._post_config({"max_file_size": bad})
self.assertEqual(resp.status, 400, bad)
self.assertIn("max_file_size", (await resp.json())["error"])
self.assertEqual(self.manager.update_calls, [])
for ok in (mb, 1024 * mb):
resp = await self._post_config({"max_file_size": ok})
self.assertEqual(resp.status, 200, ok)
self.assertEqual(
self.manager.update_calls,
[{"max_file_size": mb}, {"max_file_size": 1024 * mb}],
)
async def test_backup_count_is_bounded_to_0_to_50(self):
for bad in (-1, 51, 10**6):
resp = await self._post_config({"backup_count": bad})
self.assertEqual(resp.status, 400, bad)
self.assertIn("backup_count", (await resp.json())["error"])
self.assertEqual(self.manager.update_calls, [])
for ok in (0, 50):
resp = await self._post_config({"backup_count": ok})
self.assertEqual(resp.status, 200, ok)
async def test_body_must_be_an_object(self):
resp = await self._post_config(["level", "DEBUG"])
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_malformed_json_is_400(self):
resp = await self.client.request(
"POST",
"/prompt_manager/logs/config",
data=b"{not json",
headers={"Content-Type": "application/json"},
)
self.assertEqual(resp.status, 400)
self.assertEqual(self.manager.update_calls, [])
async def test_oversized_body_is_413(self):
body = json.dumps({"level": "x" * 1_000_100}).encode()
resp = await self.client.request(
"POST",
"/prompt_manager/logs/config",
data=body,
headers={"Content-Type": "application/json"},
)
self.assertEqual(resp.status, 413)
self.assertFalse((await resp.json())["success"])
self.assertEqual(self.manager.update_calls, [])
class TestNoAbsolutePathsInLogResponses(LoggingAPITestCase):
async def test_log_files_have_no_absolute_path(self):
(self.log_dir / "prompt_manager.log").write_text("hello")
resp = await self.client.request("GET", "/prompt_manager/logs/files")
body = await resp.text()
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
data = json.loads(body)
self.assertEqual(data["files"][0]["filename"], "prompt_manager.log")
async def test_log_stats_directory_is_not_absolute(self):
resp = await self.client.request("GET", "/prompt_manager/logs/stats")
body = await resp.text()
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
self.assertEqual(json.loads(body)["stats"]["log_directory"], "logs")
async def test_error_from_manager_does_not_leak_path(self):
def boom():
raise FileNotFoundError(
2, "No such file", os.path.join(self.tmpdir, "x.log")
)
self.manager.get_log_stats = boom
resp = await self.client.request("GET", "/prompt_manager/logs/stats")
body = await resp.text()
self.assertEqual(resp.status, 500)
self.assertNotIn(self.tmpdir, body)
self.assertIn("x.log", json.loads(body)["error"])
class TestGetLogs(LoggingAPITestCase):
def _fill(self, n):
for i in range(n):
self.manager.logs.append(
{"level": "INFO" if i % 2 else "ERROR", "message": f"m{i}"}
)
async def test_default_limit_and_count(self):
self._fill(3)
resp = await self.client.request("GET", "/prompt_manager/logs")
data = await resp.json()
self.assertEqual(resp.status, 200)
self.assertTrue(data["success"])
self.assertEqual(data["count"], 3)
self.assertEqual(data["limit"], 100)
self.assertIsNone(data["level_filter"])
async def test_level_filter_passed_through(self):
self._fill(4)
resp = await self.client.request("GET", "/prompt_manager/logs?level=error")
data = await resp.json()
self.assertEqual(data["level_filter"], "error")
self.assertEqual(data["count"], 2)
async def test_limit_is_clamped_to_500(self):
resp = await self.client.request("GET", "/prompt_manager/logs?limit=99999")
self.assertEqual((await resp.json())["limit"], 500)
async def test_limit_below_one_becomes_one(self):
self._fill(3)
resp = await self.client.request("GET", "/prompt_manager/logs?limit=0")
data = await resp.json()
self.assertEqual(data["limit"], 1)
self.assertEqual(data["count"], 1)
async def test_non_integer_limit_is_400(self):
resp = await self.client.request("GET", "/prompt_manager/logs?limit=ten")
self.assertEqual(resp.status, 400)
self.assertFalse((await resp.json())["success"])
async def test_manager_failure_is_500_with_empty_logs(self):
def boom(**kwargs):
raise RuntimeError("buffer gone")
self.manager.get_recent_logs = boom
resp = await self.client.request("GET", "/prompt_manager/logs")
data = await resp.json()
self.assertEqual(resp.status, 500)
self.assertEqual(data["logs"], [])
class TestLogFilesAndDownload(LoggingAPITestCase):
async def test_files_failure_is_500(self):
def boom():
raise RuntimeError("no dir")
self.manager.get_log_files = boom
resp = await self.client.request("GET", "/prompt_manager/logs/files")
data = await resp.json()
self.assertEqual(resp.status, 500)
self.assertEqual(data["files"], [])
async def test_download_existing_file(self):
(self.log_dir / "prompt_manager.log").write_bytes(b"line1\nline2\n")
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log"
)
self.assertEqual(resp.status, 200)
self.assertIn("attachment", resp.headers["Content-Disposition"])
self.assertEqual(await resp.text(), "line1\nline2\n")
async def test_download_missing_file_is_404(self):
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/nope.log"
)
self.assertEqual(resp.status, 404)
async def test_download_rejects_traversal(self):
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/..%2F..%2Fetc%2Fpasswd"
)
self.assertEqual(resp.status, 400)
async def test_download_rejects_backslash(self):
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/..%5Cx.log"
)
self.assertEqual(resp.status, 400)
async def test_download_read_failure_is_500(self):
(self.log_dir / "dir.log").mkdir()
resp = await self.client.request("GET", "/prompt_manager/logs/download/dir.log")
body = await resp.text()
self.assertEqual(resp.status, 500)
self.assertNotIn(self.tmpdir, body)
class TestTruncateConfigAndStats(LoggingAPITestCase):
async def test_truncate_success(self):
(self.log_dir / "prompt_manager.log").write_text("data")
resp = await self.client.request("POST", "/prompt_manager/logs/truncate")
data = await resp.json()
self.assertEqual(resp.status, 200)
self.assertEqual(data["results"]["truncated"], ["prompt_manager.log"])
self.assertEqual((self.log_dir / "prompt_manager.log").read_text(), "")
async def test_truncate_failure_is_500(self):
def boom():
raise RuntimeError("locked")
self.manager.truncate_logs = boom
resp = await self.client.request("POST", "/prompt_manager/logs/truncate")
self.assertEqual(resp.status, 500)
async def test_get_config(self):
resp = await self.client.request("GET", "/prompt_manager/logs/config")
data = await resp.json()
self.assertEqual(resp.status, 200)
self.assertEqual(data["config"]["level"], "INFO")
async def test_get_config_failure_is_500(self):
def boom():
raise RuntimeError("gone")
self.manager.get_config = boom
resp = await self.client.request("GET", "/prompt_manager/logs/config")
self.assertEqual(resp.status, 500)
async def test_update_config_manager_failure_is_500(self):
def boom(config):
raise RuntimeError("handler error")
self.manager.update_config = boom
resp = await self.client.request(
"POST", "/prompt_manager/logs/config", json={"level": "INFO"}
)
self.assertEqual(resp.status, 500)
async def test_stats_success(self):
resp = await self.client.request("GET", "/prompt_manager/logs/stats")
data = await resp.json()
self.assertEqual(resp.status, 200)
self.assertEqual(data["stats"]["current_level"], "INFO")
class TestLoggerManagerResolution(LoggingAPITestCase):
def test_default_manager_is_the_singleton(self):
from py.api.logging_routes import LoggingRoutesMixin
from utils.logging_config import get_logger_manager
manager = LoggingRoutesMixin._get_logger_manager(self.api)
self.assertIs(manager, get_logger_manager())
if __name__ == "__main__":
unittest.main()
+624
View File
@@ -0,0 +1,624 @@
"""Route tests for the LoraManager integration endpoints.
Filesystem detection and network downloads are patched out so the tests run
anywhere without a LoraManager install or a CivitAI connection.
"""
import json
import os
import shutil
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if "folder_paths" not in sys.modules:
_stub = types.ModuleType("folder_paths")
_stub.models_dir = tempfile.gettempdir()
_stub.base_path = tempfile.gettempdir()
_stub.get_output_directory = tempfile.gettempdir
_stub.get_folder_paths = lambda name: []
sys.modules["folder_paths"] = _stub
if "server" not in sys.modules:
_server = MagicMock()
_server.PromptServer.instance.routes = MagicMock()
sys.modules["server"] = _server
from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase
from database.operations import PromptDatabase
from py.api import PromptManagerAPI
from py import lora_utils
from py.config import IntegrationConfig, PromptManagerConfig
from py.lora_utils import get_trigger_cache
from utils.hashing import generate_prompt_hash
SECRET_KEY = "sk-civitai-super-secret-0123456789"
class LoraAPITestCase(AioHTTPTestCase):
"""Stands up the PromptManager routes with a temp DB and neutral LoRA config."""
async def get_application(self):
self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
self._temp_db.close()
self.tmp_root = Path(tempfile.mkdtemp())
self.custom_nodes = self.tmp_root / "custom_nodes"
self.custom_nodes.mkdir()
self._custom_nodes_patch = patch(
"py.lora_utils.custom_nodes_directories",
return_value=[self.custom_nodes],
)
self._custom_nodes_patch.start()
self._real_save = PromptManagerConfig.save_to_file
self._config_patch = patch.multiple(
IntegrationConfig,
LORA_MANAGER_ENABLED=False,
LORA_MANAGER_PATH="",
LORA_TRIGGER_WORDS_ENABLED=False,
CIVITAI_API_KEY="",
)
self._config_patch.start()
self._detect_patch = patch(
"py.lora_utils.detect_lora_manager", return_value=None
)
self.detect_mock = self._detect_patch.start()
self._save_patch = patch.object(PromptManagerConfig, "save_to_file")
self.save_mock = self._save_patch.start()
get_trigger_cache().clear()
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(self._temp_db.name)
self.api.add_routes(routes)
app.router.add_routes(routes)
return app
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
self._save_patch.stop()
self._detect_patch.stop()
self._config_patch.stop()
self._custom_nodes_patch.stop()
get_trigger_cache().clear()
self.api.db.close_all()
for path in (
self._temp_db.name,
self._temp_db.name + "-wal",
self._temp_db.name + "-shm",
):
if os.path.exists(path):
os.unlink(path)
shutil.rmtree(self.tmp_root, ignore_errors=True)
async def _post_json(self, url, body):
return await self.client.request(
"POST",
url,
data=json.dumps(body),
headers={"Content-Type": "application/json"},
)
class TestLoraStatusNeverReturnsKey(LoraAPITestCase):
async def test_status_reports_flag_without_the_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertEqual(resp.status, 200)
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertIs(data["has_civitai_api_key"], True)
self.assertNotIn("civitai_api_key", data)
self.assertNotIn(SECRET_KEY, body)
async def test_status_flag_is_false_without_a_key(self):
resp = await self.client.request("GET", "/prompt_manager/lora/status")
data = await resp.json()
self.assertIs(data["has_civitai_api_key"], False)
self.assertNotIn("civitai_api_key", data)
def _png_bytes():
from io import BytesIO
from PIL import Image
buf = BytesIO()
Image.new("RGB", (2, 2), (1, 2, 3)).save(buf, "PNG")
return buf.getvalue()
class _FakeResponse:
def __init__(self, body):
self._body = body
self._pos = 0
self.headers = {}
def read(self, amt=-1):
if amt is None or amt < 0:
amt = len(self._body)
chunk = self._body[self._pos : self._pos + amt]
self._pos += len(chunk)
return chunk
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def _make_lora_manager(where):
lm = where / "ComfyUI-Lora-Manager"
(lm / "py").mkdir(parents=True)
(lm / "__init__.py").write_text("")
return lm
def _write_metadata(directory, stem, words=None, images=None, prompt=None):
meta = {"file_name": f"{stem}.safetensors", "model_name": stem, "civitai": {}}
if words is not None:
meta["civitai"]["trainedWords"] = words
if images is not None:
meta["civitai"]["images"] = [
{"url": url, "meta": {"prompt": prompt} if prompt else {}} for url in images
]
path = directory / f"{stem}.safetensors.metadata.json"
path.write_text(json.dumps(meta))
return path
class TestLoraDetectRoute(LoraAPITestCase):
async def test_detected_path_is_public(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertTrue(data["detected"])
self.assertEqual(data["path"], self.api._public_path(str(lm)))
self.assertNotIn(str(self.tmp_root), body)
async def test_not_detected(self):
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
data = await resp.json()
self.assertFalse(data["detected"])
self.assertEqual(data["path"], "")
async def test_detection_error_is_500(self):
self.detect_mock.side_effect = RuntimeError("disk gone")
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
self.assertEqual(resp.status, 500)
self.assertFalse((await resp.json())["success"])
class TestLoraStatusRoute(LoraAPITestCase):
async def test_status_reports_detection_and_cache_with_public_paths(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(lm)
resp = await self.client.request("GET", "/prompt_manager/lora/status")
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["enabled"])
self.assertTrue(data["detected"])
public = self.api._public_path(str(lm))
self.assertEqual(data["path"], public)
self.assertEqual(data["detected_path"], public)
self.assertFalse(os.path.isabs(public))
self.assertNotIn(str(self.tmp_root), body)
self.assertFalse(data["trigger_cache_loaded"])
async def test_status_error_is_500(self):
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertEqual(resp.status, 500)
async def test_status_error_hides_the_server_path(self):
err = PermissionError(13, "Permission denied", str(self.tmp_root / "c.json"))
with patch("py.lora_utils.get_trigger_cache", side_effect=err):
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertEqual(resp.status, 500)
body = await resp.text()
self.assertNotIn(str(self.tmp_root), body)
self.assertIn("c.json", body)
class TestLoraEnableRoute(LoraAPITestCase):
async def test_enable_with_unknown_path_is_400(self):
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": "/nowhere"}
)
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.save_mock.assert_not_called()
async def test_enable_persists_and_loads_trigger_cache(self):
lm = _make_lora_manager(self.custom_nodes)
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
self.detect_mock.return_value = str(lm)
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
"enabled": True,
"path": str(lm),
"trigger_words_enabled": True,
"civitai_api_key": SECRET_KEY,
},
)
self.assertEqual(resp.status, 200)
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertEqual(data["path"], self.api._public_path(str(lm)))
self.assertNotIn(str(self.tmp_root), body)
self.assertNotIn(SECRET_KEY, body)
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertEqual(
IntegrationConfig.LORA_MANAGER_PATH,
os.path.normcase(os.path.realpath(str(lm))),
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
self.save_mock.assert_called_once_with()
self.assertTrue(get_trigger_cache().is_loaded)
self.assertEqual(get_trigger_cache().get_trigger_words("neon"), ["glow"])
async def test_enable_accepts_the_public_relative_path(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
public = self.api._public_path(str(lm))
self.assertFalse(os.path.isabs(public))
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": public}
)
self.assertEqual(resp.status, 200)
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertEqual(
IntegrationConfig.LORA_MANAGER_PATH,
os.path.normcase(os.path.realpath(str(lm))),
)
async def test_enable_with_path_outside_custom_nodes_is_400(self):
lm = _make_lora_manager(self.tmp_root)
self.detect_mock.return_value = str(lm)
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": str(lm)}
)
self.assertEqual(resp.status, 400)
body = await resp.text()
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.tmp_root), body)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.save_mock.assert_not_called()
async def test_enable_with_custom_node_lacking_lora_in_its_name_is_400(self):
lm = _make_lora_manager(self.custom_nodes).parent / "ComfyUI-Other"
(lm / "py").mkdir(parents=True)
(lm / "__init__.py").write_text("")
self.detect_mock.return_value = str(lm)
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": str(lm)}
)
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_enable_rejects_an_auto_detected_path_outside_custom_nodes(self):
self.detect_mock.return_value = str(_make_lora_manager(self.tmp_root))
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": True})
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_enable_writes_the_configured_config_file(self):
config_file = self.tmp_root / "cfg" / "config.json"
self.save_mock.side_effect = self._real_save
with patch.dict(os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(config_file)}):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
self.assertEqual(resp.status, 200)
self.save_mock.assert_called_once_with()
self.assertTrue(config_file.is_file())
saved = json.loads(config_file.read_text())["integrations"]["lora_manager"]
self.assertFalse(saved["enabled"])
self.assertTrue(saved["trigger_words_enabled"])
async def test_disable_clears_cache(self):
get_trigger_cache()._cache = {"x": ["y"]}
get_trigger_cache()._loaded = True
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
self.assertEqual(resp.status, 200)
self.assertFalse(get_trigger_cache().is_loaded)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_empty_key_keeps_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False, "civitai_api_key": ""}
)
self.assertEqual(resp.status, 200)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
async def test_absent_key_keeps_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
async def test_new_key_replaces_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "civitai_api_key": " new-key "},
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "new-key")
async def test_clear_flag_removes_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "clear_civitai_api_key": True},
)
self.assertEqual(resp.status, 200)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "")
async def test_non_string_key_is_400(self):
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False, "civitai_api_key": 123}
)
self.assertEqual(resp.status, 400)
async def test_persist_failure_is_500(self):
self.save_mock.side_effect = OSError("read-only")
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
self.assertEqual(resp.status, 500)
class TestLoraScanRoute(LoraAPITestCase):
def _enable(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(_make_lora_manager(self.custom_nodes))
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
async def test_configured_path_outside_custom_nodes_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(_make_lora_manager(self.tmp_root))
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
body = await resp.text()
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.tmp_root), body)
async def test_scan_error_hides_the_server_path(self):
self._enable()
err = PermissionError(13, "Permission denied", str(self.tmp_root / "x.db"))
with patch.object(self.api.db, "delete_prompts_by_category", side_effect=err):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
body = await resp.text()
self.assertNotIn(str(self.tmp_root), body)
self.assertIn("x.db", body)
@staticmethod
def _sse_events(text):
return [
json.loads(line[len("data: ") :])
for line in text.splitlines()
if line.startswith("data: ")
]
async def _scan(self, loras):
opener_calls = []
def fake_open(req, timeout=None):
opener_calls.append(req)
return _FakeResponse(_png_bytes())
cache_dir = self.tmp_root / "img_cache"
with (
patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=cache_dir
),
patch.object(lora_utils, "_open_url", fake_open),
):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
text = await resp.text()
return resp, self._sse_events(text), opener_calls
async def test_not_enabled_is_400(self):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
async def test_no_path_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
async def test_scan_imports_loras_with_images_and_links_existing(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(
loras,
"neon",
words=["glow"],
images=["https://civitai.com/a.png", "https://evil.example/b.png"],
prompt="neon street",
)
(loras / "neon.png").write_bytes(_png_bytes())
_write_metadata(loras, "plain")
(loras / "broken.metadata.json").write_text("{{{")
existing_id = self.api.db.save_prompt(
text="plain",
category=None,
tags=[],
rating=None,
prompt_hash=generate_prompt_hash("plain"),
)
resp, events, opener_calls = await self._scan(loras)
self.assertEqual(resp.status, 200)
done = events[-1]
self.assertEqual(done["type"], "complete")
self.assertEqual(done["total"], 3)
self.assertEqual(done["imported"], 1)
self.assertEqual(done["skipped"], 2)
self.assertEqual(
[r.full_url for r in opener_calls], ["https://civitai.com/a.png"]
)
self.assertEqual(
opener_calls[0].get_header("Authorization"), f"Bearer {SECRET_KEY}"
)
imported = self.api.db.search_prompts(category="lora-manager")
self.assertEqual(len(imported), 1)
self.assertEqual(imported[0]["text"], "neon street")
self.assertIn("lora:neon", imported[0]["tags"])
self.assertIn("glow", imported[0]["tags"])
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 2)
self.assertEqual(len(self.api.db.get_prompt_images(existing_id)), 0)
async def test_save_failure_is_counted_as_skipped(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
with patch.object(self.api.db, "save_prompt", side_effect=RuntimeError("db")):
resp, events, _ = await self._scan(loras)
self.assertEqual(events[-1]["imported"], 0)
self.assertEqual(events[-1]["skipped"], 1)
class TestLoraTriggerWordRoutes(LoraAPITestCase):
async def test_trigger_words_not_enabled_is_400(self):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=x"
)
self.assertEqual(resp.status, 400)
async def test_trigger_words_missing_name_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp = await self.client.request("GET", "/prompt_manager/lora/trigger-words")
self.assertEqual(resp.status, 400)
async def test_trigger_words_loads_cache_on_demand(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow", "bright"])
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=NEON"
)
data = await resp.json()
self.assertTrue(data["success"])
self.assertEqual(data["trigger_words"], ["glow", "bright"])
self.assertTrue(get_trigger_cache().is_loaded)
async def test_trigger_words_error_is_500(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=x"
)
self.assertEqual(resp.status, 500)
async def test_refresh_not_enabled_is_400(self):
resp = await self.client.request("POST", "/prompt_manager/lora/refresh-cache")
self.assertEqual(resp.status, 400)
async def test_refresh_without_path_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp = await self.client.request("POST", "/prompt_manager/lora/refresh-cache")
self.assertEqual(resp.status, 400)
async def test_refresh_reloads_cache(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
with patch.object(
lora_utils, "find_lora_directories", return_value=[str(loras)]
):
resp = await self.client.request(
"POST", "/prompt_manager/lora/refresh-cache"
)
data = await resp.json()
self.assertTrue(data["success"])
# keyed by both the file_name stem and the metadata-file stem
self.assertEqual(data["loras_with_trigger_words"], 2)
async def test_refresh_error_is_500(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
resp = await self.client.request(
"POST", "/prompt_manager/lora/refresh-cache"
)
self.assertEqual(resp.status, 500)
if __name__ == "__main__":
unittest.main()
+942
View File
@@ -0,0 +1,942 @@
"""
Route tests for the prompt API (py/api/prompts.py).
"""
import json
import os
import sys
import tempfile
import unittest
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.modules.setdefault("folder_paths", MagicMock())
sys.modules.setdefault("server", MagicMock())
from aiohttp import web # noqa: E402
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.api import PromptManagerAPI # noqa: E402
import py.api.prompts as prompts_module # noqa: E402
from py.api.prompts import ( # noqa: E402
MAX_PAGE_LIMIT,
MAX_PAGE_OFFSET,
safe_error_message,
)
from utils.hashing import generate_prompt_hash # noqa: E402
class PromptAPITestCase(AioHTTPTestCase):
"""App with PromptManager routes and a temporary SQLite database."""
async def get_application(self):
self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
self._temp_db.close()
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(self._temp_db.name)
self.api.add_routes(routes)
app.router.add_routes(routes)
return app
async def tearDownAsync(self):
await super().tearDownAsync() # closes the aiohttp test client
# Windows refuses to unlink a database that still has open handles.
self.api.db.close_all()
for path in (
self._temp_db.name,
self._temp_db.name + "-wal",
self._temp_db.name + "-shm",
):
if os.path.exists(path):
os.unlink(path)
def _save_prompt(self, text="Test prompt", **kwargs):
return self.api.db.save_prompt(
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
)
class TestSafeErrorMessage(unittest.TestCase):
"""Error strings returned to clients never contain absolute server paths."""
def test_oserror_uses_strerror_and_basename(self):
secret = os.path.join(os.sep, "srv", "comfy", "output", "hidden.png")
exc = FileNotFoundError(2, "No such file or directory", secret)
message = safe_error_message(exc)
self.assertNotIn(os.path.dirname(secret), message)
self.assertIn("No such file or directory", message)
self.assertIn("hidden.png", message)
def test_oserror_without_filename(self):
exc = PermissionError(13, "Permission denied")
self.assertEqual(safe_error_message(exc), "Permission denied")
def test_plain_exception_uses_str(self):
self.assertEqual(safe_error_message(ValueError("bad value")), "bad value")
def test_empty_message_falls_back_to_type_name(self):
self.assertEqual(safe_error_message(RuntimeError()), "RuntimeError")
class TestSearchBounds(PromptAPITestCase):
"""GET /prompt_manager/search clamps limit/offset and rejects junk."""
async def test_limit_is_clamped_to_max_page_limit(self):
resp = await self.client.request("GET", "/prompt_manager/search?limit=999999")
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertTrue(data["success"])
self.assertLessEqual(len(data["results"]), MAX_PAGE_LIMIT)
self.assertEqual(data["pagination"]["limit"], MAX_PAGE_LIMIT)
async def test_zero_limit_is_raised_to_one(self):
self._save_prompt("only one")
resp = await self.client.request("GET", "/prompt_manager/search?limit=0")
data = await resp.json()
self.assertEqual(data["pagination"]["limit"], 1)
self.assertEqual(len(data["results"]), 1)
async def test_non_integer_limit_is_400(self):
resp = await self.client.request("GET", "/prompt_manager/search?limit=abc")
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertIn("error", data)
async def test_non_integer_offset_is_400(self):
resp = await self.client.request("GET", "/prompt_manager/search?offset=1.5")
self.assertEqual(resp.status, 400)
async def test_negative_offset_is_clamped_to_zero(self):
# Clamped, not rejected: paging backwards past page one yields page one.
resp = await self.client.request("GET", "/prompt_manager/search?offset=-5")
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertEqual(data["pagination"]["offset"], 0)
async def test_offset_beyond_int64_is_clamped_not_500(self):
resp = await self.client.request(
"GET", "/prompt_manager/search?offset=99999999999999999999999"
)
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertEqual(data["pagination"]["offset"], MAX_PAGE_OFFSET)
async def test_offset_skips_results(self):
for i in range(3):
self._save_prompt(f"searchable {i}")
resp = await self.client.request(
"GET", "/prompt_manager/search?text=searchable&limit=2&offset=2"
)
data = await resp.json()
self.assertEqual(len(data["results"]), 1)
class TestRecentBounds(PromptAPITestCase):
"""GET /prompt_manager/recent clamps limit/offset and rejects junk."""
async def test_limit_is_clamped_to_max_page_limit(self):
resp = await self.client.request("GET", "/prompt_manager/recent?limit=999999")
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertLessEqual(len(data["results"]), MAX_PAGE_LIMIT)
self.assertEqual(data["pagination"]["limit"], MAX_PAGE_LIMIT)
async def test_non_integer_limit_is_400(self):
resp = await self.client.request("GET", "/prompt_manager/recent?limit=abc")
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
async def test_non_integer_page_is_400(self):
resp = await self.client.request("GET", "/prompt_manager/recent?page=two")
self.assertEqual(resp.status, 400)
async def test_negative_offset_is_clamped_to_zero(self):
# Clamped, not rejected: paging backwards past page one yields page one.
resp = await self.client.request("GET", "/prompt_manager/recent?offset=-5")
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertEqual(data["pagination"]["offset"], 0)
async def test_page_param_derives_offset(self):
for i in range(5):
self._save_prompt(f"paged {i}")
resp = await self.client.request("GET", "/prompt_manager/recent?limit=2&page=3")
data = await resp.json()
self.assertEqual(data["pagination"]["offset"], 4)
self.assertEqual(len(data["results"]), 1)
async def test_huge_page_number_is_clamped_not_500(self):
resp = await self.client.request(
"GET", f"/prompt_manager/recent?limit={MAX_PAGE_LIMIT}&page={10**30}"
)
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertEqual(data["pagination"]["offset"], MAX_PAGE_OFFSET)
self.assertEqual(data["results"], [])
def _raise(*_args, **_kwargs):
raise RuntimeError("boom")
class RouteCoverageCase(PromptAPITestCase):
"""Helpers for driving success and failure branches of each route."""
def _break(self, db_method):
"""Make a DB method raise so the route's 500 branch runs."""
setattr(self.api.db, db_method, _raise)
def _stub(self, db_method, value):
setattr(self.api.db, db_method, lambda *a, **k: value)
async def _json(self, method, path, **kwargs):
resp = await self.client.request(method, path, **kwargs)
return resp.status, await resp.json()
class TestListRoutesErrorPaths(RouteCoverageCase):
async def test_search_db_failure_is_500(self):
self._break("search_prompts")
status, data = await self._json("GET", "/prompt_manager/search")
self.assertEqual(status, 500)
self.assertFalse(data["success"])
self.assertEqual(data["results"], [])
async def test_search_ignores_non_integer_min_rating(self):
self._save_prompt("rated", rating=3)
status, data = await self._json("GET", "/prompt_manager/search?min_rating=x")
self.assertEqual(status, 200)
self.assertEqual(len(data["results"]), 1)
async def test_recent_db_failure_is_500(self):
self._break("get_recent_prompts")
status, data = await self._json("GET", "/prompt_manager/recent")
self.assertEqual(status, 500)
self.assertEqual(data["pagination"]["total"], 0)
async def test_categories_and_tags_db_failures_are_500(self):
self._break("get_all_categories")
self._break("get_all_tags")
c_status, c_data = await self._json("GET", "/prompt_manager/categories")
t_status, t_data = await self._json("GET", "/prompt_manager/tags")
self.assertEqual((c_status, t_status), (500, 500))
self.assertEqual(c_data["categories"], [])
self.assertEqual(t_data["tags"], [])
async def test_subfolders_success_and_failure(self):
ok_status, ok_data = await self._json(
"GET", "/prompt_manager/subfolders?include_ancestors=true"
)
self._break("get_prompt_subfolders")
bad_status, _ = await self._json("GET", "/prompt_manager/subfolders")
self.assertEqual(ok_status, 200)
self.assertEqual(ok_data["subfolders"], [])
self.assertEqual(bad_status, 500)
async def test_export_db_failure_is_500(self):
self._break("search_prompts")
status, data = await self._json("GET", "/prompt_manager/export")
self.assertEqual(status, 500)
self.assertFalse(data["success"])
class TestTagStatsAndFilters(RouteCoverageCase):
async def test_tags_stats_rejects_bad_limit(self):
status, _ = await self._json("GET", "/prompt_manager/tags/stats?limit=x")
self.assertEqual(status, 400)
async def test_tags_stats_db_failure_is_500(self):
self._break("get_tags_with_counts")
status, _ = await self._json("GET", "/prompt_manager/tags/stats")
self.assertEqual(status, 500)
async def test_tag_prompts_success(self):
self._save_prompt("P1", tags=["sky"])
status, data = await self._json("GET", "/prompt_manager/tags/sky/prompts")
self.assertEqual(status, 200)
self.assertEqual(data["tag"], "sky")
self.assertEqual(len(data["prompts"]), 1)
self.assertIn("pagination", data)
async def test_tag_prompts_bad_limit_and_db_failure(self):
bad_status, _ = await self._json(
"GET", "/prompt_manager/tags/sky/prompts?offset=z"
)
self._break("get_prompts_by_tags")
err_status, _ = await self._json("GET", "/prompt_manager/tags/sky/prompts")
self.assertEqual((bad_status, err_status), (400, 500))
async def test_tags_filter_untagged_mode(self):
self._save_prompt("no tags")
self._save_prompt("tagged", tags=["x"])
status, data = await self._json(
"GET", "/prompt_manager/tags/filter?untagged=true"
)
self.assertEqual(status, 200)
self.assertEqual(data["mode"], "untagged")
self.assertEqual(len(data["prompts"]), 1)
async def test_tags_filter_untagged_bad_limit(self):
status, _ = await self._json(
"GET", "/prompt_manager/tags/filter?untagged=true&limit=x"
)
self.assertEqual(status, 400)
async def test_tags_filter_requires_tags(self):
status, _ = await self._json("GET", "/prompt_manager/tags/filter")
self.assertEqual(status, 400)
async def test_tags_filter_and_or_modes(self):
self._save_prompt("both", tags=["a", "b"])
self._save_prompt("only a", tags=["a"])
_, and_data = await self._json(
"GET", "/prompt_manager/tags/filter?tags=a,b&mode=and"
)
_, or_data = await self._json(
"GET", "/prompt_manager/tags/filter?tags=a,b&mode=weird"
)
self.assertEqual(len(and_data["prompts"]), 1)
self.assertEqual(or_data["mode"], "and")
self.assertEqual(len(or_data["prompts"]), 1)
async def test_tags_filter_bad_limit_and_db_failure(self):
bad_status, _ = await self._json(
"GET", "/prompt_manager/tags/filter?tags=a&limit=x"
)
self._break("get_prompts_by_tags")
err_status, _ = await self._json("GET", "/prompt_manager/tags/filter?tags=a")
self.assertEqual((bad_status, err_status), (400, 500))
class TestTagSuggestRoute(RouteCoverageCase):
"""GET /prompt_manager/tags/suggest?q=&with=&limit="""
URL = "/prompt_manager/tags/suggest"
def _seed(self):
for text, tags in (
("P1", ["asian", "portrait", "smile"]),
("P2", ["asian", "portrait"]),
("P3", ["portrait", "dog"]),
):
self.api.db.save_prompt(text, "test", tags)
async def test_suggestions_narrow_by_context_tags(self):
self._seed()
status, data = await self._json("GET", f"{self.URL}?q=&with=asian")
self.assertEqual(status, 200)
self.assertTrue(data["success"])
self.assertEqual(
[(t["name"], t["count"]) for t in data["suggestions"]],
[("portrait", 2), ("smile", 1)],
)
async def test_prefix_and_messy_context_parsing(self):
self._seed()
status, data = await self._json(
"GET", f"{self.URL}?q=S&with=%20asian%2C%2Cportrait%2Casian%20"
)
self.assertEqual(status, 200)
self.assertEqual([t["name"] for t in data["suggestions"]], ["smile"])
async def test_limit_is_clamped_and_route_is_not_a_tag_name(self):
self._seed()
status, data = await self._json("GET", f"{self.URL}?limit=1")
self.assertEqual(status, 200)
self.assertEqual(len(data["suggestions"]), 1)
status, data = await self._json("GET", f"{self.URL}?limit=9999")
self.assertEqual(status, 200)
self.assertLessEqual(len(data["suggestions"]), 50)
status, data = await self._json("GET", f"{self.URL}?limit=abc")
self.assertEqual(status, 200)
async def test_db_failure_is_500(self):
self._break("suggest_tags")
status, data = await self._json("GET", self.URL)
self.assertEqual(status, 500)
self.assertFalse(data["success"])
class TestTagRoutePagingBounds(RouteCoverageCase):
"""Tag listings clamp limit/offset like every other list endpoint."""
ROUTES = (
"/prompt_manager/tags/stats",
"/prompt_manager/tags/sky/prompts",
"/prompt_manager/tags/filter?tags=sky",
"/prompt_manager/tags/filter?untagged=true",
)
@staticmethod
def _with(route, query):
return f"{route}{'&' if '?' in route else '?'}{query}"
async def test_limit_is_clamped_to_max_page_limit(self):
for route in self.ROUTES:
status, data = await self._json("GET", self._with(route, "limit=999999"))
self.assertEqual(status, 200, route)
self.assertEqual(data["pagination"]["limit"], MAX_PAGE_LIMIT, route)
async def test_zero_limit_becomes_one(self):
for route in self.ROUTES:
_, data = await self._json("GET", self._with(route, "limit=0"))
self.assertEqual(data["pagination"]["limit"], 1, route)
async def test_negative_offset_is_clamped_to_zero(self):
for route in self.ROUTES:
_, data = await self._json("GET", self._with(route, "offset=-7"))
self.assertEqual(data["pagination"]["offset"], 0, route)
async def test_huge_offset_is_clamped_not_500(self):
for route in self.ROUTES:
status, data = await self._json(
"GET", self._with(route, f"offset={10**30}")
)
self.assertEqual(status, 200, route)
self.assertEqual(data["pagination"]["offset"], MAX_PAGE_OFFSET, route)
async def test_non_integer_values_are_400(self):
for route in self.ROUTES:
status, data = await self._json("GET", self._with(route, "limit=1.5"))
self.assertEqual(status, 400, route)
self.assertFalse(data["success"], route)
class TestPublicPathsInPromptResponses(RouteCoverageCase):
"""Preview images on prompt listings never expose the server's directories."""
SECRET_DIR = os.path.join(os.sep, "srv", "secret")
async def setUpAsync(self):
await super().setUpAsync()
self.pid = self._save_prompt("has preview", tags=["sky"])
image_path = os.path.join(self.SECRET_DIR, "gen.png")
self.assertGreater(self.api.db.link_image_to_prompt(self.pid, image_path), 0)
async def _body(self, method, path, **kwargs):
resp = await self.client.request(method, path, **kwargs)
text = await resp.text()
self.assertNotIn(self.SECRET_DIR, text, path)
return resp.status, json.loads(text)
async def test_listing_routes_publish_preview_image_paths(self):
routes = {
"/prompt_manager/recent": "results",
"/prompt_manager/search?q=preview": "results",
"/prompt_manager/tags/sky/prompts": "prompts",
"/prompt_manager/tags/filter?tags=sky": "prompts",
}
for route, key in routes.items():
status, data = await self._body("GET", route)
self.assertEqual(status, 200, route)
image = data[key][0]["images"][0]
# No ComfyUI anchor is configured here, so only the name survives.
self.assertEqual(image["image_path"], "gen.png", route)
self.assertNotIn("file_path", image, route)
async def test_oserror_bodies_keep_only_the_basename(self):
def boom(*_args, **_kwargs):
raise FileNotFoundError(
2, "No such file", os.path.join(self.SECRET_DIR, "prompts.db")
)
for db_method, route in (
("get_tags_with_counts", "/prompt_manager/tags/stats"),
("search_prompts", "/prompt_manager/search"),
("get_recent_prompts", "/prompt_manager/recent"),
("get_prompts_by_tags", "/prompt_manager/tags/sky/prompts"),
):
setattr(self.api.db, db_method, boom)
status, data = await self._body("GET", route)
self.assertEqual(status, 500, route)
self.assertIn("prompts.db", data["error"], route)
class TestTagNamesAreDecodedOnce(RouteCoverageCase):
"""aiohttp already decodes match_info; a tag literally named '%41' stays '%41'."""
async def setUpAsync(self):
await super().setUpAsync()
self.pid = self._save_prompt("percent tag", tags=["%41"])
async def test_tag_prompts_lookup(self):
status, data = await self._json("GET", "/prompt_manager/tags/%2541/prompts")
self.assertEqual(status, 200)
self.assertEqual(data["tag"], "%41")
self.assertEqual(len(data["prompts"]), 1)
async def test_rename_tag(self):
status, data = await self._json(
"PUT", "/prompt_manager/tags/%2541", json={"new_name": "renamed"}
)
self.assertEqual(status, 200)
self.assertEqual(data["old_name"], "%41")
self.assertEqual(data["affected_count"], 1)
self.assertEqual(self.api.db.get_prompt_by_id(self.pid)["tags"], ["renamed"])
async def test_delete_tag(self):
status, data = await self._json("DELETE", "/prompt_manager/tags/%2541")
self.assertEqual(status, 200)
self.assertEqual(data["tag_name"], "%41")
self.assertEqual(data["affected_count"], 1)
self.assertEqual(self.api.db.get_prompt_by_id(self.pid)["tags"], [])
class TestTagInputValidation(RouteCoverageCase):
"""Tag mutations run validate_tags: bad input is a 400, never a traceback."""
def _requests(self, tag):
pid = self._save_prompt("validate me", tags=["old"])
return (
("POST", f"/prompt_manager/prompts/{pid}/tags", {"tag": tag}),
("POST", "/prompt_manager/prompts/tags", {"prompt_id": pid, "tags": [tag]}),
("POST", "/prompt_manager/bulk/tags", {"prompt_ids": [pid], "tags": [tag]}),
(
"POST",
"/prompt_manager/tags/merge",
{"source_tags": ["old"], "target_tag": tag},
),
(
"POST",
"/prompt_manager/tags/merge",
{"source_tags": [tag], "target_tag": "t"},
),
("PUT", "/prompt_manager/tags/old", {"new_name": tag}),
)
async def _assert_all_400(self, tag):
for method, route, body in self._requests(tag):
status, data = await self._json(method, route, json=body)
self.assertEqual(status, 400, (route, body))
self.assertFalse(data["success"], route)
self.assertIn("error", data, route)
async def test_non_string_tags_are_400(self):
await self._assert_all_400(123)
async def test_overlong_tags_are_400(self):
await self._assert_all_400("x" * 51)
async def test_control_characters_are_400(self):
await self._assert_all_400("bad\x00tag")
class TestExportPagesThroughAllPrompts(RouteCoverageCase):
"""GET /prompt_manager/export never truncates: it pages until exhausted."""
async def test_export_is_complete_beyond_one_page(self):
ids = {self._save_prompt(f"export {i}") for i in range(5)}
real_search = self.api.db.search_prompts
calls = []
def spy(*args, **kwargs):
calls.append((kwargs.get("limit"), kwargs.get("offset")))
return real_search(*args, **kwargs)
self.api.db.search_prompts = spy
with patch.object(prompts_module, "EXPORT_PAGE_SIZE", 2):
resp = await self.client.request("GET", "/prompt_manager/export")
self.assertEqual(resp.status, 200)
data = json.loads(await resp.text())
self.assertEqual(data["total_prompts"], 5)
self.assertEqual({p["id"] for p in data["prompts"]}, ids)
self.assertEqual(calls, [(2, 0), (2, 2), (2, 4)])
class TestTagMutationRoutes(RouteCoverageCase):
async def test_rename_rejects_bad_json_and_empty_name(self):
bad_json, _ = await self._json(
"PUT",
"/prompt_manager/tags/old",
data=b"{nope",
headers={"Content-Type": "application/json"},
)
empty, _ = await self._json(
"PUT", "/prompt_manager/tags/old", json={"new_name": " "}
)
self.assertEqual((bad_json, empty), (400, 400))
async def test_rename_reports_skipped_rows_and_db_failure(self):
self._stub("rename_tag_all_prompts", {"affected_count": 1, "skipped_count": 2})
status, data = await self._json(
"PUT", "/prompt_manager/tags/old", json={"new_name": "new"}
)
self.assertEqual(status, 200)
self.assertEqual(data["skipped_count"], 2)
self.assertIn("warning", data)
self._break("rename_tag_all_prompts")
err_status, _ = await self._json(
"PUT", "/prompt_manager/tags/old", json={"new_name": "new"}
)
self.assertEqual(err_status, 500)
async def test_delete_tag_reports_skipped_rows_and_db_failure(self):
self._stub("delete_tag_all_prompts", {"affected_count": 3, "skipped_count": 1})
status, data = await self._json("DELETE", "/prompt_manager/tags/gone")
self.assertEqual(status, 200)
self.assertEqual(data["skipped_count"], 1)
self._break("delete_tag_all_prompts")
err_status, _ = await self._json("DELETE", "/prompt_manager/tags/gone")
self.assertEqual(err_status, 500)
async def test_merge_validation_errors(self):
bad_json, _ = await self._json(
"POST",
"/prompt_manager/tags/merge",
data=b"{nope",
headers={"Content-Type": "application/json"},
)
no_source, _ = await self._json(
"POST", "/prompt_manager/tags/merge", json={"target_tag": "t"}
)
no_target, _ = await self._json(
"POST", "/prompt_manager/tags/merge", json={"source_tags": ["s"]}
)
self.assertEqual((bad_json, no_source, no_target), (400, 400, 400))
async def test_merge_reports_skipped_rows_and_db_failure(self):
self._stub(
"merge_tags",
{"affected_count": 2, "tags_merged": ["s"], "skipped_count": 1},
)
status, data = await self._json(
"POST",
"/prompt_manager/tags/merge",
json={"source_tags": ["s"], "target_tag": "t"},
)
self.assertEqual(status, 200)
self.assertEqual(data["skipped_count"], 1)
self._break("merge_tags")
err_status, _ = await self._json(
"POST",
"/prompt_manager/tags/merge",
json={"source_tags": ["s"], "target_tag": "t"},
)
self.assertEqual(err_status, 500)
class TestSaveRoute(RouteCoverageCase):
async def test_invalid_rating_is_400(self):
status, data = await self._json(
"POST", "/prompt_manager/save", json={"text": "ok", "rating": 99}
)
self.assertEqual(status, 400)
self.assertIn("Rating", data["error"])
async def test_duplicate_updates_metadata(self):
first = self._save_prompt("same text")
status, data = await self._json(
"POST",
"/prompt_manager/save",
json={"text": "same text", "category": "cat", "tags": ["t"], "rating": 4},
)
self.assertEqual(status, 200)
self.assertTrue(data["is_duplicate"])
self.assertEqual(data["prompt_id"], first)
stored = self.api.db.get_prompt_by_id(first)
self.assertEqual(stored["category"], "cat")
self.assertEqual(stored["rating"], 4)
async def test_duplicate_without_metadata_leaves_row_alone(self):
first = self._save_prompt("plain dup")
status, data = await self._json(
"POST", "/prompt_manager/save", json={"text": "plain dup"}
)
self.assertEqual(status, 200)
self.assertEqual(data["prompt_id"], first)
async def test_db_failure_is_500(self):
self._break("get_prompt_by_hash")
status, data = await self._json(
"POST", "/prompt_manager/save", json={"text": "boom"}
)
self.assertEqual(status, 500)
self.assertFalse(data["success"])
class TestDeleteUpdateRatingRoutes(RouteCoverageCase):
async def test_delete_invalid_id_and_db_failure(self):
bad, _ = await self._json("DELETE", "/prompt_manager/delete/abc")
self._break("delete_prompt")
err, _ = await self._json("DELETE", "/prompt_manager/delete/1")
self.assertEqual((bad, err), (400, 500))
async def test_update_text_validation_and_not_found(self):
pid = self._save_prompt("orig")
empty, _ = await self._json(
"PUT", f"/prompt_manager/prompts/{pid}", json={"text": " "}
)
too_long, _ = await self._json(
"PUT", f"/prompt_manager/prompts/{pid}", json={"text": "x" * 10_001}
)
missing, _ = await self._json(
"PUT", "/prompt_manager/prompts/999999", json={"text": "new"}
)
bad_id, _ = await self._json(
"PUT", "/prompt_manager/prompts/abc", json={"text": "new"}
)
self.assertEqual((empty, too_long, missing, bad_id), (400, 400, 404, 400))
async def test_update_text_db_failure(self):
pid = self._save_prompt("orig")
self._break("update_prompt_text")
status, _ = await self._json(
"PUT", f"/prompt_manager/prompts/{pid}", json={"text": "new"}
)
self.assertEqual(status, 500)
async def test_rating_validation_not_found_and_failure(self):
pid = self._save_prompt("rate me")
invalid, _ = await self._json(
"PUT", f"/prompt_manager/prompts/{pid}/rating", json={"rating": 0}
)
missing, _ = await self._json(
"PUT", "/prompt_manager/prompts/999999/rating", json={"rating": 3}
)
bad_id, _ = await self._json(
"PUT", "/prompt_manager/prompts/abc/rating", json={"rating": 3}
)
self._break("update_prompt_rating")
err, _ = await self._json(
"PUT", f"/prompt_manager/prompts/{pid}/rating", json={"rating": 3}
)
self.assertEqual((invalid, missing, bad_id, err), (400, 404, 400, 500))
class TestPromptTagRoutes(RouteCoverageCase):
async def test_add_tag_validation_paths(self):
pid = self._save_prompt("tagme", tags=["have"])
empty, _ = await self._json(
"POST", f"/prompt_manager/prompts/{pid}/tags", json={"tag": ""}
)
missing, _ = await self._json(
"POST", "/prompt_manager/prompts/999999/tags", json={"tag": "x"}
)
bad_id, _ = await self._json(
"POST", "/prompt_manager/prompts/abc/tags", json={"tag": "x"}
)
dup, dup_data = await self._json(
"POST", f"/prompt_manager/prompts/{pid}/tags", json={"tag": "have"}
)
self.assertEqual((empty, missing, bad_id, dup), (400, 404, 400, 200))
self.assertTrue(dup_data["success"])
self.assertEqual(self.api.db.get_prompt_by_id(pid)["tags"], ["have"])
async def test_add_tag_db_failure(self):
pid = self._save_prompt("tagme")
self._break("set_prompt_tags")
status, _ = await self._json(
"POST", f"/prompt_manager/prompts/{pid}/tags", json={"tag": "x"}
)
self.assertEqual(status, 500)
async def test_add_tags_bulk_to_single_prompt(self):
pid = self._save_prompt("multi", tags=["a"])
status, data = await self._json(
"POST",
"/prompt_manager/prompts/tags",
json={"prompt_id": pid, "tags": ["a", " b ", "", "c"]},
)
self.assertEqual(status, 200)
self.assertEqual(data["tags_added"], 2)
again, again_data = await self._json(
"POST",
"/prompt_manager/prompts/tags",
json={"prompt_id": pid, "tags": ["a"]},
)
self.assertEqual(again, 200)
self.assertEqual(again_data["tags_added"], 0)
self.assertIn("No new tags", again_data["message"])
async def test_add_tags_validation_and_failures(self):
no_id, _ = await self._json(
"POST", "/prompt_manager/prompts/tags", json={"tags": ["a"]}
)
no_tags, _ = await self._json(
"POST", "/prompt_manager/prompts/tags", json={"prompt_id": 1, "tags": "a"}
)
missing, _ = await self._json(
"POST",
"/prompt_manager/prompts/tags",
json={"prompt_id": 999999, "tags": ["a"]},
)
self._break("get_prompt_by_id")
err, _ = await self._json(
"POST", "/prompt_manager/prompts/tags", json={"prompt_id": 1, "tags": ["a"]}
)
self.assertEqual((no_id, no_tags, missing, err), (400, 400, 404, 500))
async def test_remove_tag_paths(self):
pid = self._save_prompt("untag", tags=["a", "b"])
removed, _ = await self._json(
"DELETE", f"/prompt_manager/prompts/{pid}/tags", json={"tag": "a"}
)
absent, _ = await self._json(
"DELETE", f"/prompt_manager/prompts/{pid}/tags", json={"tag": "zzz"}
)
missing, _ = await self._json(
"DELETE", "/prompt_manager/prompts/999999/tags", json={"tag": "a"}
)
bad_id, _ = await self._json(
"DELETE", "/prompt_manager/prompts/abc/tags", json={"tag": "a"}
)
self.assertEqual((removed, absent, missing, bad_id), (200, 200, 404, 400))
self.assertEqual(self.api.db.get_prompt_by_id(pid)["tags"], ["b"])
async def test_remove_tag_db_failure(self):
pid = self._save_prompt("untag", tags=["a"])
self._break("set_prompt_tags")
status, _ = await self._json(
"DELETE", f"/prompt_manager/prompts/{pid}/tags", json={"tag": "a"}
)
self.assertEqual(status, 500)
class TestBulkRoutes(RouteCoverageCase):
async def test_bulk_delete_paths(self):
ids = [self._save_prompt(f"bulk {i}") for i in range(3)]
empty, _ = await self._json(
"POST", "/prompt_manager/bulk/delete", json={"prompt_ids": []}
)
ok, ok_data = await self._json(
"POST", "/prompt_manager/bulk/delete", json={"prompt_ids": ids[:2]}
)
self._break("bulk_delete_prompts")
err, _ = await self._json(
"POST", "/prompt_manager/bulk/delete", json={"prompt_ids": ids}
)
self.assertEqual((empty, ok, err), (400, 200, 500))
self.assertEqual(ok_data["deleted_count"], 2)
async def test_bulk_add_tags_paths(self):
ids = [self._save_prompt(f"bulk tag {i}") for i in range(2)]
empty, _ = await self._json(
"POST", "/prompt_manager/bulk/tags", json={"prompt_ids": ids, "tags": []}
)
ok, ok_data = await self._json(
"POST", "/prompt_manager/bulk/tags", json={"prompt_ids": ids, "tags": ["z"]}
)
self._break("bulk_add_tags")
err, _ = await self._json(
"POST", "/prompt_manager/bulk/tags", json={"prompt_ids": ids, "tags": ["z"]}
)
self.assertEqual((empty, ok, err), (400, 200, 500))
self.assertEqual(ok_data["updated_count"], 2)
async def test_bulk_set_category_paths(self):
ids = [self._save_prompt(f"bulk cat {i}") for i in range(2)]
empty, _ = await self._json(
"POST", "/prompt_manager/bulk/category", json={"category": "c"}
)
ok, ok_data = await self._json(
"POST",
"/prompt_manager/bulk/category",
json={"prompt_ids": ids, "category": "c"},
)
self._break("bulk_set_category")
err, _ = await self._json(
"POST",
"/prompt_manager/bulk/category",
json={"prompt_ids": ids, "category": "c"},
)
self.assertEqual((empty, ok, err), (400, 200, 500))
self.assertEqual(ok_data["updated_count"], 2)
self.assertEqual(self.api.db.get_prompt_by_id(ids[0])["category"], "c")
if __name__ == "__main__":
unittest.main()
class TestPromptEditHash(PromptAPITestCase):
async def test_edit_moves_the_hash_and_original_text_can_be_saved_again(self):
original_id = self._save_prompt("Original landscape")
response = await self.client.put(
f"/prompt_manager/prompts/{original_id}", json={"text": "Edited portrait"}
)
self.assertEqual(response.status, 200)
row = self.api.db.get_prompt_by_id(original_id)
self.assertEqual(row["hash"], generate_prompt_hash("Edited portrait"))
self.assertIsNone(
self.api.db.get_prompt_by_hash(generate_prompt_hash("Original landscape"))
)
response = await self.client.post(
"/prompt_manager/save", json={"text": "Original landscape"}
)
self.assertEqual(response.status, 200)
self.assertNotEqual((await response.json())["prompt_id"], original_id)
response = await self.client.post(
"/prompt_manager/save", json={"text": "Edited portrait"}
)
self.assertEqual((await response.json())["prompt_id"], original_id)
async def test_duplicate_edit_returns_conflict_without_changing_either_prompt(self):
first = self._save_prompt("Original landscape", tags=["landscape"])
second = self._save_prompt("Existing portrait", tags=["portrait"])
before = [self.api.db.get_prompt_by_id(pid) for pid in (first, second)]
response = await self.client.put(
f"/prompt_manager/prompts/{first}", json={"text": " EXISTING PORTRAIT "}
)
self.assertEqual(response.status, 409)
self.assertFalse((await response.json())["success"])
self.assertEqual(
[self.api.db.get_prompt_by_id(pid) for pid in (first, second)], before
)
async def test_edit_with_same_normalized_hash_keeps_its_identity(self):
pid = self._save_prompt("Original landscape")
response = await self.client.put(
f"/prompt_manager/prompts/{pid}", json={"text": "ORIGINAL LANDSCAPE"}
)
self.assertEqual(response.status, 200)
self.assertEqual(
self.api.db.get_prompt_by_id(pid)["text"], "ORIGINAL LANDSCAPE"
)
+716
View File
@@ -0,0 +1,716 @@
"""Tests for the AutoTagService tagging pipeline with an injected fake model.
No real ML model is loaded; ``model_loader`` returns small fakes that mimic
the onnxruntime session (WD14) and the llama.cpp chat model (GGUF).
"""
import contextlib
import csv
import os
import shutil
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from PIL import Image
from py import autotag
from py.autotag import MODELS, AutoTagService
@contextlib.contextmanager
def _patch_modules(mapping):
"""Temporarily set entries of sys.modules, restoring only those keys.
unittest.mock.patch.dict restores by clearing the whole dict, which evicts
any heavy module first imported inside the window (torch cannot be
re-imported in-process), so this patcher touches only the named keys.
"""
missing = object()
saved = {name: sys.modules.get(name, missing) for name in mapping}
try:
for name, module in mapping.items():
sys.modules[name] = module
yield
finally:
for name, previous in saved.items():
if previous is missing:
sys.modules.pop(name, None)
else:
sys.modules[name] = previous
class FakeInput:
def __init__(self, name="input", shape=(1, 8, 8, 3)):
self.name = name
self.shape = list(shape)
class FakeOnnxSession:
"""Mimics onnxruntime.InferenceSession for a WD14 classifier."""
def __init__(self, probabilities):
self.probabilities = probabilities
self.run_calls = []
def get_inputs(self):
return [FakeInput()]
def get_providers(self):
return ["CPUExecutionProvider"]
def run(self, outputs, feeds):
self.run_calls.append(feeds)
return [[list(self.probabilities)]]
class FakeGgufModel:
"""Mimics llama_cpp.Llama.create_chat_completion."""
def __init__(self, text):
self.text = text
self.calls = []
def create_chat_completion(self, **kwargs):
self.calls.append(kwargs)
return {"choices": [{"message": {"content": self.text}}]}
WD14_TAGS = [
{"name": "1girl", "category": "0"},
{"name": "smile", "category": "0"},
{"name": "hatsune_miku", "category": "4"},
{"name": "rating:safe", "category": "9"},
{"name": "", "category": "0"},
{"name": "long_hair", "category": "0"},
]
def _touch_model_files(models_dir, model_type):
config = MODELS[model_type]
subdir = Path(models_dir) / config["subdir"]
subdir.mkdir(parents=True, exist_ok=True)
names = config.get("files") or [config["filename"], config["mmproj_filename"]]
for name in names:
(subdir / name).write_bytes(b"stub")
class EngineTestCase(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
self.models_dir = self.tmp / "models"
self.image_path = self.tmp / "img.png"
Image.new("RGB", (4, 6), (200, 100, 50)).save(self.image_path, "PNG")
self.loader_calls = []
def _service(self, loaded):
def loader(model_type, use_gpu):
self.loader_calls.append((model_type, use_gpu))
return loaded
return AutoTagService(models_dir=self.models_dir, model_loader=loader)
class TestModelInjection(EngineTestCase):
def test_injected_loader_is_used_and_state_tracked(self):
_touch_model_files(self.models_dir, "wd14-vit")
service = self._service(("wd14", (FakeOnnxSession([0.0] * 6), WD14_TAGS)))
self.assertTrue(service.load_model("wd14-vit", use_gpu=False))
self.assertEqual(self.loader_calls, [("wd14-vit", False)])
self.assertTrue(service.is_model_loaded())
self.assertEqual(service.get_loaded_model_type(), "wd14-vit")
def test_default_loader_is_the_lazy_dispatcher(self):
service = AutoTagService(models_dir=self.models_dir)
self.assertEqual(service._model_loader, service._load_model)
def test_importing_autotag_does_not_import_ml_backends(self):
for heavy in ("onnxruntime", "llama_cpp"):
self.assertNotIn(heavy, sys.modules)
def test_default_wd14_loader_reports_missing_onnxruntime(self):
_touch_model_files(self.models_dir, "wd14-vit")
service = AutoTagService(models_dir=self.models_dir)
if "onnxruntime" in sys.modules:
self.skipTest("onnxruntime is installed in this environment")
with self.assertRaises(RuntimeError) as ctx:
service.load_model("wd14-vit")
self.assertIn("onnxruntime", str(ctx.exception))
self.assertFalse(service.is_model_loaded())
def test_invalid_model_type_raises_value_error(self):
service = self._service(("gguf", FakeGgufModel("x")))
with self.assertRaises(ValueError):
service.load_model("nope")
def test_not_downloaded_raises_runtime_error_before_loading(self):
service = self._service(("gguf", FakeGgufModel("x")))
with self.assertRaises(RuntimeError):
service.load_model("gguf")
self.assertEqual(self.loader_calls, [])
def test_loader_failure_clears_state(self):
_touch_model_files(self.models_dir, "gguf")
def failing_loader(model_type, use_gpu):
raise OSError("cuda exploded")
service = AutoTagService(
models_dir=self.models_dir, model_loader=failing_loader
)
with self.assertRaises(RuntimeError):
service.load_model("gguf")
self.assertFalse(service.is_model_loaded())
self.assertIsNone(service.get_loaded_model_type())
def test_loading_again_unloads_first(self):
_touch_model_files(self.models_dir, "gguf")
_touch_model_files(self.models_dir, "wd14-swinv2")
service = self._service(("gguf", FakeGgufModel("x")))
service.load_model("gguf")
service.load_model("wd14-swinv2")
self.assertEqual([c[0] for c in self.loader_calls], ["gguf", "wd14-swinv2"])
self.assertEqual(service.get_loaded_model_type(), "wd14-swinv2")
def test_unload_clears_state_and_is_idempotent(self):
_touch_model_files(self.models_dir, "gguf")
service = self._service(("gguf", FakeGgufModel("x")))
service.load_model("gguf")
service.unload_model()
service.unload_model()
self.assertFalse(service.is_model_loaded())
self.assertIsNone(service.get_loaded_model_type())
class TestWd14Thresholding(EngineTestCase):
def _loaded(self, probabilities):
_touch_model_files(self.models_dir, "wd14-vit")
self.session = FakeOnnxSession(probabilities)
service = self._service(("wd14", (self.session, WD14_TAGS)))
service.load_model("wd14-vit")
return service
def test_default_thresholds_apply_per_category(self):
# 1girl 0.9 (general, keep), smile 0.3 (general, drop at 0.35),
# miku 0.8 (character, drop at 0.85), rating 0.99 (ignored),
# empty name 0.99 (ignored), long_hair 0.5 (keep)
service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5])
tags = service.generate_tags(str(self.image_path))
self.assertEqual(tags, ["1girl", "long_hair"])
def test_explicit_thresholds_override_defaults(self):
service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5])
tags = service.generate_tags(
str(self.image_path), general_threshold=0.25, character_threshold=0.75
)
self.assertEqual(tags, ["1girl", "hatsune_miku", "long_hair", "smile"])
def test_results_sorted_by_confidence(self):
service = self._loaded([0.4, 0.95, 0.0, 0.0, 0.0, 0.6])
self.assertEqual(
service.generate_tags(str(self.image_path)),
["smile", "long_hair", "1girl"],
)
def test_fewer_probabilities_than_tags_is_safe(self):
service = self._loaded([0.9, 0.9])
self.assertEqual(
service.generate_tags(str(self.image_path)), ["1girl", "smile"]
)
def test_image_is_fed_as_square_batch(self):
service = self._loaded([0.0] * 6)
service.generate_tags(str(self.image_path))
feeds = self.session.run_calls[0]
self.assertEqual(list(feeds.keys()), ["input"])
self.assertEqual(tuple(feeds["input"].shape), (1, 8, 8, 3))
def test_service_threshold_setters_clamp(self):
service = self._loaded([0.0] * 6)
service.wd14_general_threshold = 1.7
service.wd14_character_threshold = -2
self.assertEqual(service.wd14_general_threshold, 1.0)
self.assertEqual(service.wd14_character_threshold, 0.0)
class TestGgufNormalisation(EngineTestCase):
def _loaded(self, text):
_touch_model_files(self.models_dir, "gguf")
self.model = FakeGgufModel(text)
service = self._service(("gguf", self.model))
service.load_model("gguf")
return service
def test_tags_are_lowercased_underscored_deduped_and_filtered(self):
service = self._loaded(
"1Girl, Smile, smile, 'quoted tag', photo_realistic, meta:x, "
"copyright:y, prepend:z, append:w, a, , Long Hair "
)
tags = service.generate_tags(str(self.image_path))
self.assertEqual(tags, ["1girl", "smile", "quoted_tag", "long_hair"])
def test_custom_prompt_is_sent_to_the_model(self):
service = self._loaded("tag")
service.custom_prompt = "describe tersely"
service.generate_tags(str(self.image_path))
user_turn = self.model.calls[0]["messages"][1]["content"]
self.assertEqual(user_turn[0]["text"], "describe tersely")
self.assertTrue(user_turn[1]["image_url"]["url"].startswith("data:image/png"))
def test_explicit_prompt_argument_wins(self):
service = self._loaded("tag")
service.generate_tags(str(self.image_path), prompt="one-off")
self.assertEqual(
self.model.calls[0]["messages"][1]["content"][0]["text"], "one-off"
)
class TestGenerateTagsErrors(EngineTestCase):
def test_no_model_loaded(self):
service = self._service(("gguf", FakeGgufModel("x")))
with self.assertRaises(RuntimeError):
service.generate_tags(str(self.image_path))
def test_missing_image_error_does_not_embed_the_full_path(self):
_touch_model_files(self.models_dir, "gguf")
service = self._service(("gguf", FakeGgufModel("x")))
service.load_model("gguf")
missing = self.tmp / "nope.png"
with self.assertRaises(FileNotFoundError) as ctx:
service.generate_tags(str(missing))
self.assertNotIn(str(self.tmp), ctx.exception.strerror)
self.assertEqual(ctx.exception.filename, str(missing))
class TestModelsStatus(EngineTestCase):
def test_status_reflects_files_on_disk(self):
_touch_model_files(self.models_dir, "wd14-vit")
gguf_dir = self.models_dir / MODELS["gguf"]["subdir"]
gguf_dir.mkdir(parents=True)
(gguf_dir / MODELS["gguf"]["filename"]).write_bytes(b"stub")
service = AutoTagService(models_dir=self.models_dir)
status = service.get_models_status()
self.assertTrue(status["wd14-vit"]["downloaded"])
self.assertTrue(status["wd14-vit"]["model_path"].endswith("model.onnx"))
self.assertFalse(status["wd14-swinv2"]["downloaded"])
self.assertTrue(status["gguf"]["model_exists"])
self.assertFalse(status["gguf"]["mmproj_exists"])
self.assertFalse(status["gguf"]["downloaded"])
self.assertFalse(status["hf"]["downloaded"])
def test_hf_local_directory_counts_as_downloaded(self):
hf_dir = self.models_dir / MODELS["hf"]["subdir"]
hf_dir.mkdir(parents=True)
(hf_dir / "config.json").write_text("{}")
service = AutoTagService(models_dir=self.models_dir)
status = service.get_models_status()
self.assertTrue(status["hf"]["downloaded"])
self.assertEqual(
os.path.realpath(status["hf"]["model_path"]), os.path.realpath(str(hf_dir))
)
# ── Model directory resolution and singleton ──────────────────────────
class TestModelsDirResolution(unittest.TestCase):
def test_uses_folder_paths_models_dir(self):
fake = types.ModuleType("folder_paths")
fake.models_dir = os.path.join(tempfile.gettempdir(), "fp_models")
with _patch_modules({"folder_paths": fake}):
service = AutoTagService()
self.assertEqual(service.models_dir, Path(fake.models_dir) / "LLM")
def test_falls_back_without_folder_paths(self):
with _patch_modules({"folder_paths": None}):
service = AutoTagService()
self.assertEqual(service.models_dir.name, "LLM")
def test_singleton_is_created_once(self):
with patch.object(autotag, "_service_instance", None):
first = autotag.get_autotag_service()
second = autotag.get_autotag_service()
self.assertIs(first, second)
self.assertIsInstance(first, AutoTagService)
# ── HuggingFace cache lookup ──────────────────────────────────────────
def _fake_hub(snapshot_dir=None, repo_id=None, raise_scan=False):
hub = types.ModuleType("huggingface_hub")
hub.HFCacheInfo = object
def scan_cache_dir():
if raise_scan:
raise OSError("no cache")
revision = types.SimpleNamespace(snapshot_path=str(snapshot_dir))
repo = types.SimpleNamespace(
repo_id=repo_id, repo_type="model", revisions=[revision]
)
other = types.SimpleNamespace(
repo_id="someone/else", repo_type="dataset", revisions=[]
)
return types.SimpleNamespace(repos=[other, repo])
hub.scan_cache_dir = scan_cache_dir
return hub
class TestHfCacheLookup(EngineTestCase):
def test_cache_snapshot_with_config_is_found(self):
snapshot = self.tmp / "snap"
snapshot.mkdir()
(snapshot / "config.json").write_text("{}")
hub = _fake_hub(snapshot, MODELS["hf"]["repo"])
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules({"huggingface_hub": hub}):
self.assertTrue(service._check_hf_model())
status = service.get_models_status()
self.assertTrue(status["hf"]["downloaded"])
self.assertEqual(Path(status["hf"]["model_path"]), snapshot)
def test_snapshot_without_config_is_ignored(self):
snapshot = self.tmp / "snap"
snapshot.mkdir()
hub = _fake_hub(snapshot, MODELS["hf"]["repo"])
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules({"huggingface_hub": hub}):
self.assertIsNone(service._get_hf_model_path())
def test_scan_errors_and_missing_hub_give_none(self):
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules({"huggingface_hub": _fake_hub(raise_scan=True)}):
self.assertIsNone(service._get_hf_cache_path("x/y"))
with _patch_modules({"huggingface_hub": None}):
self.assertIsNone(service._get_hf_cache_path("x/y"))
# ── Downloads (hub patched, no network) ───────────────────────────────
class _FakeDownloadHub:
"""huggingface_hub stand-in that materialises files instead of fetching."""
def __init__(self):
self.module = types.ModuleType("huggingface_hub")
self.module.hf_hub_download = self.hf_hub_download
self.module.snapshot_download = self.snapshot_download
self.calls = []
def hf_hub_download(self, repo_id, filename, local_dir, **kwargs):
self.calls.append(("file", repo_id, filename))
Path(local_dir).mkdir(parents=True, exist_ok=True)
(Path(local_dir) / filename).write_bytes(b"model")
return str(Path(local_dir) / filename)
def snapshot_download(self, repo_id, local_dir, **kwargs):
self.calls.append(("snapshot", repo_id))
Path(local_dir).mkdir(parents=True, exist_ok=True)
(Path(local_dir) / "config.json").write_text("{}")
return local_dir
class TestDownloadModel(EngineTestCase):
def setUp(self):
super().setUp()
self.hub = _FakeDownloadHub()
self.service = AutoTagService(models_dir=self.models_dir)
self.progress = []
def _cb(self, status, progress):
self.progress.append((status, progress))
def test_invalid_type_raises(self):
with self.assertRaises(ValueError):
self.service.download_model("nope")
def test_missing_hub_reports_failure(self):
with _patch_modules({"huggingface_hub": None}):
self.assertFalse(self.service.download_model("gguf", self._cb))
self.assertEqual(self.progress[-1][1], 0)
self.assertIn("huggingface_hub", self.progress[-1][0])
def test_wd14_downloads_each_missing_file(self):
with _patch_modules({"huggingface_hub": self.hub.module}):
self.assertTrue(self.service.download_model("wd14-vit", self._cb))
self.assertTrue(self.service.download_model("wd14-vit", self._cb))
self.assertEqual(
[c[2] for c in self.hub.calls], ["model.onnx", "selected_tags.csv"]
)
self.assertTrue(self.service.get_models_status()["wd14-vit"]["downloaded"])
self.assertEqual(self.progress[-1], ("Download complete", 100))
def test_gguf_downloads_model_then_mmproj(self):
with _patch_modules({"huggingface_hub": self.hub.module}):
self.assertTrue(self.service.download_model("gguf", self._cb))
self.assertEqual(
[c[1] for c in self.hub.calls],
[MODELS["gguf"]["repo"], MODELS["gguf"]["mmproj_repo"]],
)
self.assertTrue(self.service.get_models_status()["gguf"]["downloaded"])
self.assertIn(("Main model ready", 50), self.progress)
def test_gguf_skips_files_already_present(self):
_touch_model_files(self.models_dir, "gguf")
with _patch_modules({"huggingface_hub": self.hub.module}):
self.assertTrue(self.service.download_model("gguf"))
self.assertEqual(self.hub.calls, [])
def test_hf_snapshot_download(self):
with _patch_modules({"huggingface_hub": self.hub.module}):
self.assertTrue(self.service.download_model("hf", self._cb))
self.assertTrue(self.service.download_model("hf"))
self.assertEqual(self.hub.calls, [("snapshot", MODELS["hf"]["repo"])])
self.assertTrue(self.service.get_models_status()["hf"]["downloaded"])
def test_download_errors_are_reported_not_raised(self):
def boom(**kwargs):
raise OSError("disk full")
self.hub.module.hf_hub_download = boom
with _patch_modules({"huggingface_hub": self.hub.module}):
self.assertFalse(self.service.download_model("wd14-vit", self._cb))
self.assertTrue(self.progress[-1][0].startswith("Error:"))
# ── Default backend loaders with the ML libraries faked ───────────────
class _FakeOrt:
def __init__(self):
self.module = types.ModuleType("onnxruntime")
self.module.InferenceSession = self.InferenceSession
self.sessions = []
def InferenceSession(self, path, providers):
session = FakeOnnxSession([0.9, 0.1])
session.path = path
session.providers = providers
self.sessions.append(session)
return session
class TestDefaultLoaders(EngineTestCase):
def test_wd14_loader_builds_session_and_reads_tags(self):
_touch_model_files(self.models_dir, "wd14-vit")
csv_path = self.models_dir / MODELS["wd14-vit"]["subdir"] / "selected_tags.csv"
with open(csv_path, "w", newline="", encoding="utf-8") as f:
writer = csv.DictWriter(f, fieldnames=["tag_id", "name", "category"])
writer.writeheader()
writer.writerow({"tag_id": 1, "name": "1girl", "category": "0"})
writer.writerow({"tag_id": 2, "name": "smile", "category": "0"})
ort = _FakeOrt()
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules({"onnxruntime": ort.module}):
service.load_model("wd14-vit", use_gpu=False)
cpu_providers = ort.sessions[-1].providers
service.load_model("wd14-vit", use_gpu=True)
gpu_providers = ort.sessions[-1].providers
self.assertEqual(cpu_providers, ["CPUExecutionProvider"])
self.assertEqual(gpu_providers[0], "CUDAExecutionProvider")
self.assertEqual(service.generate_tags(str(self.image_path)), ["1girl"])
def test_gguf_loader_wires_llama_with_mmproj(self):
_touch_model_files(self.models_dir, "gguf")
llama_cpp = types.ModuleType("llama_cpp")
chat_format = types.ModuleType("llama_cpp.llama_chat_format")
created = {}
class Llama:
def __init__(self, **kwargs):
created.update(kwargs)
def create_chat_completion(self, **kwargs):
return {"choices": [{"message": {"content": "Aa, Bb"}}]}
class Llava15ChatHandler:
def __init__(self, clip_model_path):
created["clip"] = clip_model_path
llama_cpp.Llama = Llama
chat_format.Llava15ChatHandler = Llava15ChatHandler
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules(
{"llama_cpp": llama_cpp, "llama_cpp.llama_chat_format": chat_format},
):
service.load_model("gguf", use_gpu=False)
self.assertEqual(created["n_gpu_layers"], 0)
self.assertTrue(created["clip"].endswith(MODELS["gguf"]["mmproj_filename"]))
self.assertEqual(service.generate_tags(str(self.image_path)), ["aa", "bb"])
def _fake_torch(self, cuda=False):
torch = MagicMock(name="torch")
torch.cuda.is_available.return_value = cuda
torch.float16 = "f16"
torch.bfloat16 = "bf16"
torch.inference_mode.return_value.__enter__ = lambda s: None
torch.inference_mode.return_value.__exit__ = lambda s, *a: False
torch.cuda.amp.autocast.return_value.__enter__ = lambda s: None
torch.cuda.amp.autocast.return_value.__exit__ = lambda s, *a: False
return torch
def _fake_transformers(self):
transformers = MagicMock(name="transformers")
processor = MagicMock(name="processor")
processor.apply_chat_template.return_value = "<convo>"
inputs = {"input_ids": MagicMock(shape=(1, 3)), "pixel_values": MagicMock()}
processed = MagicMock(name="inputs")
processed.to.return_value = inputs
processor.return_value = processed
processor.tokenizer.decode.return_value = "Red Hat, blue sky"
transformers.AutoProcessor.from_pretrained.return_value = processor
model = MagicMock(name="model")
model.generate.return_value = [list(range(10))]
transformers.LlavaForConditionalGeneration.from_pretrained.return_value = model
return transformers, model, processor, inputs
def test_hf_loader_8bit_on_cpu_and_generation(self):
hf_dir = self.models_dir / MODELS["hf"]["subdir"]
hf_dir.mkdir(parents=True)
(hf_dir / "config.json").write_text("{}")
torch = self._fake_torch(cuda=False)
transformers, model, processor, inputs = self._fake_transformers()
pixel = inputs["pixel_values"]
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules({"torch": torch, "transformers": transformers}):
service.load_model("hf")
tags = service.generate_tags(str(self.image_path))
self.assertEqual(tags, ["red_hat", "blue_sky"])
kwargs = (
transformers.LlavaForConditionalGeneration.from_pretrained.call_args.kwargs
)
self.assertEqual(kwargs["device_map"], "cpu")
self.assertIn("quantization_config", kwargs)
model.eval.assert_called_once()
pixel.to.assert_called_with("f16")
processor.tokenizer.decode.assert_called_once()
decoded = processor.tokenizer.decode.call_args.args[0]
self.assertEqual(decoded, list(range(3, 10)))
def test_hf_loader_without_quantization_uses_bf16_and_cuda(self):
hf_dir = self.models_dir / MODELS["hf"]["subdir"]
hf_dir.mkdir(parents=True)
(hf_dir / "config.json").write_text("{}")
torch = self._fake_torch(cuda=True)
transformers, model, processor, inputs = self._fake_transformers()
pixel = inputs["pixel_values"]
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules({"torch": torch, "transformers": transformers}):
kind, (loaded, _, device, dtype) = service._load_hf_tagger(
quantization="none"
)
service._tagger = (kind, (loaded, processor, device, dtype))
service._current_model_type = "hf"
service.generate_tags(str(self.image_path))
service.unload_model()
self.assertEqual((kind, device, dtype), ("hf", "cuda", "bf16"))
kwargs = (
transformers.LlavaForConditionalGeneration.from_pretrained.call_args.kwargs
)
self.assertNotIn("quantization_config", kwargs)
pixel.to.assert_called_with("bf16")
torch.cuda.empty_cache.assert_called_once()
def test_hf_loader_without_model_raises(self):
service = AutoTagService(models_dir=self.models_dir)
with _patch_modules(
{"torch": self._fake_torch(), "transformers": self._fake_transformers()[0]},
):
with self.assertRaises(RuntimeError):
service._load_hf_tagger()
class TestUnloadDoesNotImportTorch(EngineTestCase):
def test_unload_skips_torch_when_not_imported(self):
_touch_model_files(self.models_dir, "gguf")
service = self._service(("gguf", FakeGgufModel("x")))
service.load_model("gguf")
with _patch_modules({"torch": None}):
with patch("builtins.__import__", side_effect=ImportError("no torch")):
service.unload_model()
self.assertFalse(service.is_model_loaded())
class TestImageConversion(EngineTestCase):
def test_grayscale_images_are_converted_to_rgb(self):
_touch_model_files(self.models_dir, "wd14-vit")
gray = self.tmp / "gray.png"
Image.new("L", (3, 3), 128).save(gray, "PNG")
session = FakeOnnxSession([0.9] * 6)
service = self._service(("wd14", (session, WD14_TAGS)))
service.load_model("wd14-vit")
tags = service.generate_tags(str(gray))
self.assertIn("1girl", tags)
self.assertEqual(tuple(session.run_calls[0]["input"].shape), (1, 8, 8, 3))
if __name__ == "__main__":
unittest.main()
+17 -4
View File
@@ -12,6 +12,11 @@ import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
from utils.validators import validate_prompt_text, validate_rating, validate_tags
@@ -30,8 +35,12 @@ class TestBasicFunctionality(unittest.TestCase):
def tearDown(self):
"""Clean up test fixtures."""
# Remove temporary database
if os.path.exists(self.temp_db.name):
os.unlink(self.temp_db.name)
if getattr(self, "db", None) is not None:
self.db.close_all()
assert_closed(self, self.temp_db.name)
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.temp_db.name + suffix):
os.unlink(self.temp_db.name + suffix)
def test_prompt_hash_generation(self):
"""Test prompt hash generation."""
@@ -214,8 +223,12 @@ class TestNodeIntegration(unittest.TestCase):
def tearDown(self):
"""Clean up test fixtures."""
# Remove temporary database
if os.path.exists(self.temp_db.name):
os.unlink(self.temp_db.name)
if getattr(self, "db", None) is not None:
self.db.close_all()
assert_closed(self, self.temp_db.name)
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.temp_db.name + suffix):
os.unlink(self.temp_db.name + suffix)
@patch("prompt_manager_base.PromptDatabase")
def test_node_encode_function(self, mock_db_class):
+84
View File
@@ -65,5 +65,89 @@ class TestSaveImageNotPatched(unittest.TestCase):
self.assertEqual(integration.get_current_prompt_text("pm_1"), "hello")
import threading # noqa: E402
import unittest.mock as mock # noqa: E402
from utils import comfyui_integration # noqa: E402
from utils.comfyui_integration import get_comfyui_integration # noqa: E402
class IntegrationTestCase(unittest.TestCase):
def setUp(self):
ComfyUIMetadataIntegration._instance = None
comfyui_integration._integration_instance = None
# Patch the module's own ``time`` name rather than ``time.time`` globally:
# background threads from other tests (image monitor, watchdog) would
# otherwise tick this clock and make fresh registrations look stale.
self.clock = mock.patch.object(
comfyui_integration, "time", types.SimpleNamespace(time=self._now)
)
self.clock.start()
self.addCleanup(self.clock.stop)
self._ticks = 1000.0
self.integration = get_comfyui_integration()
def tearDown(self):
ComfyUIMetadataIntegration._instance = None
comfyui_integration._integration_instance = None
def _now(self):
self._ticks += 1.0
return self._ticks
class TestSingletonAndRegistry(IntegrationTestCase):
def test_singleton_keeps_registrations_across_constructions(self):
self.assertIs(get_comfyui_integration(), self.integration)
self.assertIs(ComfyUIMetadataIntegration(), self.integration)
self.integration.register_prompt("n1", "hello", {"category": "x"})
ComfyUIMetadataIntegration() # __init__ runs again: must not reset state
self.assertEqual(self.integration.get_current_prompt_text("n1"), "hello")
def test_latest_prompt_of_this_thread_wins_without_a_node_id(self):
self.integration.register_prompt("n1", "first", {})
self.integration.register_prompt("n2", "second", {})
self.assertEqual(self.integration.get_current_prompt_text(), "second")
self.assertEqual(self.integration.get_current_prompt_text("missing"), "second")
self.assertEqual(self.integration.get_current_prompt_text("n1"), "first")
def test_nothing_registered_returns_none(self):
self.assertIsNone(self.integration.get_current_prompt_text())
self.assertIsNone(self.integration.get_current_prompt_text("n1"))
class TestCrossThreadLookup(IntegrationTestCase):
def _register_in_thread(self, text):
worker = threading.Thread(
target=self.integration.register_prompt, args=("w1", text, {})
)
worker.start()
worker.join()
def test_recent_prompt_from_another_thread_is_used(self):
self._register_in_thread("from worker")
self.assertEqual(self.integration.get_current_prompt_text(), "from worker")
def test_stale_prompt_from_another_thread_is_ignored(self):
self._register_in_thread("from worker")
self._ticks += 400.0
self.assertIsNone(self.integration.get_current_prompt_text())
def test_global_entry_of_this_thread_is_found_without_thread_local(self):
self.integration.register_prompt("n1", "mine", {})
self.integration._thread_local = threading.local()
self.assertEqual(self.integration.get_current_prompt_text(), "mine")
def test_cleanup_removes_old_registrations_only(self):
self._register_in_thread("old")
self._ticks += 700.0
self.integration.register_prompt("n2", "new", {})
self.integration.cleanup_old_prompts(max_age_seconds=600)
texts = sorted(v["text"] for v in self.integration._current_prompts.values())
self.assertEqual(texts, ["new"])
self.integration.cleanup_old_prompts(max_age_seconds=600) # nothing to do
self.assertEqual(len(self.integration._current_prompts), 1)
if __name__ == "__main__":
unittest.main()
+114 -1
View File
@@ -15,7 +15,7 @@ from unittest.mock import MagicMock
# Mock ComfyUI's server module before importing config
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules["server"] = _mock_server
sys.modules.setdefault("server", _mock_server)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -136,6 +136,43 @@ class TestPromptManagerConfig(unittest.TestCase):
PromptManagerConfig.update_config({"performance": {"max_search_results": 50}})
self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, 50)
def test_infinite_scroll_default_off_and_only_booleans_apply(self):
self.assertIn("infinite_scroll", PromptManagerConfig.get_config()["web_ui"])
PromptManagerConfig.update_config({"web_ui": {"infinite_scroll": True}})
self.assertTrue(PromptManagerConfig.INFINITE_SCROLL)
PromptManagerConfig.update_config({"web_ui": {"infinite_scroll": "yes"}})
self.assertTrue(PromptManagerConfig.INFINITE_SCROLL)
PromptManagerConfig.update_config({"web_ui": {"infinite_scroll": False}})
self.assertFalse(PromptManagerConfig.INFINITE_SCROLL)
def test_worker_threads_default_is_half_the_cores_and_current_is_in_range(self):
# The current value may come from config.json (loaded at import), so the
# rule is checked on the constant and only the range on the live value.
cores = PromptManagerConfig.max_worker_threads()
self.assertGreaterEqual(cores, 1)
self.assertEqual(PromptManagerConfig.DEFAULT_WORKER_THREADS, max(1, cores // 2))
current = PromptManagerConfig.get_config()["performance"]["worker_threads"]
self.assertGreaterEqual(current, 1)
self.assertLessEqual(current, cores)
def test_update_worker_threads_is_clamped_to_the_machine(self):
cores = PromptManagerConfig.max_worker_threads()
PromptManagerConfig.update_config({"performance": {"worker_threads": 1}})
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
PromptManagerConfig.update_config(
{"performance": {"worker_threads": cores + 50}}
)
self.assertEqual(PromptManagerConfig.WORKER_THREADS, cores)
PromptManagerConfig.update_config({"performance": {"worker_threads": 0}})
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
def test_update_worker_threads_ignores_junk_from_a_hand_edited_file(self):
PromptManagerConfig.update_config({"performance": {"worker_threads": 1}})
PromptManagerConfig.update_config({"performance": {"worker_threads": "lots"}})
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
PromptManagerConfig.update_config({"performance": {"worker_threads": True}})
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
def test_update_propagates_to_gallery(self):
PromptManagerConfig.update_config(
{"gallery": {"monitoring": {"enabled": False}}}
@@ -193,6 +230,28 @@ class TestPromptManagerConfig(unittest.TestCase):
shutil.rmtree(tmp_dir)
def test_saved_file_is_private_to_the_user(self):
import stat
tmp_dir = tempfile.mkdtemp()
try:
path = os.path.join(tmp_dir, "config.json")
PromptManagerConfig.save_to_file(path)
mode = stat.S_IMODE(os.stat(path).st_mode)
if os.name == "posix":
self.assertEqual(mode, 0o600)
with open(path) as f:
self.assertIn("database", json.load(f))
# Overwriting an existing, more open file tightens it as well
os.chmod(path, 0o644)
PromptManagerConfig.save_to_file(path)
if os.name == "posix":
self.assertEqual(stat.S_IMODE(os.stat(path).st_mode), 0o600)
finally:
import shutil
shutil.rmtree(tmp_dir, ignore_errors=True)
def test_saved_file_is_valid_json(self):
tmp = tempfile.NamedTemporaryFile(
delete=False, suffix=".json", dir=tempfile.gettempdir()
@@ -291,5 +350,59 @@ class TestIntegrationConfig(unittest.TestCase):
self.assertEqual(config1, config2)
class TestConfigPathOverride(unittest.TestCase):
"""PROMPT_MANAGER_CONFIG_PATH relocates the persisted config.json."""
ENV = "PROMPT_MANAGER_CONFIG_PATH"
def setUp(self):
self._orig_env = os.environ.pop(self.ENV, None)
self._orig = PromptManagerConfig.get_config()
def tearDown(self):
if self._orig_env is None:
os.environ.pop(self.ENV, None)
else:
os.environ[self.ENV] = self._orig_env
PromptManagerConfig.update_config(self._orig)
def test_default_config_path_is_repo_config_json(self):
repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
expected = os.path.join(repo_root, "config.json")
self.assertEqual(
os.path.normcase(PromptManagerConfig.get_config_path()),
os.path.normcase(expected),
)
def test_env_override_is_used(self):
with tempfile.TemporaryDirectory() as tmpdir:
custom = os.path.join(tmpdir, "custom.json")
os.environ[self.ENV] = custom
self.assertEqual(PromptManagerConfig.get_config_path(), custom)
def test_save_without_path_uses_env_override(self):
with tempfile.TemporaryDirectory() as tmpdir:
custom = os.path.join(tmpdir, "nested", "custom.json")
os.environ[self.ENV] = custom
PromptManagerConfig.RESULT_TIMEOUT = 42
PromptManagerConfig.save_to_file()
with open(custom) as f:
saved = json.load(f)
self.assertEqual(saved["web_ui"]["result_timeout"], 42)
def test_load_without_path_uses_env_override(self):
with tempfile.TemporaryDirectory() as tmpdir:
custom = os.path.join(tmpdir, "custom.json")
with open(custom, "w") as f:
json.dump({"web_ui": {"result_timeout": 77}}, f)
os.environ[self.ENV] = custom
PromptManagerConfig.load_from_file()
self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 77)
if __name__ == "__main__":
unittest.main()
+776 -10
View File
@@ -6,14 +6,26 @@ search, statistics, image linking, and edge cases using
an in-memory SQLite database.
"""
import csv
import json
import os
import sqlite3
from contextlib import closing
import sys
import tempfile
import types
import unittest
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.models import PromptModel # noqa: E402
from database.operations import PromptDatabase
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
from database.operations import PromptDatabase, _resolve_db_path
from utils.hashing import generate_prompt_hash
@@ -26,13 +38,12 @@ class DatabaseTestCase(unittest.TestCase):
self.db = PromptDatabase(self.temp_db.name)
def tearDown(self):
if os.path.exists(self.temp_db.name):
os.unlink(self.temp_db.name)
wal = self.temp_db.name + "-wal"
shm = self.temp_db.name + "-shm"
for f in (wal, shm):
if os.path.exists(f):
os.unlink(f)
PromptModel.close_all_instances()
assert_closed(self, self.temp_db.name)
for suffix in ("", "-wal", "-shm"):
path = self.temp_db.name + suffix
if os.path.exists(path):
os.unlink(path)
def _save(
self, text="Test prompt", category=None, tags=None, rating=None, notes=None
@@ -51,6 +62,31 @@ class DatabaseTestCase(unittest.TestCase):
class TestPromptCRUD(DatabaseTestCase):
"""Test basic create, read, update, delete operations."""
def test_save_existing_prompt_returns_its_id_when_check_raced(self):
# Two callers check the hash, both see nothing, both insert: the second
# insert hits the UNIQUE(hash) constraint and must yield the first id.
text = "raced prompt"
prompt_hash = generate_prompt_hash(text)
original = self.db.get_prompt_by_hash
calls = []
def racing_lookup(h):
calls.append(h)
return None if len(calls) == 1 else original(h)
first_id = self._save(text)
with patch.object(self.db, "get_prompt_by_hash", racing_lookup):
self.assertIsNone(self.db.get_prompt_by_hash(prompt_hash))
second_id = self.db.save_prompt(text=text, prompt_hash=prompt_hash)
self.assertEqual(second_id, first_id)
self.assertEqual(len(self.db.search_prompts(text=text)), 1)
def test_save_without_hash_never_collides(self):
first = self.db.save_prompt(text="no hash a")
second = self.db.save_prompt(text="no hash b")
self.assertNotEqual(first, second)
def test_save_and_retrieve(self):
pid = self._save(
"A beautiful sunset", category="nature", tags=["sunset", "sky"], rating=5
@@ -118,6 +154,52 @@ class TestPromptCRUD(DatabaseTestCase):
self.assertFalse(result)
class TestTagSuggestions(DatabaseTestCase):
"""suggest_tags(prefix, with_tags, limit): autocomplete narrowed by co-occurrence.
Context tags restrict suggestions to tags seen on the same prompts.
"""
def setUp(self):
super().setUp()
self._save("P1", tags=["asian", "portrait", "smile"])
self._save("P2", tags=["asian", "portrait"])
self._save("P3", tags=["portrait", "dog"])
self._save("P4", tags=["asian", "dog"])
self._save("P5", tags=["snow_day"])
def _names(self, *args, **kwargs):
return [(t["name"], t["count"]) for t in self.db.suggest_tags(*args, **kwargs)]
def test_no_context_lists_every_tag_by_frequency_then_name(self):
self.assertEqual(
self._names("", []),
[("asian", 3), ("portrait", 3), ("dog", 2), ("smile", 1), ("snow_day", 1)],
)
def test_prefix_is_case_insensitive_and_anchored_at_the_start(self):
self.assertEqual(self._names("P", []), [("portrait", 3)])
self.assertEqual(self._names("s", []), [("smile", 1), ("snow_day", 1)])
self.assertEqual(self._names("ort", []), [])
def test_context_tags_narrow_to_co_occurring_tags_and_exclude_themselves(self):
self.assertEqual(
self._names("", ["asian"]), [("portrait", 2), ("dog", 1), ("smile", 1)]
)
self.assertEqual(self._names("", ["asian", "portrait"]), [("smile", 1)])
self.assertEqual(self._names("S", ["asian"]), [("smile", 1)])
self.assertEqual(self._names("", ["ASIAN", "dog"]), [])
def test_unknown_context_tag_yields_nothing(self):
self.assertEqual(self._names("", ["unicorn"]), [])
def test_limit_and_like_wildcards(self):
self.assertEqual(len(self._names("", [], limit=2)), 2)
self.assertEqual(self._names("%", []), [])
self.assertEqual(self._names("snow_", []), [("snow_day", 1)])
self.assertEqual(self._names("snow%", []), [])
class TestTagJunctionTables(DatabaseTestCase):
"""Test normalized tag storage via junction tables."""
@@ -313,14 +395,17 @@ class TestPagination(DatabaseTestCase):
# Use a fresh empty DB
empty_db_file = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
empty_db_file.close()
empty_db = PromptDatabase(empty_db_file.name)
try:
empty_db = PromptDatabase(empty_db_file.name)
result = empty_db.get_recent_prompts(limit=10, offset=0)
self.assertEqual(result["total"], 0)
self.assertEqual(len(result["prompts"]), 0)
self.assertFalse(result["has_more"])
finally:
os.unlink(empty_db_file.name)
empty_db.close_all()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(empty_db_file.name + suffix):
os.unlink(empty_db_file.name + suffix)
def test_total_count_is_integer(self):
result = self.db.get_recent_prompts(limit=5)
@@ -431,6 +516,139 @@ class TestImageOperations(DatabaseTestCase):
self.assertEqual(len(images), 0)
class TestImageUniquenessByPath(DatabaseTestCase):
"""ComfyUI writes same-named files into different date folders."""
def _link(self, prompt_id, path):
return self.db.link_image_to_prompt(prompt_id=prompt_id, image_path=path)
def test_same_filename_in_different_folders_both_link(self):
pid = self._save("two folders")
first = self._link(pid, os.path.join("out", "2026-01-01", "a.png"))
second = self._link(pid, os.path.join("out", "2026-01-02", "a.png"))
self.assertTrue(first)
self.assertTrue(second)
self.assertNotEqual(first, second)
self.assertEqual(len(self.db.get_prompt_images(pid)), 2)
def test_same_path_spelled_differently_links_once(self):
pid = self._save("one file")
first = self._link(pid, os.path.join("out", "x", "..", "a.png"))
second = self._link(pid, os.path.join(".", "out", "a.png"))
self.assertTrue(first)
self.assertEqual(second, 0)
images = self.db.get_prompt_images(pid)
self.assertEqual(len(images), 1)
expected = os.path.normcase(
os.path.normpath(os.path.abspath(os.path.join("out", "a.png")))
)
self.assertEqual(images[0]["file_path"], expected)
def test_same_file_can_link_to_two_prompts(self):
a = self._save("prompt a")
b = self._save("prompt b")
path = os.path.join("out", "shared.png")
self.assertTrue(self._link(a, path))
self.assertTrue(self._link(b, path))
class TestLinkImageMetadataCoercion(DatabaseTestCase):
"""file_info from the request body is coerced; junk never reaches a column."""
XSS = '"><img src=x>'
def _link(self, file_info):
pid = self._save("coerce")
path = f"/out/c{len(self.db.get_prompt_images(pid))}.png"
image_id = self.db.link_image_to_prompt(pid, path, {"file_info": file_info})
self.assertGreater(image_id, 0, "image must still link")
return self.db.get_image_by_id(image_id)
def test_string_payloads_are_stored_as_null(self):
row = self._link({"size": self.XSS, "dimensions": self.XSS, "format": self.XSS})
self.assertIsNone(row["file_size"])
self.assertIsNone(row["width"])
self.assertIsNone(row["height"])
self.assertIsNone(row["format"])
def test_numbers_and_known_formats_are_kept(self):
row = self._link({"size": 1234, "dimensions": [640, 480], "format": "png"})
self.assertEqual(row["file_size"], 1234)
self.assertEqual(row["width"], 640)
self.assertEqual(row["height"], 480)
self.assertEqual(row["format"], "PNG")
def test_numeric_strings_are_accepted(self):
row = self._link({"size": "99", "dimensions": ["8", "16"], "format": "webp"})
self.assertEqual((row["file_size"], row["width"], row["height"]), (99, 8, 16))
self.assertEqual(row["format"], "WEBP")
def test_out_of_range_values_are_dropped(self):
too_big = self._link({"size": 10**12 + 1, "dimensions": [65536, 65535]})
negative = self._link({"size": -1, "dimensions": [-1, 0]})
self.assertIsNone(too_big["file_size"])
self.assertIsNone(too_big["width"])
self.assertEqual(too_big["height"], 65535)
self.assertIsNone(negative["file_size"])
self.assertIsNone(negative["width"])
self.assertEqual(negative["height"], 0)
def test_unknown_format_and_short_dimensions_are_dropped(self):
row = self._link({"dimensions": [640], "format": "svg"})
self.assertIsNone(row["width"])
self.assertIsNone(row["height"])
self.assertIsNone(row["format"])
def test_non_dict_metadata_still_links(self):
pid = self._save("junk metadata")
first = self.db.link_image_to_prompt(pid, "/out/a.png", "not a dict")
second = self.db.link_image_to_prompt(
pid, "/out/b.png", {"file_info": ["not", "a", "dict"]}
)
self.assertGreater(first, 0)
self.assertGreater(second, 0)
for row in self.db.get_prompt_images(pid):
self.assertIsNone(row["width"])
self.assertIsNone(row["format"])
class TestSearchEscapesLikeWildcards(DatabaseTestCase):
def test_percent_is_literal(self):
self._save("100% sure")
self._save("100 percent sure")
texts = [p["text"] for p in self.db.search_prompts(text="100%")]
self.assertEqual(texts, ["100% sure"])
def test_underscore_is_literal(self):
self._save("snake_case name")
self._save("snakeXcase name")
texts = [p["text"] for p in self.db.search_prompts(text="snake_case")]
self.assertEqual(texts, ["snake_case name"])
def test_backslash_is_literal(self):
self._save("path C:\\out\\img")
self._save("path C:out img")
texts = [p["text"] for p in self.db.search_prompts(text="C:\\out")]
self.assertEqual(texts, ["path C:\\out\\img"])
def test_partial_tag_filter_is_literal(self):
self._save("tagged a", tags=["50%_off"])
self._save("tagged b", tags=["50x_off"])
texts = [
p["text"] for p in self.db.search_prompts(tags=["50%"], tag_partial=True)
]
self.assertEqual(texts, ["tagged a"])
class TestEdgeCases(DatabaseTestCase):
"""Test edge cases and boundary conditions."""
@@ -489,5 +707,553 @@ class TestPreviewImages(DatabaseTestCase):
self.assertEqual(prompt["image_count"], 5)
class TestExport(DatabaseTestCase):
def setUp(self):
super().setUp()
self.out_dir = tempfile.mkdtemp()
self.addCleanup(self._remove_out_dir)
def _remove_out_dir(self):
for name in os.listdir(self.out_dir):
os.unlink(os.path.join(self.out_dir, name))
os.rmdir(self.out_dir)
def test_json_export_writes_every_prompt_with_tags(self):
self._save("first", tags=["a", "b"], category="cat")
self._save("second")
path = os.path.join(self.out_dir, "prompts.json")
self.assertTrue(self.db.export_prompts(path, "json"))
with open(path, encoding="utf-8") as fh:
data = json.load(fh)
self.assertEqual(sorted(p["text"] for p in data), ["first", "second"])
self.assertEqual(
next(p for p in data if p["text"] == "first")["tags"], ["a", "b"]
)
def test_csv_export_joins_tags(self):
self._save("first", tags=["a", "b"])
path = os.path.join(self.out_dir, "prompts.csv")
self.assertTrue(self.db.export_prompts(path, "CSV"))
with open(path, newline="", encoding="utf-8") as fh:
rows = list(csv.DictReader(fh))
self.assertEqual(len(rows), 1)
self.assertEqual(rows[0]["text"], "first")
self.assertEqual(rows[0]["tags"], "a, b")
def test_csv_export_of_empty_database_writes_nothing(self):
path = os.path.join(self.out_dir, "empty.csv")
self.assertTrue(self.db.export_prompts(path, "csv"))
self.assertFalse(os.path.exists(path))
def test_unsupported_format_returns_false(self):
self._save("x")
path = os.path.join(self.out_dir, "prompts.xml")
self.assertFalse(self.db.export_prompts(path, "xml"))
self.assertFalse(os.path.exists(path))
def test_unwritable_destination_returns_false(self):
self._save("x")
path = os.path.join(self.out_dir, "missing", "dir", "prompts.json")
self.assertFalse(self.db.export_prompts(path, "json"))
class TestDuplicateMerging(DatabaseTestCase):
def _save_dup(self, text, created_at, **kwargs):
pid = self.db.save_prompt(text=text, prompt_hash=None, **kwargs)
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
conn.execute(
"UPDATE prompts SET created_at = ? WHERE id = ?", (created_at, pid)
)
return pid
def test_cleanup_keeps_oldest_merges_metadata_and_moves_images(self):
keep = self._save_dup(
"dup", "2026-01-01T00:00:00", tags=["a"], notes="first", rating=2
)
drop = self._save_dup(
"DUP",
"2026-02-01T00:00:00",
tags=["b"],
category="cat",
notes="second",
rating=4,
)
self.db.link_image_to_prompt(drop, "/out/moved.png")
self.db.link_image_to_prompt(keep, "/out/kept.png")
removed = self.db.cleanup_duplicates()
self.assertEqual(removed, 1)
self.assertIsNone(self.db.get_prompt_by_id(drop))
merged = self.db.get_prompt_by_id(keep)
self.assertEqual(merged["category"], "cat")
self.assertEqual(merged["rating"], 4)
self.assertEqual(merged["notes"], "first | second")
self.assertEqual(sorted(merged["tags"]), ["a", "b"])
self.assertEqual(len(self.db.get_prompt_images(keep)), 2)
def test_cleanup_takes_notes_from_duplicate_when_primary_has_none(self):
keep = self._save_dup("note dup", "2026-01-01T00:00:00")
self._save_dup("note dup", "2026-02-01T00:00:00", notes="only here")
self.db.cleanup_duplicates()
self.assertEqual(self.db.get_prompt_by_id(keep)["notes"], "only here")
def test_duplicate_scans_return_empty_on_database_error(self):
with patch.object(
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
):
self.assertEqual(self.db.cleanup_duplicates(), 0)
def test_hash_duplicates_are_impossible_under_unique_hash(self):
self._save("one")
self.assertEqual(self.db.check_hash_duplicates(), [])
def test_merge_helpers_survive_database_errors(self):
pid = self._save("helper")
broken = MagicMock()
broken.execute.side_effect = sqlite3.OperationalError("disk I/O error")
self.assertEqual(self.db._merge_duplicate_metadata(broken, pid, []), {})
self.assertEqual(self.db._transfer_images_to_primary(broken, pid, [pid]), 0)
self.db._update_primary_with_merged_metadata(broken, pid, {"tags": []})
self.assertEqual(self.db.get_prompt_by_id(pid)["text"], "helper")
class TestImageQueries(DatabaseTestCase):
METADATA = {
"file_info": {"size": 123, "dimensions": [640, 480], "format": "PNG"},
"workflow": {"nodes": [1, 2]},
"prompt": {"3": {"inputs": {"cfg": float("nan")}}},
"parameters": {"steps": 20},
}
def test_link_stores_metadata_and_cleans_nan(self):
pid = self._save("with metadata")
image_id = self.db.link_image_to_prompt(pid, "/out/meta.png", self.METADATA)
image = self.db.get_image_by_id(image_id)
self.assertEqual(image["file_size"], 123)
self.assertEqual((image["width"], image["height"]), (640, 480))
self.assertEqual(image["format"], "PNG")
self.assertEqual(image["workflow_data"], {"nodes": [1, 2]})
self.assertIsNone(image["prompt_metadata"]["3"]["inputs"]["cfg"])
self.assertEqual(image["parameters"], {"steps": 20})
def test_link_rejects_temporary_invalid_and_unknown_prompt_ids(self):
self.assertEqual(self.db.link_image_to_prompt("temp_123", "/out/a.png"), 0)
self.assertEqual(self.db.link_image_to_prompt("not-an-id", "/out/a.png"), 0)
self.assertEqual(self.db.link_image_to_prompt(99999, "/out/a.png"), 0)
self.assertEqual(self.db.link_image_to_prompt(None, "/out/a.png"), 0)
def test_link_returns_zero_on_database_error(self):
pid = self._save("err")
with patch.object(
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
):
self.assertEqual(self.db.link_image_to_prompt(pid, "/out/a.png"), 0)
def test_recent_all_and_search_images_carry_prompt_text(self):
a = self._save("alpha prompt", tags=["x"])
b = self._save("beta prompt")
self.db.link_image_to_prompt(a, "/out/a1.png")
self.db.link_image_to_prompt(b, "/out/b1.png")
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
conn.execute(
"UPDATE generated_images SET generation_time = '2026-01-01 00:00:00'"
" WHERE filename = 'a1.png'"
)
conn.execute(
"UPDATE generated_images SET generation_time = '2026-02-01 00:00:00'"
" WHERE filename = 'b1.png'"
)
recent = self.db.get_recent_images(limit=10)
self.assertEqual(
[i["prompt_text"] for i in recent], ["beta prompt", "alpha prompt"]
)
everything = self.db.get_all_images()
self.assertEqual(len(everything), 2)
by_text = {i["prompt_text"]: i for i in everything}
self.assertEqual(by_text["alpha prompt"]["prompt_tags"], ["x"])
self.assertEqual(by_text["beta prompt"]["prompt_tags"], [])
page = self.db.get_all_images(limit=1, offset=1)
self.assertEqual([i["prompt_text"] for i in page], ["alpha prompt"])
found = self.db.search_images_by_prompt("alpha")
self.assertEqual([i["filename"] for i in found], ["a1.png"])
def test_image_search_treats_like_wildcards_literally(self):
literal = self._save("50% off")
other = self._save("50 percent off")
self.db.link_image_to_prompt(literal, "/out/literal.png")
self.db.link_image_to_prompt(other, "/out/other.png")
found = self.db.search_images_by_prompt("50%")
self.assertEqual([i["filename"] for i in found], ["literal.png"])
def test_delete_image(self):
pid = self._save("del")
image_id = self.db.link_image_to_prompt(pid, "/out/del.png")
self.assertTrue(self.db.delete_image(image_id))
self.assertFalse(self.db.delete_image(image_id))
self.assertIsNone(self.db.get_image_by_id(image_id))
def test_cleanup_missing_images_keeps_files_that_exist(self):
pid = self._save("files")
fd, existing = tempfile.mkstemp(suffix=".png")
os.close(fd)
self.addCleanup(os.unlink, existing)
self.db.link_image_to_prompt(pid, existing)
self.db.link_image_to_prompt(
pid, os.path.join(tempfile.gettempdir(), "gone.png")
)
self.assertEqual(self.db.cleanup_missing_images(), 1)
images = self.db.get_prompt_images(pid)
self.assertEqual([i["image_path"] for i in images], [existing])
def test_image_prompt_info_by_exact_normalised_or_basename_path(self):
pid = self._save("lookup", tags=["t"], category="c", rating=3, notes="n")
stored = os.path.join("out", "sub", "look.png")
self.db.link_image_to_prompt(pid, stored, self.METADATA)
exact = self.db.get_image_prompt_info(stored)
self.assertEqual(exact["prompt_id"], pid)
self.assertEqual(exact["tags"], ["t"])
self.assertEqual(exact["workflow_data"], {"nodes": [1, 2]})
normalised = self.db.get_image_prompt_info(
os.path.join("out", "x", "..", "sub", "look.png")
)
self.assertEqual(normalised["prompt_id"], pid)
basename = self.db.get_image_prompt_info(os.path.join("elsewhere", "look.png"))
self.assertEqual(basename["prompt_id"], pid)
self.assertIsNone(self.db.get_image_prompt_info("nothing.png"))
self.assertEqual(self.db.get_prompt_id_for_image(stored), pid)
self.assertEqual(
self.db.get_prompt_id_for_image(
os.path.join(".", "out", "sub", "look.png")
),
pid,
)
self.assertIsNone(self.db.get_prompt_id_for_image("nothing.png"))
def test_image_row_with_broken_json_yields_empty_dicts(self):
pid = self._save("broken")
image_id = self.db.link_image_to_prompt(pid, "/out/broken.png")
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
conn.execute(
"UPDATE generated_images SET workflow_data = '{not json', "
"prompt_metadata = '{\"k\": 1}', parameters = NULL WHERE id = ?",
(image_id,),
)
image = self.db.get_image_by_id(image_id)
self.assertEqual(image["workflow_data"], {})
self.assertEqual(image["prompt_metadata"], {"k": 1})
self.assertEqual(image["parameters"], {})
info = self.db.get_image_prompt_info("/out/broken.png")
self.assertIsNone(info["workflow_data"])
self.assertEqual(info["prompt_metadata"], {"k": 1})
class TestMaintenanceOperations(DatabaseTestCase):
def test_statistics_count_everything(self):
a = self._save("a", category="cat", tags=["t1", "t2"], rating=4)
self._save("b", category=" cat ", rating=2)
self._save("c")
self.db.link_image_to_prompt(a, "/out/a.png")
self.db.link_image_to_prompt(a, "/out/b.png")
stats = self.db.get_statistics()
self.assertEqual(stats["total_prompts"], 3)
self.assertEqual(stats["total_categories"], 1)
self.assertEqual(stats["average_rating"], 3.0)
self.assertEqual(stats["total_tags"], 2)
self.assertEqual(stats["total_images"], 2)
self.assertEqual(stats["images_with_prompts"], 1)
def test_update_text_and_rating(self):
pid = self._save("old text", rating=1)
self.assertTrue(self.db.update_prompt_text(pid, "new text"))
self.assertTrue(self.db.update_prompt_rating(pid, 5))
prompt = self.db.get_prompt_by_id(pid)
self.assertEqual((prompt["text"], prompt["rating"]), ("new text", 5))
self.assertFalse(self.db.update_prompt_text(99999, "x"))
self.assertFalse(self.db.update_prompt_rating(99999, 3))
def test_update_metadata_validation(self):
pid = self._save("meta")
self.assertFalse(self.db.update_prompt_metadata(pid))
with self.assertRaises(ValueError):
self.db.update_prompt_metadata(pid, rating=7)
self.assertTrue(self.db.update_prompt_metadata(pid, tags=["only tags"]))
self.assertEqual(self.db.get_prompt_by_id(pid)["tags"], ["only tags"])
def test_save_prompt_validation(self):
with self.assertRaises(ValueError):
self.db.save_prompt(text=" ")
with self.assertRaises(ValueError):
self.db.save_prompt(text="rated", rating=0)
def test_bulk_operations(self):
a = self._save("a", tags=["keep"])
b = self._save("b")
c = self._save("c")
self.db.link_image_to_prompt(c, "/out/c.png")
self.assertEqual(self.db.bulk_add_tags([a, b, 99999], ["keep", "new"]), 2)
self.assertEqual(self.db.bulk_add_tags([a], ["keep"]), 0)
self.assertEqual(sorted(self.db.get_prompt_by_id(a)["tags"]), ["keep", "new"])
self.assertEqual(self.db.bulk_set_category([a, b, 99999], "bulk"), 2)
self.assertEqual(self.db.get_prompt_by_id(b)["category"], "bulk")
self.assertEqual(self.db.bulk_delete_prompts([b, c, 99999]), 2)
self.assertIsNone(self.db.get_prompt_by_id(c))
self.assertEqual(self.db.get_prompt_images(c), [])
def test_prune_orphaned_prompts_spares_protected_and_illustrated(self):
orphan = self._save("orphan")
protected = self._save("protected", tags=["__protected__"])
illustrated = self._save("illustrated")
self.db.link_image_to_prompt(illustrated, "/out/i.png")
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
self.assertIsNone(self.db.get_prompt_by_id(orphan))
self.assertIsNotNone(self.db.get_prompt_by_id(protected))
self.assertIsNotNone(self.db.get_prompt_by_id(illustrated))
self.assertEqual(self.db.prune_orphaned_prompts(), 0)
def test_consistency_check_reports_orphaned_rows(self):
self.assertEqual(self.db.check_consistency(), [])
pid = self._save("to orphan", tags=["t"])
self.db.link_image_to_prompt(pid, "/out/o.png")
self.db.close()
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
conn.execute("PRAGMA foreign_keys = OFF")
conn.execute("DELETE FROM prompts WHERE id = ?", (pid,))
issues = self.db.check_consistency()
self.assertEqual(len(issues), 2)
self.assertTrue(any("prompt_tags" in i for i in issues))
self.assertTrue(any(i.startswith("Image ") for i in issues))
def test_subfolders_relative_to_roots_with_ancestors(self):
pid = self._save("folders")
root = os.path.join(os.sep, "gallery")
self.db.link_image_to_prompt(
pid, os.path.join(root, "2026", "08-Aug", "2026-08-06", "a.png")
)
self.db.link_image_to_prompt(pid, os.path.join(root, "b.png"))
self.db.link_image_to_prompt(
pid, os.path.join(os.sep, "elsewhere", "deep", "c.png")
)
self.db.link_image_to_prompt(pid, "loose.png")
flat = self.db.get_prompt_subfolders(root_dirs=[root])
self.assertIn("2026/08-Aug/2026-08-06", flat)
self.assertIn(
os.path.join(os.sep, "elsewhere", "deep").replace(os.sep, "/"), flat
)
self.assertNotIn(".", flat)
nested = self.db.get_prompt_subfolders(root_dirs=[root], include_ancestors=True)
self.assertIn("2026", nested)
self.assertIn("2026/08-Aug", nested)
without_roots = self.db.get_prompt_subfolders()
self.assertIn(
os.path.join(root, "2026", "08-Aug", "2026-08-06").replace(os.sep, "/"),
without_roots,
)
def test_prompts_by_tags_and_or_modes(self):
both = self._save("both", tags=["x", "y"])
only_x = self._save("only x", tags=["x"])
self._save("none")
self.assertEqual(self.db.get_prompts_by_tags([])["total"], 0)
either = self.db.get_prompts_by_tags(["x", "y"], mode="or")
self.assertEqual(sorted(p["id"] for p in either["prompts"]), [both, only_x])
self.assertEqual(either["total"], 2)
all_of = self.db.get_prompts_by_tags(["x", "y"], mode="and", limit=1)
self.assertEqual([p["id"] for p in all_of["prompts"]], [both])
self.assertFalse(all_of["has_more"])
def test_tag_counts_with_search_and_sorts(self):
self._save("1", tags=["beta", "alpha"])
self._save("2", tags=["beta"])
by_count = self.db.get_tags_with_counts(sort="count_desc")
self.assertEqual([t["name"] for t in by_count["tags"]], ["beta", "alpha"])
ascending = self.db.get_tags_with_counts(sort="count_asc")
self.assertEqual([t["name"] for t in ascending["tags"]], ["alpha", "beta"])
reverse = self.db.get_tags_with_counts(sort="alpha_desc")
self.assertEqual([t["name"] for t in reverse["tags"]], ["beta", "alpha"])
searched = self.db.get_tags_with_counts(search="alp")
self.assertEqual(searched["total"], 1)
self.assertEqual(searched["tags"][0]["count"], 1)
def test_tag_management_edge_cases(self):
a = self._save("a", tags=["old", "target"])
b = self._save("b", tags=["old"])
merged = self.db.rename_tag_all_prompts("old", "target")
self.assertEqual(merged["affected_count"], 2)
self.assertEqual(self.db.get_prompt_by_id(a)["tags"], ["target"])
self.assertEqual(self.db.get_prompt_by_id(b)["tags"], ["target"])
self.assertEqual(
self.db.rename_tag_all_prompts("missing", "x")["affected_count"], 0
)
self.assertEqual(self.db.delete_tag_all_prompts("missing")["affected_count"], 0)
with self.assertRaises(ValueError):
self.db.rename_tag_all_prompts(" ", "x")
with self.assertRaises(ValueError):
self.db.rename_tag_all_prompts("x", "")
with self.assertRaises(ValueError):
self.db.delete_tag_all_prompts("")
with self.assertRaises(ValueError):
self.db.merge_tags([], "t")
with self.assertRaises(ValueError):
self.db.merge_tags(["a"], " ")
result = self.db.merge_tags(["missing", "target"], "final")
self.assertEqual(result["tags_merged"], 1)
self.assertEqual(result["affected_count"], 2)
self.assertEqual(self.db.get_all_tags(), ["final"])
def test_untagged_and_category_helpers(self):
self._save("tagged", tags=["t"], category=" ")
untagged = self._save("plain", category="shown")
self.assertEqual(self.db.get_untagged_prompts_count(), 1)
page = self.db.get_untagged_prompts(limit=5)
self.assertEqual([p["id"] for p in page["prompts"]], [untagged])
self.assertEqual(self.db.get_all_categories(), ["shown"])
self.assertEqual(self.db.delete_prompts_by_category("nothing"), 0)
class TestRowConversion(DatabaseTestCase):
def test_legacy_json_tags_column_is_parsed(self):
convert = self.db._row_to_dict
self.assertEqual(convert({"tags": '["a", "b"]'})["tags"], ["a", "b"])
self.assertEqual(convert({"tags": '"a, b"'})["tags"], ["a", "b"])
self.assertEqual(convert({"tags": "5"})["tags"], [])
self.assertEqual(convert({"tags": "{not json"})["tags"], [])
self.assertEqual(convert({"tags": None})["tags"], [])
self.assertEqual(convert({"_tag_list": None})["tags"], [])
def test_search_date_filters(self):
pid = self._save("dated")
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
conn.execute(
"UPDATE prompts SET created_at = '2026-05-01T00:00:00' WHERE id = ?",
(pid,),
)
self.assertEqual(
len(self.db.search_prompts(date_from="2026-04-01", date_to="2026-06-01")), 1
)
self.assertEqual(len(self.db.search_prompts(date_from="2026-06-01")), 0)
self.assertEqual(len(self.db.search_prompts(date_to="2026-04-01")), 0)
self.assertEqual(len(self.db.search_prompts(rating_max=3)), 0)
class TestResolveDbPathFromConfig(unittest.TestCase):
def test_default_path_comes_from_config_when_importable(self):
tmpdir = tempfile.mkdtemp()
self.addCleanup(os.rmdir, tmpdir)
fake = types.ModuleType("py.config")
fake.PromptManagerConfig = types.SimpleNamespace(
DEFAULT_DB_PATH=os.path.join(tmpdir, "configured.db")
)
with patch.dict(sys.modules, {"py.config": fake}):
self.assertEqual(
_resolve_db_path(None), os.path.join(tmpdir, "configured.db")
)
def test_falls_back_to_prompts_db_when_config_cannot_import(self):
# Outside ComfyUI py.config fails to import (needs the server module)
with patch.dict(sys.modules, {"py.config": None}):
resolved = _resolve_db_path(None)
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
self.assertEqual(resolved, os.path.join(root, "prompts.db"))
class TestModelHousekeeping(DatabaseTestCase):
def test_database_info_reports_counts_and_handles_errors(self):
self._save("rated", rating=4, category="c")
info = self.db.model.get_database_info()
self.assertEqual(info["total_prompts"], 1)
self.assertEqual(info["unique_categories"], 1)
self.assertEqual(info["average_rating"], 4.0)
self.assertGreater(info["database_size_bytes"], 0)
with patch.object(
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
):
self.assertEqual(self.db.model.get_database_info(), {})
def test_vacuum_logs_errors_instead_of_raising(self):
self._save("x")
self.db.model.vacuum_database()
with patch(
"database.models.sqlite3.connect", side_effect=sqlite3.OperationalError
):
self.db.model.vacuum_database() # logged, not raised
self.assertEqual(len(self.db.search_prompts()), 1)
def test_close_tolerates_a_connection_that_fails_to_close(self):
real = self.db.model.get_connection()
real.close()
failing = MagicMock()
failing.close.side_effect = sqlite3.ProgrammingError("already closed")
self.db.model._local.conn = failing
self.db.close() # logged, not raised
self.assertTrue(failing.close.called)
self.assertIsNot(self.db.model.get_connection(), failing)
if __name__ == "__main__":
unittest.main()
class TestImagePaging(DatabaseTestCase):
"""Image listings page in the database, not by over-fetching in the API."""
def _link_many(self, count):
pid = self._save("paged prompt")
image_ids = [
self.db.link_image_to_prompt(pid, f"/out/img{i:03d}.png")
for i in range(count)
]
# Stamp distinct times afterwards through the model's own connection,
# so no second connection holds a write lock while linking.
with self.db.model.get_connection() as conn:
for i, image_id in enumerate(image_ids):
conn.execute(
"UPDATE generated_images SET generation_time = ? WHERE id = ?",
(f"2026-01-01T00:00:{i:02d}", image_id),
)
return pid
def test_recent_images_offset_skips_newest_rows(self):
self._link_many(5)
newest_first = [img["filename"] for img in self.db.get_recent_images(limit=5)]
page = self.db.get_recent_images(limit=2, offset=2)
self.assertEqual([img["filename"] for img in page], newest_first[2:4])
self.assertEqual(self.db.get_recent_images(limit=2, offset=10), [])
def test_search_images_by_prompt_is_bounded_and_pages(self):
self._link_many(5)
first = self.db.search_images_by_prompt("paged", limit=2)
second = self.db.search_images_by_prompt("paged", limit=2, offset=2)
self.assertEqual(len(first), 2)
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)
+301
View File
@@ -0,0 +1,301 @@
"""Backup, verification and restore of the prompts database.
The database runs in WAL mode, so committed rows can sit in prompts.db-wal
for a long time. A backup must include them, a restore must reject files
that are not healthy SQLite databases, and the live database must stay
usable afterwards.
"""
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.models import PromptModel
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
def _count_prompts(path):
"""Row count read through a brand-new connection (closed before return)."""
conn = sqlite3.connect(path)
try:
return conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0]
finally:
conn.close()
def _texts(path):
conn = sqlite3.connect(path)
try:
return sorted(r[0] for r in conn.execute("SELECT text FROM prompts"))
finally:
conn.close()
class BackupTestCase(unittest.TestCase):
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.path = os.path.join(self.tmpdir, "prompts.db")
self.db = PromptDatabase(self.path)
self.addCleanup(self._cleanup)
def _cleanup(self):
self.db.close_all()
for name in os.listdir(self.tmpdir):
os.unlink(os.path.join(self.tmpdir, name))
os.rmdir(self.tmpdir)
def _save(self, db, text):
return db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
def _make_source_db(self, texts):
"""A second database file holding the given prompts, fully closed."""
src_path = os.path.join(self.tmpdir, "source.db")
src = PromptDatabase(src_path)
for text in texts:
self._save(src, text)
src.close_all()
return src_path
class TestBackupDatabase(BackupTestCase):
def test_backup_contains_rows_that_only_exist_in_the_wal(self):
for n in range(5):
self._save(self.db, f"wal prompt {n}")
wal = self.path + "-wal"
self.assertTrue(os.path.exists(wal) and os.path.getsize(wal) > 0)
dst = os.path.join(self.tmpdir, "backup.db")
self.assertTrue(self.db.model.backup_database(dst))
self.assertEqual(_count_prompts(dst), 5)
self.assertFalse(os.path.exists(dst + "-wal"))
# The live database is untouched and still usable
self.assertEqual(len(self.db.search_prompts(text="wal prompt")), 5)
def test_backup_overwrites_a_stale_destination(self):
self._save(self.db, "one")
dst = os.path.join(self.tmpdir, "backup.db")
with open(dst, "wb") as fh:
fh.write(b"not a database at all")
self.assertTrue(self.db.model.backup_database(dst))
self.assertEqual(_count_prompts(dst), 1)
def test_backup_onto_the_live_file_is_refused(self):
self._save(self.db, "one")
self.assertFalse(self.db.model.backup_database(self.path))
self.assertEqual(len(self.db.search_prompts(text="one")), 1)
def test_backup_to_an_unwritable_location_returns_false(self):
dst = os.path.join(self.tmpdir, "missing", "dir", "backup.db")
self.assertFalse(self.db.model.backup_database(dst))
class TestVerifyDatabaseFile(BackupTestCase):
def test_text_file_is_rejected_without_raising(self):
path = os.path.join(self.tmpdir, "notes.txt")
with open(path, "w", encoding="utf-8") as fh:
fh.write("just some text, definitely not sqlite")
ok, reason = PromptModel.verify_database_file(path)
self.assertFalse(ok)
self.assertIn("SQLite", reason)
def test_missing_file_is_rejected(self):
ok, reason = PromptModel.verify_database_file(
os.path.join(self.tmpdir, "nope.db")
)
self.assertFalse(ok)
self.assertTrue(reason)
def test_database_without_prompts_table_is_rejected(self):
path = os.path.join(self.tmpdir, "other.db")
conn = sqlite3.connect(path)
conn.execute("CREATE TABLE something (id INTEGER)")
conn.commit()
conn.close()
ok, reason = PromptModel.verify_database_file(path)
self.assertFalse(ok)
self.assertIn("prompts", reason)
def test_prompts_table_missing_required_columns_is_rejected(self):
path = os.path.join(self.tmpdir, "partial.db")
conn = sqlite3.connect(path)
conn.execute("CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT)")
conn.commit()
conn.close()
ok, reason = PromptModel.verify_database_file(path)
self.assertFalse(ok)
self.assertIn("created_at", reason)
def test_corrupted_database_fails_integrity_check(self):
path = self._make_source_db([f"prompt {n}" for n in range(40)])
with open(path, "r+b") as fh:
page_size = int.from_bytes(fh.read(18)[16:18], "big") or 4096
fh.seek(page_size * 2) # scribble over pages 3 and 4 entirely
fh.write(b"\xff" * (page_size * 2))
ok, reason = PromptModel.verify_database_file(path)
self.assertFalse(ok)
self.assertTrue(reason)
def test_valid_database_passes(self):
path = self._make_source_db(["a", "b"])
ok, reason = PromptModel.verify_database_file(path)
self.assertTrue(ok, reason)
def test_verification_does_not_lock_or_modify_the_file(self):
path = self._make_source_db(["a"])
before = os.path.getsize(path)
PromptModel.verify_database_file(path)
os.unlink(path) # would fail on Windows if a connection were left open
self.assertEqual(before, before)
class TestRestoreFromFile(BackupTestCase):
def test_restore_replaces_content_and_backs_up_the_old_database(self):
self._save(self.db, "old prompt")
src = self._make_source_db(["restored one", "restored two"])
backup_path = self.db.model.restore_from_file(src)
self.assertTrue(backup_path.startswith(self.path + ".backup_"))
self.assertTrue(os.path.exists(backup_path))
self.assertEqual(_texts(backup_path), ["old prompt"])
# The restore goes through SQLite's backup API into the live file, so
# the WAL stays valid (and shared with other instances) rather than
# being deleted under them; the content is what matters.
self.assertEqual(_texts(self.path), ["restored one", "restored two"])
def test_database_is_usable_from_every_thread_after_restore(self):
self._save(self.db, "old prompt")
src = self._make_source_db(["restored"])
ready = threading.Event()
proceed = threading.Event()
results = {}
def worker():
self.db.search_prompts(text="old") # opens this thread's connection
ready.set()
proceed.wait(5)
try:
results["texts"] = [p["text"] for p in self.db.search_prompts()]
except Exception as exc: # noqa: BLE001 - reported through results
results["error"] = exc
finally:
self.db.close()
thread = threading.Thread(target=worker)
thread.start()
ready.wait(5)
self.db.model.restore_from_file(src)
proceed.set()
thread.join(5)
self.assertNotIn("error", results)
self.assertEqual(results["texts"], ["restored"])
self.assertEqual([p["text"] for p in self.db.search_prompts()], ["restored"])
self.assertEqual(self._save(self.db, "after restore"), 2)
def test_restore_applies_migrations_to_an_old_schema(self):
src = os.path.join(self.tmpdir, "legacy.db")
conn = sqlite3.connect(src)
# The 3.0-era schema: workflow_name still present, no usage columns
conn.execute(
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT,"
" text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP,"
" category TEXT, tags TEXT, rating INTEGER, notes TEXT,"
" hash TEXT UNIQUE, workflow_name TEXT)"
)
conn.execute(
"INSERT INTO prompts (text, created_at, tags, workflow_name)"
" VALUES ('legacy', '2024-01-01T00:00:00', '[\"old\"]', 'wf')"
)
conn.commit()
conn.close()
self.db.model.restore_from_file(src)
prompt = self.db.get_prompt_by_id(1)
self.assertEqual(prompt["text"], "legacy")
self.assertEqual(prompt["tags"], ["old"])
self.assertEqual(prompt["run_count"], 1)
self.assertNotIn("workflow_name", prompt)
self.assertEqual(self.db.get_prompt_images(1), [])
def test_restore_without_a_live_database_returns_empty_backup_path(self):
src = self._make_source_db(["fresh"])
self.db.close_all()
os.unlink(self.path)
for suffix in ("-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
self.assertEqual(self.db.model.restore_from_file(src), "")
self.assertEqual(_texts(self.path), ["fresh"])
if __name__ == "__main__":
unittest.main()
class TestConnectionsAcrossInstances(BackupTestCase):
"""Several PromptModel instances share one file (nodes, API, monitor)."""
def test_close_all_closes_connections_opened_on_other_threads(self):
opened = []
def open_one():
opened.append(self.db.model.get_connection())
threads = [threading.Thread(target=open_one) for _ in range(3)]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(len(opened), 3)
self.db.close_all()
for conn in opened:
with self.assertRaises(sqlite3.ProgrammingError):
conn.execute("SELECT 1")
def test_restore_is_seen_by_another_instance_holding_an_open_connection(self):
self.db.save_prompt(text="before", prompt_hash=generate_prompt_hash("before"))
snapshot = os.path.join(self.tmpdir, "snapshot.db")
other_db = PromptDatabase(snapshot)
for i in range(3):
other_db.save_prompt(
text=f"after {i}", prompt_hash=generate_prompt_hash(f"after {i}")
)
other_db.close_all()
# A second instance on the live file, with a connection kept open on
# a worker thread for the whole restore (like the image monitor).
bystander = PromptDatabase(self.path)
ready, go, seen = threading.Event(), threading.Event(), []
def hold_and_read():
bystander.model.get_connection().execute("SELECT COUNT(*) FROM prompts")
ready.set()
go.wait(10)
seen.append(len(bystander.search_prompts(limit=50)))
worker = threading.Thread(target=hold_and_read)
worker.start()
self.assertTrue(ready.wait(10))
try:
self.db.model.restore_from_file(snapshot)
finally:
go.set()
worker.join(10)
bystander.close_all()
self.assertEqual(seen, [3])
self.assertEqual(sorted(_texts(self.path)), ["after 0", "after 1", "after 2"])
+130
View File
@@ -0,0 +1,130 @@
"""Concurrent access to PromptDatabase from several threads.
ComfyUI touches the database from the server loop, executor threads, the
watchdog image monitor and the queue hook at the same time. Every thread must
get its own connection so one thread's commit or rollback never lands on
another thread's half-done write.
"""
import os
import sys
import tempfile
import threading
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
THREADS = 8
CALLS_PER_THREAD = 50
class ConcurrencyTestCase(unittest.TestCase):
def setUp(self):
fd, self.path = tempfile.mkstemp(suffix=".db")
os.close(fd)
self.db = PromptDatabase(self.path)
self.addCleanup(self._cleanup)
def _cleanup(self):
self.db.close_all()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
def _run_threads(self, target):
errors = []
def wrapped(index):
try:
target(index)
except Exception as exc: # noqa: BLE001 - collected for the assertion
errors.append(exc)
finally:
self.db.close()
threads = [threading.Thread(target=wrapped, args=(i,)) for i in range(THREADS)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
return errors
class TestPerThreadConnections(ConcurrencyTestCase):
def test_each_thread_gets_its_own_connection(self):
main_conn = self.db.model.get_connection()
self.assertIs(self.db.model.get_connection(), main_conn)
seen = []
def grab(_index):
first = self.db.model.get_connection()
seen.append((first, self.db.model.get_connection()))
self.assertEqual(self._run_threads(grab), [])
self.assertEqual(len(seen), THREADS)
for first, second in seen:
self.assertIs(first, second)
self.assertIsNot(first, main_conn)
self.assertEqual(len({id(first) for first, _ in seen}), THREADS)
def test_close_releases_only_the_calling_threads_connection(self):
main_conn = self.db.model.get_connection()
other = []
def grab(_index):
other.append(self.db.model.get_connection())
self.db.close()
self.assertEqual(self._run_threads(grab), [])
# The main thread's connection survives closes made by other threads
self.assertIs(self.db.model.get_connection(), main_conn)
main_conn.execute("SELECT 1").fetchone()
self.db.close()
self.assertIsNot(self.db.model.get_connection(), main_conn)
class TestConcurrentWrites(ConcurrencyTestCase):
def test_distinct_prompts_saved_and_counted_from_many_threads(self):
def work(index):
for n in range(CALLS_PER_THREAD):
text = f"thread {index} prompt {n}"
pid = self.db.save_prompt(
text=text, prompt_hash=generate_prompt_hash(text)
)
self.assertTrue(self.db.record_prompt_use(pid))
self.assertEqual(self._run_threads(work), [])
info = self.db.model.get_database_info()
self.assertEqual(info["total_prompts"], THREADS * CALLS_PER_THREAD)
prompts = self.db.search_prompts(limit=THREADS * CALLS_PER_THREAD + 1)
self.assertEqual(len(prompts), THREADS * CALLS_PER_THREAD)
self.assertTrue(all(p["run_count"] == 1 for p in prompts))
def test_saving_the_same_prompt_from_many_threads_keeps_one_row(self):
text = "the same prompt from every thread"
prompt_hash = generate_prompt_hash(text)
ids = []
def work(_index):
for _ in range(CALLS_PER_THREAD):
pid = self.db.save_prompt(text=text, prompt_hash=prompt_hash)
ids.append(pid)
self.assertTrue(self.db.record_prompt_use(pid))
self.assertEqual(self._run_threads(work), [])
self.assertEqual(len(ids), THREADS * CALLS_PER_THREAD)
self.assertEqual(len(set(ids)), 1)
self.assertEqual(self.db.model.get_database_info()["total_prompts"], 1)
prompt = self.db.get_prompt_by_hash(prompt_hash)
self.assertEqual(prompt["id"], ids[0])
self.assertEqual(prompt["run_count"], THREADS * CALLS_PER_THREAD)
if __name__ == "__main__":
unittest.main()
+318
View File
@@ -0,0 +1,318 @@
"""Schema initialisation and migration behaviour of PromptModel.
Every PromptDatabase() construction used to re-run table creation and all
migrations, migrations swallowed their errors, and a rebuild that failed
half-way left a prompts_new table behind that made every later start fail.
"""
import os
import sqlite3
import sys
import tempfile
import unittest
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.models import PromptModel
from database.operations import PromptDatabase
from database.models import normalize_image_path # noqa: E402
LEGACY_PROMPTS_TABLE = (
"CREATE TABLE prompts (id INTEGER PRIMARY KEY AUTOINCREMENT,"
" text TEXT NOT NULL, created_at TIMESTAMP, updated_at TIMESTAMP,"
" category TEXT, tags TEXT, rating INTEGER, notes TEXT,"
" hash TEXT UNIQUE, workflow_name TEXT)"
)
LEGACY_IMAGES_TABLE = (
"CREATE TABLE generated_images (id INTEGER PRIMARY KEY AUTOINCREMENT,"
" prompt_id INTEGER NOT NULL, image_path TEXT NOT NULL,"
" filename TEXT NOT NULL, generation_time TIMESTAMP, file_size INTEGER,"
" width INTEGER, height INTEGER, format TEXT, workflow_data TEXT,"
" prompt_metadata TEXT, parameters TEXT,"
" FOREIGN KEY (prompt_id) REFERENCES prompts(id) ON DELETE CASCADE)"
)
class MigrationTestCase(unittest.TestCase):
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.path = os.path.join(self.tmpdir, "prompts.db")
self.dbs = []
self.addCleanup(self._cleanup)
def _cleanup(self):
for db in self.dbs:
db.close_all()
PromptModel.reset_schema_cache()
for name in os.listdir(self.tmpdir):
os.unlink(os.path.join(self.tmpdir, name))
os.rmdir(self.tmpdir)
def _open(self):
db = PromptDatabase(self.path)
self.dbs.append(db)
return db
def _legacy_db(self, statements):
"""A 3.0-era database file (workflow_name column, no usage columns)."""
conn = sqlite3.connect(self.path)
try:
conn.execute(LEGACY_PROMPTS_TABLE)
conn.execute(LEGACY_IMAGES_TABLE)
for sql in statements:
conn.execute(sql)
conn.commit()
finally:
conn.close()
def _columns(self, table):
conn = sqlite3.connect(self.path)
try:
return [r[1] for r in conn.execute(f"PRAGMA table_info({table})")]
finally:
conn.close()
class TestSchemaInitOncePerProcess(MigrationTestCase):
def _spy(self):
original = PromptModel._create_tables
calls = []
def spy(model, conn):
calls.append(model.db_path)
return original(model, conn)
return patch.object(PromptModel, "_create_tables", spy), calls
def test_second_construction_on_same_path_skips_schema_init(self):
patcher, calls = self._spy()
with patcher:
self._open()
self._open()
# Same file spelled differently still counts as the same path
self.dbs.append(
PromptDatabase(os.path.join(self.tmpdir, ".", "prompts.db"))
)
self.assertEqual(len(calls), 1)
def test_reset_schema_cache_forces_init_again(self):
patcher, calls = self._spy()
with patcher:
self._open()
PromptModel.reset_schema_cache()
self._open()
self.assertEqual(len(calls), 2)
def test_deleted_database_file_is_recreated_on_next_construction(self):
db = self._open()
db.close_all()
os.unlink(self.path)
for suffix in ("-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
reopened = self._open()
self.assertEqual(reopened.model.get_database_info()["total_prompts"], 0)
self.assertIn("prompts", self._tables())
def test_close_then_unlink_succeeds(self):
db = self._open()
db.save_prompt(text="x", prompt_hash="x")
db.close_all()
os.unlink(self.path)
self.assertFalse(os.path.exists(self.path))
def _tables(self):
conn = sqlite3.connect(self.path)
try:
return [
r[0]
for r in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table'"
)
]
finally:
conn.close()
class TestIdempotentMigrations(MigrationTestCase):
def test_stale_prompts_new_from_a_failed_rebuild_does_not_block_startup(self):
self._legacy_db(
[
"INSERT INTO prompts (text, hash, workflow_name)"
" VALUES ('a', 'h1', 'wf')",
"CREATE TABLE prompts_new (id INTEGER PRIMARY KEY, junk TEXT)",
"INSERT INTO prompts_new (junk) VALUES ('leftover')",
]
)
db = self._open()
self.assertNotIn("workflow_name", self._columns("prompts"))
self.assertEqual(db.get_prompt_by_id(1)["text"], "a")
self.assertNotIn("prompts_new", self._table_names())
def test_rebuild_keeps_linked_images(self):
# DROP TABLE prompts with foreign keys ON cascades into generated_images
self._legacy_db(
[
"INSERT INTO prompts (text, hash, workflow_name)"
" VALUES ('a', 'h1', 'wf')",
"INSERT INTO generated_images (prompt_id, image_path, filename)"
" VALUES (1, '/out/a.png', 'a.png')",
]
)
db = self._open()
self.assertNotIn("workflow_name", self._columns("prompts"))
self.assertEqual(len(db.get_prompt_images(1)), 1)
def test_migration_error_is_logged_and_raised(self):
# The legacy table has no CHECK on rating; the rebuilt one does, so the
# copy into prompts_new fails inside the migration.
self._legacy_db(
[
"INSERT INTO prompts (text, hash, rating, workflow_name)"
" VALUES ('a', 'h1', 9, 'wf')",
]
)
with self.assertLogs("prompt_manager.database.models", level="ERROR") as logs:
with self.assertRaises(sqlite3.Error):
self._open()
self.assertTrue(any("igration" in line for line in logs.output))
# Nothing was half-applied: the old table is intact and usable
self.assertIn("workflow_name", self._columns("prompts"))
self.assertNotIn("prompts_new", self._table_names())
def test_text_prompt_ids_in_images_become_integers(self):
conn = sqlite3.connect(self.path)
try:
conn.execute(LEGACY_PROMPTS_TABLE)
conn.execute(
LEGACY_IMAGES_TABLE.replace("prompt_id INTEGER", "prompt_id TEXT")
)
conn.execute(
"INSERT INTO prompts (text, hash, workflow_name)"
" VALUES ('a', 'h1', 'wf')"
)
for prompt_id, name in (
("1", "ok.png"),
("", "blank.png"),
("7", "gone.png"),
):
conn.execute(
"INSERT INTO generated_images (prompt_id, image_path, filename)"
" VALUES (?, ?, ?)",
(prompt_id, "/out/" + name, name),
)
conn.commit()
finally:
conn.close()
db = self._open()
conn = sqlite3.connect(self.path)
try:
types_ = {
r[1]: r[2] for r in conn.execute("PRAGMA table_info(generated_images)")
}
finally:
conn.close()
self.assertEqual(types_["prompt_id"], "INTEGER")
self.assertEqual([i["filename"] for i in db.get_prompt_images(1)], ["ok.png"])
self.assertEqual(db.get_statistics()["total_images"], 1)
def test_migrated_database_opens_cleanly_a_second_time(self):
self._legacy_db(
["INSERT INTO prompts (text, hash, workflow_name) VALUES ('a', 'h1', 'wf')"]
)
self._open().close()
PromptModel.reset_schema_cache()
db = self._open()
self.assertEqual(db.get_prompt_by_id(1)["text"], "a")
def test_image_table_is_rebuilt_with_uniqueness_by_full_path(self):
self._legacy_db(
[
"INSERT INTO prompts (text, hash, workflow_name)"
" VALUES ('a', 'h1', 'wf')",
# same basename in two folders: both must survive
"INSERT INTO generated_images (prompt_id, image_path, filename)"
" VALUES (1, '/out/2026-01-01/a.png', 'a.png')",
"INSERT INTO generated_images (prompt_id, image_path, filename)"
" VALUES (1, '/out/2026-01-02/a.png', 'a.png')",
# two spellings of one file: only the newest row survives
"INSERT INTO generated_images (prompt_id, image_path, filename)"
" VALUES (1, '/out/x/../b.png', 'b.png')",
"INSERT INTO generated_images (prompt_id, image_path, filename)"
" VALUES (1, '/out/b.png', 'b.png')",
]
)
db = self._open()
images = db.get_prompt_images(1)
self.assertEqual(len(images), 3)
self.assertIn("file_path", self._columns("generated_images"))
self.assertEqual(
sorted(i["file_path"] for i in images),
sorted(
normalize_image_path(p)
for p in (
"/out/2026-01-01/a.png",
"/out/2026-01-02/a.png",
"/out/b.png",
)
),
)
self.assertEqual([i["id"] for i in images if i["filename"] == "b.png"], [4])
self.assertEqual(
self._unique_columns("generated_images"), ["file_path", "prompt_id"]
)
def test_image_uniqueness_migration_is_not_undone_on_the_next_start(self):
self._legacy_db(
["INSERT INTO prompts (text, hash, workflow_name) VALUES ('a', 'h1', 'wf')"]
)
self._open().close()
PromptModel.reset_schema_cache()
db = self._open()
self.assertEqual(
self._unique_columns("generated_images"), ["file_path", "prompt_id"]
)
self.assertTrue(db.link_image_to_prompt(1, "/out/1/a.png"))
self.assertTrue(db.link_image_to_prompt(1, "/out/2/a.png"))
self.assertEqual(len(db.get_prompt_images(1)), 2)
def _unique_columns(self, table):
conn = sqlite3.connect(self.path)
try:
columns = []
for idx in conn.execute(f"PRAGMA index_list({table})"):
if idx[2] == 1:
columns.extend(
c[2] for c in conn.execute(f"PRAGMA index_info({idx[1]})")
)
return sorted(columns)
finally:
conn.close()
def _table_names(self):
conn = sqlite3.connect(self.path)
try:
return [
r[0]
for r in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table'"
)
]
finally:
conn.close()
if __name__ == "__main__":
unittest.main()
+30 -9
View File
@@ -10,6 +10,7 @@ Verifies that:
import asyncio
import json
import os
from contextlib import closing
import sys
import tempfile
import types
@@ -84,9 +85,18 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase):
def _make_api_stub(self, db_path):
"""Create a minimal object that mimics the admin mixin's self."""
from py.api import PromptManagerAPI
stub = MagicMock()
stub.db.model.db_path = db_path
stub.logger = MagicMock()
stub._public_path = PromptManagerAPI._public_path
stub._public_error = PromptManagerAPI._public_error
async def run_in_executor(func, *args, **kwargs):
return func(*args, **kwargs)
stub._run_in_executor = run_in_executor
return stub
def _run_async(self, coro):
@@ -110,12 +120,14 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase):
try:
import sqlite3
with sqlite3.connect(db_file) as conn:
with closing(sqlite3.connect(db_file)) as conn, conn:
conn.execute(
"CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT, created_at TEXT)"
"CREATE TABLE prompts "
"(id INTEGER PRIMARY KEY, text TEXT, created_at TEXT)"
)
conn.execute(
"INSERT INTO prompts (text, created_at) VALUES ('test', '2024-01-01')"
"INSERT INTO prompts (text, created_at) "
"VALUES ('test', '2024-01-01')"
)
stub = self._make_api_stub(db_file)
@@ -143,10 +155,10 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase):
body = json.loads(result.body)
self.assertTrue(body["success"])
self.assertEqual(body["diagnostics"]["database"]["status"], "error")
self.assertIn(
"/nonexistent/custom/prompts.db",
body["diagnostics"]["database"]["message"],
)
# The message names the file but never the absolute server path.
message = body["diagnostics"]["database"]["message"]
self.assertIn("prompts.db", message)
self.assertNotIn("/nonexistent/custom", message)
def test_backup_uses_model_db_path(self):
"""backup_database should read from self.db.model.db_path."""
@@ -155,20 +167,29 @@ class TestAdminEndpointsUseConfigPath(unittest.TestCase):
from py.api.admin import AdminRoutesMixin
from database.operations import PromptDatabase
with tempfile.NamedTemporaryFile(delete=False, suffix=".db") as f:
db_file = f.name
f.write(b"SQLite format 3\x00test data for backup")
# A WAL-safe backup copies pages through sqlite itself, so the
# configured path has to hold a real database, not arbitrary bytes.
real_db = PromptDatabase(db_file)
try:
stub = self._make_api_stub(db_file)
stub.db.model = real_db.model
result = self._run_async(
AdminRoutesMixin.backup_database(stub, MagicMock())
)
self.assertEqual(result.content_type, "application/octet-stream")
self.assertGreater(len(result.body), 0)
self.assertTrue(result.body.startswith(b"SQLite format 3\x00"))
finally:
os.unlink(db_file)
real_db.close_all()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(db_file + suffix):
os.unlink(db_file + suffix)
def test_backup_reports_missing_custom_path(self):
"""backup_database should return 404 when configured path doesn't exist."""
+4 -2
View File
@@ -53,7 +53,8 @@ class TestPromptManagerTextEmptyInput(unittest.TestCase):
def test_empty_text_with_prepend_append(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
"""Empty main text with prepend/append should produce single space, not double."""
"""Empty main text with prepend/append should produce single space, not
double."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
result = node.process_text(text="", prepend_text="before", append_text="after")
self.assertIsInstance(result, tuple)
@@ -128,7 +129,8 @@ class TestPromptManagerEmptyInput(unittest.TestCase):
def test_empty_text_with_prepend_append(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
"""Empty main text with prepend/append should produce single space, not double."""
"""Empty main text with prepend/append should produce single space, not
double."""
node, clip = self._make_node_and_clip(
mock_db, mock_integ, mock_tracker, mock_monitor
)
+65
View File
@@ -0,0 +1,65 @@
"""ComfyUI loads only the canvas extension, not the admin or gallery bundles.
ComfyUI serves every .js file under WEB_DIRECTORY into the canvas page, so the
directory must hold nothing but the node extension.
"""
import os
import sys
import tempfile
import unittest
TESTS_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(TESTS_DIR)
sys.path.insert(0, ROOT)
sys.path.insert(0, TESTS_DIR)
from test_init_package import close_every_database, load_package # noqa: E402
class TestExtensionWebDir(unittest.TestCase):
package_name = "prompt_manager_pkg_web"
@classmethod
def setUpClass(cls):
cls.tmp = tempfile.TemporaryDirectory()
output_dir = os.path.join(cls.tmp.name, "output")
os.makedirs(output_dir)
cls.package, _ = load_package(
cls.package_name, os.path.join(cls.tmp.name, "prompts.db"), output_dir
)
@classmethod
def tearDownClass(cls):
monitor = getattr(cls.package, "_global_image_monitor", None)
if monitor is not None:
monitor.stop_monitoring()
sys.modules.pop(cls.package_name, None)
close_every_database() # the package's database must not hold the temp dir open
cls.tmp.cleanup()
def test_web_directory_is_the_canvas_extension_folder(self):
self.assertEqual(self.package.WEB_DIRECTORY, "web/comfy")
self.assertIn("WEB_DIRECTORY", self.package.__all__)
def test_comfyui_resolves_the_folder_next_to_the_package(self):
# ComfyUI joins the module directory with WEB_DIRECTORY and checks isdir
folder = os.path.join(ROOT, self.package.WEB_DIRECTORY)
self.assertTrue(os.path.isdir(folder))
self.assertEqual(sorted(os.listdir(folder)), ["prompt_manager.js"])
def test_extension_was_moved_not_copied(self):
self.assertFalse(os.path.exists(os.path.join(ROOT, "web", "prompt_manager.js")))
def test_extension_registers_with_the_canvas_app(self):
path = os.path.join(ROOT, "web", "comfy", "prompt_manager.js")
with open(path, encoding="utf-8") as f:
source = f.read()
self.assertIn("app.registerExtension(", source)
# Served at /extensions/<name>/prompt_manager.js, so ComfyUI's app module
# is two levels up regardless of where the file lives on disk.
self.assertIn('from "../../scripts/app.js"', source)
if __name__ == "__main__":
unittest.main()
+349
View File
@@ -0,0 +1,349 @@
"""Tests for gallery root validation (GalleryConfig.validate_gallery_root).
A gallery root may only point inside one of ComfyUI's own directories
(output/input/temp/user, as reported by ``folder_paths``) or inside a
directory listed in ``PROMPT_MANAGER_EXTRA_GALLERY_ROOTS``. Anything else,
including filesystem roots and the home directory, is rejected so the
image-serving routes can never be pointed at arbitrary files.
"""
import json
import os
import shutil
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.config import GalleryConfig, PromptManagerConfig # noqa: E402
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
class TestUpdateConfigValidatesRoots(unittest.TestCase):
"""Roots hand-edited into config.json go through validate_gallery_root."""
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
self.output_dir = self.comfy_dir / "output"
self.output_dir.mkdir(parents=True)
self.outside = Path(self.tmpdir) / "outside"
self.outside.mkdir()
self._orig_folder_paths = sys.modules.get("folder_paths")
sys.modules["folder_paths"] = types.SimpleNamespace(
base_path=str(self.comfy_dir),
get_output_directory=lambda: str(self.output_dir),
)
self.addCleanup(self._restore_folder_paths)
self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
self.addCleanup(
setattr, GalleryConfig, "MONITORING_DIRECTORIES", self._orig_dirs
)
self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None)
self.addCleanup(self._restore_extra_env)
def _restore_folder_paths(self):
if self._orig_folder_paths is None:
sys.modules.pop("folder_paths", None)
else:
sys.modules["folder_paths"] = self._orig_folder_paths
def _restore_extra_env(self):
if self._orig_extra is None:
os.environ.pop(EXTRA_ROOTS_ENV, None)
else:
os.environ[EXTRA_ROOTS_ENV] = self._orig_extra
def test_invalid_roots_are_dropped_with_a_warning(self):
good = self.output_dir / "renders"
good.mkdir()
with self.assertLogs("prompt_manager.config", level="WARNING") as logs:
GalleryConfig.update_config(
{
"monitoring": {
"directories": [
str(self.outside),
str(good),
os.path.abspath(os.sep),
42,
]
}
}
)
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(good)])
self.assertEqual(len(logs.output), 3)
self.assertNotIn(str(self.outside), "".join(logs.output))
def test_non_list_directories_are_ignored(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
GalleryConfig.update_config({"monitoring": {"directories": "/etc"}})
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.output_dir)])
def test_load_from_file_drops_invalid_roots(self):
config_path = os.path.join(self.tmpdir, "config.json")
with open(config_path, "w") as f:
json.dump(
{
"gallery": {
"monitoring": {
"directories": [str(self.outside), str(self.output_dir)]
}
}
},
f,
)
PromptManagerConfig.load_from_file(config_path)
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.output_dir)])
class GalleryRootTestCase(unittest.TestCase):
"""Shared fixture: a fake ComfyUI tree exposed through a stub folder_paths."""
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.addCleanup(self._rmtree, self.tmpdir)
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
self.output_dir = self.comfy_dir / "output"
self.input_dir = self.comfy_dir / "input"
self.output_dir.mkdir(parents=True)
self.input_dir.mkdir(parents=True)
self._orig_folder_paths = sys.modules.get("folder_paths")
sys.modules["folder_paths"] = types.SimpleNamespace(
base_path=str(self.comfy_dir),
get_output_directory=lambda: str(self.output_dir),
get_input_directory=lambda: str(self.input_dir),
get_temp_directory=self._raise_runtime_error,
# get_user_directory intentionally absent
)
self.addCleanup(self._restore_folder_paths)
self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None)
self.addCleanup(self._restore_extra_env)
@staticmethod
def _raise_runtime_error():
raise RuntimeError("temp directory unavailable")
@staticmethod
def _rmtree(path):
import shutil
shutil.rmtree(path, ignore_errors=True)
def _restore_folder_paths(self):
if self._orig_folder_paths is None:
sys.modules.pop("folder_paths", None)
else:
sys.modules["folder_paths"] = self._orig_folder_paths
def _restore_extra_env(self):
if self._orig_extra is None:
os.environ.pop(EXTRA_ROOTS_ENV, None)
else:
os.environ[EXTRA_ROOTS_ENV] = self._orig_extra
class TestValidateGalleryRootRejects(GalleryRootTestCase):
def test_filesystem_root_rejected(self):
ok, reason = GalleryConfig.validate_gallery_root(os.path.abspath(os.sep))
self.assertFalse(ok)
self.assertIn("root", reason.lower())
@unittest.skipUnless(os.name == "nt", "drive roots only exist on Windows")
def test_windows_drive_root_rejected(self):
ok, _ = GalleryConfig.validate_gallery_root("C:\\")
self.assertFalse(ok)
def test_home_directory_rejected(self):
ok, reason = GalleryConfig.validate_gallery_root(os.path.expanduser("~"))
self.assertFalse(ok)
self.assertIn("home", reason.lower())
def test_directory_outside_comfyui_rejected(self):
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
ok, _ = GalleryConfig.validate_gallery_root(str(outside))
self.assertFalse(ok)
def test_nonexistent_path_rejected(self):
missing = self.output_dir / "does_not_exist"
ok, reason = GalleryConfig.validate_gallery_root(str(missing))
self.assertFalse(ok)
self.assertIn("exist", reason.lower())
def test_file_rejected(self):
file_path = self.output_dir / "image.png"
file_path.write_bytes(b"PNG")
ok, reason = GalleryConfig.validate_gallery_root(str(file_path))
self.assertFalse(ok)
self.assertIn("directory", reason.lower())
def test_empty_and_non_string_rejected(self):
self.assertFalse(GalleryConfig.validate_gallery_root("")[0])
self.assertFalse(GalleryConfig.validate_gallery_root(" ")[0])
self.assertFalse(GalleryConfig.validate_gallery_root(None)[0])
self.assertFalse(GalleryConfig.validate_gallery_root(42)[0])
def test_sibling_with_shared_prefix_rejected(self):
"""'output2' must not pass because it string-starts with 'output'."""
sibling = self.comfy_dir / "output2"
sibling.mkdir()
ok, _ = GalleryConfig.validate_gallery_root(str(sibling))
self.assertFalse(ok)
def test_traversal_out_of_output_rejected(self):
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
sneaky = os.path.join(str(self.output_dir), "..", "..", "elsewhere")
ok, _ = GalleryConfig.validate_gallery_root(sneaky)
self.assertFalse(ok)
def test_symlink_escaping_output_rejected(self):
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
link = self.output_dir / "escape"
try:
os.symlink(str(outside), str(link), target_is_directory=True)
except (OSError, NotImplementedError):
self.skipTest("symlinks not available on this platform")
ok, _ = GalleryConfig.validate_gallery_root(str(link))
self.assertFalse(ok)
def test_env_root_that_is_filesystem_root_is_ignored(self):
"""Listing '/' in the env var must not re-open the whole disk."""
os.environ[EXTRA_ROOTS_ENV] = os.path.abspath(os.sep)
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
ok, _ = GalleryConfig.validate_gallery_root(str(outside))
self.assertFalse(ok)
class TestValidateGalleryRootAccepts(GalleryRootTestCase):
def test_output_directory_itself_accepted(self):
ok, reason = GalleryConfig.validate_gallery_root(str(self.output_dir))
self.assertTrue(ok, reason)
def test_subdirectory_of_output_accepted(self):
sub = self.output_dir / "renders" / "2026"
sub.mkdir(parents=True)
ok, reason = GalleryConfig.validate_gallery_root(str(sub))
self.assertTrue(ok, reason)
def test_subdirectory_of_input_accepted(self):
sub = self.input_dir / "refs"
sub.mkdir()
ok, reason = GalleryConfig.validate_gallery_root(str(sub))
self.assertTrue(ok, reason)
def test_directory_listed_in_env_accepted(self):
extra_a = Path(self.tmpdir) / "extra_a"
extra_b = Path(self.tmpdir) / "extra_b"
extra_a.mkdir()
extra_b.mkdir()
os.environ[EXTRA_ROOTS_ENV] = os.pathsep.join([str(extra_a), str(extra_b)])
ok_a, reason_a = GalleryConfig.validate_gallery_root(str(extra_a))
ok_b, reason_b = GalleryConfig.validate_gallery_root(str(extra_b / "."))
self.assertTrue(ok_a, reason_a)
self.assertTrue(ok_b, reason_b)
def test_relative_path_resolves_against_comfyui_base_not_cwd(self):
sub = self.output_dir / "rel"
sub.mkdir()
orig_cwd = os.getcwd()
os.chdir(self.tmpdir)
self.addCleanup(os.chdir, orig_cwd)
ok, reason = GalleryConfig.validate_gallery_root("output/rel")
self.assertTrue(ok, reason)
self.assertEqual(
GalleryConfig.resolve_gallery_root("output/rel"),
os.path.normcase(os.path.realpath(sub)),
)
# A path that only exists relative to the CWD is not accepted.
(Path(self.tmpdir) / "rel").mkdir()
self.assertFalse(GalleryConfig.validate_gallery_root("rel")[0])
def test_accepts_without_folder_paths_when_env_lists_dir(self):
sys.modules.pop("folder_paths", None)
extra = Path(self.tmpdir) / "extra"
extra.mkdir()
os.environ[EXTRA_ROOTS_ENV] = str(extra)
ok, reason = GalleryConfig.validate_gallery_root(str(extra))
self.assertTrue(ok, reason)
def test_rejects_everything_without_folder_paths_and_env(self):
sys.modules.pop("folder_paths", None)
ok, _ = GalleryConfig.validate_gallery_root(str(self.output_dir))
self.assertFalse(ok)
class TestGetAllOutputDirsFiltersRoots(GalleryRootTestCase):
"""PromptManagerAPI._get_all_output_dirs() only returns validated roots."""
def setUp(self):
super().setUp()
from py.api import PromptManagerAPI
self.api = PromptManagerAPI()
self.api._cached_output_dir = None
self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
self.addCleanup(self._restore_dirs)
def _restore_dirs(self):
GalleryConfig.MONITORING_DIRECTORIES = self._orig_dirs
def _dirs(self):
return [os.path.normcase(str(p)) for p in self.api._get_all_output_dirs()]
def test_filesystem_root_is_never_served(self):
GalleryConfig.MONITORING_DIRECTORIES = [os.path.abspath(os.sep)]
dirs = self._dirs()
self.assertNotIn(os.path.normcase(os.path.abspath(os.sep)), dirs)
self.assertEqual(dirs, [os.path.normcase(os.path.realpath(self.output_dir))])
def test_outside_directory_is_dropped_and_valid_one_kept(self):
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
good = self.output_dir / "good"
good.mkdir()
GalleryConfig.MONITORING_DIRECTORIES = [str(outside), str(good)]
self.assertEqual(self._dirs(), [os.path.normcase(os.path.realpath(good))])
def test_all_rejected_falls_back_to_comfyui_output_only(self):
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
self.assertEqual(
self._dirs(), [os.path.normcase(os.path.realpath(self.output_dir))]
)
def test_no_configured_dirs_uses_auto_detection(self):
GalleryConfig.MONITORING_DIRECTORIES = []
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
self.assertEqual(self._dirs(), [os.path.normcase(str(self.output_dir))])
def test_nothing_available_returns_empty_list(self):
sys.modules.pop("folder_paths", None)
outside = Path(self.tmpdir) / "elsewhere"
outside.mkdir()
GalleryConfig.MONITORING_DIRECTORIES = [str(outside)]
self.assertEqual(self._dirs(), [])
if __name__ == "__main__":
unittest.main()
+34 -65
View File
@@ -1,7 +1,7 @@
"""
Tests for hashing utilities.
Tests generate_content_hash and is_duplicate_prompt which are
Tests is_duplicate_prompt and generate_prompt_hash, which are
currently untested. generate_prompt_hash is covered in test_basic.py
but we add edge cases here.
"""
@@ -14,83 +14,52 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.hashing import (
generate_prompt_hash,
generate_content_hash,
is_duplicate_prompt,
)
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
class TestGenerateContentHash(unittest.TestCase):
"""Test generate_content_hash function."""
def test_basic_content(self):
content = {
"text": "A prompt",
"category": "nature",
"tags": ["sky"],
"workflow_name": "wf1",
}
h = generate_content_hash(content)
self.assertEqual(len(h), 64) # SHA256 hex digest
class TestSingleHashImplementation(unittest.TestCase):
"""generate_prompt_hash is the only prompt hash; nothing re-implements it."""
def test_normalization_case_insensitive(self):
c1 = {"text": "Hello World", "category": "Test"}
c2 = {"text": "hello world", "category": "test"}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_node_base_delegates_to_generate_prompt_hash(self):
from prompt_manager_base import PromptManagerBase
def test_normalization_strips_whitespace(self):
c1 = {"text": " hello ", "category": " test "}
c2 = {"text": "hello", "category": "test"}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
node = PromptManagerBase.__new__(PromptManagerBase)
self.assertEqual(node._generate_hash(" A "), generate_prompt_hash(" A "))
self.assertEqual(node._generate_hash(" A "), generate_prompt_hash("a"))
def test_tag_order_independent(self):
c1 = {"text": "prompt", "tags": ["beta", "alpha"]}
c2 = {"text": "prompt", "tags": ["alpha", "beta"]}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_tracker_and_monitor_use_generate_prompt_hash(self):
for relative in (
os.path.join("utils", "prompt_tracker.py"),
os.path.join("utils", "image_monitor.py"),
"prompt_manager_base.py",
):
with open(os.path.join(ROOT, relative), encoding="utf-8") as f:
source = f.read()
self.assertFalse("hashlib.sha256" in source, f"inline hash in {relative}")
self.assertTrue("generate_prompt_hash" in source, relative)
def test_different_text_different_hash(self):
c1 = {"text": "prompt A"}
c2 = {"text": "prompt B"}
self.assertNotEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_tracker_finds_an_existing_prompt_by_the_shared_hash(self):
import unittest.mock as mock
def test_different_category_different_hash(self):
c1 = {"text": "same", "category": "cat1"}
c2 = {"text": "same", "category": "cat2"}
self.assertNotEqual(generate_content_hash(c1), generate_content_hash(c2))
from utils.prompt_tracker import PromptTracker
def test_different_tags_different_hash(self):
c1 = {"text": "same", "tags": ["tag1"]}
c2 = {"text": "same", "tags": ["tag2"]}
self.assertNotEqual(generate_content_hash(c1), generate_content_hash(c2))
db = mock.Mock()
db.get_prompt_by_hash.return_value = {"id": 42}
tracker = PromptTracker(db)
tracker.set_current_prompt(" Hello World ")
db.get_prompt_by_hash.assert_called_once_with(
generate_prompt_hash(" Hello World ")
)
self.assertEqual(tracker.get_current_prompt()["id"], 42)
def test_empty_dict(self):
h = generate_content_hash({})
self.assertEqual(len(h), 64)
def test_node_base_has_no_browser_opener(self):
from prompt_manager_base import PromptManagerBase
def test_missing_fields_use_defaults(self):
c1 = {"text": "hello"}
c2 = {"text": "hello", "category": None, "tags": [], "workflow_name": None}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_none_category_same_as_empty(self):
c1 = {"text": "test", "category": None}
c2 = {"text": "test", "category": ""}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_none_workflow_same_as_empty(self):
c1 = {"text": "test", "workflow_name": None}
c2 = {"text": "test", "workflow_name": ""}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_empty_tags_filtered(self):
c1 = {"text": "test", "tags": ["valid", "", " "]}
c2 = {"text": "test", "tags": ["valid"]}
self.assertEqual(generate_content_hash(c1), generate_content_hash(c2))
def test_deterministic(self):
content = {"text": "stable", "tags": ["a", "b"]}
h1 = generate_content_hash(content)
h2 = generate_content_hash(content)
self.assertEqual(h1, h2)
self.assertFalse(hasattr(PromptManagerBase, "_open_web_interface"))
self.assertFalse(os.path.exists(os.path.join(ROOT, "restart_gallery.py")))
class TestIsDuplicatePrompt(unittest.TestCase):
+248
View File
@@ -0,0 +1,248 @@
"""Image linking picks the positive prompt, preferring the image's own metadata."""
import os
import sys
import tempfile
import unittest
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
from utils.prompt_tracker import PromptTracker
LOADER = {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "m.safetensors"},
}
def graph(positive_text="a cat", negative_text="blurry"):
# Negative listed first: dict order must not decide which prompt an image gets
return {
"176": {"class_type": "PromptManager", "inputs": {"text": negative_text}},
"134": {"class_type": "PromptManager", "inputs": {"text": positive_text}},
"5": {
"class_type": "KSampler",
"inputs": {
"positive": ["134", 0],
"negative": ["176", 0],
"model": ["20", 0],
},
},
"20": LOADER,
}
class LinkingTestCase(unittest.TestCase):
def setUp(self):
fd, self.db_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
fd, self.image = tempfile.mkstemp(suffix=".png")
os.close(fd)
self.db = PromptDatabase(self.db_path)
self.tracker = PromptTracker(self.db)
self.handler = ImageGenerationHandler(self.db, self.tracker)
self.linked = []
self.handler.link_image_to_prompt = (
lambda path, prompt, meta: self.linked.append(prompt["id"])
)
def tearDown(self):
self.db.close_all() # Windows cannot unlink an open database file
for path in (
self.image,
self.db_path,
self.db_path + "-wal",
self.db_path + "-shm",
):
if os.path.exists(path):
os.unlink(path)
def _save(self, text):
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
def _process(self, metadata):
with mock.patch.object(
self.handler.metadata_extractor, "extract_metadata", return_value=metadata
):
self.handler.process_new_image(self.image)
return self.linked[-1] if self.linked else None
class TestMetadataLinking(LinkingTestCase):
def test_links_the_positive_prompt_not_the_first_prompt_manager_node(self):
positive = self._save("a cat")
self._save("blurry")
self.assertEqual(self._process({"prompt": graph()}), positive)
def test_image_metadata_wins_over_a_stale_queue_entry(self):
positive = self._save("a cat")
stale = self._save("previous run")
self.tracker.set_current_prompt("previous run", {"prompt_id": stale})
self.assertEqual(self._process({"prompt": graph()}), positive)
self.assertEqual(
self.tracker.pop_next_prompt()["id"], stale, "queue left intact"
)
def test_text_from_a_known_string_node_links_from_metadata(self):
positive = self._save("a cat")
linked_graph = graph()
linked_graph["134"]["inputs"]["text"] = ["8", 0]
linked_graph["8"] = {
"class_type": "PrimitiveString",
"inputs": {"value": "a cat"},
}
self.assertEqual(self._process({"prompt": linked_graph}), positive)
def test_batch_items_with_linked_text_use_the_queue(self):
first, second = self._save("item one"), self._save("item two")
batch_graph = graph()
batch_graph["134"]["inputs"]["text"] = ["7", 0]
batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
self.tracker.set_current_prompt("item one", {"prompt_id": first})
self.tracker.set_current_prompt("item two", {"prompt_id": second})
self.assertEqual(self._process({"prompt": batch_graph}), first)
self.assertEqual(self._process({"prompt": batch_graph}), second)
def test_api_graph_present_skips_the_role_unaware_workflow_fallback(self):
negative = self._save("blurry")
batch_graph = graph()
batch_graph["134"]["inputs"]["text"] = ["7", 0]
metadata = {
"prompt": batch_graph,
"text_encoder_nodes": [
{"type": "PromptManager", "inputs": [], "widgets_values": ["blurry"]}
],
}
self.assertEqual(
self.handler._find_prompt_from_metadata(metadata), ("unknown", None)
)
latest = self._save("most recent prompt") # recency fallback picks this
self.assertEqual(self._process(metadata), latest)
self.assertNotEqual(latest, negative)
def test_workflow_fallback_still_used_without_an_api_graph(self):
positive = self._save("from workflow only")
metadata = {
"text_encoder_nodes": [
{
"type": "PromptManager",
"inputs": [],
"widgets_values": ["from workflow only"],
}
]
}
state, prompt = self.handler._find_prompt_from_metadata(metadata)
self.assertEqual(state, "linked")
self.assertEqual(prompt["id"], positive)
def foreign_graph():
"""A workflow with no PromptManager node at all."""
return {
"9": {
"class_type": "CLIPTextEncode",
"inputs": {"text": "a cat", "clip": ["20", 1]},
},
"5": {
"class_type": "KSampler",
"inputs": {"positive": ["9", 0], "negative": ["9", 0], "model": ["20", 0]},
},
"20": LOADER,
}
class TestForeignWorkflows(LinkingTestCase):
"""Images from workflows without a positive PromptManager node are not ours."""
def _queue_batch_item(self, text):
prompt_id = self._save(text)
self.tracker.set_current_prompt(text, {"prompt_id": prompt_id})
return prompt_id
def test_foreign_image_leaves_the_batch_queue_aligned(self):
first = self._queue_batch_item("item one")
self.assertIsNone(self._process({"prompt": foreign_graph()}))
self.assertEqual(self.linked, [])
self.assertEqual(self.tracker.pop_next_prompt()["id"], first)
self.assertIsNone(self.tracker.pop_next_prompt())
def test_foreign_image_is_not_linked_by_the_later_fallbacks_either(self):
self._queue_batch_item("item one") # live tracker and recent prompt exist
self.assertIsNone(self._process({"prompt": foreign_graph()}))
self.assertEqual(self.linked, [])
def test_negative_only_workflow_is_foreign(self):
self._save("blurry")
negative_only = foreign_graph()
negative_only["176"] = {
"class_type": "PromptManager",
"inputs": {"text": "blurry"},
}
negative_only["5"]["inputs"]["negative"] = ["176", 0]
self.assertEqual(
self.handler._find_prompt_from_metadata({"prompt": negative_only}),
("foreign", None),
)
def test_workflow_only_metadata_without_prompt_manager_is_foreign(self):
metadata = {
"text_encoder_nodes": [
{"type": "CLIPTextEncode", "inputs": [], "widgets_values": ["a cat"]}
]
}
self.assertEqual(
self.handler._find_prompt_from_metadata(metadata), ("foreign", None)
)
def test_missing_metadata_is_unknown_and_uses_the_queue(self):
first = self._queue_batch_item("item one")
self.assertEqual(
self.handler._find_prompt_from_metadata(None), ("unknown", None)
)
self.assertEqual(self._process(None), first)
def test_tri_state_for_linked_unknown_and_foreign(self):
positive = self._save("a cat")
state, prompt = self.handler._find_prompt_from_metadata({"prompt": graph()})
self.assertEqual((state, prompt["id"]), ("linked", positive))
self.assertEqual(
self.handler._find_prompt_from_metadata({"prompt": foreign_graph()}),
("foreign", None),
)
self.assertEqual(
self.handler._find_prompt_from_metadata({"workflow": {}}),
("unknown", None),
)
self.assertEqual(
self.handler._find_prompt_from_metadata({"prompt": graph("not saved")}),
("unknown", None),
)
def test_unknown_image_without_any_context_falls_back_to_the_latest_prompt(self):
self._save("older")
latest = self._save("latest")
self.assertEqual(self._process(None), latest)
def test_recent_prompt_fallback_handles_an_empty_database(self):
self.assertIsNone(self.handler._get_fallback_prompt())
self.assertIsNone(self._process(None))
def test_metadata_lookup_errors_are_unknown(self):
self._save("a cat")
with mock.patch.object(
self.db, "get_prompt_by_hash", side_effect=RuntimeError("locked")
):
self.assertEqual(
self.handler._find_prompt_from_metadata({"prompt": graph()}),
("unknown", None),
)
if __name__ == "__main__":
unittest.main()
+327 -6
View File
@@ -10,19 +10,21 @@ Verifies that:
import os
import sys
import tempfile
import threading
import time
import types
import unittest
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Mock ComfyUI server before importing anything that touches config
# Mock ComfyUI server before importing anything that touches config. Only the
# first test module to run installs the stub; nothing else in sys.modules is
# replaced, so later modules (test_comfyui_integration) see the real package.
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules["server"] = _mock_server
# Mock comfyui_integration to avoid import issues
sys.modules["utils.comfyui_integration"] = MagicMock()
sys.modules.setdefault("server", _mock_server)
from utils.image_monitor import get_image_monitor
import utils.image_monitor as im_mod
@@ -49,7 +51,8 @@ class TestMonitorSingletonLifecycle(unittest.TestCase):
self.assertIs(m1, m2)
def test_singleton_survives_across_different_callers(self):
"""Different db/tracker args on subsequent calls still return the same instance."""
"""Different db/tracker args on subsequent calls still return the same
instance."""
db1, tracker1 = MagicMock(), MagicMock()
db2, tracker2 = MagicMock(), MagicMock()
@@ -132,6 +135,324 @@ class TestMonitorSingletonLifecycle(unittest.TestCase):
self.assertIs(monitor.observer, mock_observer)
class FakeObserver:
"""Records what the monitor does with the watchdog observer."""
def __init__(self):
self.scheduled = []
self.started = False
self.stopped = False
self.joined = False
def schedule(self, handler, path, recursive=False):
self.scheduled.append((handler, path, recursive))
def start(self):
self.started = True
def stop(self):
self.stopped = True
def join(self, timeout=None):
self.joined = True
def is_alive(self):
return self.started and not self.stopped
class TestStartStop(unittest.TestCase):
"""start_monitoring validates directories before it creates an observer."""
def setUp(self):
im_mod._monitor_instance = None
self.monitor = im_mod.ImageMonitor(MagicMock(), MagicMock())
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.missing = os.path.join(self.tmp.name, "does-not-exist")
def tearDown(self):
im_mod._monitor_instance = None
def test_only_nonexistent_directories_leaves_monitoring_off(self):
with patch.object(im_mod, "Observer") as observer_cls:
self.monitor.start_monitoring([self.missing])
observer_cls.assert_not_called()
self.assertFalse(self.monitor.is_monitoring)
self.assertIsNone(self.monitor.observer)
self.assertEqual(self.monitor.monitored_directories, [])
self.monitor.stop_monitoring() # must not raise
def test_stop_without_start_does_not_raise(self):
self.monitor.stop_monitoring()
self.assertFalse(self.monitor.is_monitoring)
def test_start_twice_is_a_noop(self):
fake = FakeObserver()
with patch.object(im_mod, "Observer", return_value=fake) as observer_cls:
self.monitor.start_monitoring([self.tmp.name])
self.monitor.start_monitoring([self.tmp.name])
self.assertEqual(observer_cls.call_count, 1)
self.assertEqual(len(fake.scheduled), 1)
self.assertTrue(self.monitor.is_monitoring)
def test_stop_joins_the_observer_and_clears_state(self):
fake = FakeObserver()
with patch.object(im_mod, "Observer", return_value=fake):
self.monitor.start_monitoring([self.tmp.name, self.missing])
self.assertTrue(fake.started)
self.assertEqual([p for _, p, _ in fake.scheduled], [self.tmp.name])
self.assertEqual(self.monitor.monitored_directories, [self.tmp.name])
self.assertTrue(self.monitor.get_status()["observer_alive"])
self.monitor.stop_monitoring()
self.assertTrue(fake.stopped)
self.assertTrue(fake.joined)
self.assertFalse(self.monitor.is_monitoring)
self.assertIsNone(self.monitor.handler)
self.assertEqual(self.monitor.monitored_directories, [])
self.assertFalse(self.monitor.get_status()["running"])
def test_stop_shuts_the_handler_worker_down(self):
fake = FakeObserver()
with patch.object(im_mod, "Observer", return_value=fake):
self.monitor.start_monitoring([self.tmp.name])
handler = self.monitor.handler
handler.processing_delay = 10
image = os.path.join(self.tmp.name, "w.png")
with open(image, "wb") as f:
f.write(b"x" * 16)
handler.schedule(image)
worker = handler._worker
self.assertTrue(worker.is_alive())
self.monitor.stop_monitoring()
worker.join(timeout=2)
self.assertFalse(worker.is_alive())
self.assertEqual(handler.pending_count(), 0)
# A restart gets a fresh handler and observer
fake2 = FakeObserver()
with patch.object(im_mod, "Observer", return_value=fake2):
self.monitor.start_monitoring([self.tmp.name])
self.assertTrue(self.monitor.is_monitoring)
self.assertIsNot(self.monitor.handler, handler)
self.monitor.stop_monitoring()
def test_observer_failing_to_start_leaves_monitoring_off(self):
fake = FakeObserver()
fake.start = MagicMock(side_effect=OSError("inotify limit reached"))
with patch.object(im_mod, "Observer", return_value=fake):
self.monitor.start_monitoring([self.tmp.name])
self.assertFalse(self.monitor.is_monitoring)
self.monitor.stop_monitoring()
def test_disabled_in_config_leaves_monitoring_off(self):
config = types.SimpleNamespace(
MONITORING_ENABLED=False, MONITORING_DIRECTORIES=[]
)
with patch.object(im_mod, "Observer") as observer_cls:
with patch.object(
im_mod.ImageMonitor, "_gallery_config", return_value=config
):
self.monitor.start_monitoring([self.tmp.name])
observer_cls.assert_not_called()
self.assertFalse(self.monitor.is_monitoring)
def test_directories_failing_validation_are_not_watched(self):
allowed = os.path.join(self.tmp.name, "allowed")
os.mkdir(allowed)
config = types.SimpleNamespace(
MONITORING_ENABLED=True,
MONITORING_DIRECTORIES=[self.tmp.name, allowed],
validate_gallery_root=lambda p: (p == allowed, "outside ComfyUI"),
)
fake = FakeObserver()
with patch.object(im_mod, "Observer", return_value=fake):
with patch.object(
im_mod.ImageMonitor, "_gallery_config", return_value=config
):
self.monitor.start_monitoring()
self.assertEqual([p for _, p, _ in fake.scheduled], [allowed])
self.assertEqual(self.monitor.monitored_directories, [allowed])
self.monitor.stop_monitoring()
def test_explicit_directories_are_validated_too(self):
config = types.SimpleNamespace(
MONITORING_ENABLED=True,
MONITORING_DIRECTORIES=[],
validate_gallery_root=lambda p: (False, "outside ComfyUI"),
)
with patch.object(im_mod, "Observer") as observer_cls:
with patch.object(
im_mod.ImageMonitor, "_gallery_config", return_value=config
):
self.monitor.start_monitoring([self.tmp.name])
observer_cls.assert_not_called()
self.assertFalse(self.monitor.is_monitoring)
def test_configured_directories_are_used_when_none_are_given(self):
config = types.SimpleNamespace(
MONITORING_ENABLED=True, MONITORING_DIRECTORIES=[self.tmp.name]
)
fake = FakeObserver()
with patch.object(im_mod, "Observer", return_value=fake):
with patch.object(
im_mod.ImageMonitor, "_gallery_config", return_value=config
):
self.monitor.start_monitoring()
self.assertEqual(self.monitor.monitored_directories, [self.tmp.name])
self.monitor.stop_monitoring()
class FakeEvent:
def __init__(self, path, is_directory=False):
self.src_path = path
self.is_directory = is_directory
def wait_for(predicate, timeout=5.0):
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(0.01)
return predicate()
class TestHandlerScheduling(unittest.TestCase):
"""New images are processed once each, in order, after the file settles."""
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.tracker = MagicMock()
self.tracker.get_current_prompt.return_value = {"id": 7}
self.handler = im_mod.ImageGenerationHandler(MagicMock(), self.tracker)
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.processed = []
self.handler.process_new_image = lambda path, prompt_snapshot=None: (
self.processed.append((path, prompt_snapshot))
)
def _image(self, name, content=b"x" * 16):
path = os.path.join(self.tmp.name, name)
with open(path, "wb") as f:
f.write(content)
return path
def test_same_path_is_scheduled_once(self):
path = self._image("a.png")
self.assertTrue(self.handler.schedule(path, {"id": 1}))
self.assertFalse(self.handler.schedule(path, {"id": 2}))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.processed, [(path, {"id": 1})])
def test_on_created_snapshots_the_prompt_and_skips_non_images(self):
path = self._image("b.png")
self.handler.on_created(FakeEvent(path))
self.handler.on_created(FakeEvent(os.path.join(self.tmp.name, "c.txt")))
self.handler.on_created(FakeEvent(self.tmp.name, is_directory=True))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.processed, [(path, {"id": 7})])
def test_images_are_processed_in_creation_order(self):
paths = [self._image(f"{i}.png") for i in range(5)]
for path in paths:
self.handler.schedule(path)
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual([p for p, _ in self.processed], paths)
def test_pending_set_is_bounded(self):
self.handler.processing_delay = 10 # keep the worker busy
first = self._image("first.png")
self.handler.schedule(first)
with patch.object(im_mod, "MAX_PENDING_IMAGES", 2):
self.assertTrue(self.handler.schedule(self._image("second.png")))
self.assertFalse(self.handler.schedule(self._image("third.png")))
self.assertEqual(self.handler.pending_count(), 2)
def test_file_that_never_settles_is_skipped_and_cleared(self):
path = self._image("d.png")
sizes = iter(range(1, 100))
with patch.object(im_mod.os.path, "getsize", side_effect=lambda p: next(sizes)):
self.handler._process_pending(path)
self.assertEqual(self.processed, [])
self.assertEqual(self.handler.pending_count(), 0)
def test_vanished_file_is_skipped(self):
path = os.path.join(self.tmp.name, "gone.png")
self.assertFalse(self.handler.wait_until_settled(path))
self.handler._process_pending(path)
self.assertEqual(self.processed, [])
def test_settled_file_is_detected_after_two_stable_reads(self):
path = self._image("e.png")
self.assertTrue(self.handler.wait_until_settled(path))
empty = self._image("empty.png", content=b"")
self.assertFalse(self.handler.wait_until_settled(empty))
def test_shutdown_stops_the_worker_and_clears_pending(self):
self.handler.processing_delay = 10 # keep the first image waiting
self.handler.schedule(self._image("s1.png"))
self.handler.schedule(self._image("s2.png"))
worker = self.handler._worker
self.assertTrue(worker.is_alive())
self.handler.shutdown()
worker.join(timeout=2)
self.assertFalse(worker.is_alive())
self.assertEqual(self.handler.pending_count(), 0)
def test_shutdown_is_idempotent_and_safe_before_start(self):
fresh = im_mod.ImageGenerationHandler(MagicMock(), self.tracker)
fresh.shutdown()
fresh.shutdown()
self.handler.schedule(self._image("s3.png"))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.handler.shutdown()
self.handler.shutdown()
def test_scheduling_after_shutdown_restarts_the_worker(self):
self.handler.schedule(self._image("r1.png"))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
old_worker = self.handler._worker
self.handler.shutdown()
old_worker.join(timeout=2)
self.assertFalse(old_worker.is_alive())
path = self._image("r2.png")
self.assertTrue(self.handler.schedule(path))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertIn(path, [p for p, _ in self.processed])
self.assertIsNot(self.handler._worker, old_worker)
self.handler.shutdown()
def test_processing_errors_do_not_kill_the_worker(self):
def explode(path, prompt_snapshot=None):
self.processed.append(path)
raise RuntimeError("boom")
self.handler.process_new_image = explode
first, second = self._image("f1.png"), self._image("f2.png")
self.handler.schedule(first)
self.handler.schedule(second)
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.assertEqual(self.processed, [first, second])
class TestModuleIsolation(unittest.TestCase):
"""This module must not leave mocks behind for later test modules."""
def test_comfyui_integration_is_not_replaced_by_a_mock(self):
module = sys.modules.get("utils.comfyui_integration")
self.assertFalse(
isinstance(module, MagicMock),
"utils.comfyui_integration was replaced by a MagicMock in sys.modules",
)
class TestMonitorThreadSafety(unittest.TestCase):
"""Singleton creation must be thread-safe."""
+268
View File
@@ -0,0 +1,268 @@
"""Importing the extension package wires the queue hook to the shared database.
The package is loaded the way ComfyUI loads it (``__init__`` with relative
imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed.
"""
import importlib.util
import os
import sys
import tempfile
import types
import unittest
import unittest.mock as mock
from aiohttp import web
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
try:
from tests.open_handles import open_handles_under # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import open_handles_under # noqa: E402
class FakePromptQueue:
def __init__(self):
self.results = []
def get(self, timeout=None):
return self.results.pop(0) if self.results else None
def fake_server_module():
"""A ``server`` module whose PromptServer.instance accepts routes and hooks."""
instance = types.SimpleNamespace(
routes=web.RouteTableDef(),
prompt_queue=FakePromptQueue(),
handlers=[],
)
instance.add_on_prompt_handler = instance.handlers.append
module = types.ModuleType("server")
module.PromptServer = type("PromptServer", (), {"instance": instance})
return module
def fake_folder_paths_module(output_dir):
module = types.ModuleType("folder_paths")
module.base_path = output_dir
module.models_dir = os.path.join(output_dir, "models")
module.get_output_directory = lambda: output_dir
module.get_folder_paths = lambda name: []
return module
def load_package(name, db_path, output_dir, server=None):
"""Import ROOT/__init__.py as package ``name`` with ComfyUI stubbed.
Returns (package, server instance). The caller stops the image monitor.
"""
saved = {key: sys.modules.get(key) for key in ("server", "folder_paths")}
server = server or fake_server_module()
sys.modules["server"] = server
sys.modules["folder_paths"] = fake_folder_paths_module(output_dir)
try:
from py.config import PromptManagerConfig
with mock.patch.object(PromptManagerConfig, "DEFAULT_DB_PATH", db_path):
spec = importlib.util.spec_from_file_location(
name,
os.path.join(ROOT, "__init__.py"),
submodule_search_locations=[ROOT],
)
package = importlib.util.module_from_spec(spec)
sys.modules[name] = package
try:
spec.loader.exec_module(package)
except BaseException:
sys.modules.pop(name, None)
raise
finally:
for key, module in saved.items():
if module is None:
sys.modules.pop(key, None)
else:
sys.modules[key] = module
return package, server.PromptServer.instance
def close_every_database():
"""Close connections of every PromptModel registry in the process.
A package loaded under its own name imports its own copy of
database.models, whose PromptModel class keeps its own registry.
"""
for name, module in list(sys.modules.items()):
if name.endswith("database.models") and hasattr(module, "PromptModel"):
module.PromptModel.close_all_instances()
class PackageTestCase(unittest.TestCase):
package_name = "prompt_manager_pkg"
@classmethod
def setUpClass(cls):
cls.tmp = tempfile.TemporaryDirectory()
cls.output_dir = os.path.join(cls.tmp.name, "output")
os.makedirs(cls.output_dir)
cls.db_path = cls.make_db_path(cls.tmp.name)
cls.package, cls.server = load_package(
cls.package_name, cls.db_path, cls.output_dir, server=cls.make_server()
)
@classmethod
def make_db_path(cls, tmp):
return os.path.join(tmp, "prompts.db")
@classmethod
def make_server(cls):
return fake_server_module()
@classmethod
def tearDownClass(cls):
monitor = getattr(cls.package, "_global_image_monitor", None)
if monitor is not None:
monitor.stop_monitoring()
sys.modules.pop(cls.package_name, None)
close_every_database()
hits = open_handles_under(cls.tmp.name)
cls.tmp.cleanup()
assert not hits, f"files still open before temp dir cleanup: {hits}"
class TestPackageWiring(PackageTestCase):
def test_nodes_are_registered(self):
self.assertEqual(
sorted(self.package.NODE_CLASS_MAPPINGS),
["PromptManager", "PromptManagerText", "PromptSearchList"],
)
def test_db_factory_returns_the_shared_database(self):
self.assertIsNotNone(self.package._global_db)
self.assertIs(self.package._db_factory(), self.package._global_db)
self.assertEqual(
os.path.normcase(os.path.realpath(self.package._global_db.model.db_path)),
os.path.normcase(os.path.realpath(self.db_path)),
)
def test_db_factory_raises_while_the_database_is_unavailable(self):
with mock.patch.object(self.package, "_global_db", None):
with self.assertRaisesRegex(RuntimeError, "database unavailable"):
self.package._db_factory()
def test_queue_hook_is_registered_and_the_queue_get_is_wrapped(self):
self.assertEqual(len(self.server.handlers), 1)
self.assertTrue(
getattr(self.server.prompt_queue.get, "_prompt_manager_dequeue_hook", False)
)
def test_image_monitor_watches_the_stubbed_output_directory(self):
status = self.package._global_image_monitor.get_status()
self.assertTrue(status["running"])
self.assertEqual(
[
os.path.normcase(os.path.realpath(d))
for d in status["monitored_directories"]
],
[os.path.normcase(os.path.realpath(self.output_dir))],
)
class TestPackageWithoutDatabase(PackageTestCase):
"""The database failing to initialise must not break the queue hook."""
package_name = "prompt_manager_pkg_nodb"
@classmethod
def make_db_path(cls, tmp):
blocker = os.path.join(tmp, "blocker")
with open(blocker, "w", encoding="utf-8") as f:
f.write("not a directory")
return os.path.join(blocker, "prompts.db") # parent is a file: init fails
def test_global_db_is_defined_but_unset(self):
self.assertIsNone(self.package._global_db)
with self.assertRaisesRegex(RuntimeError, "database unavailable"):
self.package._db_factory()
def test_queue_hook_still_registers_and_survives_a_dequeue(self):
self.assertEqual(len(self.server.handlers), 1)
request = self.server.handlers[0](
{"prompt": {"1": {"class_type": "PromptManager", "inputs": {"text": "x"}}}}
)
self.server.prompt_queue.results.append(
((0, request["prompt_id"], request["prompt"], {}, [], False), 1)
)
self.assertIsNotNone(self.server.prompt_queue.get())
class TestVersion(PackageTestCase):
def test_version_comes_from_pyproject(self):
import re
with open(os.path.join(ROOT, "pyproject.toml"), encoding="utf-8") as f:
expected = re.search(r'version\s*=\s*["\']([^"\']+)["\']', f.read())
self.assertEqual(self.package.get_version(), expected.group(1))
def test_unreadable_pyproject_gives_unknown(self):
with mock.patch.object(self.package, "Path", side_effect=OSError("denied")):
self.assertEqual(self.package.get_version(), "unknown")
missing = mock.Mock()
missing.parent.__truediv__ = lambda self_, name: mock.Mock(exists=lambda: False)
with mock.patch.object(self.package, "Path", return_value=missing):
self.assertEqual(self.package.get_version(), "unknown")
class TestPackageWhenHookRegistrationFails(PackageTestCase):
"""A server that rejects the hook must not break the rest of the package."""
package_name = "prompt_manager_pkg_nohook"
@classmethod
def make_server(cls):
module = fake_server_module()
instance = module.PromptServer.instance
def reject(handler):
raise RuntimeError("hooks unsupported")
instance.add_on_prompt_handler = reject
return module
def test_nodes_and_monitoring_still_come_up(self):
self.assertIn("PromptManager", self.package.NODE_CLASS_MAPPINGS)
self.assertIsNotNone(self.package._global_db)
self.assertTrue(self.package._global_image_monitor.is_monitoring)
self.assertEqual(self.server.handlers, [])
if __name__ == "__main__":
unittest.main()
class TestConsoleBanner(unittest.TestCase):
"""Loading the node must not depend on the console being able to print emoji."""
def test_package_loads_with_a_cp1252_console(self):
import contextlib
import io
tmp = tempfile.TemporaryDirectory()
self.addCleanup(tmp.cleanup)
self.addCleanup(close_every_database)
output_dir = os.path.join(tmp.name, "output")
os.makedirs(output_dir)
name = "prompt_manager_pkg_cp1252"
stream = io.TextIOWrapper(io.BytesIO(), encoding="cp1252", errors="strict")
with contextlib.redirect_stdout(stream):
package, _ = load_package(name, os.path.join(tmp.name, "p.db"), output_dir)
self.addCleanup(sys.modules.pop, name, None)
monitor = getattr(package, "_global_image_monitor", None)
if monitor is not None:
self.addCleanup(monitor.stop_monitoring)
stream.flush()
banner = stream.buffer.getvalue().decode("cp1252")
self.assertIn("Loaded:", banner)
+20
View File
@@ -26,6 +26,7 @@ class LoraDBTestCase(unittest.TestCase):
self.db = PromptDatabase(self.temp_db.name)
def tearDown(self):
self.db.close()
for suffix in ("", "-wal", "-shm"):
path = self.temp_db.name + suffix
if os.path.exists(path):
@@ -194,6 +195,25 @@ class TestGetPromptSubfolders(LoraDBTestCase):
self.assertIsInstance(folders, list)
self.assertGreater(len(folders), 0)
def test_subfolders_use_forward_slashes_on_every_platform(self):
"""Clients split folders on "/" whatever separator the OS produced."""
from unittest import mock
import database.operations as operations
pid = self._save("prompt")
self._link_image(pid, "/output/2026/08-Aug/06/img.png")
with (
mock.patch.object(
operations.os.path, "relpath", return_value="2026\\08-Aug\\06"
),
mock.patch.object(operations.os, "sep", "\\"),
):
folders = self.db.get_prompt_subfolders(root_dirs=["/output"])
self.assertEqual(folders, ["2026/08-Aug/06"])
def test_include_ancestors_adds_intermediate_paths(self):
pid = self._save("prompt")
self._link_image(pid, "/output/2026/08-Aug/2026-08-06/img.png")
+707 -1
View File
@@ -5,23 +5,35 @@ Tests metadata parsing, trigger word extraction, image URL extraction,
directory detection, TriggerWordCache, and image download logic.
"""
import email
import json
import os
import shutil
import sys
import tempfile
import threading
import types
import unittest
import urllib.parse
import urllib.request
import urllib.response
from io import BytesIO
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py import lora_utils
from py.lora_utils import (
CIVITAI_HOSTS,
MAX_CIVITAI_DOWNLOAD_BYTES,
TriggerWordCache,
download_civitai_images,
get_civitai_image_urls,
get_example_prompt_from_metadata,
get_lora_image_cache_dir,
get_trigger_words_from_metadata,
is_civitai_url,
read_lora_metadata,
)
@@ -217,13 +229,414 @@ class TestGetLoraImageCacheDir(unittest.TestCase):
def test_returns_path(self):
cache_dir = get_lora_image_cache_dir()
self.assertIsInstance(cache_dir, Path)
self.assertTrue(str(cache_dir).endswith("data/lora_images"))
self.assertEqual(cache_dir.parts[-2:], ("data", "lora_images"))
def test_directory_exists(self):
cache_dir = get_lora_image_cache_dir()
self.assertTrue(cache_dir.is_dir())
# ── CivitAI download tests (no network) ───────────────────────────────
def _png_bytes():
"""A valid 2x2 PNG so PIL can decode the fake download."""
from io import BytesIO
from PIL import Image
buf = BytesIO()
Image.new("RGB", (2, 2), (10, 20, 30)).save(buf, "PNG")
return buf.getvalue()
class _FakeResponse:
"""Minimal stand-in for the object returned by urllib.request.urlopen."""
def __init__(self, body, content_length=None):
self._body = body
self._pos = 0
self.headers = {}
if content_length is not None:
self.headers["Content-Length"] = str(content_length)
def read(self, amt=-1):
if amt is None or amt < 0:
chunk = self._body[self._pos :]
self._pos = len(self._body)
return chunk
chunk = self._body[self._pos : self._pos + amt]
self._pos += len(chunk)
return chunk
def __enter__(self):
return self
def __exit__(self, *exc):
return False
class _FakeOpener:
"""Records every request and serves a canned response."""
def __init__(self, response):
self.response = response
self.requests = []
def __call__(self, req, timeout=None):
self.requests.append(req)
return self.response
class TestCivitaiUrlPolicy(unittest.TestCase):
def test_allow_list_contains_expected_hosts(self):
self.assertEqual(
CIVITAI_HOSTS,
frozenset(
{
"civitai.com",
"www.civitai.com",
"api.civitai.com",
"image.civitai.com",
}
),
)
self.assertEqual(MAX_CIVITAI_DOWNLOAD_BYTES, 50 * 1024 * 1024)
def test_https_civitai_hosts_are_allowed(self):
for host in CIVITAI_HOSTS:
self.assertTrue(is_civitai_url(f"https://{host}/api/download/1"), host)
def test_other_hosts_and_schemes_are_refused(self):
for url in (
"https://evil.example/x",
"http://civitai.com/api/download/1",
"https://civitai.com.evil.example/x",
"https://notcivitai.com/x",
"ftp://civitai.com/x",
"not a url",
"",
):
self.assertFalse(is_civitai_url(url), url)
def test_userinfo_and_non_default_ports_are_refused(self):
for url in (
"https://evil.com\\@civitai.com/x",
"https://user@civitai.com/x",
"https://user:pw@civitai.com/x",
"https://civitai.com:8443/x",
):
self.assertFalse(is_civitai_url(url), url)
def test_explicit_default_port_is_allowed(self):
self.assertTrue(is_civitai_url("https://civitai.com:443/x"))
class TestLoraManagerPathPolicy(unittest.TestCase):
"""Only a direct child of a ComfyUI custom_nodes dir named *lora* qualifies."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
self.root = _make_comfy_root(self.tmp)
self.custom_nodes = self.root / "custom_nodes"
patcher = patch.object(
lora_utils, "custom_nodes_directories", return_value=[self.custom_nodes]
)
patcher.start()
self.addCleanup(patcher.stop)
@staticmethod
def _canonical(path):
return os.path.normcase(os.path.realpath(str(path)))
def test_direct_child_named_lora_is_accepted(self):
lm = _make_lora_manager(self.custom_nodes)
self.assertEqual(
lora_utils.resolve_lora_manager_path(str(lm)), self._canonical(lm)
)
def test_name_match_is_case_insensitive(self):
lm = _make_lora_manager(self.custom_nodes, "comfyui-LORA-manager")
self.assertEqual(
lora_utils.resolve_lora_manager_path(str(lm)), self._canonical(lm)
)
def test_outside_custom_nodes_is_rejected(self):
lm = _make_lora_manager(self.tmp)
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(lm)))
def test_child_without_lora_in_name_is_rejected(self):
other = _make_lora_manager(self.custom_nodes, "ComfyUI-Other")
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(other)))
def test_custom_nodes_itself_is_rejected(self):
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(self.custom_nodes)))
def test_nested_directory_is_rejected(self):
nested = _make_lora_manager(self.custom_nodes / "pack", "lora-sub")
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(nested)))
def test_sibling_prefix_directory_is_rejected(self):
sibling = self.root / "custom_nodes2"
lm = _make_lora_manager(sibling)
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(lm)))
def test_traversal_out_of_custom_nodes_is_rejected(self):
_make_lora_manager(self.root, "lora-escape")
candidate = os.path.join(str(self.custom_nodes), "..", "lora-escape")
self.assertIsNone(lora_utils.resolve_lora_manager_path(candidate))
def test_symlink_escaping_custom_nodes_is_rejected(self):
target = _make_lora_manager(self.tmp)
link = self.custom_nodes / "lora-link"
try:
link.symlink_to(target, target_is_directory=True)
except (OSError, NotImplementedError):
self.skipTest("symlinks not supported on this platform")
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(link)))
def test_relative_path_is_resolved_against_the_comfyui_root(self):
lm = _make_lora_manager(self.custom_nodes)
with patch.object(lora_utils, "find_comfyui_root", return_value=self.root):
resolved = lora_utils.resolve_lora_manager_path(
"custom_nodes/ComfyUI-Lora-Manager"
)
self.assertEqual(resolved, self._canonical(lm))
def test_missing_file_or_bad_input_is_rejected(self):
self.assertIsNone(
lora_utils.resolve_lora_manager_path(str(self.custom_nodes / "lora-none"))
)
plain = self.custom_nodes / "lora.txt"
plain.write_text("")
self.assertIsNone(lora_utils.resolve_lora_manager_path(str(plain)))
self.assertIsNone(lora_utils.resolve_lora_manager_path(""))
self.assertIsNone(lora_utils.resolve_lora_manager_path(None))
self.assertIsNone(lora_utils.resolve_lora_manager_path("bad\x00name"))
class TestCustomNodesDirectories(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
self.root = _make_comfy_root(self.tmp)
def test_uses_folder_paths_base_path(self):
fake = _FakeFolderPaths(str(self.root))
with patch.dict(sys.modules, {"folder_paths": fake}):
dirs = lora_utils.custom_nodes_directories()
self.assertEqual(dirs, [(self.root / "custom_nodes").resolve()])
def test_includes_folder_paths_custom_nodes_entries(self):
extra = self.tmp / "more_nodes"
extra.mkdir()
fake = types.SimpleNamespace(
base_path=str(self.root),
get_folder_paths=lambda name: (
[str(extra)] if name == "custom_nodes" else []
),
)
with patch.dict(sys.modules, {"folder_paths": fake}):
dirs = lora_utils.custom_nodes_directories()
self.assertEqual(
dirs, [(self.root / "custom_nodes").resolve(), extra.resolve()]
)
def test_falls_back_to_the_parent_of_this_package(self):
with (
patch.dict(sys.modules, {"folder_paths": None}),
patch.object(lora_utils, "find_comfyui_root", return_value=None),
):
dirs = lora_utils.custom_nodes_directories()
package_dir = Path(lora_utils.__file__).resolve().parent.parent
self.assertEqual(dirs, [package_dir.parent])
class TestDownloadOne(unittest.TestCase):
def setUp(self):
self.tmpdir = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.target = self.tmpdir / "out.jpg"
def _run(self, url, opener, api_key="secret-key"):
with patch.object(lora_utils, "_open_url", opener):
return lora_utils._download_one(url, self.target, api_key)
def test_off_host_url_is_refused_without_a_request(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
result = self._run("https://evil.example/x", opener)
self.assertIsNone(result)
self.assertEqual(opener.requests, [])
self.assertFalse(self.target.exists())
def test_plain_http_civitai_is_refused(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
result = self._run("http://civitai.com/api/download/1", opener)
self.assertIsNone(result)
self.assertEqual(opener.requests, [])
def test_civitai_download_sends_bearer_header(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
result = self._run("https://civitai.com/api/download/1", opener)
self.assertEqual(os.path.realpath(result), os.path.realpath(str(self.target)))
self.assertEqual(len(opener.requests), 1)
req = opener.requests[0]
self.assertEqual(req.get_header("Authorization"), "Bearer secret-key")
self.assertTrue(self.target.exists())
def test_no_bearer_header_without_a_key(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
self._run("https://image.civitai.com/a.png", opener, api_key="")
self.assertIsNone(opener.requests[0].get_header("Authorization"))
def test_oversized_content_length_is_refused(self):
opener = _FakeOpener(
_FakeResponse(_png_bytes(), content_length=MAX_CIVITAI_DOWNLOAD_BYTES + 1)
)
result = self._run("https://civitai.com/api/download/1", opener)
self.assertIsNone(result)
self.assertFalse(self.target.exists())
def test_reading_stops_past_cap_when_length_is_absent(self):
body = b"x" * (MAX_CIVITAI_DOWNLOAD_BYTES + 1)
response = _FakeResponse(body)
opener = _FakeOpener(response)
with patch.object(lora_utils, "MAX_CIVITAI_DOWNLOAD_BYTES", 1024):
result = self._run("https://civitai.com/api/download/1", opener)
self.assertIsNone(result)
self.assertLessEqual(response._pos, 1024 + lora_utils._DOWNLOAD_CHUNK_BYTES)
self.assertFalse(self.target.exists())
class _FakeHTTPSHandler(urllib.request.HTTPSHandler):
"""Serves canned responses through a real urllib opener, recording requests.
CivitAI hosts answer with a 302 to ``redirect_to``; every other host
answers 200 with a PNG body.
"""
def __init__(self, redirect_to):
super().__init__()
self.redirect_to = redirect_to
self.requests = []
def https_open(self, req):
self.requests.append(req)
host = (urllib.parse.urlsplit(req.full_url).hostname or "").lower()
if host in CIVITAI_HOSTS:
headers = email.message_from_string(f"Location: {self.redirect_to}\n")
resp = urllib.response.addinfourl(BytesIO(b""), headers, req.full_url, 302)
resp.msg = "Found" # urllib's error processor reads response.msg
return resp
headers = email.message_from_string("Content-Type: image/png\n")
resp = urllib.response.addinfourl(
BytesIO(_png_bytes()), headers, req.full_url, 200
)
resp.msg = "OK"
return resp
class TestRedirectsAreRefused(unittest.TestCase):
def setUp(self):
self.tmpdir = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmpdir, True)
def test_redirect_to_other_host_is_not_followed(self):
handler = _FakeHTTPSHandler("https://evil.example/steal")
opener = lora_utils._build_opener(handler)
with patch.object(lora_utils, "_opener", opener):
result = lora_utils._download_one(
"https://civitai.com/api/download/1", self.tmpdir / "x.jpg", "key"
)
self.assertIsNone(result)
self.assertEqual(
[r.full_url for r in handler.requests],
["https://civitai.com/api/download/1"],
)
for req in handler.requests:
host = urllib.parse.urlsplit(req.full_url).hostname
if host not in CIVITAI_HOSTS:
self.assertIsNone(req.get_header("Authorization"))
def test_redirect_to_civitai_is_not_followed_either(self):
handler = _FakeHTTPSHandler("https://image.civitai.com/other.png")
opener = lora_utils._build_opener(handler)
with patch.object(lora_utils, "_opener", opener):
result = lora_utils._download_one(
"https://civitai.com/api/download/1", self.tmpdir / "x.jpg", "key"
)
self.assertIsNone(result)
self.assertEqual(len(handler.requests), 1)
class TestDownloadCivitaiImages(unittest.TestCase):
def setUp(self):
self.tmpdir = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.meta_path = self.tmpdir / "lora.safetensors.metadata.json"
self.meta_path.write_text("{}")
def test_only_civitai_urls_are_fetched(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
metadata = _make_metadata(
images=[
{"url": "https://evil.example/steal"},
{"url": "https://civitai.com/api/download/1"},
{"url": "not-a-url"},
]
)
with patch.object(lora_utils, "_open_url", opener):
paths = download_civitai_images(
metadata, self.meta_path, self.tmpdir / "cache", "key"
)
self.assertEqual(len(paths), 1)
self.assertEqual(
[r.full_url for r in opener.requests],
["https://civitai.com/api/download/1"],
)
def test_cached_files_are_not_refetched(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
metadata = _make_metadata(images=[{"url": "https://civitai.com/a.png"}])
cache = self.tmpdir / "cache"
with patch.object(lora_utils, "_open_url", opener):
first = download_civitai_images(metadata, self.meta_path, cache, "")
second = download_civitai_images(metadata, self.meta_path, cache, "")
self.assertEqual(first, second)
self.assertEqual(len(opener.requests), 1)
def test_no_images_returns_empty(self):
self.assertEqual(
download_civitai_images({}, self.meta_path, self.tmpdir / "cache"), []
)
# ── TriggerWordCache tests ────────────────────────────────────────────
@@ -312,5 +725,298 @@ class TestTriggerWordCache(unittest.TestCase):
self.assertEqual(errors, [])
# ── Detection, directories and metadata helpers ───────────────────────
class _FakeFolderPaths:
def __init__(self, base_path, lora_dirs=None):
self.base_path = base_path
self._lora_dirs = lora_dirs or []
def get_folder_paths(self, name):
return list(self._lora_dirs)
def _make_comfy_root(tmp):
"""Create a fake ComfyUI tree: main.py + custom_nodes/ + models/loras."""
root = Path(tmp) / "ComfyUI"
(root / "custom_nodes").mkdir(parents=True)
(root / "models" / "loras").mkdir(parents=True)
(root / "main.py").write_text("")
return root
def _make_lora_manager(custom_nodes, name="ComfyUI-Lora-Manager"):
lm = custom_nodes / name
(lm / "py").mkdir(parents=True)
(lm / "__init__.py").write_text("")
return lm
class TestFindComfyuiRoot(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
def test_folder_paths_base_path_wins(self):
root = _make_comfy_root(self.tmp)
with patch.dict(sys.modules, {"folder_paths": _FakeFolderPaths(str(root))}):
self.assertEqual(lora_utils.find_comfyui_root(), root)
def test_walks_up_from_module_file_without_folder_paths(self):
root = _make_comfy_root(self.tmp)
fake_file = root / "custom_nodes" / "ext" / "py" / "lora_utils.py"
fake_file.parent.mkdir(parents=True)
fake_file.write_text("")
with (
patch.dict(sys.modules, {"folder_paths": None}),
patch.object(lora_utils, "__file__", str(fake_file)),
):
self.assertEqual(lora_utils.find_comfyui_root(), root.resolve())
def test_returns_none_when_no_root_found(self):
fake_file = self.tmp / "lonely" / "lora_utils.py"
fake_file.parent.mkdir()
fake_file.write_text("")
with (
patch.dict(sys.modules, {"folder_paths": None}),
patch.object(lora_utils, "__file__", str(fake_file)),
):
self.assertIsNone(lora_utils.find_comfyui_root())
class TestDetectLoraManager(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
def test_custom_path_is_used_when_it_looks_right(self):
lm = _make_lora_manager(self.tmp)
with patch.object(lora_utils, "find_comfyui_root", return_value=None):
self.assertEqual(lora_utils.detect_lora_manager(str(lm)), str(lm.resolve()))
def test_custom_path_without_init_is_ignored(self):
bogus = self.tmp / "nope"
bogus.mkdir()
with patch.object(lora_utils, "find_comfyui_root", return_value=None):
self.assertIsNone(lora_utils.detect_lora_manager(str(bogus)))
def test_auto_detects_under_custom_nodes_case_insensitively(self):
root = _make_comfy_root(self.tmp)
(root / "custom_nodes" / "other").mkdir()
lm = _make_lora_manager(root / "custom_nodes", "comfyui-LORA-manager")
with patch.object(lora_utils, "find_comfyui_root", return_value=root):
self.assertEqual(lora_utils.detect_lora_manager(), str(lm.resolve()))
def test_lora_manager_dir_variant_is_accepted(self):
lm = self.tmp / "LoraManager"
(lm / "lora_manager").mkdir(parents=True)
(lm / "__init__.py").write_text("")
self.assertTrue(lora_utils._looks_like_lora_manager(lm))
def test_nothing_found(self):
with patch.object(lora_utils, "find_comfyui_root", return_value=None):
self.assertIsNone(lora_utils.detect_lora_manager())
class TestFindLoraDirectories(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
self.root = _make_comfy_root(self.tmp)
self.lm = _make_lora_manager(self.root / "custom_nodes")
def test_collects_default_extra_runtime_and_metadata_dirs(self):
extra_abs = self.tmp / "extra_abs"
extra_abs.mkdir()
(self.root / "rel_loras").mkdir()
(self.root / "extra_model_paths.yaml").write_text(
"a:\n base_path: %s\n loras: |\n rel_loras\n %s\n\n"
"b: notadict\nc:\n base_path: x\n"
% (self.root.as_posix(), extra_abs.as_posix())
)
runtime = self.tmp / "runtime"
runtime.mkdir()
nested = self.lm / "user" / "loras"
nested.mkdir(parents=True)
(nested / "x.metadata.json").write_text("{}")
with (
patch.object(lora_utils, "find_comfyui_root", return_value=self.root),
patch.dict(
sys.modules,
{"folder_paths": _FakeFolderPaths(str(self.root), [str(runtime)])},
),
):
dirs = lora_utils.find_lora_directories(str(self.lm))
expected = {
str((self.root / "models" / "loras").resolve()),
str((self.root / "rel_loras").resolve()),
str(extra_abs.resolve()),
str(runtime.resolve()),
str(nested.resolve()),
}
self.assertEqual(set(dirs), expected)
self.assertEqual(dirs, sorted(dirs))
def test_non_dict_yaml_and_parse_errors_are_ignored(self):
(self.root / "extra_model_paths.yml").write_text("- just\n- a list\n")
self.assertEqual(lora_utils._get_extra_lora_paths(self.root), [])
(self.root / "extra_model_paths.yml").write_text("a: [unclosed\n")
self.assertEqual(lora_utils._get_extra_lora_paths(self.root), [])
def test_without_root_only_metadata_dirs_are_found(self):
with (
patch.object(lora_utils, "find_comfyui_root", return_value=None),
patch.dict(sys.modules, {"folder_paths": None}),
):
self.assertEqual(lora_utils.find_lora_directories(str(self.lm)), [])
class TestModelNameAndPreviews(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
def test_model_name_precedence(self):
get_name = lora_utils.get_model_name_from_metadata
self.assertEqual(get_name({"model_name": "top"}), "top")
self.assertEqual(get_name({"civitai": {"model": {"name": "civ"}}}), "civ")
self.assertEqual(get_name({"file_name": "f.safetensors"}), "f.safetensors")
self.assertEqual(get_name({}), "unknown")
def test_preview_images_found_by_stem(self):
meta_path = self.tmp / "mylora.safetensors.metadata.json"
meta_path.write_text("{}")
(self.tmp / "mylora.png").write_bytes(b"x")
(self.tmp / "mylora.preview.jpg").write_bytes(b"x")
found = lora_utils.get_preview_images_from_metadata({}, meta_path)
self.assertEqual(
sorted(os.path.basename(f) for f in found),
["mylora.png", "mylora.preview.jpg"],
)
self.assertEqual(
lora_utils.get_preview_image_from_metadata({}, meta_path), found[0]
)
def test_preview_uses_file_name_from_metadata(self):
meta_path = self.tmp / "other.metadata.json"
(self.tmp / "named.webp").write_bytes(b"x")
found = lora_utils.get_preview_images_from_metadata(
{"file_name": "named.safetensors"}, meta_path
)
self.assertEqual([os.path.basename(f) for f in found], ["named.webp"])
self.assertIsNone(lora_utils.get_preview_image_from_metadata({}, meta_path))
def test_example_images_dir(self):
self.assertIsNone(lora_utils.get_example_images_dir(str(self.tmp)))
nested = self.tmp / "user" / "example_images"
nested.mkdir(parents=True)
self.assertEqual(
lora_utils.get_example_images_dir(str(self.tmp)), str(nested.resolve())
)
direct = self.tmp / "example_images"
direct.mkdir()
self.assertEqual(
lora_utils.get_example_images_dir(str(self.tmp)), str(direct.resolve())
)
class TestDownloadEdgeCases(unittest.TestCase):
def setUp(self):
self.tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmp, True)
def test_malformed_url_is_refused(self):
self.assertFalse(is_civitai_url("https://[invalid"))
def test_non_numeric_content_length_is_ignored(self):
resp = _FakeResponse(b"abc")
resp.headers["Content-Length"] = "lots"
self.assertEqual(lora_utils._read_capped(resp, 10), b"abc")
def test_download_uses_metadata_filename_when_file_name_missing(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
meta_path = self.tmp / "stemmed.safetensors.metadata.json"
metadata = {"civitai": {"images": [{"url": "https://civitai.com/a.png"}]}}
with patch.object(lora_utils, "_open_url", opener):
paths = download_civitai_images(metadata, meta_path, self.tmp / "cache")
self.assertEqual(len(paths), 1)
self.assertIn("stemmed", paths[0])
def test_failed_downloads_are_dropped(self):
opener = _FakeOpener(_FakeResponse(b"not an image"))
metadata = {"civitai": {"images": [{"url": "https://civitai.com/a.png"}]}}
with patch.object(lora_utils, "_open_url", opener):
paths = download_civitai_images(
metadata, self.tmp / "x.metadata.json", self.tmp / "cache"
)
self.assertEqual(paths, [])
class TestInjectTriggerWords(unittest.TestCase):
def setUp(self):
self.cache = TriggerWordCache()
self.cache._cache = {"style": ["neon", "glow"], "empty": []}
self.cache._loaded = True
def test_not_loaded_returns_unchanged(self):
cache = TriggerWordCache()
self.assertEqual(
lora_utils.inject_trigger_words("<lora:style:1>", cache),
("<lora:style:1>", []),
)
def test_no_lora_tags_returns_unchanged(self):
self.assertEqual(
lora_utils.inject_trigger_words("plain prompt", self.cache),
("plain prompt", []),
)
def test_words_are_appended_once(self):
text = "a <lora:Style:0.8> b <lora:style:1>"
self.assertEqual(
lora_utils.inject_trigger_words(text, self.cache),
(f"{text}, neon, glow", ["neon", "glow"]),
)
def test_words_already_present_are_not_repeated(self):
text = "NEON city <lora:style:1>"
self.assertEqual(
lora_utils.inject_trigger_words(text, self.cache),
(f"{text}, glow", ["glow"]),
)
def test_lora_without_words_leaves_text_alone(self):
text = "x <lora:empty:1> <lora:unknown:1>"
self.assertEqual(lora_utils.inject_trigger_words(text, self.cache), (text, []))
def test_cache_load_skips_bad_and_wordless_metadata(self):
tmp = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, tmp, True)
(tmp / "bad.metadata.json").write_text("{{{")
(tmp / "nowords.metadata.json").write_text("{}")
(tmp / "named.metadata.json").write_text(
json.dumps({"civitai": {"trainedWords": ["w"]}})
)
cache = TriggerWordCache()
with patch.object(lora_utils, "find_lora_directories", return_value=[str(tmp)]):
self.assertEqual(cache.load(str(tmp)), 1)
self.assertEqual(cache.get_trigger_words("named"), ["w"])
if __name__ == "__main__":
unittest.main()
+313 -4
View File
@@ -8,10 +8,11 @@ Ensures all nodes are properly configured for ComfyUI's execution engine:
Regression tests for: https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/120
"""
import itertools
import os
import sys
import time
import unittest
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@@ -70,6 +71,16 @@ class TestIsChangedPromptManager(unittest.TestCase):
result = PromptManager.IS_CHANGED(clip=None, text="hello")
self.assertIsInstance(result, str)
def test_different_tags_return_different_hash(self):
result1 = PromptManager.IS_CHANGED(clip=None, text="hello", tags="a, b")
result2 = PromptManager.IS_CHANGED(clip=None, text="hello", tags="a, c")
self.assertNotEqual(result1, result2)
def test_different_category_returns_different_hash(self):
result1 = PromptManager.IS_CHANGED(clip=None, text="hello", category="art")
result2 = PromptManager.IS_CHANGED(clip=None, text="hello", category="photo")
self.assertNotEqual(result1, result2)
class TestIsChangedPromptManagerText(unittest.TestCase):
"""IS_CHANGED for PromptManagerText should behave identically to
@@ -99,6 +110,15 @@ class TestIsChangedPromptManagerText(unittest.TestCase):
result2 = PromptManagerText.IS_CHANGED(text="hello", append_text="b")
self.assertNotEqual(result1, result2)
def test_different_tags_or_category_return_different_hash(self):
base = PromptManagerText.IS_CHANGED(text="hello", category="x", tags="a")
self.assertNotEqual(
base, PromptManagerText.IS_CHANGED(text="hello", category="x", tags="b")
)
self.assertNotEqual(
base, PromptManagerText.IS_CHANGED(text="hello", category="y", tags="a")
)
def test_returns_string(self):
result = PromptManagerText.IS_CHANGED(text="hello")
self.assertIsInstance(result, str)
@@ -109,9 +129,12 @@ class TestIsChangedPromptSearchList(unittest.TestCase):
so the node re-runs every time (database contents may have changed)."""
def test_returns_different_value_on_successive_calls(self):
result1 = PromptSearchList.IS_CHANGED()
time.sleep(0.01)
result2 = PromptSearchList.IS_CHANGED()
# Mock the clock instead of sleeping: the Windows clock ticks every
# 15.6 ms, so two real time.time() reads 10 ms apart can be equal.
ticks = itertools.count(1_700_000_000.0, 1.0)
with mock.patch("time.time", side_effect=lambda: next(ticks)):
result1 = PromptSearchList.IS_CHANGED()
result2 = PromptSearchList.IS_CHANGED()
self.assertNotEqual(result1, result2)
def test_returns_numeric(self):
@@ -119,5 +142,291 @@ class TestIsChangedPromptSearchList(unittest.TestCase):
self.assertIsInstance(result, float)
import logging # noqa: E402
import tempfile # noqa: E402
import types # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from prompt_manager_base import PromptManagerBase # noqa: E402
import prompt_manager_base # noqa: E402
class FakeClip:
"""The two CLIP calls the node makes, recorded."""
def __init__(self):
self.tokenized = []
def tokenize(self, text):
self.tokenized.append(text)
return ("tokens", text)
def encode_from_tokens_scheduled(self, tokens):
return [("cond", tokens[1])]
class NodeTestCase(unittest.TestCase):
"""A node with a real temporary database and no monitor or tracker."""
node_cls = PromptManagerBase
def setUp(self):
fd, self.db_path = tempfile.mkstemp(suffix=".db")
os.close(fd)
self.db = PromptDatabase(self.db_path)
self.node = self.node_cls.__new__(self.node_cls)
self.node.db = self.db
self.node.logger = logging.getLogger("test.node")
self.node.prompt_tracker = mock.Mock()
self.node.image_monitor = mock.Mock()
self.node.comfyui_integration = mock.Mock()
def tearDown(self):
self.db.close_all() # Windows cannot unlink an open database file
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.db_path + suffix):
os.unlink(self.db_path + suffix)
class TestBaseSaveAndSearch(NodeTestCase):
def test_init_wires_database_tracker_monitor_and_integration(self):
with (
mock.patch.object(prompt_manager_base, "PromptDatabase") as db_cls,
mock.patch.object(prompt_manager_base, "get_prompt_tracker") as tracker,
mock.patch.object(prompt_manager_base, "get_image_monitor") as monitor,
mock.patch.object(
prompt_manager_base, "get_comfyui_integration"
) as integration,
):
node = PromptManagerBase()
self.assertIs(node.db, db_cls.return_value)
tracker.assert_called_once_with(db_cls.return_value)
monitor.assert_called_once_with(db_cls.return_value, tracker.return_value)
self.assertIs(node.comfyui_integration, integration.return_value)
monitor.return_value.start_monitoring.assert_called_once_with()
def test_saving_an_existing_prompt_updates_its_metadata(self):
first = self.node._save_prompt_to_database("a cat")
second = self.node._save_prompt_to_database(
"a cat", category="animals", tags=["cute"]
)
self.assertEqual(first, second)
saved = self.db.get_prompt_by_id(first)
self.assertEqual(saved["category"], "animals")
self.assertIn("cute", saved["tags"])
self.assertEqual(saved["run_count"], 0)
def test_count_run_on_an_existing_prompt(self):
prompt_id = self.node._save_prompt_to_database("a cat")
self.node._save_prompt_to_database("a cat", count_run=True)
self.assertEqual(self.db.get_prompt_by_id(prompt_id)["run_count"], 1)
def test_database_errors_are_logged_and_return_none(self):
self.node.db = mock.Mock()
self.node.db.get_prompt_by_hash.side_effect = RuntimeError("locked")
with self.assertLogs("test.node", level="ERROR"):
self.assertIsNone(self.node._save_prompt_to_database("a cat"))
def test_record_use_failure_is_a_warning(self):
self.node.db = mock.Mock()
self.node.db.record_prompt_use.side_effect = RuntimeError("locked")
with self.assertLogs("test.node", level="WARNING"):
self.node._record_use(1)
def test_parse_tags(self):
self.assertIsNone(self.node._parse_tags(""))
self.assertIsNone(self.node._parse_tags(" , ,"))
self.assertEqual(self.node._parse_tags(" a, b ,,c"), ["a", "b", "c"])
def test_search_prompts_and_api_wrappers(self):
self.node._save_prompt_to_database("a cat on a mat")
self.assertEqual(self.node._search_prompts(""), [])
self.assertEqual(self.node._search_prompts(" "), [])
self.assertEqual(len(self.node.search_prompts_api("cat")), 1)
self.assertEqual(len(self.node.get_recent_prompts_api(limit=5)["prompts"]), 1)
def test_search_and_recent_errors_return_empty_lists(self):
self.node.db = mock.Mock()
self.node.db.search_prompts.side_effect = RuntimeError("locked")
self.node.db.get_recent_prompts.side_effect = RuntimeError("locked")
with self.assertLogs("test.node", level="ERROR"):
self.assertEqual(self.node._search_prompts("cat"), [])
with self.assertLogs("test.node", level="ERROR"):
self.assertEqual(self.node.get_recent_prompts_api(), [])
def test_gallery_system_start_failure_is_logged(self):
self.node.image_monitor.start_monitoring.side_effect = OSError("no inotify")
with self.assertLogs("test.node", level="ERROR"):
self.node._start_gallery_system()
self.node.image_monitor.get_status.return_value = {"running": False}
self.node.prompt_tracker.get_status.return_value = {"active_prompts_count": 0}
self.assertEqual(
self.node.get_gallery_status(),
{
"image_monitor": {"running": False},
"prompt_tracker": {"active_prompts_count": 0},
},
)
self.node.cleanup_gallery_system()
self.node.image_monitor.stop_monitoring.assert_not_called()
def fake_integration_config(enabled, trigger_words=True, path="/loras"):
module = types.ModuleType("py.config")
module.IntegrationConfig = types.SimpleNamespace(
LORA_MANAGER_ENABLED=enabled,
LORA_TRIGGER_WORDS_ENABLED=trigger_words,
LORA_MANAGER_PATH=path,
)
return module
class TestLoraTriggerWords(NodeTestCase):
def test_disabled_integration_returns_text_unchanged(self):
with mock.patch.dict(
sys.modules, {"py.config": fake_integration_config(False)}
):
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
with mock.patch.dict(
sys.modules, {"py.config": fake_integration_config(True, False)}
):
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
def test_missing_config_or_lora_utils_returns_text_unchanged(self):
with mock.patch.dict(sys.modules, {"py.config": None}):
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
with mock.patch.dict(
sys.modules,
{"py.config": fake_integration_config(True), "py.lora_utils": None},
):
self.assertEqual(self.node._inject_lora_trigger_words("a cat"), "a cat")
def test_enabled_integration_loads_the_cache_and_injects(self):
cache = mock.Mock(is_loaded=False)
lora_utils = types.ModuleType("py.lora_utils")
lora_utils.get_trigger_cache = lambda: cache
lora_utils.inject_trigger_words = lambda text, c: (text + " trig", ["trig"])
with mock.patch.dict(
sys.modules,
{"py.config": fake_integration_config(True), "py.lora_utils": lora_utils},
):
self.assertEqual(
self.node._inject_lora_trigger_words("a cat"), "a cat trig"
)
cache.load.assert_called_once_with("/loras")
class TestPromptManagerEncoding(NodeTestCase):
node_cls = PromptManager
def setUp(self):
super().setUp()
self.node._inject_lora_trigger_words = lambda text: text
def test_missing_clip_is_a_clear_error(self):
with self.assertRaisesRegex(RuntimeError, "clip input is invalid"):
self.node.encode_prompt(None, "a cat")
def test_encodes_combined_text_and_saves_the_main_text(self):
clip = FakeClip()
conditioning, final = self.node.encode_prompt(
clip, "a cat", prepend_text=" best ", append_text="8k ", tags="x"
)
self.assertEqual(final, "best a cat 8k")
self.assertEqual(clip.tokenized, ["best a cat 8k"])
self.assertEqual(conditioning, [("cond", "best a cat 8k")])
saved = self.db.get_recent_prompts(limit=1)["prompts"][0]
self.assertEqual(saved["text"], "a cat")
self.assertIn("prepend:best", saved["tags"])
self.assertIn("append:8k", saved["tags"])
self.node.comfyui_integration.register_prompt.assert_called_once()
def test_tracking_failure_does_not_break_encoding(self):
self.node._track_prompt_execution = mock.Mock(side_effect=RuntimeError("db"))
with self.assertLogs("test.node", level="WARNING"):
_, final = self.node.encode_prompt(FakeClip(), "a cat")
self.assertEqual(final, "a cat")
def test_empty_text_is_encoded_but_not_saved(self):
_, final = self.node.encode_prompt(FakeClip(), " ")
self.assertEqual(final.strip(), "")
self.assertEqual(self.db.get_recent_prompts(limit=1)["prompts"], [])
class TestPromptManagerTextProcessing(NodeTestCase):
node_cls = PromptManagerText
def setUp(self):
super().setUp()
self.node._inject_lora_trigger_words = lambda text: text
def test_tracking_failure_does_not_break_processing(self):
self.node._track_prompt_execution = mock.Mock(side_effect=RuntimeError("db"))
with self.assertLogs("test.node", level="WARNING"):
result = self.node.process_text("a castle", prepend_text="big")
self.assertEqual(result[0], "big a castle")
class TestPromptSearchListSearch(NodeTestCase):
node_cls = PromptSearchList
def _save(self, text, **kwargs):
from utils.hashing import generate_prompt_hash
return self.db.save_prompt(
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
)
def test_init_opens_the_database(self):
with mock.patch("prompt_search_list.PromptDatabase") as db_cls:
node = PromptSearchList()
self.assertIs(node.db, db_cls.return_value)
def test_no_results_returns_a_single_empty_item(self):
result = self.node.search(text="nothing here")
self.assertEqual(result["result"], ([""], "No results found", 0))
self.assertEqual(result["ui"], {"text": ["No results found"]})
def test_filters_by_text_category_tags_and_rating(self):
self._save("a red cat", category="animals", tags=["cute"])
self._save("a blue car", category="vehicles", tags=["fast"])
self.assertEqual(self.node.search(text="cat")["result"][0], ["a red cat"])
self.assertEqual(
self.node.search(category="vehicles")["result"][0], ["a blue car"]
)
self.assertEqual(self.node.search(tags=" cute , ")["result"][0], ["a red cat"])
self.assertEqual(self.node.search(min_rating=5)["result"][2], 0)
self.assertEqual(self.node.search(limit=1)["result"][2], 1)
def test_multipart_and_lora_only_prompts_are_skipped(self):
self._save("Clip_1: a cat Clip_2: a dog")
self._save("<lora:style:0.8> <lora:other:1>")
self._save("plain prompt")
self.assertEqual(self.node.search()["result"][0], ["plain prompt"])
result = self.node.search(skip_multipart=False)["result"]
self.assertEqual(
sorted(result[0]), ["Clip_1: a cat Clip_2: a dog", "plain prompt"]
)
def test_newlines_collapse_and_preview_is_numbered_and_truncated(self):
self._save("line one\n line two")
self._save("x" * 200)
result = self.node.search()
prompts, preview, count = result["result"]
self.assertEqual(count, 2)
self.assertIn("line one line two", prompts)
lines = preview.splitlines()
self.assertTrue(lines[0].startswith("[1] "))
self.assertTrue(any(line.endswith("...") and len(line) < 140 for line in lines))
self.assertEqual(result["ui"], {"text": ["Found 2 prompts"]})
def test_database_errors_are_returned_not_raised(self):
self.node.db = mock.Mock()
self.node.db.search_prompts.side_effect = RuntimeError("locked")
result = self.node.search(text="cat")
self.assertEqual(result["result"], ([""], "Error: locked", 0))
self.assertEqual(result["ui"], {"text": ["Search error: locked"]})
if __name__ == "__main__":
unittest.main()
+83
View File
@@ -0,0 +1,83 @@
"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once."""
import os
import sys
import threading
import unittest
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from prompt_manager import PromptManager
from prompt_manager_text import PromptManagerText
from utils import usage_tracking
GRAPH = {"134": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}}
def bare(node_cls):
"""A node without its database, tracker or gallery side effects."""
node = node_cls.__new__(node_cls)
node.logger = __import__("logging").getLogger("test.wiring")
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node._track_prompt_execution = mock.Mock(return_value=42)
return node
class TestHiddenInputs(unittest.TestCase):
def test_both_nodes_request_the_queued_graph_and_their_id(self):
for node_cls in (PromptManager, PromptManagerText):
hidden = node_cls.INPUT_TYPES().get("hidden", {})
self.assertEqual(hidden.get("prompt"), "PROMPT", node_cls.__name__)
self.assertEqual(hidden.get("unique_id"), "UNIQUE_ID", node_cls.__name__)
class TestNodesForwardGraph(unittest.TestCase):
def test_prompt_manager_forwards_graph_and_texts(self):
node = bare(PromptManager)
clip = mock.Mock()
node.encode_prompt(
clip, "a cat", prepend_text="best,", prompt=GRAPH, unique_id="134"
)
kwargs = node._track_prompt_execution.call_args.kwargs
self.assertEqual(kwargs["prompt_graph"], GRAPH)
self.assertEqual(kwargs["unique_id"], "134")
self.assertEqual(kwargs["text"], "a cat")
self.assertEqual(kwargs["encoding_text"], "best, a cat")
self.assertEqual(kwargs["additional_data"]["final_text"], "best, a cat")
def test_prompt_manager_text_forwards_graph(self):
node = bare(PromptManagerText)
node.process_text("a castle", prompt=GRAPH, unique_id="7")
kwargs = node._track_prompt_execution.call_args.kwargs
self.assertEqual((kwargs["prompt_graph"], kwargs["unique_id"]), (GRAPH, "7"))
self.assertEqual(kwargs["text"], "a castle")
def test_nodes_still_work_without_graph_context(self):
node = bare(PromptManager)
node.encode_prompt(mock.Mock(), "a cat")
kwargs = node._track_prompt_execution.call_args.kwargs
self.assertIsNone(kwargs["prompt_graph"])
self.assertIsNone(kwargs["unique_id"])
class TestQueueHookRegistration(unittest.TestCase):
def test_registers_once_and_handler_passes_requests_through(self):
server = mock.Mock()
db = mock.Mock()
db.get_prompt_by_hash.return_value = None
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db))
self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db))
server.add_on_prompt_handler.assert_called_once()
handler = server.add_on_prompt_handler.call_args.args[0]
request = {"prompt": GRAPH}
self.assertIs(handler(request), request)
if __name__ == "__main__":
unittest.main()
+3 -2
View File
@@ -16,7 +16,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Mock ComfyUI's server module before importing anything that touches config
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules["server"] = _mock_server
sys.modules.setdefault("server", _mock_server)
from py.api import PromptManagerAPI
from py.config import GalleryConfig
@@ -34,7 +34,8 @@ class TestFindComfyuiOutputDir(unittest.TestCase):
GalleryConfig.MONITORING_DIRECTORIES = self._orig_dirs
def test_configured_directory_takes_priority(self):
"""When a valid directory is configured, it should be used instead of auto-detect."""
"""When a valid directory is configured, it should be used instead of
auto-detect."""
with tempfile.TemporaryDirectory() as tmpdir:
GalleryConfig.MONITORING_DIRECTORIES = [tmpdir]
self.api._cached_output_dir = None
+100
View File
@@ -0,0 +1,100 @@
"""Packaging and repository hygiene contracts.
These guard the metadata that ComfyUI-Manager, the Comfy Registry and
``pip install .`` consume, and the ignore rules that keep local secrets
(``config.json`` holds the CivitAI API key) out of the repository.
"""
import os
import re
import unittest
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# PEP 508 requirement: name, optional extras, optional version spec.
_VERSION_CLAUSE = r"[<>=!~]=?\s*[A-Za-z0-9.*+!-]+\s*"
_REQUIREMENT = re.compile(
r"^[A-Za-z0-9][A-Za-z0-9._-]*" # distribution name
r"(\[[A-Za-z0-9,._-]+\])?\s*" # optional extras
r"(" + _VERSION_CLAUSE + r"(,\s*" + _VERSION_CLAUSE + r")*)?$" # version spec
)
def _read(name):
with open(os.path.join(ROOT, name), encoding="utf-8") as fh:
return fh.read()
def _pyproject():
try:
import tomllib # Python 3.11+
except ImportError: # pragma: no cover - exercised on 3.10 only
tomllib = None
text = _read("pyproject.toml")
if tomllib is not None:
return tomllib.loads(text)
# Minimal fallback for 3.10: pull the dependency list out with a regex.
match = re.search(r"dependencies\s*=\s*\[(.*?)\]", text, re.S)
deps = re.findall(r'"([^"]*)"', match.group(1)) if match else []
return {
"project": {
"dependencies": deps,
"requires-python": (
re.search(r'requires-python\s*=\s*"([^"]+)"', text).group(1)
if "requires-python" in text
else None
),
},
"build-system": {} if "[build-system]" in text else None,
}
class PyprojectContract(unittest.TestCase):
def test_dependencies_are_valid_requirements(self):
deps = _pyproject()["project"]["dependencies"]
self.assertTrue(deps, "dependencies must not be empty")
for dep in deps:
self.assertRegex(dep, _REQUIREMENT, f"not a PEP 508 requirement: {dep!r}")
def test_dependencies_match_requirements_txt(self):
deps = set(_pyproject()["project"]["dependencies"])
wanted = {
line.split("#", 1)[0].strip()
for line in _read("requirements.txt").splitlines()
if line.split("#", 1)[0].strip()
}
self.assertEqual(deps, wanted)
def test_requires_python_floor_is_3_10(self):
self.assertEqual(_pyproject()["project"]["requires-python"], ">=3.10")
def test_build_system_is_declared(self):
self.assertIsNotNone(_pyproject().get("build-system"))
class IgnoreRulesContract(unittest.TestCase):
def test_local_secrets_and_artifacts_are_ignored(self):
patterns = {
line.strip()
for line in _read(".gitignore").splitlines()
if line.strip() and not line.startswith("#")
}
for required in (
"config.json",
"resume.sh",
"*.backup_*",
"promptmanager_logs/",
):
self.assertIn(required, patterns)
class LintConfigContract(unittest.TestCase):
def test_flake8_reads_its_own_config_file(self):
# flake8 ignores pyproject.toml, so the settings must live in .flake8.
text = _read(".flake8")
self.assertIn("max-line-length = 88", text)
self.assertIn("E203", text)
if __name__ == "__main__":
unittest.main()
-1
View File
@@ -1,6 +1,5 @@
"""Tests for path traversal and symlink security in image serving."""
import os
import tempfile
import unittest
from pathlib import Path
+494
View File
@@ -0,0 +1,494 @@
"""Tests for utils.prompt_graph: which PromptManager nodes feed a positive input."""
import os
import sys
import time
import unittest
from unittest import mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils import prompt_graph
from utils.prompt_graph import (
MAX_DEPTH,
MAX_TEXT_LENGTH,
MAX_VISITS,
_upstream_links,
follows_input,
has_sampler,
positive_prompt_nodes,
resolve_text,
run_prompt_nodes,
)
def pm(text="a cat", **inputs):
return {"class_type": "PromptManager", "inputs": {"text": text, **inputs}}
def ksampler(positive, negative, model="20"):
return {
"class_type": "KSampler",
"inputs": {
"positive": [positive, 0],
"negative": [negative, 0],
"model": [model, 0],
},
}
LOADER = {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "m.safetensors"},
}
class TestPositivePromptNodes(unittest.TestCase):
def test_positive_and_negative_prompt_manager_nodes(self):
# Negative listed first, as in the #75 / pmtest workflow
graph = {
"176": pm("blurry"),
"134": pm("a cat"),
"5": ksampler("134", "176"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["134"])
def test_negative_only_prompt_manager_is_never_positive(self):
graph = {
"176": pm("blurry"),
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}},
"5": ksampler("6", "176"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), [])
def test_through_controlnet_and_combine_pass_throughs(self):
graph = {
"1": pm("subject"),
"2": pm("style"),
"3": pm("blurry"),
"4": {
"class_type": "ConditioningCombine",
"inputs": {"conditioning_1": ["1", 0], "conditioning_2": ["2", 0]},
},
"5": {
"class_type": "ControlNetApplyAdvanced",
"inputs": {"positive": ["4", 0], "negative": ["3", 0], "strength": 1.0},
},
"6": {
"class_type": "KSampler",
"inputs": {
"positive": ["5", 0],
"negative": ["5", 1],
"model": ["20", 0],
},
},
"20": LOADER,
}
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "2"])
def test_cfg_guider_and_basic_guider(self):
graph = {
"1": pm("fox"),
"2": pm("watermark"),
"3": {
"class_type": "CFGGuider",
"inputs": {
"positive": ["1", 0],
"negative": ["2", 0],
"model": ["20", 0],
"cfg": 5,
},
},
"4": pm("owl"),
"5": {
"class_type": "BasicGuider",
"inputs": {"conditioning": ["4", 0], "model": ["20", 0]},
},
"20": LOADER,
}
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "4"])
def test_prompt_manager_text_wired_into_an_encoder(self):
graph = {
"1": {"class_type": "PromptManagerText", "inputs": {"text": "a castle"}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 0]}},
"3": {"class_type": "PromptManagerText", "inputs": {"text": "lowres"}},
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}},
"5": ksampler("2", "4"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
def test_node_used_as_both_positive_and_negative_counts_as_positive(self):
graph = {"1": pm("same"), "5": ksampler("1", "1"), "20": LOADER}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
def test_two_samplers_sharing_a_prompt_list_it_once(self):
graph = {
"1": pm("hero"),
"2": pm("blurry"),
"5": ksampler("1", "2"),
"14": ksampler("1", "2"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
def test_graph_without_sampler_or_malformed_input(self):
self.assertEqual(positive_prompt_nodes({"1": pm()}), [])
for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}):
self.assertEqual(positive_prompt_nodes(bad), [], repr(bad))
def test_cycles_terminate(self):
loop = {
"class_type": "Weird",
"inputs": {
"conditioning": ["9", 0],
"conditioning_1": ["9", 0],
"positive": ["9", 0],
},
}
graph = {"9": loop, "5": ksampler("9", "9"), "20": LOADER}
self.assertEqual(positive_prompt_nodes(graph), [])
class TestRunPromptNodes(unittest.TestCase):
"""The prompt nodes that represent a run, with a safe fallback."""
def test_uses_positive_nodes_when_found(self):
graph = {
"176": pm("blurry"),
"134": pm("a cat"),
"5": ksampler("134", "176"),
"20": LOADER,
}
self.assertEqual(run_prompt_nodes(graph), ["134"])
def test_unrecognised_sampler_falls_back_to_every_prompt_manager_node(self):
graph = {
"1": pm("a cat"),
"9": {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}},
"2": {"class_type": "PromptManagerText", "inputs": {"text": "x"}},
}
self.assertEqual(sorted(run_prompt_nodes(graph)), ["1", "2"])
def test_malformed_graph(self):
self.assertEqual(run_prompt_nodes(None), [])
def test_run_prompt_nodes_negative_only_graph_is_empty(self):
graph = {
"1": {"class_type": "PromptManager", "inputs": {"text": "bad hands"}},
"2": {
"class_type": "KSampler",
"inputs": {"negative": ["1", 0], "positive": ["9", 0]},
},
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}},
}
self.assertEqual(run_prompt_nodes(graph), [])
def test_run_prompt_nodes_falls_back_when_no_sampler(self):
graph = {"1": {"class_type": "PromptManagerText", "inputs": {"text": "x"}}}
self.assertEqual(run_prompt_nodes(graph), ["1"])
class TestInputKeys(unittest.TestCase):
"""Which input names are followed when tracing the positive prompt upstream."""
def test_negative_flavoured_keys_are_skipped(self):
for key in ("negative_prompt", "text_negative", "neg_text", "neg", "negative"):
self.assertFalse(follows_input(key), key)
self.assertEqual(_upstream_links({key: ["1", 0]}), [], key)
def test_positive_and_plain_text_keys_are_followed(self):
for key in ("text_positive", "positive", "text", "prompt", "text_g"):
self.assertTrue(follows_input(key), key)
self.assertEqual(_upstream_links({key: ["1", 0]}), [["1", 0]], key)
def test_styler_node_only_follows_its_positive_text(self):
graph = {
"1": pm("a cat"),
"2": pm("bad hands"),
"3": {
"class_type": "SDXLPromptStyler",
"inputs": {"text_positive": ["1", 0], "text_negative": ["2", 0]},
},
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}},
"5": ksampler("4", "2"),
"20": LOADER,
}
self.assertEqual(positive_prompt_nodes(graph), ["1"])
self.assertEqual(run_prompt_nodes(graph), ["1"])
self.assertEqual(resolve_text(graph, "1"), "a cat")
class TestHasSampler(unittest.TestCase):
def test_structural_sampler_and_guider_are_detected(self):
self.assertTrue(has_sampler({"5": ksampler("1", "2"), "20": LOADER}))
guider = {
"class_type": "BasicGuider",
"inputs": {"model": ["20", 0], "conditioning": ["1", 0]},
}
self.assertTrue(has_sampler({"6": guider, "20": LOADER}))
def test_sampler_without_model_link_is_detected_by_name(self):
sampler = {
"class_type": "KSamplerAdvanced",
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
}
self.assertTrue(has_sampler({"5": sampler}))
def test_non_samplers_and_malformed_graphs(self):
controlnet = {
"class_type": "ControlNetApplyAdvanced",
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
}
self.assertFalse(has_sampler({"7": controlnet}))
exotic = {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}}
self.assertFalse(has_sampler({"9": exotic}))
for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}, {}):
self.assertFalse(has_sampler(bad), repr(bad))
class TestAdversarialGraphs(unittest.TestCase):
"""Queued and embedded graphs are untrusted: never hang, never raise."""
def _assert_fast(self, graph):
started = time.monotonic()
for fn in (positive_prompt_nodes, run_prompt_nodes, has_sampler):
fn(graph)
resolve_text(graph, "1")
self.assertLess(time.monotonic() - started, 1.0)
def test_cyclic_graph_terminates(self):
graph = {
"1": pm(["2", 0]),
"2": {"class_type": "StringConcatenate", "inputs": {"string_a": ["1", 0]}},
"3": {
"class_type": "ConditioningCombine",
"inputs": {"conditioning_1": ["3", 0], "conditioning_2": ["1", 0]},
},
"5": ksampler("3", "3"),
"20": LOADER,
}
self._assert_fast(graph)
self.assertEqual(run_prompt_nodes(graph), ["1"])
def test_ten_thousand_node_fan_out_terminates(self):
graph = {"20": LOADER}
combine_inputs = {}
for i in range(10_000):
graph[str(100 + i)] = pm(f"prompt {i}")
combine_inputs[f"conditioning_{i}"] = [str(100 + i), 0]
graph["3"] = {"class_type": "ConditioningCombine", "inputs": combine_inputs}
graph["5"] = ksampler("3", "3")
self._assert_fast(graph)
self.assertTrue(run_prompt_nodes(graph))
def test_non_dict_inputs_and_odd_values(self):
graph = {
"1": {"class_type": "PromptManager", "inputs": ["not", "a", "dict"]},
"2": {"class_type": "KSampler", "inputs": 42},
"3": {"class_type": None, "inputs": {"positive": "1", "negative": [1]}},
"4": "just a string",
"5": {
"class_type": "KSampler",
"inputs": {
"positive": [None, 0],
"negative": [{"x": 1}, 0],
"model": [[], 0],
},
},
6: {"class_type": "PromptManager", "inputs": {"text": 7}},
}
self._assert_fast(graph)
for fn in (positive_prompt_nodes, run_prompt_nodes):
self.assertIsInstance(fn(graph), list)
class TestResolveText(unittest.TestCase):
"""Text known at queue time: typed, or from known pure string nodes."""
def _linked(self, source):
return {"1": pm(["2", 0]), "2": source}
def test_typed_text(self):
self.assertEqual(resolve_text({"1": pm(" a cat ")}, "1"), "a cat")
def test_core_primitive_string(self):
for cls in ("PrimitiveString", "PrimitiveStringMultiline"):
graph = self._linked({"class_type": cls, "inputs": {"value": "a cat"}})
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
def test_core_string_concatenate(self):
graph = self._linked(
{
"class_type": "StringConcatenate",
"inputs": {
"string_a": "a cat",
"string_b": "at dusk",
"delimiter": ", ",
},
}
)
self.assertEqual(resolve_text(graph, "1"), "a cat, at dusk")
def test_was_text_concatenate_skips_empty_and_cleans_whitespace(self):
graph = self._linked(
{
"class_type": "Text Concatenate",
"inputs": {
"delimiter": " ",
"clean_whitespace": "true",
"text_a": " a cat ",
"text_b": "",
"text_c": ["3", 0],
},
}
)
graph["3"] = {"class_type": "PrimitiveString", "inputs": {"value": "at dusk"}}
self.assertEqual(resolve_text(graph, "1"), "a cat at dusk")
def test_pass_through_nodes(self):
for cls in ("ShowText|pysssss", "Text Multiline"):
graph = self._linked({"class_type": cls, "inputs": {"text": "a cat"}})
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
def test_unknown_and_batch_sources_are_never_guessed(self):
for source in (
{"class_type": "PromptSearchList", "inputs": {"search_text": "cat"}},
{
"class_type": "SomeTextReplace",
"inputs": {"text": "a cat", "find": "cat", "replace": "dog"},
},
):
self.assertIsNone(
resolve_text(self._linked(source), "1"), source["class_type"]
)
def test_same_source_used_twice_resolves_both_branches(self):
# StringConcatenate(string_a=X, string_b=X): X is not a cycle, it is
# simply shared, and must resolve in both branches.
graph = {
"1": pm(["2", 0]),
"2": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["3", 0],
"string_b": ["3", 0],
"delimiter": "+",
},
},
"3": {"class_type": "PrimitiveString", "inputs": {"value": "x"}},
}
self.assertEqual(resolve_text(graph, "1"), "x+x")
def test_diamond_graph_resolves(self):
# 1 <- 2 <- (3, 4) and both 3 and 4 read from 5
graph = {
"1": pm(["2", 0]),
"2": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["3", 0],
"string_b": ["4", 0],
"delimiter": " ",
},
},
"3": {"class_type": "ShowText|pysssss", "inputs": {"text": ["5", 0]}},
"4": {"class_type": "Text Multiline", "inputs": {"text": ["5", 0]}},
"5": {"class_type": "PrimitiveString", "inputs": {"value": "cat"}},
}
self.assertEqual(resolve_text(graph, "1"), "cat cat")
def test_indirect_cycle_returns_none(self):
graph = {
"1": pm(["2", 0]),
"2": {"class_type": "ShowText|pysssss", "inputs": {"text": ["3", 0]}},
"3": {"class_type": "Text Multiline", "inputs": {"text": ["2", 0]}},
}
self.assertIsNone(resolve_text(graph, "1"))
@staticmethod
def _diamond_chain(levels, leaf):
"""Node i reads node i+1 on both inputs: 2**levels paths, levels nodes."""
graph = {"0": {"class_type": "PrimitiveString", "inputs": {"value": leaf}}}
for i in range(1, levels + 1):
prev = [str(i - 1), 0]
graph[str(i)] = {
"class_type": "StringConcatenate",
"inputs": {"string_a": prev, "string_b": prev, "delimiter": ""},
}
graph["pm"] = pm([str(levels), 0])
return graph
def _count_expansions(self, graph):
real = prompt_graph._resolve_string_node
with mock.patch.object(
prompt_graph, "_resolve_string_node", side_effect=real
) as spy:
start = time.monotonic()
text = resolve_text(graph, "pm")
elapsed = time.monotonic() - start
return text, spy.call_count, elapsed
def test_forty_level_diamond_chain_expands_each_node_once(self):
text, expansions, elapsed = self._count_expansions(self._diamond_chain(40, ""))
self.assertIsNone(text) # the leaf is empty, so the prompt is empty
self.assertLessEqual(expansions, 41)
self.assertLess(elapsed, 0.5)
def test_diamond_chain_with_text_stays_within_budget(self):
# Doubling "a" 400 times would be a 2**400-character string: the
# length cap turns it into None long before, and the visit budget
# holds regardless of how the graph is shaped.
text, expansions, elapsed = self._count_expansions(
self._diamond_chain(400, "a")
)
self.assertIsNone(text)
self.assertLessEqual(expansions, MAX_VISITS)
self.assertLess(elapsed, 0.5)
def test_shallow_diamond_chain_resolves_doubled_text(self):
text, expansions, _ = self._count_expansions(self._diamond_chain(4, "ab"))
self.assertEqual(text, "ab" * 16)
self.assertLessEqual(expansions, 5)
def test_chain_deeper_than_max_depth_is_none_without_raising(self):
graph = {"0": {"class_type": "PrimitiveString", "inputs": {"value": "a"}}}
for i in range(1, 5000):
graph[str(i)] = {
"class_type": "ShowText|pysssss",
"inputs": {"text": [str(i - 1), 0]},
}
graph["pm"] = pm(["4999", 0])
self.assertIsNone(resolve_text(graph, "pm"))
graph["pm"] = pm([str(MAX_DEPTH - 1), 0])
self.assertEqual(resolve_text(graph, "pm"), "a")
def test_resolved_text_longer_than_cap_is_none(self):
graph = {
"1": pm(["2", 0]),
"2": {
"class_type": "PrimitiveString",
"inputs": {"value": "x" * (MAX_TEXT_LENGTH + 1)},
},
}
self.assertIsNone(resolve_text(graph, "1"))
graph["2"]["inputs"]["value"] = "x" * MAX_TEXT_LENGTH
self.assertEqual(len(resolve_text(graph, "1")), MAX_TEXT_LENGTH)
def test_cycles_and_missing_links(self):
graph = {
"1": pm(["2", 0]),
"2": {
"class_type": "StringConcatenate",
"inputs": {"string_a": ["2", 0], "string_b": "x", "delimiter": ""},
},
}
self.assertIsNone(resolve_text(graph, "1"))
self.assertIsNone(resolve_text({"1": pm(["9", 0])}, "1"))
if __name__ == "__main__":
unittest.main()
+282
View File
@@ -0,0 +1,282 @@
"""Tests for prompt usage tracking: run_count, last_used_at and server-side sorting."""
import os
import sqlite3
from contextlib import closing
import sys
import tempfile
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.open_handles import assert_closed # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
from database.models import PromptModel
from database.operations import PromptDatabase
from prompt_manager_base import PromptManagerBase
from utils.hashing import generate_prompt_hash
class UsageTestCase(unittest.TestCase):
def setUp(self):
fd, self.path = tempfile.mkstemp(suffix=".db")
os.close(fd)
self.db = PromptDatabase(self.path)
def tearDown(self):
if getattr(self, "db", None) is not None:
self.db.close_all()
PromptModel.close_all_instances()
assert_closed(self, self.path)
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
def _save(self, text, **kwargs):
return self.db.save_prompt(
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
)
def _set(self, prompt_id, **columns):
with closing(sqlite3.connect(self.path)) as conn, conn:
for column, value in columns.items():
conn.execute(
f"UPDATE prompts SET {column} = ? WHERE id = ?", (value, prompt_id)
)
def _ids(self, prompts):
return [p["id"] for p in prompts]
class TestRecordPromptUse(UsageTestCase):
def test_new_prompt_starts_unused_and_unstamped(self):
# Saving is not using: only record_prompt_use stamps last_used_at
prompt = self.db.get_prompt_by_id(self._save("fresh"))
self.assertEqual(prompt["run_count"], 0)
self.assertIsNone(prompt["last_used_at"])
def test_never_run_prompts_sort_after_run_ones(self):
ran = self._save("run once")
later = self._save("saved later, never run")
self.db.record_prompt_use(ran)
last = self._save("saved last, never run")
result = self.db.get_recent_prompts(limit=10, sort="last_used_desc")
# The run prompt leads; never-run prompts follow, newest first
self.assertEqual(self._ids(result["prompts"]), [ran, last, later])
searched = self.db.search_prompts(text="run", sort="last_used_desc")
self.assertEqual(self._ids(searched)[0], ran)
def test_most_used_sort_breaks_ties_with_unstamped_last(self):
a = self._save("a")
b = self._save("b")
self._set(a, run_count=2, last_used_at="2026-01-01T00:00:00.000+00:00")
self._set(b, run_count=2) # same count, never stamped
result = self.db.get_recent_prompts(limit=10, sort="run_count_desc")
self.assertEqual(self._ids(result["prompts"]), [a, b])
def test_record_use_increments_and_touches_last_used(self):
pid = self._save("used")
self._set(pid, last_used_at="2020-01-01T00:00:00.000+00:00")
self.assertTrue(self.db.record_prompt_use(pid))
self.assertTrue(self.db.record_prompt_use(pid))
prompt = self.db.get_prompt_by_id(pid)
self.assertEqual(prompt["run_count"], 2)
self.assertGreater(prompt["last_used_at"], "2020-01-01T00:00:00.000+00:00")
def test_rerun_right_after_another_save_still_sorts_first(self):
# No hand-set timestamps: creation and re-run happen within the same millisecond
old = self._save("old")
self._save("new")
self.db.record_prompt_use(old)
result = self.db.get_recent_prompts(limit=1, sort="last_used_desc")
self.assertEqual(self._ids(result["prompts"]), [old])
def test_microsecond_timestamps_order_correctly_against_backfilled_ones(self):
backfilled = self._save("backfilled")
self._set(backfilled, last_used_at="2026-01-01T00:00:00.000+00:00")
recent = self._save("recent")
self.db.record_prompt_use(recent)
self.assertRegex(
self.db.get_prompt_by_id(recent)["last_used_at"],
r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d\.\d{6}\+00:00$",
)
result = self.db.get_recent_prompts(limit=2, sort="last_used_desc")
self.assertEqual(self._ids(result["prompts"]), [recent, backfilled])
def test_record_use_can_add_several_runs_at_once(self):
pid = self._save("several")
self.db.record_prompt_use(pid, times=3)
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 3)
def test_record_use_on_missing_prompt_returns_false(self):
self.assertFalse(self.db.record_prompt_use(999999))
class TestServerSideSort(UsageTestCase):
def setUp(self):
super().setUp()
self.old = self._save("old prompt", rating=5)
self.mid = self._save("middle prompt", rating=1)
self.new = self._save("new prompt")
self._set(self.old, last_used_at="2026-01-01T00:00:00.000+00:00", run_count=9)
self._set(self.mid, last_used_at="2026-01-02T00:00:00.000+00:00", run_count=1)
self._set(self.new, last_used_at="2026-01-03T00:00:00.000+00:00", run_count=3)
def test_rerun_bubbles_old_prompt_to_top_across_pages(self):
self.db.record_prompt_use(self.old)
page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="last_used_desc")
self.assertEqual(self._ids(page1["prompts"]), [self.old])
def test_most_used(self):
result = self.db.get_recent_prompts(limit=10, sort="run_count_desc")
self.assertEqual(self._ids(result["prompts"]), [self.old, self.new, self.mid])
def test_rating_sort_puts_unrated_last_and_spans_pages(self):
page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="rating_desc")
self.assertEqual(self._ids(page1["prompts"]), [self.old])
result = self.db.get_recent_prompts(limit=10, sort="rating_asc")
self.assertEqual(self._ids(result["prompts"]), [self.mid, self.old, self.new])
def test_text_sort(self):
result = self.db.get_recent_prompts(limit=10, sort="text_asc")
self.assertEqual(self._ids(result["prompts"]), [self.mid, self.new, self.old])
def test_default_and_unknown_sort_keep_newest_first(self):
expected = [self.new, self.mid, self.old]
self.assertEqual(
self._ids(self.db.get_recent_prompts(limit=10)["prompts"]), expected
)
injected = self.db.get_recent_prompts(limit=10, sort="id; DROP TABLE prompts")
self.assertEqual(self._ids(injected["prompts"]), expected)
def test_search_accepts_sort(self):
self.db.record_prompt_use(self.mid)
results = self.db.search_prompts(text="prompt", sort="last_used_desc")
self.assertEqual(self._ids(results)[0], self.mid)
@unittest.skipIf(
sqlite3.sqlite_version_info < (3, 35, 0),
"simulating the old schema needs ALTER TABLE DROP COLUMN (SQLite 3.35+)",
)
class TestUsageMigration(unittest.TestCase):
"""Databases from before 3.2.4 get the columns and an estimated backfill."""
def setUp(self):
fd, self.path = tempfile.mkstemp(suffix=".db")
os.close(fd)
self.dbs = []
db = self._open()
self.with_images = db.save_prompt(text="has images", prompt_hash="h1")
self.no_images = db.save_prompt(text="no images", prompt_hash="h2")
db.close_all()
with closing(sqlite3.connect(self.path)) as conn, conn:
conn.execute(
"UPDATE prompts SET created_at = '2026-01-01T10:00:00.000000+00:00'"
)
# generated_images uses SQLite's space-separated CURRENT_TIMESTAMP format
for i, ts in enumerate(("2026-03-05 08:00:00", "2026-02-01 09:00:00")):
conn.execute(
"INSERT INTO generated_images"
" (prompt_id, image_path, filename, generation_time)"
" VALUES (?, ?, ?, ?)",
(self.with_images, f"/out/{i}.png", f"{i}.png", ts),
)
# Simulate a pre-3.2.4 schema
for index in ("idx_prompts_last_used", "idx_prompts_run_count"):
conn.execute(f"DROP INDEX IF EXISTS {index}")
conn.execute("ALTER TABLE prompts DROP COLUMN last_used_at")
conn.execute("ALTER TABLE prompts DROP COLUMN run_count")
# Schema init runs once per process per path; the next PromptDatabase()
# must look at the file again to see the downgraded schema.
PromptModel.reset_schema_cache()
def _open(self):
"""Open the database under test and close it again in tearDown."""
db = PromptDatabase(self.path)
self.dbs.append(db)
return db
def tearDown(self):
for db in self.dbs:
db.close_all()
PromptModel.close_all_instances()
assert_closed(self, self.path)
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
def test_backfill_estimates_usage_from_linked_images(self):
db = self._open()
used = db.get_prompt_by_id(self.with_images)
unused = db.get_prompt_by_id(self.no_images)
self.assertEqual(used["run_count"], 2)
self.assertTrue(used["last_used_at"].startswith("2026-03-05T08:00:00"))
self.assertEqual(unused["run_count"], 1)
self.assertTrue(unused["last_used_at"].startswith("2026-01-01T10:00:00"))
def test_backfill_runs_only_when_the_columns_are_added(self):
db = self._open()
# A row whose usage is unknown (e.g. written by 3.2.3 after a downgrade)
# stays "never run": a restart must not invent a run for it, otherwise
# Recently Used could never tell saved-but-unused prompts from run ones.
with closing(sqlite3.connect(self.path)) as conn, conn:
conn.execute(
"INSERT INTO prompts (text, hash, created_at)"
" VALUES ('from 3.2.3', 'h3', '2026-04-01T00:00:00.000000+00:00')"
)
legacy_id = conn.execute(
"SELECT id FROM prompts WHERE hash = 'h3'"
).fetchone()[0]
self.assertIsNone(db.get_prompt_by_id(legacy_id)["last_used_at"])
PromptModel.reset_schema_cache()
later = self._open().get_prompt_by_id(legacy_id)
self.assertIsNone(later["last_used_at"])
self.assertEqual(later["run_count"], 0)
def test_migration_runs_once(self):
db = self._open()
db.record_prompt_use(self.no_images)
reopened = self._open()
self.assertEqual(reopened.get_prompt_by_id(self.no_images)["run_count"], 2)
class TestNodeRecordsUse(UsageTestCase):
"""Node saves count a run only when asked; re-runs are counted by the queue hook.
ComfyUI skips unchanged nodes, so counting at execution misses re-runs; see
tests/test_usage_tracking.py for the hook that counts them.
"""
def _node(self):
node = PromptManagerBase.__new__(PromptManagerBase)
node.db = self.db
node.logger = __import__("logging").getLogger("test.prompt_usage")
return node
def test_save_without_count_run_leaves_counting_to_the_queue_hook(self):
node = self._node()
first = node._save_prompt_to_database("a lighthouse at dusk")
second = node._save_prompt_to_database("a lighthouse at dusk")
self.assertEqual(first, second)
self.assertEqual(self.db.get_prompt_by_id(first)["run_count"], 0)
def test_count_run_counts_new_and_existing_prompts(self):
node = self._node()
pid = node._save_prompt_to_database("batch item", count_run=True)
node._save_prompt_to_database("batch item", count_run=True)
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2)
if __name__ == "__main__":
unittest.main()
+92
View File
@@ -0,0 +1,92 @@
"""Unit tests for py/api/scan_job.py: fan-out of scan progress to SSE subscribers."""
import asyncio
import contextlib
import os
import sys
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py.api.scan_job import ScanJob # noqa: E402
async def _collect(job, stop_after=None):
seen = []
async with contextlib.aclosing(job.events()) as events:
async for payload in events:
seen.append(payload)
if stop_after is not None and len(seen) >= stop_after:
break
return seen
class TestScanJob(unittest.IsolatedAsyncioTestCase):
async def test_new_job_is_running_with_no_event(self):
job = ScanJob()
self.assertTrue(job.running)
self.assertIsNone(job.last_event)
async def test_subscriber_receives_events_until_finish(self):
job = ScanJob()
consumer = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "progress", "progress": 10})
job.publish({"type": "complete"})
job.finish()
self.assertEqual(
await asyncio.wait_for(consumer, 2),
[{"type": "progress", "progress": 10}, {"type": "complete"}],
)
self.assertFalse(job.running)
async def test_late_subscriber_gets_latest_snapshot_then_live_events(self):
job = ScanJob()
job.publish({"type": "progress", "progress": 10})
job.publish({"type": "progress", "progress": 40})
consumer = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "complete"})
job.finish()
self.assertEqual(
await asyncio.wait_for(consumer, 2),
[{"type": "progress", "progress": 40}, {"type": "complete"}],
)
async def test_subscriber_after_finish_gets_only_the_final_event(self):
job = ScanJob()
job.publish({"type": "progress", "progress": 99})
job.publish({"type": "complete"})
job.finish()
self.assertEqual(await _collect(job), [{"type": "complete"}])
async def test_two_subscribers_both_see_every_event(self):
job = ScanJob()
first = asyncio.ensure_future(_collect(job))
second = asyncio.ensure_future(_collect(job))
await asyncio.sleep(0)
job.publish({"type": "progress", "progress": 1})
job.finish()
self.assertEqual(await first, [{"type": "progress", "progress": 1}])
self.assertEqual(await second, [{"type": "progress", "progress": 1}])
async def test_departed_subscriber_is_forgotten(self):
job = ScanJob()
job.publish({"type": "progress", "progress": 1})
await _collect(job, stop_after=1)
self.assertEqual(job.subscriber_count, 0)
job.publish({"type": "complete"}) # must not raise or queue anywhere
job.finish()
if __name__ == "__main__":
unittest.main()
+2 -1
View File
@@ -32,7 +32,8 @@ def apply_filters(texts, skip_multipart=True):
class TestNewlineCollapse(unittest.TestCase):
"""Newlines must be collapsed so StringOutputList treats each DB entry as one prompt."""
"""Newlines must be collapsed so StringOutputList treats each DB entry as one
prompt."""
def test_single_line_unchanged(self):
result = apply_filters(["a beautiful landscape"])
+39
View File
@@ -0,0 +1,39 @@
"""The UI's sort options must match the server's whitelist exactly."""
import os
import re
import sys
import unittest
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from database.operations import SORT_ORDERS
SORT_MODULE = os.path.join(ROOT, "web", "js", "prompt-list-sort.js")
NODE_EXTENSION = os.path.join(ROOT, "web", "comfy", "prompt_manager.js")
class TestSortContract(unittest.TestCase):
def test_ui_sort_options_match_server_whitelist(self):
with open(SORT_MODULE, encoding="utf-8") as fh:
source = fh.read()
ui_values = set(re.findall(r'value:\s*"([a-z_]+)"', source))
self.assertTrue(ui_values, "no sort options found in prompt-list-sort.js")
self.assertEqual(ui_values, set(SORT_ORDERS))
def test_node_recent_button_requests_recently_used(self):
# The node extension can't import the shared module (ComfyUI load order is
# not guaranteed), so it sends the key directly; it must be a server key.
with open(NODE_EXTENSION, encoding="utf-8") as fh:
source = fh.read()
match = re.search(r"/prompt_manager/recent\?[^`'\"]*sort=([a-z_]+)", source)
self.assertIsNotNone(match, "node Recent request does not send a sort")
self.assertEqual(match.group(1), "last_used_desc")
self.assertIn(match.group(1), SORT_ORDERS)
if __name__ == "__main__":
unittest.main()
+566
View File
@@ -0,0 +1,566 @@
"""Run counting and node roles: only positive PromptManager nodes count and link.
Runs are resolved when a prompt is queued (ComfyUI's on-prompt hook, which runs
before validation) and counted when ComfyUI dequeues the prompt for execution,
because ComfyUI skips unchanged nodes, so node execution can't be relied on to
count, and a rejected or cancelled queue entry must not count either.
"""
import os
import sys
import tempfile
import threading
import unittest
import unittest.mock as mock
import uuid
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.operations import PromptDatabase
from prompt_manager_base import PromptManagerBase
from utils.hashing import generate_prompt_hash
from utils.prompt_tracker import PromptTracker
from utils import usage_tracking
from utils.usage_tracking import (
PendingFirstUse,
QueuedRuns,
_wrap_queue_get,
count_queued_prompt,
handle_queued_prompt,
on_dequeue,
register_queue_hook,
)
LOADER = {
"class_type": "CheckpointLoaderSimple",
"inputs": {"ckpt_name": "m.safetensors"},
}
def workflow(positive_text="a cat", negative_text="blurry"):
"""API graph with a positive (#134) and a negative (#176) PromptManager node."""
return {
"176": {"class_type": "PromptManager", "inputs": {"text": negative_text}},
"134": {"class_type": "PromptManager", "inputs": {"text": positive_text}},
"5": {
"class_type": "KSampler",
"inputs": {
"positive": ["134", 0],
"negative": ["176", 0],
"model": ["20", 0],
},
},
"20": LOADER,
}
class FakePromptQueue:
"""Stand-in for execution.PromptQueue: get() pops preset results."""
def __init__(self):
self.results = []
self.get_calls = []
def get(self, timeout=None):
self.get_calls.append(timeout)
return self.results.pop(0) if self.results else None
class FakeServer:
"""Stand-in for PromptServer: on-prompt handlers plus an optional queue."""
def __init__(self, with_queue=True):
self.handlers = []
if with_queue:
self.prompt_queue = FakePromptQueue()
def add_on_prompt_handler(self, handler):
self.handlers.append(handler)
def on_prompt(self, json_data):
for handler in self.handlers:
json_data = handler(json_data)
return json_data
def queue_item(prompt_id, graph):
"""What PromptQueue.get returns for a prompt that passed validation."""
return ((0, prompt_id, graph, {}, [], False), 1)
class DbTestCase(unittest.TestCase):
def setUp(self):
fd, self.path = tempfile.mkstemp(suffix=".db")
os.close(fd)
self.db = PromptDatabase(self.path)
self.pending = PendingFirstUse()
self.queued = QueuedRuns()
def tearDown(self):
close = getattr(self.db, "close", None)
if callable(close):
close()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
def _save(self, text):
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
def _runs(self, prompt_id):
return self.db.get_prompt_by_id(prompt_id)["run_count"]
def _queue(self, graph):
"""Queue a graph and let ComfyUI dequeue it for execution."""
request = handle_queued_prompt({"prompt": graph, "client_id": "x"}, self.queued)
on_dequeue(
request.get("prompt_id", ""), self.queued, lambda: self.db, self.pending
)
return request
class TestQueueHook(DbTestCase):
def test_counts_the_positive_prompt_only(self):
positive, negative = self._save("a cat"), self._save("blurry")
self._queue(workflow())
self.assertEqual(self._runs(positive), 1)
self.assertEqual(self._runs(negative), 0)
def test_requeue_counts_again_even_when_comfyui_skips_the_cached_node(self):
positive = self._save("a cat")
self._queue(workflow())
self._queue(workflow()) # no node execution in between
self.assertEqual(self._runs(positive), 2)
def test_new_prompt_is_counted_when_the_node_first_saves_it(self):
self._queue(workflow(positive_text="a brand new prompt"))
self.assertTrue(
self.pending.consume(generate_prompt_hash("a brand new prompt"))
)
self.assertFalse(
self.pending.consume(generate_prompt_hash("a brand new prompt"))
)
def test_same_text_in_two_positive_nodes_counts_once(self):
positive = self._save("a cat")
graph = workflow()
graph["135"] = {"class_type": "PromptManager", "inputs": {"text": "a cat"}}
graph["4"] = {
"class_type": "ConditioningCombine",
"inputs": {"conditioning_1": ["134", 0], "conditioning_2": ["135", 0]},
}
graph["5"]["inputs"]["positive"] = ["4", 0]
self._queue(graph)
self.assertEqual(self._runs(positive), 1)
def test_text_from_a_known_string_node_is_counted_at_queue_time(self):
positive = self._save("a cat")
graph = workflow()
graph["134"]["inputs"]["text"] = ["8", 0]
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
self._queue(graph)
self._queue(graph) # cached re-run: the node won't execute
self.assertEqual(self._runs(positive), 2)
def test_new_prompt_queued_twice_before_first_run_counts_twice(self):
self._queue(workflow(positive_text="queued twice"))
self._queue(workflow(positive_text="queued twice"))
self.assertEqual(self.pending.consume(generate_prompt_hash("queued twice")), 2)
def test_linked_text_is_left_to_node_execution(self):
positive = self._save("from a batch")
graph = workflow()
graph["134"]["inputs"]["text"] = ["7", 0]
graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
self._queue(graph)
self.assertEqual(self._runs(positive), 0)
def test_returns_the_request_and_never_raises(self):
request = {"prompt": workflow(), "client_id": "x", "extra_data": {}}
self.assertIs(handle_queued_prompt(request, self.queued), request)
for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []):
self.assertIs(handle_queued_prompt(bad, self.queued), bad)
broken_db = mock.Mock()
broken_db.get_prompt_by_hash.side_effect = RuntimeError("db locked")
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
on_dequeue(
request["prompt_id"], self.queued, lambda: broken_db, self.pending
)
def test_graph_without_prompt_manager_nodes_is_not_stashed_or_tagged(self):
request = {"prompt": {"20": LOADER}}
self.assertIs(handle_queued_prompt(request, self.queued), request)
self.assertNotIn("prompt_id", request)
self.assertEqual(len(self.queued), 0)
class TestCountAtDequeue(DbTestCase):
"""End to end through the registered hook and the wrapped queue."""
def setUp(self):
super().setUp()
self.server = FakeServer()
self._register(self.server)
def _register(self, server):
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
return register_queue_hook(
server, lambda: self.db, pending=self.pending, queued=self.queued
)
def _submit(self, graph, prompt_id=None):
request = {"prompt": graph, "client_id": "x"}
if prompt_id is not None:
request["prompt_id"] = prompt_id
return self.server.on_prompt(request)
def _dequeue(self, request):
self.server.prompt_queue.results.append(
queue_item(request["prompt_id"], request["prompt"])
)
return self.server.prompt_queue.get(timeout=1.0)
def test_hook_adds_a_lowercase_uuid4_prompt_id_when_absent(self):
request = self._submit(workflow())
prompt_id = request["prompt_id"]
self.assertEqual(str(uuid.UUID(prompt_id, version=4)), prompt_id)
self.assertEqual(prompt_id, prompt_id.lower())
def test_hook_keeps_an_existing_prompt_id(self):
request = self._submit(workflow(), prompt_id="client-chosen")
self.assertEqual(request["prompt_id"], "client-chosen")
def test_hook_alone_does_not_count(self):
positive = self._save("a cat")
self._submit(workflow())
self.assertEqual(self._runs(positive), 0)
def test_dequeue_counts_exactly_once(self):
positive = self._save("a cat")
request = self._submit(workflow())
result = self._dequeue(request)
self.assertEqual(result[0][1], request["prompt_id"])
self.assertEqual(self.server.prompt_queue.get_calls, [1.0])
self.assertEqual(self._runs(positive), 1)
self._dequeue(request) # a second dequeue of the same id has nothing left
self.assertEqual(self._runs(positive), 1)
def test_two_queued_prompts_are_each_counted_once(self):
positive = self._save("a cat")
first = self._submit(workflow())
second = self._submit(workflow())
self.assertNotEqual(first["prompt_id"], second["prompt_id"])
self._dequeue(first)
self.assertEqual(self._runs(positive), 1)
self._dequeue(second)
self.assertEqual(self._runs(positive), 2)
def test_rejected_prompt_is_never_counted(self):
positive = self._save("a cat")
self._submit(workflow()) # validation fails: never dequeued
other = self._submit(workflow())
self._dequeue(other)
self.assertEqual(self._runs(positive), 1)
def test_get_returning_none_is_a_noop(self):
positive = self._save("a cat")
self._submit(workflow())
self.assertIsNone(self.server.prompt_queue.get(timeout=0.1))
self.assertEqual(self._runs(positive), 0)
def test_unknown_text_is_pending_only_after_dequeue(self):
prompt_hash = generate_prompt_hash("never saved")
request = self._submit(workflow(positive_text="never saved"))
self.assertEqual(self.pending.consume(prompt_hash), 0)
self._dequeue(request)
self.assertEqual(self.pending.consume(prompt_hash), 1)
def test_shorter_and_longer_queue_items_are_tolerated(self):
positive = self._save("a cat")
request = self._submit(workflow())
self.server.prompt_queue.results.append(
((0, request["prompt_id"], request["prompt"], {}, []), 1)
)
self.server.prompt_queue.get()
self.assertEqual(self._runs(positive), 1)
request = self._submit(workflow())
self.server.prompt_queue.results.append(
((0, request["prompt_id"], request["prompt"], {}, [], False, "x"), 2)
)
self.server.prompt_queue.get()
self.assertEqual(self._runs(positive), 2)
def test_malformed_queue_item_is_logged_and_returned(self):
item = (("not", "a", "tuple"),)
self.server.prompt_queue.results.append(item)
self.server.prompt_queue.results.append("garbage")
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
self.assertEqual(self.server.prompt_queue.get(), item)
self.assertEqual(self.server.prompt_queue.get(), "garbage")
def test_db_factory_raising_at_dequeue_is_logged_not_raised(self):
def broken():
raise RuntimeError("PromptManager database unavailable")
queued = QueuedRuns()
server = FakeServer()
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
register_queue_hook(server, broken, pending=self.pending, queued=queued)
request = server.on_prompt({"prompt": workflow()})
server.prompt_queue.results.append(queue_item(request["prompt_id"], {}))
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs:
self.assertIsNotNone(server.prompt_queue.get())
self.assertIn("database unavailable", "\n".join(logs.output))
def test_server_without_prompt_queue_counts_at_queue_time_and_warns_once(self):
positive = self._save("a cat")
server = FakeServer(with_queue=False)
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs:
self.assertTrue(self._register(server))
warnings = [line for line in logs.output if line.startswith("WARNING")]
self.assertEqual(len(warnings), 1)
server.on_prompt({"prompt": workflow()})
self.assertEqual(self._runs(positive), 1)
server.on_prompt({"prompt": workflow(positive_text="new text")})
self.assertEqual(self.pending.consume(generate_prompt_hash("new text")), 1)
def test_server_whose_queue_has_no_get_uses_legacy_counting(self):
positive = self._save("a cat")
server = FakeServer(with_queue=False)
server.prompt_queue = object()
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
self._register(server)
server.on_prompt({"prompt": workflow()})
self.assertEqual(self._runs(positive), 1)
def test_legacy_counting_calls_db_factory_inside_the_guard(self):
def broken():
raise RuntimeError("PromptManager database unavailable")
request = {"prompt": workflow()}
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
self.assertIs(count_queued_prompt(request, broken, self.pending), request)
def test_wrapping_twice_does_not_double_count(self):
positive = self._save("a cat")
prompt_queue = self.server.prompt_queue
wrapped_get = prompt_queue.get
calls = []
self.assertFalse(_wrap_queue_get(prompt_queue, calls.append))
self.assertIs(prompt_queue.get, wrapped_get)
request = self._submit(workflow())
self._dequeue(request)
self.assertEqual(self._runs(positive), 1)
self.assertEqual(calls, [])
def test_wrap_rejects_queues_without_a_callable_get(self):
self.assertFalse(_wrap_queue_get(object(), lambda prompt_id: None))
self.assertFalse(_wrap_queue_get(None, lambda prompt_id: None))
def test_wrapper_reports_the_dequeued_prompt_id(self):
prompt_queue = FakePromptQueue()
seen = []
self.assertTrue(_wrap_queue_get(prompt_queue, seen.append))
prompt_queue.results.append(queue_item("abc", {}))
self.assertEqual(prompt_queue.get(timeout=2.5), queue_item("abc", {}))
self.assertEqual(prompt_queue.get_calls, [2.5])
self.assertEqual(seen, ["abc"])
def test_dequeue_callback_errors_are_logged_not_raised(self):
prompt_queue = FakePromptQueue()
def explode(prompt_id):
raise ValueError("boom")
_wrap_queue_get(prompt_queue, explode)
prompt_queue.results.append(queue_item("abc", {}))
with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"):
self.assertIsNotNone(prompt_queue.get())
class TestQueuedRuns(unittest.TestCase):
def test_pop_returns_runs_once(self):
queued = QueuedRuns()
queued.put("p1", [("h1", "text one")])
self.assertEqual(queued.pop("p1"), [("h1", "text one")])
self.assertEqual(queued.pop("p1"), [])
self.assertEqual(queued.pop("missing"), [])
def test_bound_evicts_oldest(self):
queued = QueuedRuns(max_entries=3)
for i in range(4):
queued.put(f"p{i}", [(f"h{i}", f"text {i}")])
self.assertEqual(len(queued), 3)
self.assertEqual(queued.pop("p0"), [])
self.assertEqual(queued.pop("p3"), [("h3", "text 3")])
def test_default_bound_is_one_thousand(self):
queued = QueuedRuns()
for i in range(1001):
queued.put(str(i), [("h", "t")])
self.assertEqual(len(queued), 1000)
self.assertEqual(queued.pop("0"), [])
self.assertEqual(queued.pop("1000"), [("h", "t")])
def test_requeue_of_the_same_id_replaces_the_runs(self):
queued = QueuedRuns()
queued.put("p1", [("h1", "one")])
queued.put("p1", [("h2", "two")])
self.assertEqual(queued.pop("p1"), [("h2", "two")])
def test_empty_runs_are_not_stored(self):
queued = QueuedRuns()
queued.put("p1", [])
self.assertEqual(len(queued), 0)
class TestPendingFirstUse(unittest.TestCase):
def test_entries_expire(self):
pending = PendingFirstUse(ttl_seconds=10)
with mock.patch("utils.usage_tracking.time.monotonic", return_value=100.0):
pending.add("h")
with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0):
self.assertFalse(pending.consume("h"))
def test_consume_returns_how_many_runs_were_pending(self):
pending = PendingFirstUse()
pending.add("h")
pending.add("h")
pending.add("h")
self.assertEqual(pending.consume("h"), 3)
self.assertEqual(pending.consume("h"), 0)
def test_size_is_capped(self):
pending = PendingFirstUse(max_entries=3)
for h in ("a", "b", "c", "d"):
pending.add(h)
self.assertFalse(pending.consume("a")) # oldest evicted
self.assertTrue(pending.consume("d"))
class FakeTracker:
def __init__(self):
self.calls = []
def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True):
self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue})
return "exec-1"
class TestNodeRoles(DbTestCase):
def _node(self):
node = PromptManagerBase.__new__(PromptManagerBase)
node.db = self.db
node.prompt_tracker = FakeTracker()
node.logger = __import__("logging").getLogger("test.usage_tracking")
node.pending_first_use = self.pending
return node
def _run(self, node, graph, unique_id, text):
return node._track_prompt_execution(
text=text,
encoding_text=text,
category=None,
tags=None,
additional_data={},
prompt_graph=graph,
unique_id=unique_id,
)
def test_negative_node_saves_but_never_becomes_the_current_prompt(self):
node = self._node()
prompt_id = self._run(node, workflow(), "176", "blurry")
self.assertIsNotNone(self.db.get_prompt_by_id(prompt_id))
self.assertEqual(node.prompt_tracker.calls, [])
self.assertEqual(self._runs(prompt_id), 0)
def test_positive_typed_text_is_current_without_queue_push(self):
node = self._node()
self._run(node, workflow(), "134", "a cat")
self.assertEqual(
node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
)
def test_first_save_after_queue_hook_counts_one_run(self):
self._queue(workflow(positive_text="never seen before"))
prompt_id = self._run(
self._node(),
workflow(positive_text="never seen before"),
"134",
"never seen before",
)
self.assertEqual(self._runs(prompt_id), 1)
def test_new_prompt_queued_twice_counts_both_runs_on_first_save(self):
graph = workflow(positive_text="double queued")
self._queue(graph)
self._queue(graph)
prompt_id = self._run(self._node(), graph, "134", "double queued")
self.assertEqual(self._runs(prompt_id), 2)
def test_text_resolved_at_queue_time_is_not_counted_again_or_queued(self):
graph = workflow()
graph["134"]["inputs"]["text"] = ["8", 0]
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
positive = self._save("a cat")
self._queue(graph)
node = self._node()
self._run(node, graph, "134", "a cat")
self.assertEqual(self._runs(positive), 1)
self.assertEqual(
node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
)
def test_resolution_mismatch_falls_back_to_counting_and_queueing(self):
# Hook guessed "a cat" but the node received different text: trust the node
graph = workflow()
graph["134"]["inputs"]["text"] = ["8", 0]
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
node = self._node()
prompt_id = self._run(node, graph, "134", "actually a dog")
self.assertEqual(self._runs(prompt_id), 1)
self.assertEqual(
node.prompt_tracker.calls,
[{"text": "actually a dog", "push_to_queue": True}],
)
def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self):
positive = self._save("a cat")
self._queue(workflow())
self._run(self._node(), workflow(), "134", "a cat")
self.assertEqual(self._runs(positive), 1)
def test_positive_linked_text_pushes_queue_and_counts_each_item(self):
graph = workflow()
graph["134"]["inputs"]["text"] = ["7", 0]
graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
node = self._node()
first = self._run(node, graph, "134", "batch item one")
self._run(node, graph, "134", "batch item one")
self.assertEqual(self._runs(first), 2)
self.assertTrue(all(c["push_to_queue"] for c in node.prompt_tracker.calls))
def test_without_graph_context_keeps_legacy_behaviour(self):
node = self._node()
self._run(node, None, None, "standalone")
self.assertEqual(
node.prompt_tracker.calls, [{"text": "standalone", "push_to_queue": True}]
)
class TestTrackerQueueFlag(unittest.TestCase):
def setUp(self):
self.tracker = PromptTracker(mock.Mock())
def test_push_to_queue_false_keeps_the_batch_queue_empty(self):
self.tracker.set_current_prompt("a cat", {"prompt_id": 1}, push_to_queue=False)
self.assertIsNone(self.tracker.pop_next_prompt())
self.assertEqual(self.tracker.get_current_prompt()["id"], 1)
def test_default_still_pushes(self):
self.tracker.set_current_prompt("a cat", {"prompt_id": 1})
self.assertEqual(self.tracker.pop_next_prompt()["id"], 1)
if __name__ == "__main__":
unittest.main()
+258
View File
@@ -0,0 +1,258 @@
"""GalleryDiagnostics against a temporary database and stubbed ComfyUI paths."""
import os
import pathlib
import sqlite3
from contextlib import closing
import sys
import tempfile
import types
import unittest
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.operations import PromptDatabase
from utils import diagnostics
from utils.diagnostics import GalleryDiagnostics, run_diagnostics
from utils.hashing import generate_prompt_hash
def same_path(a, b):
return os.path.normcase(os.path.realpath(a)) == os.path.normcase(
os.path.realpath(b)
)
class DiagnosticsTestCase(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
self.addCleanup(self.tmp.cleanup)
self.db_path = os.path.join(self.tmp.name, "prompts.db")
self.diag = GalleryDiagnostics(self.db_path)
def _populate(self):
db = PromptDatabase(self.db_path)
try:
prompt_id = db.save_prompt(
text="a cat", prompt_hash=generate_prompt_hash("a cat")
)
image = os.path.join(self.tmp.name, "cat.png")
with open(image, "wb") as f:
f.write(b"x")
db.link_image_to_prompt(prompt_id=prompt_id, image_path=image, metadata={})
finally:
db.close_all() # Windows cannot delete the temp dir while it is open
return prompt_id
def _chdir(self, path):
previous = os.getcwd()
os.chdir(path)
self.addCleanup(os.chdir, previous)
class TestDatabaseChecks(DiagnosticsTestCase):
def test_missing_database_is_an_error(self):
result = self.diag.check_database()
self.assertEqual(result["status"], "error")
self.assertIn(self.db_path, result["message"])
def test_populated_database_is_ok(self):
self._populate()
result = self.diag.check_database()
self.assertEqual(
result, {"status": "ok", "prompt_count": 1, "has_images_table": True}
)
images = self.diag.check_images_table()
self.assertEqual(images["status"], "ok")
self.assertEqual(images["image_count"], 1)
self.assertEqual(images["recent_images"][0]["filename"], "cat.png")
self.assertEqual(images["recent_images"][0]["text"], "a cat")
def test_corrupt_database_is_an_error(self):
with open(self.db_path, "wb") as f:
f.write(b"not sqlite at all" * 10)
self.assertEqual(self.diag.check_database()["status"], "error")
self.assertEqual(self.diag.check_images_table()["status"], "error")
def test_database_without_images_table(self):
with closing(sqlite3.connect(self.db_path)) as conn, conn:
conn.execute("CREATE TABLE prompts (id INTEGER PRIMARY KEY, text TEXT)")
self.assertEqual(
self.diag.check_database(),
{"status": "ok", "prompt_count": 0, "has_images_table": False},
)
images = self.diag.check_images_table()
self.assertEqual(images["status"], "error")
self.assertIn("generated_images", images["message"])
class TestFileSystemCheck(DiagnosticsTestCase):
def test_writable_current_directory(self):
self._chdir(self.tmp.name)
result = self.diag.check_file_system()
self.assertEqual(result["status"], "ok")
self.assertTrue(result["can_write"])
self.assertTrue(same_path(result["current_dir"], self.tmp.name))
self.assertFalse(os.path.exists(os.path.join(self.tmp.name, "test_write.tmp")))
def test_unwritable_current_directory(self):
self._chdir(self.tmp.name)
with mock.patch("utils.diagnostics.open", create=True, side_effect=OSError):
result = self.diag.check_file_system()
self.assertEqual(result["status"], "ok")
self.assertFalse(result["can_write"])
def test_unexpected_errors_are_reported(self):
with mock.patch("utils.diagnostics.os.getcwd", side_effect=OSError("gone")):
result = self.diag.check_file_system()
self.assertEqual(result["status"], "error")
self.assertIn("gone", result["message"])
class TestComfyUIOutputCheck(DiagnosticsTestCase):
def setUp(self):
super().setUp()
self.output = os.path.join(self.tmp.name, "comfy_output")
os.makedirs(os.path.join(self.output, "sub"))
with open(os.path.join(self.output, "sub", "a.png"), "wb") as f:
f.write(b"x")
self.folder_paths = types.ModuleType("folder_paths")
self.folder_paths.get_output_directory = lambda: self.output
def _config(self, directories):
module = types.ModuleType("py.config")
module.GalleryConfig = types.SimpleNamespace(MONITORING_DIRECTORIES=directories)
return module
def test_comfyui_output_directory_is_found(self):
with mock.patch.dict(
sys.modules, {"folder_paths": self.folder_paths, "py.config": None}
):
result = self.diag.check_comfyui_output()
self.assertEqual(result["status"], "ok")
self.assertIsNone(result["message"])
self.assertEqual(len(result["output_dirs"]), 1)
self.assertTrue(same_path(result["output_dirs"][0], self.output))
def test_configured_directories_come_first_and_are_not_duplicated(self):
extra = os.path.join(self.tmp.name, "extra")
os.makedirs(extra)
config = self._config([extra, self.output, os.path.join(self.tmp.name, "nope")])
with mock.patch.dict(
sys.modules, {"folder_paths": self.folder_paths, "py.config": config}
):
result = self.diag.check_comfyui_output()
self.assertEqual(len(result["output_dirs"]), 2)
self.assertTrue(same_path(result["output_dirs"][0], extra))
self.assertTrue(same_path(result["output_dirs"][1], self.output))
def test_without_comfyui_falls_back_to_relative_paths(self):
workdir = os.path.join(self.tmp.name, "work")
os.makedirs(os.path.join(workdir, "output"))
self._chdir(workdir)
with mock.patch.dict(sys.modules, {"folder_paths": None, "py.config": None}):
result = self.diag.check_comfyui_output()
self.assertEqual(result["status"], "ok")
self.assertTrue(
same_path(result["output_dirs"][0], os.path.join(workdir, "output"))
)
def test_nothing_found_is_a_warning(self):
empty = os.path.join(self.tmp.name, "empty")
os.makedirs(empty)
self._chdir(empty)
with mock.patch.dict(sys.modules, {"folder_paths": None, "py.config": None}):
result = self.diag.check_comfyui_output()
self.assertEqual(result["status"], "warning")
self.assertEqual(result["message"], "No output directories found")
self.assertEqual(result["output_dirs"], [])
def test_scan_errors_do_not_fail_the_check(self):
with mock.patch.dict(
sys.modules, {"folder_paths": self.folder_paths, "py.config": None}
):
with mock.patch.object(
pathlib.Path, "rglob", side_effect=OSError("denied")
):
result = self.diag.check_comfyui_output()
self.assertEqual(result["status"], "ok")
class TestDependenciesAndFullRun(DiagnosticsTestCase):
def test_all_dependencies_present(self):
result = self.diag.check_dependencies()
self.assertEqual(result["status"], "ok")
self.assertIsNone(result["message"])
self.assertEqual(
result["dependencies"], {"watchdog": True, "PIL": True, "sqlite3": True}
)
def test_watchdog_without_a_version_attribute_is_still_present(self):
with mock.patch.dict(sys.modules, {"watchdog": types.ModuleType("watchdog")}):
result = self.diag.check_dependencies()
self.assertEqual(result["status"], "ok")
self.assertTrue(result["dependencies"]["watchdog"])
def test_missing_dependency_is_an_error(self):
with mock.patch.dict(sys.modules, {"watchdog": None}):
result = self.diag.check_dependencies()
self.assertEqual(result["status"], "error")
self.assertEqual(result["message"], "Missing dependencies")
self.assertFalse(result["dependencies"]["watchdog"])
def test_full_diagnostic_reports_every_category(self):
self._populate()
self._chdir(self.tmp.name)
with mock.patch.dict(sys.modules, {"folder_paths": None, "py.config": None}):
results = self.diag.run_full_diagnostic()
self.assertEqual(
sorted(results),
[
"comfyui_output",
"database",
"dependencies",
"file_system",
"images_table",
],
)
self.assertEqual(results["database"]["status"], "ok")
self.assertEqual(results["comfyui_output"]["status"], "warning")
def test_run_diagnostics_uses_the_default_instance(self):
with mock.patch.object(diagnostics, "GalleryDiagnostics") as cls:
cls.return_value.run_full_diagnostic.return_value = {
"database": {"status": "ok"}
}
self.assertEqual(run_diagnostics(), {"database": {"status": "ok"}})
class TestCreateTestImageLink(DiagnosticsTestCase):
def test_links_a_fake_image_through_the_database(self):
with mock.patch("database.operations.PromptDatabase") as db_cls:
db_cls.return_value.link_image_to_prompt.return_value = 7
result = self.diag.create_test_image_link(3)
self.assertEqual(result["status"], "ok")
self.assertEqual(result["image_id"], 7)
kwargs = db_cls.return_value.link_image_to_prompt.call_args.kwargs
self.assertEqual(kwargs["prompt_id"], "3")
self.assertEqual(kwargs["image_path"], "/fake/test/image.png")
self.assertIn("file_info", kwargs["metadata"])
def test_uses_the_given_image_path(self):
with mock.patch("database.operations.PromptDatabase") as db_cls:
self.diag.create_test_image_link(3, test_image_path="custom.png")
kwargs = db_cls.return_value.link_image_to_prompt.call_args.kwargs
self.assertEqual(kwargs["image_path"], "custom.png")
def test_failures_are_reported(self):
with mock.patch(
"database.operations.PromptDatabase", side_effect=RuntimeError("no db")
):
result = self.diag.create_test_image_link(3)
self.assertEqual(result["status"], "error")
self.assertIn("no db", result["message"])
if __name__ == "__main__":
unittest.main()
+267
View File
@@ -0,0 +1,267 @@
"""PromptManagerLogger: memory buffer, file rotation, level changes, truncation."""
import logging
import os
import pathlib
import sys
import tempfile
import threading
import unittest
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils import logging_config
from utils.logging_config import (
MemoryBufferHandler,
PromptManagerLogger,
get_logger,
get_logger_manager,
)
class IsolatedLoggerTestCase(unittest.TestCase):
"""A logger manager whose log directory is a temporary folder."""
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
package = os.path.join(self.tmp.name, "pkg")
os.makedirs(os.path.join(package, "utils"))
self.log_dir = os.path.join(package, "logs")
fake_file = os.path.join(package, "utils", "logging_config.py")
with mock.patch.object(logging_config, "__file__", fake_file):
self.manager = object.__new__(PromptManagerLogger) # bypass singleton
self.manager.__init__()
self.manager.update_config({"console_logging": False})
# Drop the "Updated logging configuration" entry.
self.manager._log_buffer.clear()
# Other tests leave daemon threads (image monitor, watchdog) logging
# under "prompt_manager.*"; only records from this test's thread may
# reach the buffer or the rotating file (a stray line after a truncate
# would roll the file over and resurrect the backups). update_config
# rebuilds the handlers, so the filter is re-applied after every call.
self._filter_to_this_thread()
original_update = self.manager.update_config
def update_and_refilter(new_config):
original_update(new_config)
self._filter_to_this_thread()
self.manager.update_config = update_and_refilter
self.logger = logging.getLogger("prompt_manager.coverage_test")
self.addCleanup(self._restore_global_logging)
def _filter_to_this_thread(self):
this_thread = threading.get_ident()
for handler in self.manager.logger.handlers:
handler.addFilter(lambda record: record.thread == this_thread)
def _restore_global_logging(self):
for handler in logging.getLogger("prompt_manager").handlers[:]:
handler.close()
get_logger_manager()._setup_loggers()
self.tmp.cleanup()
def _rotate(self, max_file_size=400, backup_count=2, messages=30):
self.manager.update_config(
{"max_file_size": max_file_size, "backup_count": backup_count}
)
for i in range(messages):
self.logger.info("message %03d %s", i, "x" * 80)
class TestSingleton(unittest.TestCase):
def test_manager_and_loggers_are_shared(self):
self.assertIs(PromptManagerLogger(), PromptManagerLogger())
self.assertIs(get_logger_manager(), get_logger_manager())
self.assertIs(
get_logger("prompt_manager.x"), logging.getLogger("prompt_manager.x")
)
self.assertEqual(get_logger().name, "prompt_manager")
def test_init_is_a_noop_on_the_existing_instance(self):
manager = get_logger_manager()
before = manager.config
manager.__init__()
self.assertIs(manager.config, before)
class TestMemoryBuffer(IsolatedLoggerTestCase):
def test_records_are_buffered_most_recent_first(self):
self.logger.info("first")
self.logger.warning("second")
self.logger.error("third")
recent = self.manager.get_recent_logs(limit=2)
self.assertEqual([r["message"] for r in recent], ["third", "second"])
entry = recent[0]
self.assertEqual(entry["level"], "ERROR")
self.assertEqual(entry["logger"], self.logger.name)
self.assertIn("third", entry["formatted"])
for key in ("timestamp", "module", "filename", "lineno", "thread", "process"):
self.assertIn(key, entry)
def test_level_filter(self):
self.logger.info("info")
self.logger.warning("warn")
self.logger.error("err")
self.assertEqual(
[r["message"] for r in self.manager.get_recent_logs(level="warning")],
["err", "warn"],
)
self.assertEqual(len(self.manager.get_recent_logs(level="NOPE")), 3)
def test_buffer_is_capped_at_the_configured_size(self):
for i in range(1005):
self.logger.info("entry %d", i)
stats = self.manager.get_log_stats()
self.assertEqual(stats["buffer_count"], 1000)
self.assertEqual(stats["level_counts"], {"INFO": 1000})
self.assertEqual(stats["current_level"], "INFO")
self.assertEqual(
os.path.normcase(os.path.realpath(stats["log_directory"])),
os.path.normcase(os.path.realpath(self.log_dir)),
)
self.assertEqual(
self.manager.get_recent_logs(limit=1)[0]["message"], "entry 1004"
)
def test_handler_swallows_buffer_errors(self):
manager = mock.Mock()
manager.add_to_buffer.side_effect = RuntimeError("boom")
handler = MemoryBufferHandler(manager)
record = logging.LogRecord("x", logging.INFO, __file__, 1, "msg", (), None)
handler.emit(record) # must not raise
manager.add_to_buffer.assert_called_once()
class TestFilesAndRotation(IsolatedLoggerTestCase):
def test_rotation_creates_backups_and_lists_them_newest_first(self):
self._rotate()
files = self.manager.get_log_files()
names = sorted(f["filename"] for f in files)
self.assertEqual(
names,
["prompt_manager.log", "prompt_manager.log.1", "prompt_manager.log.2"],
)
main = [f for f in files if f["is_main"]]
self.assertEqual(len(main), 1)
self.assertTrue(all(f["size"] > 0 for f in files if not f["is_main"]))
modified = [f["modified"] for f in files]
self.assertEqual(modified, sorted(modified, reverse=True))
self.assertEqual(self.manager.get_log_stats()["log_files_count"], 3)
def test_read_log_file_tail_and_whole(self):
self._rotate()
tail = self.manager.read_log_file("prompt_manager.log.1", lines=1)
self.assertEqual(len(tail), 1)
whole = self.manager.read_log_file("prompt_manager.log.1", lines=0)
self.assertGreater(len(whole), 1)
self.assertTrue(whole[-1].endswith(tail[0]))
def test_read_log_file_rejects_missing_and_outside_paths(self):
with self.assertRaises(FileNotFoundError):
self.manager.read_log_file("missing.log")
outside = os.path.join(self.tmp.name, "pkg", "outside.log")
with open(outside, "w", encoding="utf-8") as f:
f.write("secret\n")
with self.assertRaises(ValueError):
self.manager.read_log_file(os.path.join("..", "outside.log"))
def test_read_log_file_reraises_read_errors(self):
self.logger.info("something")
with mock.patch("builtins.open", side_effect=OSError("disk")):
with self.assertRaises(OSError):
self.manager.read_log_file("prompt_manager.log")
def test_unreadable_files_are_skipped_in_the_listing(self):
self.logger.info("something")
original_stat = pathlib.Path.stat
def stat_failing_for_logs(path, *args, **kwargs):
if path.suffix == ".log":
raise OSError("denied")
return original_stat(path, *args, **kwargs)
with mock.patch.object(pathlib.Path, "stat", autospec=True) as stat:
stat.side_effect = stat_failing_for_logs
self.assertEqual(self.manager.get_log_files(), [])
def test_truncate_does_not_roll_the_log_over_while_deleting_backups(self):
# With a file limit smaller than one record, any log line written
# while the backups are being deleted rolls the main log over and
# resurrects prompt_manager.log.1. The handler's own stream must also
# be reset, or its stale offset reports the old size, rolls over on the
# next record anyway and pads the file with zeros on POSIX.
self._rotate(max_file_size=400, backup_count=2, messages=10)
result = self.manager.truncate_logs()
self.assertEqual(result["errors"], [])
self.logger.info("after truncate")
leftovers = sorted(p.name for p in pathlib.Path(self.log_dir).iterdir())
contents = {
name: pathlib.Path(self.log_dir, name).read_bytes() for name in leftovers
}
self.assertEqual(leftovers, ["prompt_manager.log"], contents)
main_log = pathlib.Path(self.log_dir, "prompt_manager.log").read_bytes()
self.assertNotIn(b"\x00", main_log)
self.assertIn(b"after truncate", main_log)
def test_truncate_clears_main_log_backups_and_buffer(self):
self._rotate()
result = self.manager.truncate_logs()
self.assertEqual(result["errors"], [])
self.assertEqual(
sorted(result["truncated"]),
["prompt_manager.log", "prompt_manager.log.1", "prompt_manager.log.2"],
)
self.assertFalse(
os.path.exists(os.path.join(self.log_dir, "prompt_manager.log.1"))
)
# Only the "Truncated log file" lines written after clearing remain
self.assertEqual(
[
r["message"]
for r in self.manager.get_recent_logs()
if "message 0" in r["message"]
],
[],
)
def test_truncate_reports_files_it_cannot_remove(self):
self._rotate()
with mock.patch.object(pathlib.Path, "unlink", side_effect=OSError("busy")):
result = self.manager.truncate_logs()
self.assertEqual(len(result["errors"]), 2)
self.assertEqual(result["truncated"], ["prompt_manager.log"])
class TestConfiguration(IsolatedLoggerTestCase):
def test_level_change_applies_to_every_prompt_manager_logger(self):
self.manager.update_config({"level": "DEBUG"})
self.assertEqual(logging.getLogger("prompt_manager").level, logging.DEBUG)
self.assertEqual(
logging.getLogger("prompt_manager.database").level, logging.DEBUG
)
self.assertEqual(self.manager.get_config()["level"], "DEBUG")
self.logger.debug("now visible")
self.assertEqual(self.manager.get_recent_logs(limit=1)[0]["level"], "DEBUG")
self.manager.update_config({"level": "bogus"}) # falls back to INFO
self.assertEqual(logging.getLogger("prompt_manager").level, logging.INFO)
def test_get_config_returns_a_copy(self):
config = self.manager.get_config()
config["level"] = "CRITICAL"
self.assertEqual(self.manager.get_config()["level"], "INFO")
def test_handler_flags_rebuild_the_handlers(self):
root = logging.getLogger("prompt_manager")
self.assertFalse(any(type(h) is logging.StreamHandler for h in root.handlers))
self.assertTrue(any(isinstance(h, logging.FileHandler) for h in root.handlers))
self.manager.update_config({"console_logging": True, "file_logging": False})
self.assertTrue(any(type(h) is logging.StreamHandler for h in root.handlers))
self.assertFalse(any(isinstance(h, logging.FileHandler) for h in root.handlers))
self.assertTrue(any(isinstance(h, MemoryBufferHandler) for h in root.handlers))
self.manager.update_config({"console_logging": False})
if __name__ == "__main__":
unittest.main()
+61
View File
@@ -0,0 +1,61 @@
"""Tests for utils/parallel.py: bounded thread-pool map used by long jobs."""
import os
import sys
import threading
import time
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.parallel import map_parallel # noqa: E402
class TestMapParallel(unittest.TestCase):
def test_results_keep_input_order(self):
def slow_square(n):
time.sleep(0.01 * (5 - n)) # later items finish first
return n * n
self.assertEqual(
map_parallel(slow_square, [1, 2, 3, 4], workers=4), [1, 4, 9, 16]
)
def test_one_worker_runs_on_the_calling_thread(self):
names = []
map_parallel(
lambda _: names.append(threading.current_thread().name), [1, 2, 3], 1
)
self.assertEqual(set(names), {threading.current_thread().name})
def test_several_workers_use_several_threads(self):
names = set()
lock = threading.Lock()
def record(_):
with lock:
names.add(threading.current_thread().name)
time.sleep(0.05)
map_parallel(record, range(8), workers=4)
self.assertGreaterEqual(len(names), 2)
self.assertNotIn(threading.current_thread().name, names)
def test_worker_count_is_capped_by_item_count_and_floor_one(self):
self.assertEqual(map_parallel(lambda n: n + 1, [1], workers=16), [2])
self.assertEqual(map_parallel(lambda n: n + 1, [1, 2], workers=0), [2, 3])
self.assertEqual(map_parallel(lambda n: n, [], workers=4), [])
def test_exception_propagates_to_the_caller(self):
def boom(n):
if n == 2:
raise ValueError("two")
return n
with self.assertRaises(ValueError):
map_parallel(boom, [1, 2, 3], workers=3)
if __name__ == "__main__":
unittest.main()
+165
View File
@@ -0,0 +1,165 @@
"""Tests for utils/video_metadata.py: prompt metadata from video containers.
ComfyUI's video savers store ``{"prompt": ..., "workflow": ...}`` as JSON in
the container's ``comment`` tag. ffprobe reads it; when ffprobe is missing
videos are skipped without noise.
"""
import json
import os
import shutil
import subprocess
import sys
import tempfile
import unittest
from unittest import mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils import video_metadata # noqa: E402
from utils.video_metadata import ( # noqa: E402
VIDEO_EXTENSIONS,
is_video_path,
metadata_from_tags,
read_video_metadata,
)
HAVE_FFMPEG = bool(shutil.which("ffmpeg") and shutil.which("ffprobe"))
GRAPH = json.dumps({"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}})
def write_mp4(path, comment):
"""Write a tiny black video carrying ``comment`` in its container tags."""
subprocess.run(
[
"ffmpeg",
"-v",
"quiet",
"-f",
"lavfi",
"-i",
"color=c=black:s=16x16:d=0.2",
"-metadata",
f"comment={comment}",
"-y",
str(path),
],
check=True,
timeout=60,
)
class TestMetadataFromTags(unittest.TestCase):
def test_prompt_string_and_workflow_dict_become_text_chunks(self):
tags = {"comment": json.dumps({"prompt": GRAPH, "workflow": {"nodes": []}})}
meta = metadata_from_tags(tags)
self.assertEqual(meta["prompt"], GRAPH)
self.assertEqual(json.loads(meta["workflow"]), {"nodes": []})
def test_missing_or_non_json_comment_is_empty(self):
self.assertEqual(metadata_from_tags({}), {})
self.assertEqual(metadata_from_tags({"comment": "just a note"}), {})
self.assertEqual(metadata_from_tags({"comment": "[1, 2]"}), {})
self.assertEqual(metadata_from_tags({"comment": json.dumps({"other": 1})}), {})
self.assertEqual(metadata_from_tags({"comment": 7}), {})
def test_description_tag_is_a_fallback(self):
tags = {"description": json.dumps({"prompt": GRAPH})}
self.assertEqual(metadata_from_tags(tags), {"prompt": GRAPH})
class TestIsVideoPath(unittest.TestCase):
def test_extensions(self):
for ext in VIDEO_EXTENSIONS:
self.assertTrue(is_video_path(f"clip{ext.upper()}"), ext)
self.assertFalse(is_video_path("image.png"))
self.assertFalse(is_video_path("noext"))
class TestReadVideoMetadataWithoutFfprobe(unittest.TestCase):
def setUp(self):
video_metadata.ffprobe_path.cache_clear()
self.addCleanup(video_metadata.ffprobe_path.cache_clear)
def test_missing_ffprobe_returns_empty_without_spawning(self):
with mock.patch.object(video_metadata.shutil, "which", return_value=None):
with mock.patch.object(video_metadata.subprocess, "run") as run:
self.assertEqual(read_video_metadata("clip.mp4"), {})
run.assert_not_called()
def test_missing_ffprobe_is_mentioned_once(self):
video_metadata.reset_missing_notice()
with mock.patch.object(video_metadata.shutil, "which", return_value=None):
with self.assertLogs("prompt_manager.video", level="INFO") as cm:
read_video_metadata("a.mp4")
read_video_metadata("b.mp4")
self.assertEqual(len(cm.output), 1)
self.assertIn("ffprobe", cm.output[0])
def _run_returning(self, **attrs):
return mock.patch.object(
video_metadata.subprocess,
"run",
return_value=mock.Mock(**{"returncode": 0, "stdout": "{}", **attrs}),
)
def test_ffprobe_failure_timeout_and_garbage_are_empty(self):
with mock.patch.object(
video_metadata.shutil, "which", return_value="/bin/ffprobe"
):
with self._run_returning(returncode=1):
self.assertEqual(read_video_metadata("clip.mp4"), {})
with self._run_returning(stdout="not json"):
self.assertEqual(read_video_metadata("clip.mp4"), {})
with mock.patch.object(
video_metadata.subprocess,
"run",
side_effect=subprocess.TimeoutExpired("ffprobe", 1),
):
self.assertEqual(read_video_metadata("clip.mp4"), {})
with mock.patch.object(
video_metadata.subprocess, "run", side_effect=OSError
):
self.assertEqual(read_video_metadata("clip.mp4"), {})
def test_ffprobe_output_is_parsed(self):
payload = {"format": {"tags": {"comment": json.dumps({"prompt": GRAPH})}}}
with mock.patch.object(
video_metadata.shutil, "which", return_value="/bin/ffprobe"
):
with self._run_returning(stdout=json.dumps(payload)) as run:
self.assertEqual(read_video_metadata("clip.mp4"), {"prompt": GRAPH})
cmd = run.call_args.args[0]
self.assertEqual(cmd[0], "/bin/ffprobe")
self.assertEqual(cmd[-1], "clip.mp4")
self.assertIn("format_tags", cmd)
@unittest.skipUnless(HAVE_FFMPEG, "ffmpeg/ffprobe not installed")
class TestReadVideoMetadataWithRealFfprobe(unittest.TestCase):
def setUp(self):
video_metadata.ffprobe_path.cache_clear()
self.tmpdir = tempfile.mkdtemp()
self.addCleanup(shutil.rmtree, self.tmpdir, True)
def test_round_trip_through_a_real_mp4(self):
path = os.path.join(self.tmpdir, "clip.mp4")
write_mp4(path, json.dumps({"prompt": GRAPH, "workflow": {"v": 1}}))
meta = read_video_metadata(path)
self.assertEqual(meta["prompt"], GRAPH)
self.assertEqual(json.loads(meta["workflow"]), {"v": 1})
def test_video_without_a_comment_is_empty(self):
path = os.path.join(self.tmpdir, "plain.mp4")
write_mp4(path, "")
self.assertEqual(read_video_metadata(path), {})
if __name__ == "__main__":
unittest.main()
+20 -36
View File
@@ -16,10 +16,10 @@ from utils.validators import (
validate_rating,
validate_tags,
validate_category,
validate_workflow_name,
sanitize_input,
parse_tags_string,
validate_result_timeout,
validate_worker_threads,
)
@@ -77,41 +77,6 @@ class TestValidateCategory(unittest.TestCase):
validate_category("bad\ncategory")
class TestValidateWorkflowName(unittest.TestCase):
"""Test validate_workflow_name function."""
def test_none_is_valid(self):
self.assertTrue(validate_workflow_name(None))
def test_empty_string_is_valid(self):
self.assertTrue(validate_workflow_name(""))
def test_whitespace_only_is_valid(self):
self.assertTrue(validate_workflow_name(" "))
def test_normal_name(self):
self.assertTrue(validate_workflow_name("My Workflow"))
def test_name_with_special_chars(self):
self.assertTrue(validate_workflow_name("workflow-v2_final (copy)"))
def test_max_length_boundary(self):
self.assertTrue(validate_workflow_name("x" * 200))
def test_exceeds_max_length(self):
with self.assertRaises(ValueError) as ctx:
validate_workflow_name("x" * 201)
self.assertIn("200", str(ctx.exception))
def test_non_string_raises(self):
with self.assertRaises(ValueError):
validate_workflow_name(42)
def test_dict_raises(self):
with self.assertRaises(ValueError):
validate_workflow_name({"name": "workflow"})
class TestSanitizeInput(unittest.TestCase):
"""Test sanitize_input function."""
@@ -329,5 +294,24 @@ class TestValidateResultTimeout(unittest.TestCase):
validate_result_timeout(True)
class TestValidateWorkerThreads(unittest.TestCase):
"""validate_worker_threads(value, max_workers) -> int in [1, max_workers]."""
def test_accepts_ints_and_digit_strings_within_range(self):
self.assertEqual(validate_worker_threads(1, 8), 1)
self.assertEqual(validate_worker_threads(8, 8), 8)
self.assertEqual(validate_worker_threads("4", 8), 4)
def test_rejects_out_of_range_bool_and_junk(self):
for value in (0, -1, 9, True, False, "x", None, 2.5, [2]):
with self.assertRaises(ValueError, msg=repr(value)):
validate_worker_threads(value, 8)
def test_max_of_one_only_accepts_one(self):
self.assertEqual(validate_worker_threads(1, 1), 1)
with self.assertRaises(ValueError):
validate_worker_threads(2, 1)
if __name__ == "__main__":
unittest.main()

Some files were not shown because too many files have changed in this diff Show More