From 8bf759d6f83498cdad624b02e345631ca8f47832 Mon Sep 17 00:00:00 2001 From: Vito Date: Sat, 24 Oct 2026 12:54:09 -0700 Subject: [PATCH] feat(system): keep the library in ComfyUI's user folder (#318) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(system): resolve every library path in one place utils/library_location.py decides where the library lives: the PROMPT_MANAGER_DATA_DIR folder (validated: absolute, never inside a folder ComfyUI or Prompt Manager serves), an explicit database.default_path as before, a node-folder library that was not moved, else ComfyUI's user/__promptmanager. The database, diagnostics, config.json, logs and the LoRA image cache all ask it, so a new install never writes into the node folder and an existing one keeps every path until the user moves it. "Keep it here" is persisted as library.keep_in_node_folder. * feat(system): move the library as a verified Activity job POST /prompt_manager/system/library/move starts "Move library", an exclusive, non-cancellable job (database/library_move.py). It checks the target volume's free space, copies the database with SQLite's backup API from one read snapshot, verifies the copy (integrity check, every workflow document, the same rows in every table), copies the safety and pre-migration backups and config.json (owner-only, checked by size and SHA-256), then publishes under the old file's write lock and switches the live process (every PromptModel, the review and history observers, the API's plans and health report). The old prompts.db and config.json are renamed *.moved-to- only after the switch. A failure removes everything it wrote and leaves the old library in use; a marker in the new folder makes a restart after any crash pick the right library. GET /prompt_manager/system/library and the Overview's "location" say where the library lives and why; POST .../keep persists "Keep it here". * feat(system): show where the library lives and offer the move D129. System > Overview's Library health gains a "Library location" row (the database path, why it is there, the "Keep it here" choice and any PROMPT_MANAGER_DATA_DIR problem) with Move library for a node-folder library and Details otherwise. While a node-folder library is not kept there, the workspace shows the concept's .notice above every page: "Keep it here" saves the choice on the server and stops the notice; "Move library" opens the shared review (From, To, Size, what moves, what happens to the old files) and starts the exclusive job in Activity. The notice and the Overview read the location again when the job ends. The Overview comparison maps the inserted row; a new comparison holds the notice to the concept's Sources notice (markup, tokens, alignment, both themes and widths) with a mutation check. * feat(system): close the browser server's log files before its root goes The logs now live in the library's data folder, which for the browser test server is /user/__promptmanager; Windows could not remove the temp root while the server still held prompt_manager.log open. The API test's node folder is spelled in lower case because public paths are case-folded on Windows. * feat(system): show the library location on the legacy admin page A line under the admin page's statistics names the database path and why it is there (ComfyUI's user folder, PROMPT_MANAGER_DATA_DIR, config.json, or the node folder with a pointer to Move library), read once from GET /prompt_manager/system/library. * feat(system): prove a fresh host keeps its library in user/__promptmanager The execution suite starts one pinned CPU ComfyUI with no configuration at all on an OS-chosen port: the library is created in user/__promptmanager (database and logs), nothing but byte code is added to the node folder, the location route reports the default folder, and ComfyUI's /userdata API serves none of it (traversal spellings, a listing, a system-user header). * feat(system): document where the library lives and how it moves README: where the library lives, the upgrade path (Move library / Keep it here), PROMPT_MANAGER_DATA_DIR, explicit paths unchanged, and the way back to an older version. TESTING.md: D129 and the library location contract (resolution order, move steps, crash windows, coded refusals, Windows handling). * wip(system): library move fix round in progress * fix(system): retry a fallback image link that races the library switch The image monitor's fallback and recent-prompt links went through PromptDatabase.link_image_to_prompt, which logs and returns 0 on any error, so survive_switch never saw a "database is locked" or closed connection during Move library and the link was lost in both files. The monitor now links through link_image_to_prompt_or_raise (one transaction per attempt); link_image_to_prompt keeps its log-and-return-0 contract for other callers. A link that returns 0 is no longer logged as "Successfully linked image 0". A switch that fails after repointing the models now closes the old file's detached connections instead of leaving them to GC: switch() adds them to a list the mover owns. The mocked retry test is replaced by a real-lock race asserting the image row lands in the new library only. * fix(system): say which source is the default and where a symlinked library lives Sources & scans: with no source configured, ComfyUI's output folder now reads "ComfyUI/output ยท watching ยท default" on its own row instead of a separate note line. GET /settings marks each effective source with in_comfyui (its public path is relative to the ComfyUI base) so only those get the "ComfyUI/" prefix; the scan still sends the server's public path. Library location: a library in the node folder reads through the node folder's unresolved place in custom_nodes, so a node folder symlinked to a checkout outside ComfyUI shows custom_nodes/promptmanager/prompts.db instead of prompts.db (System library row, Move review, admin Library line). Never absolute or "..", else the shared _public_path, which sources and jobs keep using unchanged. --- .gitignore | 3 + README.md | 31 +- TESTING.md | 140 +++- database/execution_history.py | 11 + database/library_move.py | 768 ++++++++++++++++++ database/models.py | 85 +- database/operations.py | 265 +++--- database/review_changes.py | 11 + py/api/__init__.py | 8 + py/api/library_location.py | 420 ++++++++++ py/api/sources.py | 26 +- py/api/system.py | 5 +- py/config.py | 25 +- py/lora_utils.py | 17 +- tests/browser/server.py | 14 + tests/browser/system-concept.spec.mjs | 108 ++- tests/browser/workflows.spec.mjs | 7 + .../browser/workspace-system-sources.spec.mjs | 5 +- tests/browser/workspace-system.spec.mjs | 91 +++ tests/browser/workspace_system_fixture.py | 51 ++ tests/comfyui/library-location.spec.mjs | 73 ++ tests/js/workspace-library-location.test.js | 159 ++++ tests/js/workspace-system-sources.test.js | 28 +- tests/test_api_library_location.py | 427 ++++++++++ tests/test_ingestion_resolver_profile.py | 2 +- tests/test_library_location.py | 369 +++++++++ tests/test_library_location_wiring.py | 169 ++++ tests/test_library_move.py | 497 ++++++++++++ tests/test_library_move_review.py | 352 ++++++++ tests/test_settings_sources.py | 21 +- tests/test_system_api.py | 1 + tests/vendor/comfyui-config.test.mjs | 4 +- tools/vendor/comfyui.config.mjs | 2 +- tools/vendor/frontend-coverage-baseline.json | 6 + tools/vendor/playwright.execution.config.mjs | 2 +- utils/diagnostics.py | 20 +- utils/image_monitor.py | 49 +- utils/library_location.py | 480 +++++++++++ utils/logging_config.py | 10 +- utils/usage_tracking.py | 37 +- web/admin.html | 2 + web/js/admin.js | 24 + web/js/workspace/index.mjs | 3 + web/js/workspace/library-location.mjs | 255 ++++++ web/js/workspace/system-overview.mjs | 10 +- web/js/workspace/system-sources.mjs | 17 +- web/js/workspace/workspace.css | 6 + web/workspace.html | 2 + 48 files changed, 4911 insertions(+), 207 deletions(-) create mode 100644 database/library_move.py create mode 100644 py/api/library_location.py create mode 100644 tests/comfyui/library-location.spec.mjs create mode 100644 tests/js/workspace-library-location.test.js create mode 100644 tests/test_api_library_location.py create mode 100644 tests/test_library_location.py create mode 100644 tests/test_library_location_wiring.py create mode 100644 tests/test_library_move.py create mode 100644 tests/test_library_move_review.py create mode 100644 utils/library_location.py create mode 100644 web/js/workspace/library-location.mjs diff --git a/.gitignore b/.gitignore index 7e6132c..5773169 100644 --- a/.gitignore +++ b/.gitignore @@ -211,3 +211,6 @@ coverage/ # Disposable upstream checkout for real-host CI .ci-comfyui/ + +# Old library files a Move library renamed (never deleted; #249 S6b) +*.moved-to-* diff --git a/README.md b/README.md index b60585e..ed28d92 100644 --- a/README.md +++ b/README.md @@ -751,9 +751,34 @@ ComfyUI_PromptManager/ ## Configuration +### Where your library lives + +New installs keep the library (`prompts.db`, `config.json` with the CivitAI key, the +database backups and the logs) in ComfyUI's system user folder +`ComfyUI/user/__promptmanager/`. An update, a reinstall or ComfyUI-Manager replacing +the node folder no longer touches it, and ComfyUI's `/userdata` web API does not +serve `__` system folders. Thumbnails stay in `output/thumbnails`. + +- **Upgrading from an earlier version:** your library stays where it is, in the node + folder, until you move it; nothing is moved on its own. The workspace shows a notice + above every page and System โ€บ Overview a "Library location" row. **Move library** + copies the database with SQLite's backup API, verifies the copy (integrity check, every + workflow document, the same rows in every table), copies `config.json` and the backups, + switches to the new folder without a restart and only then renames the old files + `prompts.db.moved-to-__promptmanager` and `config.json.moved-to-__promptmanager` (never + deleted; the old backups stay too). If anything fails, the library stays in the node + folder, unchanged. **Keep it here** stops the notice; the Overview row stays. Take a + backup first (System โ€บ Backup & export or the dashboard's Backup) as with any change. +- **Your own folder:** set `PROMPT_MANAGER_DATA_DIR` to an absolute folder outside + ComfyUI's output, input and temp folders and your image sources. If it is empty and a + node-folder library exists, that library stays in use until you move it there. +- **An explicit `database.default_path`** in `config.json` and `PROMPT_MANAGER_CONFIG_PATH` + keep working as before (described below); such a library is not offered a move. + ### Database Settings -You can customize the database path by creating a `config.json` file in the extension root: +You can customize the database path by creating a `config.json` file in the library's +folder (the node folder for a library that was not moved): ```json { @@ -970,6 +995,10 @@ MIT License - see LICENSE file for details. ## Changelog +### Unreleased (Library location) + +- **๐Ÿ“ Library outside the node folder**: new installs keep `prompts.db`, `config.json`, the database backups and the logs in `ComfyUI/user/__promptmanager/`, where an update or reinstall of the node cannot delete them and ComfyUI's `/userdata` API does not reach. An existing library stays in the node folder until you press **Move library** (a verified copy, then a live switch; the old files are renamed `*.moved-to-__promptmanager`, never deleted) or **Keep it here**. `PROMPT_MANAGER_DATA_DIR` names your own folder. After moving, restart ComfyUI once so the logs follow. Going back to an older version after a move: an older version does not look in `user/__promptmanager`, so stop ComfyUI and copy `user/__promptmanager/prompts.db` (with its `-wal` file, if one is left) and `config.json` into the node folder; the `*.moved-to-__promptmanager` files are the library as it was at the move, without later changes. + ### Unreleased (Workflows stored once) - **๐Ÿ—œ๏ธ Smaller database**: each distinct workflow and prompt graph is stored once, compressed, instead of again for every image. Images recorded by this version use it from the first start; the maintenance job "Store each workflow once" rewrites the existing records and then reclaims the space. A 20,000-image test library went from 1.4 GB to 60 MB. The job always takes a verified backup first: the first time, `prompts.db.pre-workflow-storage_` beside the database, kept until you delete it; on later runs, a normal pooled safety backup. No image or video file is opened or changed. diff --git a/TESTING.md b/TESTING.md index fcd0da6..0fad7d0 100644 --- a/TESTING.md +++ b/TESTING.md @@ -2115,10 +2115,11 @@ Sources & scans (`system-sources.mjs`, S4) is the concept's "sources" body on th server's sources. `GET /prompt_manager/settings` reports them (public paths only): `sources.configured` (each `gallery_root_paths` entry with `valid`, `code` and the server's `reason`), `sources.effective` (what is scanned and served, with its -`root_index` and `watch`: `watching` only when the image monitor's observer is alive, +`root_index`, `watch` (`watching` only when the image monitor's observer is alive, scheduled that folder and monitoring is on, `stopping` when it still watches it but monitoring was switched off, `manual` when monitoring is off or no monitor runs, else -`restart_pending`), `sources.auto_detected`, `monitor` and `settings.monitoring_enabled`. +`restart_pending`) and `in_comfyui` (its public path is relative to the ComfyUI base)), +`sources.auto_detected`, `monitor` and `settings.monitoring_enabled`. `POST /prompt_manager/settings/validate-source {path}` checks one folder and writes nothing. Every caller (validate-source, the settings save, config.json at load) refuses a source through the same checks (`sources.source_problem`): its form first, then @@ -2310,7 +2311,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โ€“D67, D79โ€“D81 renumbered D92โ€“D99 and D111โ€“D113; D121โ€“D128 assigned): +Recorded deviations (plan D60โ€“D67, D79โ€“D81 renumbered D92โ€“D99 and D111โ€“D113; D121โ€“D129 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 @@ -2385,9 +2386,12 @@ Recorded deviations (plan D60โ€“D67, D79โ€“D81 renumbered D92โ€“D99 and D111โ€“D saved yet" with Remove only until Save; a configured source the server cannot use reads " ยท " with Remove only. A line under the list says when the list is not the whole story: "No source is configured. After saving, ComfyUIโ€™s output - folder is used." (the last one removed), "No source is configured; ComfyUIโ€™s output - folder is used." and "No configured source is available; ComfyUIโ€™s output folder is - used." (auto-detected rows, listed without Remove). "restart pending" is the third + folder is used." (the last one removed) and "No configured source is available; + ComfyUIโ€™s output folder is used." (every configured source unusable). Auto-detected + rows are listed without Remove and say they are the default on the row itself: + "ComfyUI/output ยท watching ยท default" ("ComfyUI/" only when `in_comfyui`; their scan + still sends the server's public path), with no line under the list when nothing is + configured. "restart pending" is the third watch state. The operation "Refresh files" ("Read every source again to find files added or removed outside ComfyUI.", button "Refresh") closes the surface: it calls the legacy gallery's `/gallery/rescan` and toasts "Files refreshed. N files in M output @@ -2458,6 +2462,20 @@ Recorded deviations (plan D60โ€“D67, D79โ€“D81 renumbered D92โ€“D99 and D111โ€“D `Watch for new outputs saves monitoring_enabled; a restart applies it` and `a save refused for one unusable source names it and marks its row` (workspace-system-sources.spec) and the system-state and system-sources node tests. +- **D129** (library location, S6b): the concept shows neither where the library lives + nor a way to move it. System โ€บ Overview's "Library health" gains a "Library location" + `.operation` row after Backup status, in that row's own markup: " ยท ", + plus "You chose to keep it here." and any `PROMPT_MANAGER_DATA_DIR` problem, with + "Move library" for a node-folder library and "Details" (a `dl.meta-grid` dialog) + otherwise. While a node-folder library is offered (`location.offer`), the shell shows + the concept's `.notice` (the Sources notice's markup: `div > h2 + p`, then + `.row-actions` with "Keep it here" and the `ca-primary` "Move library") in + `#library-banner` above `main`, in main's own width and padding, with no bottom margin + (main's padding spaces it); `workspace.css` carries the `/* D129 */` rules. Move + library opens the shared review (From, To, Size, Moves with it, Old files). Pinned by + `system shell and overview match the concept` (the mapping inserts the row) and + `library notice matches the concept notice` with its mutation check + (system-concept.spec), and the D129 tests of workspace-system.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 @@ -2829,6 +2847,116 @@ Restoring a deduped backup on a version before this one is not supported: that version's restore accepts the file and then reads rewritten images without their workflow (README changelog). +### Library location contract (#249 System S6b) + +`utils/library_location.py` is the one resolver of every data path: the database +(`database/operations.py` `_resolve_db_path`, diagnostics), config.json +(`PromptManagerConfig.get_config_path`, so the `ConfigStore`), the backups (beside the +database), the logs (`utils/logging_config.py`) and the LoRA image cache +(`py/lora_utils.py`). Order: a valid `PROMPT_MANAGER_DATA_DIR` (absolute, never inside +ComfyUI's output, input or temp folder, a public user folder, an extra gallery root or +an image source listed in any config.json the resolver may read; `config_location` and +`resolve_library` read those sources the same way, so config.json and the database +never split; an invalid one is ignored and reported as a problem code), used when +it holds a library or no node-folder library exists; an explicit +`database.default_path` other than `prompts.db` (relative to the node folder, as +before) and `PROMPT_MANAGER_CONFIG_PATH`, honoured as before; ComfyUI's +`user/__promptmanager` (`folder_paths.get_system_user_directory("promptmanager")`, or +`get_user_directory()/__promptmanager`) once its marker `library-location.json` says +`moved`; a node-folder `prompts.db` that was not moved; else `user/__promptmanager` +(new installs; outside ComfyUI the node folder). An empty environment folder never +replaces a node-folder library (`data-dir-empty`). Nothing is moved or copied at +startup (`tests/test_library_location.py`, `tests/test_library_location_wiring.py`). + +`POST /prompt_manager/system/library/move` (`py/api/library_location.py`, +`database/library_move.py`) starts "Move library": kind `library_move`, admission +`exclusive`, not cancellable. Refusals answer before any job: 400 `not-movable`, +`same-folder`, `target-invalid`; 409 `target-has-library`, `target-conflict`, `busy`; +507 `insufficient-space` (the free space of the target's volume, for the database, its +backups and config.json plus 16 MiB). The job copies with SQLite's backup API from one +read transaction (counts, document digest and `data_version` of that snapshot) into a +`.pm-move-*` temp file in the target folder (`journal_mode=DELETE`, fsync, 0600), +verifies it (`verify_database_file`, every workflow document, the same rows in every +table, the same documents), copies every `prompts.db.backup_*` and +`prompts.db.pre-workflow-storage_*` (times kept, size and SHA-256 checked), then, under +the old file's `BEGIN IMMEDIATE` and the config lock, rechecks `data_version` (a write +meanwhile copies again, up to three times, then `library-busy`), copies config.json +(0600, checked), writes the marker `pending`, gives every file its final name with +`os.replace` inside the target folder (a different drive is a copy, never a rename), +writes the marker `moved` and switches the process (`switch_process`: the resolver's +override, `PromptModel.repoint_instances` closes every connection of every model of the +old file and points it at the new one, the review and history observers are re-keyed, +the API's plans and health report are dropped). Only then are the old `prompts.db` and +`config.json` renamed `*.moved-to-` after a `wal_checkpoint(TRUNCATE)`; a file +still open elsewhere is left in place and named (`old-in-use`, `old-locked`); the old +backups stay. A failure before the marker says `moved` removes everything it wrote and +leaves the old library in use, untouched (`TestFailures`); a crash leaves the old +library in use (no or a `pending` marker; the next move replaces its own leftovers) or +the new one (`moved`) (`TestCrashWindows`). `GET /prompt_manager/system/library` and the +Overview's `location` give `{path, source, movable, keep, offer, target, problem, +size_bytes}` with paths relative to ComfyUI; a path in the node folder reads through +the node folder's own unresolved place in custom_nodes (`custom_nodes/promptmanager/ +prompts.db` even when that folder is a symlink to a checkout outside ComfyUI; never +absolute, never `..`, else the shared public path, which sources and jobs keep using +unchanged); `POST .../keep {keep}` persists +`library.keep_in_node_folder` through the `ConfigStore`. Errors are coded and name no +path (`tests/test_api_library_location.py`). The real-host `library-location.spec.mjs` +(execution suite, its own fresh host on an OS-chosen port, no configuration) proves a new +install creates `user/__promptmanager/prompts.db` and its logs, adds nothing but byte +code to the node folder, and that `/userdata` serves none of it. + +The switch runs while the copy is write-locked too (a guard transaction on it), so a +switch that fails at any point (the resolver, part of the models, the observers) is +undone whole: every model points at the old file again, the copy and its marker are +removed, the old file's connections the partial switch detached are closed (the models +reopen it on next use), and the move fails "โ€ฆ Nothing was moved." +(`TestFailedSwitchIsUndone`). A failed switch can stall up to the 30 s busy timeout when +a writer is blocked on the copy's guard lock; this is not a deadlock, and the writer +then retries on the old file. What +follows a complete switch (dropping plans, the health report, the history read fence, +which is closed first) is best effort: logged, never a failed move. A thread that +opens a connection while the model is being repointed keeps it only if the model's +generation did not change meanwhile (`PromptModel.get_connection`). + +Writes racing the switch: a background write that was waiting on the old file's lock +when the switch closed its connection fails, and run counting (`count_runs`, +`count_pair_runs`) and the image monitor's link (`ImageGenerationHandler. +link_image_to_prompt`: the resolved-pair ingest and the fallback or recent-prompt link) +run it once more through `PromptDatabase.survive_library_switch` when the model's +generation changed during the attempt; the retry opens the new library. The monitor +links through `PromptDatabase.link_image_to_prompt_or_raise`, which raises database +errors (`link_image_to_prompt` keeps its log-and-return-0 contract for other callers); +each attempt is one transaction, so a failed attempt rolls back and a committed one +does not retry, and the image is linked once. `TestBackgroundWritesSurviveTheSwitch` +races a real UPDATE (run counting) and a real fallback image link against the lock and +asserts each lands in the new file only. Any other failure is not retried. Known limit: that +blocked writer, and the mover with it, waits out the writer's busy timeout (30 s) +before its connection closes: SQLite holds the connection's mutex during the busy +wait, so neither a close nor `Connection.interrupt()` ends it sooner (measured), and +another thread cannot shorten the timeout of a connection in use. Other writes are +unaffected: threads that ask for a connection after the switch get the new file. +Writes from the HTTP routes (a prompt saved in the dashboard at that moment) are not +retried and report their error. + +A restart after a crash between the switch and the rename finishes the rename: the +marker names the old `prompts.db` and `config.json` (`from`, `config_from`), and the API +renames them `*.moved-to-` at startup when they are still there under their +own names, never deleting them and leaving a file still in use alone +(`TestRestartFinishesTheRename`). A second ComfyUI process using the same node folder +(two installs sharing one custom_nodes folder) is not switched: it keeps writing to the +old file, renamed or not, and those writes are not in the moved library; a rename while +it has the file open is refused as `old-in-use`. The Overview flags a library that stays +in the node folder through an explicit `database.default_path` or a node-folder +config.json (`inside-node-folder`: a reinstall deletes it, CivitAI key included); it is +not offered a move. + +Windows: every connection to the old file is closed before its rename (the source +snapshot connection, every model's connections, the checkpoint connection); temp files +are closed before `os.replace`; renames stay inside one folder; folder identity compares +`(st_dev, st_ino)` and `normcase(realpath)`, so case variants and links of a served +folder are refused. The browser test server closes its own log handlers before its temp +root (which now holds the logs) is removed. + ## Workspace data fields for the concept pages (#249) These fields are additive. The legacy `/prompt_manager/images/output` listing (no diff --git a/database/execution_history.py b/database/execution_history.py index d9c6cb6..6a72ae6 100644 --- a/database/execution_history.py +++ b/database/execution_history.py @@ -196,3 +196,14 @@ class ExecutionHistoryOperations: .fetchone() ) return dict(row) if row is not None else None + + +def rekey_history_restore(old_path: str, new_path: str) -> None: + """Move the observers of ``old_path`` to ``new_path`` (Move library).""" + old, new = _database_identity(old_path), _database_identity(new_path) + if old == new: + return + with _restore_lock: + observers = _restore_observers.pop(old, None) + if observers: + _restore_observers.setdefault(new, {}).update(observers) diff --git a/database/library_move.py b/database/library_move.py new file mode 100644 index 0000000..896a4bb --- /dev/null +++ b/database/library_move.py @@ -0,0 +1,768 @@ +"""Move library (#249 S6b): copy the node-folder library, verify it, switch. + +``plan_move`` checks a ``utils.library_location`` location can move and +sizes the copy; ``move_library`` does it, in this order: + +1. **Disk check** on the target's volume (the copy, the backups and + config.json, plus a margin). The source is never renamed across folders, + so another drive is simply a copy. +2. **Copy** with SQLite's backup API from a private connection holding one + read transaction: the per-table counts, the workflow documents' digest + and ``PRAGMA data_version`` describe exactly the copied snapshot. The copy + is a temp file in the target folder, made self-contained + (``journal_mode=DELETE``), fsynced and owner-only. +3. **Verify**: ``verify_database_file`` (integrity check, schema, every + workflow document decompressed against its hash, #315), the same row + count in every table and the same documents. +4. **Copy files**: every pooled safety backup and pre-migration backup, + checked by size and SHA-256 after the copy (their times are kept, so the + pool prunes in the same order). +5. **Publish** under the old database's write lock (``BEGIN IMMEDIATE``) and + the config.json lock: if any other connection wrote since the snapshot + (``data_version``), steps 2โ€“4 run again (``attempts``, then + ``library-busy``). Else config.json is copied (owner-only, checked), a + ``pending`` marker is written, the temp files take their final names + (``os.replace`` inside the target folder), the marker turns ``moved`` + and ``switch`` points the process at the new library while the old file + is still locked, so no write can land in it. +6. **Retire**: the old file's WAL is checkpointed and the old ``prompts.db`` + and ``config.json`` are renamed ``*.moved-to-`` (never deleted; + the old backups stay where they are). A file still open elsewhere is left + in place and reported (``old-in-use``): the marker already makes a restart + use the new library. + +Any failure before the marker says ``moved`` removes everything the move +wrote in the target folder and leaves the old library in use, untouched. +A crash at any point leaves either the old library in use (the marker is +absent or ``pending``; the next move replaces its own leftovers) or the new +one (``moved``). ``MoveFailed`` carries a code and a message without paths. +""" + +import contextlib +import hashlib +import json +import os +import shutil +import sqlite3 +import time +import types +import uuid +from pathlib import Path +from dataclasses import dataclass, field +from typing import Callable, Dict, List, Optional, Tuple + +from .models import BUSY_TIMEOUT_SECONDS, PromptModel + +try: + from ..utils import library_location as ll + from ..utils.logging_config import get_logger +except ImportError: + from utils import library_location as ll + from utils.logging_config import get_logger + +logger = get_logger("prompt_manager.database.library_move") + +TEMP_PREFIX = ".pm-move-" +MOVED_INFIX = ".moved-to-" +DISK_MARGIN_BYTES = 16 * 2**20 +DEFAULT_ATTEMPTS = 3 +_COPY_CHUNK = 4 * 2**20 +# Backups that live beside the database and move with it. +BACKUP_INFIXES = (".backup_", ".pre-workflow-storage_") + +MESSAGES = { + "not-movable": ( + "This library is not in the node folder, so there is nothing to move." + ), + "same-folder": "The new folder is the node folder itself.", + "target-invalid": "The new folder cannot hold the library.", + "target-has-library": ( + "The new folder already holds a Prompt Manager library. Nothing was moved." + ), + "target-conflict": ( + "The new folder already holds files with the same names. Nothing was moved." + ), + "insufficient-space": ( + "There is not enough free space for the copy. Nothing was moved." + ), + "copy-failed": "The copy could not be written or did not match. Nothing was moved.", + "verify-failed": "The copy did not pass verification. Nothing was moved.", + "library-busy": ( + "Prompt Manager kept writing to the library while it was copied. " + "Nothing was moved; try again when no workflow is running." + ), + "switch-failed": ( + "Prompt Manager could not switch to the verified copy, so the copy was " + "removed. Nothing was moved." + ), +} + + +class MoveFailed(Exception): + """A move that did not happen (or did not finish), with a coded reason.""" + + def __init__(self, code: str, message: Optional[str] = None): + self.code = code + self.message = message or MESSAGES.get(code, "The library was not moved.") + super().__init__(self.message) + + +@dataclass(frozen=True) +class MovePlan: + location: "ll.LibraryLocation" + db_path: str + source_dir: str + target_dir: str + config_path: Optional[str] + backups: Tuple[str, ...] + db_bytes: int + files_bytes: int + + @property + def required_bytes(self) -> int: + return self.db_bytes + self.files_bytes + DISK_MARGIN_BYTES + + @property + def target_location(self) -> "ll.LibraryLocation": + """Where the library is once moved (an explicit config.json stays).""" + source = self.location.target_source or ll.SOURCE_DEFAULT + config = ( + self.location.config_path + if self.location.config_explicit + else os.path.join(self.target_dir, ll.CONFIG_NAME) + ) + return ll.LibraryLocation( + source, + os.path.join(self.target_dir, ll.DATABASE_NAME), + config, + self.target_dir, + self.target_dir, + None, + source, + self.location.config_explicit, + ) + + +@dataclass +class MoveResult: + location: "ll.LibraryLocation" + counts: Dict[str, int] + files: List[str] + renamed: List[str] = field(default_factory=list) + retire_problem: Optional[str] = None + + +def _size(path: str) -> int: + try: + return os.path.getsize(path) + except OSError: + return 0 + + +def _free_bytes(path: str) -> int: + return shutil.disk_usage(path).free + + +def _existing_ancestor(path: str) -> str: + current = os.path.abspath(path) + while not os.path.isdir(current): + parent = os.path.dirname(current) + if parent == current: + break + current = parent + return current + + +def backup_names(db_path: str) -> Tuple[str, ...]: + """The pooled and pre-migration backups beside ``db_path`` (names).""" + folder = os.path.dirname(os.path.abspath(db_path)) + base = os.path.basename(db_path) + prefixes = tuple(base + infix for infix in BACKUP_INFIXES) + try: + names = os.listdir(folder) + except OSError: + return () + return tuple( + sorted( + name + for name in names + if name.startswith(prefixes) and os.path.isfile(os.path.join(folder, name)) + ) + ) + + +def plan_move(location) -> MovePlan: + """Check ``location`` can move to its target; raises MoveFailed.""" + if not location.movable: + raise MoveFailed("not-movable") + target = os.path.abspath(location.target_dir) + source_dir = os.path.dirname(os.path.abspath(location.db_path)) + if ll.same_folder(target, source_dir) or os.path.normcase( + os.path.realpath(target) + ) == os.path.normcase(os.path.realpath(source_dir)): + raise MoveFailed("same-folder") + if os.path.exists(target) and not os.path.isdir(target): + raise MoveFailed("target-invalid") + backups = backup_names(location.db_path) + # config.json moves with a node-folder library; one named by + # PROMPT_MANAGER_CONFIG_PATH is the user's and stays where it is. + config = None if location.config_explicit else location.config_path + config = config if config and os.path.isfile(config) else None + if ll.read_marker(target).get("state") != ll.MARKER_PENDING: + if os.path.exists(os.path.join(target, ll.DATABASE_NAME)): + raise MoveFailed("target-has-library") + names = [ll.CONFIG_NAME] if config else [] + if any(os.path.exists(os.path.join(target, n)) for n in names + list(backups)): + raise MoveFailed("target-conflict") + db_bytes = _size(location.db_path) + _size(location.db_path + "-wal") + files_bytes = sum(_size(os.path.join(source_dir, n)) for n in backups) + files_bytes += _size(config) if config else 0 + return MovePlan( + location, + os.path.abspath(location.db_path), + source_dir, + target, + config, + backups, + db_bytes, + files_bytes, + ) + + +def space_problem(plan: MovePlan) -> Optional[Tuple[int, int]]: + """``(required, free)`` when the target's volume is short, else None.""" + free = _free_bytes(_existing_ancestor(plan.target_dir)) + if free < plan.required_bytes: + return plan.required_bytes, free + return None + + +# โ”€โ”€ copy and verify โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def _tables(conn) -> List[str]: + return [ + row[0] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' " + "AND name NOT LIKE 'sqlite_%' ORDER BY name" + ) + ] + + +def _describe(conn) -> Tuple[Dict[str, int], str]: + """Row counts per table and a digest of the workflow documents.""" + counts = {} + for table in _tables(conn): + # A name read from sqlite_master, quoted as an identifier. + quoted = '"' + table.replace('"', '""') + '"' + sql = f"SELECT COUNT(*) FROM {quoted}" # nosec B608 + counts[table] = conn.execute(sql).fetchone()[0] + digest = hashlib.sha256() + if "workflow_documents" in counts: + for row in conn.execute( + "SELECT id, sha256, byte_length FROM workflow_documents ORDER BY id" + ): + digest.update(repr((row[0], bytes(row[1]), row[2])).encode()) + return counts, digest.hexdigest() + + +def _fsync_file(path: str) -> None: + with open(path, "rb+") as handle: + os.fsync(handle.fileno()) + + +def _fsync_dir(path: str) -> None: + if os.name != "posix": + return + fd = os.open(path, os.O_RDONLY) + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _private(path: str) -> None: + if os.name == "posix": + os.chmod(path, 0o600) + + +def _copy_snapshot(source: sqlite3.Connection, temp: str): + """Back up ``source``'s current snapshot to ``temp``; return what it holds.""" + source.execute("BEGIN") + try: + version = source.execute("PRAGMA data_version").fetchone()[0] + counts, documents = _describe(source) + destination = sqlite3.connect(temp) + try: + source.backup(destination) + finally: + destination.close() + finally: + source.execute("COMMIT") + with contextlib.closing(sqlite3.connect(temp)) as conn: + conn.execute("PRAGMA journal_mode=DELETE").fetchone() + _fsync_file(temp) + _private(temp) + return version, counts, documents + + +def _verify_copy(temp: str, counts, documents) -> None: + ok, reason = PromptModel.verify_database_file(temp) + if not ok: + logger.error(f"Library move: the copy did not verify: {reason}") + detail = reason if reason.startswith("Workflow document") else None + raise MoveFailed( + "verify-failed", + f"{MESSAGES['verify-failed']} {detail}" if detail else None, + ) + uri = Path(os.path.abspath(temp)).as_uri() + "?mode=ro" + try: + with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn: + copied = _describe(conn) + except sqlite3.Error as error: + logger.error(f"Library move: the copy could not be read: {error}") + raise MoveFailed("verify-failed") from error + if copied != (counts, documents): + logger.error("Library move: the copy's rows or documents differ") + raise MoveFailed("verify-failed") + + +def _hash_copy(source: str, destination: str) -> Tuple[int, str]: + """Copy ``source`` to ``destination`` (fsynced, owner-only, times kept).""" + digest = hashlib.sha256() + size = 0 + with open(source, "rb") as reader, open(destination, "wb") as writer: + while True: + chunk = reader.read(_COPY_CHUNK) + if not chunk: + break + digest.update(chunk) + size += len(chunk) + writer.write(chunk) + writer.flush() + os.fsync(writer.fileno()) + stat = os.stat(source) + os.utime(destination, ns=(stat.st_atime_ns, stat.st_mtime_ns)) + _private(destination) + return size, digest.hexdigest() + + +def _file_digest(path: str) -> Tuple[int, str]: + digest = hashlib.sha256() + size = 0 + with open(path, "rb") as reader: + while True: + chunk = reader.read(_COPY_CHUNK) + if not chunk: + return size, digest.hexdigest() + digest.update(chunk) + size += len(chunk) + + +def _check_copies(copies) -> None: + for source, temp, expected in copies: + if _file_digest(temp) != expected or _file_digest(source) != expected: + logger.error( + f"Library move: {os.path.basename(source)} did not copy unchanged" + ) + raise MoveFailed("copy-failed") + + +def write_marker( + directory: str, state: str, source_db: str, source_config: Optional[str] = None +) -> str: + """Atomically write the move marker; return its path. + + ``from`` and ``config_from`` name the old files, so a restart after a + crash between the switch and the rename can finish the rename. + """ + path = os.path.join(directory, ll.MARKER_NAME) + temp = os.path.join(directory, f"{TEMP_PREFIX}{uuid.uuid4().hex}.marker") + payload = json.dumps( + { + "version": 1, + "state": state, + "from": source_db, + "config_from": source_config, + "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + }, + indent=2, + ).encode("utf-8") + fd = os.open(temp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + _private(temp) + os.replace(temp, path) + except BaseException: + _remove(temp) + raise + _fsync_dir(directory) + return path + + +def _remove(*paths) -> None: + for path in paths: + try: + if path and os.path.lexists(path): + os.remove(path) + except OSError: + logger.warning(f"Library move: could not remove {os.path.basename(path)}") + + +def sweep_move_temps(directory: str) -> List[str]: + """Remove the temp files an interrupted move left in ``directory``.""" + removed = [] + try: + names = os.listdir(directory) + except OSError: + return removed + for name in names: + if name.startswith(TEMP_PREFIX): + _remove(os.path.join(directory, name)) + removed.append(name) + return removed + + +# โ”€โ”€ the move โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + +def _noop(*_args, **_kwargs): + return None + + +def switch_process(old_db: str, location, detached: list) -> None: + """Point this process at ``location``: resolver, every model, observers. + + Adds the old file's connections, detached and still open, to + ``detached`` as soon as they are detached, so the mover closes them + (``PromptModel.close_detached``) whether or not the rest succeeds. + """ + from .execution_history import rekey_history_restore + from .review_changes import rekey_review_changes + + ll.switch_to(location) + detached.extend(PromptModel.repoint_instances(old_db, location.db_path)) + rekey_review_changes(old_db, location.db_path) + rekey_history_restore(old_db, location.db_path) + + +def unswitch_process(old_db: str, new_db: str) -> list: + """Undo ``switch_process`` (any part of it); returns detached connections.""" + from .execution_history import rekey_history_restore + from .review_changes import rekey_review_changes + + ll.clear_override() + detached = PromptModel.repoint_instances(new_db, old_db) + rekey_review_changes(new_db, old_db) + rekey_history_restore(new_db, old_db) + return detached + + +class _Move: + def __init__(self, plan, *, switch, config_lock, progress, checkpoint): + self.plan = plan + self.switch = switch + self.config_lock = config_lock + self.progress = progress + self.checkpoint = checkpoint + self.created_dir = False + self.temps: List[str] = [] + self.published: List[str] = [] + self.marker: Optional[str] = None + self.moved = False + + def temp(self, suffix: str) -> str: + path = os.path.join( + self.plan.target_dir, f"{TEMP_PREFIX}{uuid.uuid4().hex}{suffix}" + ) + self.temps.append(path) + return path + + def prepare_target(self) -> None: + target = self.plan.target_dir + self.created_dir = not os.path.isdir(target) + ll.ensure_private_dir(target) + sweep_move_temps(target) + if ll.read_marker(target).get("state") == ll.MARKER_PENDING: + # An earlier move stopped before it switched: its files are ours. + for name in (ll.DATABASE_NAME, ll.CONFIG_NAME) + self.plan.backups: + _remove(os.path.join(target, name)) + for suffix in ("-wal", "-shm", "-journal"): + _remove(os.path.join(target, ll.DATABASE_NAME + suffix)) + + def copy_round(self, source): + """Steps 2โ€“4 once: the verified temp copies and the snapshot version.""" + self.progress("Copying the database") + temp_db = self.temp(".db") + version, counts, documents = _copy_snapshot(source, temp_db) + self.checkpoint("copied") + self.progress("Verifying the copy") + _verify_copy(temp_db, counts, documents) + self.checkpoint("verified") + copies = [] + for name in self.plan.backups: + src = os.path.join(self.plan.source_dir, name) + temp = self.temp(".bak") + try: + expected = _hash_copy(src, temp) + except OSError as error: + logger.error(f"Library move: copying {name} failed: {error}") + raise MoveFailed("copy-failed") from error + copies.append((src, temp, expected, name)) + if copies: + self.progress(f"Copied {len(copies)} backup file(s)") + self.checkpoint("files") + _check_copies([(s, t, e) for s, t, e, _ in copies]) + return version, counts, temp_db, copies + + def discard_round(self, temp_db, copies) -> None: + _remove(temp_db, *[temp for _s, temp, _e, _n in copies]) + + def publish(self, temp_db, copies) -> List[str]: + """Under the locks: config.json, marker, final names, marker moved.""" + target = self.plan.target_dir + files = [] + if self.plan.config_path: + temp = self.temp(".config") + try: + expected = _hash_copy(self.plan.config_path, temp) + except OSError as error: + raise MoveFailed("copy-failed") from error + _check_copies([(self.plan.config_path, temp, expected)]) + copies = copies + [(self.plan.config_path, temp, expected, ll.CONFIG_NAME)] + self.marker = write_marker( + target, ll.MARKER_PENDING, self.plan.db_path, self.plan.config_path + ) + self.checkpoint("pending") + for _src, temp, _expected, name in copies: + final = os.path.join(target, name) + os.replace(temp, final) + self.published.append(final) + files.append(name) + final_db = os.path.join(target, ll.DATABASE_NAME) + os.replace(temp_db, final_db) + self.published.append(final_db) + files.append(ll.DATABASE_NAME) + _fsync_dir(target) + self.checkpoint("published") + write_marker(target, ll.MARKER_MOVED, self.plan.db_path, self.plan.config_path) + self.moved = True + self.checkpoint("marked") + return files + + def cleanup(self) -> None: + if self.moved: + return + _remove(*self.temps, *self.published, self.marker) + if self.created_dir: + try: + os.rmdir(self.plan.target_dir) + except OSError: + pass + + +def _unique(path: str) -> str: + candidate, number = path, 1 + while os.path.lexists(candidate): + candidate = f"{path}_{number}" + number += 1 + return candidate + + +def retire_old_files(plan: MovePlan) -> Tuple[List[str], Optional[str]]: + """Rename the old database and config.json; ``(renamed, problem code)``. + + Every connection to the old file must be closed. The WAL is checkpointed + first; when it cannot be emptied (another process still has the file + open) nothing is renamed (``old-in-use``). + """ + old = plan.db_path + folder = os.path.basename(os.path.normpath(plan.target_dir)) + try: + with contextlib.closing(sqlite3.connect(old, timeout=1)) as conn: + busy = conn.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0] + except sqlite3.Error as error: + logger.warning( + f"Library move: the old database could not be checkpointed: {error}" + ) + return [], "old-in-use" + if busy or _size(old + "-wal") > 0: + return [], "old-in-use" + renamed = [] + try: + for path in (old, plan.config_path): + if not path or not os.path.isfile(path): + continue + new_name = _unique(path + MOVED_INFIX + folder) + os.rename(path, new_name) + renamed.append(os.path.basename(new_name)) + for suffix in ("-wal", "-shm"): + if os.path.isfile(old + suffix) and _size(old + suffix) == 0: + os.remove(old + suffix) + except OSError as error: + logger.warning(f"Library move: the old files could not be renamed: {error}") + return renamed, "old-locked" + return renamed, None + + +def finish_interrupted_move(location) -> List[str]: + """Rename old files a crash left between the switch and the rename. + + The marker of the folder in use says ``moved`` and names the old files; + an old ``prompts.db`` (and config.json) still there under its own name is + renamed ``*.moved-to-`` now, never deleted. Returns the new + names (nothing when there is nothing to finish or the file is in use). + """ + if location.source not in (ll.SOURCE_DEFAULT, ll.SOURCE_ENVIRONMENT): + return [] + marker = ll.read_marker(location.data_dir) + old = marker.get("from") + if marker.get("state") != ll.MARKER_MOVED or not isinstance(old, str): + return [] + if not os.path.isfile(old) or os.path.basename(old) != ll.DATABASE_NAME: + return [] + if ll.same_folder(os.path.dirname(old), location.data_dir): + return [] + config = marker.get("config_from") + if not ( + isinstance(config, str) + and os.path.basename(config) == ll.CONFIG_NAME + and os.path.isfile(config) + and os.path.normcase(os.path.abspath(config)) + != os.path.normcase(os.path.abspath(location.config_path)) + ): + config = None + pending = types.SimpleNamespace( + db_path=old, config_path=config, target_dir=location.data_dir + ) + renamed, problem = retire_old_files(pending) + if problem: + logger.warning(f"Library move: the old files are still in use ({problem})") + elif renamed: + logger.info(f"Library move finished: renamed {', '.join(renamed)}") + return renamed + + +def move_library( + plan: MovePlan, + *, + switch: Callable = switch_process, + unswitch: Callable = unswitch_process, + after_switch: Callable = _noop, + config_lock=None, + progress: Callable[[str], None] = _noop, + checkpoint: Callable[[str], None] = _noop, + attempts: int = DEFAULT_ATTEMPTS, +) -> MoveResult: + """Move the library as planned (blocking; see the module docstring). + + ``switch(old_db, location, detached)`` points the process at the copy + and adds the old connections it detached to ``detached``; if it raises, + ``unswitch(old_db, new_db)`` undoes it while both files are still + locked, every detached connection is closed, the copy is removed and + the move fails (``switch-failed``). ``after_switch()`` runs + once the switch is final; it is best effort (logged, never fails the move). + """ + short = space_problem(plan) + if short is not None: + raise MoveFailed("insufficient-space") + lock = config_lock if config_lock is not None else contextlib.nullcontext() + move = _Move( + plan, + switch=switch, + config_lock=lock, + progress=progress, + checkpoint=checkpoint, + ) + move.unswitch = unswitch + location = plan.target_location + try: + move.prepare_target() + files, counts = _copy_and_switch(move, location, attempts) + except Exception: + move.cleanup() + raise + try: + after_switch() + except Exception: + logger.exception("Library move: a follow-up after the switch failed") + from .execution_history import notify_history_restore + + notify_history_restore(location.db_path) + progress("Renaming the old files") + renamed, problem = retire_old_files(plan) + checkpoint("retired") + return MoveResult(location, counts, files, renamed, problem) + + +def _switch_guarded(move: _Move, location) -> list: + """Switch while the copy is locked too; undo it all if any part fails. + + The copy's write lock keeps every repointed thread from writing to it + until the switch is whole, so undoing a partial switch loses nothing. + Returns the old file's detached connections (still open). + """ + plan = move.plan + guard = sqlite3.connect( + location.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None + ) + try: + guard.execute("BEGIN IMMEDIATE") + detached: list = [] + try: + move.switch(plan.db_path, location, detached) + except Exception as error: + logger.exception("Library move: switching failed; undoing it") + move.moved = False # the copy and its marker are removed + try: + PromptModel.close_detached( + move.unswitch(plan.db_path, location.db_path), logger + ) + except Exception: + logger.exception("Library move: undoing the switch failed") + # The old file's connections a partial switch detached: the + # models reopen it on next use, so close them now, not at GC. + PromptModel.close_detached(detached, logger) + raise MoveFailed("switch-failed") from error + move.checkpoint("switched") + return detached + finally: + if guard.in_transaction: + guard.execute("ROLLBACK") + guard.close() + + +def _copy_and_switch(move: _Move, location, attempts: int): + plan = move.plan + for attempt in range(1, attempts + 1): + source = sqlite3.connect( + plan.db_path, timeout=BUSY_TIMEOUT_SECONDS, isolation_level=None + ) + try: + source.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_SECONDS * 1000}") + version, counts, temp_db, copies = move.copy_round(source) + source.execute("BEGIN IMMEDIATE") + try: + if source.execute("PRAGMA data_version").fetchone()[0] != version: + move.progress("The library changed during the copy; copying again") + move.discard_round(temp_db, copies) + continue + move.checkpoint("locked") + with move.config_lock: + files = move.publish(temp_db, copies) + detached = _switch_guarded(move, location) + # Still holding the old file's lock: a write that was waiting + # on it fails when its connection closes (and runs again on + # the new library) instead of landing in the old file. + PromptModel.close_detached(detached, logger) + finally: + if source.in_transaction: + source.execute("ROLLBACK") + return files, counts + finally: + source.close() + raise MoveFailed("library-busy") diff --git a/database/models.py b/database/models.py index f36fdbc..0c0210e 100644 --- a/database/models.py +++ b/database/models.py @@ -609,19 +609,37 @@ class PromptModel: if conn is not None and self._local.generation == self._generation: return conn - # check_same_thread=False only so close_all can clean up other - # threads' connections; each connection is still used by one thread. - conn = sqlite3.connect( - self.db_path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS - ) - conn.row_factory = sqlite3.Row - configure_connection(conn) + while True: + # The path and generation are read together, and the new + # connection is kept only if no Move library repointed the model + # meanwhile: otherwise it would write to the old file. + with self._conn_lock: + path, generation = self.db_path, self._generation + # check_same_thread=False only so close_all can clean up other + # threads' connections; each connection is still used by one thread. + conn = sqlite3.connect( + path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS + ) + try: + conn.row_factory = sqlite3.Row + configure_connection(conn) + except BaseException: + conn.close() + raise + with self._conn_lock: + if generation == self._generation: + self._connections.add(conn) + break + conn.close() self._local.conn = conn - self._local.generation = self._generation - with self._conn_lock: - self._connections.add(conn) + self._local.generation = generation return conn + @property + def generation(self) -> int: + """Bumped whenever every connection is closed or the model repointed.""" + return self._generation + @classmethod def close_all_instances(cls) -> None: """Close every connection of every live model in this process.""" @@ -630,22 +648,59 @@ class PromptModel: for model in models: model.close_all() + @classmethod + def repoint_instances(cls, old_path: str, new_path: str) -> list: + """Point every model of ``old_path`` at ``new_path`` (Move library). + + Every thread opens the new file on its next ``get_connection``. The + old connections are detached, not closed, and returned: the caller + closes them once nothing can block on the new file any more. Closing + a connection another thread is using waits for its statement; a write + waiting on the mover's lock then fails (and, through + ``PromptDatabase.survive_library_switch``, runs again on the new file). + """ + old_key = os.path.normcase(os.path.realpath(old_path)) + with cls._instances_lock: + models = list(cls._instances) + with cls._init_lock: + cls._initialized_paths.add(os.path.normcase(os.path.realpath(new_path))) + detached = [] + for model in models: + if model._schema_key() == old_key: + detached.extend(model._detach_connections(new_path)) + return detached + + @staticmethod + def close_detached(connections, logger=None) -> None: + """Close connections ``repoint_instances`` detached.""" + for conn in connections: + try: + conn.close() + except sqlite3.Error as e: + if logger is not None: + logger.warning(f"Error closing database connection: {e}") + def close_all(self) -> None: """Close every thread's connection; they reopen lazily on next use. Needed before deleting the database file (Windows keeps open files locked) and by tests that ran queries through executor threads. """ + self.close_detached(self._detach_connections(), self.logger) + + def _detach_connections(self, new_path: Optional[str] = None) -> list: + """Forget every connection (optionally pointing the model elsewhere). + + Returns them for the caller to close; threads reopen on next use. + """ with self._conn_lock: + if new_path is not None: + self.db_path = new_path connections = list(self._connections) self._connections.clear() self._generation += 1 self._local.conn = None - for conn in connections: - try: - conn.close() - except sqlite3.Error as e: - self.logger.warning(f"Error closing database connection: {e}") + return connections def close(self) -> None: """Close the calling thread's connection, if it has one. diff --git a/database/operations.py b/database/operations.py index aaaf2bd..1c2d89f 100644 --- a/database/operations.py +++ b/database/operations.py @@ -46,6 +46,7 @@ try: from ..utils.logging_config import get_logger from ..utils.hashing import generate_prompt_hash from ..utils.json_strict import dumps_strict + from ..utils import library_location except ImportError: import sys @@ -54,6 +55,7 @@ except ImportError: from utils.logging_config import get_logger from utils.hashing import generate_prompt_hash from utils.json_strict import dumps_strict + from utils import library_location # Subquery to fetch tags from junction table, embedded in SELECT statements @@ -99,40 +101,56 @@ def escape_like(term: str) -> str: ) -def _resolve_db_path(db_path: Optional[str] = None) -> str: - """Resolve the database path from config, falling back to defaults. +def _configured_default_path(): + """config.json's ``database.default_path``, or None. - When no explicit path is given, reads PromptManagerConfig.DEFAULT_DB_PATH - (set via config.json). Relative paths are resolved against the extension - root directory so the database location is predictable regardless of the - process working directory (important for Docker deployments). + The package-qualified config wins over a standalone ``py.config`` (a + custom-node package uses its own configuration). + """ + try: + try: + from ..py.config import PromptManagerConfig + except ImportError: + from py.config import PromptManagerConfig + except Exception: + return None + return getattr(PromptManagerConfig, "DEFAULT_DB_PATH", None) + + +def resolve_library_location(): + """Where the library lives now (``utils.library_location``, #249 S6b). + + The image sources that rule out a ``PROMPT_MANAGER_DATA_DIR`` are read by + the resolver itself, exactly as ``config_location`` reads them. + """ + return library_location.resolve_library(_configured_default_path()) + + +def _resolve_db_path(db_path: Optional[str] = None) -> str: + """Resolve the database path through the one library resolver. + + With no explicit path, ``utils.library_location.resolve_library`` decides + (``PROMPT_MANAGER_DATA_DIR``, an explicit ``database.default_path``, a + node-folder library not moved yet, else ComfyUI's ``user/__promptmanager``). + An explicit relative path is resolved against the extension root, as + before, so the location does not depend on the working directory. Args: - db_path: Explicit database path, or None to use config/default. + db_path: Explicit database path, or None to use the resolver. Returns: - Absolute path to the SQLite database file. + Absolute path to the SQLite database file; its folder exists (a new + data folder is created owner-only). """ - extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - if db_path is None: - try: - try: - from ..py.config import PromptManagerConfig - except ImportError: - from py.config import PromptManagerConfig - - db_path = PromptManagerConfig.DEFAULT_DB_PATH - except Exception: - db_path = "prompts.db" - - if not os.path.isabs(db_path): + db_path = resolve_library_location().db_path + elif not os.path.isabs(db_path): + extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) db_path = os.path.normpath(os.path.join(extension_root, db_path)) - # Ensure parent directory exists for custom paths parent = os.path.dirname(db_path) if parent: - os.makedirs(parent, exist_ok=True) + library_location.ensure_private_dir(parent) return db_path @@ -365,6 +383,25 @@ class PromptDatabase( self.model = PromptModel(db_path) self.logger.debug("Database operations initialized successfully") + def survive_library_switch(self, write): + """Run ``write()``; once more if it failed while the library moved. + + A background write (run counting, the image monitor) that was waiting + on the old file when Move library switched fails when its connection + is closed. When the model's generation changed during the attempt, + the write runs again on a fresh connection, which is the new library; + any other failure is raised as it was. ``write`` must open its own + transaction (every PromptDatabase method does). + """ + generation = self.model.generation + try: + return write() + except Exception: + if self.model.generation == generation: + raise + self.logger.info("Running a write again on the moved library") + return write() + def close(self) -> None: """Close the calling thread's database connection (see PromptModel.close).""" self.model.close() @@ -1580,103 +1617,117 @@ class PromptDatabase( metadata: Optional metadata about the image Returns: - int: The ID of the created image record + int: The ID of the created image record, or 0 when nothing was + linked (invalid or unknown prompt, already linked, or an error, + which is logged) + """ + try: + return self.link_image_to_prompt_or_raise(prompt_id, image_path, metadata) + except Exception as e: + self.logger.error(f"Error linking image to prompt {prompt_id}: {e}") + return 0 + + def link_image_to_prompt_or_raise( + self, prompt_id: str, image_path: str, metadata: Optional[Dict[str, Any]] = None + ) -> int: + """``link_image_to_prompt`` that raises database errors. + + The image monitor links through this so a write that failed because + Move library switched reaches ``survive_library_switch`` and runs + again on the new library (#249 S6b). Each call is one transaction: + a failed attempt rolls back, and a committed one does not raise. + Returns 0 for an invalid or unknown prompt or an image already linked. """ # Validate that prompt_id is a valid integer and exists in prompts table - try: - # Convert prompt_id to integer if it's a string number - if isinstance(prompt_id, str) and prompt_id.isdigit(): - prompt_id_int = int(prompt_id) - elif isinstance(prompt_id, int): - prompt_id_int = prompt_id + # Convert prompt_id to integer if it's a string number + if isinstance(prompt_id, str) and prompt_id.isdigit(): + prompt_id_int = int(prompt_id) + elif isinstance(prompt_id, int): + prompt_id_int = prompt_id + else: + # Handle temporary IDs like "temp_123456" - skip linking + if isinstance(prompt_id, str) and prompt_id.startswith("temp_"): + self.logger.debug( + f"Skipping image linking for temporary prompt ID: {prompt_id}" + ) + return 0 else: - # Handle temporary IDs like "temp_123456" - skip linking - if isinstance(prompt_id, str) and prompt_id.startswith("temp_"): - self.logger.debug( - f"Skipping image linking for temporary prompt ID: {prompt_id}" - ) - return 0 - else: - self.logger.warning( - f"Invalid prompt_id format: {prompt_id}, skipping image linking" - ) - return 0 - - # Verify the prompt exists in the database - with self.model.get_connection() as conn: - cursor = conn.execute( - "SELECT id FROM prompts WHERE id = ?", (prompt_id_int,) + self.logger.warning( + f"Invalid prompt_id format: {prompt_id}, skipping image linking" ) - if not cursor.fetchone(): - self.logger.warning( - f"Prompt ID {prompt_id_int} not found in database, skipping " - "image linking" - ) - return 0 + return 0 - # Proceed with linking - if not isinstance(metadata, dict): - metadata = {} - filename = os.path.basename(image_path) - file_path = normalize_image_path(image_path) - file_size, width, height, image_format = _coerce_file_info(metadata) - # Each workflow and prompt graph is stored once (new form). - documents = stored_values( + # Verify the prompt exists in the database + with self.model.get_connection() as conn: + cursor = conn.execute( + "SELECT id FROM prompts WHERE id = ?", (prompt_id_int,) + ) + if not cursor.fetchone(): + self.logger.warning( + f"Prompt ID {prompt_id_int} not found in database, skipping " + "image linking" + ) + return 0 + + # Proceed with linking + if not isinstance(metadata, dict): + metadata = {} + filename = os.path.basename(image_path) + file_path = normalize_image_path(image_path) + file_size, width, height, image_format = _coerce_file_info(metadata) + # Each workflow and prompt graph is stored once (new form). + documents = stored_values( + conn, + { + "workflow_data": dumps_strict(metadata.get("workflow", {})), + "prompt_metadata": dumps_strict(metadata.get("prompt", {})), + }, + ) + + # INSERT OR IGNORE skips a file already linked to this prompt + # (same prompt_id + normalised full path) + cursor = conn.execute( + """ + INSERT OR IGNORE INTO generated_images + (prompt_id, image_path, filename, file_path, file_size, width, + height, format, workflow_data, prompt_metadata, parameters, + workflow_document_id, prompt_document_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + prompt_id_int, + image_path, + filename, + file_path, + file_size, + width, + height, + image_format, + documents["workflow_data"], + documents["prompt_metadata"], + dumps_strict(metadata.get("parameters", {})), + documents["workflow_document_id"], + documents["prompt_document_id"], + ), + ) + if cursor.rowcount == 0: + # Already linked: a document made only for this row goes. + release_documents( conn, - { - "workflow_data": dumps_strict(metadata.get("workflow", {})), - "prompt_metadata": dumps_strict(metadata.get("prompt", {})), - }, - ) - - # INSERT OR IGNORE skips a file already linked to this prompt - # (same prompt_id + normalised full path) - cursor = conn.execute( - """ - INSERT OR IGNORE INTO generated_images - (prompt_id, image_path, filename, file_path, file_size, width, - height, format, workflow_data, prompt_metadata, parameters, - workflow_document_id, prompt_document_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, ( - prompt_id_int, - image_path, - filename, - file_path, - file_size, - width, - height, - image_format, - documents["workflow_data"], - documents["prompt_metadata"], - dumps_strict(metadata.get("parameters", {})), documents["workflow_document_id"], documents["prompt_document_id"], ), ) - if cursor.rowcount == 0: - # Already linked: a document made only for this row goes. - release_documents( - conn, - ( - documents["workflow_document_id"], - documents["prompt_document_id"], - ), - ) - conn.commit() + conn.commit() - if cursor.rowcount == 0: - self.logger.debug( - f"Image {file_path} already linked to prompt {prompt_id_int}" - ) - return 0 + if cursor.rowcount == 0: + self.logger.debug( + f"Image {file_path} already linked to prompt {prompt_id_int}" + ) + return 0 - return cursor.lastrowid - - except Exception as e: - self.logger.error(f"Error linking image to prompt {prompt_id}: {e}") - return 0 + return cursor.lastrowid def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]: """ diff --git a/database/review_changes.py b/database/review_changes.py index 31c54b6..d46a3f3 100644 --- a/database/review_changes.py +++ b/database/review_changes.py @@ -62,3 +62,14 @@ def notify_review_change(db_path: str) -> None: callback() except Exception: _logger.warning("Needs review change observer failed", exc_info=True) + + +def rekey_review_changes(old_path: str, new_path: str) -> None: + """Move the observers of ``old_path`` to ``new_path`` (Move library).""" + old, new = _identity(old_path), _identity(new_path) + if old == new: + return + with _lock: + observers = _observers.pop(old, None) + if observers: + _observers.setdefault(new, {}).update(observers) diff --git a/py/api/__init__.py b/py/api/__init__.py index 9ae8e08..9ce5bfe 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -14,6 +14,7 @@ Split into domain-specific mixins for maintainability: - ReviewRecordingMixin: Needs review recording for the scan and the monitor - ReviewEventsMixin: the Needs review change stream an open page re-reads on - SystemRoutesMixin: System health check, Overview and usage statistics +- LibraryLocationRoutesMixin: where the library lives, Move library, Keep it here - MaintenanceRoutesMixin: reviewed maintenance plans (preview, apply) - WorkflowStorageRoutesMixin: store each workflow once, then Reclaim space - DuplicateFilesMixin: duplicate image files as find and removal jobs @@ -51,6 +52,7 @@ from .image_tag_jobs import ImageTagJobsMixin from .thumbnail_jobs import ThumbnailJobsMixin from .thumbnail_auto import AutoThumbnailMixin from .system import SystemRoutesMixin +from .library_location import LibraryLocationRoutesMixin from .maintenance import MaintenanceRoutesMixin from .workflow_storage import WorkflowStorageRoutesMixin from .duplicate_files import DuplicateFilesMixin @@ -295,6 +297,7 @@ class PromptManagerAPI( ThumbnailJobsMixin, AutoThumbnailMixin, SystemRoutesMixin, + LibraryLocationRoutesMixin, MaintenanceRoutesMixin, WorkflowStorageRoutesMixin, DuplicateFilesMixin, @@ -376,6 +379,10 @@ class PromptManagerAPI( except Exception as e: self.logger.warning(f"Could not sweep interrupted backup copies: {e}") + # A crash between Move library's switch and its rename left the old + # node-folder files under their own names (#249 S6b): rename them now. + self.finish_interrupted_library_move() + # Run cleanup on initialization to remove any existing duplicates try: removed = self.db.cleanup_duplicates() @@ -535,6 +542,7 @@ class PromptManagerAPI( self._register_settings_routes(routes) self._register_job_routes(routes) self._register_system_routes(routes) + self._register_library_location_routes(routes) self._register_maintenance_routes(routes) self._register_duplicate_routes(routes) self._register_logging_routes(routes) diff --git a/py/api/library_location.py b/py/api/library_location.py new file mode 100644 index 0000000..8c4ec32 --- /dev/null +++ b/py/api/library_location.py @@ -0,0 +1,420 @@ +"""Library location routes (#249 S6b): where the library lives, Move, Keep. + +- ``GET /prompt_manager/system/library`` โ†’ ``{success, location}``: the + database's path (relative to ComfyUI, as the statistics show it), why it is + there (``source``: ``default``, ``environment``, ``config`` or + ``node_folder``), whether it can move and where to (``target``), the + "Keep it here" choice (``keep``), whether to offer the move (``offer``: + movable and not kept), a ``problem`` code and the database's size. The + System Overview carries the same object as ``location``. +- ``POST /prompt_manager/system/library/move`` โ†’ 202 ``{job}``: the + ``library_move`` job, ``exclusive`` (it runs alone and every other job, + legacy-page operations included, is refused while it runs), not + cancellable. It copies, verifies and switches the live library + (``database/library_move.py``) and ends with one truthful line; a failure + leaves the library in the node folder, unchanged. Refused before any job + starts: 400 ``not-movable``/``same-folder``/``target-invalid``, 409 + ``target-has-library``/``target-conflict``/``busy``, 507 + ``insufficient-space``. +- ``POST /prompt_manager/system/library/keep {keep: bool}`` persists + ``library.keep_in_node_folder`` through the config store and answers the + new location. + +Errors are ``{success:false, code, error}`` and never name a path. +""" + +import asyncio +import functools +import os + +from aiohttp import web + +from .job_routes import busy_response +from .jobs import JobBusy + +try: + from ...database import library_move as lm + from ...database.operations import resolve_library_location + from ...utils import library_location as ll +except ImportError: + from database import library_move as lm + from database.operations import resolve_library_location + from utils import library_location as ll + +LIBRARY_MOVE_JOB_KIND = "library_move" +LIBRARY_MOVE_TITLE = "Move library" +LIBRARY_MOVE_GROUPS = ("files", "errors") +KEEP_BODY_MAX = 1024 + +_REFUSAL_STATUS = { + "not-movable": 400, + "same-folder": 400, + "target-invalid": 400, + "target-has-library": 409, + "target-conflict": 409, + "insufficient-space": 507, +} +STILL_THERE = "The library is still in the node folder, unchanged." +NOT_MOVED = "The library was not moved." +OLD_LEFT = { + "old-in-use": ( + "The old database is still open elsewhere, so it was left in the node " + "folder; it is no longer used." + ), + "old-locked": ( + "The old files could not be renamed, so they were left in the node " + "folder; they are no longer used." + ), +} + + +def _failure(message, status, code): + return web.json_response( + {"success": False, "code": code, "error": message}, status=status + ) + + +def _plural(count, one, many=None): + return f"{count:,} {one if count == 1 else (many or one + 's')}" + + +def _same_file(first, second): + try: + return os.path.samefile(first, second) + except OSError: + return os.path.normcase(os.path.realpath(first)) == os.path.normcase( + os.path.realpath(second) + ) + + +CUSTOM_NODES = "custom_nodes" + + +def _inside(path, folder): + """``path`` relative to ``folder`` by text (normalised spellings), or None.""" + child, parent = os.path.normcase(path), os.path.normcase(folder) + if child == parent: + return "" + prefix = parent.rstrip(os.sep) + os.sep + return path[len(prefix) :] if child.startswith(prefix) else None + + +def _comfyui_custom_nodes(): + """ComfyUI's custom_nodes folders as it lists them (unresolved), or [].""" + try: + from ..config import _import_folder_paths + + getter = getattr(_import_folder_paths(), "get_folder_paths", None) + folders = getter(CUSTOM_NODES) if getter else [] + return [folder for folder in folders or [] if isinstance(folder, str)] + except Exception: + return [] + + +def node_folder_public_path(path, extension_root, custom_nodes): + """``custom_nodes//...`` for a path in the node folder, or None. + + The node folder is matched at its own unresolved place in a custom_nodes + folder (one of ``custom_nodes``, or the parent of ``extension_root`` when + that is named custom_nodes), so a node folder that is a symlink to a + checkout outside ComfyUI still reads relative to ComfyUI. ``path`` may be + spelled through the link or resolved. Never absolute, never ``..``: + anything else is None and the caller uses ``_public_path``. + """ + if not path or not extension_root: + return None + try: + root = os.path.normpath(os.path.abspath(extension_root)) + name = os.path.basename(root) + folders = [os.path.normpath(os.path.abspath(f)) for f in custom_nodes] + parent = os.path.dirname(root) + if os.path.normcase(os.path.basename(parent)) == CUSTOM_NODES: + folders.append(parent) + spellings = [os.path.normpath(os.path.abspath(path))] + spellings.append(os.path.realpath(spellings[0])) + except (OSError, ValueError, TypeError): + return None + for folder in folders: + node = os.path.join(folder, name) + for spelling, node_spelling in zip(spellings, (node, os.path.realpath(node))): + rest = _inside(spelling, node_spelling) + if rest is None: + continue + parts = [os.path.basename(folder), name] + ( + rest.split(os.sep) if rest else [] + ) + if all(part and part not in (".", "..") for part in parts): + return "/".join(parts) + return None + + +class LibraryLocationRoutesMixin: + """Mixin for the /prompt_manager/system/library endpoints.""" + + # Replaced in tests with a resolver over a temp tree. + _library_resolver = staticmethod(resolve_library_location) + + def _register_library_location_routes(self, routes): + @routes.get("/prompt_manager/system/library") + async def library_location_route(request): + return await self.get_library_location(request) + + @routes.post("/prompt_manager/system/library/move") + async def library_move_route(request): + return await self.start_library_move(request) + + @routes.post("/prompt_manager/system/library/keep") + async def library_keep_route(request): + return await self.keep_library(request) + + # โ”€โ”€ location โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def _library_public_path(self, path): + """The library's path for display: through the node folder's place in + custom_nodes when it is there (a symlinked node folder included), else + ``_public_path``. Only the library location uses this; sources and + jobs keep ``_public_path``, which their requests round-trip.""" + shown = node_folder_public_path( + path, self._extension_root, _comfyui_custom_nodes() + ) + return shown or self._public_path(path) + + def _live_location(self): + """The resolved location, unless the live database is another file. + + A database the API was handed (tests, the browser fixture server) or + one the resolver no longer names is reported as configured and is + never moved. + """ + location = self._library_resolver() + live = self.db.model.db_path + if _same_file(location.db_path, live): + return location + return ll.LibraryLocation( + ll.SOURCE_CONFIG, + live, + location.config_path, + os.path.dirname(live), + config_explicit=location.config_explicit, + ) + + def _location_payload(self, location): + """The public view of ``location`` (blocking: stats the database).""" + from ..config import PromptManagerConfig + + keep = bool(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + size = 0 + try: + size = os.path.getsize(location.db_path) + except OSError: + pass + return { + "path": self._library_public_path(location.db_path), + # Where Move library would put it, and whether config.json goes too. + "target_source": location.target_source if location.movable else None, + "config_moves": location.movable and not location.config_explicit, + "source": location.source, + "movable": location.movable, + "keep": keep, + "offer": location.movable and not keep, + "target": ( + self._library_public_path(location.target_dir) + if location.movable + else None + ), + "problem": location.problem or self._node_folder_risk(location), + "size_bytes": size, + } + + # The node folder (tests point it at a temp tree). + _extension_root = ll.EXTENSION_ROOT + + def _node_folder_risk(self, location): + """``inside-node-folder`` for a library that stays in the node folder. + + A configured database path, or a config.json the resolver still reads + from the node folder, is deleted with the node by a reinstall, and + it is never offered a move. + """ + if location.source == ll.SOURCE_NODE_FOLDER: + return None + paths = [location.db_path] + ( + [] if location.config_explicit else [location.config_path] + ) + root = self._extension_root + for path in paths: + parent = os.path.dirname(os.path.abspath(path)) + if ll.same_folder(parent, root) or os.path.normcase( + os.path.realpath(parent) + ).startswith(os.path.normcase(os.path.realpath(root)) + os.sep): + return "inside-node-folder" + return None + + def _library_payload(self): + return self._location_payload(self._live_location()) + + async def get_library_location(self, request): + try: + payload = await self._run_in_executor(self._library_payload) + except Exception: + self.logger.exception("Library location could not be read") + return _failure( + "The library location could not be read.", 500, "location-failed" + ) + return web.json_response( + {"success": True, "location": payload}, + headers={"Cache-Control": "no-store"}, + ) + + # โ”€โ”€ move โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + def _plan_library_move(self): + """The checked plan (blocking); raises MoveFailed.""" + plan = lm.plan_move(self._live_location()) + if lm.space_problem(plan) is not None: + raise lm.MoveFailed("insufficient-space") + return plan + + async def start_library_move(self, request): + blocker = self._jobs.blocker(LIBRARY_MOVE_JOB_KIND, admission="exclusive") + if blocker is not None: + return busy_response(JobBusy(blocker)) + try: + await self._run_in_executor(self._plan_library_move) + except lm.MoveFailed as refused: + status = _REFUSAL_STATUS.get(refused.code, 400) + return _failure(refused.message, status, refused.code) + except Exception: + self.logger.exception("Move library could not be planned") + return _failure(NOT_MOVED, 500, "move-failed") + try: + job = self._jobs.start( + LIBRARY_MOVE_JOB_KIND, + LIBRARY_MOVE_TITLE, + admission="exclusive", + cancellable=False, + scope={}, + groups=LIBRARY_MOVE_GROUPS, + runner=self._run_library_move, + ) + except JobBusy as busy: + return busy_response(busy) + return web.json_response({"success": True, "job": job.summary()}, status=202) + + def _after_library_switch(self): + """Forget what described the old file (best effort, after the switch). + + Plans, the health report and the history read fence named the old + file; the fence is closed before it is dropped. + """ + self._database_replaced() + fence, self._history_read_fence = self._history_read_fence, None + if fence is not None: + fence.close() + + def finish_interrupted_library_move(self): + """Startup: rename old files a crash left after a move (blocking).""" + try: + location = self._library_resolver() + if _same_file(location.db_path, self.db.model.db_path): + return lm.finish_interrupted_move(location) + except Exception: + self.logger.exception("Could not finish an interrupted library move") + return [] + + async def _run_library_move(self, job): + from ..config_store import _COMMIT_LOCK + + loop = asyncio.get_running_loop() + + def progress(line): + def show(): + job.log_line(line) + job.update(status_line=line) + + loop.call_soon_threadsafe(show) + + job.update(status_line="Checking the new folder") + try: + plan = await self._run_in_executor(self._plan_library_move) + target = self._library_public_path(plan.target_dir) + job.log_line(f"Moving the library to {target}") + result = await self._run_in_executor( + functools.partial( + lm.move_library, + plan, + after_switch=self._after_library_switch, + config_lock=_COMMIT_LOCK, + progress=progress, + ) + ) + except lm.MoveFailed as failed: + self.logger.error(f"Move library failed: {failed.code}") + return self._fail_library_move( + job, failed.code, f"{failed.message} {STILL_THERE}" + ) + except Exception: + self.logger.exception("Move library failed") + return self._fail_library_move(job, None, f"{NOT_MOVED} {STILL_THERE}") + await asyncio.sleep(0) # the progress lines come first + job.record_counted("files", len(result.files)) + line = ( + f"Moved the library to {target}: " + f"{_plural(result.counts.get('prompts', 0), 'prompt')} and " + f"{_plural(result.counts.get('generated_images', 0), 'image')} verified" + ) + if result.retire_problem: + line += ". " + OLD_LEFT.get(result.retire_problem, OLD_LEFT["old-locked"]) + elif result.renamed: + line += ( + f". The old files were renamed {', '.join(result.renamed)}; " + "nothing was deleted" + ) + line += ". Logs move at the next restart." + job.log_line(line) + job.update(progress=100, status_line=line) + + @staticmethod + def _fail_library_move(job, code, message): + from .maintenance import _fail_job + + _fail_job(job, code, message) + + # โ”€โ”€ keep it here โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€ + + async def keep_library(self, request): + from ..config import PromptManagerConfig + from ..config_store import get_config_store + + try: + if (request.content_length or 0) > KEEP_BODY_MAX: + raise ValueError("too large") + data = await request.json() + except Exception: + data = None + if not isinstance(data, dict) or not isinstance(data.get("keep"), bool): + return _failure("keep must be true or false.", 400, "invalid-request") + keep = data["keep"] + + def mutate(config): + library = config.get("library") + library = dict(library) if isinstance(library, dict) else {} + library["keep_in_node_folder"] = keep + config["library"] = library + return config + + def publish(_config): + PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep + + try: + await self._run_in_executor( + functools.partial( + get_config_store().commit, + mutate, + expected_revision=None, + publish=publish, + ) + ) + except Exception as exc: + return self._config_commit_failure(exc) + return await self.get_library_location(request) diff --git a/py/api/sources.py b/py/api/sources.py index 90fb18c..694780d 100644 --- a/py/api/sources.py +++ b/py/api/sources.py @@ -13,7 +13,9 @@ - ``describe_sources`` (blocking: it resolves paths) reports the configured sources with their validity, the effective sources with their listing index and watch state, and the image monitor's state. Every path in it is - a public path (``_public_path``), never an absolute one. + a public path (``_public_path``), never an absolute one; ``in_comfyui`` + says whether an effective source's public path is relative to the ComfyUI + base (``in_comfyui_base``), so the page can read it as ``ComfyUI/``. A source's watch state is ``watching`` when the monitor's observer is alive, scheduled that directory and monitoring is on; ``stopping`` when it still @@ -209,6 +211,27 @@ def watch_state(path, *, present, alive, scheduled_keys, enabled): return "restart_pending" +def in_comfyui_base(path): + """Whether ``_public_path(path)`` is relative to the ComfyUI base directory. + + The base is the first anchor ``_public_path`` tries, so this agrees with + it: the same canonical (realpath + normcase) spelling on both sides. + """ + from pathlib import Path + + from ..config import _comfyui_base_directory + + base = _comfyui_base_directory() + if not path or not base: + return False + try: + canonical = Path(os.path.normcase(os.path.realpath(str(path)))) + anchor = Path(os.path.normcase(os.path.realpath(base))) + except (OSError, ValueError): + return False + return canonical.is_relative_to(anchor) + + def describe_sources(*, configured, output_dirs, auto_detected, monitor, enabled): """``(sources, monitor)`` for GET /settings (blocking; public paths only). @@ -233,6 +256,7 @@ def describe_sources(*, configured, output_dirs, auto_detected, monitor, enabled scheduled_keys=scheduled_keys, enabled=enabled, ), + "in_comfyui": in_comfyui_base(path), } for index, path in enumerate(output_dirs) ] diff --git a/py/api/system.py b/py/api/system.py index 092042e..45d67da 100644 --- a/py/api/system.py +++ b/py/api/system.py @@ -168,9 +168,10 @@ class SystemRoutesMixin: from ..config import GalleryConfig try: - (output_dirs, auto_detected), counts = await _gather_settled( + (output_dirs, auto_detected), counts, location = await _gather_settled( self._run_in_executor(self._output_dirs_with_origin), self._run_in_executor(self.db.library_counts), + self._run_in_executor(self._library_payload), ) except Exception: self.logger.exception("System overview failed") @@ -194,6 +195,8 @@ class SystemRoutesMixin: "effective": len(output_dirs), "auto_detected": auto_detected, }, + # Where the library lives and why (py/api/library_location.py). + "location": location, }, headers={"Cache-Control": "no-store"}, ) diff --git a/py/config.py b/py/config.py index 18f94f1..11bcdf2 100644 --- a/py/config.py +++ b/py/config.py @@ -36,12 +36,14 @@ from typing import Any, Dict, List, Optional, Tuple # noqa: E402 # Import logging system try: from ..utils.logging_config import get_logger + from ..utils import library_location except ImportError: import sys current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, current_dir) from utils.logging_config import get_logger + from utils import library_location # Initialize logger for config operations config_logger = get_logger("prompt_manager.config") @@ -663,6 +665,10 @@ class PromptManagerConfig: # database first (System โ€บ Backup & export โ€บ Recovery policy, #249 S6). SAFETY_BACKUP = True + # "Keep it here" (#249 S6b): the user chose to keep a node-folder library + # where it is, so the workspace stops offering to move it. + KEEP_LIBRARY_IN_NODE_FOLDER = False + @classmethod def max_worker_threads(cls) -> int: """Largest accepted worker-thread count: the detected core count.""" @@ -718,6 +724,7 @@ class PromptManagerConfig: "worker_threads": cls.WORKER_THREADS, }, "maintenance": {"safety_backup": cls.SAFETY_BACKUP}, + "library": {"keep_in_node_folder": cls.KEEP_LIBRARY_IN_NODE_FOLDER}, "gallery": GalleryConfig.get_config(), "integrations": IntegrationConfig.get_config(), "autotag": AutotagConfig.get_config(), @@ -727,14 +734,12 @@ class PromptManagerConfig: def get_config_path(cls) -> str: """Path of the persisted config.json. - Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise the file - lives at the repository root next to ``pyproject.toml``. + Honours ``PROMPT_MANAGER_CONFIG_PATH`` when set; otherwise it lives in + the library's data folder (``utils.library_location``): the node + folder for a library that was not moved, else ComfyUI's + ``user/__promptmanager``. """ - override = os.environ.get(CONFIG_PATH_ENV, "").strip() - if override: - return override - repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - return os.path.join(repo_root, "config.json") + return library_location.config_location()[0] @classmethod def load_from_file(cls, config_path: Optional[str] = None): @@ -859,6 +864,12 @@ class PromptManagerConfig: ): cls.SAFETY_BACKUP = maintenance["safety_backup"] + library = new_config.get("library", {}) + if isinstance(library, dict) and isinstance( + library.get("keep_in_node_folder"), bool + ): + cls.KEEP_LIBRARY_IN_NODE_FOLDER = library["keep_in_node_folder"] + # Update gallery config if "gallery" in new_config: GalleryConfig.update_config(new_config["gallery"]) diff --git a/py/lora_utils.py b/py/lora_utils.py index 454c182..83091db 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -567,10 +567,21 @@ def download_civitai_images( def get_lora_image_cache_dir() -> Path: - """Get the directory used to cache downloaded LoRA example images.""" - # Store in the extension's own directory + """Get the directory used to cache downloaded LoRA example images. + + ``/data/lora_images`` (``utils.library_location``): the node + folder for a library that was not moved, else ComfyUI's + ``user/__promptmanager``. + """ + try: + from ..utils import library_location + except ImportError: + from utils import library_location + ext_root = Path(__file__).resolve().parent.parent - cache = ext_root / "data" / "lora_images" + cache = Path( + library_location.cache_dir("lora_images", extension_root=str(ext_root)) + ) cache.mkdir(parents=True, exist_ok=True) return cache diff --git a/tests/browser/server.py b/tests/browser/server.py index a04c66e..25b1368 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -3,6 +3,7 @@ import argparse import asyncio import json +import logging import os import shutil import socket @@ -656,6 +657,18 @@ def exit_when_stdin_closes(loop: asyncio.AbstractEventLoop) -> threading.Thread: return thread +def close_owned_log_files() -> None: + """Close Prompt Manager's own log handlers before the temp root goes. + + The logs live in the library's data folder (#249 S6b), which for this + server is /user/__promptmanager; Windows cannot remove a file + this process still has open. + """ + for handler in logging.getLogger("prompt_manager").handlers[:]: + if getattr(handler, "_prompt_manager_owned", False) is True: + handler.close() + + def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser() parser.add_argument( @@ -680,6 +693,7 @@ def main(argv: list[str] | None = None) -> None: finally: # create_app works inside the root; Windows cannot remove a process's cwd. os.chdir(original_cwd) + close_owned_log_files() if __name__ == "__main__": diff --git a/tests/browser/system-concept.spec.mjs b/tests/browser/system-concept.spec.mjs index f5f9e97..346452b 100644 --- a/tests/browser/system-concept.spec.mjs +++ b/tests/browser/system-concept.spec.mjs @@ -181,6 +181,17 @@ function mapOverviewConcept(reference, overview) { // D96: the dialog shows runs and ratings by category and library totals, not prompt, tag or date usage. Statistics: row => ({ ...row, description: 'Runs and ratings by category, with library totals.' }), }; + const surfaces = reference.columns.surfaces.map(surface => ({ ...surface, + rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) })); + // D129: "Library location" after Backup status, in the backup row's own markup; the fixture's database is configured + // explicitly (config.json), so it does not move and its action is Details. + const [health] = surfaces; + const backup = health.rows.find(row => row.title === 'Backup status'); + const location = { ...backup, title: 'Library location', + description: `${overview.location.path} ยท The database path set in config.json.`, + action: { ...backup.action, text: 'Details' } }; + surfaces[0] = { ...health, tree: [...health.tree, 'div.operation'], + rows: [...health.rows.slice(0, 2), location, ...health.rows.slice(2)] }; return { ...reference, // D92: no "Sample configuration" chip. @@ -188,8 +199,7 @@ function mapOverviewConcept(reference, overview) { // D122: the concept's literal label is its prototype action name ("system-tab sections"). nav: { ...reference.nav, label: 'System sections' }, metrics: { ...reference.metrics, items: reference.metrics.items.map(item => ({ ...item, value: relabel[item.label] })) }, - columns: { ...reference.columns, surfaces: reference.columns.surfaces.map(surface => ({ ...surface, - rows: surface.rows.map(row => (rows[row.title] ?? (value => value))(row)) })) }, + columns: { ...reference.columns, surfaces }, save: { ...reference.save, // D92: truthful status copy, and "Save preferences" is aria-disabled (dimmed) while nothing is unsaved. status: { ...reference.save.status, text: 'No unsaved changes.' }, @@ -232,13 +242,21 @@ for (const theme of THEMES) for (const [width, height] of SIZES) { // The recorded differences, by path: D93 the review row (its real count), the backup row, the // preferences row and the sources row, and metric 4; D96 the statistics row; D92 the chip, the status and the dimmed -// Save; D122 the nav label. (The fixture's first three metrics are the concept's own counts.) +// Save; D122 the nav label; D129 the inserted "Library location" row, which shifts the thumbnail row to the fourth place +// (row 2 takes its title, text and action; row 3 and the surface's fifth child are new). (The fixture's first three +// metrics are the concept's own counts.) +const D129_ROW_ACTION = ['backgroundColor', 'borderTopColor', 'borderTopLeftRadius', 'borderTopWidth', 'className', 'color', + 'cursor', 'fontSize', 'opacity', 'paddingLeft', 'paddingTop', 'text'].map(key => `columns.surfaces.0.rows.3.action.${key}`); const UNMAPPED_PATHS = [ 'columns.surfaces.0.rows.0.title', 'columns.surfaces.0.rows.1.description', + 'columns.surfaces.0.rows.2.action.text', 'columns.surfaces.0.rows.2.description', 'columns.surfaces.0.rows.2.title', + ...D129_ROW_ACTION, 'columns.surfaces.0.rows.3.description', 'columns.surfaces.0.rows.3.title', + 'columns.surfaces.0.rows.3.tree.0', 'columns.surfaces.0.rows.3.tree.1', 'columns.surfaces.0.rows.3.tree.length', + 'columns.surfaces.0.rows.length', 'columns.surfaces.0.tree.4', 'columns.surfaces.0.tree.length', 'columns.surfaces.1.rows.0.description', 'columns.surfaces.1.rows.1.title', 'columns.surfaces.1.rows.2.description', 'heading.actions.0', 'heading.actions.length', 'metrics.items.3.value', 'nav.label', 'save.button.cursor', 'save.button.opacity', 'save.status.text', -]; +].sort(); test('system comparison detects mutations', async ({ page, request }) => { await page.setViewportSize({ width: 1600, height: 1100 }); @@ -903,3 +921,85 @@ test('thumbnail dialog comparison detects mutations', async ({ page, request }) await undoStyle(); expect(diffInventory(expected, await page.evaluate(measureThumbnailReview))).toEqual([]); }); + +/** + * Runs in the page: one `.notice` (the concept's markup: `div > h2 + p`, then `.row-actions` buttons) with its tokens and + * geometry, and its left edge against `edgeSelector` (the content it lines up with). + */ +function measureNoticeOf([selector, edgeSelector]) { + const notice = document.querySelector(selector); + 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 button = node => ({ text: text(node), className: node.className, + ...css(node, 'backgroundColor', 'color', 'fontSize', 'paddingTop', 'paddingLeft', 'borderTopColor', 'borderTopLeftRadius') }); + const edge = document.querySelector(edgeSelector); + return { + tree: [...notice.children].map(name), + text: [...notice.querySelectorAll(':scope > div:not(.row-actions) > *')].map(node => [node.localName, text(node)]), + actions: [...notice.querySelectorAll(':scope > .row-actions > button')].map(button), + frame: css(notice, 'display', 'alignItems', 'justifyContent', 'columnGap', 'borderTopWidth', 'borderTopColor', + 'borderTopLeftRadius', 'backgroundColor', 'paddingTop', 'paddingLeft', 'marginBottom', 'flexDirection'), + h2: css(notice.querySelector('h2'), 'fontSize', 'fontWeight', 'color'), + p: css(notice.querySelector('p'), 'fontSize', 'color', 'marginTop'), + rowActions: css(notice.querySelector(':scope > .row-actions'), 'display', 'columnGap', 'flexShrink'), + alignedWith: Math.round(notice.getBoundingClientRect().left - edge.getBoundingClientRect().left), + pageScroll: document.documentElement.scrollWidth > document.documentElement.clientWidth, + }; +} + +/** D129: the concept's Sources notice as the library notice: its own copy and actions, no bottom margin (main's padding). */ +function mapLibraryNotice(reference) { + return { + ...reference, + text: [['h2', 'Your library is inside the node folder'], ['p', 'Your library is stored inside the node folder, where an ' + + 'update or reinstall can delete it. Move it to ComfyUI\'s user folder: it is copied and verified first, and nothing ' + + 'is deleted.']], + actions: reference.actions.map((action, index) => ({ ...action, text: ['Keep it here', 'Move library'][index] })), + frame: { ...reference.frame, marginBottom: '0px' }, + }; +} + +async function conceptNotice(page, theme) { + await openConcept(page, 'system', theme); + await page.locator('#main .section-nav').getByRole('button', { name: 'Sources & scans', exact: true }).click(); + await expect(page.locator('#main .settings-layout > .notice')).toBeVisible(); + await page.mouse.move(0, 0); + return page.evaluate(measureNoticeOf, ['#main .settings-layout > .notice', '#main .heading']); +} + +async function productionNotice(page, theme) { + await openProduction(page, 'system/overview', theme); + await expect(page.locator('#library-banner > .notice')).toBeVisible(); + await expect(page.locator('#workspace-main .summary-metrics')).toBeVisible(); + await page.mouse.move(0, 0); + return page.evaluate(measureNoticeOf, ['#library-banner > .notice', '#workspace-main .heading']); +} + +for (const theme of THEMES) for (const [width, height] of SIZES) { + test(`library notice matches the concept notice ${theme} ${width}x${height}`, async ({ page, request }, testInfo) => { + await page.setViewportSize({ width, height }); + await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK(); + const concept = await conceptNotice(page, theme); + const production = await productionNotice(page, theme); + await page.screenshot({ path: testInfo.outputPath(`production-library-notice-${theme}-${width}.png`), fullPage: true }); + expect(diffInventory(mapLibraryNotice(concept), production)).toEqual([]); + expect(production.alignedWith).toBe(0); + expect(production.pageScroll).toBe(false); + }); +} + +test('library notice comparison detects mutations', async ({ page, request }) => { + await page.setViewportSize({ width: 1600, height: 1100 }); + await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK(); + const expected = mapLibraryNotice(await conceptNotice(page, 'dark')); + await productionNotice(page, 'dark'); + const paths = async () => diffInventory(expected, + await page.evaluate(measureNoticeOf, ['#library-banner > .notice', '#workspace-main .heading'])).map(item => item.path); + expect(await paths()).toEqual([]); + const undoToken = await mutate(page, { css: '#library-banner { --soft-blue: rgb(255, 0, 0); }' }); + expect(await paths()).toEqual(['frame.backgroundColor']); + await undoToken(); + await page.evaluate(() => { document.querySelector('#library-banner .notice h2').textContent = 'Library'; }); + expect(await paths()).toEqual(['text.0.1']); +}); diff --git a/tests/browser/workflows.spec.mjs b/tests/browser/workflows.spec.mjs index ad42f37..633c0a0 100644 --- a/tests/browser/workflows.spec.mjs +++ b/tests/browser/workflows.spec.mjs @@ -19,6 +19,13 @@ test.afterEach(async ({ page, browserName }) => { expect(page.workflowErrors).toEqual([]); }); +// #249 S6b: the legacy admin page says where the library lives (the fixture's database is configured explicitly). +test('the admin page names the library location', async ({ page }) => { + await page.goto('/prompt_manager/admin'); + await expect(page.locator('#libraryLocation')).toHaveText('Library: browser.db (set in config.json)'); + await expect(page.locator('#totalPrompts')).not.toHaveText('-'); +}); + test('actual gallery form saves a prompt that the admin page can edit and reload', async ({ page, request }) => { await page.goto('/prompt_manager/gallery'); await page.locator('#addPromptBtn').click(); diff --git a/tests/browser/workspace-system-sources.spec.mjs b/tests/browser/workspace-system-sources.spec.mjs index 3925001..574e5a4 100644 --- a/tests/browser/workspace-system-sources.spec.mjs +++ b/tests/browser/workspace-system-sources.spec.mjs @@ -206,10 +206,9 @@ test('Remove marks the draft; removing the last entry states auto-detect', async expect(await rowsOf(page)).toEqual([['p', 'small', 'No source is configured. After saving, ComfyUIโ€™s output folder is used.']]); await saveButton(page).click(); await expect(statusOf(page)).toHaveText('Saved to the server configuration.'); - // Saved: ComfyUI's output folder is detected and listed without Remove. + // Saved: ComfyUI's output folder is detected and listed without Remove, its row saying it is the default. await expect.poll(rowsOf.bind(null, page)).toEqual([ - ['ComfyUI output', 'output ยท manual scan', ['Scan folder']], - ['p', 'small', 'No source is configured; ComfyUIโ€™s output folder is used.'], + ['ComfyUI output', 'ComfyUI/output ยท manual scan ยท default', ['Scan folder']], ]); const { sources } = await serverJson(request, SETTINGS_URL); expect([sources.configured, sources.auto_detected]).toEqual([[], true]); diff --git a/tests/browser/workspace-system.spec.mjs b/tests/browser/workspace-system.spec.mjs index 0080f28..5e09a77 100644 --- a/tests/browser/workspace-system.spec.mjs +++ b/tests/browser/workspace-system.spec.mjs @@ -68,6 +68,8 @@ test('overview shows real counts and Not checked', async ({ page, request }) => expect(await rowsOf(layout, 'Library health')).toEqual([ ['0 ambiguous image links', 'Review metadata before assigning a positive / negative pair.', 'Review metadata'], ['Backup status', 'No retained backup record.', 'View backups'], + // D129: the fixture's database is configured explicitly, so it reads as config.json's choice and does not move. + ['Library location', 'browser.db ยท The database path set in config.json.', 'Details'], ['No missing thumbnails', 'Generate only missing previews when necessary.', 'Review thumbnails'], ]); expect(await rowsOf(layout, 'Workspace preferences')).toEqual([ @@ -78,6 +80,95 @@ test('overview shows real counts and Not checked', async ({ page, request }) => expect(writes).toEqual([]); }); +// D129 (#249 S6b): where the library lives. The hook makes the fixture's database read as a node-folder library; its +// Move library is a stand-in that touches no file (workspace_system_fixture.py). +const banner = page => page.locator('#library-banner'); +const LIBRARY_URL = '/prompt_manager/system/library'; + +test('a node-folder library is offered a move above every page; Keep it here stops the offer', async ({ page, request }) => { + await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK(); + await page.goto(`${workspace}#/images`); + await expect(banner(page).locator('.notice h2')).toHaveText('Your library is inside the node folder'); + await expect(banner(page).locator('.notice p')).toHaveText(/^Your library is stored inside the node folder, where an update or reinstall can delete it\./); + await expect(banner(page).locator('.notice .row-actions button')).toHaveText(['Keep it here', 'Move library']); + await expect(banner(page).getByRole('button', { name: 'Move library', exact: true })).toHaveClass(/ca-primary/); + await page.goto(`${workspace}#/library`); + await expect(banner(page).locator('.notice')).toBeVisible(); + await banner(page).getByRole('button', { name: 'Keep it here', exact: true }).click(); + await expect(banner(page)).toBeHidden(); + const { location } = await serverJson(request, LIBRARY_URL); + expect([location.keep, location.offer, location.movable]).toEqual([true, false, true]); + await page.reload(); + const layout = await openOverview(page); + await expect(banner(page)).toBeHidden(); + const row = layout.locator('.operation').filter({ hasText: 'Library location' }); + await expect(row.locator('p')).toHaveText('browser.db ยท The node folder, where an update or reinstall can delete it. ' + + 'You chose to keep it here.'); + // The row stays: the move is still one press away. + await expect(row.locator('.row-actions button')).toHaveText('Move library'); +}); + +test('Move library is reviewed, runs as an Activity job and ends the offer', async ({ page, request }) => { + await expect(await request.post('/__test__/system/library', { data: { state: 'node_folder' } })).toBeOK(); + const layout = await openOverview(page); + const row = layout.locator('.operation').filter({ hasText: 'Library location' }); + await expect(row.locator('p')).toHaveText('browser.db ยท The node folder, where an update or reinstall can delete it.'); + const move = banner(page).getByRole('button', { name: 'Move library', exact: true }); + await move.click(); + const dialog = page.locator('dialog#modal[open]'); + await expect(dialog.locator('#modal-title')).toHaveText('Move library'); + expect(await dialog.locator('dl.meta-grid > div').evaluateAll(items => items.map(item => + [item.querySelector('dt').textContent, item.querySelector('dd').textContent]))).toEqual([ + ['From', 'browser.db'], ['To', 'user/__promptmanager/prompts.db'], ['Size', expect.stringMatching(/ (KB|MB)$/)], + ['Moves with it', 'config.json and the database backups'], ['Old files', 'Renamed *.moved-to-โ€ฆ, never deleted'], + ]); + await expect(dialog.locator('.dialog-actions button')).toHaveText(['Cancel', 'Move library']); + await dialog.getByRole('button', { name: 'Cancel', exact: true }).click(); + await expect(page.locator('dialog#modal[open]')).toHaveCount(0); + await expect(move).toBeFocused(); + await row.getByRole('button', { name: 'Move library', exact: true }).click(); + await dialog.getByRole('button', { name: 'Move library', exact: true }).click(); + await expect(page).toHaveURL(/#\/activity\/[^/]+$/); + await expect(page.locator('#workspace-main')).toContainText( + /Moved the library to user\/__promptmanager: \d+ prompts? and 0 images verified\. The old files were renamed browser\.db\.moved-to-__promptmanager; nothing was deleted\./); + // The offer ends with the job; the Overview says where the library lives now. + await expect(banner(page)).toBeHidden(); + const after = await openOverview(page); + await expect(after.locator('.operation').filter({ hasText: 'Library location' }).locator('p')) + .toHaveText('browser.db ยท ComfyUI\'s user folder, which updates and reinstalls of the node leave alone.'); +}); + +test('a move that fails says so and keeps the offer; a refused start says why', async ({ page, request }) => { + await expect(await request.post('/__test__/system/library', { data: { state: 'failing' } })).toBeOK(); + await page.goto(`${workspace}#/images`); + await banner(page).getByRole('button', { name: 'Move library', exact: true }).click(); + await page.locator('dialog#modal[open]').getByRole('button', { name: 'Move library', exact: true }).click(); + await expect(page).toHaveURL(/#\/activity\/[^/]+$/); + await expect(page.locator('#workspace-main')).toContainText( + 'The copy did not pass verification. Nothing was moved. The library is still in the node folder, unchanged.'); + await expect(banner(page).locator('.notice')).toBeVisible(); + await expect(await request.post('/__test__/system/fail', { data: { path: `${LIBRARY_URL}/move` } })).toBeOK(); + await banner(page).getByRole('button', { name: 'Move library', exact: true }).click(); + await page.locator('dialog#modal[open]').getByRole('button', { name: 'Move library', exact: true }).click(); + await expect(page.locator('#workspace-toast')).toHaveText('Move library could not start. Try again.'); +}); + +test('a library that does not move shows its details', async ({ page }) => { + const layout = await openOverview(page); + await expect(banner(page)).toBeHidden(); + const details = layout.locator('.operation').filter({ hasText: 'Library location' }).getByRole('button', { name: 'Details', exact: true }); + await details.click(); + const dialog = page.locator('dialog#modal[open]'); + await expect(dialog.locator('#modal-title')).toHaveText('Library location'); + expect(await dialog.locator('dl.meta-grid > div').evaluateAll(items => items.map(item => + [item.querySelector('dt').textContent, item.querySelector('dd').textContent.replace(/[\d.,]+ (KB|MB)$/, 'N')]))).toEqual([ + ['Database', 'browser.db'], ['Why here', 'The database path set in config.json'], + ['Beside it', 'config.json, the database backups and the logs'], ['Size', 'N'], + ]); + await dialog.getByRole('button', { name: 'Close', exact: true }).click(); + await expect(details).toBeFocused(); +}); + test('Run health checks opens Health report and updates metric 4', async ({ page, request }) => { await page.goto(`${workspace}#/system/maintenance`); const layout = layoutOf(page); diff --git a/tests/browser/workspace_system_fixture.py b/tests/browser/workspace_system_fixture.py index 95f3218..fbd4557 100644 --- a/tests/browser/workspace_system_fixture.py +++ b/tests/browser/workspace_system_fixture.py @@ -23,6 +23,12 @@ read. These per-server hooks only change what a System test needs on top: image monitor whose observer is ``alive`` (default true) and schedules the named folders of this server's root (``["output"]``), so Sources reads watching, manual scan and restart pending states (S4). +- ``POST /__test__/system/library`` ``{state}`` (S6b): ``node_folder`` makes + the live database read as a node-folder library that can move to + ``user/__promptmanager``; Move library then runs a stand-in move that + touches no file and ends with the library in the default folder. + ``failing`` is the same library whose move fails verification. Reset + restores the real resolver, the real move and "Keep it here". Delays use the shared request holds (``request_hold.py``). """ @@ -54,7 +60,12 @@ def register_system_commands(app, routes, api, root: Path): from py.config import GalleryConfig, PromptManagerConfig from utils import image_monitor + from database import library_move + from utils import library_location + configured = list(GalleryConfig.MONITORING_DIRECTORIES) + 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, @@ -92,6 +103,9 @@ def register_system_commands(app, routes, api, root: Path): GalleryConfig.MONITORING_ENABLED, ) = host_values image_monitor._monitor_instance = monitor + api.__dict__.pop("_library_resolver", None) + library_move.move_library = real_move + PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = keep if config_bytes is None: config_path.unlink(missing_ok=True) else: @@ -187,6 +201,43 @@ def register_system_commands(app, routes, api, root: Path): ) return web.json_response({"success": True}) + @routes.post("/__test__/system/library") + async def library_state(request): + data = await request.json() + state = data.get("state") + if state not in ("node_folder", "failing"): + raise web.HTTPBadRequest() + target = str(root / "user" / "__promptmanager") + + def located(source): + return library_location.LibraryLocation( + source, + api.db.model.db_path, + str(config_path), + str(root), + target if source == library_location.SOURCE_NODE_FOLDER else None, + target_source=library_location.SOURCE_DEFAULT, + ) + + current = {"source": library_location.SOURCE_NODE_FOLDER} + api._library_resolver = lambda: located(current["source"]) + + def stand_in_move(plan, *, progress, **_kwargs): + if state == "failing": + raise library_move.MoveFailed("verify-failed") + progress("Copying the database") + counts = api.db.library_counts() + current["source"] = library_location.SOURCE_DEFAULT + return library_move.MoveResult( + located(library_location.SOURCE_DEFAULT), + {"prompts": counts["positive_prompts"], "generated_images": 0}, + ["prompts.db"], + ["browser.db.moved-to-__promptmanager"], + ) + + library_move.move_library = stand_in_move + return web.json_response({"success": True}) + return reset diff --git a/tests/comfyui/library-location.spec.mjs b/tests/comfyui/library-location.spec.mjs new file mode 100644 index 0000000..ae1f969 --- /dev/null +++ b/tests/comfyui/library-location.spec.mjs @@ -0,0 +1,73 @@ +// #249 S6b: a fresh install keeps its library in ComfyUI's system user folder. A pinned CPU ComfyUI starts with Prompt +// Manager and no configuration at all (no config.json, no PROMPT_MANAGER_* paths): the library must be created in +// user/__promptmanager, nothing may be written into the node folder, and ComfyUI's /userdata API must not serve it. +import { execFileSync } from 'node:child_process'; +import { access, readdir } from 'node:fs/promises'; +import { createServer } from 'node:net'; +import { join } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs'; +import { launchSharedHost } from '../../tools/vendor/comfyui-host.mjs'; +import runtime from '../../tools/vendor/comfyui-runtime.json' with { type: 'json' }; + +const repositoryRoot = fileURLToPath(new URL('../../', import.meta.url)); + +/** An unused loopback port chosen by the OS (this host shares no port with the other suites). */ +function freePort() { + return new Promise((resolve, reject) => { + const probe = createServer(); + probe.once('error', reject); + probe.listen({ host: '127.0.0.1', port: 0, exclusive: true }, () => { + const { port } = probe.address(); + probe.close(error => error ? reject(error) : resolve(port)); + }); + }); +} + +test('a fresh host keeps the library in user/__promptmanager and serves none of it', async ({ playwright }, testInfo) => { + test.skip(testInfo.project.name !== 'classic-start', 'one fresh host is enough; the renderer does not matter here'); + test.setTimeout(180000); + const port = await freePort(); + const host = await launchSharedHost({ checkout: process.env.PM_COMFY_ROOT, repository: repositoryRoot, + python: process.env.PM_TEST_PYTHON || 'python', port, pins: runtime, forwardOutput: false, + logPath: join(repositoryRoot, 'tools/vendor/test-results/executions/fresh-host.log'), pluginName: 'promptmanager', + // A fresh install: nothing written, and no inherited path variable may choose the folder. + preparePlugin: async () => ({ PROMPT_MANAGER_CONFIG_PATH: '', PROMPT_MANAGER_DATA_DIR: '', + PROMPT_MANAGER_EXTRA_GALLERY_ROOTS: '' }) }); + const request = await playwright.request.newContext({ baseURL: `http://127.0.0.1:${port}/` }); + try { + await expect.poll(async () => { + if (host.child.exitCode !== null) throw new Error('The fresh host exited before it answered'); + try { return (await request.get('/system_stats')).ok(); } catch { return false; } + }, { timeout: 120000 }).toBe(true); + expect((await (await request.get('/system_stats')).json()).system.comfyui_version).toBe(runtime.serverVersion); + await expect(await request.get('/prompt_manager/stats')).toBeOK(); + + const dataFolder = join(host.stateDirectory, 'user', '__promptmanager'); + await access(join(dataFolder, 'prompts.db')); + await access(join(dataFolder, 'logs')); + const { location } = await (await request.get('/prompt_manager/system/library')).json(); + expect(location).toMatchObject({ path: 'user/__promptmanager/prompts.db', source: 'default', movable: false, + offer: false, problem: null }); + + // Nothing but byte code was added to the node folder: no database, config.json, logs or caches. + const installed = join(host.stateDirectory, 'custom_nodes', 'promptmanager'); + const tracked = new Set(execFileSync('git', ['-C', repositoryRoot, 'ls-files', '-z'], { encoding: 'utf8' }) + .split('\0').filter(Boolean).map(file => file.split('/')[0])); + const added = (await readdir(installed)).filter(name => !tracked.has(name) && name !== '__pycache__'); + expect(added).toEqual([]); + + // ComfyUI's /userdata serves public users only: the system user's folder is out of reach. + for (const path of ['/userdata/..%2F__promptmanager%2Fprompts.db', '/api/userdata/..%2F__promptmanager%2Fprompts.db', + '/userdata/%2E%2E%2F__promptmanager%2Fconfig.json', '/userdata/prompts.db']) { + const response = await request.get(path); + expect(response.status(), path).not.toBe(200); + } + const listing = await request.get('/userdata?dir=..%2F__promptmanager&recurse=true'); + expect(listing.status()).not.toBe(200); + const asSystemUser = await request.get('/userdata/prompts.db', { headers: { 'comfy-user': '__promptmanager' } }); + expect(asSystemUser.status()).not.toBe(200); + } finally { + try { await request.dispose(); } finally { await host.stop(); } + } +}); diff --git a/tests/js/workspace-library-location.test.js b/tests/js/workspace-library-location.test.js new file mode 100644 index 0000000..c533dc1 --- /dev/null +++ b/tests/js/workspace-library-location.test.js @@ -0,0 +1,159 @@ +// Where the library lives (#249 System S6b, D129): the Overview row's copy, the Move library review and its start, +// "Keep it here" and the Move library job's end. +const test = require('node:test'); +const assert = require('node:assert/strict'); + +const locationModule = import('../../web/js/workspace/library-location.mjs'); + +const NODE_FOLDER = Object.freeze({ + path: 'custom_nodes/ComfyUI_PromptManager/prompts.db', source: 'node_folder', movable: true, keep: false, offer: true, + target: 'user/__promptmanager', problem: null, size_bytes: 1_700_000_000, + target_source: 'default', config_moves: true, +}); + +function scope() { + const controller = new AbortController(); + return { signal: controller.signal, abort: () => controller.abort(), listen() {}, cleanup() {} }; +} + +function activityStore({ running = [] } = {}) { + const listeners = new Set(); + let jobs = []; + return { + tracked: [], loads: 0, + track(job) { this.tracked.push(job); jobs = [...jobs, job]; return job; }, + async load() { this.loads += 1; }, + running() { return running; }, + jobs() { return jobs; }, + set(next) { jobs = next; for (const listener of listeners) listener(); }, + subscribe(listener) { listeners.add(listener); return () => listeners.delete(listener); }, + }; +} + +class ApiError extends Error { + constructor(status, message, code) { super(message); this.status = status; this.code = code; } +} + +test('the row says where the library is and why, the choice and any problem', async () => { + const { locationDescription, SOURCE_TEXT } = await locationModule; + assert.equal(locationDescription(NODE_FOLDER), + 'custom_nodes/ComfyUI_PromptManager/prompts.db ยท The node folder, where an update or reinstall can delete it.'); + assert.equal(locationDescription({ ...NODE_FOLDER, keep: true }), + `${NODE_FOLDER.path} ยท ${SOURCE_TEXT.node_folder}. You chose to keep it here.`); + assert.equal(locationDescription({ path: 'user/__promptmanager/prompts.db', source: 'default', movable: false, keep: true }), + 'user/__promptmanager/prompts.db ยท ComfyUI\'s user folder, which updates and reinstalls of the node leave alone.'); + assert.match(locationDescription({ ...NODE_FOLDER, problem: 'data-dir-empty' }), /PROMPT_MANAGER_DATA_DIR has no library yet/); + assert.match(locationDescription({ path: 'prompts.db', source: 'environment', movable: false }), /PROMPT_MANAGER_DATA_DIR names/); + assert.match(locationDescription({ path: 'prompts.db', source: 'config', movable: false }), /config\.json/); +}); + +test('the notice and the review name the actual target and what moves', async () => { + const { bannerText, reviewGrid, locationDescription } = await locationModule; + assert.match(bannerText(NODE_FOLDER), /Move it to ComfyUI's user folder: it is copied/); + assert.match(bannerText({ ...NODE_FOLDER, target_source: 'environment' }), + /Move it to the folder PROMPT_MANAGER_DATA_DIR names: it is copied/); + assert.match(bannerText(null), /ComfyUI's user folder/); + assert.equal(reviewGrid({ ...NODE_FOLDER, config_moves: false })[3][1], + 'The database backups (config.json stays where PROMPT_MANAGER_CONFIG_PATH names it)'); + assert.equal(reviewGrid({ ...NODE_FOLDER, config_moves: true })[3][1], 'config.json and the database backups'); + assert.match(locationDescription({ path: 'custom_nodes/pm/data/my.db', source: 'config', movable: false, + problem: 'inside-node-folder' }), /inside the node folder, where an update or reinstall deletes it/); +}); + +test('only a well-formed location is shown', async () => { + const { isLocation } = await locationModule; + assert.equal(isLocation(NODE_FOLDER), true); + for (const bad of [null, undefined, 'x', {}, { ...NODE_FOLDER, source: 'elsewhere' }, { ...NODE_FOLDER, path: 3 }, + { ...NODE_FOLDER, movable: 'yes' }]) assert.equal(isLocation(bad), false); +}); + +test('sizes read as the server would write them', async () => { + const { sizeText } = await locationModule; + assert.equal(sizeText(1_700_000_000), '1.6 GB'); + assert.equal(sizeText(120 * 2 ** 20), '120.0 MB'); + assert.equal(sizeText(1500), '2 KB'); + assert.equal(sizeText(-1), 'Unknown'); + assert.equal(sizeText(undefined), 'Unknown'); +}); + +test('the review names both folders, the size and what happens to the old files', async () => { + const { reviewGrid } = await locationModule; + assert.deepEqual(reviewGrid(NODE_FOLDER), [ + ['From', 'custom_nodes/ComfyUI_PromptManager/prompts.db'], + ['To', 'user/__promptmanager/prompts.db'], + ['Size', '1.6 GB'], + ['Moves with it', 'config.json and the database backups'], + ['Old files', 'Renamed *.moved-to-โ€ฆ, never deleted'], + ]); +}); + +test('Move library starts the job, tracks it and opens its view', async () => { + const { startLibraryMove, MOVE_URL } = await locationModule; + const posts = []; + const api = { post: async (url, body) => { posts.push([url, body]); return { job: { id: 'j1', kind: 'library_move' } }; } }; + const activity = activityStore(); + const routes = []; + const job = await startLibraryMove({ api, activity, scope: scope(), notify: () => assert.fail('no toast'), + navigate: hash => routes.push(hash) }); + assert.equal(job.id, 'j1'); + assert.deepEqual(posts, [[MOVE_URL, {}]]); + assert.deepEqual(activity.tracked.map(item => item.id), ['j1']); + assert.deepEqual(routes, ['#/activity/j1']); +}); + +test('a busy server names the running job; another refusal says the server reason', async () => { + const { startLibraryMove, START_FAILED } = await locationModule; + const running = { id: 'scan1', title: 'Metadata scan' }; + const activity = activityStore({ running: [running] }); + const toasts = []; + const busy = { post: async () => { throw new ApiError(409, 'busy', 'busy'); } }; + assert.equal(await startLibraryMove({ api: busy, activity, scope: scope(), notify: t => toasts.push(t), + navigate: () => assert.fail('nothing opens') }), null); + assert.equal(activity.loads, 1); + assert.deepEqual(toasts.at(-1), ['Another job is running: ', { text: 'Metadata scan', href: '#/activity/scan1' }, '.']); + const short = { post: async () => { throw new ApiError(507, 'There is not enough free space for the copy.', 'insufficient-space'); } }; + await startLibraryMove({ api: short, activity, scope: scope(), notify: t => toasts.push(t), navigate: () => {} }); + assert.equal(toasts.at(-1), 'There is not enough free space for the copy.'); + const broken = { post: async () => { throw new TypeError('network'); } }; + await startLibraryMove({ api: broken, activity, scope: scope(), notify: t => toasts.push(t), navigate: () => {} }); + assert.equal(toasts.at(-1), START_FAILED); + const gone = scope(); + gone.abort(); + const before = toasts.length; + await startLibraryMove({ api: broken, activity, scope: gone, notify: t => toasts.push(t), navigate: () => {} }); + assert.equal(toasts.length, before); +}); + +test('Keep it here posts the choice and resolves the new location, or null with a toast', async () => { + const { keepLibrary, KEEP_URL, KEEP_FAILED } = await locationModule; + const posts = []; + const kept = { ...NODE_FOLDER, keep: true, offer: false }; + const api = { post: async (url, body) => { posts.push([url, body]); return { location: kept }; } }; + assert.deepEqual(await keepLibrary({ api, scope: scope() }), kept); + assert.deepEqual(posts, [[KEEP_URL, { keep: true }]]); + const toasts = []; + const failing = { post: async () => { throw new ApiError(500, 'no', 'config-write-failed'); } }; + assert.equal(await keepLibrary({ api: failing, scope: scope(), notify: t => toasts.push(t) }), null); + assert.deepEqual(toasts, [KEEP_FAILED]); + const odd = { post: async () => ({ location: { path: 1 } }) }; + assert.equal(await keepLibrary({ api: odd, scope: scope() }), null); +}); + +test('the end of a Move library job is reported once; other jobs are ignored', async () => { + const { onMoveEnded } = await locationModule; + const activity = activityStore(); + let ended = 0; + const unsubscribe = onMoveEnded(activity, () => { ended += 1; }); + activity.set([{ id: 'm', kind: 'library_move', state: 'running' }, { id: 's', kind: 'scan', state: 'running' }]); + assert.equal(ended, 0); + activity.set([{ id: 'm', kind: 'library_move', state: 'running' }, { id: 's', kind: 'scan', state: 'finished' }]); + assert.equal(ended, 0); + activity.set([{ id: 'm', kind: 'library_move', state: 'failed' }]); + assert.equal(ended, 1); + activity.set([{ id: 'm', kind: 'library_move', state: 'failed' }]); + assert.equal(ended, 1); + unsubscribe(); + activity.set([{ id: 'n', kind: 'library_move', state: 'running' }]); + activity.set([{ id: 'n', kind: 'library_move', state: 'finished' }]); + assert.equal(ended, 1); +}); diff --git a/tests/js/workspace-system-sources.test.js b/tests/js/workspace-system-sources.test.js index 53e4663..ad08662 100644 --- a/tests/js/workspace-system-sources.test.js +++ b/tests/js/workspace-system-sources.test.js @@ -72,21 +72,35 @@ test('removing the last entry states auto-detect (D98)', async () => { assert.equal(NOTES.removedAll, 'No source is configured. After saving, ComfyUIโ€™s output folder is used.'); }); -test('auto-detected sources are listed without Remove, with the line that says how they were chosen', async () => { +test('auto-detected sources are listed without Remove and read as the default: "ComfyUI/ ยท ยท default"', async () => { const { sourceRows, NOTES } = await sourcesModule; - const auto = { configured: [], effective: [{ label: 'ComfyUI output', path: 'output', root_index: 0, watch: 'manual' }], auto_detected: true }; + const auto = { configured: [], auto_detected: true, + effective: [{ label: 'ComfyUI output', path: 'output', root_index: 0, watch: 'watching', in_comfyui: true }] }; + // The row says it is the default, so no line under the list repeats it; the scan keeps the server's public path. assert.deepEqual(sourceRows({ sources: auto, saved: [], current: [] }), { rows: [ - { key: 'auto:0', title: 'ComfyUI output', description: 'output ยท manual scan', scan: { label: 'Scan folder', root: 0, path: 'output' }, - remove: null }, - ], note: NOTES.autoDetected }); - // Every configured source unusable: ComfyUI's output folder is used instead. + { key: 'auto:0', title: 'ComfyUI output', description: 'ComfyUI/output ยท watching ยท default', + scan: { label: 'Scan changes', root: 0, path: 'output' }, remove: null }, + ], note: null }); + assert.equal('autoDetected' in NOTES, false); + // An output folder outside the ComfyUI base (--output-directory) is shown by its own name. + const outside = { ...auto, effective: [{ label: 'Renders', path: 'Renders', root_index: 0, watch: 'manual', in_comfyui: false }] }; + assert.equal(sourceRows({ sources: outside, saved: [], current: [] }).rows[0].description, 'Renders ยท manual scan ยท default'); + // Every configured source unusable: ComfyUI's output folder is used instead, and the line says why. const fallback = { ...auto, configured: [{ label: 'gone', path: 'gone', valid: false, code: 'not-found', reason: 'Gallery root does not exist' }] }; const result = sourceRows({ sources: fallback, saved: ['gone'], current: ['gone'] }); - assert.deepEqual(result.rows.map(row => [row.title, row.remove]), [['gone', 'gone'], ['ComfyUI output', null]]); + assert.deepEqual(result.rows.map(row => [row.title, row.description, row.remove]), + [['gone', 'gone ยท Gallery root does not exist', 'gone'], ['ComfyUI output', 'ComfyUI/output ยท watching ยท default', null]]); assert.equal(result.note, NOTES.fallback); assert.equal(sourceRows({ sources: { configured: [], effective: [], auto_detected: false }, saved: [], current: [] }).note, NOTES.none); }); +test('a configured source keeps its plain public path even inside ComfyUI', async () => { + const { sourceRows } = await sourcesModule; + const sources = { ...STATUS, effective: STATUS.effective.map(source => ({ ...source, in_comfyui: true })) }; + const { rows } = sourceRows({ sources, saved: ['output', 'input'], current: ['output', 'input'] }); + assert.deepEqual(rows.map(row => row.description), ['output ยท watching', 'input ยท manual scan']); +}); + test('a path listed twice in the draft is one row', async () => { const { sourceRows } = await sourcesModule; assert.equal(sourceRows({ sources: STATUS, saved: ['output', 'input'], current: ['output', 'output'] }).rows.length, 1); diff --git a/tests/test_api_library_location.py b/tests/test_api_library_location.py new file mode 100644 index 0000000..7b898aa --- /dev/null +++ b/tests/test_api_library_location.py @@ -0,0 +1,427 @@ +"""Library location routes (#249 S6b): where it lives, Move library, Keep it here. + +The API's database is a node-folder library inside a fake ComfyUI tree; the +resolver is pointed at that tree (``api._library_resolver``) so nothing here +reads the real environment or the repository's node folder. +""" + +import asyncio +import json +import os +import sys +import threading +from unittest.mock import MagicMock, patch + +_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 database import library_move as lm # noqa: E402 +from database.models import PromptModel # noqa: E402 +from database.operations import PromptDatabase # noqa: E402 +from py.config import PromptManagerConfig # noqa: E402 +from utils import library_location as ll # noqa: E402 + +try: + from tests.test_api_admin import AdminAPITestCase # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from test_api_admin import AdminAPITestCase # noqa: E402 + +WAIT = 10 +URL = "/prompt_manager/system/library" + + +class LibraryAPITestCase(AdminAPITestCase): + async def get_application(self): + app = await super().get_application() + # Lower case: a public path is case-folded on Windows (normcase). + self.ext = os.path.join(str(self.comfy_dir), "custom_nodes", "pm") + os.makedirs(self.ext) + self.user = os.path.join(str(self.comfy_dir), "user") + os.makedirs(self.user) + self.target = os.path.join(self.user, "__promptmanager") + sys.modules["folder_paths"].get_user_directory = lambda: self.user + previous = self.api.db + self.api.db = PromptDatabase(os.path.join(self.ext, "prompts.db")) + previous.close_all() + self.api.db.save_prompt("a library prompt") + self.fp = sys.modules["folder_paths"] + self.environ = {} + self.api._library_resolver = lambda: ll.resolve_library( + None, environ=self.environ, extension_root=self.ext, folder_paths=self.fp + ) + ll.clear_override() + self.addCleanup(ll.clear_override) + self.addCleanup(PromptModel.close_all_instances) + self._keep = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER + self.addCleanup( + setattr, PromptManagerConfig, "KEEP_LIBRARY_IN_NODE_FOLDER", self._keep + ) + PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER = False + return app + + async def json(self, method, url, status=200, **kwargs): + response = await self.client.request(method, url, **kwargs) + text = await response.text() + self.assertEqual(response.status, status, text) + self.assertNotIn(self.tmpdir, text) + return json.loads(text) + + async def finished(self, body): + job = self.api._jobs.get(body["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + return job + + +class TestLocationRoute(LibraryAPITestCase): + async def test_node_folder_library_offers_the_move(self): + body = await self.json("GET", URL) + self.assertEqual( + body["location"], + { + "path": "custom_nodes/pm/prompts.db", + "source": "node_folder", + "movable": True, + "keep": False, + "offer": True, + "target": "user/__promptmanager", + "problem": None, + "size_bytes": os.path.getsize(self.api.db.model.db_path), + "target_source": "default", + "config_moves": True, + }, + ) + + async def test_overview_carries_the_location(self): + body = await self.json("GET", "/prompt_manager/system/overview") + self.assertEqual(body["location"]["source"], "node_folder") + self.assertTrue(body["location"]["offer"]) + + async def test_default_folder_library_offers_nothing(self): + self.api._library_resolver = lambda: ll.LibraryLocation( + ll.SOURCE_DEFAULT, + self.api.db.model.db_path, + os.path.join(self.target, "config.json"), + self.target, + self.target, + ) + location = (await self.json("GET", URL))["location"] + self.assertEqual((location["source"], location["movable"]), ("default", False)) + self.assertFalse(location["offer"]) + + async def test_a_resolved_file_that_is_not_the_live_database_is_never_moved(self): + self.api._library_resolver = lambda: ll.resolve_library( + None, + environ={}, + extension_root=os.path.join(self.tmpdir, "elsewhere"), + folder_paths=self.fp, + ) + location = (await self.json("GET", URL))["location"] + self.assertEqual(location["source"], "config") + self.assertFalse(location["movable"]) + refused = await self.json("POST", URL + "/move", 400) + self.assertEqual(refused["code"], "not-movable") + + +class TestLocationDetails(LibraryAPITestCase): + async def test_the_target_and_what_moves_follow_the_environment(self): + data = os.path.join(self.tmpdir, "pm-data") + self.environ = { + ll.DATA_DIR_ENV: data, + ll.CONFIG_PATH_ENV: os.path.join(self.tmpdir, "mine.json"), + } + location = (await self.json("GET", URL))["location"] + self.assertEqual(location["target_source"], "environment") + self.assertFalse(location["config_moves"]) + self.assertEqual(location["problem"], "data-dir-empty") + + async def test_a_configured_library_inside_the_node_folder_is_flagged(self): + self.api._library_resolver = lambda: ll.LibraryLocation( + ll.SOURCE_CONFIG, + self.api.db.model.db_path, + os.path.join(self.ext, "config.json"), + self.ext, + ) + self.api._extension_root = self.ext + location = (await self.json("GET", URL))["location"] + self.assertEqual(location["problem"], "inside-node-folder") + self.assertFalse(location["movable"]) + + async def test_a_startup_after_a_crash_finishes_the_rename(self): + target_db = os.path.join(self.target, "prompts.db") + os.makedirs(self.target) + with open(target_db, "wb") as handle: + handle.write(b"x") + lm.write_marker(self.target, ll.MARKER_MOVED, self.api.db.model.db_path) + old = self.api.db.model.db_path + self.api.db.model.close_all() + moved = ll.LibraryLocation( + ll.SOURCE_DEFAULT, + target_db, + os.path.join(self.target, "config.json"), + self.target, + self.target, + ) + self.api._library_resolver = lambda: moved + with patch.object(self.api.db.model, "db_path", target_db): + renamed = self.api.finish_interrupted_library_move() + self.assertEqual(renamed, ["prompts.db.moved-to-__promptmanager"]) + self.assertFalse(os.path.exists(old)) + + +class TestLocationThroughASymlinkedNodeFolder(LibraryAPITestCase): + """custom_nodes/ is a symlink to a checkout outside ComfyUI. + + The library location reads relative to ComfyUI through the node folder's + own (unresolved) place in custom_nodes, not as a bare file name; the + shared ``_public_path`` (sources, jobs) keeps resolving symlinks. + """ + + async def get_application(self): + app = await super().get_application() + self.real = os.path.join(self.tmpdir, "checkout", "ComfyUI_PromptManager") + os.makedirs(self.real) + self.link = os.path.join(str(self.comfy_dir), "custom_nodes", "promptmanager") + try: + os.symlink(self.real, self.link, target_is_directory=True) + except (OSError, NotImplementedError) as exc: + self.skipTest(f"symlinks not permitted here: {exc}") + previous = self.api.db + self.api.db = PromptDatabase(os.path.join(self.link, "prompts.db")) + previous.close_all() + self.api.db.save_prompt("a library prompt") + self.api._library_resolver = lambda: ll.resolve_library( + None, environ=self.environ, extension_root=self.link, folder_paths=self.fp + ) + self.api._extension_root = self.link + return app + + async def test_location_names_the_node_folder_in_custom_nodes(self): + location = (await self.json("GET", URL))["location"] + self.assertEqual(location["source"], "node_folder") + self.assertEqual(location["path"], "custom_nodes/promptmanager/prompts.db") + self.assertEqual(location["target"], "user/__promptmanager") + overview = await self.json("GET", "/prompt_manager/system/overview") + self.assertEqual( + overview["location"]["path"], "custom_nodes/promptmanager/prompts.db" + ) + + async def test_a_resolved_spelling_of_the_database_reads_the_same(self): + from py.api import library_location as routes + + resolved = os.path.join(os.path.realpath(self.real), "prompts.db") + self.assertEqual( + routes.node_folder_public_path(resolved, self.link, []), + "custom_nodes/promptmanager/prompts.db", + ) + + async def test_the_shared_public_path_is_unchanged(self): + self.assertEqual( + self.api._public_path(os.path.join(self.link, "prompts.db")), "prompts.db" + ) + + +class TestNodeFolderPublicPath(LibraryAPITestCase): + def display(self, path, root=None, custom_nodes=()): + from py.api import library_location as routes + + return routes.node_folder_public_path(path, root or self.ext, custom_nodes) + + async def test_a_plain_node_folder_reads_relative_to_custom_nodes(self): + self.assertEqual( + self.display(os.path.join(self.ext, "prompts.db")), + "custom_nodes/pm/prompts.db", + ) + self.assertEqual( + self.display(os.path.join(self.ext, "data", "prompts.db")), + "custom_nodes/pm/data/prompts.db", + ) + + async def test_a_custom_nodes_entry_from_comfyui_names_the_folder(self): + nodes = os.path.join(self.tmpdir, "more_nodes") + root = os.path.join(nodes, "pm") + self.assertEqual( + self.display(os.path.join(root, "prompts.db"), root, [nodes]), + "more_nodes/pm/prompts.db", + ) + + async def test_anything_else_is_not_a_node_folder_path(self): + # Outside the node folder, a parent that is no custom_nodes folder, or + # a spelling that climbs out: the caller falls back to _public_path. + self.assertIsNone(self.display(os.path.join(self.user, "prompts.db"))) + loose = os.path.join(self.tmpdir, "somewhere", "pm") + self.assertIsNone(self.display(os.path.join(loose, "prompts.db"), loose)) + self.assertIsNone(self.display(os.path.join(self.ext, "..", "prompts.db"))) + self.assertIsNone(self.display("")) + self.assertIsNone(self.display(None)) + + async def test_a_library_outside_the_node_folder_keeps_the_public_path(self): + self.api._library_resolver = lambda: ll.LibraryLocation( + ll.SOURCE_DEFAULT, + self.api.db.model.db_path, + os.path.join(self.target, "config.json"), + self.target, + self.target, + ) + self.api._extension_root = os.path.join(self.tmpdir, "elsewhere") + location = (await self.json("GET", URL))["location"] + self.assertEqual(location["path"], "custom_nodes/pm/prompts.db") + + +class TestMoveRoute(LibraryAPITestCase): + async def test_move_runs_as_an_exclusive_job_and_switches_the_live_library(self): + body = await self.json("POST", URL + "/move", 202) + summary = body["job"] + self.assertEqual(summary["kind"], "library_move") + self.assertEqual(summary["title"], "Move library") + self.assertEqual(summary["admission"], "exclusive") + self.assertFalse(summary["cancellable"]) + job = await self.finished(body) + self.assertEqual(job.state, "finished", job.status_line) + self.assertEqual( + self.api.db.model.db_path, os.path.join(self.target, "prompts.db") + ) + self.assertIn("user/__promptmanager", job.status_line) + self.assertIn("1 prompt", job.status_line) + self.assertNotIn(self.tmpdir, json.dumps(job.report())) + location = (await self.json("GET", URL))["location"] + self.assertEqual(location["source"], "default") + self.assertEqual(location["path"], "user/__promptmanager/prompts.db") + self.assertFalse(location["offer"]) + prompts = await self.json("GET", "/prompt_manager/stats") + self.assertEqual(prompts["stats"]["total_prompts"], 1) + self.assertTrue( + os.path.isfile( + os.path.join(self.ext, "prompts.db.moved-to-__promptmanager") + ) + ) + + async def test_refusals_are_coded_without_paths(self): + with patch.object(lm, "_free_bytes", return_value=1024): + short = await self.json("POST", URL + "/move", 507) + self.assertEqual(short["code"], "insufficient-space") + os.makedirs(self.target) + with open(os.path.join(self.target, "prompts.db"), "wb") as handle: + handle.write(b"another library") + taken = await self.json("POST", URL + "/move", 409) + self.assertEqual(taken["code"], "target-has-library") + self.assertEqual(self.api._jobs.list(), []) + + async def test_move_is_refused_while_another_job_runs_and_blocks_others(self): + gate = asyncio.Event() + + async def scanning(job): + await gate.wait() + + self.api._jobs.start( + "scan", "Metadata scan", admission="records", runner=scanning + ) + busy = await self.json("POST", URL + "/move", 409) + self.assertEqual(busy["code"], "busy") + gate.set() + await asyncio.sleep(0) + for job in self.api._jobs.list(): + if job.task is not None: + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + + release = threading.Event() + entered = threading.Event() + real = lm.move_library + + def gated(*args, **kwargs): + entered.set() + release.wait(WAIT) + return real(*args, **kwargs) + + with patch.object(lm, "move_library", side_effect=gated): + body = await self.json("POST", URL + "/move", 202) + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, entered.wait, WAIT) + refused = await self.json("POST", "/prompt_manager/system/health", 409) + self.assertEqual(refused["running"]["title"], "Move library") + release.set() + job = await self.finished(body) + self.assertEqual(job.state, "finished", job.status_line) + + async def test_a_failed_move_ends_truthfully_and_keeps_the_old_library(self): + with patch.object( + lm, "move_library", side_effect=lm.MoveFailed("verify-failed") + ): + body = await self.json("POST", URL + "/move", 202) + job = await self.finished(body) + self.assertEqual(job.state, "failed") + self.assertIn("did not pass verification", job.status_line) + self.assertIn("still in the node folder", job.status_line) + self.assertEqual( + self.api.db.model.db_path, os.path.join(self.ext, "prompts.db") + ) + with patch.object(lm, "move_library", side_effect=OSError("disk gone")): + body = await self.json("POST", URL + "/move", 202) + job = await self.finished(body) + self.assertEqual(job.state, "failed") + self.assertNotIn("disk gone", job.status_line) + self.assertIn("still in the node folder", job.status_line) + + async def test_old_file_left_in_place_is_reported(self): + real = lm.retire_old_files + with patch.object(lm, "retire_old_files", return_value=([], "old-in-use")): + body = await self.json("POST", URL + "/move", 202) + job = await self.finished(body) + self.assertEqual(job.state, "finished") + self.assertIn("left in the node folder", job.status_line) + self.assertTrue(callable(real)) + + +class TestKeepRoute(LibraryAPITestCase): + async def test_keep_it_here_is_persisted_and_stops_the_offer(self): + body = await self.json("POST", URL + "/keep", 200, json={"keep": True}) + self.assertTrue(body["location"]["keep"]) + self.assertFalse(body["location"]["offer"]) + self.assertTrue(body["location"]["movable"]) + with open(self.config_path, encoding="utf-8") as handle: + saved = json.load(handle) + self.assertEqual(saved["library"], {"keep_in_node_folder": True}) + self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + undone = await self.json("POST", URL + "/keep", 200, json={"keep": False}) + self.assertTrue(undone["location"]["offer"]) + + async def test_keep_needs_a_boolean(self): + for payload in ({"keep": "yes"}, {}, [True]): + bad = await self.json("POST", URL + "/keep", 400, json=payload) + self.assertEqual(bad["code"], "invalid-request") + self.assertFalse(os.path.exists(self.config_path)) + + async def test_unwritable_config_reports_not_saved(self): + with patch("py.config_store._write_atomically", side_effect=OSError("ro")): + body = await self.json("POST", URL + "/keep", 500, json={"keep": True}) + self.assertEqual(body["code"], "config-write-failed") + self.assertFalse(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + + +class TestLegacyRoutesDuringAMove(LibraryAPITestCase): + async def test_legacy_restore_is_refused_while_the_move_runs(self): + from aiohttp import FormData + + release = threading.Event() + entered = threading.Event() + real = lm.move_library + + def gated(*args, **kwargs): + entered.set() + release.wait(WAIT) + return real(*args, **kwargs) + + with patch.object(lm, "move_library", side_effect=gated): + body = await self.json("POST", URL + "/move", 202) + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, entered.wait, WAIT) + form = FormData() + form.add_field("database_file", b"SQLite format 3\x00", filename="x.db") + response = await self.client.request( + "POST", "/prompt_manager/restore", data=form + ) + self.assertEqual(response.status, 409, await response.text()) + release.set() + await self.finished(body) diff --git a/tests/test_ingestion_resolver_profile.py b/tests/test_ingestion_resolver_profile.py index a222849..2b96ced 100644 --- a/tests/test_ingestion_resolver_profile.py +++ b/tests/test_ingestion_resolver_profile.py @@ -66,7 +66,7 @@ class TestIngestionKeepsDefaultProfile(unittest.TestCase): handler.link_image_to_prompt("/out/f03.png", {"id": 1}, metadata) self.assertEqual(len(spy.calls), 2) spy.assert_default_and_gated(self) - handler.db_manager.link_image_to_prompt.assert_called_once() + handler.db_manager.link_image_to_prompt_or_raise.assert_called_once() def test_admin_scan_ingestion_uses_the_default_profile(self): from py.api import admin diff --git a/tests/test_library_location.py b/tests/test_library_location.py new file mode 100644 index 0000000..6a5f10c --- /dev/null +++ b/tests/test_library_location.py @@ -0,0 +1,369 @@ +"""Where the library lives (#249 S6b): the one resolver, ``utils/library_location``. + +Every case builds its own extension root (the node folder) and ComfyUI user +folder in a temp directory and hands the resolver a fake ``folder_paths``, so +no test reads the real environment, the repository's own files or ComfyUI. +""" + +import json +import os +import sys +import tempfile +import types +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils import library_location as ll # noqa: E402 + + +def _touch(path, data=b"x"): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "wb") as handle: + handle.write(data) + + +def _folder_paths(base, *, helper=True): + """A ComfyUI ``folder_paths`` stand-in rooted at ``base``.""" + module = types.SimpleNamespace(base_path=base) + for name in ("output", "input", "temp", "user"): + directory = os.path.join(base, name) + os.makedirs(directory, exist_ok=True) + setattr(module, f"get_{name}_directory", lambda d=directory: d) + if helper: + module.get_system_user_directory = lambda name="system": os.path.join( + base, "user", "__" + name + ) + return module + + +class _Case(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.base = os.path.realpath(self._tmp.name) + self.comfy = os.path.join(self.base, "ComfyUI") + self.ext = os.path.join(self.comfy, "custom_nodes", "ComfyUI_PromptManager") + os.makedirs(self.ext) + self.fp = _folder_paths(self.comfy) + self.default_dir = os.path.join(self.comfy, "user", "__promptmanager") + ll.clear_override() + self.addCleanup(ll.clear_override) + + def resolve(self, default_path=None, environ=None, **kwargs): + kwargs.setdefault("folder_paths", self.fp) + return ll.resolve_library( + default_path, environ=environ or {}, extension_root=self.ext, **kwargs + ) + + def config_location(self, environ=None, **kwargs): + kwargs.setdefault("folder_paths", self.fp) + return ll.config_location( + environ=environ or {}, extension_root=self.ext, **kwargs + ) + + +class TestDefaultDataFolder(_Case): + def test_uses_comfyui_system_user_helper(self): + self.assertEqual(ll.default_data_dir(self.fp), self.default_dir) + + def test_falls_back_to_user_directory_without_helper(self): + fp = _folder_paths(self.comfy, helper=False) + self.assertEqual(ll.default_data_dir(fp), self.default_dir) + + def test_none_outside_comfyui(self): + self.assertIsNone(ll.default_data_dir(None)) + self.assertIsNone(ll.default_data_dir(types.SimpleNamespace())) + + +class TestResolutionOrder(_Case): + def test_new_install_uses_default_folder_and_creates_nothing_in_node_folder(self): + location = self.resolve() + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db")) + self.assertEqual( + location.config_path, os.path.join(self.default_dir, "config.json") + ) + self.assertFalse(location.movable) + self.assertEqual(os.listdir(self.ext), []) + + def test_legacy_node_folder_library_stays_in_use(self): + _touch(os.path.join(self.ext, "prompts.db")) + location = self.resolve() + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db")) + self.assertEqual(location.config_path, os.path.join(self.ext, "config.json")) + self.assertEqual(location.data_dir, self.ext) + self.assertTrue(location.movable) + self.assertEqual(location.target_dir, self.default_dir) + + def test_legacy_wins_over_an_unmarked_default_folder_library(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + + def test_explicit_default_path_is_honoured_as_today(self): + _touch(os.path.join(self.ext, "prompts.db")) + elsewhere = os.path.join(self.base, "elsewhere", "pm.db") + location = self.resolve(elsewhere) + self.assertEqual(location.source, ll.SOURCE_CONFIG) + self.assertEqual(location.db_path, elsewhere) + self.assertFalse(location.movable) + relative = self.resolve("data/my.db") + self.assertEqual(relative.db_path, os.path.join(self.ext, "data", "my.db")) + self.assertEqual(relative.source, ll.SOURCE_CONFIG) + + def test_default_setting_value_is_not_explicit(self): + for value in (None, "", "prompts.db"): + self.assertEqual(self.resolve(value).source, ll.SOURCE_DEFAULT, value) + + def test_environment_folder_wins(self): + data = os.path.join(self.base, "pm-data") + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_ENVIRONMENT) + self.assertEqual(location.db_path, os.path.join(data, "prompts.db")) + self.assertEqual(location.config_path, os.path.join(data, "config.json")) + + def test_environment_folder_with_a_library_wins_over_legacy(self): + data = os.path.join(self.base, "pm-data") + _touch(os.path.join(data, "prompts.db")) + _touch(os.path.join(self.ext, "prompts.db")) + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_ENVIRONMENT) + + def test_empty_environment_folder_never_replaces_a_legacy_library(self): + data = os.path.join(self.base, "pm-data") + _touch(os.path.join(self.ext, "prompts.db")) + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db")) + self.assertEqual(location.problem, "data-dir-empty") + self.assertEqual(location.target_dir, data) + self.assertTrue(location.movable) + + def test_config_path_environment_is_honoured(self): + custom = os.path.join(self.base, "cfg", "custom.json") + location = self.resolve(environ={ll.CONFIG_PATH_ENV: custom}) + self.assertEqual(location.config_path, custom) + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + + def test_without_comfyui_the_node_folder_is_used(self): + location = self.resolve(folder_paths=None) + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.db_path, os.path.join(self.ext, "prompts.db")) + self.assertFalse(location.movable) + + def test_override_after_a_live_move_wins(self): + _touch(os.path.join(self.ext, "prompts.db")) + moved = ll.LibraryLocation( + source=ll.SOURCE_DEFAULT, + db_path=os.path.join(self.default_dir, "prompts.db"), + config_path=os.path.join(self.default_dir, "config.json"), + data_dir=self.default_dir, + target_dir=self.default_dir, + ) + ll.switch_to(moved) + self.assertEqual(self.resolve(), moved) + self.assertEqual(self.config_location(), (moved.config_path, moved.data_dir)) + ll.clear_override() + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + + +class TestMarker(_Case): + def _marker(self, directory, state): + os.makedirs(directory, exist_ok=True) + with open(os.path.join(directory, ll.MARKER_NAME), "w") as handle: + json.dump({"version": 1, "state": state}, handle) + + def test_moved_marker_wins_while_the_old_file_was_not_renamed_yet(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self._marker(self.default_dir, "moved") + location = self.resolve() + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db")) + self.assertEqual( + location.config_path, os.path.join(self.default_dir, "config.json") + ) + + def test_pending_marker_keeps_the_old_library(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self._marker(self.default_dir, "pending") + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + data = os.path.join(self.base, "pm-data") + _touch(os.path.join(data, "prompts.db")) + self._marker(data, "pending") + location = self.resolve(environ={ll.DATA_DIR_ENV: data}) + self.assertEqual(location.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(location.problem, "move-unfinished") + + def test_unreadable_marker_is_not_trusted(self): + _touch(os.path.join(self.ext, "prompts.db")) + _touch(os.path.join(self.default_dir, "prompts.db")) + _touch(os.path.join(self.default_dir, ll.MARKER_NAME), b"{not json") + self.assertEqual(self.resolve().source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(ll.read_marker(self.default_dir), {}) + + def test_renamed_old_file_is_never_picked(self): + _touch(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")) + _touch(os.path.join(self.default_dir, "prompts.db")) + self.assertEqual(self.resolve().source, ll.SOURCE_DEFAULT) + + +class TestConfigLocation(_Case): + def test_legacy_config_without_a_database_keeps_being_read(self): + _touch(os.path.join(self.ext, "config.json"), b"{}") + path, data_dir = self.config_location() + self.assertEqual(path, os.path.join(self.ext, "config.json")) + self.assertEqual(data_dir, self.ext) + + def test_legacy_config_alone_never_puts_a_new_database_in_the_node_folder(self): + _touch(os.path.join(self.ext, "config.json"), b"{}") + location = self.resolve() + self.assertEqual(location.db_path, os.path.join(self.default_dir, "prompts.db")) + self.assertEqual(location.config_path, os.path.join(self.ext, "config.json")) + self.assertFalse(location.movable) + + def test_new_install_reads_config_from_the_data_folder(self): + path, data_dir = self.config_location() + self.assertEqual(path, os.path.join(self.default_dir, "config.json")) + self.assertEqual(data_dir, self.default_dir) + + def test_logs_and_cache_follow_the_data_folder(self): + self.assertEqual( + ll.logs_dir(environ={}, extension_root=self.ext, folder_paths=self.fp), + os.path.join(self.default_dir, "logs"), + ) + _touch(os.path.join(self.ext, "prompts.db")) + self.assertEqual( + ll.logs_dir(environ={}, extension_root=self.ext, folder_paths=self.fp), + os.path.join(self.ext, "logs"), + ) + self.assertEqual( + ll.cache_dir( + "lora_images", environ={}, extension_root=self.ext, folder_paths=self.fp + ), + os.path.join(self.ext, "data", "lora_images"), + ) + + +class TestDataDirValidation(_Case): + def problem(self, path, **kwargs): + kwargs.setdefault("folder_paths", self.fp) + return ll.data_dir_problem(path, **kwargs) + + def test_valid_folder(self): + self.assertIsNone(self.problem(os.path.join(self.base, "pm-data"))) + self.assertIsNone(self.problem(self.default_dir)) + + def test_relative_and_empty_paths_are_refused(self): + self.assertEqual(self.problem("pm-data"), "data-dir-not-absolute") + self.assertEqual(self.problem(""), "data-dir-not-absolute") + + def test_served_folders_are_refused(self): + for name in ("output", "input", "temp"): + inside = os.path.join(self.comfy, name, "pm") + self.assertEqual(self.problem(inside), "data-dir-served", name) + self.assertEqual( + self.problem(os.path.join(self.comfy, "output")), "data-dir-served" + ) + + def test_public_user_folder_is_refused_but_system_users_are_fine(self): + public = os.path.join(self.comfy, "user", "default", "pm") + self.assertEqual(self.problem(public), "data-dir-served") + self.assertIsNone(self.problem(os.path.join(self.comfy, "user", "__pm"))) + + def test_gallery_source_is_refused(self): + source = os.path.join(self.base, "gallery") + os.makedirs(source) + self.assertEqual( + self.problem(os.path.join(source, "pm"), gallery_roots=[source]), + "data-dir-served", + ) + self.assertEqual( + self.problem( + os.path.join(source, "pm"), + environ={ll.EXTRA_GALLERY_ROOTS_ENV: source}, + ), + "data-dir-served", + ) + + def test_symlink_into_output_is_refused(self): + link = os.path.join(self.base, "link-to-output") + try: + os.symlink(os.path.join(self.comfy, "output"), link) + except (OSError, NotImplementedError): + self.skipTest("symlinks are not available") + self.assertEqual(self.problem(os.path.join(link, "pm")), "data-dir-served") + + def test_case_variant_spelling_of_output_is_refused_where_the_volume_folds_case( + self, + ): + upper = os.path.join(self.comfy, "OUTPUT") + if not os.path.isdir(upper): + self.skipTest("this volume is case-sensitive") + self.assertEqual(self.problem(os.path.join(upper, "pm")), "data-dir-served") + + def test_invalid_environment_folder_is_ignored_with_a_problem(self): + inside = os.path.join(self.comfy, "output", "pm") + location = self.resolve(environ={ll.DATA_DIR_ENV: inside}) + self.assertEqual(location.source, ll.SOURCE_DEFAULT) + self.assertEqual(location.problem, "data-dir-served") + self.assertFalse(os.path.exists(inside)) + relative = self.resolve(environ={ll.DATA_DIR_ENV: "relative/dir"}) + self.assertEqual(relative.problem, "data-dir-not-absolute") + + +class TestOneRuleForConfigAndDatabase(_Case): + """M1: config.json and the database never split over PROMPT_MANAGER_DATA_DIR.""" + + def test_a_source_listed_only_in_a_config_file_rules_out_the_folder_for_both(self): + source = os.path.join(self.base, "gallery") + os.makedirs(source) + _touch( + os.path.join(self.ext, "config.json"), + json.dumps({"gallery": {"monitoring": {"directories": [source]}}}).encode(), + ) + environ = {ll.DATA_DIR_ENV: os.path.join(source, "pm")} + location = self.resolve(environ=environ) + config_path, data_dir = self.config_location(environ=environ) + self.assertEqual(location.problem, "data-dir-served") + self.assertEqual( + (location.config_path, location.data_dir), (config_path, data_dir) + ) + self.assertNotEqual(location.source, ll.SOURCE_ENVIRONMENT) + + def test_a_relative_source_is_taken_from_comfyui(self): + self.fp.base_path = self.comfy + os.makedirs(os.path.join(self.comfy, "pictures")) + _touch( + os.path.join(self.ext, "config.json"), + json.dumps( + {"gallery": {"monitoring": {"directories": ["pictures"]}}} + ).encode(), + ) + environ = {ll.DATA_DIR_ENV: os.path.join(self.comfy, "pictures", "pm")} + self.assertEqual(self.resolve(environ=environ).problem, "data-dir-served") + self.assertEqual(self.config_location(environ=environ)[1], self.ext) + + def test_unreadable_or_odd_config_files_add_no_source(self): + for payload in ( + b"{not json", + b"[]", + json.dumps({"gallery": {"monitoring": {"directories": "x"}}}).encode(), + ): + _touch(os.path.join(self.ext, "config.json"), payload) + environ = {ll.DATA_DIR_ENV: os.path.join(self.base, "pm-data")} + self.assertIsNone(self.resolve(environ=environ).problem) + + +class TestSameFolder(_Case): + def test_same_folder_follows_links_and_case(self): + self.assertTrue(ll.same_folder(self.ext, self.ext + os.sep)) + self.assertFalse(ll.same_folder(self.ext, self.default_dir)) + self.assertFalse(ll.same_folder(self.ext, os.path.join(self.base, "missing"))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_library_location_wiring.py b/tests/test_library_location_wiring.py new file mode 100644 index 0000000..d2fa31a --- /dev/null +++ b/tests/test_library_location_wiring.py @@ -0,0 +1,169 @@ +"""Every data path asks the one resolver (#249 S6b). + +The database, the diagnostics database, config.json, the logs and the LoRA +image cache all follow ``utils.library_location``: a folder named by +``PROMPT_MANAGER_DATA_DIR`` (and, after a live move, the switched location). +""" + +import os +import pathlib +import sys +import tempfile +import types +import unittest +import unittest.mock as 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 utils import library_location as ll # noqa: E402 + + +class _EnvCase(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + base = os.path.realpath(self._tmp.name) + self.data = os.path.join(base, "pm-data") + patcher = mock.patch.dict( + os.environ, + {ll.DATA_DIR_ENV: self.data, ll.CONFIG_PATH_ENV: ""}, + ) + patcher.start() + self.addCleanup(patcher.stop) + # This test's own ComfyUI folders (the shared stub serves the system + # temp folder as "output", which would rightly refuse self.data). + comfy = os.path.join(base, "ComfyUI") + fp = types.ModuleType("folder_paths") + for name in ("output", "input", "temp", "user"): + setattr(fp, f"get_{name}_directory", lambda n=name: os.path.join(comfy, n)) + modules = mock.patch.dict(sys.modules, {"folder_paths": fp}) + modules.start() + self.addCleanup(modules.stop) + ll.clear_override() + self.addCleanup(ll.clear_override) + + +class TestDatabasePath(_EnvCase): + def test_database_follows_the_data_folder(self): + from database.operations import _resolve_db_path + + with mock.patch.dict(sys.modules, {"py.config": None}): + path = _resolve_db_path(None) + self.assertEqual(path, os.path.join(self.data, "prompts.db")) + self.assertTrue(os.path.isdir(self.data)) + if os.name == "posix": + self.assertEqual(os.stat(self.data).st_mode & 0o777, 0o700) + + def test_diagnostics_use_the_same_resolver(self): + from database.operations import _resolve_db_path + from utils.diagnostics import GalleryDiagnostics + + with mock.patch.dict(sys.modules, {"py.config": None}): + self.assertEqual(GalleryDiagnostics._resolve_db_path(), _resolve_db_path()) + + def test_switched_location_wins(self): + from database.operations import _resolve_db_path + + moved = ll.LibraryLocation( + ll.SOURCE_DEFAULT, + os.path.join(self.data, "moved", "prompts.db"), + os.path.join(self.data, "moved", "config.json"), + os.path.join(self.data, "moved"), + ) + ll.switch_to(moved) + self.assertEqual(_resolve_db_path(None), moved.db_path) + + +class TestConfigPath(_EnvCase): + def test_config_json_follows_the_data_folder(self): + from py.config import PromptManagerConfig + + self.assertEqual( + PromptManagerConfig.get_config_path(), + os.path.join(self.data, "config.json"), + ) + + def test_explicit_config_path_still_wins(self): + from py.config import PromptManagerConfig + + custom = os.path.join(self.data, "elsewhere.json") + with mock.patch.dict(os.environ, {ll.CONFIG_PATH_ENV: custom}): + self.assertEqual(PromptManagerConfig.get_config_path(), custom) + + def test_keep_choice_round_trips(self): + from py.config import PromptManagerConfig + + original = PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER + self.addCleanup( + setattr, PromptManagerConfig, "KEEP_LIBRARY_IN_NODE_FOLDER", original + ) + PromptManagerConfig.update_config({"library": {"keep_in_node_folder": True}}) + self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + self.assertEqual( + PromptManagerConfig.get_config()["library"], {"keep_in_node_folder": True} + ) + PromptManagerConfig.update_config({"library": {"keep_in_node_folder": "yes"}}) + self.assertTrue(PromptManagerConfig.KEEP_LIBRARY_IN_NODE_FOLDER) + + +class TestLogsAndCache(_EnvCase): + def test_logs_go_to_the_data_folder(self): + from utils import logging_config + from utils.logging_config import PromptManagerLogger + + manager = object.__new__(PromptManagerLogger) + self.addCleanup(self._close, manager) + manager.__init__() + self.assertEqual(manager.log_dir, pathlib.Path(self.data, "logs")) + self.assertTrue(manager.log_dir.is_dir()) + self.assertIs(logging_config.library_location, ll) + + @staticmethod + def _close(manager): + import logging + + for handler in logging.getLogger("prompt_manager").handlers[:]: + if getattr(handler, "_prompt_manager_owned", False): + handler.close() + try: + from utils.logging_config import get_logger_manager + + get_logger_manager()._setup_loggers() + except Exception: + pass + + def test_lora_image_cache_follows_the_data_folder(self): + from py import lora_utils + + cache = lora_utils.get_lora_image_cache_dir() + self.assertEqual(cache, pathlib.Path(self.data, "data", "lora_images")) + + +class TestNodeFolderUnchangedForLegacyInstalls(unittest.TestCase): + """A node folder that holds the library keeps every path where it was.""" + + def test_legacy_paths(self): + with tempfile.TemporaryDirectory() as tmp: + ext = os.path.join(tmp, "ext") + os.makedirs(ext) + open(os.path.join(ext, "prompts.db"), "wb").close() + fp = types.SimpleNamespace( + get_user_directory=lambda: os.path.join(tmp, "u") + ) + location = ll.resolve_library( + None, environ={}, extension_root=ext, folder_paths=fp + ) + self.assertEqual(location.db_path, os.path.join(ext, "prompts.db")) + self.assertEqual( + ll.logs_dir(environ={}, extension_root=ext, folder_paths=fp), + os.path.join(ext, "logs"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_library_move.py b/tests/test_library_move.py new file mode 100644 index 0000000..b838af3 --- /dev/null +++ b/tests/test_library_move.py @@ -0,0 +1,497 @@ +"""Move library (#249 S6b): copy, verify, switch, then retire the old files. + +Each test builds a node-folder library (a real PromptDatabase with prompts, +images and workflow documents, a config.json holding a key, a pooled safety +backup and the pre-migration backup) and a ComfyUI user folder in a temp +directory, then moves it with ``database.library_move``. Failures and crash +windows are injected through the move's ``checkpoint`` hook: an ``Exception`` +is a failure the move cleans up after, ``SimulatedCrash`` (a BaseException) +stops it dead like a killed process. +""" + +import hashlib +import json +import os +import sqlite3 +import sys +import tempfile +import threading +import types +import unittest +from unittest import mock + +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.execution_history import ( # noqa: E402 + notify_history_restore, + subscribe_history_restore, +) +from database.models import PromptModel # noqa: E402 +from database.operations import PromptDatabase # noqa: E402 +from database.review_changes import ( # noqa: E402 + notify_review_change, + subscribe_review_changes, +) +from tests.open_handles import assert_closed, open_handles_under # noqa: E402 +from utils import library_location as ll # noqa: E402 + +KEY = "civitai-secret-key" +POOL = "prompts.db.backup_20261001_101010_000001" +PRE = "prompts.db.pre-workflow-storage_20261002_101010_000002" + + +class SimulatedCrash(BaseException): + """The process died here: nothing after this point runs.""" + + +def _digest(path): + with open(path, "rb") as handle: + return hashlib.sha256(handle.read()).hexdigest() + + +def _rows(path): + """Every table's rows of ``path`` (read-only), to compare libraries.""" + uri = "file:" + path.replace(os.sep, "/") + "?mode=ro" + conn = sqlite3.connect(uri, uri=True) + try: + tables = [ + row[0] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' " + "AND name NOT LIKE 'sqlite_%' ORDER BY name" + ) + ] + return { + table: sorted( + repr(tuple(row)) for row in conn.execute(f'SELECT * FROM "{table}"') + ) + for table in tables + } + finally: + conn.close() + + +def _workflow(seed): + return { + "nodes": [ + {"id": n, "type": f"Node{seed}", "widgets": [seed, n]} for n in range(30) + ] + } + + +class _MoveCase(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.base = os.path.realpath(self._tmp.name) + self.comfy = os.path.join(self.base, "ComfyUI") + self.ext = os.path.join(self.comfy, "custom_nodes", "ComfyUI_PromptManager") + os.makedirs(self.ext) + self.fp = types.SimpleNamespace() + for name in ("output", "input", "temp", "user"): + folder = os.path.join(self.comfy, name) + os.makedirs(folder) + setattr(self.fp, f"get_{name}_directory", lambda f=folder: f) + self.target = os.path.join(self.comfy, "user", "__promptmanager") + self.old_db = os.path.join(self.ext, "prompts.db") + self.new_db = os.path.join(self.target, "prompts.db") + ll.clear_override() + self.db = PromptDatabase(self.old_db) + for index in range(3): + prompt_id = self.db.save_prompt(f"a cat number {index}", category="cats") + self.db.link_image_to_prompt( + str(prompt_id), + os.path.join(self.comfy, "output", f"cat{index}.png"), + {"workflow": _workflow(index % 2), "prompt": {"1": {"seed": index}}}, + ) + self.db.model.close_all() + self.config = os.path.join(self.ext, "config.json") + with open(self.config, "w", encoding="utf-8") as handle: + json.dump( + {"integrations": {"lora_manager": {"civitai_api_key": KEY}}}, handle + ) + for name, stamp in ((POOL, 1_700_000_000), (PRE, 1_700_000_100)): + path = os.path.join(self.ext, name) + source = sqlite3.connect(self.old_db) + copy = sqlite3.connect(path) + try: + source.backup(copy) + finally: + copy.close() + source.close() + os.utime(path, (stamp, stamp)) + self.before = { + "rows": _rows(self.old_db), + "config": _digest(self.config), + "pool": _digest(os.path.join(self.ext, POOL)), + "pre": _digest(os.path.join(self.ext, PRE)), + } + self.config_lock = threading.RLock() + + def tearDown(self): + import gc + + gc.collect() + PromptModel.close_all_instances() + ll.clear_override() + try: + assert_closed(self, self.ext, self.target) + finally: + self._tmp.cleanup() + + def location(self, **kwargs): + return ll.resolve_library( + None, environ={}, extension_root=self.ext, folder_paths=self.fp, **kwargs + ) + + def move(self, **kwargs): + plan = lm.plan_move(self.location()) + kwargs.setdefault("config_lock", self.config_lock) + return lm.move_library(plan, **kwargs) + + def assert_old_library_untouched(self): + self.assertTrue(os.path.isfile(self.old_db)) + self.assertEqual(_rows(self.old_db), self.before["rows"]) + self.assertEqual(_digest(self.config), self.before["config"]) + self.assertEqual(_digest(os.path.join(self.ext, POOL)), self.before["pool"]) + self.assertEqual(_digest(os.path.join(self.ext, PRE)), self.before["pre"]) + self.assertEqual( + os.path.normcase(self.db.model.db_path), os.path.normcase(self.old_db) + ) + self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER) + + def assert_target_clean(self): + leftovers = os.listdir(self.target) if os.path.isdir(self.target) else [] + self.assertEqual(leftovers, []) + + +class TestMove(_MoveCase): + def test_move_copies_verifies_switches_and_retires_the_old_files(self): + result = self.move() + + self.assertEqual(_rows(self.new_db), self.before["rows"]) + self.assertEqual(result.counts["prompts"], 3) + self.assertEqual(result.counts["generated_images"], 3) + self.assertEqual( + result.counts["workflow_documents"], 5 + ) # 2 workflows, 3 graphs + self.assertEqual( + sorted(result.files), sorted(["config.json", POOL, PRE, "prompts.db"]) + ) + self.assertIsNone(result.retire_problem) + # The live database object now reads and writes the new file. + self.assertEqual(self.db.model.db_path, self.new_db) + self.db.save_prompt("written after the move") + self.assertEqual(len(_rows(self.new_db)["prompts"]), 4) + image = self.db.get_prompt_images( + str(self.db.search_prompts(text="number 1")[0]["id"]) + )[0] + self.assertEqual(image["workflow_data"], _workflow(1)) + # config.json, the pooled backup and the pre-migration backup: same bytes. + self.assertEqual( + _digest(os.path.join(self.target, "config.json")), self.before["config"] + ) + self.assertEqual(_digest(os.path.join(self.target, POOL)), self.before["pool"]) + self.assertEqual(_digest(os.path.join(self.target, PRE)), self.before["pre"]) + self.assertEqual( + int(os.path.getmtime(os.path.join(self.target, POOL))), 1_700_000_000 + ) + # The old files are renamed, never deleted; the backups stay. + renamed = sorted(os.listdir(self.ext)) + self.assertIn("prompts.db.moved-to-__promptmanager", renamed) + self.assertIn("config.json.moved-to-__promptmanager", renamed) + self.assertNotIn("prompts.db", renamed) + self.assertNotIn("config.json", renamed) + self.assertFalse(any(name.endswith(("-wal", "-shm")) for name in renamed)) + self.assertEqual( + _rows(os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")), + self.before["rows"], + ) + self.assertEqual(ll.read_marker(self.target)["state"], ll.MARKER_MOVED) + + def test_new_files_are_private(self): + if os.name != "posix": + self.skipTest("POSIX permissions") + self.move() + self.assertEqual(os.stat(self.target).st_mode & 0o777, 0o700) + for name in ("prompts.db", "config.json", POOL, PRE, ll.MARKER_NAME): + mode = os.stat(os.path.join(self.target, name)).st_mode & 0o777 + self.assertEqual(mode, 0o600, name) + + def test_restart_uses_the_new_library_and_never_the_renamed_file(self): + self.move() + self.assertEqual(self.location().db_path, self.new_db) + ll.clear_override() + restarted = self.location() + self.assertEqual(restarted.source, ll.SOURCE_DEFAULT) + self.assertEqual(restarted.db_path, self.new_db) + self.assertEqual( + restarted.config_path, os.path.join(self.target, "config.json") + ) + + def test_every_open_database_object_is_repointed(self): + node = PromptDatabase(self.old_db) + monitor = PromptDatabase(self.old_db) + monitor.get_prompt_by_id(1) + other = PromptDatabase(os.path.join(self.base, "other.db")) + self.move() + self.assertEqual(node.model.db_path, self.new_db) + self.assertEqual(monitor.model.db_path, self.new_db) + self.assertEqual(monitor.get_prompt_by_id(1)["text"], "a cat number 0") + self.assertEqual(other.model.db_path, os.path.join(self.base, "other.db")) + assert_closed( + self, os.path.join(self.ext, "prompts.db.moved-to-__promptmanager") + ) + + def test_change_observers_follow_the_library(self): + reviews, restores = [], [] + unsubscribe_review = subscribe_review_changes( + self.old_db, lambda: reviews.append(1) + ) + unsubscribe_restore = subscribe_history_restore( + self.old_db, lambda: restores.append(1) + ) + self.addCleanup(unsubscribe_review) + self.addCleanup(unsubscribe_restore) + self.move() + self.assertEqual(restores, [1]) # the history readers start a new epoch + notify_review_change(self.new_db) + notify_history_restore(self.new_db) + self.assertEqual(reviews, [1]) + self.assertEqual(restores, [1, 1]) + + def test_files_are_only_renamed_inside_one_folder(self): + real_replace, real_rename = os.replace, os.rename + moves = [] + + def replace(src, dst, *args, **kwargs): + moves.append((src, dst)) + return real_replace(src, dst, *args, **kwargs) + + def rename(src, dst, *args, **kwargs): + moves.append((src, dst)) + return real_rename(src, dst, *args, **kwargs) + + with ( + mock.patch.object(lm.os, "replace", replace), + mock.patch.object(lm.os, "rename", rename), + ): + self.move() + self.assertTrue(moves) + for src, dst in moves: + self.assertEqual(os.path.dirname(src), os.path.dirname(dst), (src, dst)) + + def test_a_write_during_the_copy_is_copied_again(self): + writes = [] + + def checkpoint(step): + if step == "files" and not writes: + writes.append(PromptDatabase(self.old_db).save_prompt("late write")) + + progress = [] + self.move(checkpoint=checkpoint, progress=progress.append) + self.assertIn("late write", repr(_rows(self.new_db)["prompts"])) + self.assertTrue(any("changed" in line for line in progress), progress) + + def test_a_library_that_keeps_changing_is_not_moved(self): + def checkpoint(step): + if step == "files": + PromptDatabase(self.old_db).save_prompt(f"write {os.urandom(4).hex()}") + + with self.assertRaises(lm.MoveFailed) as caught: + self.move(checkpoint=checkpoint, attempts=2) + self.assertEqual(caught.exception.code, "library-busy") + self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER) + self.assert_target_clean() + + def test_old_file_still_open_elsewhere_is_left_in_place(self): + reader = sqlite3.connect(self.old_db) + self.addCleanup(reader.close) + reader.execute("BEGIN") + reader.execute("SELECT COUNT(*) FROM prompts").fetchone() + PromptDatabase(self.old_db).save_prompt("in the wal") + result = self.move() + self.assertEqual(result.retire_problem, "old-in-use") + self.assertTrue(os.path.isfile(self.old_db)) + self.assertEqual(self.db.model.db_path, self.new_db) + reader.rollback() + reader.close() + ll.clear_override() + self.assertEqual(self.location().source, ll.SOURCE_DEFAULT) + + +class TestFailures(_MoveCase): + def _fail_at(self, step): + def checkpoint(name): + if name == step: + raise RuntimeError(f"failure at {step}") + + return checkpoint + + def test_a_failure_at_any_step_leaves_the_old_library_in_use_and_untouched(self): + for step in ("copied", "verified", "files", "locked", "pending", "published"): + with self.subTest(step=step): + with self.assertRaises(RuntimeError): + self.move(checkpoint=self._fail_at(step)) + self.assert_old_library_untouched() + self.assert_target_clean() + assert_closed(self, self.target) + + def test_a_copy_that_does_not_verify_is_refused(self): + def checkpoint(step): + if step == "copied": + (temp,) = [ + name for name in os.listdir(self.target) if name.endswith(".db") + ] + with open(os.path.join(self.target, temp), "r+b") as handle: + handle.truncate(4096) + + with self.assertRaises(lm.MoveFailed) as caught: + self.move(checkpoint=checkpoint) + self.assertEqual(caught.exception.code, "verify-failed") + self.assertNotIn(self.base, caught.exception.message) + self.assert_old_library_untouched() + self.assert_target_clean() + + def test_a_copied_file_that_differs_is_refused(self): + def checkpoint(step): + if step == "files": + for name in os.listdir(self.target): + if name.endswith(".bak"): + with open(os.path.join(self.target, name), "ab") as handle: + handle.write(b"x") + + with self.assertRaises(lm.MoveFailed) as caught: + self.move(checkpoint=checkpoint) + self.assertEqual(caught.exception.code, "copy-failed") + self.assert_old_library_untouched() + self.assert_target_clean() + + def test_not_movable_locations_are_refused(self): + os.makedirs(self.target) + with open(self.new_db, "wb") as handle: + handle.write(b"someone else's library") + with self.assertRaises(lm.MoveFailed) as caught: + lm.plan_move(self.location()) + self.assertEqual(caught.exception.code, "target-has-library") + os.remove(self.new_db) + default = ll.LibraryLocation( + ll.SOURCE_DEFAULT, self.new_db, "c", self.target, self.target + ) + with self.assertRaises(lm.MoveFailed) as caught: + lm.plan_move(default) + self.assertEqual(caught.exception.code, "not-movable") + + def test_target_that_is_the_node_folder_is_refused(self): + link = os.path.join(self.base, "link-to-node-folder") + try: + os.symlink(self.ext, link) + except (OSError, NotImplementedError): + self.skipTest("symlinks are not available") + location = ll.LibraryLocation( + ll.SOURCE_NODE_FOLDER, self.old_db, self.config, self.ext, link + ) + with self.assertRaises(lm.MoveFailed) as caught: + lm.plan_move(location) + self.assertEqual(caught.exception.code, "same-folder") + + def test_disk_check_measures_the_target_volume(self): + measured = [] + + def free(path): + measured.append(path) + return 1024 + + with mock.patch.object(lm, "_free_bytes", free): + plan = lm.plan_move(self.location()) + with self.assertRaises(lm.MoveFailed) as caught: + lm.move_library(plan, config_lock=self.config_lock) + self.assertEqual(caught.exception.code, "insufficient-space") + self.assertTrue(measured) + for path in measured: + self.assertTrue( + os.path.normcase(self.comfy) in os.path.normcase(path) + and "custom_nodes" not in path, + path, + ) + self.assertGreater(plan.required_bytes, os.path.getsize(self.old_db)) + self.assert_old_library_untouched() + + +class TestCrashWindows(_MoveCase): + """A crash between steps never loses the library or picks a stale one.""" + + def _crash_at(self, step): + def checkpoint(name): + if name == step: + raise SimulatedCrash(step) + + return checkpoint + + def _restart(self): + PromptModel.close_all_instances() + ll.clear_override() + return self.location() + + def test_crash_before_the_switch_keeps_the_old_library_and_a_retry_succeeds(self): + for step in ("copied", "files", "pending", "published"): + with self.subTest(step=step): + with self.assertRaises(SimulatedCrash): + self.move(checkpoint=self._crash_at(step)) + restarted = self._restart() + self.assertEqual(restarted.source, ll.SOURCE_NODE_FOLDER) + self.assertEqual(restarted.db_path, self.old_db) + self.assertEqual(_rows(self.old_db), self.before["rows"]) + result = self.move() + self.assertEqual(_rows(self.new_db), self.before["rows"]) + self.assertIsNone(result.retire_problem) + self.assertFalse( + [name for name in os.listdir(self.target) if name.startswith(".pm-move-")] + ) + + def test_crash_after_the_marker_uses_the_new_library(self): + for step in ("marked", "switched"): + with self.subTest(step=step): + with self.assertRaises(SimulatedCrash): + self.move(checkpoint=self._crash_at(step)) + restarted = self._restart() + self.assertEqual(restarted.source, ll.SOURCE_DEFAULT) + self.assertEqual(restarted.db_path, self.new_db) + self.assertEqual(_rows(self.new_db), self.before["rows"]) + # The old file was not renamed yet and is never picked. + self.assertTrue(os.path.isfile(self.old_db)) + # Reset for the next window: back to the node folder only. + for name in os.listdir(self.target): + os.remove(os.path.join(self.target, name)) + + def test_crash_after_the_rename(self): + with self.assertRaises(SimulatedCrash): + self.move(checkpoint=self._crash_at("retired")) + restarted = self._restart() + self.assertEqual(restarted.db_path, self.new_db) + self.assertFalse(os.path.exists(self.old_db)) + + +class TestEnvironmentTarget(_MoveCase): + def test_move_into_the_environment_folder(self): + data = os.path.join(self.base, "pm-data") + environ = {ll.DATA_DIR_ENV: data} + location = ll.resolve_library( + None, environ=environ, extension_root=self.ext, folder_paths=self.fp + ) + self.assertEqual(location.problem, "data-dir-empty") + result = lm.move_library(lm.plan_move(location), config_lock=self.config_lock) + self.assertEqual(result.location.source, ll.SOURCE_ENVIRONMENT) + self.assertEqual(_rows(os.path.join(data, "prompts.db")), self.before["rows"]) + ll.clear_override() + self.assertEqual( + ll.resolve_library( + None, environ=environ, extension_root=self.ext, folder_paths=self.fp + ).source, + ll.SOURCE_ENVIRONMENT, + ) + self.assertEqual(open_handles_under(self.old_db), []) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_library_move_review.py b/tests/test_library_move_review.py new file mode 100644 index 0000000..1e4fa82 --- /dev/null +++ b/tests/test_library_move_review.py @@ -0,0 +1,352 @@ +"""Move library, final review round (#249 S6b PR #318). + +- I1: a switch that fails at any point is undone while both files are still + locked: the copy is removed, the old library stays in use and a restart + keeps it. A follow-up after a complete switch only logs. +- I2: a background write (run counting, the image monitor) that fails + because Move library switched runs once more, on the new library; a write + racing the switch on a real connection (a run count, an image link) lands + in the new library only. +- M2: the copy is verified through a file URI that survives ``#`` and ``%``. +- M6: a restart after a crash between the switch and the rename renames the + old files, never deletes them. +""" + +import contextlib +import json +import os +import sqlite3 +import sys +import threading +import unittest +from unittest import mock + +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.open_handles import assert_closed # noqa: E402 +from utils import library_location as ll # noqa: E402 +from utils import usage_tracking # noqa: E402 +from utils.hashing import generate_prompt_hash # noqa: E402 + +try: + from tests.test_library_move import KEY, SimulatedCrash, _MoveCase, _rows +except ImportError: # discovered with tests/ as the top-level directory + from test_library_move import KEY, SimulatedCrash, _MoveCase, _rows + +WAIT = 30 + + +def _run_count(path, text): + uri = "file:" + path.replace(os.sep, "/") + "?mode=ro" + with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn: + row = conn.execute("SELECT run_count FROM prompts WHERE text=?", (text,)) + return row.fetchone()[0] + + +class TestFailedSwitchIsUndone(_MoveCase): + """I1: whatever part of the switch fails, nothing was moved.""" + + def _assert_not_moved(self, caught): + self.assertEqual(caught.exception.code, "switch-failed") + self.assertIn("Nothing was moved", caught.exception.message) + self.assertIsNone(ll.current_override()) + self.assertEqual(self.db.model.db_path, self.old_db) + # Writes keep landing in the old library, and a restart keeps it. + self.db.save_prompt("written after the failed switch") + self.assertIn("written after the failed switch", repr(_rows(self.old_db))) + self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER) + self.assertNotEqual(ll.read_marker(self.target).get("state"), "moved") + self.assert_target_clean() + self.assertTrue(os.path.isfile(self.config)) + assert_closed(self, self.target) + + def _failing(self, point): + real_switch_to = ll.switch_to + real_repoint = models.PromptModel.repoint_instances + + def switch(old_db, location, detached): + if point == "before": + raise RuntimeError("switch failed before anything") + if point == "resolver": + real_switch_to(location) + raise RuntimeError("switch failed after the resolver") + if point == "repoint": + real_switch_to(location) + detached.extend(real_repoint(old_db, location.db_path)) + raise RuntimeError("switch failed after the models") + lm.switch_process(old_db, location, detached) + raise RuntimeError("switch failed at its very end") + + return switch + + def test_a_switch_failing_at_any_point_is_undone(self): + node = PromptDatabase(self.old_db) + node.get_prompt_by_id(1) + for point in ("before", "resolver", "repoint", "complete"): + with self.subTest(point=point): + with self.assertRaises(lm.MoveFailed) as caught: + self.move(switch=self._failing(point)) + self._assert_not_moved(caught) + self.assertEqual(node.model.db_path, self.old_db) + self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0") + + def test_a_failed_switch_closes_the_old_connections_it_detached(self): + node = PromptDatabase(self.old_db) + for point in ("repoint", "complete"): + with self.subTest(point=point): + before = node.model.get_connection() + with self.assertRaises(lm.MoveFailed): + self.move(switch=self._failing(point)) + with self.assertRaises(sqlite3.ProgrammingError): + before.execute("SELECT 1") + # The model reopens the old library on next use. + self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0") + + def test_a_failing_follow_up_after_the_switch_only_logs(self): + def follow_up(): + raise RuntimeError("the health report could not be dropped") + + with self.assertLogs("prompt_manager.database.library_move", "ERROR"): + result = self.move(after_switch=follow_up) + self.assertIsNone(result.retire_problem) + self.assertEqual(self.db.model.db_path, self.new_db) + self.assertTrue( + os.path.isfile( + os.path.join(self.ext, "prompts.db.moved-to-__promptmanager") + ) + ) + + +class _Flaky: + """A write that fails once as if Move library switched while it waited.""" + + def __init__(self, db, name): + self.db, self.name, self.calls = db, name, 0 + self.real = getattr(db, name) + + def __call__(self, *args, **kwargs): + self.calls += 1 + if self.calls == 1: + self.db.model.close_all() # what the switch does: a new generation + raise sqlite3.ProgrammingError("Cannot operate on a closed database.") + return self.real(*args, **kwargs) + + +class TestBackgroundWritesSurviveTheSwitch(_MoveCase): + """I2: run counting and the image monitor write again on a fresh connection.""" + + def test_run_counting_retries_once_after_a_switch(self): + db = PromptDatabase(self.old_db) + text = "a counted prompt" + db.save_prompt(text, prompt_hash=generate_prompt_hash(text)) + flaky = _Flaky(db, "record_prompt_use") + with mock.patch.object(db, "record_prompt_use", flaky): + pending = usage_tracking.PendingFirstUse() + usage_tracking.count_runs([(generate_prompt_hash(text), text)], db, pending) + self.assertEqual(flaky.calls, 2) + self.assertEqual(_run_count(self.old_db, text), 1) + + def test_a_failure_without_a_switch_is_not_retried(self): + db = PromptDatabase(self.old_db) + calls = [] + + def broken(): + calls.append(1) + raise sqlite3.OperationalError("disk I/O error") + + with self.assertRaises(sqlite3.OperationalError): + db.survive_library_switch(broken) + self.assertEqual(calls, [1]) + + def test_an_image_link_racing_the_switch_lands_in_the_new_library(self): + from utils.image_monitor import ImageGenerationHandler + + self.db.model.close_all() + image = os.path.join(self.comfy, "output", "late.png") + # A short busy timeout so the blocked writer gives up in seconds. + with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2): + writer_db = PromptDatabase(self.old_db) + handler = ImageGenerationHandler(writer_db, mock.MagicMock()) + started = threading.Event() + outcome = {} + + def write(): + conn = writer_db.model.get_connection() + conn.set_trace_callback( + lambda sql: ( + started.set() + if sql.lstrip().upper().startswith(("BEGIN", "INSERT")) + else None + ) + ) + try: + # No metadata: the fallback link, not the resolved pair. + handler.link_image_to_prompt(image, {"id": 1}, {}) + outcome["ok"] = True + except Exception as error: # pragma: no cover - reported below + outcome["error"] = error + finally: + writer_db.model.close() + + thread = threading.Thread(target=write) + + def checkpoint(step): + if step == "locked": + # The old file is write-locked now: start the writer and + # wait until its first write is running (it then waits on us). + thread.start() + self.assertTrue(started.wait(WAIT)) + + self.move(checkpoint=checkpoint) + thread.join(WAIT) + self.assertFalse(thread.is_alive()) + self.assertEqual(outcome, {"ok": True}) + moved = [ + row + for row in _rows(self.new_db)["generated_images"] + if "late.png" in repr(row) + ] + self.assertEqual(len(moved), 1) + renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager") + self.assertNotIn("late.png", repr(_rows(renamed).get("generated_images"))) + + def test_an_image_that_was_not_linked_is_not_reported_as_linked(self): + from utils.image_monitor import ImageGenerationHandler + + db = PromptDatabase(self.old_db) + handler = ImageGenerationHandler(db, mock.MagicMock()) + 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}, {}) + messages = [str(call) for call in log.debug.call_args_list] + self.assertFalse([m for m in messages if "Successfully linked" in m]) + self.assertNotIn("orphan.png", repr(_rows(self.old_db)["generated_images"])) + + def test_a_write_racing_the_switch_lands_in_the_new_library(self): + text = "a racing prompt" + self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text)) + self.db.model.close_all() + # A short busy timeout so the blocked writer gives up in seconds. + with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2): + writer_db = PromptDatabase(self.old_db) + started = threading.Event() + outcome = {} + + def write(): + conn = writer_db.model.get_connection() + conn.set_trace_callback( + lambda sql: started.set() if sql.startswith("UPDATE") else None + ) + try: + usage_tracking.count_runs( + [(generate_prompt_hash(text), text)], + writer_db, + usage_tracking.PendingFirstUse(), + ) + outcome["ok"] = True + except Exception as error: # pragma: no cover - reported below + outcome["error"] = error + finally: + writer_db.model.close() + + thread = threading.Thread(target=write) + + def checkpoint(step): + if step == "locked": + # The old file is write-locked now: start the writer and + # wait until its UPDATE is running (it then waits on us). + thread.start() + self.assertTrue(started.wait(WAIT)) + + self.move(checkpoint=checkpoint) + thread.join(WAIT) + self.assertFalse(thread.is_alive()) + self.assertEqual(outcome, {"ok": True}) + self.assertEqual(_run_count(self.new_db, text), 1) + renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager") + self.assertEqual(_run_count(renamed, text), 0) + + +class TestVerifyUri(unittest.TestCase): + """M2: a copy whose path holds ``#`` and ``%`` verifies.""" + + def test_copy_path_with_uri_characters_verifies(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + folder = os.path.join(tmp, "lib #1 %20") + os.makedirs(folder) + path = os.path.join(folder, "copy.db") + db = PromptDatabase(path) + db.save_prompt("a prompt") + db.model.close_all() + with contextlib.closing(sqlite3.connect(path)) as conn: + conn.execute("PRAGMA journal_mode=DELETE") + counts, documents = lm._describe(conn) + lm._verify_copy(path, counts, documents) + models.PromptModel.close_all_instances() + + +class TestRestartFinishesTheRename(_MoveCase): + """M6: the marker names the old files; a restart renames them.""" + + def _crash_after_switch(self): + def checkpoint(step): + if step == "switched": + raise SimulatedCrash(step) + + with self.assertRaises(SimulatedCrash): + self.move(checkpoint=checkpoint) + models.PromptModel.close_all_instances() + ll.clear_override() + + def test_a_restart_renames_the_old_files_never_deletes(self): + self._crash_after_switch() + self.assertTrue(os.path.isfile(self.old_db)) + restarted = self.location() + self.assertEqual(restarted.db_path, self.new_db) + marker = ll.read_marker(self.target) + self.assertEqual( + (marker["from"], marker["config_from"]), (self.old_db, self.config) + ) + renamed = lm.finish_interrupted_move(restarted) + self.assertEqual( + sorted(renamed), + [ + "config.json.moved-to-__promptmanager", + "prompts.db.moved-to-__promptmanager", + ], + ) + self.assertFalse(os.path.exists(self.old_db)) + moved_config = os.path.join(self.ext, "config.json.moved-to-__promptmanager") + with open(moved_config, encoding="utf-8") as handle: + self.assertEqual( + json.load(handle)["integrations"]["lora_manager"]["civitai_api_key"], + KEY, + ) + self.assertEqual(lm.finish_interrupted_move(self.location()), []) + + def test_nothing_is_renamed_while_the_old_file_is_in_use(self): + self._crash_after_switch() + reader = sqlite3.connect(self.old_db) + try: + reader.execute("BEGIN") + reader.execute("SELECT COUNT(*) FROM prompts").fetchone() + PromptDatabase(self.old_db).save_prompt("in the wal") + self.assertEqual(lm.finish_interrupted_move(self.location()), []) + self.assertTrue(os.path.isfile(self.old_db)) + finally: + reader.rollback() + reader.close() + models.PromptModel.close_all_instances() + + def test_a_node_folder_library_has_nothing_to_finish(self): + self.assertEqual(lm.finish_interrupted_move(self.location()), []) + self.assertTrue(os.path.isfile(self.old_db)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_settings_sources.py b/tests/test_settings_sources.py index e4958bd..b964101 100644 --- a/tests/test_settings_sources.py +++ b/tests/test_settings_sources.py @@ -150,7 +150,15 @@ class TestSourceStatus(SourcesTestCase): ) self.assertEqual( data["sources"]["effective"], - [{"label": "a", "path": "output/a", "root_index": 0, "watch": "manual"}], + [ + { + "label": "a", + "path": "output/a", + "root_index": 0, + "watch": "manual", + "in_comfyui": True, + } + ], ) self.assertFalse(data["sources"]["auto_detected"]) # The legacy keys are unchanged. @@ -183,10 +191,21 @@ class TestSourceStatus(SourcesTestCase): "path": "output", "root_index": 0, "watch": "manual", + "in_comfyui": True, } ], ) + def test_a_source_outside_comfyui_is_not_marked_in_comfyui(self): + # Sources & scans prefixes "ComfyUI/" only to a path relative to the + # ComfyUI base; a folder outside it is shown by its own name. + from py.api import sources + + self.assertTrue(sources.in_comfyui_base(str(self.a))) + self.assertTrue(sources.in_comfyui_base(str(self.comfy_dir))) + self.assertFalse(sources.in_comfyui_base(self.tmpdir)) + self.assertFalse(sources.in_comfyui_base(os.path.join(self.tmpdir, "Renders"))) + async def test_watching_requires_live_observer_with_root(self): GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.a), self._root(self.b)] diff --git a/tests/test_system_api.py b/tests/test_system_api.py index 99891bd..df7ff93 100644 --- a/tests/test_system_api.py +++ b/tests/test_system_api.py @@ -277,6 +277,7 @@ class TestOverviewRoute(SystemAPITestCase): "backup", "needs_review", "sources", + "location", }, ) diff --git a/tests/vendor/comfyui-config.test.mjs b/tests/vendor/comfyui-config.test.mjs index e13e505..dd88e12 100644 --- a/tests/vendor/comfyui-config.test.mjs +++ b/tests/vendor/comfyui-config.test.mjs @@ -25,14 +25,16 @@ test('execution script discovers only the new matrix and preserves all 78 origin assert.match(stock, /Total: 82 tests in 8 files/); // + local Inspect host case (inspect-sender.spec.mjs, 2 projects) assert.doesNotMatch(stock, /execution-history/); assert.doesNotMatch(stock, /execution-outputs/); + assert.doesNotMatch(stock, /library-location/); // it starts its own fresh host (#249 S6b) const executions = list('tools/vendor/playwright.execution.config.mjs'); assert.match(executions, /\[classic-start\]/); assert.match(executions, /\[nodes2-start\]/); - assert.match(executions, /Total: 46 tests in 2 files/); + assert.match(executions, /Total: 48 tests in 3 files/); // + the fresh-host library location case for (const project of ['classic-start', 'nodes2-start']) { const cases = executions.split('\n').filter(line => line.includes(`[${project}]`)); assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15); assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8); + assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1); } assert.doesNotMatch(executions, /prompt-manager.spec/); }); diff --git a/tools/vendor/comfyui.config.mjs b/tools/vendor/comfyui.config.mjs index 0c121a4..bc1e5bf 100644 --- a/tools/vendor/comfyui.config.mjs +++ b/tools/vendor/comfyui.config.mjs @@ -19,7 +19,7 @@ const servers = [{ export default defineConfig({ testDir: '../../tests/comfyui', - testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs'], + testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs'], outputDir: './test-results/comfyui', timeout: 60000, workers: 1, diff --git a/tools/vendor/frontend-coverage-baseline.json b/tools/vendor/frontend-coverage-baseline.json index 9942d6b..49f91a5 100644 --- a/tools/vendor/frontend-coverage-baseline.json +++ b/tools/vendor/frontend-coverage-baseline.json @@ -335,6 +335,12 @@ "functions": 98, "statements": 98 }, + "workspace/library-location.mjs": { + "lines": 94, + "branches": 91, + "functions": 92, + "statements": 94 + }, "workspace/library-state.mjs": { "lines": 98, "branches": 95, diff --git a/tools/vendor/playwright.execution.config.mjs b/tools/vendor/playwright.execution.config.mjs index 0de5792..82f1e05 100644 --- a/tools/vendor/playwright.execution.config.mjs +++ b/tools/vendor/playwright.execution.config.mjs @@ -6,7 +6,7 @@ if (host.href !== 'http://127.0.0.1:8196/') throw new Error('Execution tests req export default defineConfig({ testDir: '../../tests/comfyui', - testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs'], + testMatch: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs'], outputDir: './test-results/executions/cases', timeout: 60000, workers: 1, diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 094d81c..7d1aa76 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -69,8 +69,8 @@ class GalleryDiagnostics: Args: db_path: Path to the SQLite database file to diagnose. - If None, resolves from config.json or defaults to - prompts.db in the extension root. + If None, the library's database (see + ``utils.library_location``). """ if db_path is None: db_path = self._resolve_db_path() @@ -79,18 +79,12 @@ class GalleryDiagnostics: @staticmethod def _resolve_db_path() -> str: - """Resolve database path from config, matching PromptDatabase logic.""" - extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + """The library's database: the same resolver as PromptDatabase.""" try: - from py.config import PromptManagerConfig - - db_path = PromptManagerConfig.DEFAULT_DB_PATH - except Exception: - db_path = "prompts.db" - - if not os.path.isabs(db_path): - db_path = os.path.normpath(os.path.join(extension_root, db_path)) - return db_path + from ..database.operations import resolve_library_location + except ImportError: + from database.operations import resolve_library_location + return resolve_library_location().db_path def run_full_diagnostic(self) -> Dict[str, Any]: """Run a complete diagnostic check. diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 74676c0..b6fbc9b 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -36,6 +36,7 @@ from .metadata_extractor import ComfyUIMetadataExtractor from .logging_config import get_logger from .prompt_graph import resolve_text, run_prompt_nodes from .prompt_pair_resolver import resolve_prompt_pair +from .usage_tracking import survive_switch try: from ..database.pair_ingestion import ingest_resolved_pair @@ -587,24 +588,38 @@ class ImageGenerationHandler(FileSystemEventHandler): """ try: resolution = resolve_prompt_pair(metadata) - if resolution["state"] == "resolved": - positive = resolution["positive"] - expected = self.db_manager.get_prompt_by_hash( - generate_prompt_hash(positive.get("storage_text", positive["text"])) - ) - if expected and expected["id"] == prompt_context["id"]: - ingest_resolved_pair( - self.db_manager, image_path, metadata, resolution + + def link(): + # One write; run again on the moved library when a Move + # library switched while it waited (#249 S6b). + if resolution["state"] == "resolved": + positive = resolution["positive"] + expected = self.db_manager.get_prompt_by_hash( + generate_prompt_hash( + positive.get("storage_text", positive["text"]) + ) ) - return - image_id = self.db_manager.link_image_to_prompt( - prompt_id=prompt_context["id"], image_path=image_path, metadata=metadata - ) - fallback_note = " (fallback)" if prompt_context.get("fallback") else "" - self.logger.debug( - f"Successfully linked image {image_id} to prompt " - f"{prompt_context['id']}{fallback_note}" - ) + if expected and expected["id"] == prompt_context["id"]: + ingest_resolved_pair( + self.db_manager, image_path, metadata, resolution + ) + return None + # The raising form: a failed write reaches survive_switch. + return self.db_manager.link_image_to_prompt_or_raise( + prompt_id=prompt_context["id"], + image_path=image_path, + metadata=metadata, + ) + + image_id = survive_switch(self.db_manager, link) + if image_id is None: + return + if image_id: # 0: nothing new was linked (logged by the database) + fallback_note = " (fallback)" if prompt_context.get("fallback") else "" + self.logger.debug( + f"Successfully linked image {image_id} to prompt " + f"{prompt_context['id']}{fallback_note}" + ) except Exception as e: self.logger.error(f"Failed to link image to prompt: {e}") return diff --git a/utils/library_location.py b/utils/library_location.py new file mode 100644 index 0000000..4a69b6a --- /dev/null +++ b/utils/library_location.py @@ -0,0 +1,480 @@ +"""Where the library lives (#249 S6b): the one resolver for every data path. + +The library is ``prompts.db``, ``config.json`` (it can hold the CivitAI key), +the safety backups beside the database, the logs and the LoRA image cache. +Inside the node folder an uninstall, a reinstall or a folder-replacing update +deletes it, so new installs keep it in ComfyUI's system user folder +``user/__promptmanager`` (``folder_paths.get_system_user_directory``), which +ComfyUI's ``/userdata`` HTTP API refuses to serve. + +Resolution order (``resolve_library``): + +1. ``PROMPT_MANAGER_DATA_DIR``: an absolute folder outside every folder + ComfyUI or Prompt Manager serves. It is used when it holds a library or + when no node-folder library exists; an empty one never replaces a + node-folder library (that one stays in use, and Move library targets the + environment folder). An invalid value is ignored and reported. +2. An explicit ``database.default_path`` in config.json, honoured as before + (a relative path is relative to the node folder). ``PROMPT_MANAGER_CONFIG_PATH`` + still names config.json itself. +3. The default data folder once a move finished there (its marker says + ``moved``), even before the old file was renamed. +4. A node-folder ``prompts.db`` that was not moved: it stays in use until the + user moves it. +5. The default data folder (new installs). Outside ComfyUI (no + ``folder_paths``) the node folder, as before. + +Nothing here moves, copies or creates data: it only reads the file system. +After a live move the process uses ``switch_to``'s location until restart. +This module imports only the standard library, so the logger can use it. +""" + +import json +import os +import threading +from dataclasses import dataclass +from typing import Iterable, List, Mapping, Optional, Tuple + +DATA_DIR_ENV = "PROMPT_MANAGER_DATA_DIR" +CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH" +EXTRA_GALLERY_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" +SYSTEM_USER_NAME = "promptmanager" +SYSTEM_USER_PREFIX = "__" +DATABASE_NAME = "prompts.db" +CONFIG_NAME = "config.json" +MARKER_NAME = "library-location.json" +# config.json has always written its default; that value is not a choice. +DEFAULT_DB_SETTING = "prompts.db" + +SOURCE_ENVIRONMENT = "environment" +SOURCE_CONFIG = "config" +SOURCE_NODE_FOLDER = "node_folder" +SOURCE_DEFAULT = "default" +SOURCES = (SOURCE_ENVIRONMENT, SOURCE_CONFIG, SOURCE_NODE_FOLDER, SOURCE_DEFAULT) + +MARKER_PENDING = "pending" +MARKER_MOVED = "moved" + +EXTENSION_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +_UNSET = object() +_override_lock = threading.Lock() +_override: Optional["LibraryLocation"] = None + + +@dataclass(frozen=True) +class LibraryLocation: + """Where the library is and why. + + ``target_dir`` is where Move library would put it (None outside ComfyUI) + and ``target_source`` why (the environment folder or the default one); + ``problem`` is a code for something the user should know (an invalid or + empty ``PROMPT_MANAGER_DATA_DIR``, a move that did not finish). + ``config_explicit``: config.json is named by ``PROMPT_MANAGER_CONFIG_PATH`` + (it is the user's and never moves). + """ + + source: str + db_path: str + config_path: str + data_dir: str + target_dir: Optional[str] = None + problem: Optional[str] = None + target_source: Optional[str] = None + config_explicit: bool = False + + @property + def movable(self) -> bool: + return self.source == SOURCE_NODE_FOLDER and self.target_dir is not None + + +def import_folder_paths(): + """ComfyUI's ``folder_paths`` module, or None outside ComfyUI.""" + try: + import folder_paths + except ImportError: + return None + return folder_paths + + +def _folder_paths(value): + return import_folder_paths() if value is _UNSET else value + + +def _call(module, name, *args) -> Optional[str]: + getter = getattr(module, name, None) + if not callable(getter): + return None + try: + value = getter(*args) + except Exception: # an older or broken host: treat the folder as unknown + return None + return str(value) if value else None + + +def default_data_dir(folder_paths=_UNSET) -> Optional[str]: + """``user/__promptmanager``, or None outside ComfyUI. + + ComfyUI builds without ``get_system_user_directory`` get the same folder + joined by hand. + """ + module = _folder_paths(folder_paths) + if module is None: + return None + folder = _call(module, "get_system_user_directory", SYSTEM_USER_NAME) + if folder: + return os.path.abspath(folder) + user = _call(module, "get_user_directory") + if not user: + return None + return os.path.abspath(os.path.join(user, SYSTEM_USER_PREFIX + SYSTEM_USER_NAME)) + + +def _identity(path: str) -> Optional[Tuple[int, int]]: + try: + stat = os.stat(path) + except (OSError, ValueError): + return None + return (stat.st_dev, stat.st_ino) + + +def same_folder(first: str, second: str) -> bool: + """Whether two spellings name one existing folder (links and case folded).""" + one, two = _identity(first), _identity(second) + if one is None or two is None: + return False + if one[1] == 0: # no file id on this volume: compare the spelling + return _spelling(first) == _spelling(second) + return one == two + + +def _spelling(path: str) -> str: + return os.path.normcase(os.path.realpath(path)) + + +def _existing_ancestors(path: str) -> List[str]: + """``path`` (resolved) and its parents that exist, nearest first.""" + current = os.path.realpath(path) + found = [] + while True: + if os.path.isdir(current): + found.append(current) + parent = os.path.dirname(current) + if parent == current: + return found + current = parent + + +def _inside_any(path: str, folders: Iterable[str]) -> bool: + """Whether ``path`` is one of ``folders`` or below one. + + Compared by spelling (realpath + normcase) and by file identity of every + existing ancestor, so symlinks and case-insensitive volumes (Windows, + macOS) cannot hide a served folder. + """ + folders = [folder for folder in folders if folder] + if not folders: + return False + spelling = _spelling(path) + sep = os.path.sep + for folder in folders: + root = _spelling(folder) + if spelling == root or spelling.startswith(root.rstrip(sep) + sep): + return True + identities = {_identity(folder) for folder in folders} + identities.discard(None) + identities = {item for item in identities if item[1] != 0} + return any(_identity(a) in identities for a in _existing_ancestors(path)) + + +def _served_folders(module, environ, gallery_roots) -> List[str]: + """Folders whose files ComfyUI or Prompt Manager hand out over HTTP.""" + folders = [] + if module is not None: + for name in ( + "get_output_directory", + "get_input_directory", + "get_temp_directory", + ): + folder = _call(module, name) + if folder: + folders.append(folder) + extra = (environ or {}).get(EXTRA_GALLERY_ROOTS_ENV, "") + folders.extend(part for part in extra.split(os.pathsep) if part.strip()) + folders.extend(str(root) for root in gallery_roots or () if root) + return folders + + +def _in_public_user_folder(path: str, module) -> bool: + """Whether ``path`` lies in a ``/userdata``-served user folder (not ``__``).""" + user = _call(module, "get_user_directory") if module is not None else None + if not user: + return False + if not _inside_any(path, [user]) or same_folder(path, user): + return False + real_user = _spelling(user) + relative = os.path.relpath(_spelling(path), real_user) + first = relative.split(os.sep, 1)[0] + if first.startswith(os.pardir): + # Reached through a link or another spelling: judge by identity. + for ancestor in _existing_ancestors(path): + if same_folder(os.path.dirname(ancestor), user): + first = os.path.basename(ancestor) + break + else: + return False + return not first.startswith(SYSTEM_USER_PREFIX) + + +def data_dir_problem( + path, *, folder_paths=_UNSET, environ=None, gallery_roots=() +) -> Optional[str]: + """Why ``path`` cannot hold the library, as a code, or None. + + ``data-dir-not-absolute``: empty or relative. ``data-dir-served``: inside + ComfyUI's output, input or temp folder, a public user folder, an extra + gallery root or a gallery source (their files are served over HTTP). + """ + if not isinstance(path, str) or not path.strip() or not os.path.isabs(path): + return "data-dir-not-absolute" + module = _folder_paths(folder_paths) + if _inside_any(path, _served_folders(module, environ, gallery_roots)): + return "data-dir-served" + if _in_public_user_folder(path, module): + return "data-dir-served" + return None + + +def read_marker(directory: Optional[str]) -> dict: + """The move marker in ``directory``, or ``{}`` when absent or unreadable.""" + if not directory: + return {} + try: + with open(os.path.join(directory, MARKER_NAME), "r", encoding="utf-8") as fh: + data = json.load(fh) + except (OSError, ValueError): + return {} + if not isinstance(data, dict) or data.get("state") not in ( + MARKER_PENDING, + MARKER_MOVED, + ): + return {} + return data + + +def _marker_state(directory) -> Optional[str]: + return read_marker(directory).get("state") + + +def switch_to(location: LibraryLocation) -> None: + """Use ``location`` for the rest of this process (after a live move).""" + global _override + with _override_lock: + _override = location + + +def clear_override() -> None: + """Forget a live move's location (tests).""" + global _override + with _override_lock: + _override = None + + +def current_override() -> Optional[LibraryLocation]: + with _override_lock: + return _override + + +@dataclass(frozen=True) +class _Decision: + data_dir: str + source: str + target_dir: Optional[str] + problem: Optional[str] + db_dir: Optional[str] = None # when the database is not in data_dir + target_source: Optional[str] = None + + +def _configured_sources(module, environ, extension_root, default, raw_env): + """The image sources every config.json this resolver might read lists. + + ``config_location`` runs before any config.json is loaded, so it cannot be + handed the configured sources; both it and ``resolve_library`` read them + here, from each candidate file (the explicit one, the environment + folder's, the node folder's and the default folder's), so the two always + judge ``PROMPT_MANAGER_DATA_DIR`` alike. Relative sources are taken from + ComfyUI's base folder, as the gallery resolves them. + """ + candidates = [environ.get(CONFIG_PATH_ENV, "").strip()] + if raw_env and os.path.isabs(raw_env): + candidates.append(os.path.join(raw_env, CONFIG_NAME)) + candidates.append(os.path.join(extension_root, CONFIG_NAME)) + if default: + candidates.append(os.path.join(default, CONFIG_NAME)) + base = getattr(module, "base_path", None) if module is not None else None + sources = [] + for candidate in candidates: + if not candidate or not os.path.isfile(candidate): + continue + try: + with open(candidate, "r", encoding="utf-8") as handle: + data = json.load(handle) + directories = data["gallery"]["monitoring"]["directories"] + except (OSError, ValueError, KeyError, TypeError): + continue + if not isinstance(directories, list): + continue + for directory in directories: + if not isinstance(directory, str) or not directory.strip(): + continue + directory = directory.strip() + if not os.path.isabs(directory): + if not base: + continue + directory = os.path.join(str(base), directory) + sources.append(directory) + return sources + + +def _decide(environ, extension_root, folder_paths, gallery_roots) -> _Decision: + """Which folder holds the library (before an explicit database path).""" + module = _folder_paths(folder_paths) + environ = os.environ if environ is None else environ + legacy_db = os.path.isfile(os.path.join(extension_root, DATABASE_NAME)) + legacy = legacy_db or os.path.isfile(os.path.join(extension_root, CONFIG_NAME)) + default = default_data_dir(module) + problem = None + env_dir = None + raw = environ.get(DATA_DIR_ENV, "") + if raw.strip(): + sources = list(gallery_roots or ()) + _configured_sources( + module, environ, extension_root, default, raw.strip() + ) + problem = data_dir_problem( + raw.strip(), + folder_paths=module, + environ=environ, + gallery_roots=sources, + ) + if problem is None: + env_dir = os.path.abspath(raw.strip()) + target = env_dir or default + target_source = SOURCE_ENVIRONMENT if env_dir else SOURCE_DEFAULT + + def decision(data_dir, source, problem=None, db_dir=None): + return _Decision( + data_dir, source, target, problem, db_dir, target_source if target else None + ) + + if env_dir is not None: + state = _marker_state(env_dir) + holds = os.path.isfile(os.path.join(env_dir, DATABASE_NAME)) + if state != MARKER_PENDING and (holds or not legacy_db): + return decision(env_dir, SOURCE_ENVIRONMENT) + problem = "move-unfinished" if state == MARKER_PENDING else "data-dir-empty" + return decision(extension_root, SOURCE_NODE_FOLDER, problem) + if default is None: + return decision(extension_root, SOURCE_NODE_FOLDER, problem) + if _marker_state(default) == MARKER_MOVED and os.path.isfile( + os.path.join(default, DATABASE_NAME) + ): + return decision(default, SOURCE_DEFAULT, problem) + if legacy: + if _marker_state(default) == MARKER_PENDING: + problem = problem or "move-unfinished" + if legacy_db: + return decision(extension_root, SOURCE_NODE_FOLDER, problem) + # Only a node-folder config.json: keep reading it, but a new database + # never goes into the node folder. + return decision(extension_root, SOURCE_DEFAULT, problem, default) + return decision(default, SOURCE_DEFAULT, problem) + + +def config_location( + environ=None, extension_root=None, folder_paths=_UNSET, gallery_roots=() +) -> Tuple[str, str]: + """``(config.json path, data folder)``; needs no config value (no cycle).""" + override = current_override() + if override is not None: + return override.config_path, override.data_dir + environ = os.environ if environ is None else environ + root = extension_root or EXTENSION_ROOT + decision = _decide(environ, root, folder_paths, gallery_roots) + explicit = environ.get(CONFIG_PATH_ENV, "").strip() + config_path = explicit or os.path.join(decision.data_dir, CONFIG_NAME) + return config_path, decision.data_dir + + +def is_explicit_db_setting(value) -> bool: + """Whether config.json's ``database.default_path`` names a chosen file.""" + return isinstance(value, str) and value.strip() not in ("", DEFAULT_DB_SETTING) + + +def resolve_library( + default_path=None, + *, + environ: Optional[Mapping[str, str]] = None, + extension_root: Optional[str] = None, + folder_paths=_UNSET, + gallery_roots=(), +) -> LibraryLocation: + """The library's location now (the module docstring has the order).""" + override = current_override() + if override is not None: + return override + environ = os.environ if environ is None else environ + root = extension_root or EXTENSION_ROOT + decision = _decide(environ, root, folder_paths, gallery_roots) + explicit_config = environ.get(CONFIG_PATH_ENV, "").strip() + config_path = explicit_config or os.path.join(decision.data_dir, CONFIG_NAME) + common = { + "target_source": decision.target_source, + "config_explicit": bool(explicit_config), + } + if decision.source != SOURCE_ENVIRONMENT and is_explicit_db_setting(default_path): + db_path = default_path.strip() + if not os.path.isabs(db_path): + db_path = os.path.normpath(os.path.join(root, db_path)) + return LibraryLocation( + SOURCE_CONFIG, + db_path, + config_path, + decision.data_dir, + decision.target_dir, + decision.problem, + **common, + ) + return LibraryLocation( + decision.source, + os.path.join(decision.db_dir or decision.data_dir, DATABASE_NAME), + config_path, + decision.data_dir, + decision.target_dir, + decision.problem, + **common, + ) + + +def ensure_private_dir(path: str) -> str: + """Create ``path`` (and parents) if missing; a new folder is owner-only. + + An existing folder is left as it is (its permissions are the user's). + """ + if not os.path.isdir(path): + os.makedirs(path, mode=0o700, exist_ok=True) + if os.name == "posix": + try: + os.chmod(path, 0o700) + except OSError: + pass + return path + + +def logs_dir(**kwargs) -> str: + """The folder for Prompt Manager's log files.""" + return os.path.join(config_location(**kwargs)[1], "logs") + + +def cache_dir(name: str, **kwargs) -> str: + """A re-creatable cache folder, ``/data/``.""" + return os.path.join(config_location(**kwargs)[1], "data", name) diff --git a/utils/logging_config.py b/utils/logging_config.py index d9a6286..bb147a2 100644 --- a/utils/logging_config.py +++ b/utils/logging_config.py @@ -18,6 +18,8 @@ from datetime import datetime from typing import Dict, Any, Optional, List from pathlib import Path +from . import library_location + class _OwnedRotatingFileHandler(logging.handlers.RotatingFileHandler): """Keep a retired owned sink closed when an in-flight record resumes.""" @@ -65,8 +67,12 @@ class PromptManagerLogger: return self._initialized = True - self.log_dir = Path(__file__).parent.parent / "logs" - self.log_dir.mkdir(exist_ok=True) + # The library's data folder (#249 S6b): the node folder for a library + # that was not moved, else ComfyUI's user/__promptmanager. + self.log_dir = Path( + library_location.logs_dir(extension_root=str(Path(__file__).parent.parent)) + ) + library_location.ensure_private_dir(str(self.log_dir)) # Configuration self.config = { diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 4f0b094..0365e25 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -187,6 +187,19 @@ def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]: return row["id"] if row else None +def survive_switch(db: Any, write: Callable[[], Any]) -> Any: + """``write()`` through ``db.survive_library_switch`` when the database has it. + + A write that was waiting on the old file when Move library switched runs + again on the new library instead of being lost (#249 S6b). Stand-in + databases without the method (tests) just run ``write()``. + """ + survive = getattr(type(db), "survive_library_switch", None) + if survive is None: + return write() + return survive(db, write) + + def count_pair_runs( runs: List[Tuple[str, str, str]], db: Any, @@ -194,20 +207,28 @@ def count_pair_runs( ) -> None: """Count existing recipes at dequeue and defer new recipes until first save.""" for key, positive, negative in runs: - pair_id = find_pair_id(db, positive, negative) - if pair_id: - db.record_prompt_pair_use(pair_id) - else: + + def count(positive=positive, negative=negative): + pair_id = find_pair_id(db, positive, negative) + if pair_id: + db.record_prompt_pair_use(pair_id) + return bool(pair_id) + + if not survive_switch(db, count): pending.add(key) def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None: """Record one use per prompt in the database; remember the rest for first save.""" for prompt_hash, _text in runs: - existing = db.get_prompt_by_hash(prompt_hash) - if existing: - db.record_prompt_use(existing["id"]) - else: + + def count(prompt_hash=prompt_hash): + existing = db.get_prompt_by_hash(prompt_hash) + if existing: + db.record_prompt_use(existing["id"]) + return bool(existing) + + if not survive_switch(db, count): pending.add(prompt_hash) diff --git a/web/admin.html b/web/admin.html index cb8a6e0..bbd32c6 100644 --- a/web/admin.html +++ b/web/admin.html @@ -79,6 +79,8 @@ + + diff --git a/web/js/admin.js b/web/js/admin.js index e74d347..217aa03 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -297,6 +297,7 @@ this.loadTags(), this.loadRecentPrompts(), this.loadSettings(), + this.loadLibraryLocation(), ]); } catch (error) { this.showNotification("Failed to load initial data", "error"); @@ -344,6 +345,29 @@ this.listRequest = ApiClient.latestOnly({ lifetime: this.pageLifetime.signal }); } + /** #249 S6b: one line under the statistics saying where the library lives and why. */ + async loadLibraryLocation() { + const { signal } = this.pageLifetime; + if (signal.aborted) return; + const line = document.getElementById("libraryLocation"); + const reasons = { + default: "ComfyUI's user folder", + environment: "set by PROMPT_MANAGER_DATA_DIR", + config: "set in config.json", + node_folder: "the node folder, where an update or reinstall can delete it; move it from the workspace's System page", + }; + try { + const response = await ApiClient.fetchUrl("/prompt_manager/system/library", { signal }); + const data = response.ok ? await response.json() : null; + const location = data && data.location; + if (!line || !location || typeof location.path !== "string") return; + line.textContent = `Library: ${location.path} (${reasons[location.source] || location.source})`; + line.hidden = false; + } catch (error) { + if (!signal.aborted) console.error("Library location error:", error); + } + } + async loadStatistics() { // A leaving page starts no read; one in flight is aborted with the page and not reported. const { signal } = this.pageLifetime; diff --git a/web/js/workspace/index.mjs b/web/js/workspace/index.mjs index 219b026..c730f24 100644 --- a/web/js/workspace/index.mjs +++ b/web/js/workspace/index.mjs @@ -18,6 +18,7 @@ import { mountInspect } from './inspect.mjs'; import { mountReview } from './review.mjs'; import { createPromptSender } from './prompt-sender.mjs'; import { mountSystem } from './system.mjs'; +import { mountLibraryBanner } from './library-location.mjs'; const { createApiClient, resolveUrl } = window.ApiClient; @@ -75,6 +76,8 @@ function mountWorkspace() { const activity = createActivityStore({ api, scope, connect: url => new EventSource(resolveUrl(url)) }); mountActivityTopbar({ store: activity, scope }); activity.load(); + // D129: a node-folder library is offered a move (or "Keep it here") above every page. + mountLibraryBanner(document.getElementById('library-banner'), { api, activity, scope }); const sidebar = document.querySelector('aside.sidebar'); // The route key a nav link carries in `data-route`: a saved view's is `saved/`. const routeKey = route => route.id === 'saved' ? `saved/${route.savedId}` : route.id; diff --git a/web/js/workspace/library-location.mjs b/web/js/workspace/library-location.mjs new file mode 100644 index 0000000..447d5ae --- /dev/null +++ b/web/js/workspace/library-location.mjs @@ -0,0 +1,255 @@ +/** + * Where the library lives (#249 System S6b, D129): the workspace banner, System โ€บ Overview's "Library location" row, + * the Move library review and its start. + * + * - The server says where the database is and why (`GET /prompt_manager/system/library` and the Overview's + * `location`): `default` (ComfyUI's user folder), `environment` (PROMPT_MANAGER_DATA_DIR), `config` + * (config.json's database path) or `node_folder`, where an update or reinstall can delete it. + * - D129: while a node-folder library is not kept there, the shell shows the concept's `.notice` (`h2` + `p` + + * `.row-actions`) above every page: "Keep it here" saves the choice on the server (the banner stops; the Overview row + * stays), "Move library" opens the review. The concept has no such notice; its markup and classes are the Sources + * notice's (app.js 680). + * - The review is the shared record-changing review (`openReviewDialog`): From, To, Size, what moves with it and what + * happens to the old files. Move library starts the server's exclusive "Move library" job, tracks it in Activity and + * opens its view; a refusal (another job running, no room, a library already there) says why and nothing starts. + * - When a Move library job ends, the banner and the Overview read the location again. + */ +import { createButton } from '../shared/toolkit/core.js'; +import { busyToast } from './activity-scan.mjs'; +import { closeWorkspaceDialog, openContentDialog, openReviewDialog } from './workspace-dialogs.mjs'; +import { navigateTo } from './state.mjs'; +import { showWorkspaceToast } from './toast.mjs'; + +export const LIBRARY_URL = '/prompt_manager/system/library'; +export const MOVE_URL = `${LIBRARY_URL}/move`; +export const KEEP_URL = `${LIBRARY_URL}/keep`; +export const MOVE_KIND = 'library_move'; + +export const BANNER_TITLE = 'Your library is inside the node folder'; +const BANNER_LEAD = 'Your library is stored inside the node folder, where an update or reinstall can delete it.'; +/** Where Move library puts it, by the server's `target_source`. */ +export const TARGET_TEXT = Object.freeze({ + default: 'ComfyUI\'s user folder', + environment: 'the folder PROMPT_MANAGER_DATA_DIR names', +}); +/** The notice's sentence: it names the actual target. */ +export function bannerText(location) { + const target = TARGET_TEXT[location?.target_source] ?? TARGET_TEXT.default; + return `${BANNER_LEAD} Move it to ${target}: it is copied and verified first, and nothing is deleted.`; +} +export const BANNER_TEXT = bannerText({ target_source: 'default' }); +export const KEEP_LABEL = 'Keep it here'; +export const MOVE_LABEL = 'Move library'; +export const ROW_TITLE = 'Library location'; +export const START_FAILED = 'Move library could not start. Try again.'; +export const KEEP_FAILED = 'Your choice could not be saved. Try again.'; + +/** Why the library is where it is, by `source`. */ +export const SOURCE_TEXT = Object.freeze({ + default: 'ComfyUI\'s user folder, which updates and reinstalls of the node leave alone', + environment: 'The folder PROMPT_MANAGER_DATA_DIR names', + config: 'The database path set in config.json', + node_folder: 'The node folder, where an update or reinstall can delete it', +}); + +/** A `problem` code as a sentence. */ +export const PROBLEM_TEXT = Object.freeze({ + 'data-dir-not-absolute': 'PROMPT_MANAGER_DATA_DIR is not an absolute folder, so it was ignored.', + 'data-dir-served': 'PROMPT_MANAGER_DATA_DIR names a folder whose files are served over the web, so it was ignored.', + 'data-dir-empty': 'PROMPT_MANAGER_DATA_DIR has no library yet; Move library puts this one there.', + 'move-unfinished': 'An earlier move did not finish; this library is still the one in use.', + 'inside-node-folder': 'It is inside the node folder, where an update or reinstall deletes it, config.json and its ' + + 'CivitAI key included. Set database.default_path or PROMPT_MANAGER_DATA_DIR to a folder outside it.', +}); + +const number = value => Number(value).toLocaleString('en-US'); + +/** "1.6 GB" or "120.0 MB" or "512 KB". */ +export function sizeText(bytes) { + if (!Number.isFinite(bytes) || bytes < 0) return 'Unknown'; + if (bytes >= 2 ** 30) return `${(bytes / 2 ** 30).toLocaleString('en-US', { maximumFractionDigits: 1, minimumFractionDigits: 1 })} GB`; + if (bytes >= 2 ** 20) return `${(bytes / 2 ** 20).toLocaleString('en-US', { maximumFractionDigits: 1, minimumFractionDigits: 1 })} MB`; + return `${number(Math.ceil(bytes / 1024))} KB`; +} + +/** Whether `location` is a well-formed server answer. */ +export function isLocation(location) { + return Boolean(location) && typeof location === 'object' && typeof location.path === 'string' + && Object.hasOwn(SOURCE_TEXT, location.source) && typeof location.movable === 'boolean'; +} + +/** The Overview row's description: the path, why it is there, the user's choice and any problem. */ +export function locationDescription(location) { + const parts = [`${location.path} ยท ${SOURCE_TEXT[location.source]}.`]; + if (location.movable && location.keep) parts.push('You chose to keep it here.'); + const problem = PROBLEM_TEXT[location.problem]; + if (problem) parts.push(problem); + return parts.join(' '); +} + +/** The review's `[term, value]` rows. */ +export function reviewGrid(location) { + return [ + ['From', location.path], + ['To', `${location.target}/prompts.db`], + ['Size', sizeText(location.size_bytes)], + ['Moves with it', location.config_moves === false + ? 'The database backups (config.json stays where PROMPT_MANAGER_CONFIG_PATH names it)' + : 'config.json and the database backups'], + ['Old files', 'Renamed *.moved-to-โ€ฆ, never deleted'], + ]; +} + +export const REVIEW_SENTENCE = 'The copy is verified before Prompt Manager switches to it. No other job can run until it ' + + 'finishes, and closing this page does not stop it.'; + +function el(tag, text, className) { + const node = document.createElement(tag); + if (text !== undefined) node.textContent = text; + if (className) node.className = className; + return node; +} + +/** + * Start the Move library job, follow it in Activity and open its view. A busy server toasts the D57 link to the running + * job; any other refusal says the server's reason. Nothing is said once `scope` ended. + */ +export async function startLibraryMove({ api, activity, scope, notify = showWorkspaceToast, navigate = navigateTo }) { + try { + const { job } = await api.post(MOVE_URL, {}, { signal: scope.signal }); + if (scope.signal.aborted) return null; + const tracked = activity.track(job); + navigate(`#/activity/${tracked.id}`); + return tracked; + } catch (error) { + if (scope.signal.aborted) return null; + if (error?.status === 409 && error?.code === 'busy') { + await activity.load(); + if (!scope.signal.aborted) notify(busyToast(activity.running()[0] ?? null)); + return null; + } + notify(typeof error?.code === 'string' && error.message ? error.message : START_FAILED); + return null; + } +} + +/** The Move library review over `location`; Move library starts the job. */ +export function openMoveReview({ api, activity, scope, location, trigger = null, notify = showWorkspaceToast, + navigate = navigateTo }) { + return openReviewDialog({ scope, trigger, title: 'Move library', grid: reviewGrid(location), sentence: REVIEW_SENTENCE, + cancelLabel: 'Cancel', primary: { label: MOVE_LABEL }, + onStart: () => { startLibraryMove({ api, activity, scope, notify, navigate }); } }); +} + +/** Save "Keep it here" (`keep` true) or take it back; resolves the new location, or null when it was not saved. */ +export async function keepLibrary({ api, scope, keep = true, notify = showWorkspaceToast }) { + try { + const { location } = await api.post(KEEP_URL, { keep }, { signal: scope.signal }); + return isLocation(location) ? location : null; + } catch (error) { + if (!scope.signal.aborted) notify(KEEP_FAILED); + return null; + } +} + +/** The details of a library that does not move: where it is, why, and what lives beside it. */ +export function openLocationDetails({ scope, location, trigger = null }) { + const list = el('dl', undefined, 'meta-grid'); + for (const [term, value] of [['Database', location.path], ['Why here', SOURCE_TEXT[location.source]], + ['Beside it', 'config.json, the database backups and the logs'], ['Size', sizeText(location.size_bytes)]]) { + const entry = el('div'); + entry.append(el('dt', term), el('dd', value)); + list.append(entry); + } + const close = createButton('Close'); + scope.listen(close, 'click', () => closeWorkspaceDialog()); + const actions = el('div', undefined, 'dialog-actions'); + actions.append(close); + const nodes = [list]; + const problem = PROBLEM_TEXT[location.problem]; + if (problem) nodes.push(el('p', problem, 'small')); + nodes.push(actions); + openContentDialog({ scope, title: ROW_TITLE, trigger, nodes }); +} + +/** D129: the Overview's "Library location" `.operation` row; Move library when it can move, else Details. */ +export function locationRow(location, { api, activity, scope, notify = showWorkspaceToast, navigate = navigateTo }) { + const row = el('div', undefined, 'operation'); + const text = el('div'); + text.append(el('h3', ROW_TITLE), el('p', locationDescription(location))); + const actions = el('div', undefined, 'row-actions'); + const action = createButton(location.movable ? MOVE_LABEL : 'Details'); + scope.listen(action, 'click', () => { + if (location.movable) openMoveReview({ api, activity, scope, location, trigger: action, notify, navigate }); + else openLocationDetails({ scope, location, trigger: action }); + }); + actions.append(action); + row.append(text, actions); + return row; +} + +/** D129: the concept's `.notice` for a node-folder library: Keep it here and Move library. */ +export function libraryNotice({ scope, location = null, onKeep, onMove }) { + const notice = el('div', undefined, 'notice'); + const text = el('div'); + text.append(el('h2', BANNER_TITLE), el('p', bannerText(location))); + const keep = createButton(KEEP_LABEL); + const move = createButton(MOVE_LABEL, 'ca-primary'); + scope.listen(keep, 'click', () => onKeep(keep)); + scope.listen(move, 'click', () => onMove(move)); + const actions = el('div', undefined, 'row-actions'); + actions.append(keep, move); + notice.append(text, actions); + return notice; +} + +/** Calls `onEnded()` whenever a Move library job this page sees stops running. Returns the unsubscribe. */ +export function onMoveEnded(activity, onEnded) { + const running = new Set(); + return activity.subscribe(() => { + let ended = false; + for (const job of activity.jobs()) { + if (job.kind !== MOVE_KIND) continue; + if (job.state === 'running') running.add(job.id); + else if (running.delete(job.id)) ended = true; + } + if (ended) onEnded(); + }); +} + +/** + * Mount the banner in `container` (hidden until the server offers the move). It reads the location once, again when a + * Move library job ends, and hides for good this page once the user keeps the library where it is. + */ +export function mountLibraryBanner(container, { api, activity, scope, notify = showWorkspaceToast, navigate = navigateTo, + latestOnly = window.ApiClient.latestOnly, STALE = window.ApiClient.STALE }) { + const read = latestOnly({ lifetime: scope.signal }); + const show = location => { + container.replaceChildren(); + container.hidden = !(isLocation(location) && location.offer === true); + if (container.hidden) return; + container.append(libraryNotice({ scope, location, + onKeep: async button => { + button.setAttribute('aria-disabled', 'true'); + const saved = await keepLibrary({ api, scope, notify }); + if (scope.signal.aborted) return; + button.removeAttribute('aria-disabled'); + if (saved) show(saved); + }, + onMove: button => openMoveReview({ api, activity, scope, location, trigger: button, notify, navigate }) })); + }; + async function load() { + try { + const answer = await read(signal => api.get(LIBRARY_URL, { signal, cache: 'no-store' })); + if (answer === STALE || scope.signal.aborted) return; + show(answer.location); + } catch { + // The banner is advisory: an unreadable location shows nothing (the Overview row reports it). + if (!scope.signal.aborted) show(null); + } + } + container.hidden = true; + scope.cleanup(onMoveEnded(activity, load)); + load(); + return { reload: load }; +} diff --git a/web/js/workspace/system-overview.mjs b/web/js/workspace/system-overview.mjs index 6a66c7a..85c77f0 100644 --- a/web/js/workspace/system-overview.mjs +++ b/web/js/workspace/system-overview.mjs @@ -22,6 +22,7 @@ import { showWorkspaceToast } from './toast.mjs'; import { noun } from './plural.mjs'; import { operationRow } from './system-controls.mjs'; import { openThumbnailReview } from './thumbnail-review.mjs'; +import { isLocation, locationRow, onMoveEnded } from './library-location.mjs'; export const OVERVIEW_URL = '/prompt_manager/system/overview'; export const STATISTICS_URL = '/prompt_manager/system/statistics'; @@ -321,7 +322,7 @@ function surface(heading, rows) { } /** The Overview's nodes, and its thumbnail row (painted by its own read). */ -function overviewNodes(overview, { scope, navigate, openStatistics, openThumbnails }) { +function overviewNodes(overview, { api, activity, notify, scope, navigate, openStatistics, openThumbnails }) { const go = section => () => navigate(`#/system/${section}`); const review = operationRow({ scope, title: ambiguousTitle(overview.needs_review), description: Number.isSafeInteger(overview.needs_review?.pending) @@ -334,6 +335,9 @@ function overviewNodes(overview, { scope, navigate, openStatistics, openThumbnai review, operationRow({ scope, title: 'Backup status', description: backupText(overview.backup), label: 'View backups', onClick: go('backups') }), + // D129: where the library lives and why; Move library when it is in the node folder. + ...(isLocation(overview.location) + ? [locationRow(overview.location, { api, activity, scope, notify, navigate })] : []), // D99: Review thumbnails opens the thumbnail dialog (the concept's row starts a sample job directly). thumbnails, ]); @@ -396,7 +400,7 @@ export function mountOverview(container, { api, activity, scope, navigate = navi try { const overview = await overviewRead(signal => api.get(OVERVIEW_URL, { signal, cache: 'no-store' })); if (overview === STALE || scope.signal.aborted) return; - const { nodes, thumbnails } = overviewNodes(overview, { scope, navigate, openStatistics, + const { nodes, thumbnails } = overviewNodes(overview, { api, activity, notify, scope, navigate, openStatistics, openThumbnails: trigger => openThumbnailReview({ api, store: activity, scope, trigger, navigate }) }); thumbnailRow = thumbnails; sources = overview.sources; @@ -429,6 +433,8 @@ export function mountOverview(container, { api, activity, scope, navigate = navi paintThumbnails(); } const first = swap(loading()); + // D129: a finished Move library changes where the library lives. + scope.cleanup(onMoveEnded(activity, load)); load(); loadThumbnails(); return first; diff --git a/web/js/workspace/system-sources.mjs b/web/js/workspace/system-sources.mjs index c3b96cc..0433c86 100644 --- a/web/js/workspace/system-sources.mjs +++ b/web/js/workspace/system-sources.mjs @@ -18,7 +18,9 @@ * starts the thumbnail Activity job. * * D98: per-source Remove; an added source reads "not saved yet" until Save; the line under the list says when - * ComfyUI's output folder is used instead; "restart pending" for a source the running monitor does not watch yet; the + * ComfyUI's output folder is used instead of configured sources or after an unsaved removal of the last one; with no + * source configured, ComfyUI's output folder is the row "ComfyUI/output ยท watching ยท default" (`in_comfyui`: its public + * path is relative to the ComfyUI base) and its scan keeps the server's public path; "restart pending" for a source the running monitor does not watch yet; the * "Refresh files" operation (`/gallery/rescan`) closes the surface. D127: the add dialog's field is "Folder path", * starts empty and says the server's reason in place; its action is "Add", and Enter adds too. * D128: "watching until restart" for a source the running monitor still watches after monitoring was switched off; @@ -46,11 +48,13 @@ export const WATCH_COPY = Object.freeze({ watching: 'watching', stopping: 'watch // The states whose observer runs now: their scan only has to catch up on changes. const WATCHED = new Set(['watching', 'stopping']); const NOT_SAVED = 'not saved yet'; +// An auto-detected source (ComfyUI's output folder) says it is the default on its own row. +const DEFAULT_SOURCE = 'default'; +const COMFYUI_PREFIX = 'ComfyUI/'; export const DUPLICATE_SOURCE = 'same folder as another source'; /** D98: the line under the source list when it does not say everything by itself. */ export const NOTES = Object.freeze({ removedAll: 'No source is configured. After saving, ComfyUIโ€™s output folder is used.', - autoDetected: 'No source is configured; ComfyUIโ€™s output folder is used.', fallback: 'No configured source is available; ComfyUIโ€™s output folder is used.', none: 'No source is available.', }); @@ -95,8 +99,11 @@ const folderName = path => path.split('/').filter(Boolean).at(-1) ?? path; function effectiveRow(source, remove) { const watched = WATCHED.has(source.watch); - return { key: remove ?? `auto:${source.root_index}`, title: source.label, - description: `${source.path} ยท ${WATCH_COPY[source.watch] ?? WATCH_COPY.manual}`, + const auto = remove === null; + // The default row reads from ComfyUI's own folder ("ComfyUI/output") when its public path is relative to the base. + const shown = auto && source.in_comfyui === true ? `${COMFYUI_PREFIX}${source.path}` : source.path; + const parts = [shown, WATCH_COPY[source.watch] ?? WATCH_COPY.manual, ...(auto ? [DEFAULT_SOURCE] : [])]; + return { key: remove ?? `auto:${source.root_index}`, title: source.label, description: parts.join(' ยท '), scan: { label: watched ? 'Scan changes' : 'Scan folder', root: source.root_index, path: source.path }, remove }; } @@ -124,7 +131,7 @@ export function sourceRows({ sources, saved, current, labels = addedLabels }) { if (dirty) return { rows, note: roots.length ? null : NOTES.removedAll }; if (sources.auto_detected) { rows.push(...effective.map(source => effectiveRow(source, null))); - return { rows, note: configured.length ? NOTES.fallback : NOTES.autoDetected }; + return { rows, note: configured.length ? NOTES.fallback : null }; } return { rows, note: rows.length ? null : NOTES.none }; } diff --git a/web/js/workspace/workspace.css b/web/js/workspace/workspace.css index 7c3a5af..7d4b06b 100644 --- a/web/js/workspace/workspace.css +++ b/web/js/workspace/workspace.css @@ -414,3 +414,9 @@ pre.prompt-text { font: inherit; font-size: 12px; line-height: 1.75; margin-bloc field) beside their label, as the concept's own radio rows (s1160 `.compare-item label`). */ .field[role="radiogroup"] > label { display: flex; align-items: center; gap: 8px; /* D90 */ } .field[role="radiogroup"] input[type="radio"] { width: auto; margin: 0; /* D90 */ } +/* D129: the library-location notice sits above every page, in main's own padding and width, so its `.notice` lines + up with the page below; main's top padding then spaces it from the page. */ +.library-banner { max-width: 1900px; margin: auto; padding: 30px 30px 0; /* D129 */ } +.library-banner > .notice { margin-bottom: 0; /* D129 */ } +@media (max-width: 1250px) { .library-banner { padding: 24px 24px 0; /* D129 */ } } +@media (max-width: 700px) { .library-banner { padding: 20px 14px 0; /* D129 */ } } diff --git a/web/workspace.html b/web/workspace.html index e912c2c..c0532f6 100644 --- a/web/workspace.html +++ b/web/workspace.html @@ -40,6 +40,8 @@ ComfyUI + +