Compare commits

..
Author SHA1 Message Date
Vito Sansevero 2f19908d9a fix(test): remove /dev/shm reset scratch left by killed browser servers (#384)
A server killed outright never runs its cleanup, and tmpfs keeps the
scratch folder in memory until reboot: the Python suite's server tests
left one per killed server. Scratch folders now carry the owner's pid,
and each new one first removes those whose process is gone (POSIX only).
2026-10-27 09:57:38 -07:00
Vito Sansevero 954fee739b fix(test): drop the gallery caches only after the reset's rows are published (#384)
The seeders drop the Images listing at their end. Run on the scratch copy,
that drop came before the publish, so a background walk in between could
cache per-file values built from the live database's old rows: CI's
inspector tag chip found 0 of 0 images. The proxy now records the drop
and the seeder makes it after the publish, as the old reset did.
2026-10-27 07:47:48 -07:00
Vito Sansevero d7893b6e8e test(browser): keep the scratch on disk when /dev/shm is short, delete in place after cleanup (#384)
The largest scratch copy is about 8 MB; Docker's default /dev/shm is 64 MB
and shared with the browsers. A reset that empties a folder after the
server's cleanup closed its trash deletes in place instead of starting a
thread nobody stops.
2026-10-27 07:21:10 -07:00
Vito Sansevero f707a1abe5 test(browser): keep the reset scratch database in /dev/shm when there is one (#384)
Throttled runs still waited on the runner's disk while creating and
rewriting the scratch database, whose pages were under writeback.
2026-10-27 07:14:49 -07:00
Vito Sansevero c160b186d7 fix(test): recheck the live data version under the publish's write lock (#384)
A writer whose transaction was open while the seed ran committed before
BEGIN IMMEDIATE returned and its rows were overwritten. The publish now
rechecks data_version under the lock, and restores foreign_keys and
detaches the scratch database on every error path.
2026-10-27 07:06:33 -07:00
Vito Sansevero 2fe9686ffc test(browser): empty media folders by a rename, delete them in the background (#384)
With the database commits gone, throttled resets still stalled in rmtree:
unlinking the last fixture's files waits for their pending writeback. Each
reseed now moves the old output and input folders into <root>/reset-trash
and makes empty ones; a background thread deletes the trash, and the
server's cleanup stops it and removes the rest before the temp root goes.
2026-10-27 07:04:57 -07:00
Vito Sansevero 116a66f755 test(browser): seed resets on a scratch copy and publish them in one commit (#384)
The concept and plain resets made hundreds of committed transactions on the
live database; on a slow CI disk they queued behind WAL checkpoint syncs and
took 30+ s. Each reset now copies the live database into a scratch database
(backup API, no syncs), runs the unchanged seeder against it through an API
proxy, and replaces the rows of the tables the seed changed in one live
transaction. A commit by another live connection meanwhile fails the reset.

tests/browser/reset_equivalence.py proves every reset kind leaves the same
rows, files and response as the old path (PM_TEST_SEED_ON_LIVE=1).
2026-10-27 07:02:50 -07:00
Vito f35e010f79 feat(system): log cleanup rule, reviewed rotated-log cleanup and diagnostic report (#329, PR 1) (#367)
* fix(logs): delete and download only rotated copies under one rule

Rule R (#329): a file is a rotated copy only when it is a direct entry of
the log folder named prompt_manager.log.<n> (n >= 1, no leading zero,
compared after normcase) and a regular file that is not a symlink or a
junction. Everything else is never deleted.

- Clear (POST /logs/truncate) empties the active log and deletes only
  rotated copies; a .bak, a .1.gz, any other file and every link are kept.
  It runs off the event loop under one files lock and the file handler's
  lock, so no rollover renames a file mid-way, and logs only afterwards.
- delete_rotated(expected) deletes exactly the reviewed copies, or nothing
  when one of them changed (size, mtime_ns, inode, missing, now a link).
- The download serves only a listed regular, non-link file, opened without
  following links and read off the loop; a link or folder is a 404.
- /logs/files rows gain kind; files and stats read links with lstat and
  stats no longer count a link's target.
- The config update runs off the loop and swaps handlers under the files
  lock.

* feat(logs): review and apply rotated-log cleanup as a maintenance plan

POST /maintenance/preview {operation: "rotated_logs", options: {}} lists the
rotated copies Rule R allows, with size and time, and what it keeps (the
active log, links, other files); it writes nothing and answers the Plan
shape with basenames only, backup policy none and no disk key.

POST /maintenance/apply starts the "Rotated log copies" maintenance job
(admission read: runs beside a scan, refused beside another maintenance job
or an exclusive one). It deletes exactly the reviewed copies, or nothing and
fails plan-changed when any of them changed after the review; a copy that
cannot be deleted becomes an errors row. The token is consumed on the loop
turn the job starts. (#329)

* feat(diagnostics): add a shareable report without secrets or paths

GET /system/diagnostics builds the report on the server, off the event loop
and under a registry lease (409 busy while a restore runs): the project and
Viewer.js versions, Python and the platform, the four legacy checks (one
helper, legacy_checks, that GET /diagnostics also calls), public database
info, source counts, the log configuration and statistics, the last health
summary and the config.json-only keys from an explicit list. The CivitAI key
is reported only as has_civitai_api_key; every path is relative to ComfyUI
or reads "outside ComfyUI". A scrub replaces anything still absolute, the
key or the home folder and counts it in redacted (0 for a correct build).

The version parser moves unchanged to utils/version.py. (#329)

* fix(logs): never empty a renamed log or serve a swapped file; keep root anchors private

Review findings on #329 PR 1:
- Clear empties the active log through the handler only when the handler's
  open stream is the listed file (same inode); a log the user renamed while
  ComfyUI runs is no longer emptied.
- A download is refused when the opened file is not the one listed (inode),
  or has another name (a hard link).
- Windows reparse points that name another file (the name-surrogate bit:
  symlinks, junctions, WSL links) are links under Rule R; cloud placeholders
  stay regular files.
- The report's paths use the closest anchor that is not a filesystem root,
  and the scrub also catches the home folder without its root.

* fix(diagnostics): never anchor a report path above the home folder; serve only regular log files

Review findings on #329 PR 1:
- report_path never uses a filesystem root or a strict ancestor of the
  home folder (/home, C:\Users) as its anchor, so a path inside home needs
  an anchor at or below it and otherwise reads "outside ComfyUI"; no path
  starts with a user name. The scrub matches the home folder only as whole
  path parts (in full, without its root, and its name).
- A download serves only a regular file (a FIFO, socket or device is a
  404) and opens it non-blocking.
- One unreadable timestamp no longer breaks the log listing, stats,
  download, review or Clear.
- The report tests read a temporary log folder, never the process's own.

* fix(diagnostics): match the bare home folder name only inside paths

A home folder named like a word (/home/info) no longer removes the log
level INFO from the report; the name still goes when it is a part of a
path. Documents the accepted residual risk of another user's home. (#329)
2026-10-27 06:19:39 -07:00
Vito d868745194 fix(lora): cache LoraManager's layout and close the #355 follow-ups (#362) (#364)
* feat(lora): cache LoraManager's layout behind file signatures (#362, phase 1)

A new py/lora_layout_cache.py keeps one (install, layout) entry and reuses it
while the configured path, the plugin settings file's and the settings file's
signatures (st_mtime_ns, st_size), the install and every media root's real
path are unchanged and the entry is younger than 30 s. A load whose watched
file changed within 3 s (git's racy rule), could not be checked, or whose
settings were unreadable serves its request only. clear() bumps a generation
so a load that started earlier stores nothing. Hits take no lock.

Not wired yet; phase 2 puts it behind the image hot paths.

* fix(lora): treat an offline Windows share as unvalidated and make cache tests portable (#362)

Review findings on phase 1: Windows raises FileNotFoundError with winerror
53/67 for an offline share, which is a temporary failure, not a missing file.
Test times are whole milliseconds (NTFS keeps 100 ns steps) and a file system
that rounds them skips; cleanup no longer ignores errors; tests added for
ValueError, NotADirectoryError, a failed reload, resolve failures and the
watched-file set.

* perf(lora): serve the image hot paths from the layout cache (#362, phase 2)

_lora_image_dirs takes (install, layout) from lora_utils.cached_lora_manager_layout
instead of resolving the install and reading LoraManager's settings files on
every Library, inspect, reveal and serve request. serve_image builds the
allow-list in the executor instead of on the event loop. A finished import
(System job or legacy scan, both through replace_lora_imports) clears the
cache, so files linked under a folder LoraManager just created are served at
once. The integration-off check and the configured-path rule still run before
the cache.

* fix(lora): fingerprint every scanned sidecar, kept or skipped (#362, phase 3)

Discovery keeps every metadata file it saw (scanned_files, resolved and
sorted; scanned is now its length). The fingerprint hashes each one's path,
size, modification time and whether it was kept, so a skipped sidecar added,
removed or edited between the review and Start refuses Start with
plan-changed, as the review's counts promised.

* fix(lora): warn about refused default LoRA folders and log what failed (#362, phase 4)

ComfyUI's default LoRA folders already pass the folder guard; a refused one
that exists is now listed in the layout's refused_defaults, and the layout
cache warns once per newly refused folder. The cache writes one INFO line
when a stored layout differs from the last one logged. BAD_SIDECAR_ROOT and
the module docstring name a folder that contains the home folder. A relative
preview_url writes a DEBUG line before the suffix search; a non-OSError
failure while loading trigger words logs its traceback. Review follow-ups
harden two serve tests.

* fix(lora): skip a sidecar that is not UTF-8 instead of failing the load (#362)

Review follow-up on phase 4: read_lora_metadata caught JSONDecodeError but not
UnicodeDecodeError, so one non-UTF-8 sidecar failed the whole trigger-word
load (now logged with a traceback every retry) and the import review. It is
skipped like any unreadable sidecar. The relative preview_url DEBUG line uses
lazy formatting.

* docs(testing): the LoRA linked-only gate, layout cache and fingerprint scope (#362, phase 5)

TESTING.md's LoRA reference import contract now states what #358 and #362
shipped: the linked-only gate and _media_roots, linked_only media, the cover
link order, the guard on ComfyUI's default LoRA folders and refused_defaults,
the folder-unreadable code, the trigger cache's catch and 60 s retry, the
fingerprint over every scanned sidecar, and the layout cache's key,
signatures, racy window, maximum age, root re-check and clears. The targeted
command includes tests.test_lora_layout_cache and tests.test_api_images.

* fix(lora): a cached layout needs the same home folder; give the rebuild poll its time (#362)

Final review follow-ups: the folder guard compares against the home folder,
so a cache hit now also needs Path.home() unchanged. TESTING.md says the
configured-path rule runs on the cache's result and that a non-UTF-8
sidecar is skipped.

CI: 'Rebuild all is explicit...' (webkit) polled the job for the default
5 s while it regenerated 24 previews; the failure screenshot shows 12 of 24
done at the deadline on the busy runner (0.17 s locally). The poll now
allows 30 s for that real work.
2026-10-27 05:00:55 -07:00
Vito 1bb3678f32 fix(threads): describe the teardown as built, warn only for a live observer, and keep an interrupted teardown's result (#377) (#378)
- _stop_background_work pre-fills the started parts with False and stores
  the result and removes its atexit entry in a finally:, so a Ctrl+C between
  steps no longer leaves the result unset and the callback registered.
- The "watchdog observer still running" warning moves into _stop_observer:
  only for an observer still alive after its join, naming the join's own
  timeout (not after observer.stop() raised, and not the full deadline).
- Tests: the plain-teardown tests share one package; two teardown tests run
  with a live preview worker held in a fake encode (ffmpeg killed, step gets
  min(CLOSE_WAIT_SECONDS, time left), thumbnails False while held and True
  when the kill releases it); an interrupt-between-steps test; two warning
  tests.
- TESTING.md and the ledger fixture comment: the refusal comes first, the
  image worker and the observer share one 5 s deadline, and what an
  interrupt between steps does.

Closes #377
2026-10-27 04:22:27 -07:00
Vito 2f67b98287 fix(tagging): save WD14 confidences only on the 0.05 grid every page shows (#369) (#373)
* fix(tagging): save WD14 confidences only on the 0.05 grid every page shows (#369)

A saved threshold off the 0.05 step (for example 0.37) blocked Generate and
Start in the workspace Tag image dialogs (native step mismatch) and made the
legacy sliders send 0.35. The server now refuses such a save with a message
that names the setting, a stored one falls back on load, and an accepted one
is stored exactly on the grid. System > Auto-tagging refuses the same values
before Preview tagging review or the batch Review, and hands thresholds on
snapped to the grid. Per-request thresholds keep their own checks.

* test(browser): restore the safety backup policy on system reset

The Recovery policy test in workspace-system-backups saves safety_backup
false at runtime, and /__test__/system/reset restored config.json but not
PromptManagerConfig.SAFETY_BACKUP. A later spec on the same server (Tag
output images review) then saw "No safety backup (the policy is off)". The
new #369 test shifted the CI shards so the two met on one worker.
2026-10-27 03:26:34 -07:00
Vito f3b9f61f11 fix(threads): stop background work at exit, on cleanup and on the tests' unload (#301) (#374)
* fix(threads): stop background work at exit, on cleanup and on the tests' unload (#301)

- __init__.py: _stop_background_work() detaches the monitor's callbacks,
  stops the monitor, closes the preview worker and shuts the prompt
  cleanup thread down against one 10 s deadline, once. It runs from
  atexit (ComfyUI never runs aiohttp's cleanup hooks), from on_cleanup
  where a host runs it, and from the tests' package unload; it removes
  its own atexit callback. Registered after the history hook, so at exit
  it runs before the history dispose.
- AutoThumbnailer.close/close_auto_thumbnails take a timeout and return
  whether the worker ended.
- Tests: teardown result, once-only, on_cleanup, atexit order, a failing
  step or a second Ctrl+C never skipping the tracker, one shared
  deadline, three repeated loads leaving no thread, and a child process
  that exits normally with the image worker alive.

* fix(threads): refuse new previews before the monitor stops; isolate the exit child (#301)

Review findings:
- The teardown makes the API refuse new previews before it detaches the
  monitor's callbacks: an image in hand that already read the preview
  queue no longer starts a preview worker during exit (Python 3.12+
  refuses that thread and logs a traceback). New
  refuse_auto_thumbnails().
- The teardown imports its bounds inside each step and falls back to the
  standard logger, so a failing import never skips the other steps.
- The exit child keeps its library and log in its own folder
  (PROMPT_MANAGER_DATA_DIR), reports watchdog's threads too, fails with
  the child's output on a hang, and rejects "Error while stopping".
- TESTING.md: "Background threads at exit".
2026-10-27 02:53:02 -07:00
Vito b1e7abd380 fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370) (#372)
* fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370)

- PromptTracker: the cleanup thread waits on a stop event instead of
  time.sleep; new shutdown(timeout=2) joins it and returns whether it ended.
- ImageGenerationHandler.shutdown(timeout=5) now joins the worker and
  returns a bool; new close() is final: schedule() refuses and a sealed
  run's history re-check is dropped, so no worker restarts afterwards.
- ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then
  stops and joins the watchdog observer within one deadline.
- Every test that builds a real tracker, handler or monitor ends with its
  final stop and an assertion, before its folder and databases go.
- #370: the wrapper/session history test drains before dispose(); the
  held-write bound test uses 30 s ceilings instead of tight wall clocks.

* test(history): drain the writer before asserting a bounded close it does not test (#370)

* fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302)

Review findings on #372:
- ImageGenerationHandler keeps the workers a stop signalled that had not
  ended; every later shutdown/close joins them again, so a second call no
  longer returns True while the first worker still runs.
- ImageMonitor.stop_monitoring keeps an observer or handler it left
  running and stops and joins it again on the next call.
- A closed handler takes no awaiting history path, and a re-check that
  raced close() is not kept due.
- Tests: second stops of a busy worker and a live observer, the closed
  handler's re-check paths, the handle test asserts the worker ended and
  stops before it asserts, and a no-new-thread check instead of set
  equality.

* test(browser): the System reset restores every settings value, safety_backup included

The frontend coverage shard of #372 failed in workspace-system-tagging:
the review read "No safety backup (the policy is off)". The backups spec
saves safety_backup=false, and /__test__/system/reset restored config.json
but not PromptManagerConfig.SAFETY_BACKUP (nor INFINITE_SCROLL), so a
later spec on the same worker server saw the policy off. The reset now
captures and restores every non-tagging value POST /settings can change.
2026-10-27 02:15:23 -07:00
Vito ae755d2ce0 feat(system): Auto-tagging screen, native in System (#328, PR 3) (#368)
* feat(system): configure auto-tagging natively

* test(system): compare Auto-tagging with the concept against the real backend
2026-10-27 00:59:39 -07:00
Vito 4032f4b6ab feat(system): one tagging configuration and Tag output images as a reviewed job (#328, PR 2) (#366)
* feat(autotag): model status without paths, and a busy answer that names the operation

* feat(autotag): model downloads run as Activity jobs

* fix(autotag): keep each download route's error shape and harden the fake download

* feat(autotag): one server-side tagging configuration

* fix(autotag): the legacy batch and review work again

* feat(autotag): tag output images as a reviewed Activity job

* test(autotag): the legacy batch, review and unload page workflows

* fix(autotag): server-free tagging defaults, cancel before load, and named legacy refusals

* fix(autotag): name only the workspace batch in a legacy refusal frame
2026-10-27 00:20:13 -07:00
Vito 88fa13689a feat(system): model status, a named model gate and downloads as Activity jobs (#328, PR 1) (#365)
* feat(autotag): model status without paths, and a busy answer that names the operation

* feat(autotag): model downloads run as Activity jobs

* fix(autotag): keep each download route's error shape and harden the fake download
2026-10-26 23:49:15 -07:00
Vito cfe47d3cb0 ci: support and test Python 3.12-3.14 only (#356) (#357)
* ci: test Python 3.11-3.14 and require Python 3.11 (#356)

Python 3.10 reaches end of life in October 2026 and no PromptManager
users run it; 3.14 is current and was untested. The test matrix runs
3.11-3.14 on Linux, Windows and macOS, bandit runs on 3.11, and
requires-python is >=3.11 so the package no longer claims a version CI
does not test. The packaging test drops its 3.10-only tomllib fallback.
The required checks follow the matrix (py3.10 -> py3.14).

* ci: narrow Python support to 3.12-3.14 (#356)

ComfyUI recommends 3.13, names 3.12 as the fallback and 3.14 as working
(its README); PromptManager supports exactly those. requires-python is
>=3.12, the test matrix and the required checks drop 3.11, and bandit
runs on 3.12.
2026-10-26 22:23:06 -07:00
114 changed files with 11378 additions and 636 deletions
+254 -15
View File
@@ -275,6 +275,77 @@ Start fixing as soon as a failed log shows the cause, but push the fix only afte
the run for that commit has finished: pull-request runs use `cancel-in-progress`,
so a push cancels the running jobs and can hide a second failure.
## Background threads at exit
PromptManager's own background threads stop through one package teardown,
`_stop_background_work()` in `__init__.py` (#301, #302). It runs once and
works against one 10 s deadline. Each step gets the smaller of its own bound
and the time left. The steps, in order:
1. Refuse new previews (`refuse_auto_thumbnails()`), then detach the
monitor's callbacks (Needs review recorder, preview queue, Images listing
listener).
2. `ImageMonitor.stop_monitoring` works against one 5 s deadline. Closing the
image handler joins the image worker; watchdog's `observer.stop()` (not
bounded here) and the observer join get what is left, which can be 0 s.
3. `close_auto_thumbnails(final=True)` stops the preview worker (30 s) and
kills ffmpeg.
4. `PromptTracker.shutdown` stops the prompt cleanup thread (2 s).
A failing step, or a second Ctrl+C, is logged with its traceback and does not
skip the later steps. A clean run logs one INFO line to PromptManager's log:
`Background work stopped in 0.4 s: monitor=ok thumbnails=ok tracker=ok`. If
any part is still running, one WARNING goes to the console:
`Background work not stopped in time (10.0 s): monitor`. A Ctrl+C that lands
between two steps (each step catches its own) is re-raised after the result is
stored, with the parts not reached False, and the `atexit` callback is
removed: the later steps are skipped, no summary line is logged, and a later
call returns that result and runs nothing.
ComfyUI never runs aiohttp's `on_cleanup` hooks, because it never calls
`runner.cleanup()`. So `atexit` is the live host's path. It runs on Ctrl+C
(`main.py` catches `KeyboardInterrupt` and returns) and on a normal exit.
The teardown is registered after the execution history hook, so at exit
(newest first) it runs before the history writer's 2 s dispose.
It also runs from `on_cleanup` on hosts that do run it (our test servers) and
from the tests' package unload (`cleanup_package`), and it removes its own
`atexit` callback. Nothing runs when the process is killed: SIGTERM, a
ComfyUI-Manager restart through `os.execv`, a closed Windows console, or
ComfyUI Desktop stopping its server. SQLite's WAL journal covers those.
Stop contract: `PromptTracker.shutdown(timeout)`,
`ImageGenerationHandler.shutdown(timeout)` (restartable),
`ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a
sealed run's re-check is dropped) and `ImageMonitor.stop_monitoring(timeout)`
each signal their thread and join it within the bound. They return True only
when it ended, and log a warning naming the part when it did not. Called on
their own thread, they only signal.
Test rules:
- Every test that constructs a real `PromptTracker`, `ImageGenerationHandler`
or `ImageMonitor` ends with its final stop and an assertion
(`tests/background_stops.py`: `close_handler`, `stop_monitor`,
`stop_tracker`). The stop comes before the folder and databases are removed.
`tearDown` runs before `addCleanup` callbacks, so a class whose `tearDown`
removes files stops the thread there.
- A test waits for a state through an `Event` set by the code under test, or
by a wrapper around it, with a 30 s ceiling. It never sleeps or polls for
that state.
- A test that asserts "still running" first waits for an `entered` event
inside the held call, then stops with `timeout=0.05`.
- A test that asserts a history writer's `dispose()`/`close()` is True right
after queuing work drains first (`drain_writer`), or holds the write itself
(#370).
- Patch the module attribute of the code under test (for example
`utils.prompt_tracker.time`, or the package's `_clock`), or patch with
`wraps=`. Never replace a `time` or `atexit` function for the whole
process.
- `tests/test_package_exit.py` runs a child process that imports the
package, leaves the image worker alive and returns. Its `atexit` reporter
checks the teardown's result. This is the interpreter path ComfyUI takes
after Ctrl+C, on all three OS.
## Performance budgets
A spec performance budget gets two tests:
@@ -1504,8 +1575,8 @@ targets of the longest legal path (1024 code points, even escaped as `\uXXXX`) f
Its fields show the server's defaults (`GET /prompt_manager/autotag/models`, the
D91 Suggestions count included) once that read lands; Start waits for it. A model
the status lists as not downloaded keeps Start `aria-disabled` and
`#tag-generation-status` says "<model> is not downloaded. Download it from the
legacy Auto-tag page." (the server's own sentence), so no job starts only to fail.
`#tag-generation-status` says "<model> is not downloaded. Download it in System ›
Auto-tagging." (the server's own sentence), so no job starts only to fail.
Start validates the fields as Generate does (native validity; nothing is sent).
The concept comparison (`images-concept.spec.mjs`) uses the concept page itself as
the oracle, with Add to collection and Remove links… removed and a copy of its Add
@@ -2587,7 +2658,7 @@ the connection's `total_changes` is 0 at close and the data reads back unchanged
`-shm` bytes are not compared: every WAL reader rewrites its read marks there, so no
reader can keep it byte-identical (`test_health_check_never_writes`).
Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D131 assigned):
Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–D114; D121–D134 assigned):
- **D92** (shell): no "Sample configuration" chip. The Save status is truthful: "No
unsaved changes.", "Unsaved changes.", "Saved in this browser." / "Saved to the server
@@ -2621,7 +2692,8 @@ Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–
output folder is used." when that is the case.
- **D94** (all sections): a section not delivered yet shows the hand-off panel "Your
current tools are ready" with a link to each legacy tool that still owns it (the
plan's legacy coverage table); retired section by section, gone after S12.
plan's legacy coverage table); retired section by section, gone after S12. Since S11
(#328) only Diagnostics & logs hands off (`SECTION_HANDOFFS`, to the admin page).
- **D95** (health report): the concept's `health-check` dialog titled "Health report",
with no sample note. It lists every check in the server's order: Integrity check,
Foreign keys & pair references, Canonical hashes, Pair uniqueness (added), Missing
@@ -2997,6 +3069,59 @@ Recorded deviations (plan D60–D75, D79–D82 renumbered D92–D107 and D111–
dialog on that scope. Pinned by the thumbnail dialog comparison (mapped) and its
mutation check (system-concept.spec) and `the thumbnail row counts outdated previews`
(workspace-system.spec).
- **D132** (Auto-tagging, S11, #328, OQ1/OQ5/OQ6/OQ10): "Tagging model" (`#tag-model`)
lists the four real engines in the concept's order, each `<label> · <size_label>` from
`GET /prompt_manager/autotag/models` ("WD14 SwinV2 · ~400 MB", "WD14 ViT · ~400 MB",
"GGUF JoyCaption · ~4.5 GB", "HuggingFace JoyCaption · ~16 GB"; GGUF takes the
concept's "Ollama" slot, as D61 does). Its description is the chosen engine's live
status instead of the concept's design note: "Loading model status…" (the select is
disabled until the read answers), "Model status could not be read.", "Not downloaded.
Download size <size>.", "Downloaded. Not loaded.", "Downloaded and loaded." or "Partly
downloaded. Download again to finish." (GGUF with one of its two files). The concept's
two simulated job buttons are "Download" ("Download again" for an engine on disk),
which starts the `model_download` Activity job and opens it, and "Unload", a direct
`POST /autotag/unload` with a toast ("<label> unloaded.") and a fresh status read,
enabled only while a model is loaded; a busy model's refusal reads "The tagging model
is busy: <operation>. Try again when it finishes.". "Preview tagging review" opens the
Tag image dialog for the image selected in Images with the section's current model and
values; it is `aria-disabled` without a selection (the D115 sentence "Select an image
in Images first.") or with an invalid value (the toast names the setting); a saved
engine the section does not know reads "The saved tagging model is not one of these
engines. Choose one and save preferences.". A part-downloaded GGUF offers "Download
again"; an Unload that found nothing loaded only reads the status again. At phone width
the concept's 120px select cap cuts every engine label, so `#tag-model` alone keeps its
natural width there and its row wraps the select under the text (`workspace.css`
`/* D132 */`); Existing tags stays as the concept draws it, cut at 390 as there. The
comparison maps `selectsFit` only for the engine select and measures each row's
`flexWrap` and whether its control sits under the text. Pinned by `auto-tagging matches
the concept` (both themes and sizes), `auto-tagging at 390 wraps only the engine row`,
`auto-tagging engine select fits at 320 wide` and the mutation check (it includes a
wrapped row) (system-tagging-concept.spec), and the model, Download, Unload and Preview
tests of workspace-system-tagging.spec.
- **D133** (Auto-tagging, S11): a "Suggestions per image" row (`#tag-suggestions`,
number 1–100, host key `autotag_suggestion_limit`) follows the two confidence rows:
the persisted count the Tag image dialog's Suggestions field (D91) starts from. Pinned
by the comparison's mapping (it inserts the row) and `defaults save through Save
preferences and the Tag image dialog uses them`.
- **D134** (Auto-tagging, S11, OQ2/OQ3/OQ7/OQ8): an `.operation` "Tag output images"
with "Review" follows `.row-actions`: the legacy prompt-tag batch. Review is
`aria-disabled` (the press explains) while the chosen engine is not downloaded or a
value it uses is invalid; otherwise it posts `/autotag/batch/preview` with the drafted
engine and policy (read-only) and opens "Review tag output images" in the concept's
record-changing review (its `maintenancePlan` shell): `h3` "Tag output images", the
`dl.meta-grid` Model, Existing tags, Images linked, Prompts, Already tagged, Prompts
to tag ("up to N"), Settings ("General <g> · Character <c>" or "Template of N
characters") and Backup, the upper-bound sentence in the `p.small`, the D55 backup box
and Cancel / "Start tagging", which starts the cancellable `autotag` job and opens it.
Replace after review, a model not on disk, a short disk, no linked image and nothing to
tag each show the rows with only Cancel and the reason ("Replacing tags needs a review
of each image. Use Preview tagging review, or choose Skip or Merge for a batch run.",
"Every prompt with an output image is already tagged.", "Not enough free disk space:
the safety backup needs about X MB and Y MB is free. Nothing is written until there is
room.") in the close-only review Integrations shares (`openCloseOnlyReview`). Pinned by
`tag output images review matches the concept`, its mutation check and `the
Cancel-only review matches the concept review shell` (system-tagging-concept.spec) and
the batch tests of workspace-system-tagging.spec.
- **D124** (inspector): with "Load metadata on selection" off, the Images panel and
the viewer's details show "Metadata loads when you ask for it." and a "Load metadata"
button instead of reading the file. The concept has no such state, so no concept
@@ -3371,7 +3496,9 @@ workflow (README changelog).
### LoRA reference import contract (#221, #249 System S10)
The LoRA import is staged (`py/lora_import.py`): discover the LoraManager
`*.metadata.json` files and fingerprint them (path, size, modification time); stage each
`*.metadata.json` files and fingerprint every scanned one, kept or skipped (path, size,
modification time and whether it was kept, #362; a skipped sidecar added, removed or edited
after the review refuses Start with `plan-changed`, `test_start_refused_when_a_skipped_sidecar_changed`); stage each
record (CivitAI example images go into the image cache first, best effort, counted in
`downloaded` / `download_errors`; each image is written to a temp file beside its cache
entry and moved in place, so a half-written image is never seen as cached); the files are
@@ -3437,31 +3564,143 @@ the kept, readable LoRAs and `skipped` the rest. Media per LoRA, without repeats
preview (`preview_url`, else LoraManager's `PREVIEW_EXTENSIONS` order), then every file in
`<example-images root>[/<library>]/<sha256>/` in natural order; a leading video yields to the
first still. A file counts only with a servable extension and inside the guarded media
roots. CivitAI downloads run only for a LoRA without local media. The folder guard keeps a
settings folder only when it exists and is neither a filesystem root (`/`, `C:\`, a UNC
share root) nor the home folder or one of its ancestors (realpath + `normcase`). An
unusable layout (no install, unreadable / non-JSON / non-object / oversized settings, no
`platformdirs`, a missing or refused centralized sidecar root) gives a review without a
token ("LoraManager settings: not usable: <reason>. Nothing was changed."), a legacy scan
error with that reason, the default LoRA roots for serving, and trigger words from the
default roots in alongside mode. `tests/run_suite.py` and `tests/browser/server.py` point
roots. Only an absolute `preview_url` is used; a relative one is skipped with a DEBUG line
naming the sidecar. A sidecar that is not valid UTF-8 JSON is skipped like any unreadable
one (#362). CivitAI downloads run only for a LoRA without local media; a record's
images are linked in reverse order, so its first image (the cover) is the newest link and
the card image. The folder guard keeps a folder from the settings, and each of ComfyUI's
default LoRA folders (`models/loras`, `extra_model_paths.yaml`, `folder_paths`), only when
it exists and is neither a filesystem root (`/`, `C:\`, a UNC share root) nor the home
folder or a folder that contains it (realpath + `normcase`); a refused default folder that
exists is listed in the layout's `refused_defaults` and logged as one warning each time the
refused set changes (`RefusedDefaultsTests`). An unusable layout (no install, unreadable /
non-JSON / non-object / oversized settings, no `platformdirs`, a missing or refused
centralized sidecar root, "the sidecar folder is a filesystem root, the home folder or a
folder that contains it") gives a review without a token ("LoraManager settings: not
usable: <reason>. Nothing was changed."), a legacy scan error with that reason, the default
LoRA roots for serving (no example-images root), and trigger words from the default roots
in alongside mode. A folder discovery cannot list answers the review with 500, code
`folder-unreadable` ("A LoRA folder could not be listed (...). Nothing was changed.").
`TriggerWordCache.load` never raises: an `OSError` is a warning, anything else is logged
with its traceback (#362), both keep the previous map, and a running node retries a failed
load only after 60 s (`TRIGGER_RETRY_SECONDS`). `tests/run_suite.py` and `tests/browser/server.py` point
`LORA_MANAGER_SETTINGS_DIR` at a disposable folder, so no test reads the developer's real
LoraManager settings; tests that call `load_layout` inject their own folder too.
**Linked-only gate and the layout cache (#362).** While the integration is on, the LoRA
allow-list (`ImagesMixin._lora_image_dirs`) is the layout's `media_roots` plus the CivitAI
download cache, and `_media_roots(outputs)` appends it after the output folders for Library
previews, the Library prompt detail, the per-prompt images page, `serve_output_image` for a
root past the outputs, `inspect_image` and `reveal_image`. Under such a root only a file the
database records as a linked image is served, inspected or revealed (`_linked_only_error`,
by its given and its real path); anything else answers 403. Library responses mark that
media `linked_only`: no preview file, no favorite, no Tag button. The allow-list's layout
comes from `py/lora_layout_cache.py` (`tests/test_lora_layout_cache.py`, which injects the
clocks, `stat` and `ntpath` and never sleeps): one entry keyed by the stripped configured
path, reused while the plugin settings file's and the settings file's `(st_mtime_ns,
st_size)` (or "missing") are unchanged, the install and every media root still resolve to
themselves (a folder swapped for a symlink reloads and the guard runs again), and the entry
is under `MAX_AGE_SECONDS` (30 s, monotonic clock) old. A load is stored only when no
watched file was modified within `RACY_SECONDS` (3 s) of the load (git's racy rule: a write
during the load, LoraManager's in-place write, FAT/HFS+ coarse times), every watched file
could be checked (any `stat` error but not-found, and an offline Windows share, is "cannot
validate"), the settings were readable, and no `clear()` ran meanwhile (generation
counter). Hits take no lock; concurrent misses each load. A finished import (System job or
legacy scan, both through `replace_lora_imports`) clears it. The integration-off check runs
before it, and the configured-path rule (a configured path that is not an install under
`custom_nodes` adds no folders) on its result; a hit also needs the same home folder, which
the guard compares against. `serve_image` builds the allow-list in the
executor. The cache writes one INFO line, "LoraManager layout changed: <source> settings,
<mode> storage, <n> media folders", when a stored layout differs from the last one logged.
The review, Start, the legacy scan and the trigger cache read the layout fresh. Image route
tests clear the per-process cache in `ImageRouteCoverageCase.setUp`; `write_settings` in
`TestLoraManagerMediaServing` back-dates its file 10 s so the first load is stored.
`tests/browser/workspace_integrations_fixture.py` (`/__test__/integrations/*`) makes a fake
LoraManager install and LoRA folder under the server's temp root, holds or fails the
import, saves a blocking pair and answers CivitAI downloads with a stub opener, so no
test reaches the network.
```sh
python3 -m unittest tests.test_lora_manager_layout tests.test_lora_import tests.test_api_lora \
tests.test_lora_database tests.test_lora_utils -v
python3 -m unittest tests.test_lora_layout_cache tests.test_lora_manager_layout \
tests.test_lora_import tests.test_api_lora tests.test_lora_database tests.test_lora_utils \
tests.test_api_images -v
node --test tests/js/workspace-system-integrations.test.js tests/js/prompt-bridge-adapter.test.js
npm run --prefix tools/vendor test:workflows -- workspace-system-integrations.spec.mjs \
system-integrations-concept.spec.mjs \
--project chromium --project touch --project firefox --project webkit
```
### Auto-tagging contract (#328, #249 System S11)
System › Auto-tagging (`web/js/workspace/system-tagging.mjs`, D132–D134) edits one
server-side tagging configuration, `autotag.defaults` in config.json: `model_type`
(`wd14-swinv2`, `wd14-vit`, `gguf`, `hf`), `existing` (`skip`, `merge`,
`replace_after_review`), `general_threshold` and `character_threshold` (0–1 in
steps of 0.05, #369),
`template` (non-blank, at most 4,000 characters and 16 KiB) and `suggestion_limit`
(1–100). They are host drafts (`autotag_model_type`, `autotag_existing`,
`autotag_general_threshold`, `autotag_character_threshold`, `autotag_template`,
`autotag_suggestion_limit`) saved by the one Save row through `POST
/prompt_manager/settings` with `expected_revision`; a bad value is refused with a 400
naming the setting by its row label ("General-tag confidence must be between 0 and
1."), nothing is written, the Save row shows that sentence (`invalid-setting`, set by
`settingRefusal` in system.mjs only for a save that carried `autotag_*` keys; every
other section's plain 400 keeps the generic status) and the control is highlighted.
The template's limit counts code points, as Python's `len` does. A hand-edited bad value
falls back to its default on load, one key at a time. Every page shows the thresholds in
steps of 0.05 (the System rows, the Tag image dialogs, the legacy sliders), so a saved
threshold off that grid is refused ("General-tag confidence must be a multiple of
0.05.", `validate_saved_threshold`; `draftProblem` holds back Preview tagging review and
the batch Review with the same sentence) and falls back on load; an accepted one is
stored exactly on the grid (#369). A per-request threshold keeps its own check: `/image-annotations/suggest` and
`/jobs/autotag` take 0–1; the legacy `/autotag/single` and `/autotag/start` only require a number. `GET /autotag/models` reports the
saved values (`wd14_general_threshold`, `wd14_character_threshold`, `suggestion_limit`,
`default_prompt`, `model_type`, `existing`), so the Tag image and Auto-tag selected
dialogs start from them; a preset (the bulk review, Preview tagging review) wins.
`GET /prompt_manager/autotag/models` reads the disk in the executor, never loads or
downloads a model and carries no absolute path: per engine `name`, `label`,
`size_label`, `description`, `size`, `downloaded`, `loaded` (GGUF also `model_exists`,
`mmproj_exists`), plus `model_loaded` and `loaded_model_type`. Opening the section
makes exactly one such read per mount (`entering the section makes no model call
besides status`).
One model gate (`py/model_gate.py`) covers download, unload, both batches, the legacy
single-image tag, Generate suggestions and Auto-tag selected. A refusal is `409
model-busy` with `operation`, the holder's label: the download job's title, "Tag output
images", "Tag output images (legacy page)", "Tagging one image (legacy page)",
"Unloading the tagging model", "Auto-tag selected images", "Generate suggestions", else
"Another tagging operation". `ApiError.operation` keeps it (only on `model-busy`).
- **Download**: `POST /prompt_manager/jobs/model-download {model_type}` starts the
`model_download` job (202), title "Tagging model download · <label>", admission
`read`, not cancellable, indeterminate (Activity's generic held view: no Cancel, a
`<progress>` without a value), holding the gate until it ends. A second download is
`409 busy`; the legacy SSE route attaches to the same job and keeps its frames.
- **Unload**: `POST /prompt_manager/autotag/unload` answers `{success, message}` or
`409 model-busy`.
- **Tag output images**: `POST /prompt_manager/autotag/batch/preview {model_type,
existing}` writes nothing and opens no file; it answers `images`, `prompts`,
`already_tagged`, `to_tag` (an upper bound), `model {type, label, size_label,
downloaded}`, `existing`, `review_required`, `backup {policy}` and `disk`. `POST
/prompt_manager/jobs/autotag {model_type, existing, options}` (WD14: both thresholds;
GGUF/HF: `instructions`, the template) starts the cancellable `autotag` job (202),
title "Tag output images", admission `records`, a safety backup by the policy first.
Refusals: `400 review-required` (Replace after review never runs as a batch), `409
model-unavailable` ("<label> is not downloaded. Download it in System ›
Auto-tagging."), `507 insufficient-space`, `409 busy`, `409 model-busy`. The section
toasts `model-busy` with its operation, the other three with the server's sentence,
`busy` with the running job (D57), anything else "<what> could not start. Try again.".
Outcomes: Tagged / Unchanged / Skipped / Errors (the `image-tags` descriptor, D111).
No browser test downloads or runs a model: the injected `TestModel` answers, its
`download_model` raises unless `/__test__/tagging/download {mode: hold | succeed |
fail}` installs a fake that waits on a fixture event, and every tagging reset releases
it. Every workspace-system-tagging test asserts no `/autotag/download/` request.
Changing the routes or `py/config.py` needs a ComfyUI restart; the section itself is
HTML, CSS and JS only.
### Browser cache ownership (#222)
The legacy gallery's Clear cache used to delete every localStorage key starting `pm_` or
+170 -31
View File
@@ -22,22 +22,21 @@ Features:
- Comprehensive logging and diagnostics
"""
import re
import atexit
import threading
import time
from pathlib import Path
def get_version():
"""Parse version from pyproject.toml"""
"""Parse version from pyproject.toml (``utils/version.py``)."""
from .utils.version import project_version
try:
pyproject_path = Path(__file__).parent / "pyproject.toml"
if pyproject_path.exists():
content = pyproject_path.read_text()
match = re.search(r'version\s*=\s*["\']([^"\']+)["\']', content)
if match:
return match.group(1)
root = Path(__file__).parent
except Exception:
pass
return "unknown"
return "unknown"
return project_version(root)
from .prompt_manager import ( # noqa: E402
@@ -72,6 +71,12 @@ NODE_DISPLAY_NAME_MAPPINGS = {
# are served by our own routes instead (see py/api).
WEB_DIRECTORY = "web/comfy"
# Parts the package teardown stops (``_stop_background_work``). Each stays None
# when its start-up block below fails, and the teardown skips it.
api = None
_global_prompt_tracker = None
_global_image_monitor = None
# Add API routes (same pattern as ComfyUI_Assets)
try:
# Set extension URI
@@ -89,8 +94,6 @@ try:
api.add_routes(routes)
# Temp files of preview writes a crash interrupted (#317), off the
# import path: the walk may cover a large thumbnails folder.
import threading
threading.Thread(
target=api.sweep_thumbnail_temps,
name="PromptManagerPreviewTempSweep",
@@ -141,29 +144,16 @@ try:
# The Needs review pass for ambiguous new outputs runs through the API's
# capture store (the one the scan and the Tag dialog use): a linked output
# is never queued and its pending items leave the queue. It stops before
# that store is closed at shutdown.
_review_api = globals().get("api")
if _review_api is not None:
_global_image_monitor.set_review_recorder(_review_api.record_review_file)
# is never queued and its pending items leave the queue. The package
# teardown detaches it before that store is closed at shutdown.
if api is not None:
_global_image_monitor.set_review_recorder(api.record_review_file)
# Every settled output's preview is queued for the API's background
# worker (#249): the workspace never asks to generate thumbnails.
_global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
_global_image_monitor.set_thumbnail_queue(api.queue_capture_thumbnail)
# Every settled output walks its Images listing again (#348): the
# API hops from the monitor's thread onto the server's loop.
_global_image_monitor.set_output_listener(_review_api.notify_output_settled)
_review_host = getattr(config.server_instance, "app", None)
if _review_host is not None and hasattr(_review_host, "on_cleanup"):
async def _stop_review_recording(
app, monitor=_global_image_monitor, api=_review_api
):
monitor.set_review_recorder(None)
monitor.set_thumbnail_queue(None)
monitor.set_output_listener(None)
api.close_auto_thumbnails(final=True)
_review_host.on_cleanup.insert(0, _stop_review_recording)
_global_image_monitor.set_output_listener(api.notify_output_settled)
# Start monitoring
_global_image_monitor.start_monitoring()
@@ -233,6 +223,155 @@ except Exception:
"Failed to register execution history"
)
# ── Package teardown (#301, #302) ────────────────────────────────────────────
# ComfyUI never runs aiohttp's cleanup hooks (it never calls
# ``runner.cleanup()``), so ``atexit`` is the live host's path: Ctrl+C, and a
# normal exit. It is registered after the history hook, so at exit (newest
# first) it runs before the history dispose. Hosts that do run ``on_cleanup``
# (our test servers) run it there, and the tests' package unload calls it.
TEARDOWN_SECONDS = 10.0
_clock = time.monotonic
_background_work_lock = threading.Lock()
_background_work_result = None
def _teardown_step(name, stop, logger):
"""Run one teardown step; a failure (even Ctrl+C) is logged and is False."""
try:
return bool(stop())
except (Exception, KeyboardInterrupt):
logger.error(f"Background work: stopping {name} failed", exc_info=True)
return False
def _teardown_logger():
"""PromptManager's init logger; the standard one if that cannot load."""
try:
from .utils.logging_config import get_logger
return get_logger("prompt_manager.init")
except Exception:
import logging
return logging.getLogger("prompt_manager.init")
def _stop_background_work():
"""Stop the image monitor, the preview worker and the prompt cleanup
thread against one TEARDOWN_SECONDS deadline, once.
Returns ``{"monitor", "thumbnails", "tracker"}`` booleans (True: it
stopped), leaving out parts that never started. Later calls return a
copy of the first result. A Ctrl+C inside a step is logged and the
step is False; one between steps is re-raised after the result (parts
not reached are False) is stored and the ``atexit`` entry removed.
"""
global _background_work_result
with _background_work_lock:
if _background_work_result is not None:
return dict(_background_work_result)
monitor, thumbnails, tracker = (
_global_image_monitor,
api,
_global_prompt_tracker,
)
# A part whose global is set has started; False until its step says
# otherwise, so an interrupt between steps leaves an honest result.
started = (
("monitor", monitor),
("thumbnails", thumbnails),
("tracker", tracker),
)
result = {name: False for name, part in started if part is not None}
try:
logger = _teardown_logger()
began = _clock()
deadline = began + TEARDOWN_SECONDS
def left(own_bound):
return max(0.0, min(own_bound, deadline - _clock()))
if monitor is not None:
def stop_monitor():
from .utils.image_monitor import OBSERVER_JOIN_SECONDS
# Callbacks first: an image in hand queues no preview into a
# closing preview worker. The worker may have read the queue
# callback already, so the API refuses new previews too (at
# exit, Python 3.12+ refuses to start their worker thread).
if thumbnails is not None:
thumbnails.refuse_auto_thumbnails()
monitor.set_review_recorder(None)
monitor.set_thumbnail_queue(None)
monitor.set_output_listener(None)
return monitor.stop_monitoring(left(OBSERVER_JOIN_SECONDS))
result["monitor"] = _teardown_step(
"the image monitor", stop_monitor, logger
)
if thumbnails is not None:
def close_previews():
from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS
return thumbnails.close_auto_thumbnails(
final=True, timeout=left(CLOSE_WAIT_SECONDS)
)
result["thumbnails"] = _teardown_step(
"the preview worker", close_previews, logger
)
if tracker is not None:
def stop_tracker():
from .utils.prompt_tracker import TRACKER_JOIN_SECONDS
return tracker.shutdown(left(TRACKER_JOIN_SECONDS))
result["tracker"] = _teardown_step(
"the prompt cleanup thread", stop_tracker, logger
)
elapsed = _clock() - began
failed = [name for name, ok in result.items() if not ok]
if failed:
logger.warning(
f"Background work not stopped in time ({elapsed:.1f} s): "
+ ", ".join(failed)
)
else:
parts = " ".join(f"{name}=ok" for name in result)
logger.info(f"Background work stopped in {elapsed:.1f} s: {parts}")
finally:
_background_work_result = result
atexit.unregister(_stop_background_work)
return dict(result)
async def _stop_background_work_on_cleanup(app):
"""aiohttp ``on_cleanup``: no request is served then, so it runs inline."""
_stop_background_work()
try:
atexit.register(_stop_background_work)
from .py.config import server_instance as _teardown_server
_teardown_host = getattr(_teardown_server, "app", None)
if _teardown_host is not None and hasattr(_teardown_host, "on_cleanup"):
_teardown_host.on_cleanup.insert(0, _stop_background_work_on_cleanup)
except Exception as e:
try:
from .utils.logging_config import get_logger
get_logger("prompt_manager.init").error(
f"Failed to register the background work teardown: {e}"
)
except Exception:
pass
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
# Print startup message with loaded tools
+31
View File
@@ -1808,6 +1808,37 @@ class PromptDatabase(
reader = DocumentReader(conn)
return [self._image_row_to_dict(row, reader) for row in cursor]
def autotag_batch_candidates(self) -> List[Dict[str, Any]]:
"""The images a Tag output images batch looks at (#328).
``{id, image_path, prompt_id, prompt_tags}`` for every image with a
stored path that is linked to a prompt, newest first (the
``get_all_images`` order, ties newest row first). Never reads
``workflow_data`` or the other documents: a batch needs none of them.
"""
sql = (
"SELECT gi.id, gi.image_path, gi.prompt_id, "
"(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt "
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = gi.prompt_id) "
"AS prompt_tags "
"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 != '' "
"ORDER BY gi.generation_time DESC, gi.id DESC"
)
with self.model.get_connection() as conn:
return [
{
"id": row["id"],
"image_path": row["image_path"],
"prompt_id": row["prompt_id"],
"prompt_tags": (
row["prompt_tags"].split("|||") if row["prompt_tags"] else []
),
}
for row in conn.execute(sql)
]
def get_all_images(
self,
limit: int = 0,
+7
View File
@@ -5,6 +5,8 @@ Split into domain-specific mixins for maintainability:
- ImageRoutesMixin: gallery, thumbnails, image serving and linking
- AdminRoutesMixin: duplicates, stats, settings, diagnostics, maintenance, backup, scan
- LoggingRoutesMixin: log management endpoints
- LogCleanupMixin: the reviewed rotated-log cleanup (a maintenance plan)
- DiagnosticsReportMixin: the shareable diagnostic report, without secrets or paths
- AutotagRoutesMixin: auto-tagging model management and tagging
- JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py)
- ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs
@@ -65,6 +67,8 @@ from .duplicate_files import DuplicateFilesMixin
from .maintenance_plans import PlanStore
from .jobs import JobRegistry
from .logging_routes import LoggingRoutesMixin
from .log_cleanup import LogCleanupMixin
from .diagnostics_report import DiagnosticsReportMixin
from .autotag_routes import AutotagRoutesMixin
from .lora_integration import LoraIntegrationMixin
from .lora_import_routes import LoraImportRoutesMixin
@@ -323,6 +327,8 @@ class PromptManagerAPI(
WorkflowStorageRoutesMixin,
DuplicateFilesMixin,
LoggingRoutesMixin,
LogCleanupMixin,
DiagnosticsReportMixin,
AutotagRoutesMixin,
LoraIntegrationMixin,
LoraImportRoutesMixin,
@@ -678,6 +684,7 @@ class PromptManagerAPI(
self._register_settings_routes(routes)
self._register_job_routes(routes)
self._register_system_routes(routes)
self._register_diagnostics_report_routes(routes)
self._register_library_location_routes(routes)
self._register_maintenance_routes(routes)
self._register_backup_restore_routes(routes)
+24 -12
View File
@@ -489,9 +489,29 @@ def _diagnose_dependencies():
}
def _diagnose_output_dirs(configured_dirs):
def legacy_checks(db_path, public_path=None):
"""The four checks of ``GET /diagnostics`` (blocking).
Shared with the diagnostic report (``py/api/diagnostics_report.py``),
which passes its own ``public_path`` so a folder outside ComfyUI is never
named.
"""
from ..config import GalleryConfig
return {
"database": _diagnose_database(db_path),
"dependencies": _diagnose_dependencies(),
"comfyui_output": _diagnose_output_dirs(
GalleryConfig.MONITORING_DIRECTORIES, public_path
),
"image_monitor": _diagnose_image_monitor(public_path),
}
def _diagnose_output_dirs(configured_dirs, public_path=None):
"""Configured gallery roots, ComfyUI's output dir, or relative fallbacks."""
_public_path, _ = _public_helpers()
_public_path = public_path or _public_path
candidates = list(configured_dirs)
try:
import folder_paths
@@ -525,9 +545,10 @@ def _diagnose_output_dirs(configured_dirs):
}
def _diagnose_image_monitor():
def _diagnose_image_monitor(public_path=None):
"""Status of the running image monitor, with public directory names."""
_public_path, _public_error = _public_helpers()
_public_path = public_path or _public_path
try:
monitor = _current_image_monitor()
if monitor is None:
@@ -1055,16 +1076,7 @@ class AdminRoutesMixin:
async def run_diagnostics(self, request):
"""Run comprehensive system diagnostics and health checks."""
try:
from ..config import GalleryConfig
results = {
"database": _diagnose_database(self.db.model.db_path),
"dependencies": _diagnose_dependencies(),
"comfyui_output": _diagnose_output_dirs(
GalleryConfig.MONITORING_DIRECTORIES
),
"image_monitor": _diagnose_image_monitor(),
}
results = legacy_checks(self.db.model.db_path)
return web.json_response({"success": True, "diagnostics": results})
except Exception as e:
+17 -24
View File
@@ -10,12 +10,15 @@ filesystem or the database.
import math
WD14_MODELS = frozenset({"wd14-swinv2", "wd14-vit"})
LLM_MODELS = frozenset({"gguf", "hf"})
MODEL_TYPES = WD14_MODELS | LLM_MODELS
MAX_INSTRUCTION_CHARS = 4000
MAX_INSTRUCTION_BYTES = 16 * 1024
from ..tagging_defaults import ( # noqa: F401 - re-exported for the route
LLM_MODEL_TYPES as LLM_MODELS,
MAX_INSTRUCTION_BYTES,
MAX_INSTRUCTION_CHARS,
MODEL_TYPES,
WD14_MODEL_TYPES as WD14_MODELS,
validate_template,
validate_threshold,
)
SUGGEST_FIELDS = frozenset(
{"target", "capture_token", "expected_epoch", "model_type", "options"}
@@ -24,7 +27,7 @@ _COMMON = frozenset({"suggestion_limit", "use_gpu"})
_WD14_REQUIRED = frozenset({"general_threshold", "character_threshold"})
_LLM_REQUIRED = frozenset({"instructions"})
UNAVAILABLE_HINT = "Download it from the legacy Auto-tag page."
UNAVAILABLE_HINT = "Download it in System › Auto-tagging."
class OptionError(ValueError):
@@ -32,27 +35,17 @@ class OptionError(ValueError):
def _threshold(name: str, value) -> float:
if type(value) not in (int, float) or not math.isfinite(value):
raise OptionError(f"{name} must be a finite number")
if not 0 <= value <= 1:
raise OptionError(f"{name} must be between 0 and 1")
return float(value)
try:
return validate_threshold(name, value)
except ValueError as error:
raise OptionError(str(error)) from None
def _instructions(value) -> str:
if not isinstance(value, str) or not value.strip():
raise OptionError("instructions must be non-empty text")
if len(value) > MAX_INSTRUCTION_CHARS:
raise OptionError(
f"instructions must be at most {MAX_INSTRUCTION_CHARS} characters"
)
try:
encoded = value.encode("utf-8")
except UnicodeEncodeError:
raise OptionError("instructions must be valid Unicode text") from None
if len(encoded) > MAX_INSTRUCTION_BYTES:
raise OptionError("instructions are too long")
return value
return validate_template(value)
except ValueError as error:
raise OptionError(str(error)) from None
def validate_options(model_type, options, *, default_limit: int, max_limit: int):
+632 -89
View File
@@ -1,14 +1,26 @@
"""AutoTag API routes for PromptManager."""
import asyncio
import contextlib
import functools
import json
import os
import re
import threading
import time
from pathlib import Path
from typing import Optional
from urllib.parse import quote
from .jobs import AUTOTAG_BATCH_JOB_KIND
from .job_routes import busy_response
from .jobs import (
AUTOTAG_BATCH_JOB_KIND,
AUTOTAG_GROUP_NAMES,
MODEL_DOWNLOAD_JOB_KIND,
JobBusy,
)
from .admin import _plural, _scan_location, _scan_roots
from .maintenance import disk_plan, short_space_message
from .prompts import thumbnail_url_for
from . import thumbnail_settings
from .. import model_gate
@@ -16,20 +28,104 @@ from .. import model_gate
from aiohttp import web
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
# Tag output images (#328): the native prompt-tag batch.
AUTOTAG_JOB_TITLE = "Tag output images"
AUTOTAG_PREVIEW_LEASE = "Tag output images review"
AUTOTAG_OWNER_KIND = "autotag-job"
BATCH_PREVIEW_BODY_MAX = 1024
BATCH_JOB_BODY_MAX = 32 * 1024
REVIEW_REQUIRED = (
"Replacing tags needs a review of each image. Use Preview tagging review, "
"or choose Skip or Merge for a batch run."
)
BATCH_CANCELLED = "Cancelled. Tags already added were kept."
_BUSY_ERROR = "Auto-tagging is already running"
MODEL_BUSY_ERROR = (
"A tagging model is busy with another operation. Try again when it finishes."
)
# Who holds the model gate, as a 409 model-busy names it (``operation``). An
# owner is ``"<kind>:<uuid>"`` (``model_gate.new_owner``); a download names
# its own job, whose title says which model.
MODEL_DOWNLOAD_OWNER_KIND = "model-download"
OWNER_LABELS = {
"autotag-job": "Tag output images",
"legacy-batch": "Tag output images (legacy page)",
"legacy-single": "Tagging one image (legacy page)",
"legacy-unload": "Unloading the tagging model",
"image-tag-job": "Auto-tag selected images",
"image-annotations": "Generate suggestions",
}
ANOTHER_OPERATION = "Another tagging operation"
_DOWNLOAD_FAILED = "The download failed. Check server logs for details."
# POST /jobs/model-download carries only {model_type}.
MODEL_DOWNLOAD_BODY_MAX = 1024
# The running download's gate owner and job title. Only one download runs at a
# time and each needs the gate, so a new start replaces the entry.
_download_titles = {}
def _holder_label(holder: Optional[str]) -> str:
if not holder:
return ANOTHER_OPERATION
kind = holder.split(":", 1)[0]
if kind == MODEL_DOWNLOAD_OWNER_KIND:
return _download_titles.get(holder, ANOTHER_OPERATION)
return OWNER_LABELS.get(kind, ANOTHER_OPERATION)
def model_busy_response() -> web.Response:
"""``409 model-busy``: another operation holds the shared model gate."""
"""``409 model-busy``: another operation holds the shared model gate.
``operation`` names the holder (``OWNER_LABELS``), or "Another tagging
operation" when the gate was freed meanwhile or the kind is unknown.
"""
return web.json_response(
{"success": False, "code": "model-busy", "error": MODEL_BUSY_ERROR},
{
"success": False,
"code": "model-busy",
"error": MODEL_BUSY_ERROR,
"operation": _holder_label(model_gate.holder()),
},
status=409,
)
_SIZE = re.compile(r"^(~?[0-9.]+)\s*([KMGT]?B)$")
def size_label(size: str) -> str:
"""The engine table's size with a space before the unit (``~400 MB``)."""
match = _SIZE.match(str(size or "").strip())
return f"{match.group(1)} {match.group(2)}" if match else str(size or "")
def public_models(status: dict, loaded_type: Optional[str]) -> dict:
"""Each engine model's public entry: names and sizes from ``MODELS``, the
flags from ``status``; never a path (``model_path`` is dropped)."""
from ..autotag import MODELS
models = {}
for model_type, config in MODELS.items():
flags = status.get(model_type) or {}
entry = {
"name": config["name"],
"label": config["label"],
"description": config["description"],
"size": config["size"],
"size_label": size_label(config["size"]),
"downloaded": bool(flags.get("downloaded")),
"loaded": loaded_type == model_type,
}
for key in ("model_exists", "mmproj_exists"):
if key in flags:
entry[key] = bool(flags[key])
models[model_type] = entry
return models
# 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"}
@@ -65,6 +161,39 @@ def path_is_within(candidate, root) -> bool:
return child != parent and child.is_relative_to(parent)
def _thresholds_or_saved(general, character):
"""The request's WD14 thresholds as floats; one it omits is the saved
default (``autotag.defaults``, OQ4). Raises ValueError/TypeError."""
from ..config import AutotagConfig
return (
AutotagConfig.GENERAL_THRESHOLD if general is None else float(general),
AutotagConfig.CHARACTER_THRESHOLD if character is None else float(character),
)
async def _single(event):
"""A stream of one event."""
yield event
_monotonic = time.monotonic # the legacy frames' throttle clock (patched in tests)
# What one image came to in a batch (``on_item``): fixed strings, never an
# exception's text.
SKIP_DETAILS = {
"unlinked": "File not found",
"prompt already processed": "Prompt handled by another image",
"outside allowed directories": "Outside the output folders",
"file missing": "File not found",
"already tagged": "Prompt already tagged",
}
def _added(count):
return f"Added {count} tag" if count == 1 else f"Added {count} tags"
async def _refusal(title):
"""The single error event of an auto-tag batch refused by another job."""
message = f"{title} already running"
@@ -83,6 +212,10 @@ class AutotagRoutesMixin:
async def download_autotag_model_route(request):
return await self.download_autotag_model(request)
@routes.post("/prompt_manager/autotag/batch/preview")
async def preview_autotag_batch_route(request):
return await self.preview_autotag_batch(request)
@routes.post("/prompt_manager/autotag/start")
async def start_autotag_route(request):
return await self.start_autotag(request)
@@ -110,18 +243,24 @@ class AutotagRoutesMixin:
from ..config import AutotagConfig
service = get_autotag_service()
models_status = service.get_models_status()
# Off the loop: the HuggingFace check walks the whole HF cache.
status = await self._run_in_executor(service.get_models_status)
loaded = service.is_model_loaded()
loaded_type = service.get_loaded_model_type() if loaded else None
return web.json_response(
{
"success": True,
"models": models_status,
"default_prompt": service.default_prompt,
"model_loaded": service.is_model_loaded(),
"loaded_model_type": service.get_loaded_model_type(),
"wd14_general_threshold": service.wd14_general_threshold,
"wd14_character_threshold": service.wd14_character_threshold,
"models": public_models(status, loaded_type),
# The saved tagging defaults (#328, System › Auto-tagging).
"default_prompt": AutotagConfig.TEMPLATE,
"model_loaded": loaded,
"loaded_model_type": loaded_type,
"wd14_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
"wd14_character_threshold": AutotagConfig.CHARACTER_THRESHOLD,
"suggestion_limit": AutotagConfig.suggestion_limit(),
"model_type": AutotagConfig.MODEL_TYPE,
"existing": AutotagConfig.EXISTING,
}
)
@@ -159,55 +298,156 @@ class AutotagRoutesMixin:
return response
async def download_autotag_model(self, request):
"""Download an AutoTag model with streaming progress (POST only)."""
"""Download a tagging model, streaming the legacy events (POST only).
The download is the model_download Activity job: this request
attaches to a running download of the same model or starts one, and
streams the legacy frames until the job ends. A client disconnect
ends only the stream; the download goes on. A running download of
another model answers 409 model-busy (it holds the gate).
"""
from ..autotag import MODELS, get_autotag_service
model_type = request.match_info.get("model_type")
owner = model_gate.new_owner("legacy-download")
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 model_gate.run(
owner, 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,
}
if not model_gate.try_acquire(owner):
return model_busy_response()
if model_type not in MODELS:
invalid = {"type": "error", "message": f"Invalid model type: {model_type}"}
return await self._stream_sse(request, _single(invalid))
try:
return await self._stream_sse(request, events())
finally:
job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND)
if job is not None and job.scope.get("model_type") != model_type:
return model_busy_response()
if job is None:
job, refusal = self._start_model_download(
model_type, get_autotag_service()
)
if refusal is not None:
return refusal
except Exception:
self.logger.exception("Download model error")
failed = {"type": "error", "message": _INTERNAL_ERROR}
return await self._stream_sse(request, _single(failed))
return await self._stream_sse(request, self._download_frames(job))
@staticmethod
async def _download_frames(job):
"""The legacy download frames for job: start, then its end."""
yield {"type": "progress", "progress": 0, "status": "Starting download..."}
async with contextlib.aclosing(job.events()) as events:
async for _payload in events:
pass
if job.state == "finished":
yield {"type": "complete", "progress": 100, "status": "Download complete"}
else:
yield {"type": "error", "message": "Download failed"}
# ── model download job (#328) ───────────────────────────────────────
async def start_model_download_job(self, request):
"""``POST /jobs/model-download {model_type}``: 202 with the job.
400 for a bad body, 409 ``busy`` while a download (or an exclusive
job) runs, 409 ``model-busy`` while another operation holds the gate.
"""
data, error = await self._read_json_body(
request, max_bytes=MODEL_DOWNLOAD_BODY_MAX
)
if error is not None:
return error
if set(data) != {"model_type"}:
return self._bad_request("The body must be exactly {model_type}")
from ..autotag import MODELS, get_autotag_service
model_type = data["model_type"]
if not isinstance(model_type, str) or model_type not in MODELS:
return self._bad_request(
"model_type must be one of " + ", ".join(sorted(MODELS))
)
try:
job, refusal = self._start_model_download(model_type, get_autotag_service())
except Exception as e:
self.logger.exception("Model download start error")
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
if refusal is not None:
return refusal
return web.json_response({"success": True, "job": job.summary()}, status=202)
def _start_model_download(self, model_type, service):
"""Register the download job: ``(job, None)`` or ``(None, refusal)``.
Synchronous (no await between the checks and the start). The gate is
taken before the job registers and released when the runner ends,
also when the task is cancelled before the runner ever starts.
"""
from ..autotag import MODELS
busy = self._jobs.blocker(MODEL_DOWNLOAD_JOB_KIND, admission="read")
if busy is not None:
return None, busy_response(JobBusy(busy))
owner = model_gate.new_owner(MODEL_DOWNLOAD_OWNER_KIND)
if not model_gate.try_acquire(owner):
return None, model_busy_response()
config = MODELS[model_type]
label = config["label"]
title = f"Tagging model download · {label}"
_download_titles.clear()
_download_titles[owner] = title
try:
job = self._jobs.start(
MODEL_DOWNLOAD_JOB_KIND,
title,
admission="read",
cancellable=False,
scope={
"model_type": model_type,
"label": label,
"size_label": size_label(config["size"]),
},
runner=functools.partial(
self._run_model_download, service, model_type, owner
),
)
except BaseException as refused:
# JobBusy, or anything else the registry raised: give the gate back.
model_gate.release(owner)
_download_titles.pop(owner, None)
if isinstance(refused, JobBusy):
return None, busy_response(refused)
raise
def ended(_task):
model_gate.release(owner)
if model_gate.holder() != owner: # not while a thread still runs
_download_titles.pop(owner, None)
job.task.add_done_callback(ended)
return job, None
async def _run_model_download(self, service, model_type, owner, job):
label = job.scope["label"]
job.update(status_line=f"Downloading {label} ({job.scope['size_label']})")
job.log_line(f"Download started: {label}")
job.publish({"type": "status", "status": job.status_line})
try:
ok = await model_gate.run(owner, service.download_model, model_type, None)
except Exception:
self.logger.exception(f"Tagging model download failed: {model_type}")
ok = False
finally:
# Before the job ends, so a reader that sees the end finds the
# gate free (a running download thread defers the release).
model_gate.release(owner)
if ok:
job.update(status_line=f"Downloaded {label}")
job.log_line(f"Download finished: {label}")
job.publish({"type": "status", "status": job.status_line})
job.finish("finished")
return
job.update(status_line=_DOWNLOAD_FAILED)
job.log_line(f"Download failed: {label}")
job.publish({"type": "status", "status": job.status_line})
job.finish("failed")
async def _read_start_params(self, request) -> Optional[dict]:
"""Merge query-string and optional JSON-body parameters for autotag/start.
@@ -238,35 +478,48 @@ class AutotagRoutesMixin:
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."""
async def _apply_tags_to_prompt(self, prompt_id, tags) -> int:
"""Merge ``tags`` into the prompt; how many new tags were added."""
existing_prompt = await self._run_in_executor(
self.db.get_prompt_by_id, prompt_id
)
if not existing_prompt:
return False
return 0
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]
new_tags = list(dict.fromkeys(t for t in tags if t not in existing_tags))
if not new_tags:
return False
return 0
await self._run_in_executor(
self.db.update_prompt_metadata, prompt_id, tags=existing_tags + new_tags
)
return True
return len(new_tags)
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
"""Yield SSE progress dicts while tagging every image linked to a prompt.
Both batches run here: the legacy page's (its frames unchanged) and
the ``autotag`` job's. ``opts``: ``model_type``, ``keep_in_memory``,
``existing`` (``skip`` leaves prompts with real tags alone, ``merge``
adds to them), the two thresholds, and optionally ``gate_owner``,
``prompt`` (the template, passed to each call; the legacy page sets
``custom_prompt`` instead), ``cancelled`` (checked before each image:
a cancelled run returns without a ``complete`` frame) and
``on_item(index, total, group, image_path, detail)`` (one outcome per
image, ``SKIP_DETAILS`` for the skips).
"""
model_type = opts["model_type"]
keep_in_memory = opts["keep_in_memory"]
status = service.get_models_status()
cancelled = opts.get("cancelled") or (lambda: False)
on_item = opts.get("on_item") or (lambda *_args: None)
status = await self._run_in_executor(service.get_models_status)
if not status.get(model_type, {}).get("downloaded"):
yield {"type": "error", "message": f"Model {model_type} not downloaded"}
return
if cancelled():
return
yield {"type": "progress", "progress": 0, "status": "Loading model..."}
run = self._gated_runner(opts.get("gate_owner"))
try:
@@ -279,15 +532,21 @@ class AutotagRoutesMixin:
}
return
if opts["custom_prompt"]:
if opts.get("custom_prompt"):
service.custom_prompt = opts["custom_prompt"]
tag_kwargs = {
"general_threshold": opts["general_threshold"],
"character_threshold": opts["character_threshold"],
}
if opts.get("prompt"):
tag_kwargs["prompt"] = opts["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)
images = await self._run_in_executor(self.db.autotag_batch_candidates)
total = len(images)
if total == 0:
yield {
@@ -299,7 +558,8 @@ class AutotagRoutesMixin:
"skipped_outside_roots": 0,
"status": "No images with linked prompts found",
}
service.unload_model()
if not keep_in_memory:
service.unload_model()
return
yield {
@@ -316,7 +576,7 @@ class AutotagRoutesMixin:
"skipped_outside_roots": 0,
}
tagged_prompt_ids = set()
last_update = _time.monotonic()
last_update = _monotonic()
finished = False
def progress_event(i, message):
@@ -332,6 +592,8 @@ class AutotagRoutesMixin:
try:
for i, image_data in enumerate(images):
if cancelled():
return
image_path = image_data.get("image_path")
prompt_id = image_data.get("prompt_id")
skip_reason = None
@@ -350,7 +612,7 @@ class AutotagRoutesMixin:
counts["skipped_outside_roots"] += 1
else:
skip_reason = "file missing"
elif opts["skip_tagged"] and self._real_tags(
elif opts["existing"] == "skip" and self._real_tags(
image_data.get("prompt_tags", [])
):
tagged_prompt_ids.add(prompt_id)
@@ -358,7 +620,8 @@ class AutotagRoutesMixin:
if skip_reason is not None:
counts["skipped"] += 1
now = _time.monotonic()
on_item(i, total, "skipped", image_path, SKIP_DETAILS[skip_reason])
now = _monotonic()
if (now - last_update) >= 0.5 or i == total - 1:
yield progress_event(
i, f"Skipping {i + 1}/{total} ({skip_reason})..."
@@ -368,17 +631,17 @@ class AutotagRoutesMixin:
continue
try:
tags = await run(
service.generate_tags,
allowed_path,
general_threshold=opts["general_threshold"],
character_threshold=opts["character_threshold"],
)
tags = await run(service.generate_tags, allowed_path, **tag_kwargs)
counts["processed"] += 1
if tags and await self._apply_tags_to_prompt(prompt_id, tags):
added = 0
if tags:
added = await self._apply_tags_to_prompt(prompt_id, tags)
if added:
counts["tagged"] += 1
on_item(i, total, "tagged", image_path, _added(added))
else:
counts["skipped"] += 1
on_item(i, total, "unchanged", image_path, "No new tags")
tagged_prompt_ids.add(prompt_id)
except Exception as img_err:
self.logger.error(
@@ -387,10 +650,11 @@ class AutotagRoutesMixin:
)
counts["errors"] += 1
counts["processed"] += 1
on_item(i, total, "errors", image_path, "Tagging failed")
yield progress_event(i, f"Processing {i + 1}/{total}...")
await asyncio.sleep(0.01)
last_update = _time.monotonic()
last_update = _monotonic()
finished = True
if not keep_in_memory:
@@ -418,16 +682,22 @@ class AutotagRoutesMixin:
opts = {
"model_type": params.get("model_type", "gguf"),
"custom_prompt": params.get("prompt", ""),
"skip_tagged": str(params.get("skip_tagged", "true")).lower() == "true",
"existing": (
"skip"
if str(params.get("skip_tagged", "true")).lower() == "true"
else "merge"
),
"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])
opts["general_threshold"], opts["character_threshold"] = (
_thresholds_or_saved(
opts["general_threshold"], opts["character_threshold"]
)
)
except (ValueError, TypeError):
return self._bad_request(
"general_threshold and character_threshold must be numeric"
@@ -457,7 +727,9 @@ class AutotagRoutesMixin:
# never runs beside a scan, restore or the workspace tagging job.
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
if busy is not None:
if busy.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
# A second legacy batch keeps its JSON 409; the workspace's Tag
# output images job is named in a stream frame.
if busy.kind == AUTOTAG_BATCH_JOB_KIND and busy.title != AUTOTAG_JOB_TITLE:
return web.json_response(
{"success": False, "error": _BUSY_ERROR}, status=409
)
@@ -467,7 +739,7 @@ class AutotagRoutesMixin:
return await self._stream_sse(request, _refusal(busy.title))
async with self._jobs.hold(
AUTOTAG_BATCH_JOB_KIND,
"Auto-tag images (legacy page)",
OWNER_LABELS["legacy-batch"],
admission="records",
) as job:
response = await self._start_autotag_batch(request, opts, events)
@@ -491,6 +763,278 @@ class AutotagRoutesMixin:
finally:
_autotag_batch_lock.release()
# ── Tag output images (#328) ────────────────────────────────────────
async def _batch_body(self, request, keys, max_bytes):
"""``(data, None)`` for a JSON object with exactly ``keys``."""
data, error = await self._read_json_body(request, max_bytes=max_bytes)
if error is not None:
return None, error
if set(data) != set(keys):
return None, self._bad_request(
"The body must be exactly {" + ", ".join(keys) + "}"
)
return data, None
@staticmethod
def _batch_choice(data):
"""The model's engine entry and the policy, or a 400 message."""
from ..autotag import MODELS
from ..config import EXISTING_POLICIES
model_type, existing = data.get("model_type"), data.get("existing")
if not isinstance(model_type, str) or model_type not in MODELS:
return None, "model_type must be one of " + ", ".join(sorted(MODELS))
if not isinstance(existing, str) or existing not in EXISTING_POLICIES:
return None, "existing must be one of " + ", ".join(EXISTING_POLICIES)
return MODELS[model_type], None
async def preview_autotag_batch(self, request):
"""Review Tag output images: counts, model, backup and disk; no write.
Opens no image file. Refused (409 ``busy``) while an exclusive job runs.
"""
try:
async with self._jobs.lease(AUTOTAG_PREVIEW_LEASE):
return await self._preview_autotag_batch(request)
except JobBusy as busy:
return busy_response(busy)
async def _preview_autotag_batch(self, request):
from ..autotag import get_autotag_service
data, error = await self._batch_body(
request, ("model_type", "existing"), BATCH_PREVIEW_BODY_MAX
)
if error is not None:
return error
config, message = self._batch_choice(data)
if message:
return self._bad_request(message)
model_type, existing = data["model_type"], data["existing"]
policy = self._safety_backup_policy()
try:
candidates = await self._run_in_executor(self.db.autotag_batch_candidates)
status = await self._run_in_executor(
get_autotag_service().get_models_status
)
disk, _short, _size = await self._run_in_executor(
functools.partial(
disk_plan,
self.db.model.db_path,
backup=policy != "off",
vacuum=False,
)
)
except Exception as e:
self.logger.exception("Tag output images could not be reviewed")
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
prompts = {row["prompt_id"]: row["prompt_tags"] for row in candidates}
tagged = sum(1 for tags in prompts.values() if self._real_tags(tags))
to_tag = len(prompts) - tagged if existing == "skip" else len(prompts)
return web.json_response(
{
"success": True,
"images": len(candidates),
"prompts": len(prompts),
"already_tagged": tagged,
"to_tag": to_tag,
"model": {
"type": model_type,
"label": config["label"],
"size_label": size_label(config["size"]),
"downloaded": bool(
(status.get(model_type) or {}).get("downloaded")
),
},
"existing": existing,
"review_required": existing == "replace_after_review",
"backup": {"policy": policy},
"disk": disk,
}
)
async def start_autotag_job(self, request):
"""``POST /jobs/autotag``: start Tag output images (202), or refuse.
In order: 400 body or options, 400 ``review-required``; 409
``model-unavailable``; 507 ``insufficient-space``; 409 ``busy``
(another records job, the legacy batch's hold included); 409
``model-busy``. No await follows the busy check.
"""
from ..autotag import get_autotag_service
from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig
from .annotation_suggest import OptionError, validate_options
data, error = await self._batch_body(
request, ("model_type", "existing", "options"), BATCH_JOB_BODY_MAX
)
if error is not None:
return error
config, message = self._batch_choice(data)
if message:
return self._bad_request(message)
model_type, existing, options = (
data["model_type"],
data["existing"],
data["options"],
)
if isinstance(options, dict) and {"suggestion_limit", "use_gpu"} & set(options):
return self._bad_request("options take no suggestion_limit or use_gpu")
try:
_model, options = validate_options(
model_type,
options,
default_limit=AutotagConfig.suggestion_limit(),
max_limit=MAX_SUGGESTION_LIMIT,
)
except OptionError as error:
return self._bad_request(str(error))
if existing == "replace_after_review":
return web.json_response(
{"success": False, "code": "review-required", "error": REVIEW_REQUIRED},
status=400,
)
service = get_autotag_service()
backup_on = self._safety_backup_policy() == "on"
try:
status = await self._run_in_executor(service.get_models_status)
candidates = await self._run_in_executor(self.db.autotag_batch_candidates)
_disk, short, _size = await self._run_in_executor(
functools.partial(
disk_plan, self.db.model.db_path, backup=backup_on, vacuum=False
)
)
except Exception as e:
self.logger.exception("Tag output images could not start")
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
if not (status.get(model_type) or {}).get("downloaded"):
return web.json_response(
{
"success": False,
"code": "model-unavailable",
"error": f"{config['label']} is not downloaded. "
"Download it in System › Auto-tagging.",
},
status=409,
)
if backup_on and short is not None:
return web.json_response(
{
"success": False,
"code": "insufficient-space",
"error": short_space_message(short),
},
status=507,
)
busy = self._jobs.blocker(AUTOTAG_BATCH_JOB_KIND, admission="records")
if busy is not None:
return busy_response(JobBusy(busy))
owner = model_gate.new_owner(AUTOTAG_OWNER_KIND)
if not model_gate.try_acquire(owner):
return model_busy_response()
opts = {
"model_type": model_type,
"existing": existing,
"keep_in_memory": True,
"gate_owner": owner,
"general_threshold": options.get("general_threshold"),
"character_threshold": options.get("character_threshold"),
"prompt": options.get("instructions"),
}
try:
job = self._jobs.start(
AUTOTAG_BATCH_JOB_KIND,
AUTOTAG_JOB_TITLE,
admission="records",
cancellable=True,
scope={
"model_type": model_type,
"existing": existing,
"images": len(candidates),
"label": config["label"],
},
groups=AUTOTAG_GROUP_NAMES,
runner=functools.partial(
self._run_autotag_job, service, opts, backup_on
),
)
except BaseException as refused:
model_gate.release(owner)
if isinstance(refused, JobBusy):
return busy_response(refused)
raise
job.task.add_done_callback(lambda _task: model_gate.release(owner))
return web.json_response({"success": True, "job": job.summary()}, status=202)
async def _run_autotag_job(self, service, opts, backup_on, job):
ok, _backup = await self._safety_backup_step(job, self.db, backup_on)
if not ok:
return
try:
await self._tag_output_images(service, opts, job)
except Exception:
self.logger.exception("Tag output images failed")
job.update(status_line=_INTERNAL_ERROR)
job.log_line(f"Tagging failed · {_INTERNAL_ERROR}")
job.finish("failed")
finally:
if backup_on:
await self._prune_backup_pool(job, self.db)
async def _tag_output_images(self, service, opts, job):
roots = _scan_roots(self._get_all_output_dirs())
job.log_line(
f"Tagging started · {job.scope['label']} · "
+ ("skip prompts already tagged" if opts["existing"] == "skip" else "merge")
)
job.update(status_line="Loading the tagging model")
def on_item(index, total, group, image_path, detail):
path, root, root_index = _scan_location(image_path or "", roots)
job.update(
processed=index + 1,
total=total,
progress=100 * (index + 1) // total,
status_line=f"Tagging {index + 1} of {total}",
)
job.record(
group,
{
"path": path,
"root": root,
"root_index": root_index,
"detail": detail,
},
f"{path}: {detail}",
)
job.publish({"type": "progress", "progress": job.progress})
final = None
async for event in self._autotag_batch_events(
service,
{**opts, "cancelled": lambda: job.cancelled, "on_item": on_item},
):
if event.get("type") in ("complete", "error"):
final = event
if final is not None and final["type"] == "error":
job.update(status_line=final["message"])
job.log_line(f"Tagging failed · {final['message']}")
job.finish("failed")
elif job.cancelled: # a cancel wins over a run that then completed
job.update(status_line=BATCH_CANCELLED)
job.log_line(BATCH_CANCELLED)
job.finish("cancelled")
else:
done = f"Tagged {_plural(job.counters.get('tagged', 0), 'prompt')}"
job.update(status_line=done, progress=100)
job.log_line(done)
job.finish("finished")
def _gated_runner(self, owner):
"""Executor calls under the model gate *owner* holds (plain when None)."""
if owner is None:
@@ -607,10 +1151,9 @@ class AutotagRoutesMixin:
general_threshold = data.get("general_threshold")
character_threshold = data.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)
general_threshold, character_threshold = _thresholds_or_saved(
general_threshold, character_threshold
)
except (ValueError, TypeError):
return self._bad_request(
"general_threshold and character_threshold must be numeric"
+289
View File
@@ -0,0 +1,289 @@
"""The shareable diagnostic report (#329 System S12).
``GET /prompt_manager/system/diagnostics`` answers a report built on the
server from an explicit list of fields: the project and Viewer.js versions,
Python and the platform, the four checks of the legacy ``GET /diagnostics``
(one helper, ``admin.legacy_checks``), the database's public info, source
counts, the log configuration and statistics, the last health summary and
the configuration keys that only ``config.json`` holds. It never copies
``PromptManagerConfig.get_config()``: no path leaf is listed, and the CivitAI
key is reported only as ``has_civitai_api_key``.
Every path in the report goes through ``report_path``: relative to the
closest public anchor that is not a filesystem root, else ``"outside
ComfyUI"``, never a basename that could be a user name. As a backstop,
``scrub`` replaces any string that still looks like an absolute path, holds
the saved key or the home folder (with or without its root), and counts each
replacement in ``redacted`` (0 for a correct builder).
The report opens the database, so it takes a registry lease: it is refused
(409 ``busy``) while an ``exclusive`` job (a restore) runs, and a restore
waits for it. It is built off the event loop.
"""
import datetime
import json
import ntpath
import os
import platform
import posixpath
import re
from pathlib import Path, PurePosixPath, PureWindowsPath
from aiohttp import web
from . import admin as _admin
from .job_routes import busy_response
from .jobs import JobBusy
from .system import health_summary
try:
from ...utils.version import project_version
except ImportError: # loaded as the top-level ``py`` package (tests)
from utils.version import project_version
DIAGNOSTICS_LEASE = "Diagnostics"
_ROOT = Path(__file__).resolve().parents[2]
VIEWERJS_PROVENANCE = str(_ROOT / "web" / "lib" / "viewerjs" / "provenance.json")
OUTSIDE = "outside ComfyUI"
REMOVED = "[removed]"
# The CivitAI key is searched for only when it is long enough to be a key.
MIN_KEY_LENGTH = 8
_WINDOWS_ABSOLUTE = re.compile(r"^[A-Za-z]:[\\/]")
def anchored_path(canonical, anchors, home, pathmod=os.path):
"""``canonical`` relative to the closest usable anchor, else ``outside ComfyUI``.
All three are already canonical (``normcase``) in ``pathmod``'s spelling.
A filesystem root is never usable, nor is a strict ancestor of ``home``
(``/home``, ``C:\\Users``), whose relative paths would start with a user
name; so a path inside the home folder needs an anchor at or below it.
"""
flavor = PureWindowsPath if pathmod is ntpath else PurePosixPath
target = flavor(canonical)
home_path = flavor(home) if home else None
usable = []
for anchor in map(flavor, anchors):
if anchor.parent == anchor:
continue
if home_path is not None and home_path != anchor:
if home_path.is_relative_to(anchor):
continue
if target.is_relative_to(anchor):
usable.append(anchor)
if not usable:
return OUTSIDE
anchor = max(usable, key=lambda a: len(a.parts))
return target.relative_to(anchor).as_posix() or "."
def _canonical(path):
return os.path.normcase(os.path.realpath(str(path)))
def report_path(path):
"""``path`` relative to a public anchor (POSIX style), else ``outside ComfyUI``."""
if not path:
return ""
from ..config import GalleryConfig
try:
canonical = _canonical(path)
anchors = [_canonical(anchor) for anchor in GalleryConfig.path_anchors()]
home = os.path.expanduser("~")
home = _canonical(home) if home and home != "~" else None
except (OSError, ValueError):
return OUTSIDE
return anchored_path(canonical, anchors, home)
def _looks_absolute(text):
return (
posixpath.isabs(text)
or ntpath.isabs(text)
or _WINDOWS_ABSOLUTE.match(text) is not None
or text.startswith("\\\\")
)
# The home folder is matched only as whole path parts, never inside a word.
_PART = r"(?<![\w.-]){}(?![\w.-])"
MIN_HOME_NAME = 3
def _home_patterns():
"""``(pattern, paths_only)`` for the home folder: in full, rootless, its name.
Accepted residual risk: only this process's own home folder is known, so
another user's home (an output folder under /home/<them> while ComfyUI
runs as a service account) can still be named by a report path.
"""
home = os.path.expanduser("~")
if not home or home == "~" or len(home) < 4 or home in ("/", "\\"):
return ()
spellings = {home, home.replace("\\", "/")}
relative = os.path.splitdrive(home)[1].strip("\\/")
for spelling in (relative, relative.replace("\\", "/")):
if "/" in spelling or "\\" in spelling:
spellings.add(spelling)
patterns = [
(re.compile(_PART.format(re.escape(spelling)), re.IGNORECASE), False)
for spelling in spellings
]
name = re.split(r"[\\/]", relative)[-1] if relative else ""
if len(name) >= MIN_HOME_NAME:
# The bare name only inside a path ("info/x"), not a word ("INFO").
patterns.append(
(re.compile(_PART.format(re.escape(name)), re.IGNORECASE), True)
)
return tuple(patterns)
def scrub(report, key):
"""``report`` with leaked paths and the key replaced, counted in ``redacted``."""
key = key if isinstance(key, str) and len(key) >= MIN_KEY_LENGTH else None
homes = _home_patterns()
count = 0
def clean(value):
nonlocal count
if isinstance(value, str):
if (
_looks_absolute(value)
or (key and key in value)
or any(
pattern.search(value)
for pattern, paths_only in homes
if not paths_only or "/" in value or "\\" in value
)
):
count += 1
return REMOVED
return value
if isinstance(value, dict):
return {name: clean(item) for name, item in value.items()}
if isinstance(value, (list, tuple)):
return [clean(item) for item in value]
return value
cleaned = clean(report)
cleaned["redacted"] = int(cleaned.get("redacted") or 0) + count
return cleaned
def viewerjs_version(path=None):
"""The pinned Viewer.js version from its provenance file, or ``unknown``."""
try:
with open(path or VIEWERJS_PROVENANCE, encoding="utf-8") as handle:
version = json.load(handle).get("version")
except (OSError, ValueError, AttributeError):
return "unknown"
return version if isinstance(version, str) and version else "unknown"
def public_config():
"""The configuration keys only ``config.json`` holds, from an explicit list."""
from ..config import GalleryConfig, IntegrationConfig, PromptManagerConfig
return {
"database": {
"enable_duplicate_detection": (
PromptManagerConfig.ENABLE_DUPLICATE_DETECTION
),
"enable_auto_save": PromptManagerConfig.ENABLE_AUTO_SAVE,
},
"web_ui": {
"result_timeout": PromptManagerConfig.RESULT_TIMEOUT,
"show_test_button": PromptManagerConfig.SHOW_TEST_BUTTON,
"webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE,
"infinite_scroll": PromptManagerConfig.INFINITE_SCROLL,
},
"performance": {
"max_search_results": PromptManagerConfig.MAX_SEARCH_RESULTS,
"enable_fuzzy_search": PromptManagerConfig.ENABLE_FUZZY_SEARCH,
"auto_backup_interval": PromptManagerConfig.AUTO_BACKUP_INTERVAL,
"worker_threads": PromptManagerConfig.WORKER_THREADS,
},
"gallery": {
"processing_delay": GalleryConfig.PROCESSING_DELAY,
"prompt_timeout": GalleryConfig.PROMPT_TIMEOUT,
"cleanup_interval": GalleryConfig.CLEANUP_INTERVAL,
"auto_cleanup_missing_files": GalleryConfig.AUTO_CLEANUP_MISSING_FILES,
"max_image_age_days": GalleryConfig.MAX_IMAGE_AGE_DAYS,
"max_concurrent_processing": GalleryConfig.MAX_CONCURRENT_PROCESSING,
"metadata_extraction_timeout": GalleryConfig.METADATA_EXTRACTION_TIMEOUT,
},
"integrations": {
"lora_manager_enabled": bool(IntegrationConfig.LORA_MANAGER_ENABLED),
"has_civitai_api_key": bool(IntegrationConfig.CIVITAI_API_KEY),
},
}
def _source_counts(sources):
configured = sources.get("configured") or []
effective = sources.get("effective") or []
return {
"configured": len(configured),
"effective": len(effective),
"watching": sum(1 for row in effective if row.get("watch") == "watching"),
"unavailable": sum(1 for row in configured if row.get("valid") is False),
}
class DiagnosticsReportMixin:
"""Mixin providing ``GET /prompt_manager/system/diagnostics``."""
def _register_diagnostics_report_routes(self, routes):
@routes.get("/prompt_manager/system/diagnostics")
async def diagnostics_report_route(request):
return await self.get_diagnostics_report(request)
async def get_diagnostics_report(self, request):
"""The report, built off the loop under a lease (409 ``busy`` in a restore)."""
from ..config import IntegrationConfig
try:
async with self._jobs.lease(DIAGNOSTICS_LEASE):
report = await self._run_in_executor(self.build_report)
except JobBusy as busy:
return busy_response(busy)
except Exception:
self.logger.exception("The diagnostic report could not be built")
return web.json_response(
{
"success": False,
"code": None,
"error": "The diagnostic report could not be built.",
},
status=500,
)
report = scrub(report, IntegrationConfig.CIVITAI_API_KEY)
return web.json_response(report, headers={"Cache-Control": "no-store"})
def build_report(self):
"""The report's fields (blocking: reads the database and the log folder)."""
db = self.db
manager = self._get_logger_manager()
stats = dict(manager.get_log_stats())
stats["log_directory"] = report_path(stats.get("log_directory"))
info = dict(db.model.get_database_info() or {})
if "database_path" in info:
info["database_path"] = report_path(info["database_path"])
sources, _monitor = self._source_status()
return {
"success": True,
"generated_at": datetime.datetime.now().replace(microsecond=0).isoformat(),
"version": project_version(_ROOT),
"python": platform.python_version(),
"platform": {"system": platform.system(), "release": platform.release()},
"viewerjs": viewerjs_version(VIEWERJS_PROVENANCE),
"checks": _admin.legacy_checks(db.model.db_path, report_path),
"database_info": info,
"sources": _source_counts(sources),
"logging": {"config": dict(manager.get_config()), "stats": stats},
"health": health_summary(self._current_health_report()),
"config": public_config(),
"redacted": 0,
}
+11 -14
View File
@@ -2272,7 +2272,8 @@ class ImageRoutesMixin:
)
image_path = Path(image["image_path"])
allowed_dirs = list(self._get_all_output_dirs()) + self._lora_image_dirs()
lora_dirs = await self._run_in_executor(self._lora_image_dirs)
allowed_dirs = list(self._get_all_output_dirs()) + lora_dirs
denied = _media_access_error(image_path, allowed_dirs)
if denied is not None:
@@ -2320,7 +2321,10 @@ class ImageRoutesMixin:
LoraManager's guarded LoRA roots, sidecar root and example-images
root (``lora_manager_layout.media_roots``, #355) and PromptManager's
CivitAI download cache. An unusable LoraManager layout gives the
default LoRA roots and the cache.
default LoRA roots and the cache. The layout comes from the layout
cache (``lora_utils.cached_lora_manager_layout``, #362), which reloads
when LoraManager's settings files change; blocking, so the routes
call it in the executor.
"""
try:
from ..config import IntegrationConfig
@@ -2328,25 +2332,18 @@ class ImageRoutesMixin:
if not IntegrationConfig.LORA_MANAGER_ENABLED:
return []
from ..lora_utils import (
cached_lora_manager_layout,
get_lora_image_cache_dir,
layout_for_install,
lora_manager_layout,
)
lora_root = IntegrationConfig.LORA_MANAGER_PATH or ""
if lora_root.strip():
configured = IntegrationConfig.LORA_MANAGER_PATH or ""
install, layout = cached_lora_manager_layout(configured)
if configured.strip() and install is None:
# A configured path comes from config.json, which can be hand
# edited: only a LoraManager install under custom_nodes may
# widen the serving allow-list. An empty path keeps the
# auto-detection, which itself only scans custom_nodes.
from ..lora_utils import resolve_lora_manager_path
lora_root = resolve_lora_manager_path(lora_root)
if not lora_root:
return []
layout = layout_for_install(lora_root)
else:
layout = lora_manager_layout("")
return []
dirs = [Path(d) for d in layout.media_roots]
# Named, not made: listing and serving never create the cache.
+8
View File
@@ -130,6 +130,14 @@ class JobRoutesMixin:
async def start_image_tag_job_route(request):
return await self.start_image_tag_job(request)
@routes.post("/prompt_manager/jobs/model-download")
async def start_model_download_job_route(request):
return await self.start_model_download_job(request)
@routes.post("/prompt_manager/jobs/autotag")
async def start_autotag_job_route(request):
return await self.start_autotag_job(request)
@routes.post("/prompt_manager/jobs/thumbnails")
async def start_thumbnail_job_route(request):
return await self.start_thumbnail_job(request)
+4
View File
@@ -87,6 +87,8 @@ SCAN_GROUPS = {
# Record-changing: it writes image tag sets and new vocabulary names.
IMAGE_TAG_JOB_KIND = "image-tags"
IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
# Tag output images (#328) counts its prompts the same way.
AUTOTAG_GROUP_NAMES = IMAGE_TAG_GROUP_NAMES
# Thumbnail generation (#249 System S5): previews written, already fresh, failed.
THUMBNAIL_GROUP_NAMES = ("generated", "skipped", "errors")
@@ -109,6 +111,8 @@ DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
THUMBNAILS_JOB_KIND = "thumbnails"
LORA_IMPORT_JOB_KIND = "lora_import"
AUTOTAG_BATCH_JOB_KIND = "autotag"
# A tagging model download (#328): read admission, not cancellable.
MODEL_DOWNLOAD_JOB_KIND = "model_download"
# Duplicate image files (#249 System S6, py/api/duplicate_files.py): the
# find is a cancellable ``read`` job; the reviewed removal is ``records``.
DUPLICATES_JOB_KIND = "duplicates"
+207
View File
@@ -0,0 +1,207 @@
"""Reviewed rotated-log cleanup as a maintenance plan (#329 System S12).
``POST /prompt_manager/maintenance/preview {operation: "rotated_logs",
options: {}}`` lists the rotated copies of the active log that Rule R allows
(``utils/logging_config.py``: ``prompt_manager.log.<n>`` regular files that
are not links, directly in the log folder) and writes nothing. The plan keeps
each copy's name, size, ``mtime_ns`` and inode under its token; the response
is the maintenance ``Plan`` shape with basenames only, a ``none`` backup
policy (no library record changes) and no ``disk`` key.
``POST /prompt_manager/maintenance/apply {token}`` starts a ``maintenance``
job titled "Rotated log copies" with admission ``read``: it may run beside a
scan and is refused only beside another maintenance job or an ``exclusive``
one. The job deletes exactly the reviewed copies under the logger's files
lock and the file handler's lock, or nothing, failing ``plan-changed``, when
any of them is missing, is no longer a rotated copy or differs in size,
``mtime_ns`` or inode. A copy that cannot be deleted (open in another program
on Windows) is an ``errors`` row; the others are deleted.
"""
import datetime
import json
from aiohttp import web
try:
from ...utils.logging_config import RotatedChanged, rotated_sort_key
except ImportError: # loaded as the top-level ``py`` package (tests)
from utils.logging_config import RotatedChanged, rotated_sort_key
from .jobs import MAINTENANCE_JOB_KIND
from .maintenance_plans import ScopeTooLarge
ROTATED_LOGS = "rotated_logs"
LOG_CLEANUP_TITLE = "Rotated log copies"
LOGS_CHANGED = (
"The log files changed since the review; nothing was deleted. Review again."
)
LOG_CLEANUP_FAILED = "The log cleanup failed."
TOO_MANY_LOG_FILES = "Too many log files to review at once."
LOG_CLEANUP_MEDIA = (
"Deletes only the listed log files. No image, video or library record changes."
)
_GROUPS = ("changed", "protected", "errors")
def size_label(size):
"""``1000 bytes``, ``2.0 KB``, ``10.0 MB`` or ``1.2 GB`` (1 KB = 1024 bytes)."""
size = max(0, int(size or 0))
if size < 1024:
return f"{size} {'byte' if size == 1 else 'bytes'}"
for unit, scale in (("KB", 2**10), ("MB", 2**20)):
if size < scale * 1024:
return f"{size / scale:.1f} {unit}"
return f"{size / 2**30:.1f} GB"
def time_label(mtime_ns):
"""The local modification time as ``YYYY-MM-DD HH:MM``."""
try:
moment = datetime.datetime.fromtimestamp(mtime_ns / 1e9)
except (OSError, OverflowError, ValueError): # a time before 1970 on Windows
return "unknown time"
return moment.strftime("%Y-%m-%d %H:%M")
def _count(count, singular, plural):
return f"{count} {singular if count == 1 else plural}"
def _protected_reasons(rows):
"""The plan's Protected reasons: what the cleanup keeps, zero counts left out."""
active = sum(1 for row in rows if row["kind"] == "active")
links = sum(1 for row in rows if row["kind"] == "link")
other = sum(1 for row in rows if row["kind"] == "other")
reasons = []
if active:
reasons.append("Kept: the active log")
if links:
reasons.append(f"Kept: {_count(links, 'link', 'links')}")
if other:
reasons.append(f"Kept: {_count(other, 'other file', 'other files')}")
return active + links + other, reasons
def _scope_label(files):
if not files:
return "No rotated log copies"
total = sum(row["size"] for row in files)
return (
f"{_count(len(files), 'rotated log copy', 'rotated log copies')}"
f" · {size_label(total)}"
)
class LogCleanupMixin:
"""The ``rotated_logs`` maintenance plan and its job.
The logger manager is reached only through ``_get_logger_manager``
(``logging_routes.py``), so the tests can point it at a temporary folder.
"""
async def preview_log_cleanup(self):
"""The ``rotated_logs`` plan (runs inside the preview's lease)."""
from . import maintenance
manager = self._get_logger_manager()
try:
rows = await self._run_in_executor(manager.scan_log_folder)
except Exception:
self.logger.exception("Rotated log review failed")
return maintenance._failure("Failed to review maintenance", 500)
rotated = sorted(
(row for row in rows if row["kind"] == "rotated"),
key=rotated_sort_key,
)
files = [
{
"name": row["name"],
"size": row["size"],
"mtime_ns": row["mtime_ns"],
"ino": row["ino"],
}
for row in rotated
]
protected, reasons = _protected_reasons(rows)
store = self._maintenance_plans
try:
token = store.issue(
ROTATED_LOGS,
json.dumps(files, sort_keys=True),
{"files": files, "protected": protected},
(),
)
except ScopeTooLarge:
return maintenance._failure(TOO_MANY_LOG_FILES, 413, "too-large")
plan = store.peek(token, ROTATED_LOGS)
return web.json_response(
{
"success": True,
"token": token,
"operation": ROTATED_LOGS,
"expires_at": maintenance._iso(plan.expires_at),
"scope": {"count": len(files), "label": _scope_label(files)},
"protected": {"count": protected, "reasons": reasons},
"backup": {
"policy": "none",
"destination_label": None,
"expected_bytes": None,
},
"media": {"deletes_files": True, "label": LOG_CLEANUP_MEDIA},
"details": [
{
"label": row["name"],
"detail": (
f"{size_label(row['size'])} · {time_label(row['mtime_ns'])}"
),
}
for row in files
],
},
headers={"Cache-Control": "no-store"},
)
def _start_log_cleanup(self, plan):
"""Register the cleanup job (on the loop); raises JobBusy when refused."""
return self._jobs.start(
MAINTENANCE_JOB_KIND,
LOG_CLEANUP_TITLE,
admission="read",
cancellable=False,
scope={"operation": ROTATED_LOGS},
groups=_GROUPS,
runner=lambda job: self._run_log_cleanup(job, plan),
)
async def _run_log_cleanup(self, job, plan):
from . import maintenance
files = plan.payload["files"]
sizes = {row["name"]: row["size"] for row in files}
job.update(total=len(files), status_line="Deleting rotated log copies")
manager = self._get_logger_manager()
try:
result = await self._run_in_executor(manager.delete_rotated, files)
except RotatedChanged:
return maintenance._fail_job(job, "plan-changed", LOGS_CHANGED)
except Exception:
self.logger.exception("Rotated log cleanup failed")
return maintenance._fail_job(job, None, LOG_CLEANUP_FAILED)
for name in result["removed"]:
detail = f"Deleted · {size_label(sizes.get(name, 0))}"
job.record(
"changed", {"path": name, "detail": detail}, f"- {name} · {detail}"
)
for name, error in result["failed"]:
detail = self._public_error(error)
job.record(
"errors", {"path": name, "detail": detail}, f"! {name} · {detail}"
)
job.record_counted("protected", plan.payload["protected"])
removed, failed = len(result["removed"]), len(result["failed"])
line = f"Deleted {_count(removed, 'rotated log copy', 'rotated log copies')}"
if failed:
line += f", {failed} could not be deleted"
job.log_line(line)
job.update(processed=len(files), progress=100, status_line=line)
+51 -9
View File
@@ -1,6 +1,13 @@
"""Logging API routes for PromptManager."""
"""Logging API routes for PromptManager.
Every route that names a log file goes through Rule R (#329,
``utils/logging_config.py``): a download is served only for a regular,
non-link file listed directly in the log folder, and Clear deletes only
rotated copies. File work runs off the event loop.
"""
import os
import stat
from aiohttp import web
@@ -85,6 +92,41 @@ def _validate_log_config(data):
return clean, None
def _read_log_bytes(path):
"""The whole file at ``path``, opened without following a link (blocking).
The file is checked again through the open handle: a regular file with
one name and, when ``path`` carries it, the inode it was listed with. A
name swapped for a link or another file after it was listed, or a hard
link to a file elsewhere, is refused as missing. The handle is closed
before this returns.
"""
flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_NOFOLLOW", 0)
# A FIFO swapped in after the listing must not block the executor thread.
flags |= getattr(os, "O_NONBLOCK", 0)
try:
fd = os.open(str(path), flags)
except OSError as error:
if getattr(os, "O_NOFOLLOW", 0) and os.path.islink(str(path)):
raise FileNotFoundError("Log file not found") from error
raise
listed_ino = getattr(path, "listed_ino", 0)
with os.fdopen(fd, "rb") as handle:
opened = os.fstat(handle.fileno())
if (
not stat.S_ISREG(opened.st_mode)
or opened.st_nlink > 1 # a hard link to a file elsewhere
or (listed_ino and opened.st_ino and opened.st_ino != listed_ino)
):
raise FileNotFoundError("Log file not found")
return handle.read()
def _download(logger_manager, filename):
"""``bytes`` of a downloadable log file; FileNotFoundError otherwise."""
return _read_log_bytes(logger_manager.log_file_entry(filename))
class LoggingRoutesMixin:
"""Mixin providing logging-related API endpoints."""
@@ -195,16 +237,16 @@ class LoggingRoutesMixin:
{"success": False, "error": "Invalid filename"}, status=400
)
log_file_path = logger_manager.log_dir / filename
if not log_file_path.exists():
try:
file_content = await self._run_in_executor(
_download, logger_manager, filename
)
except FileNotFoundError:
# Missing, a link, a folder or not listed: all read as absent.
return web.json_response(
{"success": False, "error": "Log file not found"}, status=404
)
with open(log_file_path, "rb") as f:
file_content = f.read()
return web.Response(
body=file_content,
content_type="text/plain",
@@ -224,7 +266,7 @@ class LoggingRoutesMixin:
"""Truncate all log files."""
try:
logger_manager = self._get_logger_manager()
results = logger_manager.truncate_logs()
results = await self._run_in_executor(logger_manager.truncate_logs)
return web.json_response(
{
@@ -266,7 +308,7 @@ class LoggingRoutesMixin:
return web.json_response({"success": False, "error": error}, status=400)
logger_manager = self._get_logger_manager()
logger_manager.update_config(clean)
await self._run_in_executor(logger_manager.update_config, clean)
return web.json_response(
{
+9 -2
View File
@@ -194,8 +194,13 @@ class LoraImportRoutesMixin:
return False
async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty):
"""Swap the collection for ``staged.records`` in one transaction."""
return await self._run_in_executor(
"""Swap the collection for ``staged.records`` in one transaction.
Then the image routes' layout cache is dropped (#362): the import may
have linked files under a folder LoraManager created after the cache
stored its layout (its first example-image download creates one).
"""
result = await self._run_in_executor(
lambda: db.replace_category_prompts(
CATEGORY,
staged.records,
@@ -203,6 +208,8 @@ class LoraImportRoutesMixin:
allow_empty=allow_empty,
)
)
lora_utils.clear_layout_cache()
return result
# ── preview ───────────────────────────────────────────────────────
+19 -2
View File
@@ -2,7 +2,9 @@
``POST /prompt_manager/maintenance/preview {operation, options:{}}`` reviews
one operation (``duplicate_prompts``, ``optimize``, ``unused_components`` or
``missing_links``; ``database/maintenance.py``) off the event loop and
``missing_links``; ``database/maintenance.py``; ``duplicate_files``,
``py/api/duplicate_files.py``; ``rotated_logs``, the rotated log copies,
``py/api/log_cleanup.py``, which has its own apply job) off the event loop and
answers a ``Plan``: ``{success, token, operation, expires_at, scope:{count,
label}, protected:{count, reasons}, backup:{policy, destination_label},
media:{deletes_files, label}, details:[≤50 {label, detail}]}``. The preview
@@ -54,6 +56,7 @@ from aiohttp import web
from .admin import _prune_safety_backups, _scan_location, _scan_roots
from .job_routes import busy_response
from .jobs import MAINTENANCE_JOB_KIND, JobBusy
from .log_cleanup import ROTATED_LOGS
from .maintenance_plans import PlanExpired, PlanUsed, ScopeTooLarge
try:
@@ -83,7 +86,8 @@ MAINTENANCE_PREVIEW_LEASE = "Maintenance preview"
DUPLICATE_PREVIEW_BODY_MAX = 1024 * 1024
DUPLICATE_FILES = "duplicate_files"
WORKFLOW_STORAGE = "workflow_storage" # applied by py/api/workflow_storage.py
PREVIEW_OPERATIONS = MAINTENANCE_OPERATIONS + (DUPLICATE_FILES,)
# rotated_logs (#329, py/api/log_cleanup.py) reviews log files, not records.
PREVIEW_OPERATIONS = MAINTENANCE_OPERATIONS + (DUPLICATE_FILES, ROTATED_LOGS)
# Room left over beyond what a step writes, so a disk is never filled to 0.
DISK_MARGIN_BYTES = 16 * 2**20
# The concept's Maintenance row titles, which name each operation's job.
@@ -391,6 +395,8 @@ class MaintenanceRoutesMixin:
return await self.preview_duplicate_files(options)
if not isinstance(options, dict) or options:
return _invalid(f"{operation} takes no options")
if operation == ROTATED_LOGS:
return await self.preview_log_cleanup()
db = self.db
identity = self._database_identity(db)
try:
@@ -492,6 +498,17 @@ class MaintenanceRoutesMixin:
if refused is not None:
return refused
store = self._maintenance_plans
if plan.operation == ROTATED_LOGS:
# Log files only: no database identity, backup or disk check.
try:
job = self._start_log_cleanup(plan)
except JobBusy as busy:
return busy_response(busy) # the token stays usable
# Same loop turn as the peek and the start: the runner has not run.
store.consume(token, plan.operation, job_id=job.id)
return web.json_response(
{"success": True, "job": job.summary()}, status=202
)
db = self.db
if plan.db_identity != self._database_identity(db):
store.discard(token)
+84 -1
View File
@@ -47,11 +47,23 @@ _SETTINGS_FILE_KEYS = {
"THUMBNAIL_FORMAT": ("thumbnails", "format"),
"THUMBNAIL_QUALITY": ("thumbnails", "quality"),
"THUMBNAIL_EDGE": ("thumbnails", "edge"),
# The one tagging configuration (#328), published to AutotagConfig.
"AUTOTAG_MODEL_TYPE": ("autotag", "defaults", "model_type"),
"AUTOTAG_EXISTING": ("autotag", "defaults", "existing"),
"AUTOTAG_GENERAL_THRESHOLD": ("autotag", "defaults", "general_threshold"),
"AUTOTAG_CHARACTER_THRESHOLD": ("autotag", "defaults", "character_threshold"),
"AUTOTAG_TEMPLATE": ("autotag", "defaults", "template"),
"AUTOTAG_SUGGESTION_LIMIT": ("autotag", "defaults", "suggestion_limit"),
}
# The thumbnail settings a running Convert outdated job fixes (#317).
_THUMBNAIL_SETTINGS = ("THUMBNAIL_FORMAT", "THUMBNAIL_QUALITY", "THUMBNAIL_EDGE")
# Attributes that live on GalleryConfig; the others on PromptManagerConfig.
_GALLERY_SETTINGS = frozenset({"MONITORING_ENABLED"})
# Attributes that live on AutotagConfig, named without the AUTOTAG_ prefix.
_AUTOTAG_SETTINGS = frozenset(
name for name in _SETTINGS_FILE_KEYS if name.startswith("AUTOTAG_")
)
_AUTOTAG_PREFIX = "AUTOTAG_"
SOURCE_CHANGE_JOB_KIND = "source_change"
SOURCE_CHANGE_JOB_TITLE = "Source change"
@@ -151,6 +163,63 @@ def _set_nested(config, sections, key, value):
config[sections[0]] = merged
def _validate_autotag_settings(data):
"""The ``autotag_*`` settings as ``{AUTOTAG_ATTR: value}``.
The first bad value raises ValueError with a message that names the
setting by its label in System › Auto-tagging.
"""
from ..config import (
EXISTING_POLICIES,
MAX_SUGGESTION_LIMIT,
MODEL_TYPES,
validate_saved_threshold,
validate_template,
)
updates = {}
if "autotag_model_type" in data:
value = data["autotag_model_type"]
if not isinstance(value, str) or value not in MODEL_TYPES:
raise ValueError(
"Tagging model must be one of " + ", ".join(sorted(MODEL_TYPES)) + "."
)
updates["AUTOTAG_MODEL_TYPE"] = value
if "autotag_existing" in data:
value = data["autotag_existing"]
if not isinstance(value, str) or value not in EXISTING_POLICIES:
raise ValueError(
"Existing tags must be skip, merge or replace_after_review."
)
updates["AUTOTAG_EXISTING"] = value
for key, label in (
("autotag_general_threshold", "General-tag confidence"),
("autotag_character_threshold", "Character-tag confidence"),
):
if key in data:
try:
value = validate_saved_threshold(label, data[key])
except ValueError as error:
raise ValueError(f"{error}.") from None
updates[key.upper()] = value
if "autotag_template" in data:
try:
updates["AUTOTAG_TEMPLATE"] = validate_template(
data["autotag_template"], "LLM tag template"
)
except ValueError as error:
raise ValueError(f"{error}.") from None
if "autotag_suggestion_limit" in data:
value = data["autotag_suggestion_limit"]
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
raise ValueError(
"Suggestions per image must be a whole number from 1 to "
f"{MAX_SUGGESTION_LIMIT}."
)
updates["AUTOTAG_SUGGESTION_LIMIT"] = value
return updates
async def _read_json(request):
"""Parse a JSON object body through the package's capped reader."""
from . import _read_json_body
@@ -177,7 +246,7 @@ class SettingsRoutesMixin:
async def get_settings(self, request):
"""Get current settings."""
try:
from ..config import PromptManagerConfig, GalleryConfig
from ..config import AutotagConfig, PromptManagerConfig, GalleryConfig
root_paths = [
self._public_path(d) for d in GalleryConfig.MONITORING_DIRECTORIES
@@ -203,6 +272,14 @@ class SettingsRoutesMixin:
"thumbnail_format": PromptManagerConfig.THUMBNAIL_FORMAT,
"thumbnail_quality": PromptManagerConfig.THUMBNAIL_QUALITY,
"thumbnail_edge": PromptManagerConfig.THUMBNAIL_EDGE,
"autotag_model_type": AutotagConfig.MODEL_TYPE,
"autotag_existing": AutotagConfig.EXISTING,
"autotag_general_threshold": AutotagConfig.GENERAL_THRESHOLD,
"autotag_character_threshold": (
AutotagConfig.CHARACTER_THRESHOLD
),
"autotag_template": AutotagConfig.TEMPLATE,
"autotag_suggestion_limit": AutotagConfig.SUGGESTION_LIMIT,
"cpu_count": PromptManagerConfig.max_worker_threads(),
"gallery_root_paths": root_paths,
"gallery_root_path": root_paths[0] if root_paths else "",
@@ -425,7 +502,12 @@ class SettingsRoutesMixin:
"""
from ..config import GalleryConfig, PromptManagerConfig
from ..config import AutotagConfig
for name, value in updates.items():
if name in _AUTOTAG_SETTINGS:
setattr(AutotagConfig, name[len(_AUTOTAG_PREFIX) :], value)
continue
owner = GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig
setattr(owner, name, value)
restart_fields = []
@@ -564,6 +646,7 @@ class SettingsRoutesMixin:
updates["THUMBNAIL_EDGE"] = thumbnail_settings.validate_edge(
data["thumbnail_edge"]
)
updates.update(_validate_autotag_settings(data))
return updates
def _parse_gallery_roots(self, data):
+31 -12
View File
@@ -488,10 +488,16 @@ class AutoThumbnailer:
f"Automatic preview failed for {src.name}", exc_info=True
)
def close(self):
def close(self, timeout=None):
"""Stop the worker at its next file boundary (queued files are
dropped), kill a running ffmpeg and wait up to CLOSE_WAIT_SECONDS;
a worker still running after that is logged."""
dropped), kill a running ffmpeg and wait up to ``timeout`` seconds
(None: CLOSE_WAIT_SECONDS, read now); a worker still running after
that is logged.
Returns True when there was no worker or it ended; False when it is
still running, or when called on the worker itself (it only signals).
"""
wait = CLOSE_WAIT_SECONDS if timeout is None else timeout
with self._changed:
self._closed = True
self._queue.clear()
@@ -504,14 +510,18 @@ class AutoThumbnailer:
process.kill()
except OSError:
pass
if worker is None or worker is threading.current_thread():
return
worker.join(CLOSE_WAIT_SECONDS)
if worker is None:
return True
if worker is threading.current_thread():
return False
worker.join(wait)
if worker.is_alive():
self._logger.warning(
"The automatic preview worker is still running after "
f"{CLOSE_WAIT_SECONDS} s; it stops at its next file"
f"{wait} s; it stops at its next file"
)
return False
return True
# ── waits (tests and shutdown; never timers) ──────────────────────
@@ -609,12 +619,19 @@ class AutoThumbnailMixin:
)
return self._auto_thumbnailer
def close_auto_thumbnails(self, final=False):
def refuse_auto_thumbnails(self):
"""Refuse every later capture without stopping the worker yet (the
package teardown, before ``close_auto_thumbnails(final=True)``)."""
with AutoThumbnailMixin._auto_thumbnailer_lock:
self._auto_thumbnails_shut = True
def close_auto_thumbnails(self, final=False, timeout=None):
"""Stop the capture worker and the view pool, and forget them.
Bounded (see AutoThumbnailer.close). ``final`` (the app's cleanup)
also keeps a late monitor event from starting a worker again; a
test reset passes False and the next use starts fresh ones.
Bounded by ``timeout`` (see AutoThumbnailer.close). ``final`` (the
package teardown) also keeps a late monitor event from starting a
worker again; a test reset passes False and the next use starts
fresh ones. Returns True when no capture worker is left running.
"""
with AutoThumbnailMixin._auto_thumbnailer_lock:
if final:
@@ -622,12 +639,14 @@ class AutoThumbnailMixin:
thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None
pool, self._preview_pool_obj = self._preview_pool_obj, None
self._preview_gate_obj = None
closed = True
if thumbnailer is not None:
thumbnailer.close()
closed = thumbnailer.close(timeout)
if pool is not None:
pool.shutdown(wait=False, cancel_futures=True)
if final:
shutdown_encoders() # kill JPEG XL / ffmpeg encoders still running
return closed
def _auto_thumbnail_failures(self):
thumbnailer = self._auto_thumbnailer
+30 -28
View File
@@ -25,6 +25,23 @@ except ImportError:
from database.annotations import MAX_IMAGE_TAGS
from database.vocabulary import validate_vocabulary_tag_name
# The tag template and thresholds live with the saved tagging defaults
# (#328), so config.py never imports PIL.
try:
from .tagging_defaults import DEFAULT_TAG_TEMPLATE
except ImportError:
from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
def _autotag_config():
"""The saved defaults; imported late, as config needs ComfyUI's server."""
try:
from .config import AutotagConfig
except ImportError:
from py.config import AutotagConfig
return AutotagConfig
# Try to import logging from utils, fallback to standard logging
try:
from ..utils.logging_config import get_logger
@@ -45,10 +62,12 @@ except ImportError:
return logger
# Model configurations - same as standalone_tagger.py
# Model configurations - same as standalone_tagger.py. ``label`` is the short
# name System, Activity and a 409 model-busy show (#328).
MODELS = {
"gguf": {
"name": "GGUF (Recommended)",
"label": "GGUF JoyCaption",
"description": "Quantized model, ~4.5GB, faster loading",
"size": "~4.5GB",
"repo": "mradermacher/llama-joycaption-beta-one-hf-llava-GGUF",
@@ -59,6 +78,7 @@ MODELS = {
},
"hf": {
"name": "HuggingFace",
"label": "HuggingFace JoyCaption",
"description": "Full precision model, ~16GB, higher quality",
"size": "~16GB",
"repo": "fancyfeast/llama-joycaption-beta-one-hf-llava",
@@ -66,6 +86,7 @@ MODELS = {
},
"wd14-swinv2": {
"name": "WD14 SwinV2",
"label": "WD14 SwinV2",
"description": "Fast classifier tagger, ~400MB, no prompt needed",
"size": "~400MB",
"repo": "SmilingWolf/wd-swinv2-tagger-v3",
@@ -75,6 +96,7 @@ MODELS = {
},
"wd14-vit": {
"name": "WD14 ViT",
"label": "WD14 ViT",
"description": "Fast classifier tagger, ~400MB, ViT architecture",
"size": "~400MB",
"repo": "SmilingWolf/wd-vit-tagger-v3",
@@ -85,17 +107,7 @@ MODELS = {
}
# Default prompts
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."
)
DEFAULT_PROMPT = DEFAULT_TAG_TEMPLATE
# Upper bound of one suggestion request: the image tag set cap.
MAX_SUGGESTION_LIMIT = MAX_IMAGE_TAGS
@@ -179,10 +191,6 @@ class AutoTagService:
self._current_model_type: Optional[str] = None
self._custom_prompt: str = DEFAULT_PROMPT
# WD14 threshold defaults
self._wd14_general_threshold: float = 0.35
self._wd14_character_threshold: float = 0.85
@property
def models_config(self) -> Dict[str, Dict[str, Any]]:
"""Get the models configuration dictionary."""
@@ -205,19 +213,13 @@ class AutoTagService:
@property
def wd14_general_threshold(self) -> float:
return self._wd14_general_threshold
@wd14_general_threshold.setter
def wd14_general_threshold(self, value: float):
self._wd14_general_threshold = max(0.0, min(1.0, value))
"""The saved General-tag confidence (``autotag.defaults``)."""
return _autotag_config().GENERAL_THRESHOLD
@property
def wd14_character_threshold(self) -> float:
return self._wd14_character_threshold
@wd14_character_threshold.setter
def wd14_character_threshold(self, value: float):
self._wd14_character_threshold = max(0.0, min(1.0, value))
"""The saved Character-tag confidence (``autotag.defaults``)."""
return _autotag_config().CHARACTER_THRESHOLD
def get_models_status(self) -> Dict[str, Dict[str, Any]]:
"""Get availability status for all model types.
@@ -860,12 +862,12 @@ class AutoTagService:
gen_thresh = (
general_threshold
if general_threshold is not None
else self._wd14_general_threshold
else _autotag_config().GENERAL_THRESHOLD
)
char_thresh = (
character_threshold
if character_threshold is not None
else self._wd14_character_threshold
else _autotag_config().CHARACTER_THRESHOLD
)
return self._generate_wd14(
session, tags_list, image, gen_thresh, char_thresh
+98 -16
View File
@@ -563,18 +563,92 @@ class IntegrationConfig:
DEFAULT_SUGGESTION_LIMIT = 10
MAX_SUGGESTION_LIMIT = 100 # pinned to MAX_IMAGE_TAGS by tests/test_config.py
# The tagging engines, options, template and validators (server-free, #328).
try:
from .tagging_defaults import ( # noqa: F401 - re-exported
DEFAULT_TAG_TEMPLATE,
EXISTING_POLICIES,
LLM_MODEL_TYPES,
MAX_INSTRUCTION_BYTES,
MAX_INSTRUCTION_CHARS,
MODEL_TYPES,
WD14_MODEL_TYPES,
validate_saved_threshold,
validate_template,
validate_threshold,
)
except ImportError:
from py.tagging_defaults import ( # noqa: F401 - re-exported
DEFAULT_TAG_TEMPLATE,
EXISTING_POLICIES,
LLM_MODEL_TYPES,
MAX_INSTRUCTION_BYTES,
MAX_INSTRUCTION_CHARS,
MODEL_TYPES,
WD14_MODEL_TYPES,
validate_saved_threshold,
validate_template,
validate_threshold,
)
def _valid_model_type(value: Any) -> str:
if not isinstance(value, str) or value not in MODEL_TYPES:
raise ValueError("must be one of " + ", ".join(sorted(MODEL_TYPES)))
return value
def _valid_existing(value: Any) -> str:
if not isinstance(value, str) or value not in EXISTING_POLICIES:
raise ValueError("must be one of " + ", ".join(EXISTING_POLICIES))
return value
def _valid_suggestion_limit(value: Any) -> int:
if type(value) is not int or not 1 <= value <= MAX_SUGGESTION_LIMIT:
raise ValueError(f"must be an integer from 1 to {MAX_SUGGESTION_LIMIT}")
return value
class AutotagConfig:
"""Persisted tagging defaults under ``autotag.defaults`` in config.json.
``suggestion_limit`` is how many suggestions Generate returns (the top N
by score), an integer in ``1..MAX_SUGGESTION_LIMIT``. A hand-edited value
outside that range, or of another type, falls back to the default with a
warning rather than failing the load.
One tagging configuration (#328): the model, what a batch does with
prompts already tagged, the two WD14 thresholds, the language-model
template and how many suggestions Generate returns. Each key is checked
alone on load; a hand-edited value that fails its check falls back to
its default with one warning rather than failing the load.
"""
MODEL_TYPE = "wd14-swinv2"
EXISTING = "skip"
GENERAL_THRESHOLD = 0.35
CHARACTER_THRESHOLD = 0.85
TEMPLATE = DEFAULT_TAG_TEMPLATE
SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
# config.json key -> (attribute, validator, default)
_KEYS = {
"model_type": ("MODEL_TYPE", _valid_model_type, "wd14-swinv2"),
"existing": ("EXISTING", _valid_existing, "skip"),
"general_threshold": (
"GENERAL_THRESHOLD",
lambda v: validate_saved_threshold("value", v),
0.35,
),
"character_threshold": (
"CHARACTER_THRESHOLD",
lambda v: validate_saved_threshold("value", v),
0.85,
),
"template": ("TEMPLATE", validate_template, DEFAULT_TAG_TEMPLATE),
"suggestion_limit": (
"SUGGESTION_LIMIT",
_valid_suggestion_limit,
DEFAULT_SUGGESTION_LIMIT,
),
}
@classmethod
def suggestion_limit(cls) -> int:
return cls.SUGGESTION_LIMIT
@@ -587,7 +661,14 @@ class AutotagConfig:
@classmethod
def get_config(cls) -> Dict[str, Any]:
return {"defaults": {"suggestion_limit": cls.SUGGESTION_LIMIT}}
"""The saved defaults; the template only when it is not the built-in
one, so a later built-in template still reaches who never set one."""
defaults = {
key: getattr(cls, attr) for key, (attr, _v, _d) in cls._KEYS.items()
}
if defaults["template"] == DEFAULT_TAG_TEMPLATE:
del defaults["template"]
return {"defaults": defaults}
@classmethod
def update_config(cls, new_config: Any) -> None:
@@ -598,17 +679,18 @@ class AutotagConfig:
if not isinstance(defaults, dict):
config_logger.warning("Ignoring autotag.defaults that is not an object")
return
if "suggestion_limit" not in defaults:
return
value = defaults["suggestion_limit"]
if cls.valid_suggestion_limit(value):
cls.SUGGESTION_LIMIT = value
return
config_logger.warning(
"autotag.defaults.suggestion_limit must be an integer from 1 to "
f"{MAX_SUGGESTION_LIMIT}; using {DEFAULT_SUGGESTION_LIMIT}"
)
cls.SUGGESTION_LIMIT = DEFAULT_SUGGESTION_LIMIT
for key, (attr, validate, default) in cls._KEYS.items():
if key not in defaults:
continue
try:
setattr(cls, attr, validate(defaults[key]))
except ValueError as error:
config_logger.warning(
f"autotag.defaults.{key} {error}; using {default!r}"
if key != "template"
else f"autotag.defaults.template {error}; using the built-in one"
)
setattr(cls, attr, default)
class PromptManagerConfig:
+13 -10
View File
@@ -12,9 +12,10 @@ last writes to the database:
the sidecars (``discover``). In ``alongside`` mode a sidecar beside its
model is kept when ``<base>.safetensors`` exists; in ``centralized`` mode
every sidecar under the sidecar root is kept when its ``file_path`` names
a model inside a LoRA root. ``fingerprint`` covers the kept sidecars (path,
size and modification time) and the layout; ``preview`` adds the counts a
review shows. Nothing is written and nothing is requested.
a model inside a LoRA root. ``fingerprint`` covers every scanned sidecar,
kept or skipped (path, size, modification time, kept or not, #362) and the
layout; ``preview`` adds the counts a review shows. Nothing is written
and nothing is requested.
2. **stage**: read each sidecar and build its record. Its LoraManager preview
and example images are linked where they are (``local_media``); only a
LoRA without any local image downloads its CivitAI example images into the
@@ -69,9 +70,10 @@ def discover(lora_manager_path: Optional[str]) -> Discovery:
def fingerprint(discovery: Discovery) -> str:
"""What a review saw: the layout, then each kept sidecar's path, size and
modification time (a LoraManager save that changes unrelated settings
leaves it as it was)."""
"""What a review saw: the layout, then every scanned sidecar, kept or
skipped, with its path, size, modification time and whether it was kept
(#362). A LoraManager save that changes unrelated settings leaves it as
it was."""
layout = discovery.layout
digest = hashlib.sha256()
for value in (
@@ -87,13 +89,14 @@ def fingerprint(discovery: Discovery) -> str:
digest.update(f"{value}\n".encode("utf-8", "surrogatepass"))
for name in discovery.unreadable:
digest.update(f"unreadable\0{name}\n".encode("utf-8", "surrogatepass"))
for sidecar in discovery.sidecars:
path = sidecar.metadata_path
kept = {str(sidecar.metadata_path) for sidecar in discovery.sidecars}
for path in discovery.scanned_files:
flag = 1 if str(path) in kept else 0
try:
stat = path.stat()
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\n"
entry = f"{path}\0{stat.st_size}\0{stat.st_mtime_ns}\0{flag}\n"
except OSError:
entry = f"{path}\0missing\n"
entry = f"{path}\0missing\0{flag}\n"
digest.update(entry.encode("utf-8", "surrogatepass"))
return "sha256:" + digest.hexdigest()
+217
View File
@@ -0,0 +1,217 @@
"""A per-process cache of LoraManager's layout for the image hot paths (#362).
Every Library list, prompt detail, inspect, reveal and serve request needs the
LoRA allow-list (``ImagesMixin._lora_image_dirs``). Building it means finding
the LoraManager install, reading its plugin settings file and its settings
file, and resolving every folder. ``LayoutCache`` keeps the last result and
reuses it while all of these hold:
- the key (the configured install path, stripped) and the home folder the
folder guard compares against are the same;
- the entry is younger than ``MAX_AGE_SECONDS`` (monotonic clock);
- the plugin settings file and the settings file have the same signature
(``st_mtime_ns``, ``st_size``, or ``MISSING``) as when it was stored;
- the install and every stored media root still resolve to themselves (a
folder swapped for a symlink reloads, so the folder guard runs again).
A result is stored only when no watched file is "racy" (modified within
``RACY_SECONDS`` of the load, git's rule for racily clean index entries),
every watched file could be checked, the settings file was readable, and no
``clear()`` ran during the load. What each rule covers: signatures catch a
settings change on the next request; the racy rule catches a write during the
load, a half-written file and coarse timestamps; the maximum age bounds the
inputs no signature sees (``extra_model_paths.yaml``, environment, temporary
failures); ``clear()`` makes a finished import's new folders servable at once.
Hits and loads take no lock: a slow ``os.stat`` on a network share holds up
only its own request, and two threads that miss together both load (both
loads only read). A short lock guards the compare-and-store and the logging
state; no file I/O runs under it. Standard library only, so it is tested
without ComfyUI or LoraManager.
"""
import logging
import os
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional, Tuple
from .lora_manager_layout import UNREADABLE, LoraManagerLayout, _canonical
# A watched file modified this recently is not trusted: LoraManager writes its
# settings in place, and FAT/exFAT store modification times in 2 s steps.
RACY_SECONDS = 3
# Bounds what no signature sees (unwatched inputs, temporary failures).
MAX_AGE_SECONDS = 30
# The signature of a watched file that does not exist (a valid state).
MISSING = object()
# ERROR_BAD_NETPATH and ERROR_BAD_NET_NAME, which Windows raises as
# FileNotFoundError.
_OFFLINE_SHARE_ERRORS = frozenset({53, 67})
logger = logging.getLogger("prompt_manager.lora_layout_cache")
@dataclass(frozen=True)
class _Entry:
key: str
home: str
install: Optional[str]
layout: LoraManagerLayout
signatures: Tuple[Tuple[str, object], ...]
stored_at: float
class _CannotValidate(Exception):
"""A watched file could not be checked (not merely missing)."""
class LayoutCache:
"""At most one ``(install, layout)`` for the active LoraManager config."""
def __init__(
self,
resolve: Callable[[str], Optional[str]],
load: Callable[[Optional[str]], LoraManagerLayout],
*,
wall_ns: Callable[[], int] = time.time_ns,
monotonic: Callable[[], float] = time.monotonic,
stat=os.stat,
pathmod=os.path,
home: Callable[[], str] = lambda: str(Path.home()),
):
self._resolve = resolve
self._load = load
self._wall_ns = wall_ns
self._monotonic = monotonic
self._stat = stat
self._pathmod = pathmod
self._home = home
self._entry: Optional[_Entry] = None
self._generation = 0
self._lock = threading.Lock()
# What was last logged; survives clear() so a reload of the same
# layout stays quiet.
self._logged_layout: Optional[LoraManagerLayout] = None
self._logged_refused: Tuple[str, ...] = ()
def get(self, configured: str) -> Tuple[Optional[str], LoraManagerLayout]:
"""``(install, layout)`` for the configured path ("" auto-detects)."""
key = (configured or "").strip()
home = self._home()
entry = self._entry
if entry is not None and self._valid(entry, key, home):
return entry.install, entry.layout
with self._lock:
if entry is not None and self._entry is entry:
self._entry = None
generation = self._generation
install = self._resolve(key)
layout = self._load(install)
self._log_refused(layout)
try:
signatures = tuple(
(path, self._signature(path)) for path in self._watched(install, layout)
)
except _CannotValidate:
return install, layout
if layout.error == UNREADABLE or self._racy(signatures, self._wall_ns()):
return install, layout
stored = _Entry(key, home, install, layout, signatures, self._monotonic())
changed = False
with self._lock:
if self._generation == generation:
self._entry = stored
changed = layout != self._logged_layout
if changed:
self._logged_layout = layout
if changed:
logger.info(
f"LoraManager layout changed: {layout.source or 'no'} settings, "
f"{layout.mode} storage, {len(layout.media_roots)} media folders"
)
return install, layout
def clear(self) -> None:
"""Drop the entry; a load that started before this stores nothing."""
with self._lock:
self._entry = None
self._generation += 1
def _log_refused(self, layout: LoraManagerLayout) -> None:
"""One warning per default LoRA folder newly refused by the guard."""
refused = tuple(layout.refused_defaults)
with self._lock:
if refused == self._logged_refused:
return
new = [folder for folder in refused if folder not in self._logged_refused]
self._logged_refused = refused
for folder in new:
logger.warning(
f"ComfyUI LoRA folder '{folder}' is neither scanned nor served: "
"it is a filesystem root, the home folder or a folder that "
"contains it"
)
# ── validation ──────────────────────────────────────────────────
def _valid(self, entry: _Entry, key: str, home: str) -> bool:
if entry.key != key or entry.home != home:
return False
if self._monotonic() - entry.stored_at >= MAX_AGE_SECONDS:
return False
try:
for path, signature in entry.signatures:
if self._signature(path) != signature:
return False
except _CannotValidate:
return False
folders = entry.layout.media_roots
if entry.install is not None:
folders = (entry.install,) + tuple(folders)
for folder in folders:
try:
if _canonical(folder, self._pathmod) != folder:
return False
except (OSError, ValueError):
return False
return True
def _watched(self, install, layout) -> Tuple[str, ...]:
"""The plugin settings file, then the settings file, each once."""
if install is None:
return ()
paths = [self._pathmod.join(install, "settings.json")]
if layout.settings_path is not None:
paths.append(str(layout.settings_path))
watched, seen = [], set()
for path in paths:
identity = self._pathmod.normcase(self._pathmod.abspath(path))
if identity not in seen:
seen.add(identity)
watched.append(path)
return tuple(watched)
def _signature(self, path: str):
try:
result = self._stat(path)
except (FileNotFoundError, NotADirectoryError) as error:
# Windows reports an offline share as "not found": that is a
# temporary failure, not a missing file.
if getattr(error, "winerror", None) in _OFFLINE_SHARE_ERRORS:
raise _CannotValidate(path) from error
return MISSING
except (OSError, ValueError) as error:
raise _CannotValidate(path) from error
return (result.st_mtime_ns, result.st_size)
@staticmethod
def _racy(signatures, now_ns: int) -> bool:
threshold = now_ns - RACY_SECONDS * 10**9
return any(
signature is not MISSING and signature[0] > threshold
for _path, signature in signatures
)
+65 -14
View File
@@ -12,8 +12,10 @@ server share one definition:
plugin file, or the user config folder, in LoraManager's order), keeps only
the seven keys the import needs, and resolves the LoRA roots, the sidecar
root (centralized storage), the example-images root and the active
LoraManager library. Every folder from the settings passes the folder
guard (exists, not a filesystem root, not the home folder itself).
LoraManager library. Every folder from the settings, and each of
ComfyUI's default LoRA folders, passes the folder guard: it must exist and
must not be a filesystem root, the home folder or a folder that contains
it. Default folders the guard refuses are listed in ``refused_defaults``.
- ``discover_sidecars`` finds every sidecar whose model file exists inside a
LoRA root: beside the models in ``alongside`` mode, under the sidecar root
(matched to its model by ``file_path``) in ``centralized`` mode.
@@ -25,6 +27,7 @@ it is testable outside ComfyUI.
"""
import json
import logging
import os
import re
from dataclasses import dataclass, replace
@@ -98,7 +101,12 @@ NOT_OBJECT = "the settings file is not a JSON object"
TOO_LARGE = "the settings file is larger than 5 MB"
NO_PLATFORMDIRS = "platformdirs is not available"
NO_SIDECAR_ROOT = "the sidecar folder does not exist"
BAD_SIDECAR_ROOT = "the sidecar folder is a filesystem root or the home folder"
BAD_SIDECAR_ROOT = (
"the sidecar folder is a filesystem root, the home folder or a folder "
"that contains it"
)
logger = logging.getLogger("prompt_manager.lora_manager_layout")
@dataclass(frozen=True)
@@ -114,6 +122,9 @@ class LoraManagerLayout:
example_images_root: Optional[str]
library: str
media_roots: Tuple[str, ...]
# ComfyUI's default LoRA folders that exist but that the folder guard
# refused (canonical): the layout cache logs them (#362).
refused_defaults: Tuple[str, ...] = ()
@dataclass(frozen=True)
@@ -126,13 +137,20 @@ class Sidecar:
@dataclass(frozen=True)
class Discovery:
"""The sidecars kept, how many metadata files were seen, and those unread."""
"""The sidecars kept, every metadata file seen, and those unread."""
layout: LoraManagerLayout
sidecars: List[Sidecar]
scanned: int
# Every ``*.metadata.json`` seen, kept or skipped: resolved, once each,
# sorted by ``str(path)``. The import's fingerprint covers all of them.
scanned_files: Tuple[Path, ...]
unreadable: List[str]
@property
def scanned(self) -> int:
"""How many metadata files were seen."""
return len(self.scanned_files)
class _SettingsError(Exception):
"""The settings file cannot be used, for the reason it carries."""
@@ -315,6 +333,23 @@ def _guarded(paths, home: str) -> Tuple[str, ...]:
return tuple(kept)
def _refused(paths, home: str) -> Tuple[str, ...]:
"""The folders in ``paths`` that exist but that the guard refuses."""
refused = []
for path in paths:
absolute = _absolute(path)
if not absolute or not os.path.isdir(absolute):
continue
if _guard_folder(absolute, home) is None:
try:
canonical = _canonical(absolute)
except (OSError, ValueError):
continue
if canonical not in refused:
refused.append(canonical)
return tuple(refused)
def _example_images_root(
settings: Mapping, library: str, libraries: int, home: str
) -> Optional[str]:
@@ -351,7 +386,9 @@ def load_layout(
default LoRA roots, so serving and trigger words keep working.
"""
home = str(Path.home())
defaults = _guarded(default_roots(), home)
default_list = list(default_roots())
defaults = _guarded(default_list, home)
refused_defaults = _refused(default_list, home)
fallback = LoraManagerLayout(
settings_path=None,
source=None,
@@ -362,6 +399,7 @@ def load_layout(
example_images_root=None,
library="default",
media_roots=defaults,
refused_defaults=refused_defaults,
)
if not lora_manager_path:
return fallback
@@ -377,10 +415,14 @@ def load_layout(
except _SettingsError as error:
return replace(fallback, error=str(error))
settings = _settings_subset(raw or {})
return _layout_from_settings(settings, settings_path, source, defaults, home)
return _layout_from_settings(
settings, settings_path, source, defaults, home, refused_defaults
)
def _layout_from_settings(settings, settings_path, source, defaults, home):
def _layout_from_settings(
settings, settings_path, source, defaults, home, refused_defaults=()
):
library, entry, libraries = _active_library(settings)
lora_roots = _media_roots(
defaults,
@@ -400,6 +442,7 @@ def _layout_from_settings(settings, settings_path, source, defaults, home):
example_images_root=example_root,
library=library,
media_roots=_media_roots(lora_roots, [example_root]),
refused_defaults=refused_defaults,
)
if mode != MODE_CENTRALIZED:
return layout
@@ -482,13 +525,13 @@ def discover_sidecars(layout: LoraManagerLayout) -> Discovery:
cannot see every folder must not replace the collection with part of it.
"""
if layout.error is not None:
return Discovery(layout, [], 0, [])
return Discovery(layout, [], (), [])
if layout.mode == MODE_CENTRALIZED:
folders = [layout.sidecar_root]
else:
folders = list(layout.lora_roots)
seen = set()
scanned = 0
scanned: List[Path] = []
unreadable: List[str] = []
sidecars = []
for folder in folders:
@@ -498,12 +541,13 @@ def discover_sidecars(layout: LoraManagerLayout) -> Discovery:
if key in seen:
continue
seen.add(key)
scanned += 1
scanned.append(resolved)
model = _model_for(resolved, layout, unreadable)
if model is not None:
sidecars.append(Sidecar(resolved, Path(os.path.realpath(model))))
sidecars.sort(key=lambda sidecar: str(sidecar.metadata_path))
return Discovery(layout, sidecars, scanned, unreadable)
scanned.sort(key=str)
return Discovery(layout, sidecars, tuple(scanned), unreadable)
# ── images ──────────────────────────────────────────────────────────
@@ -526,8 +570,15 @@ def _example_files(folder: str) -> List[str]:
def _preview_candidates(sidecar: Sidecar, metadata: Dict) -> Iterator[str]:
"""The paths a preview is looked for at, in LoraManager's order."""
url = metadata.get("preview_url")
if isinstance(url, str) and url.strip() and os.path.isabs(url.strip()):
yield url.strip()
if isinstance(url, str) and url.strip():
if os.path.isabs(url.strip()):
yield url.strip()
else:
# LoraManager writes absolute paths; it would resolve a relative
# one against its working directory, not the sidecar folder.
logger.debug(
"Relative preview_url skipped in %s", sidecar.metadata_path.name
)
base = sidecar.metadata_path.name[: -len(METADATA_SUFFIX)]
for suffix in PREVIEW_EXTENSIONS:
yield str(sidecar.metadata_path.parent / f"{base}{suffix}")
+40 -6
View File
@@ -30,6 +30,7 @@ except ImportError:
from dataclasses import replace
from .lora_layout_cache import LayoutCache
from .lora_manager_layout import (
MODE_ALONGSIDE,
METADATA_SUFFIX,
@@ -295,10 +296,35 @@ def layout_for_install(install: Optional[str]) -> LoraManagerLayout:
def lora_manager_layout(lora_manager_path: Optional[str]) -> LoraManagerLayout:
"""Where LoraManager keeps its files (``py/lora_manager_layout.py``)."""
"""Where LoraManager keeps its files (``py/lora_manager_layout.py``).
A fresh read: the import review, Start, the legacy scan and the trigger
word cache use it. The image hot paths use ``cached_lora_manager_layout``.
"""
return layout_for_install(lora_manager_install(lora_manager_path))
# The image hot paths' layout cache (#362). The lambdas look the functions up
# when called, so a test that patches them here is seen by the cache.
_layout_cache = LayoutCache(
lambda path: lora_manager_install(path),
lambda install: layout_for_install(install),
)
def cached_lora_manager_layout(
configured: Optional[str],
) -> Tuple[Optional[str], LoraManagerLayout]:
"""``(install, layout)`` for the configured path, from the layout cache
while LoraManager's settings files are unchanged (``py/lora_layout_cache.py``)."""
return _layout_cache.get(configured or "")
def clear_layout_cache() -> None:
"""Drop the cached layout (after an import linked files in new folders)."""
_layout_cache.clear()
def _get_extra_lora_paths(comfyui_root: Path) -> List[Path]:
"""Parse extra_model_paths.yaml for additional LoRA directories."""
results = []
@@ -342,7 +368,7 @@ def read_lora_metadata(metadata_path: Path) -> Optional[Dict]:
try:
with open(metadata_path, "r", encoding="utf-8") as f:
return json.load(f)
except (json.JSONDecodeError, OSError) as e:
except (ValueError, OSError) as e: # JSONDecodeError, UnicodeDecodeError
logger.debug(f"Failed to read {metadata_path}: {e}")
return None
@@ -857,11 +883,13 @@ class TriggerWordCache:
"""
try:
new_cache = self._build(lora_manager_path)
except Exception as error:
except OSError as error:
logger.warning(f"Trigger word cache not refreshed: {error}")
with self._lock:
self._failed_at = time.monotonic()
return len(self._cache)
return self._failed()
except Exception:
# Not a disk error: keep the traceback (#362).
logger.exception("Trigger word cache not refreshed")
return self._failed()
with self._lock:
self._cache = new_cache
@@ -873,6 +901,12 @@ class TriggerWordCache:
)
return len(new_cache)
def _failed(self) -> int:
"""Keep the previous map and note when this load failed."""
with self._lock:
self._failed_at = time.monotonic()
return len(self._cache)
@staticmethod
def _build(lora_manager_path: str) -> Dict[str, List[str]]:
layout = lora_manager_layout(lora_manager_path)
+74
View File
@@ -0,0 +1,74 @@
"""The tagging configuration's constants and validators (#328).
Server-free on purpose: ``py.api`` must import without ComfyUI's ``server``
module, which ``py/config.py`` needs; config re-exports all of this.
"""
import math
from typing import Any
# The tagging engines (``py/autotag.py`` MODELS) and the options they take.
WD14_MODEL_TYPES = frozenset({"wd14-swinv2", "wd14-vit"})
LLM_MODEL_TYPES = frozenset({"gguf", "hf"})
MODEL_TYPES = WD14_MODEL_TYPES | LLM_MODEL_TYPES
MAX_INSTRUCTION_CHARS = 4000
MAX_INSTRUCTION_BYTES = 16 * 1024
# What a batch does with a prompt that already has tags (#328).
EXISTING_POLICIES = ("skip", "merge", "replace_after_review")
# The language-model tag template every page starts from until one is saved.
DEFAULT_TAG_TEMPLATE = (
"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."
)
def validate_threshold(name: str, value: Any) -> float:
"""A WD14 confidence threshold in [0, 1], or ValueError naming ``name``."""
if type(value) not in (int, float) or not math.isfinite(value):
raise ValueError(f"{name} must be a finite number")
if not 0 <= value <= 1:
raise ValueError(f"{name} must be between 0 and 1")
return float(value)
# Every page shows the thresholds in steps of 0.05 (the concept's System rows,
# the Tag image dialogs and the legacy sliders), so a saved one must be on that
# grid (#369). Per-request thresholds keep their own checks.
THRESHOLD_STEP = 0.05
_STEPS_PER_UNIT = round(1 / THRESHOLD_STEP)
# A computed float can miss the grid by a hair: (0.1 + 0.2) * 20 is 6.000000000000001.
_STEP_TOLERANCE = 1e-12 # in steps; refuses 0.35000000001
def validate_saved_threshold(name: str, value: Any) -> float:
"""A saved WD14 threshold: in [0, 1] and a multiple of ``THRESHOLD_STEP``,
returned exactly on the grid; else ValueError naming ``name``."""
value = validate_threshold(name, value)
steps = value * _STEPS_PER_UNIT
if abs(steps - round(steps)) > _STEP_TOLERANCE:
raise ValueError(f"{name} must be a multiple of {THRESHOLD_STEP}")
return round(steps) / _STEPS_PER_UNIT
def validate_template(value: Any, name: str = "instructions") -> str:
"""Non-empty language-model instructions within the character and byte
caps, or ValueError naming ``name``."""
if not isinstance(value, str) or not value.strip():
raise ValueError(f"{name} must be non-empty text")
if len(value) > MAX_INSTRUCTION_CHARS:
raise ValueError(f"{name} must be at most {MAX_INSTRUCTION_CHARS} characters")
try:
encoded = value.encode("utf-8")
except UnicodeEncodeError:
raise ValueError(f"{name} must be valid Unicode text") from None
if len(encoded) > MAX_INSTRUCTION_BYTES:
raise ValueError(f"{name} {'are' if name == 'instructions' else 'is'} too long")
return value
+27
View File
@@ -0,0 +1,27 @@
"""Final stops for the background threads a test starts (#301, #302).
Every real ``ImageGenerationHandler``, ``ImageMonitor`` or ``PromptTracker``
a test constructs ends with its final stop and an assertion, before the
test's folder and databases go: ``close()`` for a handler,
``stop_monitoring()`` for a monitor, ``shutdown()`` for a tracker.
"""
from utils.prompt_tracker import PromptTracker
def stop_tracker(testcase, tracker):
"""Shut a real tracker's cleanup thread down; mocks are left alone."""
if isinstance(tracker, PromptTracker):
testcase.assertTrue(tracker.shutdown(), "prompt cleanup thread still running")
def close_handler(testcase, handler):
"""Close ``handler`` for good, then stop the tracker it was built with."""
testcase.assertTrue(handler.close(), "image worker still running")
stop_tracker(testcase, getattr(handler, "prompt_tracker", None))
def stop_monitor(testcase, monitor):
"""Stop ``monitor`` (its observer and worker), then its tracker."""
testcase.assertTrue(monitor.stop_monitoring(), "image monitor still running")
stop_tracker(testcase, getattr(monitor, "prompt_tracker", None))
+283
View File
@@ -0,0 +1,283 @@
"""Run a browser-test fixture seeder off the live database, then publish once (#384).
The seeders make hundreds of small commits. On a slow CI disk the live
database's commits queue behind WAL checkpoint syncs and a reset took over
30 s. ``ScratchSeeder.run`` instead:
1. copies the live database into a scratch database (SQLite backup API) whose
connections never sync (a crash loses only the copy);
2. runs the unchanged seeder against an API proxy whose ``db`` is the scratch
database; everything else (files, thumbnails, caches, command state) is
the live API's, written exactly as before;
3. replaces the rows of every table the seed changed in the live database, in
one transaction on the live connection. Tables the seed left alone are not
touched, so state another fixture owns stays as it was.
A commit by any other live connection while the seed ran fails the reset
instead of being overwritten (``PRAGMA data_version``).
"""
from __future__ import annotations
import hashlib
import os
import shutil
import sqlite3
import tempfile
import weakref
from pathlib import Path
from typing import Any, Callable, TypeVar
T = TypeVar("T")
# Set to 1 to run seeders on the live database as before #384 (the equivalence
# proof, tests/browser/reset_equivalence.py, compares both ways).
SEED_ON_LIVE_ENV = "PM_TEST_SEED_ON_LIVE"
SCRATCH_SCHEMA = "reset_seed"
# API calls a seed makes that must see the published rows (see _SeedApi).
DEFERRED_CALLS = frozenset({"invalidate_gallery_cache"})
# Linux's shared-memory tmpfs (present in the Playwright image and on runners).
MEMORY_ROOT = Path("/dev/shm")
# The largest fixtures' scratch copy is about 8 MB (database plus a WAL that
# autocheckpoints near 4 MB). Docker's default /dev/shm is 64 MB, shared with
# the browsers and the other workers; with less free than this, the scratch
# copy goes on disk instead.
MEMORY_FREE_BYTES = 32 * 1024 * 1024
# Scratch folders are named <prefix><owner pid>-<random>.
SCRATCH_PREFIX = "pm-reset-"
def _quote(name: str) -> str:
return '"' + name.replace('"', '""') + '"'
def _tables(conn: sqlite3.Connection, schema: str = "main") -> list[str]:
rows = conn.execute(
f"SELECT name, sql FROM {schema}.sqlite_master WHERE type = 'table' "
"ORDER BY name"
).fetchall()
for name, sql in rows:
if sql and sql.lstrip().upper().startswith("CREATE VIRTUAL"):
raise NotImplementedError(f"virtual table {name} cannot be copied")
return [name for name, _sql in rows]
def _columns(conn: sqlite3.Connection, table: str, schema: str = "main") -> list[str]:
"""Writable columns (hidden and generated ones excluded), plus ``rowid``
when the table has a rowid that no INTEGER PRIMARY KEY column carries."""
info = conn.execute(f"PRAGMA {schema}.table_xinfo({_quote(table)})").fetchall()
columns = [row[1] for row in info if row[6] == 0]
keys = [row for row in info if row[5]]
sql = conn.execute(
f"SELECT sql FROM {schema}.sqlite_master WHERE type = 'table' AND name = ?",
(table,),
).fetchone()[0]
without_rowid = sql is not None and "WITHOUT ROWID" in sql.upper()
integer_key = len(keys) == 1 and str(keys[0][2]).upper() == "INTEGER"
if not without_rowid and not integer_key and table != "sqlite_sequence":
columns.insert(0, "rowid")
return columns
def table_digests(conn: sqlite3.Connection, schema: str = "main") -> dict[str, str]:
"""A digest of every table's rows (rowids included), by table name."""
result = {}
for table in _tables(conn, schema):
selected = ", ".join(
_quote(c) if c != "rowid" else "rowid"
for c in _columns(conn, table, schema)
)
digest = hashlib.sha256()
rows = conn.execute(f"SELECT {selected} FROM {schema}.{_quote(table)}")
for row in sorted(repr(tuple(row)) for row in rows):
digest.update(row.encode("utf-8", "surrogatepass"))
digest.update(b"\n")
result[table] = digest.hexdigest()
return result
def _process_alive(pid: int) -> bool:
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
except OSError:
return True # it exists but belongs to someone else
return True
def remove_orphaned_scratch(root: Path) -> None:
"""Remove scratch folders of servers that died without cleaning up.
A server killed outright (SIGKILL, a crash) never runs ``close``, and
tmpfs keeps its folder in memory until reboot. The folder name carries
the owner's process id; a folder whose process is gone is removed.
"""
if os.name != "posix":
return # os.kill(pid, 0) would terminate the process on Windows
for folder in root.glob(f"{SCRATCH_PREFIX}*-*"):
pid = folder.name[len(SCRATCH_PREFIX) :].split("-", 1)[0]
if pid.isdigit() and not _process_alive(int(pid)):
shutil.rmtree(folder, ignore_errors=True)
class _SeedApi:
"""The live API, except that ``db`` is the scratch database.
Cache drops the seed asks for (``DEFERRED_CALLS``) are recorded and made
after the publish: dropped earlier, a request in between would cache the
live database's old rows again.
"""
def __init__(self, api: Any, db: Any) -> None:
object.__setattr__(self, "_api", api)
object.__setattr__(self, "db", db)
object.__setattr__(self, "deferred", [])
def __getattr__(self, name: str) -> Any:
if name in DEFERRED_CALLS:
def record(*args: Any, **kwargs: Any) -> None:
self.deferred.append((name, args, kwargs))
return record
return getattr(self._api, name)
def __setattr__(self, name: str, value: Any) -> None:
setattr(self._api, name, value)
class ScratchSeeder:
"""Seed into a scratch copy and publish the changed tables in one commit."""
def __init__(
self,
api: Any,
directory: Path,
*,
on_live: bool = False,
memory_root: Path | None = MEMORY_ROOT,
) -> None:
self._api = api
self._fallback = Path(directory)
self._memory_root = memory_root
self._directory: Path | None = None
self._on_live = on_live
self._scratch = None
self._unsynced: "weakref.WeakSet[sqlite3.Connection]" = weakref.WeakSet()
def run(self, seed: Callable[[Any], T]) -> T:
"""Run ``seed(api)`` (blocking; call it in an executor thread)."""
if self._on_live:
return seed(self._api)
live = self._api.db.model.get_connection()
if live.in_transaction:
raise RuntimeError("the live connection has an open transaction")
version = live.execute("PRAGMA data_version").fetchone()[0]
scratch = self._scratch_database()
target = self._scratch_connection(scratch)
with target.using():
live.backup(target)
before = table_digests(target)
target_api = _SeedApi(self._api, scratch)
result = seed(target_api)
target = self._scratch_connection(scratch)
if target.in_transaction:
raise RuntimeError("the seed left a scratch transaction open")
after = table_digests(target)
changed = [table for table in after if after[table] != before.get(table)]
if set(before) != set(after):
raise RuntimeError("the seed changed the database schema")
if live.execute("PRAGMA data_version").fetchone()[0] != version:
raise RuntimeError("the live database changed while the seed ran")
self._publish(live, scratch.model.db_path, changed, version)
for name, args, kwargs in target_api.deferred:
getattr(self._api, name)(*args, **kwargs)
return result
def _scratch_database(self):
if self._scratch is None:
from database.operations import PromptDatabase
self._directory = self._scratch_directory()
self._scratch = PromptDatabase(str(self._directory / "seed.db"))
return self._scratch
def _scratch_connection(self, scratch) -> sqlite3.Connection:
"""The calling thread's scratch connection, set never to sync."""
conn = scratch.model.get_connection()
if conn not in self._unsynced:
conn.execute("PRAGMA synchronous = OFF")
self._unsynced.add(conn)
return conn
@staticmethod
def _publish(
live: sqlite3.Connection, path: str, tables: list[str], version: int
) -> None:
"""Replace ``tables`` with the scratch rows in one live transaction.
``version`` is the live ``data_version`` read before the seed. It is
checked again under the write lock: a writer whose transaction was
open during the seed commits before ``BEGIN IMMEDIATE`` returns, and
its rows must fail the reset instead of being overwritten.
"""
if not tables:
return
foreign_keys = live.execute("PRAGMA foreign_keys").fetchone()[0]
# Rows are copied as they are; no cascade may reach a table left alone.
live.execute("PRAGMA foreign_keys = OFF")
try:
live.execute(f"ATTACH DATABASE ? AS {SCRATCH_SCHEMA}", (path,))
try:
live.execute("BEGIN IMMEDIATE")
try:
if live.execute("PRAGMA data_version").fetchone()[0] != version:
raise RuntimeError(
"the live database changed while the seed ran"
)
for table in tables:
columns = ", ".join(
_quote(c) if c != "rowid" else "rowid"
for c in _columns(live, table)
)
name = _quote(table)
live.execute(f"DELETE FROM main.{name}")
live.execute(
f"INSERT INTO main.{name} ({columns}) "
f"SELECT {columns} FROM {SCRATCH_SCHEMA}.{name}"
)
live.commit()
except BaseException:
live.rollback()
raise
finally:
live.execute(f"DETACH DATABASE {SCRATCH_SCHEMA}")
finally:
live.execute(f"PRAGMA foreign_keys = {int(foreign_keys)}")
def close(self) -> None:
"""Close every scratch connection, then remove the scratch files."""
if self._on_live:
return
if self._scratch is not None:
self._scratch.close_all()
self._scratch = None
if self._directory is not None:
shutil.rmtree(self._directory, ignore_errors=True)
self._directory = None
def _scratch_directory(self) -> Path:
"""A folder in memory (tmpfs) where there is one, else ``directory``.
Even unsynced, scratch writes on the runner's disk wait whenever the
kernel is writing back the pages they rewrite.
"""
root = self._memory_root
if root is not None and root.is_dir() and os.access(root, os.W_OK):
remove_orphaned_scratch(root)
if shutil.disk_usage(root).free >= MEMORY_FREE_BYTES:
prefix = f"{SCRATCH_PREFIX}{os.getpid()}-"
return Path(tempfile.mkdtemp(prefix=prefix, dir=root))
self._fallback.mkdir(parents=True, exist_ok=True)
return self._fallback
+78
View File
@@ -0,0 +1,78 @@
"""Empty a browser-test media folder without waiting for the disk (#384).
A reset used to ``shutil.rmtree`` the output and input folders the last
fixture wrote. Unlinking a file whose pages are still being written back
waits for that write, so on a slow CI disk a reset stalled in ``rmtree``.
``replace_folder`` renames the old folder into ``<parent>/reset-trash`` (one
metadata change) and makes a new empty one; a background thread deletes the
trash, and ``close_trash`` stops that thread and removes what is left before
the server's temp root goes.
"""
from __future__ import annotations
import queue
import shutil
import threading
import uuid
from pathlib import Path
TRASH_NAME = "reset-trash"
class _Trash:
def __init__(self, directory: Path) -> None:
self.directory = directory
self._queue: "queue.SimpleQueue[Path | None]" = queue.SimpleQueue()
self._thread = threading.Thread(
target=self._work, name="pm-reset-trash", daemon=True
)
self._thread.start()
def _work(self) -> None:
while (folder := self._queue.get()) is not None:
shutil.rmtree(folder, ignore_errors=True)
def discard(self, folder: Path) -> None:
self.directory.mkdir(exist_ok=True)
moved = self.directory / f"{folder.name}-{uuid.uuid4().hex}"
folder.rename(moved)
self._queue.put(moved)
def close(self) -> None:
self._queue.put(None)
self._thread.join()
shutil.rmtree(self.directory, ignore_errors=True)
_trashes: dict[Path, _Trash] = {}
# Trash folders whose server has cleaned up: a reset still running then
# deletes in place instead of starting a thread nobody would stop.
_closed: set[Path] = set()
_lock = threading.Lock()
def replace_folder(folder: Path) -> None:
"""Make ``folder`` an empty folder; its old content is deleted later."""
folder = Path(folder)
if folder.exists():
directory = folder.parent / TRASH_NAME
with _lock:
trash = _trashes.get(directory)
if trash is None and directory not in _closed:
trash = _trashes[directory] = _Trash(directory)
if trash is not None:
trash.discard(folder)
if folder.exists():
shutil.rmtree(folder)
folder.mkdir()
def close_trash(parent: Path) -> None:
"""Stop deleting under ``parent``'s trash and remove it (blocking)."""
directory = Path(parent) / TRASH_NAME
with _lock:
_closed.add(directory)
trash = _trashes.pop(directory, None)
if trash is not None:
trash.close()
+152
View File
@@ -0,0 +1,152 @@
// #328: the legacy admin and gallery auto-tag tools keep working beside System › Auto-tagging. The model download is
// the `model_download` Activity job; the pages' stream attaches to it. No model is ever downloaded: the fixture's
// fake download waits until the test releases it (`/__test__/tagging/download`).
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
import { test, expect } from './fixtures.mjs';
const TITLE = 'Tagging model download · WD14 ViT';
async function post(request, url, data) {
await expect(await request.post(url, data === undefined ? {} : { data })).toBeOK();
}
async function downloads(request) {
const response = await request.get('/prompt_manager/jobs');
await expect(response).toBeOK();
return (await response.json()).jobs.filter(job => job.kind === 'model_download');
}
test.beforeEach(async ({ page, request, browserName }) => {
await startCoverage(page, browserName);
// The 3.2.0 update notice would cover the admin page's toolbar.
await page.addInitScript(() => { localStorage.setItem('pm_update_notice_dismissed', '3.2.0'); });
await post(request, '/__test__/jobs/restart');
await post(request, '/__test__/tagging/reset');
await post(request, '/__test__/tagging/models', { unavailable: ['wd14-vit'] });
await post(request, '/__test__/tagging/download', { mode: 'hold' });
});
test.afterEach(async ({ page, request, browserName }) => {
await saveCoverage(page, browserName);
// Releases a download a failed test left held, and puts the raising default back.
await post(request, '/__test__/tagging/reset');
await post(request, '/__test__/jobs/restart');
});
async function startDownload(page, request) {
await page.locator('#autoTagBtn').click();
const status = page.locator('#wd14VitModelStatus');
await status.getByRole('button', { name: 'Download', exact: true }).click();
const modal = page.locator('#autoTagDownloadModal');
await expect(modal).toBeVisible();
await expect(page.locator('#downloadStatus')).toHaveText('Downloading…');
await expect(page.locator('#downloadProgressBar')).toHaveClass(/animate-pulse/);
await expect(modal).toContainText('The download continues in Activity.');
await expect(page.locator('#cancelDownloadBtn')).toHaveText('Close');
const running = (await downloads(request)).filter(job => job.state === 'running');
expect(running.map(job => [job.title, job.cancellable])).toEqual([[TITLE, false]]);
return { modal, status };
}
test('admin download modal closes on complete and Activity lists the job without Cancel', async ({ page, request }) => {
await page.goto('/prompt_manager/admin');
const { modal, status } = await startDownload(page, request);
await page.goto('/prompt_manager/workspace#/activity');
const item = page.locator('#workspace-main section.surface > .activity-item')
.filter({ has: page.getByRole('link', { name: new RegExp(`${TITLE}$`) }) });
await expect(item.locator('.chip')).toHaveText('Running');
await expect(item.locator('progress')).not.toHaveAttribute('value');
await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
// Back on the admin page a second download of the same model attaches to the running job.
await page.goto('/prompt_manager/admin');
await startDownload(page, request);
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
await expect(modal).toBeHidden();
await expect(status).toContainText('Downloaded');
expect((await downloads(request)).map(job => [job.title, job.state])).toEqual([[TITLE, 'finished']]);
});
test('gallery download modal closes on complete', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
const { modal, status } = await startDownload(page, request);
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
await expect(modal).toBeHidden();
await expect(page.locator('#autoTagModal')).toBeVisible();
await expect(status).toContainText('Downloaded');
expect((await downloads(request)).map(job => job.state)).toEqual(['finished']);
});
test('a failed download closes the modal and says so', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
const { modal, status } = await startDownload(page, request);
await post(request, '/__test__/tagging/download', { mode: 'fail' });
await expect(modal).toBeHidden();
await expect(page.getByText('Download failed: Download failed')).toBeVisible();
await expect(status.getByRole('button', { name: 'Download', exact: true })).toBeVisible();
expect((await downloads(request)).map(job => job.state)).toEqual(['failed']);
});
// (wd14-vit is not downloaded in these tests: beforeEach marks it so the download tests can offer it.)
// The gallery's batch reader keyed on `data.status` values the server never sends, so its modal never moved or
// closed (#328 Problem 5). It now reads `type` frames as the admin page does.
async function watchProgress(page) {
await page.evaluate(() => {
window.progressSeen = [];
const node = document.getElementById('autoTagProgressPercent');
new MutationObserver(() => window.progressSeen.push(node.textContent))
.observe(node, { childList: true, characterData: true, subtree: true });
});
}
async function runGalleryBatch(page) {
await page.locator('#autoTagBtn').click();
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
await watchProgress(page);
const done = page.waitForResponse(response => response.url().includes('/prompt_manager/autotag/start'));
await page.locator('#startAutoTagBtn').click();
await done;
await expect(page.locator('#autoTagProgressModal')).toBeHidden();
await expect(page.getByText(/^Auto tagging complete! Applied tags to \d+ prompts\.$/)).toBeVisible();
}
test('gallery batch modal counts progress and closes on complete', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
await runGalleryBatch(page);
const seen = await page.evaluate(() => window.progressSeen);
// The fixture's prompts are already tagged, so the batch skips them; its progress frames still move the bar.
expect(seen).toContain('95%');
expect((await downloads(request)).length).toBe(0);
});
test('admin review opens the first image', async ({ page }) => {
await page.goto('/prompt_manager/admin');
await page.locator('#autoTagBtn').click();
await page.locator('input[name="autoTagModel"][value="wd14-swinv2"]').check();
await page.locator('#startReviewBtn').click();
// The first image's prompt is already tagged: confirm a new pass, then the review opens on that image.
await page.locator('#retagConfirmBtn').click();
await expect(page.locator('#autoTagReviewModal')).toBeVisible();
await expect(page.locator('#reviewImageName')).not.toHaveText(/^(image_name\.png|Unknown|)$/);
await expect(page.getByText('Failed to start review mode')).toHaveCount(0);
});
test('gallery unload works and is refused while the gate is held', async ({ page, request }) => {
await page.goto('/prompt_manager/gallery');
await runGalleryBatch(page); // keeps the model in memory
await page.locator('#autoTagBtn').click();
const unload = page.locator('#unloadModelBtn');
await expect(unload).toBeVisible();
// A held download owns the model gate: Unload is refused with the server's error.
const held = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'gguf' } });
expect(held.status()).toBe(202);
await unload.click();
await expect(page.getByText(/^Failed to unload model: A tagging model is busy/)).toBeVisible();
await post(request, '/__test__/tagging/download', { mode: 'succeed' });
await expect.poll(async () => (await downloads(request)).map(job => job.state)).toEqual(['finished']);
await unload.click();
await expect(page.getByText('WD14-SWINV2 model unloaded successfully')).toBeVisible();
});
+279
View File
@@ -0,0 +1,279 @@
"""Prove the scratch-seeded resets leave the state the old resets left (#384).
Run by hand on Linux (the browser server's platform):
python -m tests.browser.reset_equivalence
It starts the browser server three times: twice seeding on the live database
as before #384 (``PM_TEST_SEED_ON_LIVE=1``) and once with the scratch seeder.
Each server gets the same sequence: every reset kind of ``/__test__/reset``
and ``/__test__/concept/reset``, with other fixtures' resets and API writes
between them. After every reset it reads every table's rows (by rowid), the
output and input file trees (size, SHA-256, mtime) and the response body.
Values that differ between the two old runs are wall-clock values (created_at,
an unpinned concept "now", file mtimes); they are masked. Everything else must
be equal between the old runs and the new one. Exit status 0 means equal.
"""
from __future__ import annotations
import hashlib
import json
import os
import sqlite3
import subprocess
import sys
import tempfile
import time
import urllib.request
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
READY_TIMEOUT_S = 60
CLOCK = "<clock>"
RESETS = [
("/__test__/reset", {}),
("/__test__/reset", {"extra_images": 2}),
("/__test__/reset", {"extra_images": 108}),
("/__test__/reset", {"include_video": True}),
("/__test__/reset", {"prompt_inspector": True}),
("/__test__/reset", {"native_library": True}),
("/__test__/reset", {"native_library": True, "extra_prompts": 125}),
("/__test__/reset", {"native_library": True, "extra_variants": 41}),
("/__test__/reset", {"workspace_visual": True}),
("/__test__/reset", {"full_inspector": True}),
(
"/__test__/reset",
{"full_inspector": True, "include_video": True, "extra_images": 2},
),
("/__test__/reset", {"chronological_history": True}),
("/__test__/reset", {"chronological_history": True, "history_runs": 85}),
("/__test__/concept/reset", {}),
("/__test__/concept/reset", {"extra_images": 100}),
("/__test__/concept/reset", {"seedless_images": 2, "short_seedless": True}),
("/__test__/concept/reset", {"now": 1_780_000_000}),
]
# Writes by other fixtures and the API between resets; the resets must keep or
# replace them exactly as before.
DIRTY = [
("/__test__/tagging/reset", {}),
("/__test__/system/reset", None),
(
"/__test__/saved-views/raw",
{"page": "images", "name": "dirty", "state": "{}"},
),
]
class Server:
def __init__(self, on_live: bool) -> None:
self.tmp = tempfile.TemporaryDirectory(prefix="pm-equivalence-")
env = {**os.environ, "TMPDIR": self.tmp.name}
env.update(TEMP=self.tmp.name, TMP=self.tmp.name)
env.pop("PM_TEST_SEED_ON_LIVE", None)
if on_live:
env["PM_TEST_SEED_ON_LIVE"] = "1"
self.process = subprocess.Popen(
[
sys.executable,
"-m",
"tests.browser.server",
"--port",
"0",
"--exit-on-stdin-eof",
],
cwd=REPO_ROOT,
env=env,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
)
line = self.process.stdout.readline().decode()
if not line.startswith("PM_TEST_SERVER_PORT="):
raise RuntimeError(f"no port line: {line!r}")
self.base = f"http://127.0.0.1:{int(line.split('=', 1)[1])}"
deadline = time.monotonic() + READY_TIMEOUT_S
while True:
try:
self.post_json("GET", "/prompt_manager/test", None)
break
except OSError:
if time.monotonic() > deadline or self.process.poll() is not None:
raise
time.sleep(0.1)
(self.root,) = Path(self.tmp.name).glob("pm-browser-*")
def post_json(self, method: str, path: str, data) -> dict:
body = None if data is None else json.dumps(data).encode()
request = urllib.request.Request(
self.base + path,
data=body if method == "POST" else None,
method=method,
headers={"Content-Type": "application/json"} if body else {},
)
with urllib.request.urlopen(request, timeout=120) as response: # nosec B310
text = response.read().decode()
try:
return json.loads(text)
except ValueError:
return {"text": text[:200]}
def stop(self) -> None:
self.process.stdin.close()
self.process.wait(timeout=READY_TIMEOUT_S)
self.process.stdout.close()
self.tmp.cleanup()
def read_tables(path: Path) -> dict:
conn = sqlite3.connect(path.as_uri() + "?mode=ro", uri=True)
try:
tables = {}
names = [
row[0]
for row in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
)
]
for name in names:
quoted = '"' + name.replace('"', '""') + '"'
columns = [row[1] for row in conn.execute(f"PRAGMA table_info({quoted})")]
try:
cursor = conn.execute(f"SELECT rowid, * FROM {quoted} ORDER BY rowid")
columns = ["rowid", *columns]
except sqlite3.OperationalError:
cursor = conn.execute(f"SELECT * FROM {quoted} ORDER BY 1")
rows = [
[value.hex() if isinstance(value, bytes) else value for value in row]
for row in cursor
]
tables[name] = {"columns": columns, "rows": rows}
return tables
finally:
conn.close()
def read_files(root: Path) -> dict:
files = {}
for top in ("output", "input"):
for path in sorted((root / top).rglob("*")):
if path.is_file():
stat = path.stat()
files[path.relative_to(root).as_posix()] = [
stat.st_size,
hashlib.sha256(path.read_bytes()).hexdigest(),
stat.st_mtime_ns,
]
return files
def flatten(value, prefix="") -> dict:
if isinstance(value, dict):
out = {}
for key, item in value.items():
out.update(flatten(item, f"{prefix}/{key}"))
return out
if isinstance(value, list):
out = {f"{prefix}#len": len(value)}
for index, item in enumerate(value):
out.update(flatten(item, f"{prefix}[{index}]"))
return out
return {prefix: value}
def snapshot(server: Server, response: dict) -> dict:
flat = {}
for table, data in read_tables(server.root / "browser.db").items():
flat[f"table {table}#rows"] = len(data["rows"])
for index, row in enumerate(data["rows"]):
for column, value in zip(data["columns"], row):
flat[f"table {table}[{index}].{column}"] = value
for name, (size, digest, mtime) in read_files(server.root).items():
flat[f"file {name}"] = [size, digest]
flat[f"file {name}#mtime"] = mtime
flat["files#count"] = sum(
1 for key in flat if key.startswith("file ") and "#" not in key
)
for key, value in flatten(response).items():
flat[f"response{key}"] = value
# Each server has its own temp root and a random storage epoch; name them
# instead of masking every column that carries them.
epoch = flat.get("table execution_history_metadata[0].storage_epoch")
names = [(str(server.root), "<root>")]
if isinstance(epoch, str) and epoch:
names.append((epoch, "<epoch>"))
for key, value in flat.items():
if isinstance(value, str):
for text, name in names:
value = value.replace(text, name)
flat[key] = value
return flat
def run(on_live: bool) -> list[dict]:
server = Server(on_live)
try:
states = []
for index, (path, data) in enumerate(RESETS * 2):
dirty_path, dirty_data = DIRTY[index % len(DIRTY)]
server.post_json("POST", dirty_path, dirty_data)
started = time.monotonic()
response = server.post_json("POST", path, data)
elapsed = time.monotonic() - started
print(f" {'old' if on_live else 'new'} {path} {data}: {elapsed:.2f}s")
states.append(snapshot(server, response))
return states
finally:
server.stop()
def column_key(key: str) -> str:
"""``table t[3].created_at`` -> ``table t.created_at`` (mask whole columns)."""
if key.startswith("table ") and "[" in key:
head, _, tail = key.partition("[")
return head + tail.partition("]")[2]
return key
def main() -> int:
old_a, old_b, new = run(True), run(True), run(False)
failures = 0
masked_all = set()
for step, (path, data) in enumerate(RESETS * 2):
a, b, n = old_a[step], old_b[step], new[step]
masked = {
column_key(key)
for key in set(a) | set(b)
if a.get(key, CLOCK) != b.get(key, CLOCK)
}
masked_all |= masked
keys = set(a) | set(n)
diffs = [
key
for key in sorted(keys)
if column_key(key) not in masked
and a.get(key, "<missing>") != n.get(key, "<missing>")
]
tables = len(
{k.split("[")[0].split("#")[0] for k in a if k.startswith("table ")}
)
status = "EQUAL" if not diffs else f"DIFFERENT ({len(diffs)} values)"
print(
f"{status}: {path} {json.dumps(data)}: {tables} tables, "
f"{a['files#count']} files, {len(keys)} values, "
f"{len(masked)} masked columns"
)
for key in diffs[:20]:
old, now = a.get(key, "<missing>"), n.get(key, "<missing>")
print(f" {key}: old {old!r} new {now!r}")
failures += bool(diffs)
print("masked (differ between the two old runs):")
for key in sorted(masked_all):
print(f" {key}")
return 1 if failures else 0
if __name__ == "__main__":
sys.exit(main())
+65 -39
View File
@@ -19,8 +19,10 @@ from aiohttp.web_runner import GracefulExit
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from tests.browser.folder_trash import close_trash, replace_folder
from tests.browser.request_hold import register_request_holds
from tests.browser import slow_reset_watch
from tests.browser.fast_reset import SEED_ON_LIVE_ENV, ScratchSeeder
from tests.browser.workspace_activity_fixture import register_activity_commands
from tests.browser.workspace_system_fixture import register_system_commands
from tests.browser.workspace_integrations_fixture import (
@@ -72,6 +74,8 @@ class TestModel:
# Scored names the concept comparisons answer instead of the fixed list
# (``/__test__/tagging/suggestions``); None keeps the fixed list.
suggestions: list | None = None
# The engine's model table (``py/autotag.py`` MODELS), set at startup.
models_config: dict = {}
def get_models_status(self) -> dict:
return {
@@ -79,7 +83,7 @@ class TestModel:
"downloaded": name not in self.unavailable,
"loaded": self.loaded == name,
}
for name in ("gguf", "hf", "wd14-vit", "wd14-convnext", "wd14-swinv2")
for name in self.models_config
}
def is_model_loaded(self) -> bool:
@@ -207,6 +211,7 @@ def create_app(root: Path) -> web.Application:
from py import autotag
from utils.hashing import generate_prompt_hash
TestModel.models_config = autotag.MODELS
model = TestModel()
autotag.get_autotag_service = lambda: model
api = PromptManagerAPI()
@@ -222,6 +227,10 @@ def create_app(root: Path) -> web.Application:
register_integrations_commands(app, routes, api, root)
request_holds = register_request_holds(app, routes)
slow_reset_watch.install(app, api)
# Resets seed a scratch copy and publish it in one commit (#384).
seeder = ScratchSeeder(
api, root / "reset-scratch", on_live=os.environ.get(SEED_ON_LIVE_ENV) == "1"
)
async def canonical_asset(request: web.Request) -> web.Response:
filename, content_type = CANONICAL_ASSETS[request.path]
@@ -235,6 +244,7 @@ def create_app(root: Path) -> web.Application:
routes.get(path)(canonical_asset)
def seed(
target,
extra_images: int = 0,
include_video: bool = False,
prompt_inspector: bool = False,
@@ -242,7 +252,7 @@ def create_app(root: Path) -> web.Application:
extra_prompts: int = 0,
extra_variants: int = 0,
) -> None:
with api.db.model.get_connection() as conn:
with target.db.model.get_connection() as conn:
for table in (
"saved_views",
"image_favorites",
@@ -261,8 +271,7 @@ def create_app(root: Path) -> web.Application:
):
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
for index, (directory, text, tag) in enumerate(
[
("output", "Cat portrait", "cats"),
@@ -305,28 +314,30 @@ def create_app(root: Path) -> web.Application:
image.save(image_path, pnginfo=info)
stamp = time.time() - index * 10
os.utime(image_path, (stamp, stamp))
pid = api.db.save_prompt(
pid = target.db.save_prompt(
text, tags=[tag], prompt_hash=generate_prompt_hash(text)
)
api.db.link_image_to_prompt(pid, str(image_path))
target.db.link_image_to_prompt(pid, str(image_path))
if index == 0:
video_prompt_id = pid
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))
target.db.link_image_to_prompt(pid, str(extra))
if native_library:
positive = api.db.get_prompt_by_hash(generate_prompt_hash("Cat portrait"))
positive = target.db.get_prompt_by_hash(
generate_prompt_hash("Cat portrait")
)
positive_id = positive["id"]
api.db.update_prompt_metadata(
target.db.update_prompt_metadata(
positive_id, category="Animals", tags=["cats", "studio"]
)
variants = []
for index, negative in enumerate(
["blurry, watermark", "low quality", "low contrast", ""]
):
pair_id = api.db.save_prompt_pair(
pair_id = target.db.save_prompt_pair(
"Cat portrait",
negative,
title=f"Portrait variant {index + 1}",
@@ -334,28 +345,30 @@ def create_app(root: Path) -> web.Application:
tags=["studio"] if index == 0 else ["outdoors"],
)
variants.append(pair_id)
with api.db.model.get_connection() as conn:
with target.db.model.get_connection() as conn:
images = conn.execute(
"SELECT id FROM generated_images WHERE prompt_id=? ORDER BY id",
(positive_id,),
).fetchall()
for index, image in enumerate(images[:2]):
api.db.link_image_to_pair(
target.db.link_image_to_pair(
variants[index], image["id"], {"source": "fixture", "anchor": index}
)
api.db.link_image_to_pair(
target.db.link_image_to_pair(
variants[0], images[0]["id"], {"source": "fixture second evidence"}
)
api.db.record_prompt_pair_use(variants[0], times=3)
target.db.record_prompt_pair_use(variants[0], times=3)
for number in range(extra_variants):
api.db.save_prompt_pair("Cat portrait", f"Negative extra {number:04}")
target.db.save_prompt_pair(
"Cat portrait", f"Negative extra {number:04}"
)
legacy = root / "output" / "legacy-unassigned.png"
shutil.copyfile(root / "output" / "shared" / "same.png", legacy)
api.db.link_image_to_prompt(positive_id, str(legacy))
api.db.save_prompt("Dog landscape", category="Animals", tags=["dogs"])
api.db.save_prompt_pair("Dual role portrait", "Cat portrait")
target.db.link_image_to_prompt(positive_id, str(legacy))
target.db.save_prompt("Dog landscape", category="Animals", tags=["dogs"])
target.db.save_prompt_pair("Dual role portrait", "Cat portrait")
for number in range(extra_prompts):
api.db.save_prompt(
target.db.save_prompt(
f"Library extra {number:04}", category="Bulk", tags=["bulk"]
)
if prompt_inspector:
@@ -376,8 +389,8 @@ def create_app(root: Path) -> web.Application:
from py.api.thumbnail_settings import ThumbnailSettings
# 150 px previews in the default format (#317 thumbnail settings).
api._generate_gallery_thumbnails_sync(
api._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150)
target._generate_gallery_thumbnails_sync(
target._get_all_output_dirs(), ThumbnailSettings("webp", 80, 150)
)
if include_video:
# Seed after thumbnail generation to exercise a real unthumbnailed video.
@@ -385,8 +398,8 @@ def create_app(root: Path) -> web.Application:
shutil.copyfile(
Path(__file__).resolve().parents[1] / "fixtures" / "tiny.webm", video
)
api.db.link_image_to_prompt(video_prompt_id, str(video))
api.invalidate_gallery_cache()
target.db.link_image_to_prompt(video_prompt_id, str(video))
target.invalidate_gallery_cache()
history_fixture = None
@@ -431,7 +444,10 @@ def create_app(root: Path) -> web.Application:
raise web.HTTPBadRequest()
if chronological_history:
manifest = await api._run_in_executor(
lambda: seed_workspace_history(api, root, history_runs=history_runs)
seeder.run,
lambda target: seed_workspace_history(
target, root, history_runs=history_runs
),
)
history_fixture = manifest
return web.json_response({"success": True, **manifest})
@@ -489,25 +505,31 @@ def create_app(root: Path) -> web.Application:
raise web.HTTPBadRequest()
if workspace_visual:
history_fixture = None
await api._run_in_executor(seed_workspace_visual, api, root)
await api._run_in_executor(
seeder.run, lambda target: seed_workspace_visual(target, root)
)
return web.json_response({"success": True})
if full_inspector:
history_fixture = None
await api._run_in_executor(
lambda: seed_workspace_inspector(
api, root, extra_images=extra, include_video=include_video
)
seeder.run,
lambda target: seed_workspace_inspector(
target, root, extra_images=extra, include_video=include_video
),
)
return web.json_response({"success": True})
history_fixture = None
await api._run_in_executor(
seed,
extra,
include_video,
prompt_inspector,
native_library,
extra_prompts,
extra_variants,
seeder.run,
lambda target: seed(
target,
extra,
include_video,
prompt_inspector,
native_library,
extra_prompts,
extra_variants,
),
)
return web.json_response({"success": True})
@@ -547,14 +569,15 @@ def create_app(root: Path) -> web.Application:
reveal_stub.reset()
history_fixture = None
manifest = await api._run_in_executor(
lambda: seed_workspace_concept(
api,
seeder.run,
lambda target: seed_workspace_concept(
target,
root,
extra_images=data.get("extra_images", 0),
seedless_images=data.get("seedless_images", 0),
short_seedless=data.get("short_seedless", False),
now=data.get("now"),
)
),
)
return web.json_response({"success": True, **manifest})
@@ -622,9 +645,12 @@ def create_app(root: Path) -> web.Application:
async def cleanup(application: web.Application) -> None:
# A scan of an ambiguous fixture file captures it for Needs review.
api.close_annotation_captures()
seeder.close()
api.db.close_all()
# Old media folders still waiting for deletion go before the temp root.
close_trash(root)
seed()
seed(api)
# Open Needs review change streams end first, as on the real host.
app.on_shutdown.append(api.close_review_events)
app.on_cleanup.append(cleanup)
@@ -0,0 +1,385 @@
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
import { test, expect } from './fixtures.mjs';
import { SIZES, THEMES, seedConcept, openConcept, openProduction, diffInventory, mutate } from './concept-inventory.mjs';
// #328 System S11 Task 8: System › Auto-tagging (app.js renderSystem 685–686) and the Tag output images review against
// the immutable concept, on the concept-shaped fixture with the fake tagging model (every engine downloaded, none
// loaded). Each concept measurement is mapped only by the recorded deviations (D132, D133, D134), each commented where
// it applies; after that the difference must be empty, and a mutated label, option, order, class, select width or
// token must be noticed. The concept has no batch review, so the Tag output images review is compared with the
// concept's own record-changing review, `maintenancePlan` (953–958), whose shell, grid, note and actions it shares.
test.use({ locale: 'en-US', timezoneId: 'UTC' });
test.beforeEach(async ({ page, request, browserName }) => {
await startCoverage(page, browserName);
// The tagging reset restores the fake model (nothing loaded, every engine on disk) and the saved tagging defaults;
// the concept fixture then replaces its rows and files.
await expect(await request.post('/__test__/tagging/reset', { data: {} })).toBeOK();
await seedConcept(request);
await expect(await request.post('/__test__/system/reset')).toBeOK();
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
page.conceptErrors = [];
page.on('pageerror', error => page.conceptErrors.push(error.message));
});
test.afterEach(async ({ page, request, browserName }) => {
await saveCoverage(page, browserName);
await request.post('/__test__/jobs/restart');
await request.post('/__test__/system/reset');
await request.post('/__test__/tagging/reset', { data: {} });
expect(page.conceptErrors).toEqual([]);
});
const PREVIEW_URL = '/prompt_manager/autotag/batch/preview';
const MODEL_OPTIONS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB'];
const SUGGESTIONS = 'From 1 to 100: the most suggestions Generate shows for one image.';
const BATCH = 'Add tags to each prompt from its linked output images with the model above. Review the scope first.';
/** Runs in the page: Auto-tagging's surface under `mainSelector`: rows, field, actions, operation, geometry and tokens. */
function measureTagging(mainSelector) {
const layout = document.querySelector(`${mainSelector} .settings-layout`);
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
const text = node => node?.textContent.trim() ?? null;
const control = row => {
const select = row.querySelector(':scope > select');
if (select) return { kind: 'select', id: select.id, options: [...select.options].map(option => [option.textContent, option.disabled]),
selected: select.selectedOptions[0]?.textContent ?? null, disabled: select.disabled };
const input = row.querySelector(':scope > input');
return input ? { kind: `input.${input.type}`, id: input.id, value: input.value, disabled: input.disabled } : null;
};
const surface = layout.querySelector(':scope > section.surface');
const firstRow = surface.querySelector('.setting-row');
const field = surface.querySelector(':scope > .field');
const actions = surface.querySelector(':scope > .row-actions');
return {
layout: [...layout.children].map(name),
tree: [...surface.children].map(name),
heading: text(surface.querySelector(':scope > h2')),
intro: text(surface.querySelector(':scope > p')),
// Per row: whether it may wrap, and whether its control sits under the text (a wrapped row) or beside it.
rows: [...surface.querySelectorAll(':scope > .setting-row')].map(row => ({ tree: [...row.children].map(name),
text: [...row.firstElementChild.children].map(node => [name(node), text(node), node.localName === 'label' ? node.htmlFor : null]),
control: control(row), flexWrap: getComputedStyle(row).flexWrap,
controlBelow: row.lastElementChild.getBoundingClientRect().top >= row.firstElementChild.getBoundingClientRect().bottom - 0.5 })),
field: { tree: [...field.children].map(name), label: [text(field.querySelector('label')), field.querySelector('label').htmlFor],
textarea: { id: field.querySelector('textarea').id, value: field.querySelector('textarea').value,
...css(field.querySelector('textarea'), 'fontSize', 'minHeight', 'backgroundColor', 'color', 'borderTopColor') },
...css(field, 'display', 'rowGap', 'fontSize', 'marginTop', 'marginBottom') },
actions: { buttons: [...actions.children].map(item => [text(item), item.className]),
...css(actions, 'display', 'flexWrap', 'columnGap', 'justifyContent'),
primary: css(actions.querySelector('.ca-primary'), 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft'),
plain: css(actions.querySelector('button:not(.ca-primary)'), 'backgroundColor', 'color', 'fontSize', 'borderTopColor') },
operations: [...surface.querySelectorAll(':scope > .operation')].map(row => ({ tree: [...row.children].map(name),
title: text(row.querySelector('h3')), description: text(row.querySelector('p')),
actions: [...row.querySelectorAll('.row-actions > button')].map(item => [text(item), item.className]) })),
surface: css(surface, 'paddingTop', 'paddingLeft', 'backgroundColor', 'borderTopWidth', 'borderTopColor', 'borderTopLeftRadius'),
row: css(firstRow, 'display', 'alignItems', 'justifyContent', 'columnGap', 'paddingTop', 'paddingBottom', 'borderTopWidth',
'borderTopColor', 'marginTop'),
label: css(firstRow.querySelector('label'), 'fontSize', 'fontWeight', 'color'),
description: css(firstRow.querySelector('p'), 'fontSize', 'color', 'marginTop', 'maxWidth'),
select: css(firstRow.querySelector('select'), 'fontSize', 'minWidth', 'backgroundColor', 'color', 'borderTopColor'),
number: css(surface.querySelector('.setting-row input[type="number"]'), 'fontSize', 'minWidth', 'maxWidth', 'backgroundColor',
'color', 'borderTopColor'),
// Whether each select is as wide as the browser lays out a copy holding only its chosen option, i.e. it shows
// that option whole (the browser's own select sizing, arrow included).
selectsFit: [...surface.querySelectorAll('.setting-row > select')].map(select => {
const copy = select.cloneNode(false);
copy.removeAttribute('id');
copy.append(new Option(select.selectedOptions[0]?.textContent ?? ''));
// Important inline values: the phone cap is an !important rule (concept.css), which a plain inline style loses to.
for (const [key, value] of [['position', 'absolute'], ['visibility', 'hidden'], ['width', 'max-content'], ['min-width', '0'],
['max-width', 'none'], ['flex', 'none']]) copy.style.setProperty(key, value, 'important');
select.parentElement.append(copy);
const needed = copy.getBoundingClientRect().width;
copy.remove();
return select.getBoundingClientRect().width + 0.5 >= needed;
}),
pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth,
};
}
/**
* The concept's Auto-tagging with the recorded deviations applied, for the fixture (all engines on disk, none loaded).
* `phone` is a viewport at the concept's 700px breakpoint or below.
*/
function mapTaggingConcept(reference, { template, phone }) {
const [model, existing, general, character] = reference.rows;
const [preview] = reference.actions.buttons;
return {
...reference,
// D133 adds a row after the thresholds; D134 adds the operation after the actions.
tree: [...reference.tree.slice(0, 6), 'div.setting-row', ...reference.tree.slice(6), 'div.operation'],
rows: [
// D132 (OQ5, OQ6): the four real engines with their download size, and the chosen one's live status.
{ ...model, text: [model.text[0], ['p', 'Downloaded. Not loaded.', null]],
control: { ...model.control, options: MODEL_OPTIONS.map(label => [label, false]), selected: MODEL_OPTIONS[0] },
// D132: at phone width the row wraps, so the engine select sits whole under its text.
...(phone ? { flexWrap: 'wrap', controlBelow: true } : {}) },
existing, general, character,
// D133: the persisted suggestion count, on the concept's number row.
{ ...general, text: [['label', 'Suggestions per image', 'tag-suggestions'], ['p', SUGGESTIONS, null]],
control: { ...general.control, id: 'tag-suggestions', value: '10' } },
],
// The saved template (the server's default until one is saved) replaces the concept's sample text.
field: { ...reference.field, textarea: { ...reference.field.textarea, value: template } },
// D132: Download (for an engine already on disk, "Download again") and Unload replace the two simulated jobs.
actions: { ...reference.actions, buttons: [preview, ['Download again', ''], ['Unload', '']] },
// D134: the prompt-tag batch, reviewed first.
operations: [{ tree: ['div', 'div.row-actions'], title: 'Tag output images', description: BATCH, actions: [['Review', '']] }],
// D132: the engine select shows its whole chosen option at every width (at 390 the concept's 120px cap cuts it;
// `#tag-model` alone may grow there). Existing tags keeps the concept's own fit, cut or not.
selectsFit: reference.selectsFit.map((fit, index) => (index === 0 ? true : fit)),
};
}
async function conceptTagging(page, theme) {
await openConcept(page, 'system', theme);
await page.locator('#main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
await expect(page.locator('#main #tag-template')).toBeVisible();
await page.mouse.move(0, 0);
return page.evaluate(measureTagging, '#main');
}
async function productionTagging(page, theme) {
await openProduction(page, 'system/tagging', theme);
// Compare once the settings and the model status answered.
await expect(page.locator('#workspace-main #tag-model')).toBeEnabled();
await expect(page.locator('#workspace-main .setting-row').first().locator('p')).toHaveText('Downloaded. Not loaded.');
await page.mouse.move(0, 0);
return page.evaluate(measureTagging, '#workspace-main');
}
async function savedTemplate(request, page) {
const response = await request.get('/prompt_manager/settings');
await expect(response).toBeOK();
return { template: (await response.json()).settings.autotag_template, phone: page.viewportSize().width <= 700 };
}
for (const theme of THEMES) for (const [width, height] of SIZES) {
test(`auto-tagging matches the concept ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => {
await page.setViewportSize({ width, height });
const concept = await conceptTagging(page, theme);
await page.screenshot({ path: testInfo.outputPath(`concept-tagging-${theme}-${width}.png`), fullPage: true });
const production = await productionTagging(page, theme);
await page.screenshot({ path: testInfo.outputPath(`production-tagging-${theme}-${width}.png`), fullPage: true });
expect(diffInventory(mapTaggingConcept(concept, await savedTemplate(request, page)), production)).toEqual([]);
expect(production.pageScroll).toBe(false);
});
}
// D132 below the concept's smallest size: on a 320px phone the engine select still shows its whole option.
for (const theme of THEMES) test(`auto-tagging engine select fits at 320 wide ${theme}`, async ({ page }) => {
await page.setViewportSize({ width: 320, height: 700 });
const production = await productionTagging(page, theme);
for (const model of ['hf', 'gguf', 'wd14-swinv2']) {
await page.locator('#workspace-main #tag-model').selectOption(model);
expect((await page.evaluate(measureTagging, '#workspace-main')).selectsFit[0], model).toBe(true);
}
expect(production.pageScroll).toBe(false);
});
test('auto-tagging comparison detects mutations', async ({ page, request }) => {
await page.setViewportSize({ width: 1600, height: 1100 });
const concept = await conceptTagging(page, 'dark');
const saved = await savedTemplate(request, page);
const expected = mapTaggingConcept(concept, saved);
const production = await productionTagging(page, 'dark');
expect(diffInventory(expected, production)).toEqual([]);
// Unmapped, the section differs only where D132, D133 and D134 are recorded (and the saved template's text).
expect(diffInventory(concept, production).map(item => item.path.split('.').slice(0, 3).join('.'))
.filter((path, index, all) => all.indexOf(path) === index).sort()).toEqual([
'actions.buttons.1', 'actions.buttons.2', 'field.textarea.value', 'operations.0.actions', 'operations.0.description',
'operations.0.title', 'operations.0.tree', 'operations.length', 'rows.0.control', 'rows.0.text', 'rows.4.control', 'rows.4.controlBelow',
'rows.4.flexWrap', 'rows.4.text', 'rows.4.tree', 'rows.length', 'tree.6', 'tree.7', 'tree.8', 'tree.9', 'tree.length'].sort());
const measure = () => page.evaluate(measureTagging, '#workspace-main');
const paths = async () => diffInventory(expected, await measure()).map(item => item.path);
// A removed row.
const undoRow = await mutate(page, { remove: ['#workspace-main section.surface > .setting-row:nth-of-type(2)'] });
expect(await paths()).toEqual(expect.arrayContaining(['rows.length', 'tree.length']));
await undoRow();
// A renamed label.
await page.evaluate(() => { document.querySelector('#workspace-main label[for="tag-existing"]').textContent = 'Existing'; });
expect(await paths()).toEqual(['rows.1.text.0.1']);
await page.evaluate(() => { document.querySelector('#workspace-main label[for="tag-existing"]').textContent = 'Existing tags'; });
// A reordered option.
await page.evaluate(() => { const select = document.querySelector('#workspace-main #tag-existing'); select.append(select.options[0]); });
expect(await paths()).toEqual(expect.arrayContaining(['rows.1.control.options.0.0']));
await page.evaluate(() => { const select = document.querySelector('#workspace-main #tag-existing'); select.prepend(select.options[2]); });
// ca-primary moved to another action.
await page.evaluate(() => {
const [first, second] = document.querySelectorAll('#workspace-main section.surface > .row-actions > button');
first.classList.remove('ca-primary');
second.classList.add('ca-primary');
});
expect(await paths()).toEqual(expect.arrayContaining(['actions.buttons.0.1', 'actions.buttons.1.1']));
await page.evaluate(() => {
const [first, second] = document.querySelectorAll('#workspace-main section.surface > .row-actions > button');
second.classList.remove('ca-primary');
first.classList.add('ca-primary');
});
// A select squeezed until its chosen option is cut.
const undoSqueeze = await mutate(page, { css: '#workspace-main #tag-model { max-width: 60px !important; min-width: 0 !important; }' });
expect(await paths()).toEqual(expect.arrayContaining(['selectsFit.0']));
await undoSqueeze();
// A row that wraps where the concept's does not.
const undoWrap = await mutate(page, { css: '#workspace-main .setting-row:has(> #tag-existing) { flex-wrap: wrap; }' });
expect(await paths()).toEqual(['rows.1.flexWrap']);
await undoWrap();
// A swapped token.
const undoToken = await mutate(page, { css: '#workspace-main .settings-layout { --line: rgb(255, 0, 0); }' });
expect(await paths()).toEqual(expect.arrayContaining(['row.borderTopColor']));
await undoToken();
expect(diffInventory(expected, await measure())).toEqual([]);
});
// D132 at phone width, unmapped: only the engine row wraps; Existing tags keeps the concept's cut select.
test('auto-tagging at 390 wraps only the engine row', async ({ page, request }) => {
await page.setViewportSize({ width: 390, height: 844 });
const concept = await conceptTagging(page, 'dark');
const production = await productionTagging(page, 'dark');
expect(diffInventory(mapTaggingConcept(concept, await savedTemplate(request, page)), production)).toEqual([]);
expect(production.rows.map(row => row.flexWrap)).toEqual(['wrap', 'nowrap', 'nowrap', 'nowrap', 'nowrap']);
expect(production.selectsFit[1]).toBe(concept.selectsFit[1]);
});
/** Runs in the page: the open review dialog: its title, body and the concept's tokens. */
function measureReview() {
const dialog = document.querySelector('dialog#modal[open]');
const css = (node, ...keys) => node ? Object.fromEntries(keys.map(key => [key, getComputedStyle(node)[key]])) : null;
const name = node => [node.localName, ...[...node.classList].sort()].join('.');
const text = node => node?.textContent.trim() ?? null;
const body = dialog.querySelector('#modal-body');
const grid = body.querySelector(':scope > dl.meta-grid');
const note = [...body.querySelectorAll(':scope > p.small')].at(-1);
const actions = body.querySelector(':scope > .dialog-actions');
return {
title: text(dialog.querySelector('#modal-title')),
tree: [...body.children].map(name),
heading: { text: text(body.querySelector(':scope > h3')),
...css(body.querySelector(':scope > h3'), 'fontSize', 'fontWeight', 'color', 'marginBottom') },
grid: { rows: [...grid.children].map(row => [text(row.querySelector('dt')), text(row.querySelector('dd'))]),
...css(grid, 'display', 'gridTemplateColumns', 'columnGap', 'rowGap'),
dt: css(grid.querySelector('dt'), 'fontSize', 'color', 'textTransform'), dd: css(grid.querySelector('dd'), 'fontSize', 'color') },
backup: body.querySelector(':scope > .notice') && { tree: [...body.querySelector(':scope > .notice').children].map(name),
text: [...body.querySelectorAll(':scope > .notice > div:first-child > *')].map(node => [node.localName, text(node)]),
buttons: [...body.querySelectorAll(':scope > .notice button')].map(text) },
note: { text: text(note), ...css(note, 'fontSize', 'color', 'marginTop') },
actions: { tree: [...actions.children].map(name), ...css(actions, 'display', 'justifyContent', 'columnGap', 'marginTop'),
buttons: [...actions.children].map(item => ({ text: text(item), className: item.className,
...css(item, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor') })) },
};
}
const number = value => value.toLocaleString('en-US');
/**
* D134: the server's batch preview in the concept's review: "Review tag output images" with "Tag output images" as its
* `h3`, the model, policy, counts and settings in the concept's `dl.meta-grid` (Backup in the concept's wording), the
* upper-bound sentence in the concept's `p.small`, the D55 backup box and Cancel / "Start tagging".
*/
function mapBatchConcept(reference, plan) {
const [cancel, start] = reference.actions.buttons;
return {
...reference,
title: 'Review tag output images',
tree: ['h3', 'dl.meta-grid', 'p.small', 'div.notice', 'div.dialog-actions'],
heading: { ...reference.heading, text: 'Tag output images' },
grid: { ...reference.grid, rows: [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'],
['Images linked', number(plan.images)], ['Prompts', number(plan.prompts)], ['Already tagged', number(plan.already_tagged)],
['Prompts to tag', `up to ${number(plan.to_tag)}`], ['Settings', 'General 0.35 · Character 0.85'],
['Backup', 'Create and verify before changes']] },
backup: { tree: ['div', 'div.row-actions'], text: [['h3', 'Back up the database first (recommended)'],
['p', 'Downloads a copy of the prompts database before this operation changes any records.']], buttons: ['Back up now'] },
note: { ...reference.note, text: 'Nothing is written until you start. Missing files and files outside the output folders '
+ 'are counted, so the prompts to tag are an upper bound; Activity reports what was tagged.' },
actions: { ...reference.actions, buttons: [cancel, { ...start, text: 'Start tagging' }] },
};
}
async function conceptReview(page, theme) {
await openConcept(page, 'system', theme);
await page.locator('#main .section-nav').getByRole('button', { name: 'Maintenance', exact: true }).click();
await page.locator('#main [data-plan="duplicate prompts"]').click();
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review maintenance plan');
await page.mouse.move(0, 0);
return page.evaluate(measureReview);
}
async function productionReview(page, theme) {
await productionTagging(page, theme);
// Every concept-fixture prompt with an image is already tagged, so under the saved Skip there is nothing to start
// (the Cancel-only review). Merge reviews a batch with work: the full review the concept's shell is compared with.
await page.locator('#workspace-main #tag-existing').selectOption('merge');
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click();
const plan = await (await answer).json();
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review tag output images');
await page.mouse.move(0, 0);
return { plan, production: await page.evaluate(measureReview) };
}
for (const theme of THEMES) for (const [width, height] of SIZES) {
test(`tag output images review matches the concept ${theme} ${width}x${height}`, async ({ page }, testInfo) => {
await page.setViewportSize({ width, height });
const concept = await conceptReview(page, theme);
expect(concept.tree).toEqual(['p.demo-note', 'h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions']);
await page.screenshot({ path: testInfo.outputPath(`concept-plan-review-${theme}-${width}.png`) });
const { plan, production } = await productionReview(page, theme);
// Under Merge every prompt with a linked image is tagged again: the batch has work to review.
expect(plan.to_tag).toBeGreaterThan(0);
await page.screenshot({ path: testInfo.outputPath(`production-tagging-review-${theme}-${width}.png`) });
expect(diffInventory(mapBatchConcept(concept, plan), production)).toEqual([]);
});
}
/**
* D134 (CUJ 5): Replace after review never runs as a batch: the same review shell with the counts, the reason as its
* `p.small`, no backup box and only Cancel (the shared close-only review Integrations also uses).
*/
function mapCancelOnlyConcept(reference, plan) {
const [cancel] = reference.actions.buttons;
const mapped = mapBatchConcept(reference, plan);
return {
...mapped,
tree: ['h3', 'dl.meta-grid', 'p.small', 'div.dialog-actions'],
grid: { ...mapped.grid, rows: mapped.grid.rows.map(([term, value]) => [term, term === 'Existing tags' ? 'Replace after review' : value]) },
backup: null,
note: { ...reference.note, text: 'Replacing tags needs a review of each image. Use Preview tagging review, or choose Skip or '
+ 'Merge for a batch run.' },
actions: { ...reference.actions, tree: [reference.actions.tree[0]], buttons: [cancel] },
};
}
test('the Cancel-only review matches the concept review shell', async ({ page }) => {
await page.setViewportSize({ width: 1600, height: 1100 });
const concept = await conceptReview(page, 'dark');
await productionTagging(page, 'dark');
await page.locator('#workspace-main #tag-existing').selectOption('replace_after_review');
const answer = page.waitForResponse(response => response.url().endsWith(PREVIEW_URL));
await page.locator('#workspace-main .operation').getByRole('button', { name: 'Review', exact: true }).click();
const plan = await (await answer).json();
expect(plan.review_required).toBe(true);
await expect(page.locator('dialog#modal[open] #modal-title')).toHaveText('Review tag output images');
await page.mouse.move(0, 0);
expect(diffInventory(mapCancelOnlyConcept(concept, plan), await page.evaluate(measureReview))).toEqual([]);
});
test('tag output images review comparison detects mutations', async ({ page }) => {
await page.setViewportSize({ width: 1600, height: 1100 });
const concept = await conceptReview(page, 'dark');
const { plan } = await productionReview(page, 'dark');
const expected = mapBatchConcept(concept, plan);
const paths = async () => diffInventory(expected, await page.evaluate(measureReview)).map(item => item.path);
expect(await paths()).toEqual([]);
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'To tag'; });
expect(await paths()).toEqual(['grid.rows.5.0']);
await page.evaluate(() => { document.querySelectorAll('dialog#modal[open] dl.meta-grid dt')[5].textContent = 'Prompts to tag'; });
const undoGrid = await mutate(page, { css: 'dialog#modal dl.meta-grid { column-gap: 3px; }' });
expect(await paths()).toEqual(expect.arrayContaining(['grid.columnGap']));
await undoGrid();
const undoBox = await mutate(page, { remove: ['dialog#modal[open] #modal-body > .notice'] });
expect(await paths()).toEqual(expect.arrayContaining(['backup', 'tree.length']));
await undoBox();
expect(await paths()).toEqual([]);
});
+1 -1
View File
@@ -236,7 +236,7 @@ test('paired Tag dialog captures', async ({ page, request }, testInfo) => {
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-swinv2'] } })).toBeOK();
await openTagDialog(page, 'production', 'dark');
await generateWith(page, 'production').catch(() => {});
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.');
await expect(status()).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
await shot('production-unavailable');
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: [] } })).toBeOK();
+30
View File
@@ -1088,6 +1088,36 @@ test('a held legacy job and the health check show no tiles, no Results and no ma
await expect(toast(page)).toHaveText('Health check finished.');
});
// #328: a tagging model download is a job of its own kind with no byte progress and no Cancel; Activity shows it
// through the generic held view (its title, status line and an indeterminate bar), then its end toast.
test('a held model download shows its title, an indeterminate bar and no Cancel; its end toast follows', async ({ page, request }) => {
await page.setViewportSize({ width: 1600, height: 1100 });
const title = 'Tagging model download · WD14 ViT';
await expect(await request.post('/__test__/tagging/reset')).toBeOK();
await expect(await request.post('/__test__/tagging/download', { data: { mode: 'hold' } })).toBeOK();
try {
const response = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'wd14-vit' } });
expect(response.status()).toBe(202);
const job = (await response.json()).job;
expect([job.kind, job.title, job.cancellable]).toEqual(['model_download', title, false]);
await page.goto(`${workspace}#/activity`);
const item = items(page).filter({ has: page.getByRole('link', { name: new RegExp(`${title}$`) }) });
await expect(item.locator('progress')).not.toHaveAttribute('value');
await expect(item.getByRole('button', { name: 'Cancel job', exact: true })).toHaveCount(0);
await item.getByRole('link').click();
await expect(heading(page).locator('h1')).toHaveText(title);
await expect(progressSection(page).locator('.panel-heading > span.small')).toHaveText('Downloading WD14 ViT (~400 MB)');
await expect(progressSection(page).locator('progress')).not.toHaveAttribute('value');
await expect(progressSection(page).locator('dl.stat-row')).toHaveCount(0);
await expect(heading(page).getByRole('button', { name: /^Cancel/ })).toHaveCount(0);
await expect(await request.post('/__test__/tagging/download', { data: { mode: 'succeed' } })).toBeOK();
await expect(toast(page)).toHaveText(`${title} finished.`);
await expect(heading(page).locator('p.muted')).toHaveText(/^Finished \d\d:\d\d$/);
} finally {
await request.post('/__test__/tagging/reset');
}
});
// S5 review I3: a thumbnail job with nothing to generate ends on the server's final line ("No previews to generate · N
// already up to date"), never on a stale "Processing file 0 of 0" under the Completed chip.
test('a finished thumbnail job with nothing to do says so', async ({ page, request }) => {
@@ -429,7 +429,9 @@ test('Rebuild all is explicit and the plan line says every preview is regenerate
await expect(page).toHaveURL(/#\/activity\/[0-9a-f]{32}$/);
expect(seen.starts).toEqual([{ mode: 'rebuild' }]);
const id = page.url().split('/').pop();
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).toBe('finished');
// Rebuild regenerates all 24 previews with real image work: on a busy CI
// runner that took over 5 s (12 of 24 done at the default poll budget).
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state, { timeout: 30000 }).toBe('finished');
});
test('a plan that cannot be read says so; Start still starts', async ({ page }) => {
@@ -0,0 +1,385 @@
import { startCoverage, saveCoverage } from '../../tools/vendor/browser-coverage.mjs';
import { test, expect } from './fixtures.mjs';
// System › Auto-tagging (#328, System S11 Tasks 6–7) on the tagging fixture (workspace_tagging_fixture.py) with the
// injected fake model: no model file exists and nothing is downloaded or run for real. The fake download waits on the
// fixture (`/__test__/tagging/download {mode}`), and every reset releases it. Every test asserts that no legacy
// download stream (`/autotag/download/`) was requested: the section downloads only through the Activity job.
const workspace = '/prompt_manager/workspace';
const MODELS_URL = '/prompt_manager/autotag/models';
const PREVIEW_URL = '/prompt_manager/autotag/batch/preview';
const SETTINGS_URL = '/prompt_manager/settings';
const MODEL_LABELS = ['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB'];
const FIRST = 'tagging/dunes_00001.png';
const SECOND = 'tagging/dunes_00002.png';
test.beforeEach(async ({ page, request, browserName }) => {
await startCoverage(page, browserName);
await expect(await request.post('/__test__/tagging/reset', { data: {} })).toBeOK();
await expect(await request.post('/__test__/system/reset')).toBeOK();
await expect(await request.post('/__test__/jobs/restart')).toBeOK();
page.systemErrors = [];
page.on('pageerror', error => page.systemErrors.push(error.message));
page.downloads = [];
page.modelCalls = [];
page.on('request', item => {
const { pathname } = new URL(item.url());
if (pathname.includes('/autotag/download/')) page.downloads.push(pathname);
if (pathname.startsWith('/prompt_manager/autotag/') || pathname.startsWith('/prompt_manager/image-annotations/suggest')
|| pathname.startsWith('/prompt_manager/jobs/model-download')) page.modelCalls.push(`${item.method()} ${pathname}`);
});
});
test.afterEach(async ({ page, request, browserName }) => {
await saveCoverage(page, browserName);
// The reset releases a held fake download, so no runner outlives the test.
await request.post('/__test__/tagging/reset', { data: {} });
await request.post('/__test__/jobs/restart');
await request.post('/__test__/system/reset');
expect(page.systemErrors).toEqual([]);
expect(page.downloads).toEqual([]);
});
const toast = page => page.locator('#workspace-toast');
const layoutOf = page => page.locator('#workspace-main .settings-layout');
const dialog = page => page.locator('dialog#modal[open]');
const status = page => page.locator('#settings-status');
const save = page => page.getByRole('button', { name: 'Save preferences', exact: true });
const surface = page => layoutOf(page).locator(':scope > section.surface');
const settingOf = (page, label) => layoutOf(page).locator('.setting-row').filter({ has: page.locator(`label:text-is("${label}")`) });
const modelRow = page => settingOf(page, 'Tagging model');
const actionOf = (page, name) => surface(page).locator(':scope > .row-actions').getByRole('button', { name, exact: true });
const reviewButton = page => surface(page).locator('.operation').getByRole('button', { name: 'Review', exact: true });
const card = (page, path) => page.locator(`article[data-image-key="0:${path}"]`);
const tagDialog = page => page.getByRole('dialog', { name: 'Tag image', exact: true });
const gridOf = page => dialog(page).locator('dl.meta-grid > div').evaluateAll(items => items.map(item =>
[item.querySelector('dt').textContent, item.querySelector('dd').textContent]));
async function serverJson(request, url) {
const response = await request.get(url);
await expect(response).toBeOK();
return response.json();
}
async function fakeDownload(request, mode) {
await expect(await request.post('/__test__/tagging/download', { data: { mode } })).toBeOK();
}
async function unavailable(request, models) {
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: models } })).toBeOK();
}
async function ended(request, id) {
await expect.poll(async () => (await serverJson(request, `/prompt_manager/jobs/${id}`)).job.state).not.toBe('running');
return (await serverJson(request, `/prompt_manager/jobs/${id}`)).job;
}
/** System › Auto-tagging on a fresh page (the drafts live for the page session, Q10), settled on the model status. */
async function openTagging(page) {
if (page.url().endsWith('#/system/tagging')) await page.reload();
else await page.goto(`${workspace}#/system/tagging`);
await expect(surface(page)).toHaveCount(1);
await expect(page.locator('#tag-model')).toBeEnabled();
await expect(modelRow(page).locator('p').first()).not.toHaveText('Loading model status…');
}
/** Select `path` in Images (the concept's `state.selectedImage`), then open System › Auto-tagging from the sidebar. */
async function selectThenOpen(page, path) {
await page.goto(`${workspace}#/images`);
await card(page, path).getByRole('button', { name: `Inspect ${path.split('/').pop()}`, exact: true }).click();
await expect(card(page, path)).toHaveClass(/\bselected\b/);
await page.locator('aside.sidebar a[href="#/system"]').click();
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
await expect(page.locator('#tag-model')).toBeEnabled();
await expect(modelRow(page).locator('p').first()).not.toHaveText('Loading model status…');
}
test('model options are the four real engines with size and the description shows live status', async ({ page, request }) => {
await unavailable(request, ['wd14-vit', 'hf']);
await openTagging(page);
expect(await surface(page).locator(':scope > h2').textContent()).toBe('Auto-tagging');
expect(await page.locator('#tag-model option').allTextContents()).toEqual(MODEL_LABELS);
await expect(page.locator('#tag-model')).toHaveValue('wd14-swinv2');
const line = modelRow(page).locator('p').first();
await expect(line).toHaveText('Downloaded. Not loaded.');
await expect(actionOf(page, 'Download again')).toBeVisible();
// Choosing another engine drafts it and tells its own status; nothing is downloaded or loaded.
await page.locator('#tag-model').selectOption('wd14-vit');
await expect(line).toHaveText('Not downloaded. Download size ~400 MB.');
await expect(actionOf(page, 'Download')).toBeVisible();
await page.locator('#tag-model').selectOption('hf');
await expect(line).toHaveText('Not downloaded. Download size ~16 GB.');
await expect(status(page)).toHaveText('Unsaved changes.');
// The other rows keep the concept's copy and the saved values (D133 adds Suggestions per image).
expect(await surface(page).locator(':scope > .setting-row label').allTextContents()).toEqual(['Tagging model', 'Existing tags',
'General-tag confidence', 'Character-tag confidence', 'Suggestions per image']);
await expect(page.locator('#tag-existing')).toHaveValue('skip');
await expect(page.locator('#threshold-general')).toHaveValue('0.35');
await expect(page.locator('#threshold-character')).toHaveValue('0.85');
await expect(page.locator('#tag-suggestions')).toHaveValue('10');
const settings = await serverJson(request, SETTINGS_URL);
await expect(page.locator('#tag-template')).toHaveValue(settings.settings.autotag_template);
// Unload is offered only while a model is loaded (none is).
await expect(actionOf(page, 'Unload')).toBeDisabled();
});
test('entering the section makes no model call besides status', async ({ page }) => {
await openTagging(page);
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Integrations', exact: true }).click();
await page.locator('#workspace-main .section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
await expect(page.locator('#tag-model')).toBeEnabled();
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
// One status read per mount (two mounts), and nothing else that could load, run or download a model.
expect(page.modelCalls).toEqual([`GET ${MODELS_URL}`, `GET ${MODELS_URL}`]);
const models = await (await page.request.get(MODELS_URL)).json();
expect(models.model_loaded).toBe(false);
});
test('defaults save through Save preferences and the Tag image dialog uses them', async ({ page, request }) => {
await openTagging(page);
await page.locator('#tag-existing').selectOption('merge');
await page.locator('#threshold-general').fill('0.4');
await page.locator('#threshold-general').press('Tab');
await page.locator('#tag-suggestions').fill('15');
await page.locator('#tag-suggestions').press('Tab');
const posted = page.waitForRequest(item => item.method() === 'POST' && new URL(item.url()).pathname === SETTINGS_URL);
await save(page).click();
const body = (await posted).postDataJSON();
expect(body).toMatchObject({ autotag_existing: 'merge', autotag_general_threshold: 0.4, autotag_suggestion_limit: 15 });
expect(typeof body.expected_revision).toBe('string');
await expect(status(page)).toHaveText('Saved to the server configuration.');
const saved = (await serverJson(request, SETTINGS_URL)).settings;
expect([saved.autotag_existing, saved.autotag_general_threshold, saved.autotag_character_threshold, saved.autotag_suggestion_limit])
.toEqual(['merge', 0.4, 0.85, 15]);
const models = await serverJson(request, MODELS_URL);
expect([models.existing, models.wd14_general_threshold, models.suggestion_limit]).toEqual(['merge', 0.4, 15]);
// The Tag image dialog in a fresh page starts from them.
await page.goto(`${workspace}#/images`);
await card(page, FIRST).getByRole('button', { name: `Tag or auto-tag ${FIRST.split('/').pop()}`, exact: true }).click();
await expect(tagDialog(page)).toBeVisible();
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.4');
await expect(tagDialog(page).locator('#image-character-threshold')).toHaveValue('0.85');
await expect(tagDialog(page).locator('#image-suggestion-limit')).toHaveValue('15');
});
test('a refused value shows the server\'s message', async ({ page, request }) => {
const before = (await serverJson(request, SETTINGS_URL)).settings;
// An image is selected, so Preview tagging review is held back only by the invalid value below.
await selectThenOpen(page, FIRST);
await page.locator('#threshold-general').fill('1.5');
await page.locator('#threshold-general').press('Tab');
await save(page).click();
await expect(status(page)).toHaveText('General-tag confidence must be between 0 and 1.');
await expect(page.locator('#threshold-general')).toHaveAttribute('aria-invalid', 'true');
// Nothing was written; the draft is kept for the next Save.
expect((await serverJson(request, SETTINGS_URL)).settings.autotag_general_threshold).toBe(before.autotag_general_threshold);
await expect(page.locator('#threshold-general')).toHaveValue('1.5');
// While the value is invalid, Preview tagging review is aria-disabled and explains instead of opening (Playwright
// waits on an aria-disabled control, so it is pressed from the keyboard, as a user can).
await expect(actionOf(page, 'Preview tagging review')).toHaveAttribute('aria-disabled', 'true');
await actionOf(page, 'Preview tagging review').focus();
await page.keyboard.press('Enter');
await expect(toast(page)).toHaveText('General-tag confidence must be between 0 and 1.');
await expect(tagDialog(page)).toHaveCount(0);
await page.locator('#threshold-general').fill('0.5');
await page.locator('#threshold-general').press('Tab');
await expect(page.locator('#threshold-general')).not.toHaveAttribute('aria-invalid', 'true');
await expect(status(page)).toHaveText('Unsaved changes.');
await expect(actionOf(page, 'Preview tagging review')).not.toHaveAttribute('aria-disabled', 'true');
});
test('a value off the 0.05 step is refused and never reaches the Tag image dialog (#369)', async ({ page, request }) => {
const before = (await serverJson(request, SETTINGS_URL)).settings;
const STEP = 'General-tag confidence must be a multiple of 0.05.';
await selectThenOpen(page, FIRST);
await page.locator('#threshold-general').fill('0.37');
await page.locator('#threshold-general').press('Tab');
// Preview tagging review explains instead of opening a dialog whose Generate the browser would block.
await expect(actionOf(page, 'Preview tagging review')).toHaveAttribute('aria-disabled', 'true');
await actionOf(page, 'Preview tagging review').focus();
await page.keyboard.press('Enter');
await expect(toast(page)).toHaveText(STEP);
await expect(tagDialog(page)).toHaveCount(0);
// The server refuses the save with the same sentence and writes nothing.
await save(page).click();
await expect(status(page)).toHaveText(STEP);
expect((await serverJson(request, SETTINGS_URL)).settings.autotag_general_threshold).toBe(before.autotag_general_threshold);
// A value on the step opens the dialog, and its Generate passes the browser's own check and sends.
await page.locator('#threshold-general').fill('0.4');
await page.locator('#threshold-general').press('Tab');
await expect(actionOf(page, 'Preview tagging review')).not.toHaveAttribute('aria-disabled', 'true');
await actionOf(page, 'Preview tagging review').click();
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.4');
expect(await tagDialog(page).locator('#image-general-threshold').evaluate(input => input.checkValidity())).toBe(true);
const suggested = page.waitForResponse(item => item.url().endsWith('/image-annotations/suggest'));
await tagDialog(page).getByRole('button', { name: 'Generate suggestions', exact: true }).click();
const answer = await suggested;
expect(answer.request().postDataJSON().options.general_threshold).toBe(0.4);
});
test('Download starts the job and Activity shows it without Cancel', async ({ page, request }) => {
await unavailable(request, ['wd14-vit']);
await fakeDownload(request, 'hold');
await openTagging(page);
await page.locator('#tag-model').selectOption('wd14-vit');
const started = page.waitForResponse(item => new URL(item.url()).pathname === '/prompt_manager/jobs/model-download');
await actionOf(page, 'Download').click();
const answer = await started;
expect(answer.status()).toBe(202);
expect(answer.request().postDataJSON()).toEqual({ model_type: 'wd14-vit' });
const { job } = await answer.json();
expect(job.title).toBe('Tagging model download · WD14 ViT');
await expect(toast(page)).toHaveText('Tagging model download · WD14 ViT started.');
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
const main = page.locator('#workspace-main');
await expect(main.locator('h1')).toHaveText('Tagging model download · WD14 ViT');
await expect(main.getByRole('button', { name: /^Cancel/ })).toHaveCount(0);
// Indeterminate: the server knows no byte progress.
await expect(main.locator('progress').first()).toBeVisible();
expect(await main.locator('progress').first().evaluate(node => node.hasAttribute('value'))).toBe(false);
await fakeDownload(request, 'succeed');
expect((await ended(request, job.id)).state).toBe('finished');
await expect(toast(page)).toHaveText('Tagging model download · WD14 ViT finished.');
// Back in the section, the engine reads as on disk (the fake marks it downloaded on success).
await page.goto(`${workspace}#/system/tagging`);
await expect(page.locator('#tag-model')).toBeEnabled();
await page.locator('#tag-model').selectOption('wd14-vit');
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
expect(page.modelCalls.filter(call => call.startsWith('POST'))).toEqual(['POST /prompt_manager/jobs/model-download']);
});
test('Unload is enabled only while a model is loaded, and a held gate\'s refusal names the operation', async ({ page, request }) => {
await openTagging(page);
await expect(actionOf(page, 'Unload')).toBeDisabled();
// Load WD14 SwinV2 the way a user does: Generate suggestions in the Tag image dialog.
await page.goto(`${workspace}#/images`);
await card(page, FIRST).getByRole('button', { name: `Tag or auto-tag ${FIRST.split('/').pop()}`, exact: true }).click();
await expect(tagDialog(page)).toBeVisible();
const suggested = page.waitForResponse(item => item.url().endsWith('/image-annotations/suggest'));
await tagDialog(page).getByRole('button', { name: 'Generate suggestions', exact: true }).click();
expect((await suggested).status()).toBe(200);
await tagDialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
await openTagging(page);
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded and loaded.');
await expect(actionOf(page, 'Unload')).toBeEnabled();
// A download holds the model gate: Unload is refused and the toast names it; nothing is unloaded.
await unavailable(request, ['wd14-vit']);
await fakeDownload(request, 'hold');
const held = await request.post('/prompt_manager/jobs/model-download', { data: { model_type: 'wd14-vit' } });
expect(held.status()).toBe(202);
const { job } = await held.json();
const refused = page.waitForResponse(item => item.url().endsWith('/autotag/unload'));
await actionOf(page, 'Unload').click();
const refusal = await refused;
expect(refusal.status()).toBe(409);
expect(await refusal.json()).toMatchObject({ code: 'model-busy', operation: 'Tagging model download · WD14 ViT' });
await expect(toast(page)).toHaveText('The tagging model is busy: Tagging model download · WD14 ViT. Try again when it finishes.');
expect((await serverJson(request, MODELS_URL)).loaded_model_type).toBe('wd14-swinv2');
await expect(actionOf(page, 'Unload')).toBeEnabled();
// Once the download ended, Unload works, says so and the row follows.
await fakeDownload(request, 'succeed');
await ended(request, job.id);
await actionOf(page, 'Unload').click();
await expect(toast(page)).toHaveText('WD14 SwinV2 unloaded.');
await expect(modelRow(page).locator('p').first()).toHaveText('Downloaded. Not loaded.');
await expect(actionOf(page, 'Unload')).toBeDisabled();
});
test('Preview tagging review opens the Tag image dialog for the selected image', async ({ page }) => {
await selectThenOpen(page, SECOND);
// The section's current (unsaved) values are what the dialog starts from.
await page.locator('#threshold-character').fill('0.6');
await page.locator('#threshold-character').press('Tab');
await page.locator('#tag-model').selectOption('wd14-vit');
const preview = actionOf(page, 'Preview tagging review');
await expect(preview).toHaveClass(/ca-primary/);
await expect(preview).not.toHaveAttribute('aria-disabled', 'true');
const observed = page.waitForResponse(item => item.url().endsWith('/image-annotations/observe'));
await preview.click();
expect((await observed).status()).toBe(200);
await expect(tagDialog(page)).toBeVisible();
await expect(tagDialog(page).locator('.tag-image-heading p.small')).toHaveText('dunes_00002.png');
await expect(tagDialog(page).getByRole('combobox', { name: 'Tagging model', exact: true })).toHaveValue('wd14-vit');
await expect(tagDialog(page).locator('#image-character-threshold')).toHaveValue('0.6');
await expect(tagDialog(page).locator('#image-general-threshold')).toHaveValue('0.35');
});
test('Preview tagging review without a selection toasts the D115 sentence', async ({ page }) => {
await openTagging(page);
const preview = actionOf(page, 'Preview tagging review');
await expect(preview).toHaveAttribute('aria-disabled', 'true');
// An aria-disabled control still takes a press; it only explains.
await preview.focus();
await page.keyboard.press('Enter');
await expect(toast(page)).toHaveText('Select an image in Images first.');
await expect(tagDialog(page)).toHaveCount(0);
});
test('Tag output images reviews the scope and starts a batch job', async ({ page, request }) => {
const before = await serverJson(request, '/__test__/tagging/snapshot');
await openTagging(page);
// The fixture's one prompt is already tagged: under Merge it is tagged again with what its images give.
await page.locator('#tag-existing').selectOption('merge');
const previewed = page.waitForResponse(item => new URL(item.url()).pathname === PREVIEW_URL);
await reviewButton(page).click();
const answer = await previewed;
expect(answer.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'merge' });
const plan = await answer.json();
await expect(dialog(page).locator('#modal-title')).toHaveText('Review tag output images');
await expect(dialog(page).locator('#modal-body > h3')).toHaveText('Tag output images');
expect(await gridOf(page)).toEqual([['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Merge with existing'],
['Images linked', String(plan.images)], ['Prompts', String(plan.prompts)], ['Already tagged', String(plan.already_tagged)],
['Prompts to tag', `up to ${plan.to_tag}`], ['Settings', 'General 0.35 · Character 0.85'],
['Backup', 'Create and verify before changes']]);
await expect(dialog(page).locator('#modal-body > .notice h3')).toHaveText('Back up the database first (recommended)');
// The review wrote nothing.
expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before);
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel', 'Start tagging']);
const started = page.waitForResponse(item => new URL(item.url()).pathname === '/prompt_manager/jobs/autotag');
await dialog(page).getByRole('button', { name: 'Start tagging', exact: true }).click();
const start = await started;
expect(start.request().postDataJSON()).toEqual({ model_type: 'wd14-swinv2', existing: 'merge',
options: { general_threshold: 0.35, character_threshold: 0.85 } });
const { job } = await start.json();
await expect(page).toHaveURL(`${workspace}#/activity/${job.id}`);
await expect(page.locator('#workspace-main h1')).toHaveText('Tag output images');
const finished = await ended(request, job.id);
expect(finished.state).toBe('finished');
const after = await serverJson(request, '/__test__/tagging/snapshot');
expect(after.prompt_tags.length).toBeGreaterThan(before.prompt_tags.length);
});
test('Replace after review shows only Cancel with the reason', async ({ page, request }) => {
const before = await serverJson(request, '/__test__/tagging/snapshot');
await openTagging(page);
await page.locator('#tag-existing').selectOption('replace_after_review');
await reviewButton(page).click();
await expect(dialog(page).locator('#modal-title')).toHaveText('Review tag output images');
expect((await gridOf(page)).map(([term]) => term)).toEqual(['Model', 'Existing tags', 'Images linked', 'Prompts', 'Already tagged',
'Prompts to tag', 'Settings', 'Backup']);
await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Replacing tags needs a review of each image. '
+ 'Use Preview tagging review, or choose Skip or Merge for a batch run.');
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']);
await dialog(page).getByRole('button', { name: 'Cancel', exact: true }).click();
await expect(dialog(page)).toHaveCount(0);
// Under Skip the fixture's only prompt is already tagged: nothing to start either.
await page.locator('#tag-existing').selectOption('skip');
await reviewButton(page).click();
await expect(dialog(page).locator('#modal-body > p.small')).toHaveText('Every prompt with an output image is already tagged.');
await expect(dialog(page).locator('.dialog-actions button')).toHaveText(['Cancel']);
expect(await serverJson(request, '/__test__/tagging/snapshot')).toEqual(before);
expect((await serverJson(request, '/prompt_manager/jobs')).jobs.filter(item => item.kind === 'autotag')).toEqual([]);
});
test('a model that is not downloaded cannot be reviewed for a batch; the button says why', async ({ page, request }) => {
await unavailable(request, ['gguf']);
await openTagging(page);
await page.locator('#tag-model').selectOption('gguf');
await expect(reviewButton(page)).toHaveAttribute('aria-disabled', 'true');
// Playwright waits on an aria-disabled control: press it from the keyboard, as a user can.
await reviewButton(page).focus();
await page.keyboard.press('Enter');
await expect(toast(page)).toHaveText('GGUF JoyCaption is not downloaded. Download it first.');
await expect(dialog(page)).toHaveCount(0);
expect(page.modelCalls.filter(call => call.includes('/batch/preview'))).toEqual([]);
});
+3 -3
View File
@@ -412,11 +412,11 @@ test('out-of-order overview responses never repaint a newer section', async ({ p
await nav.getByRole('button', { name: 'Appearance', exact: true }).click();
await nav.getByRole('button', { name: 'Overview', exact: true }).click();
await expect.poll(held).toBe(2);
// Auto-tagging still hands off (S11), so its body is one known panel.
await nav.getByRole('button', { name: 'Auto-tagging', exact: true }).click();
// Diagnostics still hands off (S12), so its body is one known panel.
await nav.getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
await expect(await request.post('/__test__/release')).toBeOK();
await expect.poll(held).toBe(0);
// Both late answers reached a section that had left: Auto-tagging keeps its own body.
// Both late answers reached a section that had left: Diagnostics keeps its own body.
await expect(layout.locator(':scope > section.handoff')).toBeVisible();
await page.evaluate(() => new Promise(resolve => requestAnimationFrame(() => requestAnimationFrame(resolve))));
await expect(layout.locator('.summary-metrics')).toHaveCount(0);
@@ -17,7 +17,7 @@ const chips = page => modal(page).locator('#tag-suggestions .chip.tag-suggestion
const actionButtons = page => modal(page).locator('.dialog-actions button');
const panel = page => page.getByRole('complementary', { name: 'Selected image details', exact: true });
const tagRow = page => panel(page).locator('.detail-inner > .tags').last();
const NOT_DOWNLOADED = 'WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.';
const NOT_DOWNLOADED = 'WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.';
test.beforeEach(async ({ page, request, browserName }) => {
await startCoverage(page, browserName);
+1 -1
View File
@@ -246,7 +246,7 @@ test('unavailable model shows inline status and never downloads', async ({ page,
await expect(await request.post('/__test__/tagging/models', { data: { unavailable: ['wd14-swinv2'] } })).toBeOK();
await openTag(page);
expect((await generated(page)).status()).toBe(409);
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.');
await expect(status(page)).toHaveText('WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.');
await expect(chips(page)).toHaveCount(0);
await expect(dialog(page).locator('#suggestion-count')).toHaveText('Not generated');
// D116: Retry leads the actions while the controller allows it, and generates again.
+7 -5
View File
@@ -413,7 +413,7 @@ for (const width of [1600, 390]) test(`overview and handoff render with concept
await sidebarLink(page, 'System').click();
// System's Overview is native (S2 Task 6); a section not delivered yet still hands off (D94).
await main.locator('.section-nav').getByRole('button', { name: 'Auto-tagging', exact: true }).click();
await main.locator('.section-nav').getByRole('button', { name: 'Diagnostics & logs', exact: true }).click();
const handoff = main.locator('section.handoff');
await expect(handoff).toHaveClass(/\bsurface\b/);
expect(await handoff.evaluate(surfaceStyle)).toMatchObject(SURFACE);
@@ -465,8 +465,9 @@ test('System is native: the concept shell, section navigation, sections, hand-of
// S2 Task 6: Overview is native (workspace-system.spec.mjs), and so are Appearance (S3 Task 8) and Sources & scans
// (S4 Task 11, workspace-system-sources.spec.mjs, its update notice shown in a fresh browser), Integrations (S10
// Task 27, workspace-system-integrations.spec.mjs), Maintenance (S7 Task 19,
// workspace-system-maintenance.spec.mjs) and Backup & export (S9 Task 24, workspace-system-backups.spec.mjs);
// every other section hands off.
// workspace-system-maintenance.spec.mjs), Backup & export (S9 Task 24, workspace-system-backups.spec.mjs) and
// Auto-tagging (S11 Task 6, workspace-system-tagging.spec.mjs); every other section hands off.
if (hash === 'tagging') await expect(layout.locator('#tag-model')).toBeVisible();
if (hash === 'overview') await expect(layout.locator(':scope > .summary-metrics')).toBeVisible();
if (hash === 'sources') await expect(layout.locator('#source-list .operation')).toHaveCount(2);
expect(await layout.evaluate(node => [...node.children].map(child => `${child.localName}.${child.className}`)))
@@ -474,9 +475,10 @@ test('System is native: the concept shell, section navigation, sections, hand-of
sources: ['div.notice', 'section.surface', 'div.save-row'],
integrations: ['section.surface', 'section.surface', 'div.save-row'],
maintenance: ['div.notice', 'section.surface', 'section.surface', 'div.save-row'],
backups: ['div.two-col', 'section.surface', 'div.save-row'] }[hash]
backups: ['div.two-col', 'section.surface', 'div.save-row'],
tagging: ['section.surface', 'div.save-row'] }[hash]
?? ['section.surface handoff', 'div.save-row']);
if (['overview', 'appearance', 'sources', 'integrations', 'maintenance', 'backups'].includes(hash)) continue;
if (['overview', 'appearance', 'sources', 'integrations', 'tagging', 'maintenance', 'backups'].includes(hash)) continue;
const handoff = layout.locator('section.handoff');
await expect(handoff.locator('h2')).toHaveText('Your current tools are ready');
expect(await handoff.locator('a.workspace-link').evaluateAll(links => links.map(link => link.getAttribute('href'))))
+2 -3
View File
@@ -16,7 +16,6 @@ import json
import math
import os
from pathlib import Path
import shutil
import time
from types import SimpleNamespace
from typing import TYPE_CHECKING
@@ -34,6 +33,7 @@ from utils.execution_contexts import (
from utils.execution_history_events import make_finish, make_start
from utils.execution_outputs import OutputBundle, OutputReference, snapshot_output_roots
from utils.pair_execution import EXECUTION_KEY
from tests.browser.folder_trash import replace_folder
if TYPE_CHECKING:
from py.api import PromptManagerAPI
@@ -300,8 +300,7 @@ def _clear(api: PromptManagerAPI, root: Path) -> None:
for table in _TABLES:
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
def _record_run(
@@ -3,11 +3,11 @@
from __future__ import annotations
from pathlib import Path
import shutil
from typing import TYPE_CHECKING
from aiohttp import web
from PIL import Image
from tests.browser.folder_trash import replace_folder
if TYPE_CHECKING:
from py.api import PromptManagerAPI
@@ -51,8 +51,7 @@ def seed_library_editor(api: PromptManagerAPI, root: Path) -> dict:
for table in _TABLES:
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
epoch = db.get_execution_epoch()
receipt = db.create_library_recipe(
EDITOR_POSITIVE,
+2 -2
View File
@@ -52,6 +52,7 @@ from typing import TYPE_CHECKING
from aiohttp import web
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from tests.browser.folder_trash import replace_folder
if TYPE_CHECKING:
from py.api import PromptManagerAPI
@@ -225,8 +226,7 @@ def seed_review(
conn.execute(f"DELETE FROM {table}")
api.close_annotation_captures()
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
shutil.rmtree(root / "temp" / "review", ignore_errors=True)
output = root / "output"
files = []
+20 -17
View File
@@ -97,15 +97,25 @@ def register_system_commands(app, routes, api, root: Path):
real_move = library_move.move_library
keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER
config_bytes = config_path.read_bytes() if config_path.exists() else None
host_values = (
PromptManagerConfig.RESULT_TIMEOUT,
PromptManagerConfig.WEBUI_DISPLAY_MODE,
PromptManagerConfig.WORKER_THREADS,
GalleryConfig.MONITORING_ENABLED,
PromptManagerConfig.THUMBNAIL_FORMAT,
PromptManagerConfig.THUMBNAIL_QUALITY,
PromptManagerConfig.THUMBNAIL_EDGE,
from py.api.settings_routes import (
_AUTOTAG_SETTINGS,
_GALLERY_SETTINGS,
_SETTINGS_FILE_KEYS,
)
# Every value POST /settings can change, so a reset restores each one
# (a saved safety_backup=false once leaked into a later spec's review).
# The tagging settings are the tagging fixture's to reset.
host_values = {
(GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig, name): (
getattr(
GalleryConfig if name in _GALLERY_SETTINGS else PromptManagerConfig,
name,
)
)
for name in _SETTINGS_FILE_KEYS
if name not in _AUTOTAG_SETTINGS
}
from py.api import thumbnail_encode
monitor = getattr(image_monitor, "_monitor_instance", None)
@@ -135,15 +145,8 @@ def register_system_commands(app, routes, api, root: Path):
unplugged.rmdir()
GalleryConfig.MONITORING_DIRECTORIES = list(configured)
api._health_report = None
(
PromptManagerConfig.RESULT_TIMEOUT,
PromptManagerConfig.WEBUI_DISPLAY_MODE,
PromptManagerConfig.WORKER_THREADS,
GalleryConfig.MONITORING_ENABLED,
PromptManagerConfig.THUMBNAIL_FORMAT,
PromptManagerConfig.THUMBNAIL_QUALITY,
PromptManagerConfig.THUMBNAIL_EDGE,
) = host_values
for (owner, name), value in host_values.items():
setattr(owner, name, value)
thumbnail_encode.reset_encoders()
image_monitor._monitor_instance = monitor
api.__dict__.pop("_library_resolver", None)
+65 -4
View File
@@ -12,12 +12,14 @@ import json
import os
from pathlib import Path
import shutil
import threading
import time
from typing import TYPE_CHECKING
from aiohttp import web
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from tests.browser.folder_trash import replace_folder
if TYPE_CHECKING:
from py.api import PromptManagerAPI
@@ -97,8 +99,7 @@ def seed_tagging(api: PromptManagerAPI, root: Path, extra_images: int = 0) -> di
for table in _TABLES:
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
output = root / "output"
_png(output / FILES[0][0], (196, 140, 92))
_png(output / FILES[1][0], (150, 110, 80))
@@ -181,6 +182,28 @@ def register_tagging_commands(
routes: web.RouteTableDef, api: PromptManagerAPI, root: Path, model
) -> None:
"""Register fixed seed, snapshot and fake-model commands; no paths or SQL input."""
# The fake download (``/__test__/tagging/download``): waits on ``event``,
# then answers ``result``. Every reset sets the event and puts back the
# model's raising default, so no test can leave a worker hanging.
# ``generation`` grows on every reset: a download a reset released (or one
# that timed out) fails without touching the next test's model state.
download = {"event": threading.Event(), "result": True, "generation": 0}
def restore_download() -> None:
download["generation"] += 1
download["event"].set()
model.__dict__.pop("download_model", None)
def fake_download(model_type: str, progress_callback=None) -> bool:
event, generation = download["event"], download["generation"]
released = event.wait(DOWNLOAD_HOLD_LIMIT_S)
if not released or download["generation"] != generation:
return False
if download["result"]:
model.unavailable = model.unavailable - {model_type}
return download["result"]
_set_defaults(model, DEFAULTS) # the server starts on the standard ones
@routes.post("/__test__/tagging/reset")
async def reset(request: web.Request) -> web.Response:
@@ -195,6 +218,7 @@ def register_tagging_commands(
model.unavailable = set()
model.suggestions = None
model.loaded = None
restore_download()
_set_defaults(model, DEFAULTS)
# A test that ends with the dialog open leaves its capture behind; drop
# every capture so the next test starts below the store's limits.
@@ -223,6 +247,33 @@ def register_tagging_commands(
model.unavailable = set(data["unavailable"])
return web.json_response({"success": True})
@routes.post("/__test__/tagging/download")
async def download_mode(request: web.Request) -> web.Response:
"""Switch the model's download to a fake (never a real download).
``{"mode": "hold"}``: the next download waits until a later
``succeed`` or ``fail`` (or a reset) releases it. ``succeed`` and
``fail`` release a held download with that result and answer every
later one at once. A successful fake marks the model downloaded.
"""
data = await request.json() if request.can_read_body else None
if (
request.query
or not isinstance(data, dict)
or set(data) != {"mode"}
or data["mode"] not in ("hold", "succeed", "fail")
):
raise web.HTTPBadRequest()
if data["mode"] == "hold":
download["event"].set()
download["event"] = threading.Event()
download["result"] = True
else:
download["result"] = data["mode"] == "succeed"
download["event"].set()
model.download_model = fake_download
return web.json_response({"success": True})
@routes.post("/__test__/tagging/suggestions")
async def suggestions(request: web.Request) -> web.Response:
"""Answer Generate with fixed scored names (the concept's own sample).
@@ -273,18 +324,28 @@ def register_tagging_commands(
return web.json_response({"captures": count})
# A held fake download gives up after this long, whatever the test does.
DOWNLOAD_HOLD_LIMIT_S = 60
# The standard tagging defaults: WD14 general and character thresholds, and
# the suggestion count (``AutotagConfig``).
DEFAULTS = (0.35, 0.85, 10)
def _set_defaults(model, values) -> None:
"""Set the saved tagging defaults (``AutotagConfig``, #328); a reset also
puts back the saved model, batch policy and template (the fake model's
short instructions, which the concept comparisons show)."""
from py.config import AutotagConfig
general, character, limit = values
model.wd14_general_threshold = general
model.wd14_character_threshold = character
AutotagConfig.GENERAL_THRESHOLD = general
AutotagConfig.CHARACTER_THRESHOLD = character
AutotagConfig.SUGGESTION_LIMIT = limit
if values == DEFAULTS:
AutotagConfig.MODEL_TYPE = "wd14-swinv2"
AutotagConfig.EXISTING = "skip"
AutotagConfig.TEMPLATE = model.default_prompt
def _valid_defaults(data: object) -> bool:
+2 -3
View File
@@ -5,13 +5,13 @@ from __future__ import annotations
import json
import os
from pathlib import Path
import shutil
from typing import TYPE_CHECKING
from PIL import Image, ImageDraw
from PIL.PngImagePlugin import PngInfo
from utils.hashing import generate_prompt_hash
from tests.browser.folder_trash import replace_folder
if TYPE_CHECKING:
from py.api import PromptManagerAPI
@@ -129,8 +129,7 @@ def seed_workspace_visual(api: PromptManagerAPI, root: Path) -> None:
):
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
for parent_index, (positive, negatives) in enumerate(zip(_POSITIVES, _NEGATIVES)):
category = ("Architecture", "Landscape", "Abstract")[parent_index % 3]
tags = ["geometric", ("warm", "cool")[parent_index % 2]]
@@ -7,10 +7,10 @@ in concept order, so tag creation order and category first use match the mock.
from __future__ import annotations
from pathlib import Path
import shutil
from typing import TYPE_CHECKING
from aiohttp import web
from tests.browser.folder_trash import replace_folder
if TYPE_CHECKING:
from py.api import PromptManagerAPI
@@ -120,8 +120,7 @@ def seed_vocabulary(
for table in _TABLES:
conn.execute(f"DELETE FROM {table}")
for name in ("output", "input"):
shutil.rmtree(root / name)
(root / name).mkdir()
replace_folder(root / name)
epoch = db.get_execution_epoch()
recipes = []
for recipe in CONCEPT_RECIPES:
+4 -1
View File
@@ -1,5 +1,6 @@
"""Deterministic nodes and queue controls installed only in an owned test host."""
import asyncio
import os
import threading
import time
@@ -291,7 +292,9 @@ async def image_monitor(request: web.Request) -> web.Response:
config.MONITORING_ENABLED = True
monitor.start_monitoring([folder_paths.get_output_directory()])
else:
monitor.stop_monitoring()
# Joins the image worker and the observer (one shared 5 s deadline),
# so it runs on a worker thread, not on the event loop.
await asyncio.to_thread(monitor.stop_monitoring)
config.MONITORING_ENABLED = False
return web.json_response({"monitoring": monitor.is_monitoring})
+10
View File
@@ -388,6 +388,16 @@ test("ApiError keeps a validated running job on busy and a validated source on i
assert.deepEqual([new ApiError(409, "old caller").running, new ApiError(409, "old caller").source], [null, null]);
});
test("ApiError keeps a validated operation on model-busy only (#328)", () => {
const body = { operation: "Tagging model download · WD14 ViT" };
assert.equal(new ApiError(409, "busy", "model-busy", null, body).operation, "Tagging model download · WD14 ViT");
assert.equal(new ApiError(409, "busy", "busy", null, body).operation, null);
for (const operation of [null, 7, "", "x".repeat(201)]) {
assert.equal(new ApiError(409, "busy", "model-busy", null, { operation }).operation, null);
}
assert.equal(new ApiError(409, "busy", "model-busy").operation, null);
});
test("ApiError says whether the server itself refused (a parsed success:false body; S7 review I1)", () => {
assert.equal(new ApiError(400, "bad", null, null, { success: false, error: "bad" }).refused, true);
assert.equal(new ApiError(409, "busy", "busy", null, { success: false, code: "busy" }).refused, true);
+23
View File
@@ -225,3 +225,26 @@ test('a failed duplicate cleanup says the server reason (S7 review M3)', async (
assert.equal(finishToast(failed), `Duplicate cleanup failed · ${reason}`);
assert.equal(finishToast({ ...failed, status_line: '' }), 'Duplicate cleanup failed · 0 removed · 1 error');
});
// #328: Tag output images (`autotag`) reuses the image tagging view: Tagged / Unchanged / Skipped / Errors and Cancel.
// The legacy page's hold of the same kind counts nothing and is not cancellable, so it keeps the generic held view.
test('autotag gets Tagged / Unchanged / Skipped / Errors and Cancel; its legacy hold keeps the generic view', async () => {
const { kindOf, viewOf, canCancel } = await source;
const { startToast, finishToast } = await store;
const counters = { tagged: 3, unchanged: 1, skipped: 2, errors: 0 };
const batch = job({ kind: 'autotag', title: 'Tag output images', counters, scope: { images: 6 } });
const view = viewOf(batch);
assert.deepEqual(view.tiles.map(([group, label]) => [group, label]),
[['tagged', 'Tagged'], ['unchanged', 'Unchanged'], ['skipped', 'Skipped'], ['errors', 'Errors']]);
assert.equal(kindOf('autotag').pillVerb, 'Tagging');
assert.equal(kindOf('autotag').cancelLabel, 'Cancel tagging');
assert.equal(canCancel(batch), true);
assert.equal(startToast(batch), 'Tag output images started.');
assert.match(finishToast({ ...batch, state: 'finished', finished_at: 2 }), /^Tag output images finished/);
const held = job({ kind: 'autotag', title: 'Tag output images (legacy page)', counters: {}, cancellable: false,
progress: 0, processed: 0, total: 0 });
const legacy = viewOf(held);
assert.deepEqual([legacy.tiles, legacy.outcomes, [...legacy.openable]], [[], [], []]);
assert.equal(canCancel(held), false);
});
+47 -1
View File
@@ -197,6 +197,52 @@ test('instructions default to the server prompt', async () => {
assert.equal(c.state.options.suggestion_limit, 7, 'the suggestion count survives a model change');
});
// #328: one tagging configuration. The dialogs start on the model saved in System › Auto-tagging, unless a preset
// (a review session's choice) or the user already picked one; an unknown saved model keeps the first option.
test('dialog starts on the saved model unless a preset or the user chose one', async () => {
const saved = { [`GET ${MODELS}`]: () => response(200, models({ model_type: 'gguf' })) };
const { c } = await opened(saved);
assert.equal(c.state.model, 'gguf');
assert.equal(c.state.options.instructions, DEFAULT_PROMPT);
const { c: preset } = await opened(saved, { preset: { model: 'wd14-vit', options: {} } });
assert.equal(preset.state.model, 'wd14-vit');
const { createImageTaggingController, createTaggingOptionsController } = await load();
const h = harness(defaultRoutes(saved));
const picked = createImageTaggingController({ api: h.api, target: TARGET, source: 'images' });
picked.setModel('hf');
await picked.open();
await settle();
assert.equal(picked.state.model, 'hf', 'the user picked a model before the read landed');
// A pick made while the models read is still in flight wins over the saved model that lands after it.
const gate = deferred();
const late = harness(defaultRoutes({ [`GET ${MODELS}`]: () => gate.promise }));
const racing = createImageTaggingController({ api: late.api, target: TARGET, source: 'images' });
const opening = racing.open();
await settle();
assert.equal(racing.state.phase, 'observing');
racing.setModel('wd14-vit');
gate.resolve(response(200, models({ model_type: 'gguf' })));
await opening;
await settle();
assert.equal(racing.state.model, 'wd14-vit');
racing.dispose();
const { c: unknown } = await opened({ [`GET ${MODELS}`]: () => response(200, models({ model_type: 'wd14-convnext' })) });
assert.equal(unknown.state.model, 'wd14-swinv2');
const options = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api });
await options.open();
assert.equal(options.state.model, 'gguf');
const chosen = createTaggingOptionsController({ api: harness(defaultRoutes(saved)).api });
chosen.setModel('wd14-vit');
await chosen.open();
assert.equal(chosen.state.model, 'wd14-vit');
for (const controller of [c, preset, picked, unknown, options, chosen]) controller.dispose();
});
test('an unreadable or malformed models read falls back to the documented defaults', async () => {
for (const route of [() => fail(500, null), () => response(200, models({ default_prompt: 4, suggestion_limit: 0, wd14_general_threshold: 2, wd14_character_threshold: 'x' }))]) {
const { c } = await opened({ [`GET ${MODELS}`]: route });
@@ -542,7 +588,7 @@ test('model-unavailable and model-busy set the status line only', async () => {
await c.generate();
const before = c.state;
for (const [code, message] of [
['model-unavailable', 'WD14 SwinV2 is not downloaded. Download it from the legacy Auto-tag page.'],
['model-unavailable', 'WD14 SwinV2 is not downloaded. Download it in System › Auto-tagging.'],
['model-busy', 'A tagging model is busy with another operation. Try again when it finishes.'],
['unsupported-media', 'Images larger than 4096 px per side or 16,777,216 px are not supported.'],
]) {
@@ -129,9 +129,9 @@ test('a model that is not downloaded is named with the server sentence', () => {
const { modelUnavailable } = state;
assert.equal(modelUnavailable(null, 'wd14-vit'), '');
assert.equal(modelUnavailable({ 'wd14-vit': { downloaded: true } }, 'wd14-vit'), '');
assert.equal(modelUnavailable({ 'wd14-vit': { downloaded: false } }, 'wd14-vit'), 'WD14 ViT is not downloaded. Download it from the legacy Auto-tag page.');
assert.equal(modelUnavailable({ 'wd14-vit': { downloaded: false } }, 'wd14-vit'), 'WD14 ViT is not downloaded. Download it in System › Auto-tagging.');
assert.equal(modelUnavailable({ gguf: { downloaded: false, name: 'GGUF (Recommended)' } }, 'gguf'),
'GGUF (Recommended) is not downloaded. Download it from the legacy Auto-tag page.');
'GGUF (Recommended) is not downloaded. Download it in System › Auto-tagging.');
assert.throws(() => modelUnavailable({}, 'nope'), TypeError);
});
+16
View File
@@ -341,6 +341,22 @@ test('a host save names the settings that apply after a restart (D128)', async (
assert.equal(drafts.statusText(), 'Unsaved changes.');
});
test('a host save refused for one setting says the server reason and names no source (S11 CUJ 2a)', async () => {
const { INVALID_SETTING_CODE } = await source;
const domains = adapters({ host: { autotag_general_threshold: 0.35 } });
const drafts = await loaded(domains);
const message = 'General-tag confidence must be between 0 and 1.';
domains.outcome.host = () => { throw Object.assign(new Error(message), { code: INVALID_SETTING_CODE }); };
drafts.set('host', 'autotag_general_threshold', 1.5);
assert.deepEqual(await drafts.save(), { browser: 'clean', host: 'failed', runtime: 'clean', invalid: { source: null, message } });
assert.equal(drafts.statusText(), message);
assert.equal(drafts.invalidSource(), null);
assert.deepEqual([...drafts.failed('host')], ['autotag_general_threshold']);
assert.equal(drafts.get('host', 'autotag_general_threshold'), 1.5);
drafts.set('host', 'autotag_general_threshold', 0.4);
assert.equal(drafts.statusText(), 'Unsaved changes.');
});
test('a host save refused for one source says the server reason and names it until the next edit', async () => {
const domains = adapters({ host: { gallery_root_paths: ['output', 'output/gone'] } });
const drafts = await loaded(domains);
+600
View File
@@ -0,0 +1,600 @@
// System › Auto-tagging (#328, System S11 Tasks 6–7): the concept's rows with real data (D132 model options and live
// status, D133 suggestions per image), host drafts saved by the one Save row, and the actions: Download (an Activity
// job), Unload (direct), Preview tagging review (D115) and the reviewed Tag output images batch (D134).
const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const path = require('node:path');
const tagging = import('../../web/js/workspace/system-tagging.mjs');
const stateModule = import('../../web/js/workspace/system-state.mjs');
const coreModule = import('../../web/js/shared/toolkit/core.js');
const CONCEPT = fs.readFileSync(path.join(__dirname, '../../docs/concepts/workspace/app.js'), 'utf8');
const SOURCE = fs.readFileSync(path.join(__dirname, '../../web/js/workspace/system-tagging.mjs'), 'utf8');
/** A minimal DOM whose elements count their live listeners (a listener's signal aborts it, as in a browser). */
function fakeDOM() {
const previous = global.document;
const elements = [];
class Element {
constructor(tag) {
this.tagName = tag.toUpperCase();
this.children = [];
this.attributes = {};
this.listeners = [];
this.dataset = {};
this.className = '';
this.id = '';
this.value = '';
this.disabled = false;
this._text = '';
elements.push(this);
}
get localName() { return this.tagName.toLowerCase(); }
get options() { return this.children.filter(child => child.tagName === 'OPTION'); }
get textContent() { return this._text + this.children.map(child => child.textContent).join(''); }
set textContent(text) { this.children = []; this._text = String(text); }
append(...nodes) {
for (const node of nodes) {
node.parentNode?.children && (node.parentNode.children = node.parentNode.children.filter(child => child !== node));
node.parentNode = this;
this.children.push(node);
}
}
remove() { if (this.parentNode) this.parentNode.children = this.parentNode.children.filter(child => child !== this); this.parentNode = null; }
setAttribute(name, value) { this.attributes[name] = String(value); }
getAttribute(name) { return Object.hasOwn(this.attributes, name) ? this.attributes[name] : null; }
removeAttribute(name) { delete this.attributes[name]; }
addEventListener(type, callback, options = {}) {
if (options.signal?.aborted) return;
const entry = { type, callback };
this.listeners.push(entry);
options.signal?.addEventListener('abort', () => { this.listeners = this.listeners.filter(item => item !== entry); });
}
removeEventListener(type, callback) { this.listeners = this.listeners.filter(item => item.type !== type || item.callback !== callback); }
dispatch(type) { return Promise.all(this.listeners.filter(item => item.type === type).map(({ callback }) => callback({ type, target: this }))); }
}
global.document = { createElement: tag => new Element(tag), activeElement: null };
const walk = (node, found = []) => { found.push(node); node.children.forEach(child => walk(child, found)); return found; };
return {
listenerCount: () => elements.reduce((sum, element) => sum + element.listeners.length, 0),
find: (root, predicate) => walk(root).find(predicate),
restore() { global.document = previous; },
};
}
const HOST = { autotag_model_type: 'wd14-swinv2', autotag_existing: 'skip', autotag_general_threshold: 0.35,
autotag_character_threshold: 0.85, autotag_suggestion_limit: 10, autotag_template: 'Describe the image as tags.' };
const entry = (label, size, extra = {}) => ({ name: label, label, size_label: size, downloaded: true, loaded: false, ...extra });
const MODELS = {
success: true,
models: { 'wd14-swinv2': entry('WD14 SwinV2', '~400 MB'), 'wd14-vit': entry('WD14 ViT', '~400 MB', { downloaded: false }),
gguf: entry('GGUF JoyCaption', '~4.5 GB', { downloaded: false, model_exists: true, mmproj_exists: false }),
hf: entry('HuggingFace JoyCaption', '~16 GB', { downloaded: false }) },
model_loaded: false, loaded_model_type: null,
};
const tick = () => new Promise(resolve => setImmediate(resolve));
const IMAGE = { filename: 'dunes_00001.png' };
async function mounted({ host = HOST, models = MODELS, modelsFail = false, image = IMAGE, start, post, hostFails = false } = {}) {
const { mountTagging } = await tagging;
const { createSystemDrafts } = await stateModule;
const { createScope } = await coreModule;
const dom = fakeDOM();
const drafts = createSystemDrafts({
readBrowser: () => ({}), readRuntime: () => ({}),
readHost: async () => { if (hostFails) throw new Error('offline'); return { values: host, revision: 'sha256:one' }; },
writeBrowser: async () => ({ persisted: true }), writeHost: async () => ({ success: true, persisted: true }),
writeRuntime: async () => ({}),
});
const scope = createScope();
const calls = { get: [], post: [], toasts: [], navigated: [], opened: [], shown: [], started: [] };
let answer = models;
const api = {
get: async (url, options) => {
calls.get.push(url);
assert.equal(options.cache, 'no-store');
if (modelsFail) throw new Error('offline');
return answer;
},
post: async (url, body) => { calls.post.push([url, body]); return post ? post(url, body) : { success: true }; },
};
const activity = { start: async (kind, body) => { calls.started.push([kind, body]); return start ? start(kind, body) : { job: { id: 'job1' } }; } };
const [surface] = mountTagging(null, { api, activity, scope, drafts, notify: text => calls.toasts.push(text),
navigate: hash => calls.navigated.push(hash), selected: () => image, openTagging: options => calls.opened.push(options),
show: (plan, options) => calls.shown.push([plan, options]) });
const byId = id => dom.find(surface, node => node.id === id);
const button = label => dom.find(surface, node => node.tagName === 'BUTTON' && node.textContent === label);
const description = () => byId('tag-model').parentNode.children[0].children[1].textContent;
await drafts.reloadHost();
await tick();
return { dom, drafts, scope, surface, byId, button, description, calls, setModels: value => { answer = value; } };
}
test('the section keeps the concept surface, rows, labels and actions in order', () => {
const start = CONCEPT.indexOf('case "tagging":');
const body = CONCEPT.slice(start, CONCEPT.indexOf('break;', start));
const rows = [...body.matchAll(/setting\("([^"]+)", "([^"]+)", "([^"]+)"/g)].map(([, id, label, description]) => [id, label, description]);
assert.deepEqual(rows.map(([id]) => id), ['tag-model', 'tag-existing', 'threshold-general', 'threshold-character']);
for (const [id, label] of rows) assert.ok(SOURCE.includes(`id: '${id}', label: '${label}'`), `${id} keeps "${label}"`);
// OQ6: the true descriptions are kept verbatim; the model's design note is replaced by its live status (D132).
for (const [id, , description] of rows.slice(1)) assert.ok(SOURCE.includes(`'${description}'`), `${id} description`);
assert.ok(!SOURCE.includes(rows[0][2]));
for (const option of ['Skip already tagged', 'Merge with existing', 'Replace after review']) assert.ok(body.includes(option) && SOURCE.includes(`'${option}'`));
assert.ok(body.includes('<h2>Auto-tagging</h2>') && SOURCE.includes("'h2', 'Auto-tagging'"));
assert.ok(body.includes('One tagging configuration for images and prompt-library actions. Every result can be reviewed before applying.'));
assert.ok(body.includes('LLM tag template') && SOURCE.includes("'LLM tag template'"));
assert.ok(body.includes('"Preview tagging review", "tag-review", "", "ca-primary"'));
// Prototype copy is never shipped.
for (const sample of ['Simulate download', 'Simulate unload', 'Ollama']) assert.ok(!SOURCE.includes(sample), sample);
});
test('model options are the four real engines with their size, in the concept order (D132, OQ5)', async () => {
const { modelOptions, MODEL_IDS, modelLabel, modelSize } = await tagging;
assert.deepEqual(MODEL_IDS, ['wd14-swinv2', 'wd14-vit', 'gguf', 'hf']);
assert.deepEqual(modelOptions(MODELS.models).map(([, label]) => label),
['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB']);
// Before the status answered (or when it cannot be read) the documented labels stand in.
assert.deepEqual(modelOptions(null), modelOptions(MODELS.models));
assert.equal(modelLabel('hf', { label: '' }), 'HuggingFace JoyCaption');
assert.equal(modelLabel('other', null), 'other');
assert.equal(modelSize('other', null), '');
});
test('the Tagging model description is the chosen model’s live status (D132)', async () => {
const { modelDescription, readStatus, STATUS_COPY } = await tagging;
const status = readStatus(MODELS);
assert.equal(modelDescription({ phase: 'loading' }, 'wd14-vit'), 'Loading model status…');
assert.equal(modelDescription({ phase: 'failed' }, 'wd14-vit'), 'Model status could not be read.');
assert.equal(modelDescription(status, 'wd14-swinv2'), 'Downloaded. Not loaded.');
assert.equal(modelDescription(status, 'wd14-vit'), 'Not downloaded. Download size ~400 MB.');
assert.equal(modelDescription(status, 'hf'), 'Not downloaded. Download size ~16 GB.');
assert.equal(modelDescription(status, 'gguf'), 'Partly downloaded. Download again to finish.');
const both = readStatus({ ...MODELS, models: { ...MODELS.models, gguf: entry('GGUF JoyCaption', '~4.5 GB', { downloaded: false,
model_exists: false, mmproj_exists: false }) } });
assert.equal(modelDescription(both, 'gguf'), 'Not downloaded. Download size ~4.5 GB.');
const loaded = readStatus({ ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2',
models: { ...MODELS.models, 'wd14-swinv2': entry('WD14 SwinV2', '~400 MB', { loaded: true }) } });
assert.equal(loaded.loadedType, 'wd14-swinv2');
assert.equal(modelDescription(loaded, 'wd14-swinv2'), STATUS_COPY.loaded);
assert.equal(modelDescription(status, 'unknown'), STATUS_COPY.failed);
assert.deepEqual(readStatus(null), { phase: 'failed', models: null, loadedType: null });
assert.equal(readStatus({ models: {}, model_loaded: true, loaded_model_type: '/abs/path' }).loadedType, null);
});
test('controls draft their host keys; numbers as numbers, other text as typed', async () => {
const { SETTING_KEYS, coerceNumber, draftValue } = await tagging;
assert.deepEqual(SETTING_KEYS, { 'tag-model': 'autotag_model_type', 'tag-existing': 'autotag_existing',
'threshold-general': 'autotag_general_threshold', 'threshold-character': 'autotag_character_threshold',
'tag-suggestions': 'autotag_suggestion_limit', 'tag-template': 'autotag_template' });
assert.deepEqual(['0.4', ' 1.5 ', '15', '', 'abc', '1e-1'].map(coerceNumber), [0.4, 1.5, 15, '', 'abc', 0.1]);
assert.equal(draftValue('autotag_suggestion_limit', '15'), 15);
assert.equal(draftValue('autotag_template', ' 12 '), ' 12 ');
assert.equal(draftValue('autotag_existing', undefined), '');
const { dom, drafts, scope, byId } = await mounted();
try {
assert.equal(byId('tag-model').value, 'wd14-swinv2');
assert.equal(byId('threshold-general').value, '0.35');
assert.equal(byId('tag-template').value, HOST.autotag_template);
assert.deepEqual([byId('threshold-general').min, byId('threshold-general').max, byId('threshold-general').step], ['0', '1', '0.05']);
assert.deepEqual([byId('tag-suggestions').min, byId('tag-suggestions').max], ['1', '100']);
byId('tag-existing').value = 'merge'; await byId('tag-existing').dispatch('change');
byId('threshold-general').value = '0.4'; await byId('threshold-general').dispatch('change');
byId('tag-suggestions').value = '15'; await byId('tag-suggestions').dispatch('change');
byId('tag-template').value = 'Tags only.'; await byId('tag-template').dispatch('input');
assert.deepEqual(drafts.changes('host'), { autotag_existing: 'merge', autotag_general_threshold: 0.4,
autotag_suggestion_limit: 15, autotag_template: 'Tags only.' });
// Every control names its stored key, so a refused save highlights it (D92).
for (const id of Object.keys(SETTING_KEYS)) assert.equal(byId(id).dataset.key, SETTING_KEYS[id]);
assert.equal(byId('tag-template').dataset.setting, 'tag-template');
} finally { scope.dispose(); assert.equal(dom.listenerCount(), 0); dom.restore(); }
});
test('opening the section reads the model status once; the select waits for it', async () => {
const { dom, scope, byId, description, calls, surface } = await mounted();
try {
assert.deepEqual(calls.get, ['/prompt_manager/autotag/models']);
assert.deepEqual(calls.post, []);
assert.equal(byId('tag-model').disabled, false);
assert.equal(description(), 'Downloaded. Not loaded.');
assert.deepEqual(byId('tag-model').options.map(option => option.textContent),
['WD14 SwinV2 · ~400 MB', 'WD14 ViT · ~400 MB', 'GGUF JoyCaption · ~4.5 GB', 'HuggingFace JoyCaption · ~16 GB']);
assert.deepEqual(surface.children.map(node => node.className),
['', '', 'setting-row', 'setting-row', 'setting-row', 'setting-row', 'setting-row', 'field', 'row-actions', 'operation']);
} finally { scope.dispose(); dom.restore(); }
const failed = await mounted({ modelsFail: true });
try {
assert.equal(failed.description(), 'Model status could not be read.');
assert.equal(failed.button('Download').getAttribute('aria-disabled'), 'true');
assert.equal(failed.button('Unload').disabled, true);
await failed.button('Download').dispatch('click');
assert.deepEqual(failed.calls.toasts, ['Model status could not be read.']);
assert.deepEqual(failed.calls.started, []);
} finally { failed.scope.dispose(); failed.dom.restore(); }
});
test('server settings that cannot be read leave the rows disabled and say so', async () => {
const { dom, scope, byId, button, description, calls } = await mounted({ hostFails: true });
try {
assert.equal(byId('tag-existing').disabled, true);
assert.equal(byId('tag-existing').parentNode.children[0].children.at(-1).textContent, 'Server settings could not be read.');
assert.equal(description(), 'Model status could not be read.');
await button('Preview tagging review').dispatch('click');
assert.deepEqual(calls.toasts, ['Server settings could not be read.']);
} finally { scope.dispose(); dom.restore(); }
});
test('Download starts the model download job and opens it; the label follows the disk (OQ10)', async () => {
const { dom, drafts, scope, button, calls, byId } = await mounted();
try {
assert.ok(button('Download again'));
byId('tag-model').value = 'wd14-vit'; await byId('tag-model').dispatch('change');
assert.equal(drafts.get('host', 'autotag_model_type'), 'wd14-vit');
await button('Download').dispatch('click');
assert.deepEqual(calls.started, [['model-download', { model_type: 'wd14-vit' }]]);
assert.deepEqual(calls.navigated, ['#/activity/job1']);
} finally { scope.dispose(); dom.restore(); }
});
test('startJob toasts by code and never navigates on a refusal', async () => {
const { startJob, startRefusal } = await tagging;
const { createScope } = await coreModule;
const error = (code, message = 'Server says.', extra = {}) => Object.assign(new Error(message), { code, ...extra });
assert.equal(startRefusal(error('model-busy', 'x', { operation: 'Tagging model download · WD14 ViT' }), 'X'),
'The tagging model is busy: Tagging model download · WD14 ViT. Try again when it finishes.');
assert.equal(startRefusal(error('model-busy', 'x'), 'X'),
'The tagging model is busy: Another tagging operation. Try again when it finishes.');
for (const code of ['model-unavailable', 'review-required', 'insufficient-space']) assert.equal(startRefusal(error(code), 'X'), 'Server says.');
assert.equal(startRefusal(error('weird'), 'Tag output images'), 'Tag output images could not start. Try again.');
assert.equal(startRefusal(new Error('offline'), 'Tag output images'), 'Tag output images could not start. Try again.');
const scope = createScope();
const toasts = [];
const navigated = [];
const run = start => startJob({ activity: { start }, scope, kind: 'autotag', body: {}, what: 'Tag output images',
notify: text => toasts.push(text), navigate: hash => navigated.push(hash) });
assert.equal(await run(async () => { throw error('model-busy', 'busy', { operation: 'Generate suggestions' }); }), null);
assert.equal(await run(async () => ({ busy: { id: 'a1', title: 'Scan for metadata' } })), null);
assert.equal(await run(async () => ({ busy: null })), null);
assert.deepEqual(toasts, ['The tagging model is busy: Generate suggestions. Try again when it finishes.',
['Another job is running: ', { text: 'Scan for metadata', href: '#/activity/a1' }, '.'],
'Another job is running. Try again when it ends.']);
assert.deepEqual(await run(async () => ({ job: { id: 'j2' } })), { id: 'j2' });
assert.deepEqual(navigated, ['#/activity/j2']);
scope.dispose();
assert.equal(await run(async () => { throw error('weird'); }), null);
assert.deepEqual(await run(async () => ({ job: { id: 'j3' } })), { id: 'j3' });
assert.deepEqual(navigated, ['#/activity/j2']);
assert.equal(toasts.length, 3);
});
test('Unload is enabled only while a model is loaded; a busy refusal names the operation', async () => {
const loadedModels = { ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2' };
let refuse = null;
const post = url => {
if (refuse) throw refuse;
return url.endsWith('/unload') ? { success: true, message: 'WD14SWINV2 model unloaded successfully', model_loaded: false } : {};
};
const { dom, scope, button, calls, setModels } = await mounted({ models: loadedModels, post });
try {
assert.equal(button('Unload').disabled, false);
refuse = Object.assign(new Error('A tagging model is busy'), { code: 'model-busy', operation: 'Tagging model download · WD14 ViT' });
await button('Unload').dispatch('click');
assert.deepEqual(calls.toasts, ['The tagging model is busy: Tagging model download · WD14 ViT. Try again when it finishes.']);
refuse = new Error('offline');
await button('Unload').dispatch('click');
assert.equal(calls.toasts.at(-1), 'The tagging model could not be unloaded. Try again.');
refuse = null;
setModels(MODELS);
await button('Unload').dispatch('click');
assert.equal(calls.toasts.at(-1), 'WD14 SwinV2 unloaded.');
assert.deepEqual(calls.post.map(([url]) => url), Array(3).fill('/prompt_manager/autotag/unload'));
// Each attempt reads the status again; nothing is loaded now.
assert.equal(calls.get.length, 4);
assert.equal(button('Unload').disabled, true);
} finally { scope.dispose(); dom.restore(); }
});
test('Preview tagging review opens the Tag image dialog with the section’s values, else says why (D115)', async () => {
const { NO_SELECTION, taggingPreset, draftProblem } = await tagging.then(async module => ({ ...module,
NO_SELECTION: (await import('../../web/js/workspace/vocabulary.mjs')).NO_SELECTION }));
assert.equal(NO_SELECTION, 'Select an image in Images first.');
assert.deepEqual(taggingPreset(HOST), { model: 'wd14-swinv2', options: { general_threshold: 0.35, character_threshold: 0.85,
suggestion_limit: 10 } });
assert.deepEqual(taggingPreset({ ...HOST, autotag_model_type: 'gguf' }), { model: 'gguf',
options: { instructions: HOST.autotag_template, suggestion_limit: 10 } });
assert.equal(draftProblem(HOST), null);
assert.deepEqual(draftProblem({ ...HOST, autotag_general_threshold: 1.5 }),
{ key: 'autotag_general_threshold', message: 'General-tag confidence must be between 0 and 1.' });
assert.equal(draftProblem({ ...HOST, autotag_character_threshold: '' }).key, 'autotag_character_threshold');
// #369: every page shows the thresholds in steps of 0.05, so a value off that grid cannot be saved or previewed.
assert.deepEqual(draftProblem({ ...HOST, autotag_general_threshold: 0.37 }),
{ key: 'autotag_general_threshold', message: 'General-tag confidence must be a multiple of 0.05.' });
assert.deepEqual(draftProblem({ ...HOST, autotag_character_threshold: 0.33 }),
{ key: 'autotag_character_threshold', message: 'Character-tag confidence must be a multiple of 0.05.' });
assert.equal(draftProblem({ ...HOST, autotag_general_threshold: 0.15, autotag_character_threshold: 0.95 }), null);
assert.equal(draftProblem({ ...HOST, autotag_general_threshold: 0, autotag_character_threshold: 1 }), null);
// A computed value a hair off the grid passes and is handed on snapped to it, as the server stores it.
assert.equal(draftProblem({ ...HOST, autotag_general_threshold: 0.1 + 0.2 }), null);
assert.equal(taggingPreset({ ...HOST, autotag_general_threshold: 0.1 + 0.2 }).options.general_threshold, 0.3);
assert.equal(draftProblem({ ...HOST, autotag_suggestion_limit: 2.5 }).key, 'autotag_suggestion_limit');
assert.equal(draftProblem({ ...HOST, autotag_suggestion_limit: 0 }, { limit: false }), null);
assert.deepEqual(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: ' ' }),
{ key: 'autotag_template', message: 'LLM tag template must be non-empty text.' });
assert.deepEqual(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: 'x'.repeat(4001) }),
{ key: 'autotag_template', message: 'LLM tag template must be at most 4000 characters.' });
// Characters are code points, as Python's len counts them: 4000 emoji (8000 UTF-16 units) are allowed.
assert.equal(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: '🏜'.repeat(4000) }), null);
assert.equal(draftProblem({ ...HOST, autotag_model_type: 'hf', autotag_template: '🏜'.repeat(4001) }).key, 'autotag_template');
// A WD14 run never reads the template, nor a language model the thresholds.
assert.equal(draftProblem({ ...HOST, autotag_template: '' }), null);
assert.equal(draftProblem({ ...HOST, autotag_model_type: 'gguf', autotag_general_threshold: 7 }), null);
const none = await mounted({ image: null });
try {
const preview = none.button('Preview tagging review');
assert.equal(preview.getAttribute('aria-disabled'), 'true');
assert.equal(preview.className, 'ca-primary');
await preview.dispatch('click');
assert.deepEqual(none.calls.toasts, ['Select an image in Images first.']);
assert.deepEqual(none.calls.opened, []);
} finally { none.scope.dispose(); none.dom.restore(); }
const { dom, scope, button, byId, calls } = await mounted();
try {
await button('Preview tagging review').dispatch('click');
assert.equal(calls.opened.length, 1);
assert.equal(calls.opened[0].image, IMAGE);
assert.equal(calls.opened[0].source, 'tags');
assert.deepEqual(calls.opened[0].preset, taggingPreset(HOST));
byId('threshold-general').value = '1.5'; await byId('threshold-general').dispatch('change');
assert.equal(button('Preview tagging review').getAttribute('aria-disabled'), 'true');
await button('Preview tagging review').dispatch('click');
assert.deepEqual(calls.toasts, ['General-tag confidence must be between 0 and 1.']);
assert.equal(calls.opened.length, 1);
byId('threshold-general').value = '0.37'; await byId('threshold-general').dispatch('change');
assert.equal(button('Preview tagging review').getAttribute('aria-disabled'), 'true');
await button('Preview tagging review').dispatch('click');
assert.deepEqual(calls.toasts.at(-1), 'General-tag confidence must be a multiple of 0.05.');
assert.equal(calls.opened.length, 1);
} finally { scope.dispose(); dom.restore(); }
});
const PLAN = { images: 3, prompts: 2, already_tagged: 1, to_tag: 1,
model: { type: 'wd14-swinv2', label: 'WD14 SwinV2', size_label: '~400 MB', downloaded: true }, existing: 'skip',
review_required: false, backup: { policy: 'on' }, disk: { free_bytes: 10 * 2 ** 30, required_bytes: 2 ** 20, sufficient: true } };
test('the batch review grid, body and Cancel-only reasons (D134, CUJ 4–5)', async () => {
const { batchGrid, batchBody, batchCloseNote, existingLabel, REPLACE_NOTE, NOTHING_TO_TAG, NO_LINKED_IMAGES } = await tagging;
assert.deepEqual(batchGrid(PLAN, HOST), [['Model', 'WD14 SwinV2 · ~400 MB'], ['Existing tags', 'Skip already tagged'],
['Images linked', '3'], ['Prompts', '2'], ['Already tagged', '1'], ['Prompts to tag', 'up to 1'],
['Settings', 'General 0.35 · Character 0.85'], ['Backup', 'Create and verify before changes']]);
const llm = { ...HOST, autotag_model_type: 'gguf' };
assert.deepEqual(batchGrid({ ...PLAN, model: { type: 'gguf' }, backup: { policy: 'off' } }, llm).slice(-2),
[['Settings', 'Template of 27 characters'], ['Backup', 'No safety backup (the policy is off)']]);
assert.equal(batchGrid({}, HOST)[2][1], '0');
assert.equal(existingLabel('odd'), 'odd');
assert.deepEqual(batchBody(HOST), { model_type: 'wd14-swinv2', existing: 'skip',
options: { general_threshold: 0.35, character_threshold: 0.85 } });
// #369: a threshold a hair off the grid is sent snapped to it.
assert.equal(batchBody({ ...HOST, autotag_character_threshold: 0.1 + 0.2 }).options.character_threshold, 0.3);
assert.deepEqual(batchBody(llm), { model_type: 'gguf', existing: 'skip', options: { instructions: HOST.autotag_template } });
assert.equal(batchCloseNote(PLAN), null);
assert.equal(batchCloseNote({ ...PLAN, review_required: true }), REPLACE_NOTE);
assert.equal(REPLACE_NOTE, 'Replacing tags needs a review of each image. Use Preview tagging review, or choose Skip or Merge for a batch run.');
assert.equal(batchCloseNote({ ...PLAN, model: { ...PLAN.model, downloaded: false } }), 'WD14 SwinV2 is not downloaded. Download it first.');
assert.equal(batchCloseNote({ ...PLAN, disk: { free_bytes: 2 ** 20, required_bytes: 2 ** 21, sufficient: false } }),
'Not enough free disk space: the safety backup needs about 2.0 MB and 1.0 MB is free. Nothing is written until there is room.');
assert.equal(batchCloseNote({ ...PLAN, prompts: 0, to_tag: 0 }), NO_LINKED_IMAGES);
assert.equal(batchCloseNote({ ...PLAN, to_tag: 0 }), NOTHING_TO_TAG);
assert.equal(NOTHING_TO_TAG, 'Every prompt with an output image is already tagged.');
});
test('showBatchPlan opens the review with the backup box, or the Cancel-only review', async () => {
const { showBatchPlan, REVIEW_TITLE, BATCH_SENTENCE } = await tagging;
const { createScope } = await coreModule;
const scope = createScope();
const reviews = [];
const closed = [];
const started = [];
const options = { values: HOST, scope, activity: { start: async (kind, body) => { started.push([kind, body]); return { job: { id: 'b1' } }; } },
notify: () => {}, navigate: () => {}, openReview: review => reviews.push(review), openCloseOnly: review => closed.push(review) };
showBatchPlan(PLAN, options);
assert.equal(reviews.length, 1);
const [review] = reviews;
assert.equal(review.title, REVIEW_TITLE);
assert.equal(review.title, 'Review tag output images');
assert.equal(review.heading, 'Tag output images');
assert.equal(review.sentence, BATCH_SENTENCE);
assert.equal(review.primary.label, 'Start tagging');
assert.equal(review.cancelLabel, 'Cancel');
assert.equal(review.backup.copy.heading, 'Back up the database first (recommended)');
await review.onStart();
assert.deepEqual(started, [['autotag', { model_type: 'wd14-swinv2', existing: 'skip',
options: { general_threshold: 0.35, character_threshold: 0.85 } }]]);
showBatchPlan({ ...PLAN, backup: { policy: 'off' } }, options);
assert.equal(reviews[1].backup, null);
showBatchPlan({ ...PLAN, review_required: true }, options);
assert.equal(closed.length, 1);
assert.equal(closed[0].note.startsWith('Replacing tags needs a review'), true);
assert.equal(closed[0].grid[0][0], 'Model');
scope.dispose();
});
test('Tag output images › Review previews the drafts, else says why; a refused preview is toasted', async () => {
let refuse = null;
const post = url => {
if (refuse) throw refuse;
assert.equal(url, '/prompt_manager/autotag/batch/preview');
return PLAN;
};
const { dom, scope, button, byId, calls } = await mounted({ post });
try {
const review = button('Review');
assert.equal(review.getAttribute('aria-disabled'), null);
await review.dispatch('click');
await tick();
assert.deepEqual(calls.post, [['/prompt_manager/autotag/batch/preview', { model_type: 'wd14-swinv2', existing: 'skip' }]]);
assert.equal(calls.shown.length, 1);
assert.equal(calls.shown[0][0], PLAN);
assert.deepEqual(calls.shown[0][1].values, HOST);
assert.equal(review.textContent, 'Review');
refuse = Object.assign(new Error('busy'), { code: 'busy', running: { id: 'r1', title: 'Optimize database' } });
await review.dispatch('click');
await tick();
assert.deepEqual(calls.toasts.at(-1), ['Another job is running: ', { text: 'Optimize database', href: '#/activity/r1' }, '.']);
refuse = new Error('offline');
await review.dispatch('click');
await tick();
assert.equal(calls.toasts.at(-1), 'The plan could not be reviewed. Nothing was changed.');
// A model that is not on disk cannot be reviewed: the button explains instead.
byId('tag-model').value = 'wd14-vit'; await byId('tag-model').dispatch('change');
assert.equal(review.getAttribute('aria-disabled'), 'true');
await review.dispatch('click');
assert.equal(calls.toasts.at(-1), 'WD14 ViT is not downloaded. Download it first.');
assert.equal(calls.post.length, 3);
} finally { scope.dispose(); dom.restore(); }
});
test('settingRefusal marks only an Auto-tagging save\'s plain 400 as invalid-setting (CUJ 2a)', async () => {
const { settingRefusal } = await import('../../web/js/workspace/system.mjs');
const refusal = (extra = {}) => Object.assign(new Error('General-tag confidence must be between 0 and 1.'),
{ status: 400, refused: true, code: null, ...extra });
const tagging = { autotag_general_threshold: 1.5, expected_revision: 'sha256:one' };
assert.equal(settingRefusal(refusal(), tagging).code, 'invalid-setting');
// Another section's refusal keeps the generic copy, even with a plain 400.
const other = refusal({ message: 'Result timeout must be between 0 and 300.' });
assert.equal(settingRefusal(other, { result_timeout: 999, expected_revision: 'sha256:one' }).code, null);
assert.equal(settingRefusal(refusal(), undefined).code, null);
for (const extra of [{ status: 409 }, { refused: false }, { code: 'jxl-unavailable' }]) {
assert.notEqual(settingRefusal(refusal(extra), tagging).code, 'invalid-setting');
}
assert.equal(settingRefusal(undefined, tagging), undefined);
});
test('a part-downloaded GGUF offers Download again', async () => {
const { downloadLabel, partlyDownloaded } = await tagging;
const partial = MODELS.models.gguf;
assert.equal(partlyDownloaded('gguf', partial), true);
assert.equal(downloadLabel('gguf', partial), 'Download again');
assert.equal(downloadLabel('gguf', { ...partial, model_exists: false }), 'Download');
assert.equal(downloadLabel('hf', MODELS.models.hf), 'Download');
assert.equal(downloadLabel('wd14-swinv2', MODELS.models['wd14-swinv2']), 'Download again');
assert.equal(partlyDownloaded('hf', { model_exists: true, mmproj_exists: false }), false);
const { dom, scope, byId, button, description } = await mounted();
try {
byId('tag-model').value = 'gguf'; await byId('tag-model').dispatch('change');
assert.equal(description(), 'Partly downloaded. Download again to finish.');
assert.ok(button('Download again'));
} finally { scope.dispose(); dom.restore(); }
});
test('the disk note says what is short in the server\'s words, and guesses no unknown byte counts', async () => {
const { diskShortNote, DISK_SHORT_UNKNOWN } = await tagging;
assert.equal(diskShortNote({ required_bytes: 3 * 2 ** 20, free_bytes: 2 ** 19, sufficient: false }),
'Not enough free disk space: the safety backup needs about 3.0 MB and 0.5 MB is free. Nothing is written until there is room.');
for (const disk of [{ required_bytes: NaN, free_bytes: 1 }, { required_bytes: 1, free_bytes: Infinity }, { required_bytes: -1, free_bytes: 1 },
{ required_bytes: '1', free_bytes: 1 }, null]) {
assert.equal(diskShortNote(disk), DISK_SHORT_UNKNOWN);
}
assert.equal(DISK_SHORT_UNKNOWN, 'Not enough free disk space for the safety backup. Nothing is written until there is room.');
assert.ok(!diskShortNote({ required_bytes: 1, free_bytes: 0 }).includes('apply'));
});
test('a drafts repaint never re-enables Download while its start is in flight', async () => {
let answer;
const start = () => new Promise(resolve => { answer = resolve; });
const { dom, scope, byId, button, calls } = await mounted({ start });
try {
const download = button('Download again');
const pressed = download.dispatch('click');
await tick();
assert.equal(download.getAttribute('aria-disabled'), 'true');
byId('tag-existing').value = 'merge'; await byId('tag-existing').dispatch('change');
assert.equal(download.getAttribute('aria-disabled'), 'true');
// A second press while it starts is ignored.
await download.dispatch('click');
assert.equal(calls.started.length, 1);
answer({ job: { id: 'd1' } });
await pressed;
assert.equal(download.getAttribute('aria-disabled'), null);
assert.deepEqual(calls.navigated, ['#/activity/d1']);
} finally { scope.dispose(); dom.restore(); }
});
test('an Unload that found nothing loaded says nothing and reads the status again', async () => {
const { dom, scope, button, calls } = await mounted({ models: { ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2' },
post: () => ({ success: true, message: 'No model was loaded' }) });
try {
await button('Unload').dispatch('click');
assert.deepEqual(calls.toasts, []);
assert.equal(calls.get.length, 2);
} finally { scope.dispose(); dom.restore(); }
});
test('an unknown saved engine has its own reason', async () => {
const { UNKNOWN_MODEL } = await tagging;
const { dom, scope, button, calls } = await mounted({ host: { ...HOST, autotag_model_type: 'wd14-convnext' } });
try {
await button('Preview tagging review').dispatch('click');
await button('Download').dispatch('click');
await button('Review').dispatch('click');
assert.deepEqual(calls.toasts, [UNKNOWN_MODEL, UNKNOWN_MODEL, UNKNOWN_MODEL]);
assert.equal(UNKNOWN_MODEL, 'The saved tagging model is not one of these engines. Choose one and save preferences.');
assert.deepEqual(calls.started, []);
} finally { scope.dispose(); dom.restore(); }
});
test('a review whose dialog fails to open is toasted, not lost', async () => {
const { dom, scope, button, calls } = await mounted({ post: () => PLAN });
try {
calls.shown.push = () => { throw new Error('dialog failed'); };
await button('Review').dispatch('click');
await tick();
assert.equal(calls.toasts.at(-1), 'The plan could not be reviewed. Nothing was changed.');
assert.equal(button('Review').getAttribute('aria-disabled'), null);
} finally { scope.dispose(); dom.restore(); }
});
test('a status read, an unload and a preview that finish after the section left do nothing', async () => {
const { mountTagging } = await tagging;
const { createSystemDrafts } = await stateModule;
const { createScope } = await coreModule;
const dom = fakeDOM();
const pending = [];
const later = value => new Promise(resolve => pending.push(() => resolve(value)));
const drafts = createSystemDrafts({ readBrowser: () => ({}), readRuntime: () => ({}),
readHost: async () => ({ values: HOST, revision: 'sha256:one' }), writeBrowser: async () => ({}), writeHost: async () => ({}),
writeRuntime: async () => ({}) });
await drafts.load('host');
const scope = createScope();
const toasts = [];
const shown = [];
let gets = 0;
const loaded = { ...MODELS, model_loaded: true, loaded_model_type: 'wd14-swinv2',
models: { ...MODELS.models, 'wd14-swinv2': entry('WD14 SwinV2', '~400 MB', { loaded: true }) } };
const api = { get: () => { gets += 1; return gets === 1 ? Promise.resolve(loaded) : later(MODELS); }, post: url => later(url.endsWith('/unload') ? { success: true, model_loaded: false } : PLAN) };
const [surface] = mountTagging(null, { api, activity: {}, scope, drafts, notify: text => toasts.push(text), navigate: () => {},
selected: () => IMAGE, openTagging: () => {}, show: plan => shown.push(plan) });
try {
await tick();
const find = label => dom.find(surface, node => node.tagName === 'BUTTON' && node.textContent === label);
const description = () => dom.find(surface, node => node.id === 'tag-model').parentNode.children[0].children[1].textContent;
assert.equal(description(), 'Downloaded and loaded.');
const unloading = find('Unload').dispatch('click');
const reviewing = find('Review').dispatch('click');
await tick();
scope.dispose();
for (const settle of pending.splice(0)) settle();
await unloading;
await reviewing;
await tick();
for (const settle of pending.splice(0)) settle();
await tick();
assert.deepEqual(toasts, []);
assert.deepEqual(shown, []);
// The status read after the unload never reached the page either.
assert.equal(description(), 'Downloaded and loaded.');
assert.equal(dom.listenerCount(), 0);
} finally { dom.restore(); }
});
+1 -1
View File
@@ -36,7 +36,7 @@ SUGGEST = "/prompt_manager/image-annotations/suggest"
WAIT = 30.0
WD14 = {"general_threshold": 0.35, "character_threshold": 0.85}
LLM = {"instructions": "Generate tags."}
UNAVAILABLE = "Download it from the legacy Auto-tag page."
UNAVAILABLE = "Download it in System › Auto-tagging."
SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [
(f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150)
+2
View File
@@ -1145,8 +1145,10 @@ class TestEnvelopes(AnnotationAPITestCase):
sorted(await self.json_of(resp)),
[
"default_prompt",
"existing", # #328 adds the saved batch policy and model
"loaded_model_type",
"model_loaded",
"model_type",
"models",
"success",
"suggestion_limit", # B2 adds it (D68); nothing else changes
+1066 -3
View File
File diff suppressed because it is too large Load Diff
+409
View File
@@ -0,0 +1,409 @@
"""The shareable diagnostic report (#329, py/api/diagnostics_report.py).
``GET /prompt_manager/system/diagnostics`` builds the report on the server
from an explicit list of fields: the project version, the four legacy checks,
public paths only and no secret. A scrub replaces anything that still looks
like an absolute path or holds the CivitAI key and counts it in ``redacted``;
these tests require ``redacted == 0`` so a builder mistake fails instead of
being hidden.
"""
import json
import ntpath
import os
import posixpath
import re
import sqlite3
import sys
import unittest
from unittest import mock
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.api import admin, diagnostics_report # noqa: E402
from py.api.diagnostics_report import OUTSIDE, report_path, scrub # noqa: E402
from py.config import (
GalleryConfig,
IntegrationConfig,
PromptManagerConfig,
) # noqa: E402
try:
from tests.test_api_log_cleanup import LogCleanupTestCase # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from test_api_log_cleanup import LogCleanupTestCase # noqa: E402
REPORT = "/prompt_manager/system/diagnostics"
KEY = "civitai-secret-0123456789abcdef"
SEEDED_PATHS = (
"/home/someone/private/LoraManager",
"/srv/someone/renders",
"C:\\Users\\someone\\Documents\\renders",
"D:/Users/someone/prompts.db",
"\\\\nas\\share\\someone\\renders",
)
_WINDOWS_ABSOLUTE = re.compile(r"^[A-Za-z]:[\\/]")
TOP_KEYS = {
"success",
"generated_at",
"version",
"python",
"platform",
"viewerjs",
"checks",
"database_info",
"sources",
"logging",
"health",
"config",
"redacted",
}
CONFIG = {
"database": {"enable_duplicate_detection", "enable_auto_save"},
"web_ui": {
"result_timeout",
"show_test_button",
"webui_display_mode",
"infinite_scroll",
},
"performance": {
"max_search_results",
"enable_fuzzy_search",
"auto_backup_interval",
"worker_threads",
},
"gallery": {
"processing_delay",
"prompt_timeout",
"cleanup_interval",
"auto_cleanup_missing_files",
"max_image_age_days",
"max_concurrent_processing",
"metadata_extraction_timeout",
},
"integrations": {"lora_manager_enabled", "has_civitai_api_key"},
}
def _strings(value):
if isinstance(value, str):
yield value
elif isinstance(value, dict):
for key, item in value.items():
yield str(key)
yield from _strings(item)
elif isinstance(value, (list, tuple)):
for item in value:
yield from _strings(item)
class DiagnosticsReportTestCase(LogCleanupTestCase):
"""The admin app with its logger manager on a temporary log folder."""
async def get_application(self):
app = await super().get_application()
saved = IntegrationConfig.get_config()
self.addCleanup(IntegrationConfig.update_config, saved)
return app
async def report(self, status=200):
response = await self.client.request("GET", REPORT)
text = await response.text()
self.assertEqual(response.status, status, text)
return json.loads(text), text
def seed_secrets(self):
IntegrationConfig.update_config(
{
"lora_manager": {
"enabled": True,
"path": SEEDED_PATHS[0],
"civitai_api_key": KEY,
}
}
)
GalleryConfig.MONITORING_DIRECTORIES = list(SEEDED_PATHS[1:3]) + [
SEEDED_PATHS[4]
]
PromptManagerConfig.update_config(
{"database": {"default_path": SEEDED_PATHS[3]}}
)
def assert_no_paths_or_key(self, report, text):
self.assertNotIn(KEY, text)
self.assertNotIn(self.tmpdir, text)
self.assertNotIn(self.tmpdir.replace("\\", "/"), text)
for seeded in SEEDED_PATHS:
self.assertNotIn(seeded, text)
self.assertNotIn(json.dumps(seeded)[1:-1], text)
self.assertNotIn("someone", text)
for value in _strings(report):
self.assertFalse(posixpath.isabs(value), value)
self.assertFalse(ntpath.isabs(value), value)
self.assertIsNone(_WINDOWS_ABSOLUTE.match(value), value)
self.assertFalse(value.startswith("\\\\"), value)
class TestReport(DiagnosticsReportTestCase):
async def test_report_has_version_viewerjs_and_the_four_checks(self):
report, _text = await self.report()
self.assertEqual(set(report), TOP_KEYS)
self.assertTrue(report["success"])
with open(
os.path.join(os.path.dirname(os.path.dirname(__file__)), "pyproject.toml"),
encoding="utf-8",
) as handle:
version = re.search(r'version\s*=\s*["\']([^"\']+)["\']', handle.read())
self.assertEqual(report["version"], version.group(1))
self.assertRegex(report["viewerjs"], r"^\d+\.\d+\.\d+$")
self.assertRegex(report["python"], r"^3\.\d+\.\d+")
self.assertEqual(set(report["platform"]), {"system", "release"})
self.assertEqual(
set(report["checks"]),
{"database", "dependencies", "comfyui_output", "image_monitor"},
)
self.assertEqual(report["checks"]["database"]["status"], "ok")
self.assertEqual(
set(report["sources"]),
{"configured", "effective", "watching", "unavailable"},
)
self.assertEqual(set(report["logging"]), {"config", "stats"})
self.assertIsNone(report["health"])
self.assertEqual(report["redacted"], 0)
self.assertRegex(report["generated_at"], r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d$")
async def test_report_has_no_secret_or_absolute_path(self):
self.seed_secrets()
with mock.patch.object(
admin.sqlite3,
"connect",
side_effect=sqlite3.DatabaseError("file is not a database"),
):
report, text = await self.report()
self.assertEqual(report["checks"]["database"]["status"], "error")
self.assertTrue(report["config"]["integrations"]["has_civitai_api_key"])
self.assertTrue(report["config"]["integrations"]["lora_manager_enabled"])
self.assertEqual(report["sources"]["configured"], 3)
self.assert_no_paths_or_key(report, text)
self.assertEqual(report["redacted"], 0)
async def test_report_leaves_get_config_unchanged(self):
self.seed_secrets()
before = json.dumps(PromptManagerConfig.get_config(), sort_keys=True)
directories = GalleryConfig.MONITORING_DIRECTORIES
await self.report()
self.assertEqual(
json.dumps(PromptManagerConfig.get_config(), sort_keys=True), before
)
self.assertIs(GalleryConfig.MONITORING_DIRECTORIES, directories)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, KEY)
def test_outside_path_reads_outside_comfyui(self):
inside = os.path.join(str(self.comfy_dir), "user", "__promptmanager", "logs")
self.assertEqual(report_path(inside), "user/__promptmanager/logs")
self.assertEqual(
report_path(os.path.join(str(self.outside_dir), "vito")), "outside ComfyUI"
)
self.assertEqual(report_path(""), "")
self.assertEqual(report_path(None), "")
def test_a_filesystem_root_anchor_never_makes_a_path_public(self):
# --output-directory /renders makes "/" an anchor; a path outside
# ComfyUI must still read "outside ComfyUI", never "home/<user>/...".
root = os.path.abspath(os.sep)
anchors = [str(self.comfy_dir), root]
with mock.patch.object(GalleryConfig, "path_anchors", return_value=anchors):
self.assertEqual(
report_path(os.path.join(str(self.outside_dir), "x")), OUTSIDE
)
self.assertEqual(
report_path(os.path.join(str(self.comfy_dir), "output")), "output"
)
def test_scrub_removes_the_home_folder_without_its_root(self):
home = os.path.expanduser("~")
relative = os.path.splitdrive(home)[1].lstrip("\\/").replace("\\", "/")
cleaned = scrub({"a": f"{relative}/area/prompts.db", "redacted": 0}, None)
self.assertEqual(cleaned, {"a": "[removed]", "redacted": 1})
def test_the_home_folder_is_matched_as_a_whole_path_part(self):
with mock.patch.object(
diagnostics_report.os.path, "expanduser", return_value="/root"
):
cleaned = scrub(
{
"a": "output/rootfiles",
"b": "output/root/x",
"c": "see /root/x",
"redacted": 0,
},
None,
)
self.assertEqual(cleaned["a"], "output/rootfiles")
# The home folder's name is a whole path part here: removed.
self.assertEqual(cleaned["b"], "[removed]")
self.assertEqual(cleaned["c"], "[removed]")
self.assertEqual(cleaned["redacted"], 2)
async def test_report_reads_only_the_temporary_log_folder(self):
self.write("prompt_manager.log.1", b"x" * 1234)
report, _text = await self.report()
stats = report["logging"]["stats"]
self.assertEqual(stats["log_files_count"], 2)
self.assertEqual(
stats["total_log_size"],
1234 + os.path.getsize(self.path("prompt_manager.log")),
)
def test_an_ancestor_of_home_is_never_an_anchor_on_linux(self):
anchored = diagnostics_report.anchored_path
home = "/home/someone"
anchors = ["/home", "/opt/ComfyUI"]
self.assertEqual(
anchored("/home/someone/renders/a.png", anchors, home, posixpath), OUTSIDE
)
self.assertEqual(
anchored("/home/other/renders", anchors, home, posixpath), OUTSIDE
)
self.assertEqual(
anchored("/opt/ComfyUI/output", anchors, home, posixpath), "output"
)
def test_an_ancestor_of_home_is_never_an_anchor_on_windows(self):
anchored = diagnostics_report.anchored_path
home = ntpath.normcase("C:\\Users\\someone")
anchors = [ntpath.normcase(a) for a in ("C:\\Users", "C:\\", "D:\\ComfyUI")]
path = ntpath.normcase("C:\\Users\\someone\\renders\\a.png")
self.assertEqual(anchored(path, anchors, home, ntpath), OUTSIDE)
inside = ntpath.normcase("D:\\ComfyUI\\output\\x")
self.assertEqual(anchored(inside, anchors, home, ntpath), "output/x")
def test_an_anchor_that_is_home_names_no_user(self):
anchored = diagnostics_report.anchored_path
home = "/home/someone"
self.assertEqual(
anchored("/home/someone/ComfyUI/output", [home], home, posixpath),
"ComfyUI/output",
)
self.assertEqual(
anchored(
"/home/someone/ComfyUI/output",
[home, "/home/someone/ComfyUI"],
home,
posixpath,
),
"output",
)
def test_the_home_folder_name_is_removed_only_from_paths(self):
with mock.patch.object(
diagnostics_report.os.path, "expanduser", return_value="/home/info"
):
cleaned = scrub({"level": "INFO", "p": "logs/info/x", "redacted": 0}, None)
self.assertEqual(cleaned, {"level": "INFO", "p": "[removed]", "redacted": 1})
def test_scrub_removes_the_home_folder_name(self):
with mock.patch.object(
diagnostics_report.os.path, "expanduser", return_value="/home/someone"
):
cleaned = scrub(
{"a": "someone/renders", "b": "someones", "redacted": 0}, None
)
self.assertEqual(cleaned["a"], "[removed]")
self.assertEqual(cleaned["b"], "someones")
def test_scrub_removes_a_leaked_path_and_counts_it(self):
leaked = {
"a": "/home/someone/x",
"b": ["C:\\Users\\someone", "D:/data", "\\\\nas\\share", "ok"],
"c": {"d": f"token {KEY} here", "e": 3, "f": None, "g": "relative/path"},
"redacted": 0,
}
cleaned = scrub(leaked, KEY)
self.assertEqual(cleaned["redacted"], 5)
self.assertEqual(cleaned["a"], "[removed]")
self.assertEqual(cleaned["b"], ["[removed]"] * 3 + ["ok"])
self.assertEqual(
cleaned["c"], {"d": "[removed]", "e": 3, "f": None, "g": "relative/path"}
)
# A short key is not searched for (it would match ordinary words).
self.assertEqual(scrub({"x": "abc", "redacted": 0}, "abc")["redacted"], 0)
async def test_report_refused_while_exclusive_job_runs(self):
async with self.api._jobs.hold(
"restore", "Database restore", admission="exclusive"
) as restore:
report, _text = await self.report(status=409)
self.assertEqual(report["code"], "busy")
self.assertEqual(
report["running"], {"id": restore.id, "title": "Database restore"}
)
async def test_report_lists_config_only_keys(self):
self.seed_secrets()
report, _text = await self.report()
self.assertEqual({k: set(v) for k, v in report["config"].items()}, CONFIG)
self.assertNotIn("default_path", json.dumps(report["config"]))
self.assertNotIn("directories", json.dumps(report["config"]))
self.assertNotIn(
"civitai_api_key",
json.dumps(report["config"]).replace("has_civitai_api_key", ""),
)
async def test_report_logging_has_public_directory(self):
report, _text = await self.report()
stats = report["logging"]["stats"]
self.assertFalse(posixpath.isabs(stats["log_directory"]))
self.assertIn("level", report["logging"]["config"])
async def test_report_survives_an_unreadable_viewerjs_provenance(self):
with mock.patch.object(
diagnostics_report,
"VIEWERJS_PROVENANCE",
os.path.join(self.tmpdir, "nope.json"),
):
report, _text = await self.report()
self.assertEqual(report["viewerjs"], "unknown")
async def test_a_build_failure_is_coded_without_details(self):
with mock.patch.object(
self.api, "build_report", side_effect=RuntimeError(self.tmpdir)
):
with self.assertLogs("prompt_manager", "ERROR"):
report, text = await self.report(status=500)
self.assertNotIn(self.tmpdir, text)
self.assertFalse(report["success"])
class TestLegacyDiagnostics(DiagnosticsReportTestCase):
async def test_legacy_diagnostics_keeps_its_four_checks(self):
response = await self.client.request("GET", "/prompt_manager/diagnostics")
data = await response.json()
self.assertEqual(response.status, 200)
self.assertEqual(set(data), {"success", "diagnostics"})
self.assertEqual(
set(data["diagnostics"]),
{"database", "dependencies", "comfyui_output", "image_monitor"},
)
report, _text = await self.report()
self.assertEqual(report["checks"]["database"], data["diagnostics"]["database"])
self.assertEqual(
report["checks"]["dependencies"], data["diagnostics"]["dependencies"]
)
async def test_both_routes_build_their_checks_with_one_helper(self):
with mock.patch.object(
admin, "legacy_checks", wraps=admin.legacy_checks
) as spy:
await self.client.request("GET", "/prompt_manager/diagnostics")
await self.report()
self.assertEqual(spy.call_count, 2)
if __name__ == "__main__":
unittest.main()
+93 -18
View File
@@ -9,9 +9,11 @@ import asyncio
import json
import os
import pathlib
import shutil
import sys
import tempfile
import threading
import time
import types
import unittest
from contextlib import ExitStack
@@ -36,6 +38,8 @@ from PIL import Image # noqa: E402
from database.operations import PromptDatabase # noqa: E402
import py.api.images as images_module
from py.api import gallery_listing # noqa: E402
from py import lora_utils # noqa: E402
from py import lora_manager_layout as layout_module # noqa: E402
from py.api.thumbnail_settings import ThumbnailSettings # noqa: E402
from py.api import PromptManagerAPI # noqa: E402
from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402
@@ -111,13 +115,25 @@ class ImageAPITestCase(AioHTTPTestCase):
self.assertGreater(image_id, 0, "fixture: image link failed")
return image_id
def _make_symlink(self, target, link):
def _make_symlink(self, target, link, target_is_directory=False):
"""Create a symlink or skip the test when the platform refuses."""
try:
os.symlink(str(target), str(link))
os.symlink(str(target), str(link), target_is_directory=target_is_directory)
except (OSError, NotImplementedError, AttributeError) as exc:
self.skipTest(f"symlinks unavailable: {exc}")
def _spy_executor(self):
"""Record the __name__ of every callable sent to _run_in_executor."""
original = self.api._run_in_executor
calls = []
async def spy(func, *args, **kwargs):
calls.append(getattr(func, "__name__", repr(func)))
return await original(func, *args, **kwargs)
self.api._run_in_executor = spy
return calls
class TestServeOutputImage(ImageAPITestCase):
"""GET /prompt_manager/images/serve/{filepath}"""
@@ -388,18 +404,6 @@ class ThumbnailTestCase(ImageAPITestCase):
PROGRESS_URL = "/prompt_manager/images/generate-thumbnails/progress"
def _spy_executor(self):
"""Record the __name__ of every callable sent to _run_in_executor."""
original = self.api._run_in_executor
calls = []
async def spy(func, *args, **kwargs):
calls.append(getattr(func, "__name__", repr(func)))
return await original(func, *args, **kwargs)
self.api._run_in_executor = spy
return calls
def _thumb_path(self, name, suffix=".webp"):
return self.output_dir / "thumbnails" / f"{name}_thumb{suffix}"
@@ -1216,6 +1220,13 @@ def _raise(*_args, **_kwargs):
class ImageRouteCoverageCase(ImageAPITestCase):
"""Helpers for driving success and failure branches of each route."""
def setUp(self):
# #362: the LoraManager layout cache is per process; every test here
# starts and ends with it empty (runs before the subclasses' setUpAsync).
super().setUp()
lora_utils.clear_layout_cache()
self.addCleanup(lora_utils.clear_layout_cache)
def _break_db(self, db_method):
setattr(self.api.db, db_method, _raise)
@@ -1327,6 +1338,9 @@ class TestServeErrorBranches(ImageRouteCoverageCase):
)
self.assertEqual(resp.status, 200)
# The cache does not watch default_lora_roots (#362): without a
# clear the next request would reuse the folders above.
lora_utils.clear_layout_cache()
with (
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
@@ -1419,32 +1433,39 @@ class TestLoraManagerMediaServing(ImageRouteCoverageCase):
self.cache = root / "cache"
for folder in (self.sidecars, self.examples, self.cache):
folder.mkdir()
self.root = root
self.write_settings(str(self.sidecars))
self.enabled = patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True)
self.default_roots = patch("py.lora_utils.default_lora_roots", return_value=[])
for item in (
self.enabled,
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", str(self.lm)),
patch(
"py.lora_utils.custom_nodes_directories", return_value=[custom_nodes]
),
patch("py.lora_utils.default_lora_roots", return_value=[]),
self.default_roots,
patch("py.lora_utils.get_lora_image_cache_dir", return_value=self.cache),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
):
item.start()
self.addCleanup(item.stop)
def write_settings(self, sidecar_root):
(self.settings / "settings.json").write_text(
def write_settings(self, sidecar_root, examples=None):
path = self.settings / "settings.json"
path.write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": sidecar_root,
"example_images_path": str(self.examples),
"example_images_path": str(examples or self.examples),
"civitai_api_key": "never-served",
}
)
)
# Ten seconds old, so the layout cache (real clock here) stores the
# first load instead of treating the file as being written (#362).
stamp = time.time_ns() - 10 * 10**9
os.utime(path, ns=(stamp, stamp))
async def _status(self, image_id):
resp = await self.client.request(
@@ -1464,11 +1485,65 @@ class TestLoraManagerMediaServing(ImageRouteCoverageCase):
for image_id in ids:
self.assertEqual(await self._status(image_id), 200)
# A warm layout cache does not outlive the integration (#362).
self.assertIsNotNone(lora_utils._layout_cache._entry)
self.enabled.stop()
for image_id in ids:
self.assertEqual(await self._status(image_id), 403)
self.enabled.start()
def _example(self, folder=None, name="1.mp4"):
path = (folder or self.examples) / "ab12" / name
path.parent.mkdir(exist_ok=True)
path.write_bytes(b"video")
return path
async def test_unchanged_settings_reuse_the_cached_layout(self):
image_id = self._link_image(self._example())
self.assertEqual(await self._status(image_id), 200)
self.assertIsNotNone(lora_utils._layout_cache._entry)
self.assertEqual(await self._status(image_id), 200)
self.assertEqual(lora_utils.default_lora_roots.call_count, 1)
async def test_a_settings_change_is_seen_without_a_clear(self):
image_id = self._link_image(self._example())
self.assertEqual(await self._status(image_id), 200)
self.assertIsNotNone(lora_utils._layout_cache._entry)
moved = self.root / "examples-moved"
moved.mkdir()
self.write_settings(str(self.sidecars), examples=moved)
self.assertEqual(await self._status(image_id), 403)
async def test_a_cached_folder_swapped_for_a_symlink_to_home_is_refused(self):
# CUJ 7: the example-images folder becomes a symlink to a folder that
# holds the home folder; the same recorded file is refused.
outer = self.root / "outer"
(outer / "home").mkdir(parents=True)
self._example(outer)
image_id = self._link_image(self._example())
with patch.object(layout_module.Path, "home", return_value=outer / "home"):
self.assertEqual(await self._status(image_id), 200)
self.assertIsNotNone(lora_utils._layout_cache._entry)
# A skip after the rmtree is safe: setUp builds a fresh tree.
shutil.rmtree(self.examples)
self._make_symlink(outer, self.examples, target_is_directory=True)
self.assertEqual(await self._status(image_id), 403)
async def test_a_configured_path_outside_custom_nodes_adds_nothing(self):
from py.config import IntegrationConfig
self.assertEqual(await self._status(self._link_image(self._example())), 200)
self.assertIsNotNone(lora_utils._layout_cache._entry)
with patch.object(IntegrationConfig, "LORA_MANAGER_PATH", str(self.sidecars)):
self.assertEqual(self.api._lora_image_dirs(), [])
# The warm entry for the real install was never handed out for it.
self.assertIsNone(lora_utils.cached_lora_manager_layout(str(self.sidecars))[0])
async def test_serve_image_builds_the_allow_list_in_the_executor(self):
calls = self._spy_executor()
self.assertEqual(await self._status(self._link_image(self._example())), 200)
self.assertIn("_lora_image_dirs", calls)
async def test_a_filesystem_root_as_sidecar_folder_adds_nothing(self):
self.write_settings(os.path.abspath(os.sep))
dirs = self.api._lora_image_dirs()
+419
View File
@@ -0,0 +1,419 @@
"""Reviewed rotated-log cleanup as a maintenance plan (#329, py/api/log_cleanup.py).
``POST /maintenance/preview {operation: "rotated_logs"}`` lists exactly the
rotated copies Rule R allows and writes nothing; ``POST /maintenance/apply``
deletes exactly those as a ``maintenance`` job (admission ``read``), or
nothing, failing ``plan-changed``, when any of them changed after the review.
Every test runs on a real ``PromptManagerLogger`` whose log folder is a
temporary folder, never the real one.
"""
import asyncio
import json
import logging
import ntpath
import os
import posixpath
import re
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 py.api.maintenance_plans import PlanStore # noqa: E402
from utils import logging_config # noqa: E402
from utils.logging_config import PromptManagerLogger, get_logger_manager # noqa: E402
try:
from tests.open_handles import assert_closed # noqa: E402
from tests.test_api_maintenance import MaintenanceAPITestCase # noqa: E402
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed # noqa: E402
from test_api_maintenance import MaintenanceAPITestCase # noqa: E402
WAIT = 10
PREVIEW = "/prompt_manager/maintenance/preview"
APPLY = "/prompt_manager/maintenance/apply"
CHANGED = "The log files changed since the review; nothing was deleted. Review again."
MEDIA = "Deletes only the listed log files. No image, video or library record changes."
_WINDOWS_ABSOLUTE = re.compile(r"^[A-Za-z]:[\\/]")
def _strings(value):
if isinstance(value, str):
yield value
elif isinstance(value, dict):
for item in value.values():
yield from _strings(item)
elif isinstance(value, (list, tuple)):
for item in value:
yield from _strings(item)
class LogCleanupTestCase(MaintenanceAPITestCase):
async def get_application(self):
app = await super().get_application()
package = os.path.join(self.tmpdir, "pkg")
os.makedirs(os.path.join(package, "utils"))
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.addCleanup(self._restore_global_logging)
self.manager.update_config({"console_logging": False})
self._only_this_thread()
self.log_dir = str(self.manager.log_dir)
self.app_logger = logging.getLogger("prompt_manager.log_cleanup_test")
self.app_logger.info("active log line")
self.api._get_logger_manager = lambda: self.manager
self.outside = os.path.join(self.tmpdir, "link-targets")
os.makedirs(self.outside)
return app
def _only_this_thread(self):
"""Only this test's records reach the temporary log (not other threads')."""
this_thread = threading.get_ident()
for handler in self.manager.logger.handlers:
handler.filters.clear()
handler.addFilter(lambda record: record.thread == this_thread)
def _restore_global_logging(self):
for handler in logging.getLogger("prompt_manager").handlers[:]:
handler.close()
assert_closed(self, self.log_dir)
get_logger_manager()._setup_loggers()
def path(self, name):
return os.path.join(self.log_dir, name)
def write(self, name, data=b"rotated\n"):
with open(self.path(name), "wb") as handle:
handle.write(data)
def symlink(self, name):
target = os.path.join(self.outside, "secret.txt")
with open(target, "wb") as handle:
handle.write(b"secret " * 100)
try:
os.symlink(target, self.path(name))
except (OSError, NotImplementedError): # pragma: no cover - no privilege
return False
return True
def names(self):
return sorted(os.listdir(self.log_dir))
def snapshot(self):
"""Every entry but the active log, with its bytes (None for a folder)."""
state = {}
for name in self.names():
if name == "prompt_manager.log":
continue
path = self.path(name)
if os.path.islink(path):
state[name] = ("link", os.readlink(path))
elif os.path.isdir(path):
state[name] = None
else:
with open(path, "rb") as handle:
state[name] = handle.read()
return state
def seed(self):
"""CUJ 5's folder: three rotated copies, a .bak, a folder and a link."""
self.write("prompt_manager.log.1", b"a" * 1000)
self.write("prompt_manager.log.2", b"b" * 2000)
self.write("prompt_manager.log.3", b"c" * 3000)
self.write("prompt_manager.log.bak", b"mine")
os.makedirs(self.path("prompt_manager.log.9"))
return self.symlink("prompt_manager.log.4")
async def preview(self, status=200, **body):
return await self.post(
PREVIEW, {"operation": "rotated_logs", "options": {}, **body}, status
)
async def apply(self, token, status=202):
return await self.post(APPLY, {"token": token}, status)
async def run_apply(self, plan):
body = await self.apply(plan["token"])
job = self.api._jobs.get(body["job"]["id"])
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
return body["job"], job
class TestPreview(LogCleanupTestCase):
async def test_preview_lists_only_rotated_copies(self):
linked = self.seed()
plan = await self.preview()
self.assertEqual(
set(plan),
{
"success",
"token",
"operation",
"expires_at",
"scope",
"protected",
"backup",
"media",
"details",
},
)
self.assertEqual(plan["operation"], "rotated_logs")
self.assertEqual(
plan["scope"], {"count": 3, "label": "3 rotated log copies · 5.9 KB"}
)
self.assertEqual(
[row["label"] for row in plan["details"]],
["prompt_manager.log.1", "prompt_manager.log.2", "prompt_manager.log.3"],
)
self.assertRegex(
plan["details"][2]["detail"], r"^2\.9 KB · \d{4}-\d\d-\d\d \d\d:\d\d$"
)
reasons = ["Kept: the active log"]
if linked:
reasons.append("Kept: 1 link")
reasons.append("Kept: 2 other files")
self.assertEqual(
plan["protected"], {"count": len(reasons) + 1, "reasons": reasons}
)
self.assertEqual(
plan["backup"],
{"policy": "none", "destination_label": None, "expected_bytes": None},
)
self.assertEqual(plan["media"], {"deletes_files": True, "label": MEDIA})
self.assertTrue(plan["expires_at"].endswith("+00:00"))
self.assertNotIn("mtime_ns", json.dumps(plan))
async def test_preview_writes_nothing(self):
self.seed()
before = self.snapshot()
await self.preview()
self.assertEqual(self.snapshot(), before)
async def test_preview_has_no_paths(self):
self.seed()
plan = await self.preview()
for text in _strings(plan):
self.assertFalse(posixpath.isabs(text), text)
self.assertFalse(ntpath.isabs(text), text)
self.assertIsNone(_WINDOWS_ABSOLUTE.match(text), text)
async def test_preview_rejects_non_empty_options(self):
for options in ({"all": True}, [], "x"):
with self.subTest(options=options):
reply = await self.preview(status=400, options=options)
self.assertEqual(reply["code"], "invalid-request")
async def test_preview_too_large_is_413(self):
self.seed()
self.api._maintenance_plans = PlanStore(max_bytes=10)
reply = await self.preview(status=413)
self.assertEqual(
reply,
{
"success": False,
"code": "too-large",
"error": "Too many log files to review at once.",
},
)
async def test_empty_folder_plan_has_scope_zero(self):
plan = await self.preview()
self.assertEqual(plan["scope"], {"count": 0, "label": "No rotated log copies"})
self.assertEqual(plan["details"], [])
self.assertEqual(
plan["protected"], {"count": 1, "reasons": ["Kept: the active log"]}
)
async def test_preview_refused_while_an_exclusive_job_runs(self):
async with self.api._jobs.hold(
"restore", "Database restore", admission="exclusive"
):
reply = await self.preview(status=409)
self.assertEqual(reply["code"], "busy")
class TestApply(LogCleanupTestCase):
async def test_apply_deletes_exactly_the_reviewed_files(self):
linked = self.seed()
plan = await self.preview()
summary, job = await self.run_apply(plan)
self.assertEqual(summary["kind"], "maintenance")
self.assertEqual(summary["title"], "Rotated log copies")
self.assertEqual(summary["admission"], "read")
self.assertFalse(summary["cancellable"])
self.assertEqual(summary["scope"], {"operation": "rotated_logs"})
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(
job.counters,
{"changed": 3, "protected": plan["protected"]["count"], "errors": 0},
)
expected = [
"prompt_manager.log",
"prompt_manager.log.9",
"prompt_manager.log.bak",
]
if linked:
expected.append("prompt_manager.log.4")
self.assertEqual(self.names(), sorted(expected))
self.assertEqual(job.status_line, "Deleted 3 rotated log copies")
report = json.dumps(job.report())
self.assertNotIn(self.tmpdir, report)
async def test_apply_rows_name_each_file(self):
self.seed()
plan = await self.preview()
_summary, job = await self.run_apply(plan)
self.assertEqual(
job.outcomes["changed"],
[
{
"path": "prompt_manager.log.1",
"root": "",
"root_index": None,
"detail": "Deleted · 1000 bytes",
},
{
"path": "prompt_manager.log.2",
"root": "",
"root_index": None,
"detail": "Deleted · 2.0 KB",
},
{
"path": "prompt_manager.log.3",
"root": "",
"root_index": None,
"detail": "Deleted · 2.9 KB",
},
],
)
log = "\n".join(job.log)
for n in (1, 2, 3):
self.assertIn(f"prompt_manager.log.{n}", log)
async def test_apply_after_rollover_fails_plan_changed_and_deletes_nothing(self):
self.seed()
plan = await self.preview()
handler = next(
h
for h in self.manager.logger.handlers
if isinstance(h, logging.FileHandler)
)
handler.acquire()
try:
handler.doRollover() # every listed name now holds another file
finally:
handler.release()
before = self.snapshot()
_summary, job = await self.run_apply(plan)
self.assertEqual(job.state, "failed")
self.assertEqual(job.scope["failure"], "plan-changed")
self.assertEqual(job.status_line, CHANGED)
self.assertEqual(job.counters["changed"], 0)
self.assertEqual(self.snapshot(), before)
async def test_apply_reports_a_file_in_use(self):
self.seed()
plan = await self.preview()
real_unlink = os.unlink
held = self.path("prompt_manager.log.2")
def unlink(path, *args, **kwargs):
if os.fspath(path) == held:
error = PermissionError(
13,
"The process cannot access the file because it is being "
"used by another process",
held,
)
error.winerror = 32
raise error
return real_unlink(path, *args, **kwargs)
with mock.patch.object(logging_config.os, "unlink", unlink):
_summary, job = await self.run_apply(plan)
self.assertEqual(job.counters["changed"], 2)
self.assertEqual(job.counters["errors"], 1)
[row] = job.outcomes["errors"]
self.assertEqual(row["path"], "prompt_manager.log.2")
self.assertIn("being used by another process", row["detail"])
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
self.assertIn("prompt_manager.log.2", self.names())
self.assertEqual(
job.status_line, "Deleted 2 rotated log copies, 1 could not be deleted"
)
async def test_token_is_single_use(self):
self.seed()
plan = await self.preview()
summary, _job = await self.run_apply(plan)
again = await self.apply(plan["token"], 409)
self.assertEqual(again["code"], "plan-used")
self.assertEqual(
again["job"], {"id": summary["id"], "title": "Rotated log copies"}
)
async def test_apply_runs_beside_a_scan(self):
self.seed()
plan = await self.preview()
async with self.api._jobs.hold("scan", "Metadata scan", admission="records"):
_summary, job = await self.run_apply(plan)
self.assertEqual(job.state, "finished", job.log)
self.assertNotIn("prompt_manager.log.1", self.names())
async def test_apply_refused_while_restore_runs(self):
self.seed()
plan = await self.preview()
async with self.api._jobs.hold(
"restore", "Database restore", admission="exclusive"
) as restore:
busy = await self.apply(plan["token"], 409)
self.assertEqual(busy["code"], "busy")
self.assertEqual(
busy["running"], {"id": restore.id, "title": "Database restore"}
)
self.assertIn("prompt_manager.log.1", self.names())
# The refused apply kept the token.
_summary, job = await self.run_apply(plan)
self.assertEqual(job.state, "finished", job.log)
async def test_apply_refused_beside_another_maintenance_job(self):
self.seed()
plan = await self.preview()
async with self.api._jobs.hold(
"maintenance", "Optimize database", admission="records"
):
busy = await self.apply(plan["token"], 409)
self.assertEqual(busy["code"], "busy")
async def test_active_log_survives_and_keeps_growing(self):
self.seed()
plan = await self.preview()
await self.run_apply(plan)
self.app_logger.info("written after the cleanup")
with open(self.path("prompt_manager.log"), "rb") as handle:
content = handle.read()
self.assertIn(b"active log line", content)
self.assertIn(b"written after the cleanup", content)
async def test_an_unexpected_failure_fails_the_job_without_details(self):
self.seed()
plan = await self.preview()
with mock.patch.object(
self.manager, "delete_rotated", side_effect=RuntimeError(self.tmpdir)
):
with self.assertLogs("prompt_manager", "ERROR"):
_summary, job = await self.run_apply(plan)
self.assertEqual(job.state, "failed")
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
self.assertEqual(job.status_line, "The log cleanup failed.")
if __name__ == "__main__":
unittest.main()
+183 -11
View File
@@ -5,9 +5,10 @@ import os
import shutil
import sys
import tempfile
import threading
import unittest
from pathlib import Path
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
@@ -20,7 +21,9 @@ 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 import logging_routes # noqa: E402
from py.api.logging_routes import LOG_CONFIG_KEYS # noqa: E402
from utils import logging_config # noqa: E402
class FakeLoggerManager:
@@ -38,11 +41,16 @@ class FakeLoggerManager:
}
self.update_calls = []
self.logs = []
self.threads = {}
def update_config(self, new_config):
self.threads["update_config"] = threading.current_thread()
self.update_calls.append(dict(new_config))
self.config.update(new_config)
def log_file_entry(self, name):
return logging_config.log_file_entry(self.log_dir, name)
def get_config(self):
return dict(self.config)
@@ -54,23 +62,36 @@ class FakeLoggerManager:
def get_log_files(self):
files = []
for path in sorted(self.log_dir.glob("*.log*")):
for row in sorted(
logging_config.scan_log_folder(self.log_dir), key=lambda r: r["name"]
):
if row["_dir"]:
continue
files.append(
{
"filename": path.name,
"path": str(path),
"size": path.stat().st_size,
"filename": row["name"],
"path": str(self.log_dir / row["name"]),
"size": row["size"],
"modified": "2026-10-07T00:00:00",
"is_main": path.name == "prompt_manager.log",
"is_main": row["kind"] == "active",
"kind": row["kind"],
}
)
return files
def truncate_logs(self):
"""Rule R: empty the active log, delete rotated copies, keep the rest."""
self.threads["truncate_logs"] = threading.current_thread()
truncated = []
for path in self.log_dir.glob("prompt_manager.log*"):
path.write_text("")
truncated.append(path.name)
for row in logging_config.scan_log_folder(self.log_dir):
path = self.log_dir / row["name"]
if row["kind"] == "active":
path.write_text("")
elif row["kind"] == "rotated":
path.unlink()
else:
continue
truncated.append(row["name"])
return {"truncated": truncated, "errors": []}
def get_log_stats(self):
@@ -378,12 +399,129 @@ class TestLogFilesAndDownload(LoggingAPITestCase):
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")
# Rewritten (#329): a directory is now a 404; inject a read error.
(self.log_dir / "prompt_manager.log").write_bytes(b"x")
secret = os.path.join(self.tmpdir, "logs", "prompt_manager.log")
def fail(path):
raise PermissionError(13, "Permission denied", secret)
with patch.object(logging_routes, "_read_log_bytes", fail):
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log"
)
body = await resp.text()
self.assertEqual(resp.status, 500)
self.assertNotIn(self.tmpdir, body)
async def test_directory_download_is_404(self):
(self.log_dir / "dir.log").mkdir()
resp = await self.client.request("GET", "/prompt_manager/logs/download/dir.log")
self.assertEqual(resp.status, 404)
async def test_download_refuses_a_symlink(self):
secret = Path(self.tmpdir) / "secret.txt"
secret.write_text("the key")
try:
os.symlink(secret, self.log_dir / "prompt_manager.log.3")
except (OSError, NotImplementedError): # pragma: no cover
self.skipTest("this account cannot create symlinks")
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log.3"
)
body = await resp.text()
self.assertEqual(resp.status, 404)
self.assertNotIn("the key", body)
async def test_download_refuses_a_hard_link(self):
secret = Path(self.tmpdir) / "secret.txt"
secret.write_text("the key")
try:
os.link(secret, self.log_dir / "prompt_manager.log.3")
except (OSError, NotImplementedError): # pragma: no cover
self.skipTest("this filesystem cannot create hard links")
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log.3"
)
self.assertEqual(resp.status, 404)
self.assertNotIn("the key", await resp.text())
async def test_download_refuses_a_file_swapped_after_it_was_listed(self):
(self.log_dir / "prompt_manager.log.2").write_text("listed")
other = Path(self.tmpdir) / "other.txt"
other.write_text("swapped in")
real_entry = self.manager.log_file_entry
def entry_then_swap(name):
path = real_entry(name)
os.replace(other, path) # another file now has the listed name
return path
self.manager.log_file_entry = entry_then_swap
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log.2"
)
self.assertEqual(resp.status, 404)
self.assertNotIn("swapped in", await resp.text())
@unittest.skipUnless(hasattr(os, "mkfifo"), "FIFOs are POSIX only")
async def test_download_refuses_a_fifo_without_blocking(self):
os.mkfifo(self.log_dir / "pipe.log")
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/pipe.log"
)
self.assertEqual(resp.status, 404)
async def test_download_reads_off_the_loop(self):
(self.log_dir / "prompt_manager.log").write_bytes(b"line\n")
seen = []
real = logging_routes._read_log_bytes
def read(path):
seen.append(threading.current_thread())
return real(path)
with patch.object(logging_routes, "_read_log_bytes", read):
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log"
)
self.assertEqual(resp.status, 200)
self.assertEqual(len(seen), 1)
self.assertIsNot(seen[0], threading.current_thread())
async def test_download_keeps_its_response_shape(self):
(self.log_dir / "prompt_manager.log.2").write_bytes(b"a\nb\n")
resp = await self.client.request(
"GET", "/prompt_manager/logs/download/prompt_manager.log.2"
)
self.assertEqual(resp.status, 200)
self.assertEqual(resp.content_type, "text/plain")
self.assertEqual(
resp.headers["Content-Disposition"],
'attachment; filename="prompt_manager.log.2"',
)
self.assertEqual(resp.headers["Content-Length"], "4")
self.assertEqual(await resp.read(), b"a\nb\n")
async def test_log_files_rows_carry_their_kind(self):
for name in (
"prompt_manager.log",
"prompt_manager.log.1",
"prompt_manager.log.bak",
):
(self.log_dir / name).write_text("x")
resp = await self.client.request("GET", "/prompt_manager/logs/files")
data = await resp.json()
kinds = {row["filename"]: row["kind"] for row in data["files"]}
self.assertEqual(
kinds,
{
"prompt_manager.log": "active",
"prompt_manager.log.1": "rotated",
"prompt_manager.log.bak": "other",
},
)
class TestTruncateConfigAndStats(LoggingAPITestCase):
@@ -395,6 +533,40 @@ class TestTruncateConfigAndStats(LoggingAPITestCase):
self.assertEqual(data["results"]["truncated"], ["prompt_manager.log"])
self.assertEqual((self.log_dir / "prompt_manager.log").read_text(), "")
async def test_truncate_keeps_other_files(self):
for name in (
"prompt_manager.log",
"prompt_manager.log.1",
"prompt_manager.log.bak",
):
(self.log_dir / name).write_text("data")
resp = await self.client.request("POST", "/prompt_manager/logs/truncate")
data = await resp.json()
self.assertEqual(resp.status, 200)
self.assertEqual(set(data), {"success", "message", "results"})
self.assertEqual(
sorted(data["results"]["truncated"]),
["prompt_manager.log", "prompt_manager.log.1"],
)
self.assertEqual(
sorted(p.name for p in self.log_dir.iterdir()),
["prompt_manager.log", "prompt_manager.log.bak"],
)
async def test_truncate_runs_off_the_loop(self):
resp = await self.client.request("POST", "/prompt_manager/logs/truncate")
self.assertEqual(resp.status, 200)
self.assertIsNot(
self.manager.threads["truncate_logs"], threading.current_thread()
)
async def test_update_config_runs_off_the_loop(self):
resp = await self._post_config({"level": "DEBUG"})
self.assertEqual(resp.status, 200)
self.assertIsNot(
self.manager.threads["update_config"], threading.current_thread()
)
async def test_truncate_failure_is_500(self):
def boom():
raise RuntimeError("locked")
+12
View File
@@ -876,6 +876,18 @@ class TestLoraScanRoute(LoraAPITestCase):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
async def test_a_finished_scan_clears_the_layout_cache(self):
# #362: the legacy scan replaces through the same method.
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "plain")
(loras / "plain.png").write_bytes(_png_bytes())
with patch("py.lora_utils.clear_layout_cache") as clear:
resp, events, _calls = await self._scan(loras)
self.assertEqual(events[-1]["type"], "complete", events)
clear.assert_called_once_with()
async def test_scan_imports_loras_with_images_and_links_existing(self):
self._enable()
loras = self.tmp_root / "loras"
+29 -12
View File
@@ -19,8 +19,15 @@ 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
# py.autotag reads its saved defaults from py.config, which imports ComfyUI's
# ``server`` at load: stand in for it when this module runs alone.
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
from py import autotag # noqa: E402
from py.autotag import MODELS, AutoTagService # noqa: E402
from py.config import AutotagConfig # noqa: E402
@contextlib.contextmanager
@@ -254,12 +261,19 @@ class TestWd14Thresholding(EngineTestCase):
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)
def test_service_thresholds_are_the_saved_defaults(self):
# #328: one configuration; the service has no thresholds of its own.
service = self._loaded([0.9, 0.3, 0.8, 0.99, 0.99, 0.5])
with (
patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.25),
patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.75),
):
self.assertEqual(
(service.wd14_general_threshold, service.wd14_character_threshold),
(0.25, 0.75),
)
tags = service.generate_tags(str(self.image_path))
self.assertEqual(tags, ["1girl", "hatsune_miku", "long_hair", "smile"])
class TestGgufNormalisation(EngineTestCase):
@@ -369,10 +383,13 @@ class TestTagSuggestions(EngineTestCase):
def test_thresholds_default_to_the_service_values(self):
service = self._wd14()
service.wd14_general_threshold = 0.95
service.wd14_character_threshold = 0.95
suggestions = service.generate_tag_suggestions(str(self.image_path), limit=10)
with (
patch.object(AutotagConfig, "GENERAL_THRESHOLD", 0.95),
patch.object(AutotagConfig, "CHARACTER_THRESHOLD", 0.95),
):
suggestions = service.generate_tag_suggestions(
str(self.image_path), limit=10
)
self.assertEqual(suggestions, [("sandstone", 0.98), ("desert", 0.96)])
+335
View File
@@ -0,0 +1,335 @@
"""Browser-test resets seed a scratch copy and publish it in one commit (#384)."""
import hashlib
import os
import sqlite3
import subprocess
import sys
import tempfile
import threading
import time
import types
import unittest
from pathlib import Path
from unittest import mock
try:
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.operations import PromptDatabase
from tests.browser import fast_reset
from tests.browser.fast_reset import ScratchSeeder, table_digests
def digests(db: PromptDatabase) -> dict:
return table_digests(db.model.get_connection())
def add_prompt(db: PromptDatabase, text: str) -> int:
with db.model.get_connection() as conn:
return conn.execute(
"INSERT INTO prompts (text, hash, created_at, updated_at) "
"VALUES (?, ?, '2026-01-01 00:00:00', '2026-01-01 00:00:00')",
(text, hashlib.sha256(text.encode()).hexdigest()),
).lastrowid
def add_image(db: PromptDatabase, prompt_id: int, name: str) -> int:
with db.model.get_connection() as conn:
return conn.execute(
"INSERT INTO generated_images (prompt_id, image_path, filename, "
"generation_time) VALUES (?, ?, ?, '2026-01-01 00:00:00')",
(prompt_id, f"/out/{name}", name),
).lastrowid
def reseed(api) -> str:
"""A seeder in the old style: clears its tables, then many small commits."""
db = api.db
with db.model.get_connection() as conn:
for table in ("generated_images", "prompt_tags", "tags", "prompts"):
conn.execute(f"DELETE FROM {table}")
for index in range(30):
prompt_id = add_prompt(db, f"seeded {index}")
add_image(db, prompt_id, f"seeded-{index}.png")
with db.model.get_connection() as conn:
tag_id = conn.execute(
"INSERT INTO tags (name) VALUES (?)", (f"tag{index}",)
).lastrowid
conn.execute(
"INSERT INTO prompt_tags (prompt_id, tag_id) VALUES (?, ?)",
(prompt_id, tag_id),
)
return "manifest"
class TestScratchSeeder(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.root = Path(self.tmp.name)
self.addCleanup(assert_closed, self, self.tmp.name)
self.live = self.database("live.db")
self.api = types.SimpleNamespace(db=self.live, marker="live api")
self.seeder = ScratchSeeder(self.api, self.root / "scratch", memory_root=None)
self.addCleanup(self.seeder.close)
# Rows from "earlier tests": AUTOINCREMENT positions have moved on.
for index in range(5):
add_image(self.live, add_prompt(self.live, f"old {index}"), f"o{index}")
with self.live.model.get_connection() as conn:
conn.execute("DELETE FROM prompts WHERE text = 'old 4'")
def database(self, name: str) -> PromptDatabase:
db = PromptDatabase(str(self.root / name))
self.addCleanup(db.close_all)
return db
def oracle(self, seed):
"""Run ``seed`` the old way, on a copy of the live database."""
path = self.root / "oracle.db"
target = sqlite3.connect(path)
try:
self.live.model.get_connection().backup(target)
finally:
target.close()
db = self.database("oracle.db")
seed(types.SimpleNamespace(db=db))
return digests(db)
def test_publishes_exactly_what_the_seed_writes_on_live(self):
expected = self.oracle(reseed)
self.assertEqual(self.seeder.run(reseed), "manifest")
self.assertEqual(digests(self.live), expected)
with self.live.model.get_connection() as conn:
first = conn.execute("SELECT MIN(id) FROM prompts").fetchone()[0]
self.assertEqual(first, 6, "AUTOINCREMENT continues as on live")
def test_runs_twice_from_the_state_the_first_run_left(self):
self.seeder.run(reseed)
expected = self.oracle(reseed)
self.seeder.run(reseed)
self.assertEqual(digests(self.live), expected)
def test_tables_the_seed_did_not_change_keep_their_rows(self):
before = digests(self.live)
def rename(api):
with api.db.model.get_connection() as conn:
conn.execute("UPDATE prompts SET text = 'renamed' WHERE id = 1")
self.seeder.run(rename)
after = digests(self.live)
changed = {table for table in after if after[table] != before[table]}
self.assertEqual(changed, {"prompts"})
with self.live.model.get_connection() as conn:
images = conn.execute("SELECT COUNT(*) FROM generated_images").fetchone()
# Replacing prompts' rows never cascades into generated_images.
self.assertEqual(images[0], 4)
def test_the_seed_sees_the_scratch_database_and_the_live_api(self):
seen = {}
def seed(api):
seen["db"] = api.db
seen["marker"] = api.marker
api.note = "set by the seed"
self.seeder.run(seed)
self.assertIsNot(seen["db"], self.live)
self.assertNotEqual(seen["db"].model.db_path, self.live.model.db_path)
self.assertEqual(seen["marker"], "live api")
self.assertEqual(self.api.note, "set by the seed")
self.assertIs(self.api.db, self.live)
def test_cache_drops_wait_until_the_rows_are_published(self):
seen = []
def invalidate_gallery_cache(retire=False):
with self.live.model.get_connection() as conn:
count = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0]
seen.append((count, retire))
self.api.invalidate_gallery_cache = invalidate_gallery_cache
def seed(api):
reseed(api)
api.invalidate_gallery_cache(retire=True)
self.assertEqual(seen, [], "not dropped while live has the old rows")
self.seeder.run(seed)
self.assertEqual(seen, [(30, True)])
def test_a_failing_seed_drops_no_cache(self):
seen = []
self.api.invalidate_gallery_cache = lambda: seen.append(True)
def broken(api):
api.invalidate_gallery_cache()
raise ValueError("seed failed")
with self.assertRaises(ValueError):
self.seeder.run(broken)
self.assertEqual(seen, [])
def test_a_failing_seed_leaves_live_untouched(self):
before = digests(self.live)
def broken(api):
reseed(api)
raise ValueError("seed failed")
with self.assertRaisesRegex(ValueError, "seed failed"):
self.seeder.run(broken)
self.assertEqual(digests(self.live), before)
self.assertIs(self.api.db, self.live)
def test_a_concurrent_live_commit_fails_the_reset(self):
other = self.database("live.db")
def racing(api):
reseed(api)
add_prompt(other, "written to live meanwhile")
with self.assertRaisesRegex(RuntimeError, "changed while"):
self.seeder.run(racing)
with self.live.model.get_connection() as conn:
texts = [row[0] for row in conn.execute("SELECT text FROM prompts")]
self.assertIn("written to live meanwhile", texts)
self.assertNotIn("seeded 0", texts)
def test_a_write_open_during_the_seed_and_committed_later_fails_it(self):
other = self.database("live.db")
began, finished = threading.Event(), []
def write():
with other.model.get_connection() as conn:
conn.execute("BEGIN IMMEDIATE")
conn.execute(
"INSERT INTO prompts (text, hash) VALUES ('concurrent', 'c')"
)
began.set()
time.sleep(0.5) # the publish waits on this lock meanwhile
other.model.close()
finished.append(True)
writer = threading.Thread(target=write)
def racing(api):
reseed(api)
writer.start()
self.assertTrue(began.wait(10))
with self.assertRaisesRegex(RuntimeError, "changed while"):
self.seeder.run(racing)
writer.join()
self.assertEqual(finished, [True])
with self.live.model.get_connection() as conn:
texts = [row[0] for row in conn.execute("SELECT text FROM prompts")]
self.assertIn("concurrent", texts)
self.assertNotIn("seeded 0", texts)
self.assert_live_connection_restored()
def test_a_failed_publish_leaves_live_and_its_connection_as_they_were(self):
before = digests(self.live)
columns = fast_reset._columns
def broken_columns(conn, table, *schema):
# The publish asks for the live table's columns without a schema.
return columns(conn, table, *schema) if schema else ["no_such_column"]
with mock.patch.object(fast_reset, "_columns", broken_columns):
with self.assertRaises(sqlite3.OperationalError):
self.seeder.run(reseed)
self.assertEqual(digests(self.live), before)
self.assert_live_connection_restored()
def assert_live_connection_restored(self):
conn = self.live.model.get_connection()
self.assertFalse(conn.in_transaction)
self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone()[0], 1)
names = [row[1] for row in conn.execute("PRAGMA database_list")]
self.assertEqual(names, ["main"])
def test_the_live_connection_is_left_as_it_was(self):
self.seeder.run(reseed)
conn = self.live.model.get_connection()
self.assertFalse(conn.in_transaction)
self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone()[0], 1)
names = [row[1] for row in conn.execute("PRAGMA database_list")]
self.assertEqual(names, ["main"])
self.assertEqual(conn.execute("PRAGMA foreign_key_check").fetchall(), [])
def test_on_live_mode_runs_the_seed_on_the_live_api(self):
seeder = ScratchSeeder(
self.api, self.root / "unused", on_live=True, memory_root=self.root
)
self.addCleanup(seeder.close)
expected = self.oracle(reseed)
seen = []
self.assertEqual(
seeder.run(lambda api: seen.append(api) or reseed(api)), "manifest"
)
self.assertIs(seen[0], self.api)
self.assertEqual(digests(self.live), expected)
self.assertFalse((self.root / "unused").exists())
def test_the_scratch_lives_under_the_memory_root_when_there_is_one(self):
memory = self.root / "memory"
memory.mkdir()
seeder = ScratchSeeder(self.api, self.root / "fallback", memory_root=memory)
self.addCleanup(seeder.close)
expected = self.oracle(reseed)
seeder.run(reseed)
self.assertEqual(digests(self.live), expected)
self.assertEqual(len(list(memory.glob("pm-reset-*/seed.db"))), 1)
self.assertFalse((self.root / "fallback").exists())
seeder.close()
self.assertEqual(list(memory.iterdir()), [])
def test_a_full_memory_root_puts_the_scratch_on_disk(self):
memory = self.root / "memory"
memory.mkdir()
seeder = ScratchSeeder(self.api, self.root / "fallback", memory_root=memory)
self.addCleanup(seeder.close)
full = mock.patch.object(
fast_reset.shutil,
"disk_usage",
return_value=types.SimpleNamespace(free=fast_reset.MEMORY_FREE_BYTES - 1),
)
with full:
seeder.run(reseed)
self.assertEqual(list(memory.iterdir()), [])
self.assertTrue((self.root / "fallback" / "seed.db").exists())
@unittest.skipUnless(os.name == "posix", "process probe is POSIX only")
def test_scratch_folders_of_dead_servers_are_removed(self):
memory = self.root / "memory"
memory.mkdir()
dead = subprocess.Popen([sys.executable, "-c", "pass"])
dead.wait()
orphan = memory / f"pm-reset-{dead.pid}-abc"
alive = memory / f"pm-reset-{os.getpid()}-def"
other = memory / "pm-reset-notapid"
for folder in (orphan, alive, other):
folder.mkdir()
(folder / "seed.db").write_bytes(b"x")
seeder = ScratchSeeder(self.api, self.root / "fallback", memory_root=memory)
self.addCleanup(seeder.close)
seeder.run(reseed)
self.assertFalse(orphan.exists())
self.assertTrue(alive.exists())
self.assertTrue(other.exists())
own = [p for p in memory.glob(f"pm-reset-{os.getpid()}-*") if p != alive]
self.assertEqual(len(own), 1)
def test_close_releases_the_scratch_files(self):
self.seeder.run(reseed)
self.seeder.close()
assert_closed(self, self.root / "scratch")
self.seeder.close() # idempotent
if __name__ == "__main__":
unittest.main()
+67
View File
@@ -0,0 +1,67 @@
"""Browser-test media folders are emptied by a rename, deleted later (#384)."""
import tempfile
import threading
import unittest
from pathlib import Path
from tests.browser.folder_trash import TRASH_NAME, close_trash, replace_folder
class TestFolderTrash(unittest.TestCase):
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.root = Path(self.tmp.name)
self.addCleanup(close_trash, self.root)
def trash_threads(self):
return [t for t in threading.enumerate() if t.name == "pm-reset-trash"]
def test_replaces_the_folder_with_an_empty_one(self):
folder = self.root / "output"
(folder / "nested").mkdir(parents=True)
(folder / "nested" / "image.png").write_bytes(b"png")
replace_folder(folder)
self.assertTrue(folder.is_dir())
self.assertEqual(list(folder.iterdir()), [])
close_trash(self.root)
self.assertFalse((self.root / TRASH_NAME).exists())
self.assertEqual(self.trash_threads(), [])
def test_a_missing_folder_is_created_without_a_trash(self):
replace_folder(self.root / "input")
self.assertTrue((self.root / "input").is_dir())
self.assertFalse((self.root / TRASH_NAME).exists())
self.assertEqual(self.trash_threads(), [])
def test_repeated_replacements_share_one_thread_and_all_are_removed(self):
folder = self.root / "output"
for number in range(5):
folder.mkdir(exist_ok=True)
(folder / f"{number}.png").write_bytes(b"png")
replace_folder(folder)
self.assertEqual(len(self.trash_threads()), 1)
close_trash(self.root)
self.assertEqual(list(folder.iterdir()), [])
self.assertFalse((self.root / TRASH_NAME).exists())
close_trash(self.root) # idempotent
class TestFolderTrashAfterClose(unittest.TestCase):
def test_a_replacement_after_close_deletes_in_place_without_a_thread(self):
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
close_trash(root)
folder = root / "output"
folder.mkdir()
(folder / "late.png").write_bytes(b"png")
replace_folder(folder)
self.assertEqual(list(folder.iterdir()), [])
self.assertFalse((root / TRASH_NAME).exists())
names = [t.name for t in threading.enumerate()]
self.assertNotIn("pm-reset-trash", names)
if __name__ == "__main__":
unittest.main()
+3 -3
View File
@@ -221,12 +221,12 @@ class TestSlowResetOnTheServer(unittest.TestCase):
self.assertRegex(
text, r"kernel: asyncio_\d+ \(0x[0-9a-f]{16}\) state \w, waiting in "
)
# The plain reset's seed() waits on its first DELETE.
self.assertRegex(text, r"line \d+ in seed\n")
# The plain reset seeds a scratch copy (#384), then waits to publish it.
self.assertRegex(text, r"line \d+ in _publish\n")
self.assertIn("POST /__test__/reset answered 200 after", text)
self.assertRegex(
text,
r"executor call create_app.<locals>.seed: queued \d+\.\d+s, ran \d+\.\d+s",
r"executor call ScratchSeeder.run: queued \d+\.\d+s, ran \d+\.\d+s",
)
+125 -2
View File
@@ -413,9 +413,21 @@ class TestAutotagConfig(unittest.TestCase):
self.config_class = getattr(config_module, "AutotagConfig", None)
self.assertIsNotNone(self.config_class, "AutotagConfig is not implemented")
self._orig = PromptManagerConfig.get_config()
self._defaults = {
attr: getattr(self.config_class, attr, None)
for attr in (
"MODEL_TYPE",
"EXISTING",
"GENERAL_THRESHOLD",
"CHARACTER_THRESHOLD",
"TEMPLATE",
"SUGGESTION_LIMIT",
)
}
def tearDown(self):
self.config_class.SUGGESTION_LIMIT = 10
for attr, value in self._defaults.items():
setattr(self.config_class, attr, value)
PromptManagerConfig.update_config(self._orig)
def _load(self, value):
@@ -446,8 +458,23 @@ class TestAutotagConfig(unittest.TestCase):
self.assertEqual(self.config_class.suggestion_limit(), 10)
self.assertEqual(
PromptManagerConfig.get_config()["autotag"],
{"defaults": {"suggestion_limit": 10}},
{
"defaults": {
"model_type": "wd14-swinv2",
"existing": "skip",
"general_threshold": 0.35,
"character_threshold": 0.85,
"suggestion_limit": 10,
}
},
)
# The built-in template is not written; a saved one is.
self.config_class.TEMPLATE = "Saved."
self.assertEqual(
PromptManagerConfig.get_config()["autotag"]["defaults"]["template"],
"Saved.",
)
self.config_class.TEMPLATE = config_module.DEFAULT_TAG_TEMPLATE
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
PromptManagerConfig.update_config(
@@ -466,6 +493,102 @@ class TestAutotagConfig(unittest.TestCase):
)
self.assertEqual(self.config_class.suggestion_limit(), value)
def test_autotag_defaults_fall_back_per_key_on_load(self):
"""#328: each saved tagging default is checked alone on load; a bad
hand-edited value falls back to its default with one warning and
leaves the other keys as saved."""
good = {
"model_type": "gguf",
"existing": "merge",
"general_threshold": 0.4,
"character_threshold": 0.7,
"template": "Tags only.",
"suggestion_limit": 15,
}
fallback = {
"model_type": "wd14-swinv2",
"existing": "skip",
"general_threshold": 0.35,
"character_threshold": 0.85,
"template": config_module.DEFAULT_TAG_TEMPLATE,
"suggestion_limit": 10,
}
bad_values = {
"model_type": ["wd14-convnext", 3, None],
"existing": ["replace", "", True],
# #369: a saved threshold off the 0.05 step no page can show.
"general_threshold": [1.5, -0.1, float("nan"), "0.4", True, 0.37],
"character_threshold": [2, float("inf"), None, 0.33],
# 4000 characters are at most 16,000 UTF-8 bytes, so the 16 KiB
# byte cap can only be met past the character cap.
"template": ["", " ", "x" * 4001, "\U0001f600" * 4097, "\ud800", 7],
}
for key, values in bad_values.items():
for value in values:
with self.subTest(key=key, value=value):
saved = dict(good, **{key: value})
with tempfile.TemporaryDirectory() as tmpdir:
path = os.path.join(tmpdir, "config.json")
with open(path, "w") as f:
json.dump({"autotag": {"defaults": saved}}, f)
with self.assertLogs(
"prompt_manager.config", level="WARNING"
) as logs:
PromptManagerConfig.load_from_file(path)
self.assertEqual(len(logs.records), 1, logs.output)
self.assertIn(f"autotag.defaults.{key}", logs.output[0])
loaded = {
k: getattr(self.config_class, attr)
for k, (attr, _v, _d) in self.config_class._KEYS.items()
}
self.assertEqual(loaded, dict(good, **{key: fallback[key]}))
def test_tagging_validators_moved_with_their_messages(self):
validate_threshold = config_module.validate_threshold
validate_template = config_module.validate_template
self.assertEqual(validate_threshold("general_threshold", 1), 1.0)
for value, message in (
(float("nan"), "general_threshold must be a finite number"),
("0.4", "general_threshold must be a finite number"),
(1.01, "general_threshold must be between 0 and 1"),
):
with self.assertRaises(ValueError) as caught:
validate_threshold("general_threshold", value)
self.assertEqual(str(caught.exception), message)
self.assertEqual(validate_template("Tags."), "Tags.")
for value, message in (
(" ", "instructions must be non-empty text"),
("x" * 4001, "instructions must be at most 4000 characters"),
("\ud800", "instructions must be valid Unicode text"),
):
with self.assertRaises(ValueError) as caught:
validate_template(value)
self.assertEqual(str(caught.exception), message)
self.assertEqual(
config_module.MODEL_TYPES, {"gguf", "hf", "wd14-swinv2", "wd14-vit"}
)
def test_saved_threshold_is_a_multiple_of_the_step(self):
"""#369: every page shows the thresholds in steps of 0.05, so a saved
one must be on that grid; it is stored exactly on it."""
from py.tagging_defaults import validate_saved_threshold as validate
for value in (0, 0.05, 0.15, 0.35, 0.85, 1):
with self.subTest(value=value):
self.assertEqual(validate("x", value), value)
self.assertEqual(validate("x", 0.1 + 0.2), 0.3)
for value, message in (
(0.37, "x must be a multiple of 0.05"),
(0.351, "x must be a multiple of 0.05"),
(0.35000000001, "x must be a multiple of 0.05"),
(1.5, "x must be between 0 and 1"),
("0.4", "x must be a finite number"),
):
with self.subTest(value=value):
with self.assertRaises(ValueError) as caught:
validate("x", value)
self.assertEqual(str(caught.exception), message)
def test_suggestion_limit_setting_validated_on_load(self):
for value in (0, 101, "10", float("nan"), 2.5, True, None, [10]):
with self.subTest(value=value):
+44
View File
@@ -1398,6 +1398,50 @@ class TestModelHousekeeping(DatabaseTestCase):
self.assertIsNot(self.db.model.get_connection(), failing)
class TestAutotagBatchCandidates(DatabaseTestCase):
"""#328: the batch's candidate read never loads workflow documents."""
def test_autotag_candidates_skip_workflow_data(self):
from utils.hashing import generate_prompt_hash
first = self.db.save_prompt(
text="a", tags=["one", "two"], prompt_hash=generate_prompt_hash("a")
)
second = self.db.save_prompt(text="b", prompt_hash=generate_prompt_hash("b"))
old = self.db.link_image_to_prompt(
first, "/out/a.png", {"workflow": {"big": "x" * 1000}}
)
new = self.db.link_image_to_prompt(second, "/out/b.png")
statements = []
conn = self.db.model.get_connection() # this thread's own connection
conn.set_trace_callback(statements.append)
try:
rows = self.db.autotag_batch_candidates()
finally:
conn.set_trace_callback(None)
self.assertEqual(
rows[0],
{
"id": new,
"image_path": "/out/b.png",
"prompt_id": second,
"prompt_tags": [],
},
)
self.assertEqual(
{**rows[1], "prompt_tags": sorted(rows[1]["prompt_tags"])},
{
"id": old,
"image_path": "/out/a.png",
"prompt_id": first,
"prompt_tags": ["one", "two"],
},
)
self.assertEqual(len(rows), 2)
self.assertTrue(statements)
self.assertFalse(any("workflow_data" in sql for sql in statements))
if __name__ == "__main__":
unittest.main()
+10 -3
View File
@@ -28,6 +28,8 @@ except ImportError:
Status = namedtuple("ExecutionStatus", "status_str completed messages")
JOB = "a0c26da3-eb26-405c-b143-c662a84c3112"
# A cross-thread wait or a duration check fails only past this (#370).
CEILING_SECONDS = 30.0
class HostQueue:
@@ -209,6 +211,9 @@ class HistoryHooksTests(unittest.TestCase):
self.queue.get = unrelated_after
self.dequeue(0)
self.queue.task_done(0, {}, status())
# The writer drains first: this test checks wrappers and the session,
# not the 2 s bound, which a slow runner can spend on the write (#370).
self.drain(first)
self.assertTrue(first.dispose())
self.assertTrue(first.dispose())
self.assertIs(self.queue.get, unrelated_after)
@@ -339,15 +344,17 @@ class HistoryHooksTests(unittest.TestCase):
with patch.object(self.db, "write_execution_start", blocked):
registration = self.register()
self.dequeue(0)
self.assertTrue(entered.wait(1))
# Ceilings, not tight wall clocks (#370): the return values show
# the bound; a ceiling still fails a call that waits unbounded.
self.assertTrue(entered.wait(CEILING_SECONDS))
began = time.monotonic()
self.queue.task_done(0, {}, status())
self.dequeue(1)
self.assertLess(time.monotonic() - began, 0.2)
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
began = time.monotonic()
self.assertFalse(registration.dispose())
self.assertGreaterEqual(time.monotonic() - began, 1.8)
self.assertLess(time.monotonic() - began, 2.3)
self.assertLess(time.monotonic() - began, CEILING_SECONDS)
release.set()
self.assertTrue(registration.dispose(2))
+2
View File
@@ -642,6 +642,8 @@ class HistoryWriterTests(unittest.TestCase):
== "succeeded"
)
self.assertTrue(writer.submit(self.start_command(2)))
# Drain first: the 2 s bound is not under test here (#370).
drain_writer(writer)
self.assertTrue(writer.close(2))
self.assertEqual(len(set(threads)), 1)
self.assertNotEqual(threads[0], threading.get_ident())
+2
View File
@@ -245,6 +245,8 @@ class OutputHooksTests(fixtures.unittest.TestCase):
registration = self.register()
self.dequeue(0)
late = registration._task_done
# Drain first: the 2 s bound is not under test here (#370).
self.drain(registration)
self.assertTrue(registration.dispose())
thread = threading.Thread(
target=lambda: late((0, history(), fixtures.status()), {})
+3
View File
@@ -17,6 +17,7 @@ from utils.execution_contexts import MAX_ATTEMPT_BYTES
from utils.execution_output_capture import seal_outputs
from utils.execution_outputs import snapshot_output_roots
from utils.execution_history_writer import HistoryWriter
from tests.history_waits import drain_writer
try:
from tests import test_execution_history_writer as fixtures
@@ -161,6 +162,8 @@ class OutputWriterTests(fixtures.unittest.TestCase):
start = replace(self.start_command(), run_id=str(fixtures.uuid4()))
self.assertTrue(writer.submit(start))
self.finalization(writer, start)
# Drain first: the 2 s bound is not under test here (#370).
drain_writer(writer)
self.assertTrue(writer.close(2))
self.assertEqual(
self.db.get_execution_run(start.run_id)["outcome"], "succeeded"
+2
View File
@@ -12,6 +12,7 @@ import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import stop_tracker # noqa: E402
from utils.hashing import (
generate_prompt_hash,
is_duplicate_prompt,
@@ -49,6 +50,7 @@ class TestSingleHashImplementation(unittest.TestCase):
db = mock.Mock()
db.get_prompt_by_hash.return_value = {"id": 42}
tracker = PromptTracker(db)
self.addCleanup(stop_tracker, self, tracker)
tracker.set_current_prompt(" Hello World ")
db.get_prompt_by_hash.assert_called_once_with(
generate_prompt_hash(" Hello World ")
+2
View File
@@ -8,6 +8,7 @@ import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import close_handler # noqa: E402
from database.operations import PromptDatabase
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
@@ -52,6 +53,7 @@ class LinkingTestCase(unittest.TestCase):
)
def tearDown(self):
close_handler(self, self.handler) # and the tracker's cleanup thread
self.db.close_all() # Windows cannot unlink an open database file
for path in (
self.image,
+259 -12
View File
@@ -136,13 +136,21 @@ class TestMonitorSingletonLifecycle(unittest.TestCase):
class FakeObserver:
"""Records what the monitor does with the watchdog observer."""
"""Records what the monitor does with the watchdog observer.
def __init__(self):
``calls`` lists ``stop`` and ``join`` in order (shared with a test's own
recorders when given); ``stay_alive`` keeps it alive after ``join``;
``stop_error`` makes ``stop`` raise.
"""
def __init__(self, calls=None, stay_alive=False, stop_error=None):
self.scheduled = []
self.started = False
self.stopped = False
self.joined = False
self.calls = calls if calls is not None else []
self.stay_alive = stay_alive
self.stop_error = stop_error
def schedule(self, handler, path, recursive=False):
self.scheduled.append((handler, path, recursive))
@@ -151,12 +159,18 @@ class FakeObserver:
self.started = True
def stop(self):
self.calls.append(("stop",))
if self.stop_error is not None:
raise self.stop_error
self.stopped = True
def join(self, timeout=None):
self.calls.append(("join", timeout))
self.joined = True
def is_alive(self):
if self.stay_alive:
return self.started
return self.started and not self.stopped
@@ -171,6 +185,8 @@ class TestStartStop(unittest.TestCase):
self.missing = os.path.join(self.tmp.name, "does-not-exist")
def tearDown(self):
# Before the folder is removed (addCleanup runs after tearDown).
self.assertTrue(self.monitor.stop_monitoring())
im_mod._monitor_instance = None
def test_only_nonexistent_directories_leaves_monitoring_off(self):
@@ -225,9 +241,8 @@ class TestStartStop(unittest.TestCase):
worker = handler._worker
self.assertTrue(worker.is_alive())
self.monitor.stop_monitoring()
self.assertTrue(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
@@ -303,6 +318,89 @@ class TestStartStop(unittest.TestCase):
self.assertEqual(self.monitor.monitored_directories, [self.tmp.name])
self.monitor.stop_monitoring()
def _start_with(self, fake):
with patch.object(im_mod, "Observer", return_value=fake):
self.monitor.start_monitoring([self.tmp.name])
handler = self.monitor.handler
real_close = handler.close
def close(timeout=im_mod.WORKER_JOIN_SECONDS):
fake.calls.append(("close",))
return real_close(timeout)
handler.close = close
return handler
def test_stop_monitoring_joins_observer_and_worker(self):
fake = FakeObserver()
handler = self._start_with(fake)
self.assertTrue(self.monitor.stop_monitoring(timeout=3.0))
self.assertEqual([c[0] for c in fake.calls], ["close", "stop", "join"])
join_timeout = fake.calls[2][1]
self.assertIsNotNone(join_timeout)
self.assertLessEqual(join_timeout, 3.0)
self.assertTrue(handler._closed)
def test_stop_monitoring_reports_an_observer_still_running(self):
fake = FakeObserver(stay_alive=True)
self._start_with(fake)
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
self.assertIn("watchdog observer still running", "\n".join(logs.output))
self.assertFalse(self.monitor.is_monitoring)
fake.stay_alive = False # it ends; tearDown's stop waits for it
def test_a_second_stop_monitoring_waits_for_the_observer(self):
fake = FakeObserver(stay_alive=True)
self._start_with(fake)
with self.assertLogs("prompt_manager.image_monitor", "WARNING"):
self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
self.assertFalse(self.monitor.stop_monitoring(timeout=0.05))
fake.stay_alive = False
self.assertTrue(self.monitor.stop_monitoring())
self.assertFalse(self.monitor._unfinished)
def test_stop_monitoring_reports_an_observer_stop_that_raised(self):
fake = FakeObserver(stop_error=OSError("descriptor gone"))
self._start_with(fake)
fake.stopped = True # its thread has ended; stop() still raises
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.assertFalse(self.monitor.stop_monitoring())
text = "\n".join(logs.output)
self.assertIn("Error while stopping image monitoring", text)
self.assertNotIn("watchdog observer still running", text)
fake.stop_error = None # tearDown's stop calls stop() again
def test_the_still_running_warning_names_the_time_the_join_got(self):
fake = FakeObserver(stay_alive=True)
handler = self._start_with(fake)
now = [100.0]
recorded_close = handler.close
def close_that_uses_the_deadline(timeout=im_mod.WORKER_JOIN_SECONDS):
ok = recorded_close(timeout)
now[0] += 5.0 # a busy worker took the whole deadline
return ok
handler.close = close_that_uses_the_deadline
with (
patch.object(
im_mod, "time", types.SimpleNamespace(monotonic=lambda: now[0])
),
self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs,
):
self.assertFalse(self.monitor.stop_monitoring(timeout=5.0))
self.assertEqual(fake.calls[-1], ("join", 0.0))
self.assertIn(
"watchdog observer still running after 0.0 s", "\n".join(logs.output)
)
fake.stay_alive = False # it ends; tearDown's stop waits for it
def test_stop_monitoring_without_an_observer_returns_true(self):
self.assertTrue(self.monitor.stop_monitoring())
class FakeEvent:
def __init__(self, path, is_directory=False):
@@ -335,6 +433,10 @@ class TestHandlerScheduling(unittest.TestCase):
self.processed.append((path, prompt_snapshot))
)
def tearDown(self):
# Before the folder is removed (addCleanup runs after tearDown).
self.assertTrue(self.handler.close())
def _image(self, name, content=b"x" * 16):
path = os.path.join(self.tmp.name, name)
with open(path, "wb") as f:
@@ -399,23 +501,127 @@ class TestHandlerScheduling(unittest.TestCase):
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):
def test_shutdown_joins_the_worker_within_the_bound(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()
self.assertTrue(self.handler.shutdown())
worker.join(timeout=2)
self.assertFalse(worker.is_alive())
self.assertEqual(self.handler.pending_count(), 0)
def test_shutdown_reports_a_worker_still_busy(self):
entered, release = threading.Event(), threading.Event()
def busy(path, prompt_snapshot=None):
entered.set()
release.wait(30)
self.handler.process_new_image = busy
self.handler.schedule(self._image("busy.png"))
worker = self.handler._worker
self.addCleanup(worker.join, 30)
self.addCleanup(release.set)
self.assertTrue(entered.wait(30), "the worker never took the image")
with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs:
self.assertFalse(self.handler.shutdown(timeout=0.05))
self.assertIn("image worker still running", "\n".join(logs.output))
release.set()
worker.join(30)
self.assertFalse(worker.is_alive())
def test_a_second_stop_waits_for_a_worker_still_finishing(self):
entered, release = threading.Event(), threading.Event()
def busy(path, prompt_snapshot=None):
entered.set()
release.wait(30)
self.handler.process_new_image = busy
self.handler.schedule(self._image("twice.png"))
worker = self.handler._worker
self.addCleanup(worker.join, 30)
self.addCleanup(release.set)
self.assertTrue(entered.wait(30), "the worker never took the image")
with self.assertLogs("prompt_manager.image_monitor", "WARNING"):
self.assertFalse(self.handler.close(timeout=0.05))
self.assertFalse(self.handler.close(timeout=0.05))
release.set()
self.assertTrue(self.handler.close(timeout=30))
self.assertFalse(worker.is_alive())
def test_shutdown_on_the_worker_itself_only_signals(self):
entered, result = threading.Event(), {}
def stop_from_inside(path, prompt_snapshot=None):
result["value"] = self.handler.shutdown()
entered.set()
self.handler.process_new_image = stop_from_inside
self.handler.schedule(self._image("self.png"))
worker = self.handler._worker
self.assertTrue(entered.wait(30), "the worker never took the image")
worker.join(30)
self.assertFalse(result["value"])
self.assertFalse(worker.is_alive())
def _live_workers(self):
return {
t for t in threading.enumerate() if t.name == "PromptManagerImageWorker"
}
def _seal(self, handler, name):
handler._await_history(os.path.join(self.tmp.name, name))
handler._on_outputs_sealed(
"run-1",
[{"folder_type": "output", "subfolder": "", "filename": name}],
)
def test_close_refuses_new_work(self):
path = self._image("late.png")
with patch.object(im_mod, "host_output_directory", return_value=self.tmp.name):
self.assertTrue(self.handler.close())
before = self._live_workers()
self.assertFalse(self.handler.schedule(path))
# The image in hand can still ask for a history re-check and its
# run can be sealed afterwards, on the history writer's thread.
self.handler._await_history(path)
self.assertEqual(self.handler.awaiting_history(), [])
self._seal(self.handler, "late.png")
self.handler._schedule_recheck(path)
self.assertIsNone(self.handler._worker)
self.assertEqual(list(self.handler._history_due), [])
self.assertEqual(self._live_workers() - before, set())
def test_a_sealed_run_starts_a_worker_on_an_open_handler(self):
self._image("open.png")
with patch.object(im_mod, "host_output_directory", return_value=self.tmp.name):
self._seal(self.handler, "open.png")
self.assertIsNotNone(self.handler._worker)
self.assertTrue(self.handler.close())
def test_close_is_idempotent(self):
self.handler.schedule(self._image("c1.png"))
self.assertTrue(self.handler.close())
self.assertTrue(self.handler.close())
def test_on_created_after_close_takes_no_lock(self):
self.assertTrue(self.handler.close())
self.handler.on_created(FakeEvent(self._image("after.png")))
self.tracker.keyed_snapshot.assert_not_called()
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.assertTrue(fresh.shutdown())
self.assertTrue(fresh.shutdown())
self.assertTrue(fresh.close())
self.handler.schedule(self._image("s3.png"))
self.assertTrue(wait_for(lambda: self.handler.pending_count() == 0))
self.handler.shutdown()
@@ -434,7 +640,7 @@ class TestHandlerScheduling(unittest.TestCase):
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()
self.assertTrue(self.handler.close())
def test_processing_errors_do_not_kill_the_worker(self):
def explode(path, prompt_snapshot=None):
@@ -449,6 +655,47 @@ class TestHandlerScheduling(unittest.TestCase):
self.assertEqual(self.processed, [first, second])
class TestCloseReleasesHandles(unittest.TestCase):
"""``close`` joins the worker, so the test folder can be removed (#302)."""
def test_close_releases_handles(self):
from PIL import Image
from database.operations import PromptDatabase
from tests.open_handles import assert_closed
from utils.prompt_tracker import PromptTracker
with tempfile.TemporaryDirectory() as folder:
db = PromptDatabase(os.path.join(folder, "prompts.db"))
tracker = PromptTracker(db)
handler = im_mod.ImageGenerationHandler(db, tracker)
handler.processing_delay = 0
handler.settle_interval = 0.01
done = threading.Event()
real_process = handler.process_new_image
def process(path, prompt_snapshot=None):
try:
real_process(path, prompt_snapshot=prompt_snapshot)
finally:
done.set()
handler.process_new_image = process
path = os.path.join(folder, "real.png")
Image.new("RGB", (8, 8), "red").save(path)
try:
self.assertTrue(handler.schedule(path))
worker = handler._worker
self.assertTrue(done.wait(30), "the worker never processed the image")
finally:
closed, shut = handler.close(), tracker.shutdown()
db.model.close_all()
self.assertTrue(closed)
self.assertTrue(shut)
self.assertFalse(worker.is_alive())
assert_closed(self, folder)
class TestOutputListener(unittest.TestCase):
"""Every settled output is reported to the Images listing (#348, PR 2)."""
@@ -456,7 +703,7 @@ class TestOutputListener(unittest.TestCase):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.handler = im_mod.ImageGenerationHandler(MagicMock(), MagicMock())
self.addCleanup(self.handler.shutdown)
self.addCleanup(lambda: self.assertTrue(self.handler.close()))
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.handler.process_new_image = lambda path, prompt_snapshot=None: None
@@ -549,7 +796,7 @@ class TestOutputListener(unittest.TestCase):
self.assertIsNone(monitor.output_listener)
self.assertIsNone(monitor.handler.output_listener)
finally:
monitor.stop_monitoring()
self.assertTrue(monitor.stop_monitoring())
class TestModuleIsolation(unittest.TestCase):
+1 -1
View File
@@ -310,7 +310,7 @@ class TestAutoTag(ImageTagJobTestCase):
self.assertTrue(
any(
line.endswith(
"is not downloaded. Download it from the legacy Auto-tag page."
"is not downloaded. Download it in System › Auto-tagging."
)
for line in job.log
),
+439 -12
View File
@@ -4,6 +4,8 @@ The package is loaded the way ComfyUI loads it (``__init__`` with relative
imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed.
"""
import asyncio
import atexit
import contextlib
import importlib.util
import io
@@ -16,6 +18,7 @@ import threading
import types
import unittest
import unittest.mock as mock
from pathlib import Path
from aiohttp import web
@@ -40,13 +43,20 @@ class FakePromptQueue:
return self.results.pop(0) if self.results else None
def fake_server_module():
"""A ``server`` module whose PromptServer.instance accepts routes and hooks."""
def fake_server_module(app=None, queue_task_done=False):
"""A ``server`` module whose PromptServer.instance accepts routes and hooks.
``app`` replaces the default app (one without aiohttp hooks);
``queue_task_done`` gives the queue a ``task_done``, so history registers.
"""
queue = FakePromptQueue()
if queue_task_done:
queue.task_done = lambda *a, **k: None
instance = types.SimpleNamespace(
routes=web.RouteTableDef(),
prompt_queue=FakePromptQueue(),
prompt_queue=queue,
handlers=[],
app=types.SimpleNamespace(middlewares=[]),
app=app if app is not None else types.SimpleNamespace(middlewares=[]),
)
instance.add_on_prompt_handler = instance.handlers.append
module = types.ModuleType("server")
@@ -92,7 +102,7 @@ def _import_package(name, db_path, output_dir, server=None):
try:
spec.loader.exec_module(package)
except BaseException:
cleanup_package(package)
cleanup_package(package, require_stopped=False)
raise
finally:
for key, module in saved.items():
@@ -103,20 +113,34 @@ def _import_package(name, db_path, output_dir, server=None):
return package, server.PromptServer.instance
def cleanup_package(package):
"""Stop fixture-owned workers before releasing its other resources."""
def cleanup_package(package, require_stopped=True):
"""Stop fixture-owned workers before releasing its other resources.
The package's own teardown runs first (the reload case of #301); with
``require_stopped`` every part it reports must have stopped. The direct
stops after it also cover an import that failed before the teardown
existed, and a test that made a teardown step fail on purpose.
"""
teardown = getattr(package, "_stop_background_work", None)
if callable(teardown):
result = teardown()
if require_stopped:
assert all(result.values()), f"package teardown left work running: {result}"
monitor = getattr(package, "_global_image_monitor", None)
if monitor is not None:
assert monitor.stop_monitoring(), "package image monitor still running"
tracker = getattr(package, "_global_prompt_tracker", None)
if tracker is not None:
assert tracker.shutdown(), "package prompt cleanup thread still running"
registration = getattr(package, "_history_registration", None)
if registration is not None:
assert registration.dispose(3), "package history worker leaked"
monitor = getattr(package, "_global_image_monitor", None)
if monitor is not None:
monitor.stop_monitoring()
close_every_database()
sys.modules.pop(package.__name__, None)
@contextlib.contextmanager
def load_package(name, db_path, output_dir, server=None):
def load_package(name, db_path, output_dir, server=None, require_stopped=True):
"""Dispose package workers before closing its isolated logging handlers."""
root_logger = logging.getLogger("prompt_manager")
loggers = [root_logger] + [
@@ -144,7 +168,7 @@ def load_package(name, db_path, output_dir, server=None):
try:
yield package, instance
finally:
cleanup_package(package)
cleanup_package(package, require_stopped=require_stopped)
finally:
current_loggers = [root_logger] + [
logger
@@ -556,3 +580,406 @@ class TestPackageLoggerIsolation(unittest.TestCase):
self.assert_logging_restored(
captured["handlers"], captured["package"]._history_registration
)
PACKAGE_THREAD_NAMES = ("PromptManagerPromptCleanup", "PromptManagerImageWorker")
def package_threads():
"""Live threads this design owns: the tracker, the image worker, and the
watchdog observer and emitters (which have no name of their own)."""
from watchdog.utils import BaseThread
return {
thread
for thread in threading.enumerate()
if thread.name in PACKAGE_THREAD_NAMES or isinstance(thread, BaseThread)
}
def hooked_app():
return types.SimpleNamespace(
middlewares=[], on_cleanup=[], on_startup=[], on_shutdown=[]
)
WAIT_SECONDS = 30 # ceiling for every Event wait (TESTING.md test rules)
def start_busy_worker(package, folder, name="busy.png"):
"""The image worker alive in its processing delay."""
handler = package._global_image_monitor.handler
if handler is None:
raise AssertionError("the package's monitor did not start")
handler.processing_delay = 3600
# Outside the watched folder, so the observer does not schedule it first.
path = os.path.join(folder, name)
with open(path, "wb") as image:
image.write(b"x" * 16)
if not handler.schedule(path) or not handler._worker.is_alive():
raise AssertionError("the image worker did not start")
return handler._worker
class PackageTeardownCase(unittest.TestCase):
"""Each test loads its own package (#301): the teardown runs once."""
def setUp(self):
self.tmp = tempfile.TemporaryDirectory()
self.addCleanup(self.tmp.cleanup)
self.addCleanup(lambda: self.assertEqual(open_handles_under(self.tmp.name), []))
self.output_dir = os.path.join(self.tmp.name, "output")
os.makedirs(self.output_dir)
self.db_path = os.path.join(self.tmp.name, "prompts.db")
def load(self, server=None, require_stopped=True, suffix=""):
name = f"pm_teardown_{self._testMethodName}{suffix}"
stack = contextlib.ExitStack()
self.addCleanup(stack.close)
stack.enter_context(contextlib.redirect_stdout(io.StringIO()))
package, _ = stack.enter_context(
load_package(
name,
self.db_path,
self.output_dir,
server,
require_stopped=require_stopped,
)
)
return package
def busy_worker(self, package, name="busy.png"):
return start_busy_worker(package, self.tmp.name, name)
def hold_a_preview(self, package, *, advance, kill_releases):
"""A preview worker held inside a fake encode that tracks a fake
ffmpeg, and a package clock that the monitor step moves by
``advance`` seconds. Cleanups release the encode first."""
# Not the watched folder. Real path for both the file and the root:
# macOS /var vs /private/var, Windows short names.
previews = os.path.realpath(os.path.join(self.tmp.name, "previews"))
os.makedirs(previews)
from PIL import Image
image = os.path.join(previews, "held.png")
Image.new("RGB", (8, 8), "red").save(image)
held = types.SimpleNamespace(
entered=threading.Event(),
release=threading.Event(),
killed=threading.Event(),
stop_seen=[],
timeouts=[],
now=[100.0],
)
class FakeFfmpeg:
def poll(self): # _track keeps only processes still running
return None
def kill(self):
held.killed.set()
if kill_releases:
held.release.set()
def held_encode(src, dst, settings, is_video, *, stop, guard, track):
track(FakeFfmpeg())
held.entered.set()
held.release.wait(WAIT_SECONDS)
held.stop_seen.append(stop())
return {"action": "skipped"}
api = package.api
monitor = package._global_image_monitor
real_stop = monitor.stop_monitoring
real_close = api.close_auto_thumbnails
def stop_then_advance(timeout):
ok = real_stop(timeout)
held.now[0] += advance
return ok
def recorded_close(final=False, timeout=None):
held.timeouts.append(timeout)
return real_close(final=final, timeout=timeout)
for target, name, value in (
(api, "_get_all_output_dirs", lambda: [Path(previews)]),
(api, "_generate_one", held_encode),
(api, "close_auto_thumbnails", recorded_close),
(monitor, "stop_monitoring", stop_then_advance),
(package, "_clock", lambda: held.now[0]),
):
patcher = mock.patch.object(target, name, value)
patcher.start()
self.addCleanup(patcher.stop)
self.assertEqual(api.queue_capture_thumbnail(image), "queued")
self.addCleanup(held.release.set)
self.assertTrue(held.entered.wait(WAIT_SECONDS), "the encode never started")
held.thumbnailer = api._auto_thumbnailer # the teardown drops the API's
held.worker = held.thumbnailer._worker
# Cleanups run newest first: release the encode, then join the worker
# (even when an assertion failed first), then end the patches.
self.addCleanup(held.worker.join, WAIT_SECONDS)
self.addCleanup(held.release.set)
self.assertTrue(held.worker.is_alive())
return held
class TestPackageTeardown(PackageTeardownCase):
def test_an_image_in_hand_queues_no_preview_during_teardown(self):
package = self.load()
monitor = package._global_image_monitor
real_stop = monitor.stop_monitoring
image = os.path.join(self.output_dir, "in_hand.png")
from PIL import Image
Image.new("RGB", (8, 8), "red").save(image)
seen = {}
def stop_with_an_image_in_hand(timeout):
# The worker read ``thumbnail_queue`` before the detach and now
# hands its finished image over (at exit, Python 3.12+ refuses a
# new thread here and logs a traceback).
seen["queued"] = package.api.queue_capture_thumbnail(image)
seen["worker"] = package.api._auto_thumbnailer
return real_stop(timeout)
with (
mock.patch.object(
package.api,
"_get_all_output_dirs",
return_value=[Path(os.path.realpath(self.output_dir))],
),
mock.patch.object(
monitor, "stop_monitoring", side_effect=stop_with_an_image_in_hand
),
):
result = package._stop_background_work()
self.assertEqual(seen["queued"], "skipped")
self.assertIsNone(seen["worker"])
self.assertTrue(all(result.values()), result)
def test_teardown_kills_ffmpeg_and_stops_the_preview_worker(self):
package = self.load()
held = self.hold_a_preview(package, advance=4.0, kill_releases=True)
result = package._stop_background_work()
self.assertTrue(held.killed.is_set())
self.assertEqual(held.timeouts, [6.0]) # min(30, 10 - 4)
self.assertTrue(result["thumbnails"], result)
self.assertFalse(held.worker.is_alive())
self.assertEqual(held.stop_seen, [True])
def test_on_cleanup_runs_the_teardown(self):
app = hooked_app()
package = self.load(fake_server_module(app=app))
self.assertIs(app.on_cleanup[0], package._stop_background_work_on_cleanup)
self.assertIsNone(package._background_work_result)
async def run_cleanup():
for callback in list(app.on_cleanup):
await callback(app)
asyncio.run(run_cleanup())
self.assertEqual(
package._background_work_result,
{"monitor": True, "thumbnails": True, "tracker": True},
)
def test_teardown_runs_before_history_dispose_at_exit(self):
with mock.patch("atexit.register", wraps=atexit.register) as register:
package = self.load(fake_server_module(queue_task_done=True))
registered = [call.args[0] for call in register.call_args_list]
registration = package._history_registration
self.assertIsNotNone(registration, "history did not register")
self.assertIn(registration.dispose, registered)
self.assertIn(package._stop_background_work, registered)
# atexit runs newest first: the teardown before the history dispose.
self.assertGreater(
registered.index(package._stop_background_work),
registered.index(registration.dispose),
)
class ListHandler(logging.Handler):
def __init__(self):
super().__init__(logging.INFO)
self.lines = []
def emit(self, record):
self.lines.append(f"{record.levelname}:{record.getMessage()}")
class TestPackageTeardownOnce(PackageTestCase):
"""The tests that only observe one plain teardown share one package.
The teardown runs once, so it runs here in ``setUpClass`` (per-test
``assertLogs`` cannot wrap it); the tests read what it recorded.
"""
package_name = "pm_teardown_once"
@classmethod
def setUpClass(cls):
with contextlib.redirect_stdout(io.StringIO()):
super().setUpClass()
package = cls.package
cls.monitor = package._global_image_monitor
cls.observer = cls.monitor.observer
cls.tracker_thread = package._global_prompt_tracker.cleanup_thread
cls.worker = start_busy_worker(package, cls.tmp.name)
cls.recorder_before = cls.monitor.review_recorder
logger = logging.getLogger("prompt_manager.init")
handler, level = ListHandler(), logger.level
logger.addHandler(handler)
logger.setLevel(logging.INFO) # otherwise set by the config
try:
with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister:
cls.result = package._stop_background_work()
finally:
logger.removeHandler(handler)
logger.setLevel(level)
cls.log = "\n".join(handler.lines)
cls.unregister_calls = unregister.call_args_list
def test_teardown_stops_every_part_once(self):
expected = {"monitor": True, "thumbnails": True, "tracker": True}
self.assertEqual(self.result, expected)
self.assertEqual(list(self.result), ["monitor", "thumbnails", "tracker"])
self.assertIn("Background work stopped in", self.log)
self.assertIn("monitor=ok thumbnails=ok tracker=ok", self.log)
self.assertFalse(self.worker.is_alive())
self.assertFalse(self.observer.is_alive())
self.assertFalse(self.tracker_thread.is_alive())
copy = self.package._stop_background_work()
copy["monitor"] = False # a copy: the stored result is unchanged
with mock.patch.object(self.monitor, "stop_monitoring") as again:
self.assertEqual(self.package._stop_background_work(), expected)
again.assert_not_called()
def test_teardown_detaches_the_monitor_callbacks(self):
self.assertIsNotNone(self.recorder_before)
self.assertIsNone(self.monitor.review_recorder)
self.assertIsNone(self.monitor.thumbnail_queue)
self.assertIsNone(self.monitor.output_listener)
def test_teardown_removes_its_atexit_callback(self):
self.assertEqual(
self.unregister_calls, [mock.call(self.package._stop_background_work)]
)
class TestPackageTeardownFailures(PackageTeardownCase):
"""A failing step never skips the later ones (design D14, R2)."""
def assert_tracker_stops_after(self, error):
package = self.load(require_stopped=False)
monitor = package._global_image_monitor
tracker_thread = package._global_prompt_tracker.cleanup_thread
with (
mock.patch.object(monitor, "stop_monitoring", side_effect=error),
self.assertLogs("prompt_manager.init", "WARNING") as logs,
):
result = package._stop_background_work()
self.assertFalse(result["monitor"])
self.assertTrue(result["tracker"])
self.assertFalse(tracker_thread.is_alive())
text = "\n".join(logs.output)
self.assertIn("Background work not stopped in time", text)
self.assertRegex(text, r"not stopped in time \([0-9.]+ s\): monitor$")
self.assertIn("Traceback", text)
def test_teardown_step_failure_does_not_skip_later_steps(self):
self.assert_tracker_stops_after(RuntimeError("stop failed"))
def test_second_interrupt_does_not_skip_the_tracker(self):
self.assert_tracker_stops_after(KeyboardInterrupt())
def test_an_interrupt_between_steps_still_records_the_teardown(self):
package = self.load(require_stopped=False)
real_step = package._teardown_step
def interrupted_before_the_previews(name, stop, logger):
if name == "the preview worker":
raise KeyboardInterrupt # outside the step's own try
return real_step(name, stop, logger)
with (
mock.patch.object(
package, "_teardown_step", interrupted_before_the_previews
),
mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister,
self.assertRaises(KeyboardInterrupt),
):
package._stop_background_work()
expected = {"monitor": True, "thumbnails": False, "tracker": False}
self.assertEqual(package._background_work_result, expected)
unregister.assert_called_once_with(package._stop_background_work)
with mock.patch.object(
package._global_image_monitor, "stop_monitoring"
) as again:
self.assertEqual(package._stop_background_work(), expected)
again.assert_not_called()
def test_teardown_reports_a_preview_worker_still_held(self):
package = self.load(require_stopped=False)
held = self.hold_a_preview(package, advance=9.95, kill_releases=False)
with self.assertLogs("prompt_manager.init", "WARNING") as logs:
result = package._stop_background_work()
self.assertTrue(held.killed.is_set())
self.assertEqual(len(held.timeouts), 1)
self.assertAlmostEqual(held.timeouts[0], 0.05, places=6)
self.assertTrue(held.thumbnailer.closed)
self.assertFalse(result["thumbnails"], result)
self.assertRegex("\n".join(logs.output), r"not stopped in time .*thumbnails")
held.release.set()
held.worker.join(WAIT_SECONDS)
self.assertFalse(held.worker.is_alive())
self.assertEqual(held.stop_seen, [True])
def test_teardown_shares_one_deadline(self):
package = self.load(require_stopped=False)
now = [100.0]
seen = []
def step(timeout, **_):
seen.append(timeout)
now[0] += 6.0
return True
with (
mock.patch.object(package, "_clock", lambda: now[0]),
mock.patch.object(
package._global_image_monitor, "stop_monitoring", side_effect=step
),
mock.patch.object(
package.api,
"close_auto_thumbnails",
side_effect=lambda final=False, timeout=None: step(timeout),
),
mock.patch.object(
package._global_prompt_tracker, "shutdown", side_effect=step
),
):
package._stop_background_work()
self.assertEqual(seen, [5.0, 4.0, 0.0])
class TestRepeatedLoads(PackageTeardownCase):
def test_repeated_loads_leave_no_threads(self):
for round_ in range(3):
before = package_threads()
with contextlib.redirect_stdout(io.StringIO()):
with load_package(
f"pm_teardown_repeat_{round_}",
self.db_path,
self.output_dir,
) as (package, _):
self.busy_worker(package, f"busy_{round_}.png")
started = package_threads() - before
self.assertTrue(started, "the package started no thread")
left = {t for t in package_threads() - before if t.is_alive()}
self.assertEqual(left, set(), f"round {round_}")
+1 -1
View File
@@ -460,7 +460,7 @@ class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase):
async def test_second_legacy_autotag_batch_keeps_its_json_409(self):
from py.api.autotag_routes import _BUSY_ERROR
self._run_job("autotag", "Auto-tag images (legacy page)")
self._run_job("autotag", "Tag output images (legacy page)")
resp = await self.client.request(
"POST", "/prompt_manager/autotag/start", json={}
)
+3
View File
@@ -29,6 +29,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database import library_move as lm # noqa: E402
from database import models # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from tests.background_stops import close_handler # noqa: E402
from tests.open_handles import assert_closed # noqa: E402
from utils import library_location as ll # noqa: E402
from utils import usage_tracking # noqa: E402
@@ -204,6 +205,7 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
image = os.path.join(self.comfy, "output", "late.png")
writer_db = PromptDatabase(self.old_db)
handler = ImageGenerationHandler(writer_db, mock.MagicMock())
self.addCleanup(close_handler, self, handler)
started = threading.Event()
probe = _CrossThreadCloseProbe()
outcome = {}
@@ -256,6 +258,7 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
db = PromptDatabase(self.old_db)
handler = ImageGenerationHandler(db, mock.MagicMock())
self.addCleanup(close_handler, self, handler)
image = os.path.join(self.comfy, "output", "orphan.png")
with mock.patch.object(handler, "logger") as log:
handler.link_image_to_prompt(image, {"id": 99999}, {})
+95
View File
@@ -529,6 +529,16 @@ class LoraImportRouteTests(LoraImportAPITestCase):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "plan-used"))
async def test_a_finished_import_clears_the_layout_cache(self):
# #362: the import may link files under a folder LoraManager created
# after the image routes cached its layout.
self._two_loras()
with patch("py.lora_utils.clear_layout_cache") as clear:
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "finished", job.log)
clear.assert_called_once_with()
async def test_discovery_failure_keeps_previous_imports(self):
self._two_loras()
before = _snapshot(self.api.db)
@@ -936,6 +946,35 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
self.assertEqual((job.state, job.scope["failure"]), ("failed", "plan-changed"))
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_refused_import_keeps_the_layout_cache(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
self.api.db.save_prompt(
text="late", category=CATEGORY, prompt_hash=generate_prompt_hash("late")
)
with patch("py.lora_utils.clear_layout_cache") as clear:
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(job.scope["failure"], "plan-changed")
clear.assert_not_called()
async def test_start_refused_when_a_skipped_sidecar_changed(self):
# #362 CUJ 4: the review counted the orphan as skipped.
self._two_loras()
orphan = self.loras / "orphan.metadata.json"
orphan.write_text(json.dumps({"model_name": "orphan"}))
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual((plan["lora_files"], plan["skipped"]), (3, 1))
before = _snapshot(self.api.db)
orphan.write_text(orphan.read_text() + " ")
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 409, body)
self.assertEqual(body["code"], "plan-changed")
self.assertEqual(self.api._jobs.list(), [])
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_pair_saved_after_the_review_blocks_the_job(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
@@ -1187,6 +1226,62 @@ class LoraImportStageTests(unittest.TestCase):
)
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
# #362: a skipped sidecar is part of what the review counted.
def _assert_fingerprint_follows(self, change):
first = lora_import.fingerprint(self.discover())
self.assertEqual(lora_import.fingerprint(self.discover()), first)
change()
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
def _assert_add_edit_remove_change_it(self, path, text):
self._assert_fingerprint_follows(lambda: path.write_text(text))
self._assert_fingerprint_follows(
lambda: path.write_text(path.read_text() + " ")
)
self._assert_fingerprint_follows(path.unlink)
def _centralized(self):
sidecars = self.tmp / "sidecars"
sidecars.mkdir()
self.settings_file(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
}
)
model = self.loras / "kept.safetensors"
model.write_bytes(b"m")
(sidecars / "kept.metadata.json").write_text(
json.dumps({"file_path": str(model)})
)
return sidecars
def test_fingerprint_follows_a_centralized_sidecar_whose_model_is_missing(self):
sidecars = self._centralized()
gone = self.loras / "gone.safetensors"
self._assert_add_edit_remove_change_it(
sidecars / "gone.metadata.json", json.dumps({"file_path": str(gone)})
)
self.assertEqual(len(self.discover().sidecars), 1)
def test_fingerprint_follows_a_checkpoint_sidecar_outside_the_lora_roots(self):
sidecars = self._centralized()
checkpoints = self.tmp / "checkpoints"
checkpoints.mkdir()
ckpt = checkpoints / "base.safetensors"
ckpt.write_bytes(b"c")
self._assert_add_edit_remove_change_it(
sidecars / "base.metadata.json", json.dumps({"file_path": str(ckpt)})
)
def test_fingerprint_follows_an_alongside_sidecar_without_its_model(self):
_write_metadata(self.loras, "kept")
self._assert_add_edit_remove_change_it(
self.loras / "orphan.metadata.json", json.dumps({"model_name": "orphan"})
)
self.assertEqual(len(self.discover().sidecars), 1)
def test_fingerprint_follows_a_missing_file_between_stat_calls(self):
meta = _write_metadata(self.loras, "one")
discovery = self.discover()
+590
View File
@@ -0,0 +1,590 @@
"""The LoraManager layout cache for the image hot paths (#362).
``py/lora_layout_cache.py``. Every test injects the install resolver, the
layout loader (both counting calls), the wall clock, the monotonic clock and,
where needed, ``stat`` and ``pathmod``, over a temporary folder. No test
sleeps: modification times are set with ``os.utime`` and the clocks are
plain attributes the test moves.
"""
import ntpath
import os
import shutil
import sys
import tempfile
import threading
import time
import unittest
from dataclasses import replace
from pathlib import Path, PureWindowsPath
from types import SimpleNamespace
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from py import lora_layout_cache as cache_module
from py.lora_layout_cache import (
MAX_AGE_SECONDS,
MISSING,
RACY_SECONDS,
LayoutCache,
)
from py.lora_manager_layout import UNREADABLE, LoraManagerLayout, _canonical
SECOND_NS = 10**9
# NTFS keeps modification times in 100 ns steps and FAT in 2 s steps: every
# time a test sets is a whole millisecond, and a test whose file system rounds
# it anyway skips (``set_mtime``).
TICK_NS = 10**6
def make_layout(settings_path=None, media_roots=(), error=None, source="user config"):
return LoraManagerLayout(
settings_path=settings_path,
source=source,
error=error,
mode="alongside",
lora_roots=tuple(media_roots),
sidecar_root=None,
example_images_root=None,
library="default",
media_roots=tuple(media_roots),
)
class CacheTestCase(unittest.TestCase):
"""An install, its plugin settings file, a settings file and a LoRA root."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(shutil.rmtree, self.tmp)
self.install_dir = self.tmp / "custom_nodes" / "ComfyUI-Lora-Manager"
self.install_dir.mkdir(parents=True)
self.plugin_settings = self.install_dir / "settings.json"
self.settings = self.tmp / "user-config" / "settings.json"
self.settings.parent.mkdir()
self.root_dir = self.tmp / "loras"
self.root_dir.mkdir()
# Files are written ten minutes before the injected wall clock, so
# nothing is racy unless a test sets a recent modification time.
self.now_ns = time.time_ns() // TICK_NS * TICK_NS
self.mono = 1000.0
self.resolved = []
self.loaded = []
self.install = _canonical(str(self.install_dir))
self.layout = make_layout(
settings_path=self.settings,
media_roots=(_canonical(str(self.root_dir)),),
)
self.write(self.settings, "{}")
def write(self, path, text, mtime_ns=None):
Path(path).write_text(text, encoding="utf-8")
stamp = mtime_ns if mtime_ns is not None else self.now_ns - 600 * SECOND_NS
self.set_mtime(path, stamp)
def set_mtime(self, path, stamp):
os.utime(path, ns=(stamp, stamp))
if os.stat(path).st_mtime_ns != stamp:
self.skipTest("this file system rounds modification times")
def resolve(self, key):
self.resolved.append(key)
return self.install
def load(self, install):
self.loaded.append(install)
return self.layout
def cache(self, **kwargs):
kwargs.setdefault("wall_ns", lambda: self.now_ns)
kwargs.setdefault("monotonic", lambda: self.mono)
return LayoutCache(
kwargs.pop("resolve", self.resolve), kwargs.pop("load", self.load), **kwargs
)
class ValidationTests(CacheTestCase):
def test_hit_reuses_the_entry_without_resolving_or_loading(self):
cache = self.cache()
first = cache.get("")
second = cache.get("")
self.assertEqual(len(self.resolved), 1)
self.assertEqual(len(self.loaded), 1)
self.assertIs(first[1], second[1])
self.assertIs(first[0], second[0])
self.assertEqual(first, (self.install, self.layout))
def test_settings_size_change_reloads_once(self):
cache = self.cache()
cache.get("")
stamp = self.settings.stat().st_mtime_ns
self.write(self.settings, '{"a": 1}', mtime_ns=stamp)
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_settings_mtime_change_reloads_once(self):
cache = self.cache()
cache.get("")
self.set_mtime(self.settings, self.settings.stat().st_mtime_ns + TICK_NS)
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_settings_deletion_reloads_once(self):
cache = self.cache()
cache.get("")
self.settings.unlink()
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_settings_creation_reloads_once(self):
self.settings.unlink()
cache = self.cache()
cache.get("")
self.write(self.settings, "{}")
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_plugin_settings_change_reloads_once(self):
cache = self.cache()
cache.get("")
self.write(self.plugin_settings, '{"use_portable_settings": true}')
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_a_different_configured_path_reloads(self):
cache = self.cache()
cache.get("")
cache.get("other")
self.assertEqual(self.resolved, ["", "other"])
self.assertEqual(len(self.loaded), 2)
def test_a_different_home_folder_reloads(self):
# The folder guard compares against the home folder: a cached layout
# must not outlive a change of it.
home = {"value": str(self.tmp / "home-a")}
cache = self.cache(home=lambda: home["value"])
cache.get("")
cache.get("")
home["value"] = str(self.tmp / "home-b")
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_configured_path_is_stripped_for_the_key_and_the_resolver(self):
cache = self.cache()
cache.get(" x ")
cache.get("x")
self.assertEqual(self.resolved, ["x"])
self.assertEqual(len(self.loaded), 1)
def test_a_recently_modified_file_is_not_stored(self):
self.write(self.settings, "{}", mtime_ns=self.now_ns - SECOND_NS)
cache = self.cache()
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
self.assertIsNone(cache._entry)
self.now_ns += RACY_SECONDS * SECOND_NS
cache.get("")
self.assertIsNotNone(cache._entry)
cache.get("")
self.assertEqual(len(self.loaded), 3)
def test_racy_boundary(self):
# Exactly RACY_SECONDS old is trusted; anything newer is not.
boundary = self.now_ns - RACY_SECONDS * SECOND_NS
self.write(self.settings, "{}", mtime_ns=boundary + TICK_NS)
cache = self.cache()
cache.get("")
self.assertIsNone(cache._entry)
self.set_mtime(self.settings, boundary)
cache.get("")
self.assertIsNotNone(cache._entry)
def test_a_modification_time_in_the_future_is_racy(self):
self.write(self.plugin_settings, "{}", mtime_ns=self.now_ns + 60 * SECOND_NS)
cache = self.cache()
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 2)
self.assertIsNone(cache._entry)
def test_a_missing_file_is_not_racy(self):
self.assertFalse(self.plugin_settings.exists())
self.settings.unlink()
cache = self.cache()
cache.get("")
self.assertIsNotNone(cache._entry)
signatures = dict(cache._entry.signatures)
self.assertTrue(all(value is MISSING for value in signatures.values()))
self.assertEqual(len(signatures), 2)
def test_entry_is_reused_below_the_maximum_age_and_reloaded_at_it(self):
cache = self.cache()
cache.get("")
stored_at = self.mono
self.mono = stored_at + MAX_AGE_SECONDS - 0.5
cache.get("")
self.assertEqual(len(self.loaded), 1)
self.mono = stored_at + MAX_AGE_SECONDS
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_a_wall_clock_moved_backward_does_not_extend_the_age(self):
cache = self.cache()
cache.get("")
self.now_ns -= 3600 * SECOND_NS
self.mono += MAX_AGE_SECONDS
# The old entry is too old; the new load sees files "from the future"
# (racy) and is not stored, but it does load.
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_stat_failure_on_a_hit_reloads(self):
failing = {"on": False}
def stat(path):
if failing["on"]:
raise PermissionError(13, "sharing violation", path)
return os.stat(path)
cache = self.cache(stat=stat)
cache.get("")
failing["on"] = True
cache.get("")
self.assertEqual(len(self.loaded), 2)
self.assertIsNone(cache._entry)
def test_stat_failure_on_a_miss_leaves_the_layout_unstored(self):
def stat(path):
raise PermissionError(13, "sharing violation", path)
cache = self.cache(stat=stat)
self.assertEqual(cache.get(""), (self.install, self.layout))
self.assertIsNone(cache._entry)
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_an_unreadable_settings_layout_is_not_stored(self):
self.layout = make_layout(settings_path=self.settings, error=UNREADABLE)
cache = self.cache()
self.assertEqual(cache.get("")[1].error, UNREADABLE)
cache.get("")
self.assertEqual(len(self.loaded), 2)
self.assertIsNone(cache._entry)
def test_other_layout_errors_are_stored(self):
self.layout = make_layout(
settings_path=self.settings, error="the settings file is not valid JSON"
)
cache = self.cache()
cache.get("")
cache.get("")
self.assertEqual(len(self.loaded), 1)
def _symlink(self, link, target):
try:
os.symlink(target, link, target_is_directory=True)
except (OSError, NotImplementedError) as error:
self.skipTest(f"symlinks are not available here: {error}")
def test_a_media_root_swapped_for_a_symlink_reloads(self):
elsewhere = self.tmp / "elsewhere"
elsewhere.mkdir()
cache = self.cache()
cache.get("")
shutil.rmtree(self.root_dir)
self._symlink(self.root_dir, elsewhere)
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_an_install_swapped_for_a_symlink_reloads(self):
elsewhere = self.tmp / "elsewhere-lm"
elsewhere.mkdir()
cache = self.cache()
cache.get("")
moved = self.tmp / "moved-lm"
self.install_dir.rename(moved)
self._symlink(self.install_dir, elsewhere)
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_a_root_whose_realpath_fails_reloads(self):
cache = self.cache()
cache.get("")
with patch.object(
cache_module, "_canonical", side_effect=OSError("realpath failed")
):
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_a_stat_value_error_cannot_validate(self):
def stat(path):
raise ValueError("embedded null character in path")
cache = self.cache(stat=stat)
cache.get("")
self.assertIsNone(cache._entry)
def test_not_a_directory_is_missing(self):
def stat(path):
raise NotADirectoryError(20, "Not a directory", path)
cache = self.cache(stat=stat)
cache.get("")
self.assertTrue(
all(value is MISSING for _path, value in cache._entry.signatures)
)
cache.get("")
self.assertEqual(len(self.loaded), 1)
def test_an_offline_windows_share_cannot_validate(self):
def stat(path):
error = FileNotFoundError(2, "The network path was not found", path)
error.winerror = 53
raise error
cache = self.cache(stat=stat)
cache.get("")
self.assertIsNone(cache._entry)
def test_a_failed_reload_leaves_no_stale_entry(self):
failing = {"on": False}
def load(install):
self.loaded.append(install)
if failing["on"]:
raise RuntimeError("boom")
return self.layout
cache = self.cache(load=load)
cache.get("")
self.mono += MAX_AGE_SECONDS
failing["on"] = True
with self.assertRaises(RuntimeError):
cache.get("")
self.assertIsNone(cache._entry)
def test_a_resolve_failure_propagates_and_stores_nothing(self):
def resolve(key):
raise RuntimeError("boom")
cache = self.cache(resolve=resolve)
with self.assertRaises(RuntimeError):
cache.get("")
self.assertIsNone(cache._entry)
def test_without_a_settings_path_only_the_plugin_file_is_watched(self):
self.layout = make_layout(media_roots=self.layout.media_roots)
cache = self.cache()
cache.get("")
self.assertEqual(
[path for path, _sig in cache._entry.signatures],
[os.path.join(self.install, "settings.json")],
)
def test_portable_settings_are_watched_once(self):
self.layout = make_layout(settings_path=self.install_dir / "settings.json")
cache = self.cache()
cache.get("")
self.assertEqual(len(cache._entry.signatures), 1)
def test_a_load_failure_propagates_and_stores_nothing(self):
calls = []
def load(install):
calls.append(install)
if len(calls) == 1:
raise RuntimeError("boom")
return self.layout
cache = self.cache(load=load)
with self.assertRaises(RuntimeError):
cache.get("")
self.assertIsNone(cache._entry)
self.assertEqual(cache.get(""), (self.install, self.layout))
self.assertEqual(len(calls), 2)
def test_without_an_install_nothing_is_watched_and_the_entry_ages_out(self):
stats = []
def stat(path):
stats.append(path)
return os.stat(path)
self.layout = make_layout(media_roots=self.layout.media_roots, source=None)
cache = self.cache(resolve=lambda key: self.resolved.append(key), stat=stat)
self.assertEqual(cache.get(""), (None, self.layout))
self.assertEqual(cache._entry.signatures, ())
cache.get("")
self.assertEqual(len(self.loaded), 1)
self.assertEqual(stats, [])
self.mono += MAX_AGE_SECONDS
cache.get("")
self.assertEqual(len(self.loaded), 2)
def test_windows_paths_differing_only_in_case_are_watched_once(self):
stats = []
def stat(path):
stats.append(path)
return SimpleNamespace(st_mtime_ns=self.now_ns - 600 * SECOND_NS, st_size=2)
install = _canonical("C:\\ComfyUI\\custom_nodes\\LM", ntpath)
# The hit check needs ntpath's realpath to be idempotent here.
self.assertEqual(install, _canonical(install, ntpath))
self.layout = make_layout(
settings_path=PureWindowsPath(
"C:\\ComfyUI\\custom_nodes\\lm\\SETTINGS.JSON"
)
)
cache = self.cache(resolve=lambda key: install, stat=stat, pathmod=ntpath)
cache.get("")
self.assertEqual(len(cache._entry.signatures), 1)
self.assertEqual(len(stats), 1)
cache.get("")
self.assertEqual(len(self.loaded), 1)
self.assertEqual(len(stats), 2)
class ConcurrencyTests(CacheTestCase):
def test_a_load_that_started_before_clear_stores_nothing(self):
started = threading.Event()
release = threading.Event()
def load(install):
self.loaded.append(install)
if len(self.loaded) == 1:
started.set()
self.assertTrue(release.wait(10))
return self.layout
cache = self.cache(load=load)
results = []
worker = threading.Thread(target=lambda: results.append(cache.get("")))
worker.start()
self.assertTrue(started.wait(10))
cache.clear()
release.set()
worker.join(10)
self.assertFalse(worker.is_alive())
self.assertEqual(results, [(self.install, self.layout)])
self.assertIsNone(cache._entry)
cache.get("")
self.assertEqual(len(self.loaded), 2)
self.assertIsNotNone(cache._entry)
def test_concurrent_misses_all_get_a_layout_and_one_entry_is_stored(self):
cache = self.cache()
count = 6
barrier = threading.Barrier(count)
results = []
lock = threading.Lock()
def worker():
barrier.wait(10)
result = cache.get("")
with lock:
results.append(result)
threads = [threading.Thread(target=worker) for _ in range(count)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(10)
self.assertEqual(results, [(self.install, self.layout)] * count)
self.assertIsNotNone(cache._entry)
loads = len(self.loaded)
cache.get("")
self.assertEqual(len(self.loaded), loads)
def test_a_valid_hit_takes_no_lock(self):
cache = self.cache()
cache.get("")
results = []
with cache._lock:
worker = threading.Thread(target=lambda: results.append(cache.get("")))
worker.start()
worker.join(10)
alive = worker.is_alive()
worker.join(10)
self.assertFalse(alive)
self.assertEqual(results, [(self.install, self.layout)])
def test_clear_forces_a_reload(self):
cache = self.cache()
cache.get("")
cache.clear()
self.assertIsNone(cache._entry)
cache.get("")
self.assertEqual(len(self.loaded), 2)
class LoggingTests(CacheTestCase):
LOGGER = "prompt_manager.lora_layout_cache"
def get(self, cache):
with self.assertLogs(self.LOGGER, "DEBUG") as logs:
cache.get("")
# assertLogs needs one record; this marker is not counted.
cache_module.logger.debug("marker")
return [r for r in logs.records if r.getMessage() != "marker"]
def test_a_new_layout_logs_one_info_line(self):
cache = self.cache()
records = self.get(cache)
self.assertEqual([r.levelname for r in records], ["INFO"])
self.assertEqual(
records[0].getMessage(),
"LoraManager layout changed: user config settings, alongside storage, "
"1 media folders",
)
self.mono += MAX_AGE_SECONDS
self.assertEqual(self.get(cache), []) # an equal layout, reloaded
cache.clear()
self.assertEqual(self.get(cache), [])
self.layout = replace(self.layout, media_roots=(), lora_roots=())
self.mono += MAX_AGE_SECONDS
self.assertEqual(len(self.get(cache)), 1)
def test_no_settings_source_prints_no(self):
self.layout = replace(self.layout, source=None)
records = self.get(self.cache())
self.assertIn("layout changed: no settings,", records[0].getMessage())
def test_an_unstored_load_logs_no_info_line(self):
self.write(self.settings, "{}", mtime_ns=self.now_ns)
self.assertEqual(
[r for r in self.get(self.cache()) if r.levelname == "INFO"], []
)
def test_a_newly_refused_default_folder_logs_one_warning(self):
self.layout = replace(self.layout, refused_defaults=("/home",))
cache = self.cache()
records = self.get(cache)
warnings = [r for r in records if r.levelname == "WARNING"]
self.assertEqual(len(warnings), 1)
self.assertIn("/home", warnings[0].getMessage())
self.mono += MAX_AGE_SECONDS
self.assertEqual(self.get(cache), [])
cache.clear()
self.assertEqual(self.get(cache), [])
self.layout = replace(self.layout, refused_defaults=("/home", "/"))
self.mono += MAX_AGE_SECONDS
warnings = [r for r in self.get(cache) if r.levelname == "WARNING"]
self.assertEqual([w.getMessage().count("'/'") for w in warnings], [1])
class ModuleTests(unittest.TestCase):
def test_constants(self):
self.assertEqual(cache_module.RACY_SECONDS, 3)
self.assertEqual(cache_module.MAX_AGE_SECONDS, 30)
if __name__ == "__main__":
unittest.main()
+114 -3
View File
@@ -33,6 +33,10 @@ MEDIA = frozenset(
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".mp4", ".webm", ".mov", ".bmp"}
)
MARKER = "marker-from-unkept-settings-0123"
BAD_SIDECAR_REASON = (
"the sidecar folder is a filesystem root, the home folder or a folder "
"that contains it"
)
class LayoutTestCase(unittest.TestCase):
@@ -261,11 +265,26 @@ class UnusableLayoutTests(LayoutTestCase):
{"sidecar_storage_mode": "centralized", "sidecar_storage_path": "/"},
)
layout = self.load()
self.assertEqual(
layout.error, "the sidecar folder is a filesystem root or the home folder"
)
self.assertEqual(layout.error, BAD_SIDECAR_REASON)
self.assertNotIn(os.path.normcase(os.path.realpath("/")), layout.media_roots)
def test_a_centralized_sidecar_root_that_contains_home(self):
home = self.tmp / "outer" / "me"
home.mkdir(parents=True)
self.write_settings(
self.user_config,
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(self.tmp / "outer"),
},
)
with patch.object(layout_module.Path, "home", return_value=home):
layout = self.load()
self.assertEqual(layout.error, BAD_SIDECAR_REASON)
self.assertEqual(
layout.media_roots, (os.path.normcase(str(self.default_loras)),)
)
class SecretTests(LayoutTestCase):
@@ -501,6 +520,31 @@ class CentralizedDiscoveryTests(LayoutTestCase):
self.assertEqual(discovery.scanned, 3)
self.assertEqual(discovery.sidecars, [])
def test_every_scanned_file_is_listed_resolved_and_sorted(self):
# #362: the fingerprint covers every scanned sidecar, kept or skipped.
ckpt = self.model(self.checkpoints, "base")
alpha = self.model(self.loras, "alpha")
self.sidecar(self.sidecars / "z", "base", {"file_path": str(ckpt)})
self.sidecar(self.sidecars / "a", "alpha", {"file_path": str(alpha)})
self.sidecar(self.sidecars / "m", "nopath", {"model_name": "x"})
discovery = discover_sidecars(self.load())
files = discovery.scanned_files
self.assertIsInstance(files, tuple)
self.assertEqual(len(files), 3)
self.assertEqual(discovery.scanned, len(files))
self.assertEqual([str(f) for f in files], sorted(str(f) for f in files))
self.assertEqual(
[f.name for f in files],
[
f"alpha{METADATA_SUFFIX}",
f"nopath{METADATA_SUFFIX}",
f"base{METADATA_SUFFIX}",
],
)
for path in files:
self.assertEqual(str(path), os.path.realpath(path))
self.assertEqual(len(discovery.sidecars), 1)
def test_leftover_sidecars_beside_models_are_not_scanned(self):
model = self.model(self.loras, "alpha")
self.sidecar(self.loras, "alpha", {"file_path": str(model)})
@@ -611,6 +655,16 @@ class LocalMediaTests(LayoutTestCase):
self.touch(self.default_loras / "alpha.webp")
self.assertEqual(self.media({"preview_url": preview}), [preview])
def test_a_relative_preview_url_is_logged_and_skipped(self):
fallback = self.touch(self.default_loras / "alpha.webp")
self.touch(self.default_loras / "relative.png")
with self.assertLogs("prompt_manager.lora_manager_layout", "DEBUG") as logs:
media = self.media({"preview_url": "relative.png"})
self.assertEqual(media, [fallback])
self.assertEqual(len(logs.records), 1)
self.assertEqual(logs.records[0].levelname, "DEBUG")
self.assertIn(self.meta.name, logs.output[0])
def test_preview_url_outside_media_roots_or_not_servable_is_ignored(self):
outside = self.touch(self.tmp / "elsewhere" / "p.png")
avif = self.touch(self.default_loras / "p.avif")
@@ -810,6 +864,63 @@ class FolderGuardTests(unittest.TestCase):
self.assertFalse(layout_module._inside(str(self.tmp / "other.png"), (root,)))
class RefusedDefaultsTests(LayoutTestCase):
"""#362: ComfyUI's default LoRA folders pass the same guard; the refused
ones that exist are listed so the layout cache can warn about them."""
def setUp(self):
super().setUp()
self.outer = self.tmp / "outer"
self.home = self.outer / "me"
self.home.mkdir(parents=True)
self.fs_root = os.path.abspath(os.sep)
self.defaults = [
str(self.home),
str(self.outer),
self.fs_root,
str(self.tmp / "missing"),
str(self.default_loras),
]
self.expected = (
layout_module._canonical(str(self.home)),
layout_module._canonical(str(self.outer)),
layout_module._canonical(self.fs_root),
)
home_patch = patch.object(layout_module.Path, "home", return_value=self.home)
home_patch.start()
self.addCleanup(home_patch.stop)
def load_defaults(self, path="__plugin__"):
return self.load(path, default_roots=lambda: list(self.defaults))
def test_refused_existing_defaults_are_listed_on_a_usable_layout(self):
self.write_settings(self.user_config, {})
layout = self.load_defaults()
self.assertIsNone(layout.error)
self.assertEqual(layout.refused_defaults, self.expected)
self.assertEqual(
layout.lora_roots, (os.path.normcase(str(self.default_loras)),)
)
def test_they_are_listed_on_an_unusable_centralized_layout(self):
self.write_settings(
self.user_config,
{"sidecar_storage_mode": "centralized", "sidecar_storage_path": "/"},
)
layout = self.load_defaults()
self.assertEqual(layout.error, BAD_SIDECAR_REASON)
self.assertEqual(layout.refused_defaults, self.expected)
def test_they_are_listed_without_an_install(self):
layout = self.load_defaults(None)
self.assertEqual(layout.error, "no LoraManager install found")
self.assertEqual(layout.refused_defaults, self.expected)
def test_nothing_is_listed_when_every_default_is_accepted(self):
self.write_settings(self.user_config, {})
self.assertEqual(self.load().refused_defaults, ())
class SummaryTests(LayoutTestCase):
def test_summary_names_what_the_import_read(self):
+36
View File
@@ -224,6 +224,14 @@ class TestReadLoraMetadata(unittest.TestCase):
result = read_lora_metadata(Path("/nonexistent/file.metadata.json"))
self.assertIsNone(result)
def test_returns_none_for_a_file_that_is_not_utf8(self):
# #362: one malformed sidecar is skipped; it must not fail the whole
# trigger-word load (now logged with a traceback) or the review.
with tempfile.TemporaryDirectory() as folder:
path = Path(folder) / "bad.metadata.json"
path.write_bytes(b'{"model_name": "\xff\xfe"}')
self.assertIsNone(read_lora_metadata(path))
class TestGetLoraImageCacheDir(unittest.TestCase):
"""Test get_lora_image_cache_dir — returns and creates cache path."""
@@ -1791,6 +1799,34 @@ class TestTriggerWordCacheLayout(unittest.TestCase):
self.assertTrue(cache.is_loaded)
self.assertTrue(cache.retry_due())
def test_an_os_error_is_a_warning_without_a_traceback(self):
cache = TriggerWordCache()
cache._cache = {"old": ["word"]}
with (
patch.object(TriggerWordCache, "_build", side_effect=OSError("denied")),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual([r.levelname for r in logs.records], ["WARNING"])
self.assertIsNone(logs.records[0].exc_info)
self.assertEqual(cache.get_trigger_words("old"), ["word"])
self.assertFalse(cache.retry_due())
def test_an_unexpected_failure_logs_its_traceback(self):
# #362: a programming error must not look like a disk error.
cache = TriggerWordCache()
cache._cache = {"old": ["word"]}
with (
patch.object(TriggerWordCache, "_build", side_effect=ValueError("bug")),
self.assertLogs("prompt_manager.lora_utils", level="WARNING") as logs,
):
self.assertEqual(cache.load(str(self.lm)), 1)
self.assertEqual(len(logs.records), 1)
self.assertIsNotNone(logs.records[0].exc_info)
self.assertIs(logs.records[0].exc_info[0], ValueError)
self.assertEqual(cache.get_trigger_words("old"), ["word"])
self.assertFalse(cache.retry_due())
def test_any_other_failure_never_raises(self):
cache = TriggerWordCache()
with (
+2
View File
@@ -24,6 +24,7 @@ from unittest.mock import patch
from PIL import Image, PngImagePlugin
from database.operations import PromptDatabase
from tests.background_stops import close_handler
from tests.history_waits import drain_writer
from tests.test_api_admin import AdminAPITestCase
from tests.test_api_images import ImageAPITestCase
@@ -224,6 +225,7 @@ class TestMonitorIngestsNonFiniteMetadata(StrictRows, unittest.TestCase):
self.handler.logger.disabled = True
def tearDown(self):
close_handler(self, self.handler)
self.db.close_all()
assert_closed(self, self.temp.name)
self.temp.cleanup()
+142
View File
@@ -0,0 +1,142 @@
"""The package teardown runs at a normal interpreter exit (#301, #302).
ComfyUI catches Ctrl+C in ``main.py`` and returns, so its exit is a normal
return: ``atexit`` callbacks then run, newest first. A child process imports
the package with ComfyUI stubbed, leaves the image worker alive in its
processing delay and returns. Its own ``atexit`` reporter, registered before
the package, runs after the teardown and prints the result as one JSON line.
"""
import json
import os
import subprocess
import sys
import tempfile
import textwrap
import unittest
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# The parent never waits longer than this for the child (the hang watchdog
# allows 120 s per test).
CHILD_CEILING_SECONDS = 60
REPORT_PREFIX = "PM_EXIT_REPORT "
CHILD = textwrap.dedent("""
import atexit
import contextlib
import io
import json
import os
import sys
import threading
root, tmp = sys.argv[1], sys.argv[2]
sys.path.insert(0, root)
state = {}
def report():
from watchdog.utils import BaseThread
package = state.get("package")
result = getattr(package, "_background_work_result", None)
# The image worker, the tracker, and watchdog's observer and emitters.
alive = sorted(
t.name
for t in threading.enumerate()
if t.name in ("PromptManagerImageWorker", "PromptManagerPromptCleanup")
or isinstance(t, BaseThread)
)
print("PM_EXIT_REPORT " + json.dumps({"result": result, "alive": alive}))
sys.stdout.flush()
# First registered, so it runs last: after the package's teardown and
# after the history dispose.
atexit.register(report)
from tests.test_init_package import _import_package, fake_server_module
output = os.path.join(tmp, "output")
with contextlib.redirect_stdout(io.StringIO()):
package, _ = _import_package(
"pm_exit_child",
os.path.join(tmp, "prompts.db"),
output,
fake_server_module(queue_task_done=True),
)
state["package"] = package
monitor = package._global_image_monitor
handler = monitor.handler if monitor is not None else None
if handler is None:
sys.exit(3)
handler.processing_delay = 3600
path = os.path.join(tmp, "busy.png") # outside the watched folder
with open(path, "wb") as image:
image.write(b"x" * 16)
if not handler.schedule(path) or not handler._worker.is_alive():
sys.exit(4)
# Return normally with the worker alive, as ComfyUI does after Ctrl+C.
""")
class TestPackageExit(unittest.TestCase):
def test_teardown_runs_at_a_normal_exit_with_the_worker_alive(self):
with tempfile.TemporaryDirectory() as tmp:
os.makedirs(os.path.join(tmp, "output"))
script = os.path.join(tmp, "child.py")
with open(script, "w", encoding="utf-8") as handle:
handle.write(CHILD)
env = dict(os.environ)
env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
# The child's library and log file live in its own folder, never
# in the checkout's logs/ (which this runner may hold open) or a
# data folder the caller set.
env["PROMPT_MANAGER_DATA_DIR"] = tmp
env["PYTHONIOENCODING"] = "utf-8"
try:
completed = subprocess.run(
[sys.executable, script, ROOT, tmp],
cwd=ROOT,
env=env,
capture_output=True,
timeout=CHILD_CEILING_SECONDS,
encoding="utf-8",
errors="replace",
)
except subprocess.TimeoutExpired as hung:
self.fail(
f"the child did not exit within {CHILD_CEILING_SECONDS} s\n"
f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}"
)
# Leaving the block removes the folder: on Windows that fails
# if the child left a handle open, which it cannot once it exited,
# so the check here is that the child exited at all.
detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
self.assertEqual(completed.returncode, 0, detail)
reports = [
line[len(REPORT_PREFIX) :]
for line in completed.stdout.splitlines()
if line.startswith(REPORT_PREFIX)
]
self.assertEqual(len(reports), 1, detail)
report = json.loads(reports[0])
# "thumbnails" is True here because no preview worker started; its
# close result is tested in test_thumbnail_auto and test_init_package.
self.assertEqual(
report["result"],
{"monitor": True, "thumbnails": True, "tracker": True},
detail,
)
self.assertEqual(report["alive"], [], detail)
# The teardown's own warnings. Import-time warnings (the API's
# output-folder search finds no ComfyUI here) are not about exit.
for text in (
"Traceback",
"Background work not stopped in time",
"still running after",
"Error while stopping",
):
self.assertNotIn(text, completed.stderr, detail)
if __name__ == "__main__":
unittest.main()
+2
View File
@@ -23,6 +23,7 @@ try:
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from tests.background_stops import close_handler # noqa: E402
from database.operations import PromptDatabase
from prompt_manager import PromptManager
from prompt_manager_pair import PromptManagerPair
@@ -278,6 +279,7 @@ class MonitorCase(DatabaseCase):
fd, self.image = tempfile.mkstemp(suffix=".png", dir=self.tmp.name)
os.close(fd)
self.handler = ImageGenerationHandler(self.db, self.tracker)
self.addCleanup(close_handler, self, self.handler)
self.linked = []
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
self.linked.append(prompt["id"])
+90
View File
@@ -5,10 +5,12 @@ Tests PromptTracker, PromptExecutionContext, and the singleton
get_prompt_tracker. Uses a mock db_manager to avoid database deps.
"""
import logging
import os
import sys
import time
import threading
import types
import unittest
from unittest.mock import MagicMock, patch
@@ -312,5 +314,93 @@ class TestPromptQueue(unittest.TestCase):
self.assertEqual(popped_ids, ids)
def _real_tracker():
"""A PromptTracker built by its own ``__init__``, so its cleanup thread runs."""
mock_db = MagicMock()
mock_db.get_prompt_by_hash.return_value = None
return PromptTracker(mock_db)
class TestTrackerShutdown(unittest.TestCase):
"""The cleanup thread stops on a signal and is joined within a bound (#301)."""
def test_shutdown_stops_the_cleanup_thread_within_the_bound(self):
tracker = _real_tracker()
thread = tracker.cleanup_thread
self.assertTrue(thread.is_alive())
self.assertEqual(thread.name, "PromptManagerPromptCleanup")
self.assertTrue(tracker.shutdown())
self.assertFalse(thread.is_alive())
def test_shutdown_is_idempotent(self):
tracker = _real_tracker()
self.assertTrue(tracker.shutdown())
self.assertTrue(tracker.shutdown())
def test_cleanup_error_wait_is_interruptible(self):
import utils.prompt_tracker as tracker_mod
real_time = time.time
calls = {"n": 0}
def failing_once():
calls["n"] += 1
if calls["n"] == 1:
raise RuntimeError("clock unavailable")
return real_time()
stub = types.SimpleNamespace(
time=failing_once, sleep=time.sleep, monotonic=time.monotonic
)
logged = threading.Event()
logger = logging.getLogger("prompt_manager.prompt_tracker")
real_error = logger.error
def error(*args, **kwargs):
real_error(*args, **kwargs)
logged.set()
with (
patch.object(tracker_mod, "time", stub),
patch.object(logger, "error", side_effect=error),
):
tracker = _real_tracker()
self.addCleanup(tracker.shutdown)
self.assertTrue(logged.wait(30), "the cleanup error was never logged")
# The thread is now in its 60 s error wait.
self.assertTrue(tracker.shutdown())
self.assertFalse(tracker.cleanup_thread.is_alive())
def test_tracker_works_after_shutdown(self):
tracker = _real_tracker()
self.assertTrue(tracker.shutdown())
tracker.set_current_prompt("still tracked", {"prompt_id": 9})
self.assertEqual(tracker.pop_next_prompt()["id"], 9)
def test_shutdown_of_a_tracker_built_without_init(self):
self.assertTrue(_make_tracker().shutdown())
def test_shutdown_on_the_cleanup_thread_only_signals(self):
tracker = _real_tracker()
self.addCleanup(lambda: self.assertTrue(tracker.shutdown()))
result = {}
with patch.object(tracker, "cleanup_thread", threading.current_thread()):
result["value"] = tracker.shutdown()
self.assertFalse(result["value"])
self.assertTrue(tracker._stop.is_set())
def test_shutdown_reports_a_thread_still_running(self):
tracker = _real_tracker()
self.assertTrue(tracker.shutdown())
busy = MagicMock()
busy.is_alive.return_value = True
tracker.cleanup_thread = busy
with self.assertLogs("prompt_manager.prompt_tracker", "WARNING") as logs:
self.assertFalse(tracker.shutdown(timeout=0.05))
busy.join.assert_called_once_with(0.05)
self.assertIn("prompt cleanup thread still running", "\n".join(logs.output))
if __name__ == "__main__":
unittest.main()
+4 -1
View File
@@ -27,6 +27,7 @@ except ImportError: # discovered with tests/ as the top-level directory
from PIL import Image, PngImagePlugin
from tests.background_stops import close_handler, stop_monitor
from database.operations import PromptDatabase
from database.review_schema import canonical_evidence, evidence_digest
from py.api import PromptManagerAPI
@@ -790,6 +791,7 @@ class TestReviewEdges(ReviewTestCase):
class TestMonitorRecording(ReviewTestCase):
def _handler(self):
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
return handler
def _linked_prompt(self, path):
@@ -953,7 +955,7 @@ class TestMonitorRecording(ReviewTestCase):
monitor.set_review_recorder(None)
self.assertIsNone(monitor.handler.review_recorder)
finally:
monitor.stop_monitoring()
stop_monitor(self, monitor)
if __name__ == "__main__":
@@ -1092,6 +1094,7 @@ class TestCaptureStoreShutdown(ReviewTestCase):
self._scan_record(path)
# The monitor's pass records nothing, so it needs no store.
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.review_recorder = self.api.record_review_file
handler._record_review(str(path))
self.assertEqual(self.api.record_review_file(str(path)), "skipped")
+3 -1
View File
@@ -17,6 +17,7 @@ except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.operations import PromptDatabase
from tests.background_stops import close_handler
from tests.test_api_images import ImageAPITestCase
from tests.test_prompt_attribution import issue_graph, pair_evidence
from tools import review_image_links as tool
@@ -158,8 +159,8 @@ class ReviewToolTests(ToolFixture):
def test_conflicting_sources_are_left_to_the_maintainer(self):
db = PromptDatabase(self.db_path)
handler = tool.ImageGenerationHandler(db, None)
try:
handler = tool.ImageGenerationHandler(db, None)
row = {
"id": 1,
"prompt_id": self.coffee,
@@ -170,6 +171,7 @@ class ReviewToolTests(ToolFixture):
):
verdict = tool.judge(db, handler, row, self.output)
finally:
close_handler(self, handler)
db.close_all()
self.assertEqual(verdict["verdict"], "conflict")
+3
View File
@@ -28,6 +28,7 @@ except ImportError: # discovered with tests/ as the top-level directory
from test_api_review import ReviewRoutesTestCase
from test_review import unlinkable_graph, write_png
from tests.background_stops import close_handler
from utils import inspection_semantics, review_capture
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
@@ -437,6 +438,7 @@ class TestApplyInspectionCandidate(InspectionTestCase):
class TestMonitorNeverQueues(InspectionTestCase):
def _handler(self):
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.review_recorder = self.api.record_review_file
return handler
@@ -760,6 +762,7 @@ class TestRecordingEdges(InspectionTestCase):
text="latest", prompt_hash=generate_prompt_hash("latest")
)
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.link_image_to_prompt(
str(path),
{"id": prompt_id, "text": "latest", "fallback": True},
+2
View File
@@ -38,6 +38,7 @@ except ImportError: # discovered with tests/ as the top-level directory
write_png,
)
from tests.background_stops import close_handler
from utils.hashing import generate_prompt_hash
from utils.image_monitor import ImageGenerationHandler
from utils.prompt_tracker import PromptTracker
@@ -299,6 +300,7 @@ class TestQueueSelfHeals(OutcomeTestCase):
class TestMonitorRule(OutcomeTestCase):
def _handler(self):
handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db))
self.addCleanup(close_handler, self, handler)
handler.review_recorder = self.api.record_review_file
return handler
+4 -2
View File
@@ -11,6 +11,7 @@ from unittest import mock
from database.run_attribution import stored_run_prompt
from tests import test_execution_history_database as history_tests
from tests.background_stops import close_handler
from tests.test_api_admin import AdminAPITestCase
from tests.test_execution_context_database import snapshot
from tests.test_prompt_attribution import issue_graph, llm_text_graph, make_tracker
@@ -151,6 +152,7 @@ class TestStoredRunPrompt(RecordedRuns, unittest.TestCase):
prompt_id = self.save(LLM_TEXT)
self.record_run(self.db, self.output, [("sub", "img_00001_.png")], [LLM_TEXT])
handler = ImageGenerationHandler(self.db, make_tracker(self.db))
self.addCleanup(close_handler, self, handler)
linked = []
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
prompt["id"]
@@ -332,7 +334,7 @@ class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase):
self.handler = ImageGenerationHandler(self.db, make_tracker(self.db))
self.handler.processing_delay = 0
self.handler.settle_interval = 0.01
self.addCleanup(self.handler.shutdown)
self.addCleanup(close_handler, self, self.handler)
self.linked = []
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
self.linked.append((os.path.basename(path), prompt["id"]))
@@ -650,7 +652,7 @@ class TestLinkedOutputRoot(RecordedRuns, unittest.TestCase):
handler = ImageGenerationHandler(db, make_tracker(db))
handler.processing_delay = 0
handler.settle_interval = 0.01
self.addCleanup(handler.shutdown)
self.addCleanup(close_handler, self, handler)
linked = []
handler.link_image_to_prompt = lambda path, prompt, meta: linked.append(
prompt["id"]
+153
View File
@@ -976,3 +976,156 @@ class TestSourceFormError(unittest.TestCase):
if __name__ == "__main__":
unittest.main()
class TestAutotagDefaultsSettings(SourcesTestCase):
"""#328: the tagging defaults save through POST /settings like the others."""
ATTRS = (
"MODEL_TYPE",
"EXISTING",
"GENERAL_THRESHOLD",
"CHARACTER_THRESHOLD",
"TEMPLATE",
"SUGGESTION_LIMIT",
)
SAVED = {
"autotag_model_type": "gguf",
"autotag_existing": "merge",
"autotag_general_threshold": 0.4,
"autotag_character_threshold": 0.8,
"autotag_template": "Tags only.",
"autotag_suggestion_limit": 15,
}
async def asyncSetUp(self):
await super().asyncSetUp()
from py.config import AutotagConfig
self.autotag = AutotagConfig
for attr in self.ATTRS:
self.addCleanup(setattr, AutotagConfig, attr, getattr(AutotagConfig, attr))
def _runtime(self):
return {attr: getattr(self.autotag, attr) for attr in self.ATTRS}
async def test_autotag_defaults_save_round_trip(self):
self._write_config({"custom": {"keep": True}})
revision = (await self._get())["revision"]
status, data = await self._post(
"/prompt_manager/settings", {"expected_revision": revision, **self.SAVED}
)
self.assertEqual(status, 200, data)
self.assertEqual(data["restart_required_fields"], [])
saved = self._read_config()
self.assertEqual(
saved["autotag"]["defaults"],
{key[len("autotag_") :]: value for key, value in self.SAVED.items()},
)
self.assertEqual(saved["custom"], {"keep": True})
self.assertEqual(
self._runtime(),
{
"MODEL_TYPE": "gguf",
"EXISTING": "merge",
"GENERAL_THRESHOLD": 0.4,
"CHARACTER_THRESHOLD": 0.8,
"TEMPLATE": "Tags only.",
"SUGGESTION_LIMIT": 15,
},
)
settings = (await self._get())["settings"]
self.assertEqual({key: settings[key] for key in self.SAVED}, self.SAVED)
async def test_threshold_near_the_grid_is_saved_on_it(self):
"""#369: a computed value a hair off the 0.05 grid is stored on it, so
every page's step check accepts what /settings and the file return."""
revision = (await self._get())["revision"]
status, data = await self._post(
"/prompt_manager/settings",
{"expected_revision": revision, "autotag_general_threshold": 0.1 + 0.2},
)
self.assertEqual(status, 200, data)
self.assertEqual(
repr(self._read_config()["autotag"]["defaults"]["general_threshold"]),
"0.3",
)
self.assertEqual(
repr((await self._get())["settings"]["autotag_general_threshold"]), "0.3"
)
async def test_stale_revision_changes_no_default(self):
revision = (await self._get())["revision"]
self._write_config({"custom": {"edited": "elsewhere"}})
before = (self._config_bytes(), self._runtime())
status, data = await self._post(
"/prompt_manager/settings", {"expected_revision": revision, **self.SAVED}
)
self.assertEqual(status, 409, data)
self.assertEqual(data["code"], "settings-changed")
self.assertEqual((self._config_bytes(), self._runtime()), before)
async def test_bad_default_message_names_the_setting(self):
cases = (
(
"autotag_general_threshold",
1.5,
"General-tag confidence must be between 0 and 1.",
),
(
"autotag_general_threshold",
0.37,
"General-tag confidence must be a multiple of 0.05.",
),
(
"autotag_character_threshold",
0.33,
"Character-tag confidence must be a multiple of 0.05.",
),
(
"autotag_character_threshold",
"0.8",
"Character-tag confidence must be a finite number.",
),
(
"autotag_character_threshold",
True,
"Character-tag confidence must be a finite number.",
),
(
"autotag_model_type",
"wd14-convnext",
"Tagging model must be one of gguf, hf, wd14-swinv2, wd14-vit.",
),
(
"autotag_existing",
"replace",
"Existing tags must be skip, merge or replace_after_review.",
),
("autotag_template", " ", "LLM tag template must be non-empty text."),
(
"autotag_template",
"x" * 4001,
"LLM tag template must be at most 4000 characters.",
),
(
"autotag_suggestion_limit",
0,
"Suggestions per image must be a whole number from 1 to 100.",
),
(
"autotag_suggestion_limit",
2.5,
"Suggestions per image must be a whole number from 1 to 100.",
),
)
self._write_config({"custom": {"keep": True}})
before = (self._config_bytes(), self._runtime())
for key, value, message in cases:
with self.subTest(key=key, value=value):
status, data = await self._post(
"/prompt_manager/settings", {**self.SAVED, key: value}
)
self.assertEqual(status, 400, data)
self.assertEqual(data["error"], message)
self.assertEqual((self._config_bytes(), self._runtime()), before)
+20 -2
View File
@@ -32,6 +32,7 @@ from PIL import Image # noqa: E402
import py.api.images as images_module # noqa: E402
import py.api.thumbnail_auto as auto # noqa: E402
from tests.background_stops import close_handler, stop_monitor # noqa: E402
from py.api.jobs import THUMBNAILS_JOB_KIND # noqa: E402
SERVE = "/prompt_manager/images/serve/"
@@ -737,11 +738,27 @@ class TestShutdown(AutoThumbnailCase):
patch.object(auto, "CLOSE_WAIT_SECONDS", 0),
self.assertLogs("prompt_manager.api", "WARNING") as logs,
):
self.api.close_auto_thumbnails()
self.assertFalse(self.api.close_auto_thumbnails())
process.kill.assert_called_once()
self.assertIn("still running", "\n".join(logs.output))
release.set()
async def test_close_reports_an_ended_worker(self):
path = make_png(self.output_dir / "a.png", size=(64, 64))
self.api.queue_capture_thumbnail(str(path))
worker = self.api._auto_thumbnails()._worker
await self._idle()
loop = asyncio.get_running_loop()
closed = await loop.run_in_executor(
None, lambda: self.api.close_auto_thumbnails(timeout=WAIT)
)
self.assertTrue(closed)
self.assertFalse(worker.is_alive())
def test_close_without_a_thumbnailer_returns_true(self):
self.assertIsNone(self.api._auto_thumbnailer)
self.assertTrue(self.api.close_auto_thumbnails())
class TestHistorySharesTheAdmission(AutoThumbnailCase):
async def _shared_flight(self, name, is_video, first):
@@ -980,6 +997,7 @@ class TestMonitorQueuesThumbnails(AutoThumbnailCase):
from utils.image_monitor import ImageGenerationHandler
handler = ImageGenerationHandler(self.api.db, unittest.mock.MagicMock())
self.addCleanup(close_handler, self, handler)
handler.processing_delay = 0
handler.settle_interval = 0
return handler
@@ -1022,7 +1040,7 @@ class TestMonitorQueuesThumbnails(AutoThumbnailCase):
monitor.set_thumbnail_queue(None)
self.assertIsNone(monitor.handler.thumbnail_queue)
finally:
monitor.stop_monitoring()
stop_monitor(self, monitor)
class TestScanHandsOffThumbnails(AutoThumbnailCase):
+2
View File
@@ -16,6 +16,7 @@ import uuid
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import stop_tracker # noqa: E402
from database.operations import PromptDatabase
from prompt_manager_base import PromptManagerBase
from utils.hashing import generate_prompt_hash
@@ -596,6 +597,7 @@ class TestNodeRoles(DbTestCase):
class TestTrackerQueueFlag(unittest.TestCase):
def setUp(self):
self.tracker = PromptTracker(mock.Mock())
self.addCleanup(stop_tracker, self, self.tracker)
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)
+539 -10
View File
@@ -1,8 +1,10 @@
"""PromptManagerLogger: memory buffer, file rotation, level changes, truncation."""
import logging
import ntpath
import os
import pathlib
import stat
import sys
import tempfile
import threading
@@ -11,14 +13,20 @@ 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 (
from utils import logging_config # noqa: E402
from utils.logging_config import ( # noqa: E402
MemoryBufferHandler,
PromptManagerLogger,
RotatedChanged,
get_logger,
get_logger_manager,
)
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
class IsolatedLoggerTestCase(unittest.TestCase):
"""A logger manager whose log directory is a temporary folder."""
@@ -32,6 +40,7 @@ class IsolatedLoggerTestCase(unittest.TestCase):
with mock.patch.object(logging_config, "__file__", fake_file):
self.manager = object.__new__(PromptManagerLogger) # bypass singleton
self.manager.__init__()
self.allowed_thread_names = set()
self.manager.update_config({"console_logging": False})
# Drop the "Updated logging configuration" entry.
self.manager._log_buffer.clear()
@@ -53,12 +62,21 @@ class IsolatedLoggerTestCase(unittest.TestCase):
def _filter_to_this_thread(self):
this_thread = threading.get_ident()
# A test may let named helper threads of its own through too.
allowed_names = self.allowed_thread_names
def allowed(record):
return record.thread == this_thread or record.threadName in allowed_names
for handler in self.manager.logger.handlers:
handler.addFilter(lambda record: record.thread == this_thread)
handler.filters.clear()
handler.addFilter(allowed)
def _restore_global_logging(self):
for handler in logging.getLogger("prompt_manager").handlers[:]:
handler.close()
# Windows refuses to delete an open log: prove every handle is closed.
assert_closed(self, self.tmp.name)
get_logger_manager()._setup_loggers()
self.tmp.cleanup()
@@ -175,15 +193,14 @@ class TestFilesAndRotation(IsolatedLoggerTestCase):
def test_unreadable_files_are_skipped_in_the_listing(self):
self.logger.info("something")
original_stat = pathlib.Path.stat
original_lstat = os.lstat
def stat_failing_for_logs(path, *args, **kwargs):
if path.suffix == ".log":
def lstat_failing_for_logs(path, *args, **kwargs):
if str(path).endswith(".log"):
raise OSError("denied")
return original_stat(path, *args, **kwargs)
return original_lstat(path, *args, **kwargs)
with mock.patch.object(pathlib.Path, "stat", autospec=True) as stat:
stat.side_effect = stat_failing_for_logs
with mock.patch.object(logging_config.os, "lstat", lstat_failing_for_logs):
self.assertEqual(self.manager.get_log_files(), [])
def test_truncate_does_not_roll_the_log_over_while_deleting_backups(self):
@@ -228,10 +245,522 @@ class TestFilesAndRotation(IsolatedLoggerTestCase):
def test_truncate_reports_files_it_cannot_remove(self):
self._rotate()
with mock.patch.object(pathlib.Path, "unlink", side_effect=OSError("busy")):
def busy(path, *args, **kwargs):
raise PermissionError(13, "busy", str(path))
with mock.patch.object(logging_config.os, "unlink", busy):
result = self.manager.truncate_logs()
self.assertEqual(len(result["errors"]), 2)
self.assertEqual(result["truncated"], ["prompt_manager.log"])
for error in result["errors"]:
self.assertNotIn(self.tmp.name, error)
self.assertRegex(error, r"prompt_manager\.log\.[12]")
class TestRotatedCopyRule(IsolatedLoggerTestCase):
"""Rule R (#329): only prompt_manager.log.<n> regular files are deletable."""
def setUp(self):
super().setUp()
self.logger.info("active log line") # the active log exists
self.outside = os.path.join(self.tmp.name, "outside")
os.makedirs(self.outside)
def _path(self, name):
return os.path.join(self.log_dir, name)
def _write(self, name, data=b"rotated\n"):
with open(self._path(name), "wb") as handle:
handle.write(data)
return self._path(name)
def _symlink(self, name, target_name="secret.txt", data=b"secret " * 1000):
"""A symlink in the log folder to a file outside it, or None."""
target = os.path.join(self.outside, target_name)
with open(target, "wb") as handle:
handle.write(data)
try:
os.symlink(target, self._path(name))
except (OSError, NotImplementedError): # pragma: no cover - no privilege
return None
return target
def _require_symlink(self, name, **kwargs):
target = self._symlink(name, **kwargs)
if target is None: # pragma: no cover - Windows without the privilege
self.skipTest("this account cannot create symlinks")
return target
def _kinds(self):
return {row["name"]: row["kind"] for row in self.manager.scan_log_folder()}
def _rotated_rows(self):
return [r for r in self.manager.scan_log_folder() if r["kind"] == "rotated"]
def _names(self):
return sorted(os.listdir(self.log_dir))
def _file_handler(self):
return next(
h
for h in self.manager.logger.handlers
if isinstance(h, logging.FileHandler)
)
def test_scan_classes_rotated_active_link_and_other(self):
for name in ("prompt_manager.log.1", "prompt_manager.log.2"):
self._write(name)
self._write("prompt_manager.log.3", b"x" * 3000)
for name in (
"prompt_manager.log.0",
"prompt_manager.log.01",
"prompt_manager.log.bak",
"prompt_manager.log.1.gz",
"prompt_manager.log.-1",
"notes.txt",
):
self._write(name)
os.makedirs(self._path("prompt_manager.log.5"))
self._require_symlink("prompt_manager.log.4")
kinds = self._kinds()
self.assertEqual(
kinds,
{
"prompt_manager.log": "active",
"prompt_manager.log.1": "rotated",
"prompt_manager.log.2": "rotated",
"prompt_manager.log.3": "rotated",
"prompt_manager.log.0": "other",
"prompt_manager.log.01": "other",
"prompt_manager.log.bak": "other",
"prompt_manager.log.1.gz": "other",
"prompt_manager.log.-1": "other",
"notes.txt": "other",
"prompt_manager.log.5": "other",
"prompt_manager.log.4": "link",
},
)
row = next(
r
for r in self.manager.scan_log_folder()
if r["name"] == "prompt_manager.log.3"
)
self.assertEqual(
set(row), {"name", "kind", "size", "mtime_ns", "ino", "modified"}
)
stat = os.lstat(self._path("prompt_manager.log.3"))
self.assertEqual(row["size"], 3000)
self.assertEqual(row["mtime_ns"], stat.st_mtime_ns)
self.assertEqual(row["ino"], stat.st_ino)
def test_scan_matches_names_case_insensitively_on_windows(self):
self._write("PROMPT_MANAGER.LOG.2")
if os.name == "nt": # pragma: no cover - native case folding
self.assertEqual(self._kinds()["PROMPT_MANAGER.LOG.2"], "rotated")
return
# POSIX names are case-sensitive: another file, never deleted.
self.assertEqual(self._kinds()["PROMPT_MANAGER.LOG.2"], "other")
with mock.patch.object(logging_config, "_name_key", ntpath.normcase):
self.assertEqual(self._kinds()["PROMPT_MANAGER.LOG.2"], "rotated")
result = self.manager.delete_rotated(self._rotated_rows())
self.assertEqual(result["removed"], ["PROMPT_MANAGER.LOG.2"])
def test_a_junction_is_classed_as_a_link_and_kept(self):
# Portable control: whatever os.path.isjunction reports is a link.
self._write("prompt_manager.log.4")
self._write("prompt_manager.log.1")
junction = self._path("prompt_manager.log.4")
real_isjunction = os.path.isjunction
def isjunction(path):
return os.fspath(path) == junction or real_isjunction(path)
with mock.patch.object(logging_config.os.path, "isjunction", isjunction):
self.assertEqual(self._kinds()["prompt_manager.log.4"], "link")
self.manager.delete_rotated(self._rotated_rows())
self.manager.truncate_logs()
self.assertIn("prompt_manager.log.4", self._names())
self.assertNotIn("prompt_manager.log.1", self._names())
@unittest.skipUnless(os.name == "nt", "junctions exist only on Windows")
def test_junction_is_never_deleted_on_windows(self): # pragma: no cover
import _winapi
target = os.path.join(self.outside, "kept")
os.makedirs(target)
with open(os.path.join(target, "inside.txt"), "wb") as handle:
handle.write(b"keep me")
_winapi.CreateJunction(target, self._path("prompt_manager.log.4"))
self._write("prompt_manager.log.1")
self.assertEqual(self._kinds()["prompt_manager.log.4"], "link")
self.manager.delete_rotated(self._rotated_rows())
self.manager.truncate_logs()
self.assertIn("prompt_manager.log.4", self._names())
self.assertTrue(os.path.exists(os.path.join(target, "inside.txt")))
with self.assertRaises(FileNotFoundError):
self.manager.log_file_entry("prompt_manager.log.4")
os.rmdir(self._path("prompt_manager.log.4")) # removes the junction only
def test_delete_rotated_removes_exactly_the_listed_files(self):
for n in (1, 2, 3):
self._write(f"prompt_manager.log.{n}", b"r" * n)
self._write("prompt_manager.log.bak")
os.makedirs(self._path("prompt_manager.log.7"))
target = self._symlink("prompt_manager.log.4")
listed = [
r for r in self._rotated_rows() if r["name"] != "prompt_manager.log.3"
]
result = self.manager.delete_rotated(listed)
self.assertEqual(
result,
{"removed": ["prompt_manager.log.1", "prompt_manager.log.2"], "failed": []},
)
expected = [
"prompt_manager.log",
"prompt_manager.log.3",
"prompt_manager.log.7",
"prompt_manager.log.bak",
]
if target is not None:
expected.append("prompt_manager.log.4")
self.assertTrue(os.path.exists(target))
self.assertEqual(self._names(), sorted(expected))
def test_delete_rotated_refuses_all_when_one_changed(self):
def change_size(path):
with open(path, "ab") as handle:
handle.write(b"more")
def change_mtime(path):
stat = os.lstat(path)
os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns + 5_000_000_000))
def change_inode(path):
stat = os.lstat(path)
with open(path, "rb") as handle:
data = handle.read()
other = path + ".new"
with open(other, "wb") as handle:
handle.write(data)
os.utime(other, ns=(stat.st_atime_ns, stat.st_mtime_ns))
os.replace(other, path)
if os.lstat(path).st_ino in (0, stat.st_ino): # pragma: no cover
self.skipTest("this filesystem reports no inode numbers")
def remove(path):
os.unlink(path)
def replace_by_link(path):
os.unlink(path)
target = os.path.join(self.outside, "elsewhere.txt")
with open(target, "wb") as handle:
handle.write(b"rotated\n")
try:
os.symlink(target, path)
except (OSError, NotImplementedError): # pragma: no cover
self.skipTest("this account cannot create symlinks")
changes = {
"size": change_size,
"mtime": change_mtime,
"inode": change_inode,
"missing": remove,
"link": replace_by_link,
}
for label, change in changes.items():
with self.subTest(change=label):
for n in (1, 2, 3):
self._write(f"prompt_manager.log.{n}", b"rotated\n")
reviewed = self._rotated_rows()
change(self._path("prompt_manager.log.2"))
before = self._names()
with self.assertRaises(RotatedChanged):
self.manager.delete_rotated(reviewed)
self.assertEqual(self._names(), before)
self.assertIn("prompt_manager.log.1", before)
self.assertIn("prompt_manager.log.3", before)
for name in self._names():
if name.startswith("prompt_manager.log."):
os.unlink(self._path(name))
def test_rollover_waits_while_cleanup_holds_the_handler_lock(self):
# The file is already over its limit: the next record rolls it over,
# renaming .1 to .2 and the active log to .1.
self.manager.update_config({"max_file_size": 64, "backup_count": 3})
self.logger.info("x" * 100)
self._write("prompt_manager.log.1", b"first copy\n")
self._write("prompt_manager.log.2", b"second copy\n")
reviewed = self._rotated_rows()
holding = threading.Event()
release = threading.Event()
logged = threading.Event()
real_unlink = os.unlink
self.allowed_thread_names.add("pm-test-rollover")
def log_from_another_thread():
self.logger.info("rolls the log over")
logged.set()
worker = threading.Thread(
target=log_from_another_thread, name="pm-test-rollover", daemon=True
)
def slow_unlink(path, *args, **kwargs):
if not holding.is_set():
worker.start() # its record must wait for the cleanup's lock
holding.set()
if not release.wait(10):
raise RuntimeError("the test never released the cleanup")
return real_unlink(path, *args, **kwargs)
result = {}
def cleanup():
result.update(self.manager.delete_rotated(reviewed))
with mock.patch.object(logging_config.os, "unlink", slow_unlink):
cleaner = threading.Thread(target=cleanup, daemon=True)
cleaner.start()
self.assertTrue(holding.wait(10), "the cleanup never started deleting")
# The record waits on the handler lock the cleanup holds.
self.assertFalse(logged.wait(0.3), "a record was written mid-cleanup")
release.set()
cleaner.join(10)
worker.join(10)
self.assertFalse(cleaner.is_alive())
self.assertTrue(logged.is_set())
self.assertEqual(
result,
{"removed": ["prompt_manager.log.1", "prompt_manager.log.2"], "failed": []},
)
# The rollover ran after the cleanup: only the old active log moved.
self.assertEqual(self._names(), ["prompt_manager.log", "prompt_manager.log.1"])
with open(self._path("prompt_manager.log.1"), "rb") as handle:
self.assertIn(b"x" * 100, handle.read())
def test_open_rotated_copy_is_reported_not_deleted(self):
for n in (1, 2, 3):
self._write(f"prompt_manager.log.{n}")
real_unlink = os.unlink
held = self._path("prompt_manager.log.2")
def unlink(path, *args, **kwargs):
if os.fspath(path) == held:
error = PermissionError(
13,
"The process cannot access the file because it is being "
"used by another process",
held,
)
error.winerror = 32
raise error
return real_unlink(path, *args, **kwargs)
with mock.patch.object(logging_config.os, "unlink", unlink):
result = self.manager.delete_rotated(self._rotated_rows())
self.assertEqual(
result["removed"], ["prompt_manager.log.1", "prompt_manager.log.3"]
)
[(name, error)] = result["failed"]
self.assertEqual(name, "prompt_manager.log.2")
self.assertIsInstance(error, PermissionError)
self.assertIn("prompt_manager.log.2", self._names())
@unittest.skipUnless(os.name == "nt", "only Windows refuses to unlink an open file")
def test_open_rotated_copy_is_reported_not_deleted_on_windows(
self,
): # pragma: no cover
for n in (1, 2, 3):
self._write(f"prompt_manager.log.{n}")
reviewed = self._rotated_rows()
handle = open(self._path("prompt_manager.log.2"), "rb")
try:
result = self.manager.delete_rotated(reviewed)
finally:
handle.close()
self.assertEqual(
result["removed"], ["prompt_manager.log.1", "prompt_manager.log.3"]
)
[(name, error)] = result["failed"]
self.assertEqual(name, "prompt_manager.log.2")
self.assertIsInstance(error, PermissionError)
self.assertIn("prompt_manager.log.2", self._names())
def test_cleanup_never_logs_while_holding_the_handler_lock(self):
for n in (1, 2):
self._write(f"prompt_manager.log.{n}")
handler = self._file_handler()
depth = 0
calls = []
real_acquire, real_release = handler.acquire, handler.release
def acquire():
nonlocal depth
real_acquire()
depth += 1
def release():
nonlocal depth
depth -= 1
real_release()
real_call_handlers = logging.Logger.callHandlers
def call_handlers(logger_self, record):
calls.append((record.getMessage(), depth))
return real_call_handlers(logger_self, record)
with (
mock.patch.object(handler, "acquire", acquire),
mock.patch.object(handler, "release", release),
mock.patch.object(logging.Logger, "callHandlers", call_handlers),
):
self.manager.delete_rotated(self._rotated_rows())
self.assertEqual(depth, 0)
self._write("prompt_manager.log.1")
self.manager.truncate_logs()
self.assertEqual(depth, 0)
self.assertTrue(calls, "the cleanup logged no summary")
self.assertEqual([c for c in calls if c[1]], [])
def test_truncate_keeps_other_files_and_links(self):
# Corrected (#329): the legacy glob deleted every prompt_manager.log*.
self._rotate()
self._write("prompt_manager.log.bak", b"mine")
self._write("prompt_manager.log.1.gz", b"zip")
target = self._symlink("prompt_manager.log.4")
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"],
)
names = self._names()
self.assertIn("prompt_manager.log.bak", names)
self.assertIn("prompt_manager.log.1.gz", names)
self.assertNotIn("prompt_manager.log.1", names)
if target is not None:
self.assertIn("prompt_manager.log.4", names)
self.assertTrue(os.path.islink(self._path("prompt_manager.log.4")))
with open(target, "rb") as handle:
self.assertEqual(handle.read(), b"secret " * 1000)
@unittest.skipIf(os.name == "nt", "Windows cannot rename a log that is open")
def test_truncate_never_empties_a_renamed_log_the_handler_still_writes(self):
# The user renamed the active log to keep it; the handler still has it
# open. Clear must empty the new prompt_manager.log, not the kept copy.
self.logger.info("worth keeping")
kept = os.path.join(self.log_dir, "saved_copy.txt")
os.rename(self._path("prompt_manager.log"), kept)
with open(kept, "rb") as handle:
before = handle.read()
self._write("prompt_manager.log", b"fresh active log\n")
result = self.manager.truncate_logs()
self.assertEqual(result["errors"], [])
with open(kept, "rb") as handle:
self.assertTrue(handle.read().startswith(before))
with open(self._path("prompt_manager.log"), "rb") as handle:
self.assertEqual(handle.read(), b"")
def test_name_surrogate_reparse_points_are_links_and_placeholders_are_not(self):
regular = stat.S_IFREG | 0o644
path = self._write("prompt_manager.log.3")
for tag, is_link in (
(0xA000001D, True), # WSL symlink
(0xA0000003, True), # mount point (junction)
(0x9000001A, False), # OneDrive cloud placeholder: a regular file
(0, False),
):
with self.subTest(tag=hex(tag)):
st = mock.Mock(st_mode=regular, st_reparse_tag=tag)
self.assertIs(logging_config._is_link(path, st), is_link)
def test_truncate_with_file_logging_off_never_creates_the_active_log(self):
self.manager.update_config({"file_logging": False})
os.unlink(self._path("prompt_manager.log"))
self._write("prompt_manager.log.1")
result = self.manager.truncate_logs()
self.assertEqual(result, {"truncated": ["prompt_manager.log.1"], "errors": []})
self.assertEqual(self._names(), [])
def test_truncate_never_empties_a_linked_active_log(self):
self.manager.update_config({"file_logging": False})
os.unlink(self._path("prompt_manager.log"))
target = self._require_symlink("prompt_manager.log", data=b"keep")
self.manager.truncate_logs()
with open(target, "rb") as handle:
self.assertEqual(handle.read(), b"keep")
def test_log_files_mark_kind_and_lstat_links(self):
self._write("prompt_manager.log.1", b"r" * 10)
self._write("prompt_manager.log.bak", b"b")
target = self._require_symlink("prompt_manager.log.4")
rows = {row["filename"]: row for row in self.manager.get_log_files()}
self.assertEqual(rows["prompt_manager.log"]["kind"], "active")
self.assertTrue(rows["prompt_manager.log"]["is_main"])
self.assertEqual(rows["prompt_manager.log.1"]["kind"], "rotated")
self.assertEqual(rows["prompt_manager.log.bak"]["kind"], "other")
link = rows["prompt_manager.log.4"]
self.assertEqual(link["kind"], "link")
self.assertEqual(
link["size"], os.lstat(self._path("prompt_manager.log.4")).st_size
)
self.assertNotEqual(link["size"], os.path.getsize(target))
for row in rows.values():
self.assertEqual(
set(row), {"filename", "path", "size", "modified", "is_main", "kind"}
)
def test_an_unreadable_timestamp_does_not_break_the_scan(self):
self._write("prompt_manager.log.1")
real = logging_config.datetime
class Clock:
@staticmethod
def fromtimestamp(value):
raise OSError(22, "Invalid argument")
with mock.patch.object(logging_config, "datetime", Clock):
rows = {r["name"]: r for r in self.manager.scan_log_folder()}
self.assertIs(logging_config.datetime, real)
self.assertEqual(rows["prompt_manager.log.1"]["kind"], "rotated")
self.assertEqual(rows["prompt_manager.log.1"]["modified"], "")
def test_stats_exclude_links(self):
self._write("prompt_manager.log.1", b"r" * 10)
before = self.manager.get_log_stats()
self._require_symlink("prompt_manager.log.4")
after = self.manager.get_log_stats()
self.assertEqual(after["log_files_count"], before["log_files_count"])
self.assertEqual(after["total_log_size"], before["total_log_size"])
def test_log_file_entry_accepts_only_listed_regular_files(self):
self._write("prompt_manager.log.1")
self._write("prompt_manager.log.bak")
os.makedirs(self._path("folder.log"))
self._symlink("prompt_manager.log.4")
self.assertEqual(
os.path.normcase(str(self.manager.log_file_entry("prompt_manager.log.1"))),
os.path.normcase(self._path("prompt_manager.log.1")),
)
self.manager.log_file_entry("prompt_manager.log")
self.manager.log_file_entry("prompt_manager.log.bak")
for name in (
"prompt_manager.log.4",
"folder.log",
"missing.log",
"",
".",
"..",
"../outside/secret.txt",
"..\\outside\\secret.txt",
"prompt_manager.log:stream",
"prompt_manager.log\x00",
):
with self.subTest(name=name):
with self.assertRaises(FileNotFoundError):
self.manager.log_file_entry(name)
class TestConfiguration(IsolatedLoggerTestCase):
+6
View File
@@ -545,6 +545,12 @@
"functions": 98,
"statements": 98
},
"workspace/system-tagging.mjs": {
"lines": 80,
"branches": 70,
"functions": 75,
"statements": 80
},
"workspace/system-thumbnails.mjs": {
"lines": 98,
"branches": 95,
+137 -21
View File
@@ -70,6 +70,11 @@ MAX_AWAITING_HISTORY = 1000
AWAIT_HISTORY_SECONDS = 3600
# ``schedule`` snapshot meaning "history-only re-check" (not a new image)
HISTORY_ONLY = object()
# How long a stop waits for the image worker, and for the watchdog observer
# after its stop (#302). Every wait before the image in hand is interrupted
# at once, so a bound is reached only while one image is still being linked.
WORKER_JOIN_SECONDS = 5.0
OBSERVER_JOIN_SECONDS = 5.0
class ImageGenerationHandler(FileSystemEventHandler):
@@ -143,6 +148,11 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._history_due: "OrderedDict[str, None]" = OrderedDict()
self._awaiting_lock = threading.Lock()
self._sealed_subscribed = False
# Set by ``close``: no new work, no new worker, ever (#302).
self._closed = False
# Workers a stop signalled that had not ended yet: every later stop
# waits for them again, so it never reports True too early.
self._finishing: List[threading.Thread] = []
def on_created(self, event):
"""Handle filesystem creation events.
@@ -157,6 +167,10 @@ class ImageGenerationHandler(FileSystemEventHandler):
Args:
event: FileSystemEvent object containing event details
"""
# Closed: return before any log or tracker lock, so watchdog's
# ``observer.stop()`` never waits on this dispatch.
if self._closed:
return
if event.is_directory or not self.is_image_file(event.src_path):
return
self.logger.info(f"New image detected: {event.src_path}")
@@ -180,9 +194,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
and once its size has stopped changing.
Returns:
False when the path is already pending or the pending set is full
False when the path is already pending, the pending set is full or
the handler is closed
"""
with self._pending_lock:
if self._closed:
return False
if image_path in self._pending:
self.logger.debug(f"Already pending, not rescheduled: {image_path}")
return False
@@ -203,12 +220,19 @@ class ImageGenerationHandler(FileSystemEventHandler):
with self._pending_lock:
return len(self._pending)
def shutdown(self) -> None:
"""Stop the worker thread and drop every image still waiting.
def shutdown(self, timeout: float = WORKER_JOIN_SECONDS) -> bool:
"""Stop the worker thread, drop every image still waiting, and wait
for the worker within ``timeout`` seconds.
The worker wakes from any delay or settle wait, finishes nothing more
and exits; the next ``schedule`` starts a fresh worker with its own
queue. Safe to call repeatedly and before anything was scheduled.
The worker wakes from any delay or settle wait, finishes at most the
image in hand and exits; the next ``schedule`` starts a fresh worker
with its own queue (use ``close`` for a final stop). Safe to call
repeatedly and before anything was scheduled.
Returns:
True when no worker is left running. False when the worker is
still finishing its image after ``timeout`` (a warning is
logged), or when called on the worker itself (it only signals).
"""
with self._awaiting_lock:
self._awaiting.clear()
@@ -218,9 +242,40 @@ class ImageGenerationHandler(FileSystemEventHandler):
worker, self._worker = self._worker, None
work, self._work = self._work, queue.Queue()
stop, self._stop = self._stop, threading.Event()
if worker is not None:
self._finishing.append(worker)
finishing = [w for w in self._finishing if w.is_alive()]
self._finishing = finishing
stop.set()
if worker is not None and worker.is_alive():
work.put(None) # sentinel: wake a worker idle on an empty queue
if threading.current_thread() in finishing:
self.logger.debug("shutdown called on the image worker itself")
return False
deadline = time.monotonic() + timeout
for thread in finishing:
thread.join(max(0.0, deadline - time.monotonic()))
with self._pending_lock:
self._finishing = [w for w in self._finishing if w.is_alive()]
still_running = bool(self._finishing)
if still_running:
self.logger.warning(
f"image worker still running after {timeout} s; "
"it stops after the image in hand"
)
return False
return True
def close(self, timeout: float = WORKER_JOIN_SECONDS) -> bool:
"""Final stop: ``shutdown`` and refuse every later ``schedule``.
After ``close`` no worker starts again, not even from a sealed run's
history re-check on the history writer's thread. Returns what
``shutdown`` returns.
"""
with self._pending_lock:
self._closed = True
return self.shutdown(timeout)
def _ensure_worker(self) -> None:
"""Start the single daemon worker thread if it isn't running.
@@ -521,7 +576,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._awaiting.popitem(last=False)
def _await_history(self, image_path: str) -> None:
"""Re-check ``image_path`` once, history only, when its run is sealed."""
"""Re-check ``image_path`` once, history only, when its run is sealed.
A closed handler re-checks nothing.
"""
if self._closed:
return
if not self._sealed_subscribed:
subscribe_outputs_sealed(self._on_outputs_sealed)
self._sealed_subscribed = True
@@ -570,10 +630,19 @@ class ImageGenerationHandler(FileSystemEventHandler):
self._schedule_recheck(image_path)
def _schedule_recheck(self, image_path: str) -> None:
"""Schedule the history-only re-check, or keep it due for later."""
"""Schedule the history-only re-check, or keep it due for later.
A closed handler drops it.
"""
if self._closed:
return
if self.schedule(image_path, HISTORY_ONLY):
return
with self._awaiting_lock:
# ``close`` sets the flag before ``shutdown`` clears this set, so
# a close that raced the refusal above is seen here.
if self._closed:
return
self._history_due.pop(image_path, None)
self._history_due[image_path] = None
while len(self._history_due) > MAX_AWAITING_HISTORY:
@@ -872,6 +941,9 @@ class ImageMonitor:
self.observer = None
self.handler = None
self.monitored_directories = []
# (observer, handler) pairs a stop left running: the next stop waits
# for them again instead of reporting True.
self._unfinished: List[Tuple[Any, Any]] = []
self.logger = get_logger("prompt_manager.image_monitor")
self.review_recorder: Optional[Callable[[str], str]] = None
self.thumbnail_queue: Optional[Callable[[str], str]] = None
@@ -999,29 +1071,73 @@ class ImageMonitor:
"""Whether a filesystem observer is running."""
return self.observer is not None
def stop_monitoring(self):
"""Stop the image monitoring system.
def stop_monitoring(self, timeout: float = OBSERVER_JOIN_SECONDS) -> bool:
"""Stop the image monitoring system within one ``timeout`` deadline.
Cleanly shuts down the filesystem watcher and clears all monitoring state.
This method should be called before program exit to ensure proper cleanup.
Safe to call when monitoring never started.
Closes the handler first (its worker is joined and its dispatch
returns at once), then stops the watchdog observer and joins it. The
worker and the observer share the one deadline: the observer's join
gets what the worker left, possibly 0 s.
watchdog's ``observer.stop()`` joins its emitters without a limit;
it counts against the deadline but is not bounded here. A part an
earlier call left running is waited for again. Safe to call when
monitoring never started and more than once.
Returns:
True when every worker and observer it stopped has ended.
"""
deadline = time.monotonic() + timeout
observer, self.observer = self.observer, None
handler, self.handler = self.handler, None
self.monitored_directories = []
pairs, self._unfinished = self._unfinished, []
if observer is not None:
pairs.append((observer, handler))
if not pairs:
return True
def remaining() -> float:
return max(0.0, deadline - time.monotonic())
for pair_observer, pair_handler in pairs:
handler_ok = self._close_handler(pair_handler, remaining())
observer_ok = self._stop_observer(pair_observer, remaining())
if not (handler_ok and observer_ok):
self._unfinished.append(
(
None if observer_ok else pair_observer,
None if handler_ok else pair_handler,
)
)
self.logger.debug("Image monitoring stopped")
return not self._unfinished
def _close_handler(self, handler, timeout: float) -> bool:
if handler is None:
return True
try:
return handler.close(timeout)
except Exception as e:
self.logger.warning(f"Error while stopping the image worker: {e}")
return False
def _stop_observer(self, observer, timeout: float) -> bool:
"""``observer.stop()`` (safe to repeat), then a join of ``timeout``
seconds; a join that times out is logged with the time it got."""
if observer is None:
return
return True
try:
observer.stop()
observer.join()
observer.join(timeout)
except Exception as e:
self.logger.warning(f"Error while stopping image monitoring: {e}")
if handler is not None:
try:
handler.shutdown()
except Exception as e:
self.logger.warning(f"Error while stopping the image worker: {e}")
self.logger.debug("Image monitoring stopped")
return False
if observer.is_alive():
self.logger.warning(
f"watchdog observer still running after {timeout:.1f} s"
)
return False
return True
def detect_comfyui_output_dirs(self) -> list:
"""Auto-detect ComfyUI output directories.
+304 -45
View File
@@ -10,9 +10,12 @@ This module provides centralized logging configuration with support for:
"""
import collections
import fnmatch
import logging
import logging.handlers
import os
import re
import stat
import threading
from datetime import datetime
from typing import Dict, Any, Optional, List
@@ -20,6 +23,160 @@ from pathlib import Path
from . import library_location
ACTIVE_LOG_NAME = "prompt_manager.log"
# Rule R (#329): a rotated copy is prompt_manager.log.<n>, n >= 1, no leading 0.
_ROTATED_NAME = re.compile(r"prompt_manager\.log\.([1-9][0-9]*)")
# The legacy listing's pattern (GET /logs/files and /logs/stats).
_LISTED_PATTERN = "*.log*"
LOG_FILE_KINDS = ("active", "rotated", "link", "other")
class RotatedChanged(Exception):
"""A reviewed rotated copy changed since the review; nothing was deleted."""
def __init__(self, name):
super().__init__(f"{name} changed since the review")
self.name = name
def _name_key(name: str) -> str:
"""How a name compares: case-folded where the filesystem folds it (Windows)."""
return os.path.normcase(name)
def _rotated_number(name: str) -> Optional[int]:
match = _ROTATED_NAME.fullmatch(_name_key(name))
return int(match.group(1)) if match else None
# Windows reparse tags with the name-surrogate bit name another file or
# folder (symlinks, junctions, WSL symlinks); cloud placeholders do not.
_NAME_SURROGATE = 0x20000000
def _is_link(path: str, st: os.stat_result) -> bool:
if getattr(st, "st_reparse_tag", 0) & _NAME_SURROGATE:
return True
return stat.S_ISLNK(st.st_mode) or os.path.isjunction(path)
class LogFilePath(type(Path())):
"""A listed log file's path, with the inode it had when it was listed."""
listed_ino = 0
def _kind(name: str, path: str, st: os.stat_result) -> str:
"""Rule R: ``active``, ``rotated``, ``link`` or ``other`` (never deleted)."""
if _is_link(path, st):
return "link"
if not stat.S_ISREG(st.st_mode):
return "other"
if _name_key(name) == _name_key(ACTIVE_LOG_NAME):
return "active"
if _rotated_number(name) is not None:
return "rotated"
return "other"
def scan_log_folder(log_dir) -> List[Dict[str, Any]]:
"""Every direct entry of ``log_dir`` with its Rule R kind (blocking).
Rows are ``{name, kind, size, mtime_ns, ino, modified}``, read with
``lstat`` so a link is described, never followed. An entry that cannot
be read is left out. Never logs: the cleanup calls it while it holds the
file handler's lock.
"""
rows = []
try:
names = [entry.name for entry in os.scandir(log_dir)]
except OSError:
return rows
for name in names:
path = os.path.join(log_dir, name)
try:
st = os.lstat(path)
kind = _kind(name, path, st)
except OSError:
continue
try:
modified = datetime.fromtimestamp(st.st_mtime).isoformat()
except (OSError, OverflowError, ValueError): # a time before 1970 on Windows
modified = ""
rows.append(
{
"name": name,
"kind": kind,
"size": st.st_size,
"mtime_ns": st.st_mtime_ns,
"ino": st.st_ino,
"modified": modified,
"_dir": stat.S_ISDIR(st.st_mode) and kind != "link",
"_regular": stat.S_ISREG(st.st_mode) and kind != "link",
}
)
return rows
def public_row(row: Dict[str, Any]) -> Dict[str, Any]:
"""A scan row without its private keys."""
return {key: value for key, value in row.items() if not key.startswith("_")}
def rotated_sort_key(row: Dict[str, Any]) -> int:
return _rotated_number(row["name"]) or 0
def log_file_entry(log_dir, name) -> Path:
"""The path of a downloadable log file named ``name`` (blocking).
Only a regular, non-link file listed directly in ``log_dir`` under that
exact name; a separator, ``..``, a stream suffix or any unlisted name
raises FileNotFoundError, as a missing file does.
"""
if not isinstance(name, str) or not name or name in (".", ".."):
raise FileNotFoundError("Log file not found")
if any(ch in name for ch in ("/", "\\", "\x00", ":")):
raise FileNotFoundError("Log file not found")
for row in scan_log_folder(log_dir):
if row["name"] == name and row["kind"] in ("active", "rotated", "other"):
if not row["_regular"]: # a folder, FIFO, socket or device
break
entry = LogFilePath(log_dir, row["name"])
entry.listed_ino = row["ino"]
return entry
raise FileNotFoundError("Log file not found")
def _stream_is(stream, ino: int) -> bool:
"""Whether ``stream`` is open on the file with inode ``ino`` (0: unknown)."""
try:
open_ino = os.fstat(stream.fileno()).st_ino
except (OSError, ValueError, AttributeError):
return False
return not (ino and open_ino and open_ino != ino)
def _error_text(error: BaseException) -> str:
"""An OSError's reason and file name, without the folder it is in."""
if isinstance(error, OSError):
parts = [error.strerror or type(error).__name__]
if error.filename:
parts.append(os.path.basename(str(error.filename)))
return ": ".join(parts)
return str(error) or type(error).__name__
def _matches_review(row: Optional[Dict[str, Any]], reviewed: Dict[str, Any]) -> bool:
if row is None or row["kind"] != "rotated":
return False
if row["size"] != reviewed.get("size") or row["mtime_ns"] != reviewed.get(
"mtime_ns"
):
return False
ino = reviewed.get("ino") or 0
return not (ino and row["ino"] and row["ino"] != ino)
class _OwnedRotatingFileHandler(logging.handlers.RotatingFileHandler):
"""Keep a retired owned sink closed when an in-flight record resumes."""
@@ -87,6 +244,9 @@ class PromptManagerLogger:
# Memory buffer for recent logs (for web viewer)
self._log_buffer = collections.deque(maxlen=self.config["buffer_size"])
self._buffer_lock = threading.Lock()
# Held while a handler is replaced or log files are deleted (#329):
# taken before any file handler's own lock, never the other way round.
self._files_lock = threading.Lock()
# Initialize loggers
self._setup_loggers()
@@ -238,6 +398,14 @@ class PromptManagerLogger:
# Return most recent first
return list(reversed(logs[-limit:]))
def scan_log_folder(self) -> List[Dict[str, Any]]:
"""Every entry of the log folder with its Rule R kind (blocking)."""
return [public_row(row) for row in scan_log_folder(self.log_dir)]
def log_file_entry(self, name: str) -> Path:
"""A downloadable (regular, non-link) log file; FileNotFoundError else."""
return log_file_entry(self.log_dir, name)
def get_log_files(self) -> List[Dict[str, Any]]:
"""Get information about available log files.
@@ -245,26 +413,27 @@ class PromptManagerLogger:
List of dictionaries containing file information:
- filename: Name of the log file
- path: Full path to the log file
- size: File size in bytes
- size: File size in bytes (a link's own size: links are not followed)
- modified: ISO formatted modification timestamp
- is_main: True if this is the main log file
- kind: ``active``, ``rotated``, ``link`` or ``other`` (Rule R)
"""
log_files = []
for log_file in self.log_dir.glob("*.log*"):
try:
stat = log_file.stat()
log_files.append(
{
"filename": log_file.name,
"path": str(log_file),
"size": stat.st_size,
"modified": datetime.fromtimestamp(stat.st_mtime).isoformat(),
"is_main": log_file.name == "prompt_manager.log",
}
)
except Exception as e:
self.logger.error(f"Error getting info for log file {log_file}: {e}")
for row in scan_log_folder(self.log_dir):
if row["_dir"]:
continue
if not fnmatch.fnmatchcase(_name_key(row["name"]), _LISTED_PATTERN):
continue
log_files.append(
{
"filename": row["name"],
"path": str(self.log_dir / row["name"]),
"size": row["size"],
"modified": row["modified"],
"is_main": row["kind"] == "active",
"kind": row["kind"],
}
)
# Sort by modification time, newest first
log_files.sort(key=lambda x: x["modified"], reverse=True)
@@ -302,14 +471,44 @@ class PromptManagerLogger:
self.logger.error(f"Error reading log file {filename}: {e}")
raise
def _truncate_main_log(self, log_file: Path) -> None:
def _file_handlers(self) -> List[logging.Handler]:
return [h for h in self.logger.handlers if isinstance(h, logging.FileHandler)]
def _acquire_file_handlers(self) -> List[logging.Handler]:
"""Take every file handler's lock: no rollover renames a file meanwhile.
The caller holds ``_files_lock`` (so no handler is replaced) and must
not log until it released them: the lock is re-entrant, so a record
logged here could roll the log over under the caller's feet.
"""
held = []
try:
for handler in self._file_handlers():
handler.acquire()
held.append(handler)
except BaseException:
self._release_handlers(held)
raise
return held
@staticmethod
def _release_handlers(held) -> None:
for handler in reversed(held):
handler.release()
def _truncate_main_log(self, log_file: Path, expected_ino: int = 0) -> None:
"""Empty the main log through the handler that writes to it.
Opening the file a second time with ``"w"`` leaves the handler's own
stream at its old offset: later records land past the new end
(zero-padded on POSIX) and ``tell()`` keeps reporting the old size, so
the next record rolls the file over. Truncating the handler's stream
keeps position and size in step.
keeps position and size in step. Without a handler on it (file logging
off) the existing file is emptied in place, never through a link.
A handler whose open stream is not the listed file (``expected_ino``:
the user renamed the log, which the handler still writes to) is left
alone: the file under the active name is emptied instead.
"""
handled = False
target = os.path.normcase(os.path.realpath(str(log_file)))
@@ -320,6 +519,8 @@ class PromptManagerLogger:
continue
if os.path.normcase(os.path.realpath(base)) != target:
continue
if not _stream_is(stream, expected_ino):
continue
handler.acquire()
try:
stream.flush()
@@ -329,14 +530,65 @@ class PromptManagerLogger:
finally:
handler.release()
if not handled:
with open(log_file, "w", encoding="utf-8", errors="replace") as f:
f.write("")
flags = os.O_WRONLY | os.O_TRUNC | getattr(os, "O_NOFOLLOW", 0)
os.close(os.open(str(log_file), flags))
def delete_rotated(self, expected: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Delete exactly the reviewed rotated copies, or none of them (blocking).
``expected`` are scan rows (``name``, ``size``, ``mtime_ns``,
``ino``). Under ``_files_lock`` and every file handler's lock the
folder is scanned again; when any expected file is missing, is no
longer a rotated copy (Rule R) or differs in size, ``mtime_ns`` or a
non-zero ``ino``, RotatedChanged is raised with nothing deleted.
Otherwise each is unlinked in ``<n>`` order.
Returns ``{"removed": [names], "failed": [(name, OSError)]}``. One
summary line is logged once both locks are released.
"""
removed, failed = [], []
changed = None
with self._files_lock:
held = self._acquire_file_handlers()
try:
current = {
_name_key(row["name"]): row for row in scan_log_folder(self.log_dir)
}
targets = []
for reviewed in sorted(expected, key=rotated_sort_key):
row = current.get(_name_key(str(reviewed.get("name", ""))))
if not _matches_review(row, reviewed):
changed = str(reviewed.get("name", ""))
break
targets.append(row["name"])
if changed is None:
for name in targets:
try:
os.unlink(os.path.join(self.log_dir, name))
removed.append(name)
except OSError as error:
failed.append((name, error))
finally:
self._release_handlers(held)
if changed is not None:
self.logger.warning(
"Rotated log cleanup refused: %s changed since the review", changed
)
raise RotatedChanged(changed)
self.logger.info(
"Deleted %d rotated log copies%s",
len(removed),
f", {len(failed)} could not be deleted" if failed else "",
)
return {"removed": removed, "failed": failed}
def truncate_logs(self) -> Dict[str, Any]:
"""Truncate all log files.
"""Empty the active log and delete its rotated copies (Rule R).
Clears the main log file and deletes rotated log files.
Also clears the memory buffer.
Every other file in the log folder, and every link, is kept. Also
clears the memory buffer. Runs under ``_files_lock`` and the file
handlers' locks, so no rollover recreates a just-deleted copy, and
logs only once they are released.
Returns:
Dictionary with:
@@ -345,28 +597,31 @@ class PromptManagerLogger:
"""
results = {"truncated": [], "errors": []}
# Collect first and stay silent until every file has been handled:
# a log line written mid-way can roll the just-emptied main log over
# and recreate the backup that was deleted a moment earlier.
log_files = sorted(self.log_dir.glob("prompt_manager.log*"))
for log_file in log_files:
with self._files_lock:
held = self._acquire_file_handlers()
try:
if log_file.name == "prompt_manager.log":
self._truncate_main_log(log_file)
else:
# For rotated files, delete them
log_file.unlink()
results["truncated"].append(log_file.name)
except Exception as e:
results["errors"].append(
f"Failed to truncate {log_file.name}: {str(e)}"
rows = scan_log_folder(self.log_dir)
active = [row for row in rows if row["kind"] == "active"]
rotated = sorted(
(row for row in rows if row["kind"] == "rotated"),
key=rotated_sort_key,
)
# Clear memory buffer
with self._buffer_lock:
self._log_buffer.clear()
for row in active + rotated:
path = os.path.join(self.log_dir, row["name"])
try:
if row["kind"] == "active":
self._truncate_main_log(Path(path), row["ino"])
else:
os.unlink(path)
results["truncated"].append(row["name"])
except OSError as error:
results["errors"].append(
f"Failed to truncate {row['name']}: {_error_text(error)}"
)
with self._buffer_lock:
self._log_buffer.clear()
finally:
self._release_handlers(held)
if results["truncated"]:
self.logger.info("Truncated log files: %s", ", ".join(results["truncated"]))
@@ -409,7 +664,9 @@ class PromptManagerLogger:
"backup_count",
]
):
self._setup_loggers()
# No cleanup may hold the old handler while it is replaced.
with self._files_lock:
self._setup_loggers()
self.logger.info(f"Updated logging configuration: {new_config}")
@@ -440,7 +697,9 @@ class PromptManagerLogger:
level = entry["level"]
level_counts[level] = level_counts.get(level, 0) + 1
log_files = self.get_log_files()
# A link's target is not a log file of this folder: neither counted
# nor sized.
log_files = [f for f in self.get_log_files() if f["kind"] != "link"]
total_log_size = sum(f["size"] for f in log_files)
return {
+39 -6
View File
@@ -44,6 +44,9 @@ except ImportError:
from utils.hashing import generate_prompt_hash
from utils.logging_config import get_logger
# How long ``PromptTracker.shutdown`` waits for the cleanup thread (#301).
TRACKER_JOIN_SECONDS = 2.0
class PromptTracker:
"""Thread-safe tracking of current prompt executions.
@@ -92,9 +95,12 @@ class PromptTracker:
self.cleanup_interval = 300 # 5 minutes
self.prompt_timeout = 600 # 10 minutes
# Start cleanup thread
# Start cleanup thread; ``shutdown`` sets ``_stop`` to end it.
self._stop = threading.Event()
self.cleanup_thread = threading.Thread(
target=self._cleanup_expired_prompts, daemon=True
target=self._cleanup_expired_prompts,
name="PromptManagerPromptCleanup",
daemon=True,
)
self.cleanup_thread.start()
self.logger.debug("PromptTracker initialization completed")
@@ -464,9 +470,10 @@ class PromptTracker:
Runs continuously as a daemon thread, periodically removing expired
prompt contexts from the global tracking dictionary. This prevents
memory leaks from accumulating old prompt data.
memory leaks from accumulating old prompt data. Every wait ends at
once when ``shutdown`` sets the stop event.
"""
while True:
while not self._stop.is_set():
try:
current_time = time.time()
expired_ids = []
@@ -497,11 +504,37 @@ class PromptTracker:
if pruned:
self.logger.debug(f"Pruned {pruned} expired queue entries")
time.sleep(self.cleanup_interval)
if self._stop.wait(self.cleanup_interval):
break
except Exception as e:
self.logger.error(f"Error in cleanup thread: {e}")
time.sleep(60) # Wait a minute before retrying
if self._stop.wait(60): # Wait a minute before retrying
break
def shutdown(self, timeout: float = TRACKER_JOIN_SECONDS) -> bool:
"""Stop the cleanup thread and wait for it within ``timeout`` seconds.
Returns True when the thread has ended (or never existed). The tracker
keeps working afterwards; only the background pruning of expired
entries stops. Safe to call more than once. Called on the cleanup
thread itself it only signals and returns False.
"""
stop = getattr(self, "_stop", None)
thread = getattr(self, "cleanup_thread", None)
if stop is None or thread is None:
return True # built without __init__: no thread was started
stop.set()
if thread is threading.current_thread():
self.logger.debug("shutdown called on the prompt cleanup thread")
return False
thread.join(timeout)
if thread.is_alive():
self.logger.warning(
f"prompt cleanup thread still running after {timeout} s"
)
return False
return True
def get_active_prompts(self) -> Dict[str, Dict[str, Any]]:
"""
+18
View File
@@ -0,0 +1,18 @@
"""The project version, read from ``pyproject.toml`` (#329)."""
import re
from pathlib import Path
def project_version(root):
"""The version in ``root``'s ``pyproject.toml``; ``"unknown"`` on failure."""
try:
pyproject_path = Path(root) / "pyproject.toml"
if pyproject_path.exists():
content = pyproject_path.read_text()
match = re.search(r'version\s*=\s*["\']([^"\']+)["\']', content)
if match:
return match.group(1)
except Exception: # nosec B110 - an unreadable file reads as unknown
pass
return "unknown"

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