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 + +