diff --git a/__init__.py b/__init__.py index 86d9d6e7..2e070920 100644 --- a/__init__.py +++ b/__init__.py @@ -111,6 +111,25 @@ except Exception as e: except Exception: print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}") +# Count prompt runs when workflows are queued (ComfyUI skips unchanged nodes, so +# node execution alone misses re-runs). Separate from monitoring so either can fail. +try: + from .py.config import server_instance + from .utils.usage_tracking import register_queue_hook + + register_queue_hook(server_instance, lambda: _global_db) +except Exception as e: + try: + from .utils.logging_config import get_logger + + get_logger("prompt_manager.init").error(f"Failed to register usage hook: {e}") + except Exception: + import logging + + logging.getLogger("prompt_manager.init").error( + f"Failed to register usage hook: {e}" + ) + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] # Print startup message with loaded tools diff --git a/database/models.py b/database/models.py index cee627a8..5dc2cfeb 100644 --- a/database/models.py +++ b/database/models.py @@ -2,6 +2,7 @@ Database schema and models for KikoTextEncode prompt storage. """ +import datetime import sqlite3 import os import threading @@ -17,6 +18,18 @@ except ImportError: sys.path.insert(0, current_dir) from utils.logging_config import get_logger +# last_used_at is UTC ISO text so ORDER BY can compare it as text. The backfill +# writes milliseconds (SQLite's %f); live writes use microseconds so runs in the +# same millisecond still order correctly. Both forms sort correctly together. +ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00" + + +def utc_now_iso() -> str: + """Current UTC time as ISO text with microseconds, e.g. 2026-10-07T18:39:06.758123+00:00.""" + return datetime.datetime.now(datetime.timezone.utc).isoformat( + timespec="microseconds" + ) + class PromptModel: """Database model for prompt storage and schema management.""" @@ -78,7 +91,9 @@ class PromptModel: tags TEXT, rating INTEGER CHECK(rating >= 1 AND rating <= 5), notes TEXT, - hash TEXT UNIQUE + hash TEXT UNIQUE, + last_used_at TIMESTAMP, + run_count INTEGER NOT NULL DEFAULT 0 ) """) @@ -131,6 +146,9 @@ class PromptModel: # Migrate JSON tags to normalized junction tables self._migrate_json_tags_to_junction(conn) + # Usage tracking columns (last, after any migration that rebuilds prompts) + self._migrate_add_usage_columns(conn) + def _create_indexes(self, conn: sqlite3.Connection) -> None: """ Create indexes for better query performance. @@ -150,6 +168,8 @@ class PromptModel: "CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)", "CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)", "CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)", + "CREATE INDEX IF NOT EXISTS idx_prompts_last_used ON prompts(last_used_at)", + "CREATE INDEX IF NOT EXISTS idx_prompts_run_count ON prompts(run_count)", "CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)", "CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)", "CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)", @@ -440,6 +460,50 @@ class PromptModel: except Exception as e: self.logger.error(f"Tag junction migration error: {e}") + def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None: + """ + Add last_used_at and run_count (3.2.4) and fill in any prompt without usage data. + + The backfill runs on every start but only touches rows whose last_used_at + is NULL, so it is a no-op once healed. That also repairs rows written by an + older version after a downgrade, or left unfilled by an interrupted upgrade. + + The estimate comes from linked images: run_count is the image count (at + least 1) and last_used_at the newest image time, falling back to created_at. + Timestamps are compared with julianday() because prompts use ISO 'T' + timestamps while generated_images uses SQLite's CURRENT_TIMESTAMP format. + """ + try: + columns = {row[1] for row in conn.execute("PRAGMA table_info(prompts)")} + if "last_used_at" not in columns: + self.logger.info("Adding prompt usage column last_used_at") + conn.execute("ALTER TABLE prompts ADD COLUMN last_used_at TIMESTAMP") + if "run_count" not in columns: + self.logger.info("Adding prompt usage column run_count") + conn.execute( + "ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL DEFAULT 0" + ) + + cursor = conn.execute(f""" + UPDATE prompts SET + run_count = MAX(run_count, 1, ( + SELECT COUNT(*) FROM generated_images gi + WHERE gi.prompt_id = prompts.id + )), + last_used_at = strftime('{ISO_TIMESTAMP_FORMAT}', MAX( + COALESCE(julianday(created_at), julianday('now')), + COALESCE(( + SELECT MAX(julianday(gi.generation_time)) + FROM generated_images gi WHERE gi.prompt_id = prompts.id + ), 0) + )) + WHERE last_used_at IS NULL + """) + if cursor.rowcount: + self.logger.info(f"Estimated usage for {cursor.rowcount} prompts") + except Exception as e: + self.logger.error(f"Usage tracking migration error: {e}") + def migrate_database(self) -> None: """ Apply any pending database migrations. diff --git a/database/operations.py b/database/operations.py index bceb9adb..f3c0765b 100644 --- a/database/operations.py +++ b/database/operations.py @@ -8,7 +8,7 @@ import datetime import os from typing import Optional, List, Dict, Any, Union -from .models import PromptModel +from .models import PromptModel, utc_now_iso # Import logging system try: @@ -28,6 +28,25 @@ TAG_SUBQUERY = ( "WHERE pt.prompt_id = prompts.id) AS _tag_list" ) +# Server-side sort orders, keyed by the value the UI sends. Only these strings ever +# reach ORDER BY; unknown keys fall back to DEFAULT_SORT. +SORT_ORDERS = { + "last_used_desc": "last_used_at DESC, id DESC", + "created_desc": "created_at DESC, id DESC", + "created_asc": "created_at ASC, id ASC", + "run_count_desc": "run_count DESC, last_used_at DESC, id DESC", + "rating_desc": "rating IS NULL, rating DESC, id DESC", + "rating_asc": "rating IS NULL, rating ASC, id DESC", + "text_asc": "text COLLATE NOCASE ASC, id ASC", + "text_desc": "text COLLATE NOCASE DESC, id DESC", +} +DEFAULT_SORT = "created_desc" + + +def order_by_clause(sort: Optional[str]) -> str: + """Return a whitelisted ORDER BY expression for a UI sort key.""" + return SORT_ORDERS.get(sort or DEFAULT_SORT, SORT_ORDERS[DEFAULT_SORT]) + def _resolve_db_path(db_path: Optional[str] = None) -> str: """Resolve the database path from config, falling back to defaults. @@ -122,8 +141,9 @@ class PromptDatabase: cursor = conn.execute( """ INSERT INTO prompts ( - text, category, tags, rating, notes, hash, created_at, updated_at - ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?) + text, category, tags, rating, notes, hash, created_at, updated_at, + last_used_at + ) VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?) """, ( text.strip(), @@ -133,6 +153,7 @@ class PromptDatabase: prompt_hash, datetime.datetime.now(datetime.timezone.utc).isoformat(), datetime.datetime.now(datetime.timezone.utc).isoformat(), + utc_now_iso(), ), ) prompt_id = cursor.lastrowid @@ -142,6 +163,26 @@ class PromptDatabase: self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}") return prompt_id + def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool: + """ + Count runs of a prompt and mark it as most recently used. + + Args: + prompt_id: The prompt ID + times: Number of runs to add (queued runs counted on first save) + + Returns: + True if the prompt exists and was updated + """ + with self.model.get_connection() as conn: + cursor = conn.execute( + "UPDATE prompts SET run_count = run_count + ?, last_used_at = ?" + " WHERE id = ?", + (max(1, int(times)), utc_now_iso(), prompt_id), + ) + conn.commit() + return cursor.rowcount > 0 + def get_prompt_by_id(self, prompt_id: int) -> Optional[Dict[str, Any]]: """ Get a prompt by its ID. @@ -191,6 +232,7 @@ class PromptDatabase: limit: int = 100, offset: int = 0, tag_partial: bool = False, + sort: Optional[str] = None, ) -> List[Dict[str, Any]]: """ Search prompts with various filters. @@ -207,6 +249,7 @@ class PromptDatabase: limit: Maximum number of results offset: Number of results to skip tag_partial: Use LIKE matching for tags instead of exact match + sort: Sort key from SORT_ORDERS (default: newest first) Returns: List of dictionaries containing prompt data @@ -275,7 +318,7 @@ class PromptDatabase: query_parts.append("AND created_at <= ?") params.append(date_to) - query_parts.append("ORDER BY created_at DESC LIMIT ? OFFSET ?") + query_parts.append(f"ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?") params.extend([limit, offset]) query = " ".join(query_parts) @@ -291,13 +334,16 @@ class PromptDatabase: return prompts - def get_recent_prompts(self, limit: int = 10, offset: int = 0) -> Dict[str, Any]: + def get_recent_prompts( + self, limit: int = 10, offset: int = 0, sort: Optional[str] = None + ) -> Dict[str, Any]: """ Get the most recent prompts with pagination support. Args: limit: Maximum number of prompts to return offset: Number of prompts to skip (for pagination) + sort: Sort key from SORT_ORDERS (default: newest first) Returns: Dictionary containing prompt data and pagination info @@ -310,7 +356,8 @@ class PromptDatabase: # Get paginated results cursor = conn.execute( - f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts ORDER BY created_at DESC LIMIT ? OFFSET ?", + f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts" + f" ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?", (limit, offset), ) rows = cursor.fetchall() diff --git a/prompt_manager.py b/prompt_manager.py index ffbec23a..9637c7c5 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -4,7 +4,7 @@ with persistent prompt storage and search capabilities. """ import time -from typing import Any, Tuple +from typing import Any, Dict, Optional, Tuple try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict @@ -94,6 +94,12 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): }, ), }, + "hidden": { + # The queued graph and this node's id: used to tell the positive + # prompt of a run from negative-prompt PromptManager nodes + "prompt": "PROMPT", + "unique_id": "UNIQUE_ID", + }, } RETURN_TYPES = (IO.CONDITIONING, IO.STRING) @@ -119,6 +125,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): search_text: str = "", prepend_text: str = "", append_text: str = "", + prompt: Optional[Dict[str, Any]] = None, + unique_id: Optional[str] = None, ) -> Tuple[Any]: """ Encode the text prompt and save it to the database. @@ -181,30 +189,24 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): extended_tags.append(f"append:{append_text.strip()[:50]}") try: - prompt_id = self._save_prompt_to_database( + prompt_id = self._track_prompt_execution( text=storage_text.strip(), + encoding_text=encoding_text.strip(), category=category.strip() if category else None, tags=extended_tags if extended_tags else None, + additional_data={ + "category": category.strip() if category else None, + "tags": extended_tags, + "prepend_text": ( + prepend_text.strip() if prepend_text else None + ), + "append_text": append_text.strip() if append_text else None, + "final_text": encoding_text.strip(), + }, + prompt_graph=prompt, + unique_id=unique_id, ) - - # Set current prompt for image tracking - if prompt_id: - execution_id = self.prompt_tracker.set_current_prompt( - prompt_text=encoding_text.strip(), - additional_data={ - "category": category.strip() if category else None, - "tags": extended_tags, - "prompt_id": prompt_id, - "prepend_text": ( - prepend_text.strip() if prepend_text else None - ), - "append_text": append_text.strip() if append_text else None, - "final_text": encoding_text.strip(), - }, - ) - self.logger.debug( - f"Set execution context: {execution_id} for prompt ID: {prompt_id}" - ) + self.logger.debug(f"Tracked prompt ID: {prompt_id}") except Exception as e: # Log error but don't fail the encoding diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 43a63264..ee190d56 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -24,6 +24,9 @@ try: from .utils.comfyui_integration import get_comfyui_integration from .utils.image_monitor import get_image_monitor from .utils.prompt_tracker import get_prompt_tracker + from .utils.hashing import generate_prompt_hash + from .utils.prompt_graph import resolve_text, run_prompt_nodes + from .utils.usage_tracking import PENDING_FIRST_USE except ImportError: import sys @@ -32,6 +35,9 @@ except ImportError: from utils.comfyui_integration import get_comfyui_integration from utils.image_monitor import get_image_monitor from utils.prompt_tracker import get_prompt_tracker + from utils.hashing import generate_prompt_hash + from utils.prompt_graph import resolve_text, run_prompt_nodes + from utils.usage_tracking import PENDING_FIRST_USE class PromptManagerBase: @@ -56,14 +62,23 @@ class PromptManagerBase: self.logger.debug(f"{self.__class__.__name__} node initialization completed") def _save_prompt_to_database( - self, text: str, category: Optional[str] = None, tags: Optional[list] = None + self, + text: str, + category: Optional[str] = None, + tags: Optional[list] = None, + count_run: bool = False, ) -> Optional[int]: """Save the prompt to the SQLite database. + Runs are normally counted by the queue hook (utils/usage_tracking.py), so a + save only counts one when asked to (batch items whose text the hook could not + see) or when the hook queued this prompt before it existed. + Args: text: The prompt text category: Optional category tags: List of tags + count_run: Count this execution as a run of the prompt Returns: The prompt ID if saved successfully, None otherwise @@ -82,6 +97,8 @@ class PromptManagerBase: prompt_id=existing["id"], category=category, tags=tags ) self.logger.debug("Updated metadata for existing prompt") + if count_run: + self._record_use(existing["id"]) return existing["id"] self.logger.debug( @@ -93,6 +110,10 @@ class PromptManagerBase: if prompt_id: self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}") + pending = getattr(self, "pending_first_use", PENDING_FIRST_USE) + runs = pending.consume(prompt_hash) + (1 if count_run else 0) + if runs: + self._record_use(prompt_id, times=runs) else: self.logger.warning("Failed to save prompt - no ID returned") @@ -102,6 +123,68 @@ class PromptManagerBase: self.logger.error(f"Error saving prompt to database: {e}") return None + def _prompt_role( + self, prompt_graph: Any, unique_id: Any, text: str + ) -> Dict[str, bool]: + """How this node's execution relates to the run, from the queued graph. + + Returns flags: + is_run_prompt: feeds a sampler's positive input (images belong to it) + seen_at_queue: the queue hook resolved exactly this text, so it already + counted the run and the image's metadata can identify the prompt + known: the graph and node id were available + """ + if not isinstance(prompt_graph, dict) or unique_id is None: + return {"is_run_prompt": True, "seen_at_queue": False, "known": False} + node_id = str(unique_id) + resolved = resolve_text(prompt_graph, node_id) + return { + "is_run_prompt": node_id in run_prompt_nodes(prompt_graph), + "seen_at_queue": resolved is not None + and generate_prompt_hash(resolved) == generate_prompt_hash(text), + "known": True, + } + + def _track_prompt_execution( + self, + text: str, + encoding_text: str, + category: Optional[str], + tags: Optional[list], + additional_data: Dict[str, Any], + prompt_graph: Any = None, + unique_id: Any = None, + ) -> Optional[int]: + """Save the prompt and, for the run's positive prompt, register it for linking. + + Negative-prompt nodes are saved but never become the prompt their images are + linked to. Text the queue hook could resolve is counted there and linked from + the image's own metadata; anything else (PromptSearchList batches, unknown + text nodes, or a resolution that doesn't match) is counted here and pushed + onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept. + """ + role = self._prompt_role(prompt_graph, unique_id, text) + batch_item = ( + role["known"] and role["is_run_prompt"] and not role["seen_at_queue"] + ) + prompt_id = self._save_prompt_to_database( + text=text, category=category, tags=tags, count_run=batch_item + ) + if prompt_id and role["is_run_prompt"]: + self.prompt_tracker.set_current_prompt( + prompt_text=encoding_text, + additional_data={**additional_data, "prompt_id": prompt_id}, + push_to_queue=batch_item or not role["known"], + ) + return prompt_id + + def _record_use(self, prompt_id: int, times: int = 1) -> None: + """Count runs; a tracking failure must never block generation.""" + try: + self.db.record_prompt_use(prompt_id, times=times) + except Exception as e: + self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}") + def _inject_lora_trigger_words(self, text: str) -> str: """Append LoRA trigger words if integration is enabled. diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 73bc061b..37df8564 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -4,7 +4,7 @@ without CLIP encoding, while maintaining all database and search features. """ import time -from typing import Tuple +from typing import Any, Dict, Optional, Tuple try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict @@ -94,6 +94,12 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): }, ), }, + "hidden": { + # The queued graph and this node's id: used to tell the positive + # prompt of a run from negative-prompt PromptManager nodes + "prompt": "PROMPT", + "unique_id": "UNIQUE_ID", + }, } RETURN_TYPES = (IO.STRING,) @@ -117,6 +123,8 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): search_text: str = "", prepend_text: str = "", append_text: str = "", + prompt: Optional[Dict[str, Any]] = None, + unique_id: Optional[str] = None, ) -> Tuple[str]: """ Process the text prompt and save it to the database. @@ -162,29 +170,23 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): extended_tags.append(f"append:{append_text.strip()[:50]}") try: - prompt_id = self._save_prompt_to_database( + prompt_id = self._track_prompt_execution( text=storage_text.strip(), + encoding_text=final_text.strip(), category=category.strip() if category else None, tags=extended_tags if extended_tags else None, + additional_data={ + "category": category.strip() if category else None, + "tags": extended_tags, + "prepend_text": ( + prepend_text.strip() if prepend_text else None + ), + "append_text": append_text.strip() if append_text else None, + }, + prompt_graph=prompt, + unique_id=unique_id, ) - - # Set current prompt for image tracking - if prompt_id: - execution_id = self.prompt_tracker.set_current_prompt( - prompt_text=final_text.strip(), - additional_data={ - "category": category.strip() if category else None, - "tags": extended_tags, - "prompt_id": prompt_id, - "prepend_text": ( - prepend_text.strip() if prepend_text else None - ), - "append_text": append_text.strip() if append_text else None, - }, - ) - self.logger.debug( - f"Set execution context: {execution_id} for prompt ID: {prompt_id}" - ) + self.logger.debug(f"Tracked prompt ID: {prompt_id}") except Exception as e: # Log error but don't fail the processing diff --git a/py/api/prompts.py b/py/api/prompts.py index 4a47eeda..5d0efdfc 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -132,6 +132,8 @@ class PromptRoutesMixin: limit = int(request.query.get("limit", 50)) folder = request.query.get("folder", "").strip() or None + # Validated against a whitelist in the database layer + sort = request.query.get("sort") or None tags = None if tags_str: @@ -150,6 +152,7 @@ class PromptRoutesMixin: rating_min=min_rating, limit=limit, folder=folder, + sort=sort, ) self._enrich_prompt_images(results) @@ -182,7 +185,7 @@ class PromptRoutesMixin: return web.json_response({"success": False, "error": str(e)}, status=500) async def get_recent_prompts(self, request): - """Retrieve recently created prompts with pagination support.""" + """Retrieve prompts with pagination and an optional sort (default newest first).""" try: limit = int(request.query.get("limit", 50)) page = int(request.query.get("page", 1)) @@ -196,8 +199,9 @@ class PromptRoutesMixin: elif limit < 1: limit = 1 + sort = request.query.get("sort") or None results = await self._run_in_executor( - self.db.get_recent_prompts, limit=limit, offset=offset + self.db.get_recent_prompts, limit=limit, offset=offset, sort=sort ) self._enrich_prompt_images(results["prompts"]) diff --git a/pyproject.toml b/pyproject.toml index a53431b7..188eef9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "promptmanager" description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite." -version = "3.2.3" +version = "3.2.4" license = {file = "LICENSE"} dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"] diff --git a/tests/js/prompt-list-sort.test.js b/tests/js/prompt-list-sort.test.js new file mode 100644 index 00000000..c530ad04 --- /dev/null +++ b/tests/js/prompt-list-sort.test.js @@ -0,0 +1,66 @@ +// Run with: node --test tests/js/ +const test = require("node:test"); +const assert = require("node:assert/strict"); +const sort = require("../../web/js/prompt-list-sort.js"); + +test("default sort is recently used and is a listed option", () => { + assert.equal(sort.DEFAULT_SORT, "last_used_desc"); + assert.ok(sort.SORT_OPTIONS.some((o) => o.value === sort.DEFAULT_SORT)); +}); + +test("options include Recently used and Most used, each with a label", () => { + const values = sort.SORT_OPTIONS.map((o) => o.value); + assert.ok(values.includes("last_used_desc")); + assert.ok(values.includes("run_count_desc")); + for (const option of sort.SORT_OPTIONS) assert.ok(option.label.trim(), option.value); + assert.equal(new Set(values).size, values.length, "values are unique"); +}); + +test("options cannot be mutated by callers", () => { + assert.throws(() => sort.SORT_OPTIONS.push({ value: "x", label: "x" })); + // Non-strict code fails silently on frozen objects, so check the value survived + const first = sort.SORT_OPTIONS[0].value; + sort.SORT_OPTIONS[0].value = "hacked"; + assert.equal(sort.SORT_OPTIONS[0].value, first); + assert.ok(sort.SORT_OPTIONS.every(Object.isFrozen)); +}); + +test("normalizeSort keeps known values and falls back for anything else", () => { + assert.equal(sort.normalizeSort("run_count_desc"), "run_count_desc"); + for (const bad of [undefined, null, "", "created_at; DROP TABLE prompts", 42, "LAST_USED_DESC"]) { + assert.equal(sort.normalizeSort(bad), sort.DEFAULT_SORT, String(bad)); + } +}); + +test("buildRecentUrl computes offset from page and always sends a valid sort", () => { + assert.equal( + sort.buildRecentUrl({ page: 3, limit: 25, sort: "rating_desc" }), + "/prompt_manager/recent?limit=25&offset=50&page=3&sort=rating_desc", + ); + assert.equal( + sort.buildRecentUrl({ sort: "bogus" }), + "/prompt_manager/recent?limit=50&offset=0&page=1&sort=last_used_desc", + ); +}); + +test("withSort returns a copy with the sort set and leaves the input untouched", () => { + const params = new URLSearchParams({ text: "castle", sort: "text_asc" }); + const result = sort.withSort(params, "run_count_desc"); + assert.equal(result.get("sort"), "run_count_desc"); + assert.equal(result.get("text"), "castle"); + assert.equal(params.get("sort"), "text_asc"); + assert.equal(sort.withSort(new URLSearchParams(), "nope").get("sort"), sort.DEFAULT_SORT); +}); + +test("formatRunCount labels zero, one and many runs", () => { + assert.equal(sort.formatRunCount(0), "Never run"); + assert.equal(sort.formatRunCount(1), "1 run"); + assert.equal(sort.formatRunCount(642), "642 runs"); + assert.equal(sort.formatRunCount(1234), "1,234 runs"); +}); + +test("formatRunCount hides the badge for missing or invalid counts", () => { + for (const bad of [undefined, null, -1, 1.5, "3", NaN]) { + assert.equal(sort.formatRunCount(bad), "", String(bad)); + } +}); diff --git a/tests/test_api.py b/tests/test_api.py index f9853537..d7bc4164 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -238,6 +238,65 @@ class TestUpdatePrompt(APITestCase): self.assertTrue(data["success"]) +class TestUsageSorting(APITestCase): + """Server-side sort for /recent and /search (3.2.4, #90.1).""" + + def _use(self, prompt_id, times=1): + for _ in range(times): + self.api.db.record_prompt_use(prompt_id) + + async def _ids(self, path): + resp = await self.client.request("GET", path) + self.assertEqual(resp.status, 200) + data = await resp.json() + self.assertTrue(data["success"]) + return [p["id"] for p in data["results"]] + + async def test_recent_sorted_by_last_used_bubbles_rerun_prompt(self): + old = self._save_prompt("Old prompt") + self._save_prompt("New prompt") + self._use(old) + ids = await self._ids("/prompt_manager/recent?limit=1&sort=last_used_desc") + self.assertEqual(ids, [old]) + + async def test_recent_sorted_by_run_count(self): + # Most-run prompt is the oldest, so newest-first would give the opposite order + most = self._save_prompt("Run thrice") + least = self._save_prompt("Run once") + self._use(most, times=3) + self._use(least) + ids = await self._ids("/prompt_manager/recent?sort=run_count_desc") + self.assertEqual(ids, [most, least]) + + async def test_recent_default_is_newest_first(self): + first = self._save_prompt("First") + second = self._save_prompt("Second") + self._use(first) + self.assertEqual(await self._ids("/prompt_manager/recent"), [second, first]) + + async def test_unknown_sort_falls_back_instead_of_erroring(self): + self._save_prompt("Only") + ids = await self._ids( + "/prompt_manager/recent?sort=created_at;DROP%20TABLE%20prompts" + ) + self.assertEqual(len(ids), 1) + + async def test_search_respects_sort(self): + a = self._save_prompt("castle at dawn") + self._save_prompt("castle at dusk") + self._use(a) + ids = await self._ids("/prompt_manager/search?text=castle&sort=last_used_desc") + self.assertEqual(ids[0], a) + + async def test_results_include_usage_fields(self): + pid = self._save_prompt("Tracked") + self._use(pid, times=2) + resp = await self.client.request("GET", "/prompt_manager/recent") + prompt = (await resp.json())["results"][0] + self.assertEqual(prompt["run_count"], 2) + self.assertTrue(prompt["last_used_at"]) + + class TestTags(APITestCase): async def test_get_all_tags(self): diff --git a/tests/test_image_linking.py b/tests/test_image_linking.py new file mode 100644 index 00000000..9c2827a9 --- /dev/null +++ b/tests/test_image_linking.py @@ -0,0 +1,141 @@ +"""Image linking picks the positive prompt, preferring the image's own metadata.""" + +import os +import sys +import tempfile +import unittest +import unittest.mock as mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash +from utils.image_monitor import ImageGenerationHandler +from utils.prompt_tracker import PromptTracker + +LOADER = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "m.safetensors"}, +} + + +def graph(positive_text="a cat", negative_text="blurry"): + # Negative listed first: dict order must not decide which prompt an image gets + return { + "176": {"class_type": "PromptManager", "inputs": {"text": negative_text}}, + "134": {"class_type": "PromptManager", "inputs": {"text": positive_text}}, + "5": { + "class_type": "KSampler", + "inputs": { + "positive": ["134", 0], + "negative": ["176", 0], + "model": ["20", 0], + }, + }, + "20": LOADER, + } + + +class LinkingTestCase(unittest.TestCase): + def setUp(self): + fd, self.db_path = tempfile.mkstemp(suffix=".db") + os.close(fd) + fd, self.image = tempfile.mkstemp(suffix=".png") + os.close(fd) + self.db = PromptDatabase(self.db_path) + self.tracker = PromptTracker(self.db) + self.handler = ImageGenerationHandler(self.db, self.tracker) + self.linked = [] + self.handler.link_image_to_prompt = ( + lambda path, prompt, meta: self.linked.append(prompt["id"]) + ) + + def tearDown(self): + for path in ( + self.image, + self.db_path, + self.db_path + "-wal", + self.db_path + "-shm", + ): + if os.path.exists(path): + os.unlink(path) + + def _save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + def _process(self, metadata): + with mock.patch.object( + self.handler.metadata_extractor, "extract_metadata", return_value=metadata + ): + self.handler.process_new_image(self.image) + return self.linked[-1] if self.linked else None + + +class TestMetadataLinking(LinkingTestCase): + def test_links_the_positive_prompt_not_the_first_prompt_manager_node(self): + positive = self._save("a cat") + self._save("blurry") + self.assertEqual(self._process({"prompt": graph()}), positive) + + def test_image_metadata_wins_over_a_stale_queue_entry(self): + positive = self._save("a cat") + stale = self._save("previous run") + self.tracker.set_current_prompt("previous run", {"prompt_id": stale}) + + self.assertEqual(self._process({"prompt": graph()}), positive) + self.assertEqual( + self.tracker.pop_next_prompt()["id"], stale, "queue left intact" + ) + + def test_text_from_a_known_string_node_links_from_metadata(self): + positive = self._save("a cat") + linked_graph = graph() + linked_graph["134"]["inputs"]["text"] = ["8", 0] + linked_graph["8"] = { + "class_type": "PrimitiveString", + "inputs": {"value": "a cat"}, + } + self.assertEqual(self._process({"prompt": linked_graph}), positive) + + def test_batch_items_with_linked_text_use_the_queue(self): + first, second = self._save("item one"), self._save("item two") + batch_graph = graph() + batch_graph["134"]["inputs"]["text"] = ["7", 0] + batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} + self.tracker.set_current_prompt("item one", {"prompt_id": first}) + self.tracker.set_current_prompt("item two", {"prompt_id": second}) + + self.assertEqual(self._process({"prompt": batch_graph}), first) + self.assertEqual(self._process({"prompt": batch_graph}), second) + + def test_api_graph_present_skips_the_role_unaware_workflow_fallback(self): + negative = self._save("blurry") + batch_graph = graph() + batch_graph["134"]["inputs"]["text"] = ["7", 0] + metadata = { + "prompt": batch_graph, + "text_encoder_nodes": [ + {"type": "PromptManager", "inputs": [], "widgets_values": ["blurry"]} + ], + } + self.assertIsNone(self.handler._find_prompt_from_metadata(metadata)) + self.assertNotEqual(self._process(metadata), negative) + + def test_workflow_fallback_still_used_without_an_api_graph(self): + positive = self._save("from workflow only") + metadata = { + "text_encoder_nodes": [ + { + "type": "PromptManager", + "inputs": [], + "widgets_values": ["from workflow only"], + } + ] + } + self.assertEqual( + self.handler._find_prompt_from_metadata(metadata)["id"], positive + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_node_tracking_wiring.py b/tests/test_node_tracking_wiring.py new file mode 100644 index 00000000..e880cf5c --- /dev/null +++ b/tests/test_node_tracking_wiring.py @@ -0,0 +1,83 @@ +"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once.""" + +import os +import sys +import threading +import unittest +import unittest.mock as mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from prompt_manager import PromptManager +from prompt_manager_text import PromptManagerText +from utils import usage_tracking + +GRAPH = {"134": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}} + + +def bare(node_cls): + """A node without its database, tracker or gallery side effects.""" + node = node_cls.__new__(node_cls) + node.logger = __import__("logging").getLogger("test.wiring") + node.comfyui_integration = mock.Mock() + node._inject_lora_trigger_words = lambda text: text + node._track_prompt_execution = mock.Mock(return_value=42) + return node + + +class TestHiddenInputs(unittest.TestCase): + def test_both_nodes_request_the_queued_graph_and_their_id(self): + for node_cls in (PromptManager, PromptManagerText): + hidden = node_cls.INPUT_TYPES().get("hidden", {}) + self.assertEqual(hidden.get("prompt"), "PROMPT", node_cls.__name__) + self.assertEqual(hidden.get("unique_id"), "UNIQUE_ID", node_cls.__name__) + + +class TestNodesForwardGraph(unittest.TestCase): + def test_prompt_manager_forwards_graph_and_texts(self): + node = bare(PromptManager) + clip = mock.Mock() + node.encode_prompt( + clip, "a cat", prepend_text="best,", prompt=GRAPH, unique_id="134" + ) + + kwargs = node._track_prompt_execution.call_args.kwargs + self.assertEqual(kwargs["prompt_graph"], GRAPH) + self.assertEqual(kwargs["unique_id"], "134") + self.assertEqual(kwargs["text"], "a cat") + self.assertEqual(kwargs["encoding_text"], "best, a cat") + self.assertEqual(kwargs["additional_data"]["final_text"], "best, a cat") + + def test_prompt_manager_text_forwards_graph(self): + node = bare(PromptManagerText) + node.process_text("a castle", prompt=GRAPH, unique_id="7") + + kwargs = node._track_prompt_execution.call_args.kwargs + self.assertEqual((kwargs["prompt_graph"], kwargs["unique_id"]), (GRAPH, "7")) + self.assertEqual(kwargs["text"], "a castle") + + def test_nodes_still_work_without_graph_context(self): + node = bare(PromptManager) + node.encode_prompt(mock.Mock(), "a cat") + kwargs = node._track_prompt_execution.call_args.kwargs + self.assertIsNone(kwargs["prompt_graph"]) + self.assertIsNone(kwargs["unique_id"]) + + +class TestQueueHookRegistration(unittest.TestCase): + def test_registers_once_and_handler_passes_requests_through(self): + server = mock.Mock() + db = mock.Mock() + db.get_prompt_by_hash.return_value = None + with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()): + self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db)) + self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db)) + + server.add_on_prompt_handler.assert_called_once() + handler = server.add_on_prompt_handler.call_args.args[0] + request = {"prompt": GRAPH} + self.assertIs(handler(request), request) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prompt_graph.py b/tests/test_prompt_graph.py new file mode 100644 index 00000000..b0b61857 --- /dev/null +++ b/tests/test_prompt_graph.py @@ -0,0 +1,244 @@ +"""Tests for utils.prompt_graph: which PromptManager nodes feed a positive input.""" + +import os +import sys +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils.prompt_graph import ( + positive_prompt_nodes, + resolve_text, + run_prompt_nodes, +) + + +def pm(text="a cat", **inputs): + return {"class_type": "PromptManager", "inputs": {"text": text, **inputs}} + + +def ksampler(positive, negative, model="20"): + return { + "class_type": "KSampler", + "inputs": { + "positive": [positive, 0], + "negative": [negative, 0], + "model": [model, 0], + }, + } + + +LOADER = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "m.safetensors"}, +} + + +class TestPositivePromptNodes(unittest.TestCase): + def test_positive_and_negative_prompt_manager_nodes(self): + # Negative listed first, as in the #75 / pmtest workflow + graph = { + "176": pm("blurry"), + "134": pm("a cat"), + "5": ksampler("134", "176"), + "20": LOADER, + } + self.assertEqual(positive_prompt_nodes(graph), ["134"]) + + def test_negative_only_prompt_manager_is_never_positive(self): + graph = { + "176": pm("blurry"), + "6": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}, + "5": ksampler("6", "176"), + "20": LOADER, + } + self.assertEqual(positive_prompt_nodes(graph), []) + + def test_through_controlnet_and_combine_pass_throughs(self): + graph = { + "1": pm("subject"), + "2": pm("style"), + "3": pm("blurry"), + "4": { + "class_type": "ConditioningCombine", + "inputs": {"conditioning_1": ["1", 0], "conditioning_2": ["2", 0]}, + }, + "5": { + "class_type": "ControlNetApplyAdvanced", + "inputs": {"positive": ["4", 0], "negative": ["3", 0], "strength": 1.0}, + }, + "6": { + "class_type": "KSampler", + "inputs": { + "positive": ["5", 0], + "negative": ["5", 1], + "model": ["20", 0], + }, + }, + "20": LOADER, + } + self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "2"]) + + def test_cfg_guider_and_basic_guider(self): + graph = { + "1": pm("fox"), + "2": pm("watermark"), + "3": { + "class_type": "CFGGuider", + "inputs": { + "positive": ["1", 0], + "negative": ["2", 0], + "model": ["20", 0], + "cfg": 5, + }, + }, + "4": pm("owl"), + "5": { + "class_type": "BasicGuider", + "inputs": {"conditioning": ["4", 0], "model": ["20", 0]}, + }, + "20": LOADER, + } + self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "4"]) + + def test_prompt_manager_text_wired_into_an_encoder(self): + graph = { + "1": {"class_type": "PromptManagerText", "inputs": {"text": "a castle"}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 0]}}, + "3": {"class_type": "PromptManagerText", "inputs": {"text": "lowres"}}, + "4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}}, + "5": ksampler("2", "4"), + "20": LOADER, + } + self.assertEqual(positive_prompt_nodes(graph), ["1"]) + + def test_node_used_as_both_positive_and_negative_counts_as_positive(self): + graph = {"1": pm("same"), "5": ksampler("1", "1"), "20": LOADER} + self.assertEqual(positive_prompt_nodes(graph), ["1"]) + + def test_two_samplers_sharing_a_prompt_list_it_once(self): + graph = { + "1": pm("hero"), + "2": pm("blurry"), + "5": ksampler("1", "2"), + "14": ksampler("1", "2"), + "20": LOADER, + } + self.assertEqual(positive_prompt_nodes(graph), ["1"]) + + def test_graph_without_sampler_or_malformed_input(self): + self.assertEqual(positive_prompt_nodes({"1": pm()}), []) + for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}): + self.assertEqual(positive_prompt_nodes(bad), [], repr(bad)) + + def test_cycles_terminate(self): + loop = { + "class_type": "Weird", + "inputs": { + "conditioning": ["9", 0], + "conditioning_1": ["9", 0], + "positive": ["9", 0], + }, + } + graph = {"9": loop, "5": ksampler("9", "9"), "20": LOADER} + self.assertEqual(positive_prompt_nodes(graph), []) + + +class TestRunPromptNodes(unittest.TestCase): + """The prompt nodes that represent a run, with a safe fallback.""" + + def test_uses_positive_nodes_when_found(self): + graph = { + "176": pm("blurry"), + "134": pm("a cat"), + "5": ksampler("134", "176"), + "20": LOADER, + } + self.assertEqual(run_prompt_nodes(graph), ["134"]) + + def test_unrecognised_sampler_falls_back_to_every_prompt_manager_node(self): + graph = { + "1": pm("a cat"), + "9": {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}}, + "2": {"class_type": "PromptManagerText", "inputs": {"text": "x"}}, + } + self.assertEqual(sorted(run_prompt_nodes(graph)), ["1", "2"]) + + def test_malformed_graph(self): + self.assertEqual(run_prompt_nodes(None), []) + + +class TestResolveText(unittest.TestCase): + """Text known at queue time: typed, or from known pure string nodes.""" + + def _linked(self, source): + return {"1": pm(["2", 0]), "2": source} + + def test_typed_text(self): + self.assertEqual(resolve_text({"1": pm(" a cat ")}, "1"), "a cat") + + def test_core_primitive_string(self): + for cls in ("PrimitiveString", "PrimitiveStringMultiline"): + graph = self._linked({"class_type": cls, "inputs": {"value": "a cat"}}) + self.assertEqual(resolve_text(graph, "1"), "a cat", cls) + + def test_core_string_concatenate(self): + graph = self._linked( + { + "class_type": "StringConcatenate", + "inputs": { + "string_a": "a cat", + "string_b": "at dusk", + "delimiter": ", ", + }, + } + ) + self.assertEqual(resolve_text(graph, "1"), "a cat, at dusk") + + def test_was_text_concatenate_skips_empty_and_cleans_whitespace(self): + graph = self._linked( + { + "class_type": "Text Concatenate", + "inputs": { + "delimiter": " ", + "clean_whitespace": "true", + "text_a": " a cat ", + "text_b": "", + "text_c": ["3", 0], + }, + } + ) + graph["3"] = {"class_type": "PrimitiveString", "inputs": {"value": "at dusk"}} + self.assertEqual(resolve_text(graph, "1"), "a cat at dusk") + + def test_pass_through_nodes(self): + for cls in ("ShowText|pysssss", "Text Multiline"): + graph = self._linked({"class_type": cls, "inputs": {"text": "a cat"}}) + self.assertEqual(resolve_text(graph, "1"), "a cat", cls) + + def test_unknown_and_batch_sources_are_never_guessed(self): + for source in ( + {"class_type": "PromptSearchList", "inputs": {"search_text": "cat"}}, + { + "class_type": "SomeTextReplace", + "inputs": {"text": "a cat", "find": "cat", "replace": "dog"}, + }, + ): + self.assertIsNone( + resolve_text(self._linked(source), "1"), source["class_type"] + ) + + def test_cycles_and_missing_links(self): + graph = { + "1": pm(["2", 0]), + "2": { + "class_type": "StringConcatenate", + "inputs": {"string_a": ["2", 0], "string_b": "x", "delimiter": ""}, + }, + } + self.assertIsNone(resolve_text(graph, "1")) + self.assertIsNone(resolve_text({"1": pm(["9", 0])}, "1")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py new file mode 100644 index 00000000..5a0ce6ca --- /dev/null +++ b/tests/test_prompt_usage.py @@ -0,0 +1,230 @@ +"""Tests for prompt usage tracking: run_count, last_used_at and server-side sorting.""" + +import os +import sqlite3 +import sys +import tempfile +import unittest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase +from prompt_manager_base import PromptManagerBase +from utils.hashing import generate_prompt_hash + + +class UsageTestCase(unittest.TestCase): + def setUp(self): + fd, self.path = tempfile.mkstemp(suffix=".db") + os.close(fd) + self.db = PromptDatabase(self.path) + + def tearDown(self): + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.path + suffix): + os.unlink(self.path + suffix) + + def _save(self, text, **kwargs): + return self.db.save_prompt( + text=text, prompt_hash=generate_prompt_hash(text), **kwargs + ) + + def _set(self, prompt_id, **columns): + with sqlite3.connect(self.path) as conn: + for column, value in columns.items(): + conn.execute( + f"UPDATE prompts SET {column} = ? WHERE id = ?", (value, prompt_id) + ) + + def _ids(self, prompts): + return [p["id"] for p in prompts] + + +class TestRecordPromptUse(UsageTestCase): + def test_new_prompt_starts_unused_but_timestamped(self): + prompt = self.db.get_prompt_by_id(self._save("fresh")) + self.assertEqual(prompt["run_count"], 0) + self.assertTrue(prompt["last_used_at"]) + + def test_record_use_increments_and_touches_last_used(self): + pid = self._save("used") + self._set(pid, last_used_at="2020-01-01T00:00:00.000+00:00") + + self.assertTrue(self.db.record_prompt_use(pid)) + self.assertTrue(self.db.record_prompt_use(pid)) + + prompt = self.db.get_prompt_by_id(pid) + self.assertEqual(prompt["run_count"], 2) + self.assertGreater(prompt["last_used_at"], "2020-01-01T00:00:00.000+00:00") + + def test_rerun_right_after_another_save_still_sorts_first(self): + # No hand-set timestamps: creation and re-run happen within the same millisecond + old = self._save("old") + self._save("new") + self.db.record_prompt_use(old) + result = self.db.get_recent_prompts(limit=1, sort="last_used_desc") + self.assertEqual(self._ids(result["prompts"]), [old]) + + def test_microsecond_timestamps_order_correctly_against_backfilled_ones(self): + backfilled = self._save("backfilled") + self._set(backfilled, last_used_at="2026-01-01T00:00:00.000+00:00") + recent = self._save("recent") + self.db.record_prompt_use(recent) + self.assertRegex( + self.db.get_prompt_by_id(recent)["last_used_at"], + r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d\.\d{6}\+00:00$", + ) + result = self.db.get_recent_prompts(limit=2, sort="last_used_desc") + self.assertEqual(self._ids(result["prompts"]), [recent, backfilled]) + + def test_record_use_can_add_several_runs_at_once(self): + pid = self._save("several") + self.db.record_prompt_use(pid, times=3) + self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 3) + + def test_record_use_on_missing_prompt_returns_false(self): + self.assertFalse(self.db.record_prompt_use(999999)) + + +class TestServerSideSort(UsageTestCase): + def setUp(self): + super().setUp() + self.old = self._save("old prompt", rating=5) + self.mid = self._save("middle prompt", rating=1) + self.new = self._save("new prompt") + self._set(self.old, last_used_at="2026-01-01T00:00:00.000+00:00", run_count=9) + self._set(self.mid, last_used_at="2026-01-02T00:00:00.000+00:00", run_count=1) + self._set(self.new, last_used_at="2026-01-03T00:00:00.000+00:00", run_count=3) + + def test_rerun_bubbles_old_prompt_to_top_across_pages(self): + self.db.record_prompt_use(self.old) + page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="last_used_desc") + self.assertEqual(self._ids(page1["prompts"]), [self.old]) + + def test_most_used(self): + result = self.db.get_recent_prompts(limit=10, sort="run_count_desc") + self.assertEqual(self._ids(result["prompts"]), [self.old, self.new, self.mid]) + + def test_rating_sort_puts_unrated_last_and_spans_pages(self): + page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="rating_desc") + self.assertEqual(self._ids(page1["prompts"]), [self.old]) + result = self.db.get_recent_prompts(limit=10, sort="rating_asc") + self.assertEqual(self._ids(result["prompts"]), [self.mid, self.old, self.new]) + + def test_text_sort(self): + result = self.db.get_recent_prompts(limit=10, sort="text_asc") + self.assertEqual(self._ids(result["prompts"]), [self.mid, self.new, self.old]) + + def test_default_and_unknown_sort_keep_newest_first(self): + expected = [self.new, self.mid, self.old] + self.assertEqual( + self._ids(self.db.get_recent_prompts(limit=10)["prompts"]), expected + ) + injected = self.db.get_recent_prompts(limit=10, sort="id; DROP TABLE prompts") + self.assertEqual(self._ids(injected["prompts"]), expected) + + def test_search_accepts_sort(self): + self.db.record_prompt_use(self.mid) + results = self.db.search_prompts(text="prompt", sort="last_used_desc") + self.assertEqual(self._ids(results)[0], self.mid) + + +@unittest.skipIf( + sqlite3.sqlite_version_info < (3, 35, 0), + "simulating the old schema needs ALTER TABLE DROP COLUMN (SQLite 3.35+)", +) +class TestUsageMigration(unittest.TestCase): + """Databases from before 3.2.4 get the columns and an estimated backfill.""" + + def setUp(self): + fd, self.path = tempfile.mkstemp(suffix=".db") + os.close(fd) + db = PromptDatabase(self.path) + self.with_images = db.save_prompt(text="has images", prompt_hash="h1") + self.no_images = db.save_prompt(text="no images", prompt_hash="h2") + with sqlite3.connect(self.path) as conn: + conn.execute( + "UPDATE prompts SET created_at = '2026-01-01T10:00:00.000000+00:00'" + ) + # generated_images uses SQLite's space-separated CURRENT_TIMESTAMP format + for i, ts in enumerate(("2026-03-05 08:00:00", "2026-02-01 09:00:00")): + conn.execute( + "INSERT INTO generated_images (prompt_id, image_path, filename, generation_time)" + " VALUES (?, ?, ?, ?)", + (self.with_images, f"/out/{i}.png", f"{i}.png", ts), + ) + # Simulate a pre-3.2.4 schema + for index in ("idx_prompts_last_used", "idx_prompts_run_count"): + conn.execute(f"DROP INDEX IF EXISTS {index}") + conn.execute("ALTER TABLE prompts DROP COLUMN last_used_at") + conn.execute("ALTER TABLE prompts DROP COLUMN run_count") + + def tearDown(self): + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.path + suffix): + os.unlink(self.path + suffix) + + def test_backfill_estimates_usage_from_linked_images(self): + db = PromptDatabase(self.path) + used = db.get_prompt_by_id(self.with_images) + unused = db.get_prompt_by_id(self.no_images) + + self.assertEqual(used["run_count"], 2) + self.assertTrue(used["last_used_at"].startswith("2026-03-05T08:00:00")) + self.assertEqual(unused["run_count"], 1) + self.assertTrue(unused["last_used_at"].startswith("2026-01-01T10:00:00")) + + def test_rows_written_while_downgraded_are_healed_on_next_start(self): + db = PromptDatabase(self.path) + # 3.2.3 does not know the columns: its inserts leave last_used_at NULL + with sqlite3.connect(self.path) as conn: + conn.execute( + "INSERT INTO prompts (text, hash, created_at)" + " VALUES ('from 3.2.3', 'h3', '2026-04-01T00:00:00.000000+00:00')" + ) + legacy_id = conn.execute( + "SELECT id FROM prompts WHERE hash = 'h3'" + ).fetchone()[0] + self.assertIsNone(db.get_prompt_by_id(legacy_id)["last_used_at"]) + + healed = PromptDatabase(self.path).get_prompt_by_id(legacy_id) + self.assertTrue(healed["last_used_at"].startswith("2026-04-01T00:00:00")) + self.assertEqual(healed["run_count"], 1) + + def test_migration_runs_once(self): + db = PromptDatabase(self.path) + db.record_prompt_use(self.no_images) + reopened = PromptDatabase(self.path) + self.assertEqual(reopened.get_prompt_by_id(self.no_images)["run_count"], 2) + + +class TestNodeRecordsUse(UsageTestCase): + """Node saves count a run only when asked; re-runs are counted by the queue hook. + + ComfyUI skips unchanged nodes, so counting at execution misses re-runs; see + tests/test_usage_tracking.py for the hook that counts them. + """ + + def _node(self): + node = PromptManagerBase.__new__(PromptManagerBase) + node.db = self.db + node.logger = __import__("logging").getLogger("test.prompt_usage") + return node + + def test_save_without_count_run_leaves_counting_to_the_queue_hook(self): + node = self._node() + first = node._save_prompt_to_database("a lighthouse at dusk") + second = node._save_prompt_to_database("a lighthouse at dusk") + + self.assertEqual(first, second) + self.assertEqual(self.db.get_prompt_by_id(first)["run_count"], 0) + + def test_count_run_counts_new_and_existing_prompts(self): + node = self._node() + pid = node._save_prompt_to_database("batch item", count_run=True) + node._save_prompt_to_database("batch item", count_run=True) + self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_sort_contract.py b/tests/test_sort_contract.py new file mode 100644 index 00000000..103815fe --- /dev/null +++ b/tests/test_sort_contract.py @@ -0,0 +1,39 @@ +"""The UI's sort options must match the server's whitelist exactly.""" + +import os +import re +import sys +import unittest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) + +from database.operations import SORT_ORDERS + +SORT_MODULE = os.path.join(ROOT, "web", "js", "prompt-list-sort.js") +NODE_EXTENSION = os.path.join(ROOT, "web", "prompt_manager.js") + + +class TestSortContract(unittest.TestCase): + def test_ui_sort_options_match_server_whitelist(self): + with open(SORT_MODULE, encoding="utf-8") as fh: + source = fh.read() + ui_values = set(re.findall(r'value:\s*"([a-z_]+)"', source)) + + self.assertTrue(ui_values, "no sort options found in prompt-list-sort.js") + self.assertEqual(ui_values, set(SORT_ORDERS)) + + def test_node_recent_button_requests_recently_used(self): + # The node extension can't import the shared module (ComfyUI load order is + # not guaranteed), so it sends the key directly; it must be a server key. + with open(NODE_EXTENSION, encoding="utf-8") as fh: + source = fh.read() + match = re.search(r"/prompt_manager/recent\?[^`'\"]*sort=([a-z_]+)", source) + + self.assertIsNotNone(match, "node Recent request does not send a sort") + self.assertEqual(match.group(1), "last_used_desc") + self.assertIn(match.group(1), SORT_ORDERS) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py new file mode 100644 index 00000000..d5740f28 --- /dev/null +++ b/tests/test_usage_tracking.py @@ -0,0 +1,283 @@ +"""Run counting and node roles: only positive PromptManager nodes count and link. + +Runs are counted when a prompt is queued (ComfyUI's on-prompt hook), because +ComfyUI skips unchanged nodes, so node execution can't be relied on to count. +""" + +import os +import sys +import tempfile +import unittest +import unittest.mock as mock + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from database.operations import PromptDatabase +from prompt_manager_base import PromptManagerBase +from utils.hashing import generate_prompt_hash +from utils.prompt_tracker import PromptTracker +from utils.usage_tracking import PendingFirstUse, handle_queued_prompt + +LOADER = { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "m.safetensors"}, +} + + +def workflow(positive_text="a cat", negative_text="blurry"): + """API graph with a positive (#134) and a negative (#176) PromptManager node.""" + return { + "176": {"class_type": "PromptManager", "inputs": {"text": negative_text}}, + "134": {"class_type": "PromptManager", "inputs": {"text": positive_text}}, + "5": { + "class_type": "KSampler", + "inputs": { + "positive": ["134", 0], + "negative": ["176", 0], + "model": ["20", 0], + }, + }, + "20": LOADER, + } + + +class DbTestCase(unittest.TestCase): + def setUp(self): + fd, self.path = tempfile.mkstemp(suffix=".db") + os.close(fd) + self.db = PromptDatabase(self.path) + self.pending = PendingFirstUse() + + def tearDown(self): + for suffix in ("", "-wal", "-shm"): + if os.path.exists(self.path + suffix): + os.unlink(self.path + suffix) + + def _save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + def _runs(self, prompt_id): + return self.db.get_prompt_by_id(prompt_id)["run_count"] + + def _queue(self, graph): + return handle_queued_prompt( + {"prompt": graph, "client_id": "x"}, self.db, self.pending + ) + + +class TestQueueHook(DbTestCase): + def test_counts_the_positive_prompt_only(self): + positive, negative = self._save("a cat"), self._save("blurry") + self._queue(workflow()) + self.assertEqual(self._runs(positive), 1) + self.assertEqual(self._runs(negative), 0) + + def test_requeue_counts_again_even_when_comfyui_skips_the_cached_node(self): + positive = self._save("a cat") + self._queue(workflow()) + self._queue(workflow()) # no node execution in between + self.assertEqual(self._runs(positive), 2) + + def test_new_prompt_is_counted_when_the_node_first_saves_it(self): + self._queue(workflow(positive_text="a brand new prompt")) + self.assertTrue( + self.pending.consume(generate_prompt_hash("a brand new prompt")) + ) + self.assertFalse( + self.pending.consume(generate_prompt_hash("a brand new prompt")) + ) + + def test_same_text_in_two_positive_nodes_counts_once(self): + positive = self._save("a cat") + graph = workflow() + graph["135"] = {"class_type": "PromptManager", "inputs": {"text": "a cat"}} + graph["4"] = { + "class_type": "ConditioningCombine", + "inputs": {"conditioning_1": ["134", 0], "conditioning_2": ["135", 0]}, + } + graph["5"]["inputs"]["positive"] = ["4", 0] + self._queue(graph) + self.assertEqual(self._runs(positive), 1) + + def test_text_from_a_known_string_node_is_counted_at_queue_time(self): + positive = self._save("a cat") + graph = workflow() + graph["134"]["inputs"]["text"] = ["8", 0] + graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}} + self._queue(graph) + self._queue(graph) # cached re-run: the node won't execute + self.assertEqual(self._runs(positive), 2) + + def test_new_prompt_queued_twice_before_first_run_counts_twice(self): + self._queue(workflow(positive_text="queued twice")) + self._queue(workflow(positive_text="queued twice")) + self.assertEqual(self.pending.consume(generate_prompt_hash("queued twice")), 2) + + def test_linked_text_is_left_to_node_execution(self): + positive = self._save("from a batch") + graph = workflow() + graph["134"]["inputs"]["text"] = ["7", 0] + graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} + self._queue(graph) + self.assertEqual(self._runs(positive), 0) + + def test_returns_the_request_unchanged_and_never_raises(self): + request = {"prompt": workflow(), "client_id": "x", "extra_data": {}} + self.assertIs(handle_queued_prompt(request, self.db, self.pending), request) + for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []): + self.assertIs(handle_queued_prompt(bad, self.db, self.pending), bad) + broken_db = mock.Mock() + broken_db.get_prompt_by_hash.side_effect = RuntimeError("db locked") + self.assertIs(handle_queued_prompt(request, broken_db, self.pending), request) + + +class TestPendingFirstUse(unittest.TestCase): + def test_entries_expire(self): + pending = PendingFirstUse(ttl_seconds=10) + with mock.patch("utils.usage_tracking.time.monotonic", return_value=100.0): + pending.add("h") + with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0): + self.assertFalse(pending.consume("h")) + + def test_consume_returns_how_many_runs_were_pending(self): + pending = PendingFirstUse() + pending.add("h") + pending.add("h") + pending.add("h") + self.assertEqual(pending.consume("h"), 3) + self.assertEqual(pending.consume("h"), 0) + + def test_size_is_capped(self): + pending = PendingFirstUse(max_entries=3) + for h in ("a", "b", "c", "d"): + pending.add(h) + self.assertFalse(pending.consume("a")) # oldest evicted + self.assertTrue(pending.consume("d")) + + +class FakeTracker: + def __init__(self): + self.calls = [] + + def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True): + self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue}) + return "exec-1" + + +class TestNodeRoles(DbTestCase): + def _node(self): + node = PromptManagerBase.__new__(PromptManagerBase) + node.db = self.db + node.prompt_tracker = FakeTracker() + node.logger = __import__("logging").getLogger("test.usage_tracking") + node.pending_first_use = self.pending + return node + + def _run(self, node, graph, unique_id, text): + return node._track_prompt_execution( + text=text, + encoding_text=text, + category=None, + tags=None, + additional_data={}, + prompt_graph=graph, + unique_id=unique_id, + ) + + def test_negative_node_saves_but_never_becomes_the_current_prompt(self): + node = self._node() + prompt_id = self._run(node, workflow(), "176", "blurry") + self.assertIsNotNone(self.db.get_prompt_by_id(prompt_id)) + self.assertEqual(node.prompt_tracker.calls, []) + self.assertEqual(self._runs(prompt_id), 0) + + def test_positive_typed_text_is_current_without_queue_push(self): + node = self._node() + self._run(node, workflow(), "134", "a cat") + self.assertEqual( + node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}] + ) + + def test_first_save_after_queue_hook_counts_one_run(self): + self._queue(workflow(positive_text="never seen before")) + prompt_id = self._run( + self._node(), + workflow(positive_text="never seen before"), + "134", + "never seen before", + ) + self.assertEqual(self._runs(prompt_id), 1) + + def test_new_prompt_queued_twice_counts_both_runs_on_first_save(self): + graph = workflow(positive_text="double queued") + self._queue(graph) + self._queue(graph) + prompt_id = self._run(self._node(), graph, "134", "double queued") + self.assertEqual(self._runs(prompt_id), 2) + + def test_text_resolved_at_queue_time_is_not_counted_again_or_queued(self): + graph = workflow() + graph["134"]["inputs"]["text"] = ["8", 0] + graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}} + positive = self._save("a cat") + self._queue(graph) + node = self._node() + self._run(node, graph, "134", "a cat") + self.assertEqual(self._runs(positive), 1) + self.assertEqual( + node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}] + ) + + def test_resolution_mismatch_falls_back_to_counting_and_queueing(self): + # Hook guessed "a cat" but the node received different text: trust the node + graph = workflow() + graph["134"]["inputs"]["text"] = ["8", 0] + graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}} + node = self._node() + prompt_id = self._run(node, graph, "134", "actually a dog") + self.assertEqual(self._runs(prompt_id), 1) + self.assertEqual( + node.prompt_tracker.calls, + [{"text": "actually a dog", "push_to_queue": True}], + ) + + def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self): + positive = self._save("a cat") + self._queue(workflow()) + self._run(self._node(), workflow(), "134", "a cat") + self.assertEqual(self._runs(positive), 1) + + def test_positive_linked_text_pushes_queue_and_counts_each_item(self): + graph = workflow() + graph["134"]["inputs"]["text"] = ["7", 0] + graph["7"] = {"class_type": "PromptSearchList", "inputs": {}} + node = self._node() + first = self._run(node, graph, "134", "batch item one") + self._run(node, graph, "134", "batch item one") + self.assertEqual(self._runs(first), 2) + self.assertTrue(all(c["push_to_queue"] for c in node.prompt_tracker.calls)) + + def test_without_graph_context_keeps_legacy_behaviour(self): + node = self._node() + self._run(node, None, None, "standalone") + self.assertEqual( + node.prompt_tracker.calls, [{"text": "standalone", "push_to_queue": True}] + ) + + +class TestTrackerQueueFlag(unittest.TestCase): + def setUp(self): + self.tracker = PromptTracker(mock.Mock()) + + def test_push_to_queue_false_keeps_the_batch_queue_empty(self): + self.tracker.set_current_prompt("a cat", {"prompt_id": 1}, push_to_queue=False) + self.assertIsNone(self.tracker.pop_next_prompt()) + self.assertEqual(self.tracker.get_current_prompt()["id"], 1) + + def test_default_still_pushes(self): + self.tracker.set_current_prompt("a cat", {"prompt_id": 1}) + self.assertEqual(self.tracker.pop_next_prompt()["id"], 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 892cec81..efcfb7f7 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -34,6 +34,7 @@ from watchdog.events import FileSystemEventHandler from .metadata_extractor import ComfyUIMetadataExtractor from .logging_config import get_logger +from .prompt_graph import resolve_text, run_prompt_nodes class ImageGenerationHandler(FileSystemEventHandler): @@ -146,21 +147,26 @@ class ImageGenerationHandler(FileSystemEventHandler): except Exception as meta_error: self.logger.warning(f"Metadata extraction failed: {meta_error}") - # Strategy 1: Pop from batch queue (most reliable for batch workflows). - # Prompts are queued during CLIP encoding in order; images save in - # the same order, so FIFO pop gives the correct prompt per image. - current_prompt = self.prompt_tracker.pop_next_prompt() + # Strategy 1: The image's own metadata. It records the positive prompt + # that produced this image, independent of node caching and of any + # negative-prompt PromptManager nodes in the workflow. + current_prompt = self._find_prompt_from_metadata(metadata) if current_prompt: self.logger.info( - f"Queue match: prompt {current_prompt['id']} for " + f"Metadata match: prompt {current_prompt['id']} for " f"{os.path.basename(image_path)}" ) - # Strategy 2: Find prompt from image metadata + # Strategy 2: Pop from batch queue. Batch items (text supplied by + # PromptSearchList) share one graph, so metadata can't tell them apart; + # they are queued in encoding order and images save in the same order. if not current_prompt: - current_prompt = self._find_prompt_from_metadata(metadata) + current_prompt = self.prompt_tracker.pop_next_prompt() if current_prompt: - self.logger.info(f"Metadata match: prompt {current_prompt['id']}") + self.logger.info( + f"Queue match: prompt {current_prompt['id']} for " + f"{os.path.basename(image_path)}" + ) # Strategy 3: Use snapshot captured at file-creation time if not current_prompt and prompt_snapshot: @@ -231,21 +237,19 @@ class ImageGenerationHandler(FileSystemEventHandler): prompt_text = None - # Try to find PromptManager node in the prompt execution data + # The executed graph: use the PromptManager node feeding the positive input prompt_data = metadata.get("prompt") if isinstance(prompt_data, dict): - for node_id, node_info in prompt_data.items(): - if not isinstance(node_info, dict): - continue - class_type = node_info.get("class_type", "") - if class_type in ("PromptManager", "PromptManagerText"): - inputs = node_info.get("inputs", {}) - text = inputs.get("text", "") - if text and isinstance(text, str) and text.strip(): - prompt_text = text.strip() - break + for node_id in run_prompt_nodes(prompt_data): + prompt_text = resolve_text(prompt_data, node_id) + if prompt_text: + break + if not prompt_text: + # Positive text can't be known from the graph (batch item or unknown + # node): leave it to the queue rather than guess from workflow widgets + return None - # Fallback: check text_encoder_nodes from workflow, but only if + # Fallback (no executed graph): check text_encoder_nodes from workflow, but only if # the text input is NOT connected (connected inputs override widget values, # so the widget value would be stale in batch workflows). if not prompt_text: diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py new file mode 100644 index 00000000..fbb0df02 --- /dev/null +++ b/utils/prompt_graph.py @@ -0,0 +1,178 @@ +"""Find which PromptManager nodes feed a sampler's positive input in a ComfyUI graph. + +Works on the API-format prompt graph ({node_id: {"class_type", "inputs"}}) that +ComfyUI queues and embeds in saved images. Usage counting and image linking use it +so that negative-prompt PromptManager nodes never count as "the prompt" of a run. + +Pure functions with no ComfyUI imports, so they are unit-testable. Mirrors the +link-tracing rules of web/js/comfy-metadata.js. +""" + +import re +from typing import Any, Dict, List, Optional + +PROMPT_MANAGER_TYPES = frozenset({"PromptManager", "PromptManagerText"}) + +# Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted +MAX_VISITS = 500 + +# Inputs that carry positive conditioning or prompt text towards a sampler +_CONDITIONING_KEY = re.compile(r"^conditioning(_\w+)?$") +_TEXT_KEY = re.compile(r"^(text|string|prompt)(_?[a-z0-9]+)?$", re.IGNORECASE) + + +def _is_link(value: Any) -> bool: + return ( + isinstance(value, (list, tuple)) + and len(value) == 2 + and isinstance(value[0], (str, int)) + and isinstance(value[1], int) + ) + + +def _inputs(node: Any) -> Dict[str, Any]: + inputs = node.get("inputs") if isinstance(node, dict) else None + return inputs if isinstance(inputs, dict) else {} + + +def _positive_roots(graph: Dict[str, Any]) -> List[Any]: + """Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...).""" + roots = [] + for node in graph.values(): + inputs = _inputs(node) + if not _is_link(inputs.get("model")): + continue + if _is_link(inputs.get("positive")) and _is_link(inputs.get("negative")): + roots.append(inputs["positive"]) + elif ( + _is_link(inputs.get("conditioning")) + and "guider" in str(node.get("class_type", "")).lower() + ): + roots.append(inputs["conditioning"]) + return roots + + +def _upstream_links(inputs: Dict[str, Any]) -> List[Any]: + """Inputs to follow from a pass-through or encoder node, never `negative`.""" + links = [] + for key, value in inputs.items(): + if not _is_link(value): + continue + if key == "positive" or _CONDITIONING_KEY.match(key) or _TEXT_KEY.match(key): + links.append(value) + return links + + +def positive_prompt_nodes(graph: Any) -> List[str]: + """Ids of PromptManager/PromptManagerText nodes that feed a positive input. + + Ordered by discovery, each id once. Malformed graphs yield an empty list. + """ + if not isinstance(graph, dict): + return [] + + found: List[str] = [] + seen = set() + stack = list(reversed(_positive_roots(graph))) + while stack and len(seen) < MAX_VISITS: + node_id = str(stack.pop()[0]) + if node_id in seen: + continue + seen.add(node_id) + node = graph.get(node_id) + if not isinstance(node, dict): + continue + if node.get("class_type") in PROMPT_MANAGER_TYPES: + found.append(node_id) + continue + stack.extend(reversed(_upstream_links(_inputs(node)))) + return found + + +def run_prompt_nodes(graph: Any) -> List[str]: + """Prompt nodes that represent a run: the positive ones, if any can be found. + + When no PromptManager node can be traced to a positive input (an unrecognised + custom sampler, say), every PromptManager node counts, as before 3.2.4, so an + exotic workflow never silently loses usage counting or image linking. + """ + positive = positive_prompt_nodes(graph) + if positive or not isinstance(graph, dict): + return positive + return [ + str(node_id) + for node_id, node in graph.items() + if isinstance(node, dict) and node.get("class_type") in PROMPT_MANAGER_TYPES + ] + + +def _join_text_inputs( + inputs: Dict[str, Any], keys: List[str], resolve +) -> Optional[str]: + """WAS-style join: resolve each key, optionally strip, skip empty, join by delimiter.""" + delimiter = inputs.get("delimiter", " ") + if not isinstance(delimiter, str): + return None + clean = str(inputs.get("clean_whitespace", "true")).lower() == "true" + parts = [] + for key in keys: + value = resolve(inputs.get(key, "")) + if value is None: + return None + value = value.strip() if clean else value + if value: + parts.append(value) + return delimiter.join(parts) + + +def _resolve_string_node( + class_type: str, inputs: Dict[str, Any], resolve +) -> Optional[str]: + """Output of a known pure string node, or None when the type isn't known.""" + if class_type in ("PrimitiveString", "PrimitiveStringMultiline"): + return resolve(inputs.get("value")) + if class_type in ("ShowText|pysssss", "Text Multiline", "String Literal"): + return resolve(inputs.get("text", inputs.get("string"))) + if class_type == "StringConcatenate": + a, b = resolve(inputs.get("string_a", "")), resolve(inputs.get("string_b", "")) + delimiter = inputs.get("delimiter", "") + if a is None or b is None or not isinstance(delimiter, str): + return None + return f"{a}{delimiter}{b}" + if class_type == "Text Concatenate": + keys = sorted(k for k in inputs if re.fullmatch(r"text_[a-z]", k)) + return _join_text_inputs(inputs, keys, resolve) + return None + + +def resolve_text(graph: Any, node_id: str) -> Optional[str]: + """A prompt node's `text` as it will be at run time, if knowable from the graph. + + Typed text, or text produced by known pure string nodes (core primitives and + StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText). Anything + else, including PromptSearchList batches, returns None rather than a guess. + Callers compare the result with the text the node actually received. + """ + if not isinstance(graph, dict): + return None + seen = set() + + def resolve(value: Any) -> Optional[str]: + if isinstance(value, str): + return value + if not _is_link(value) or len(seen) >= MAX_VISITS: + return None + source_id = str(value[0]) + if source_id in seen: + return None + seen.add(source_id) + source = graph.get(source_id) + if not isinstance(source, dict): + return None + return _resolve_string_node( + str(source.get("class_type", "")), _inputs(source), resolve + ) + + seen.add(str(node_id)) + text = resolve(_inputs(graph.get(str(node_id))).get("text")) + return text.strip() if isinstance(text, str) and text.strip() else None diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index 2c32298c..91242dd8 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -99,7 +99,10 @@ class PromptTracker: self.logger.debug("PromptTracker initialization completed") def set_current_prompt( - self, prompt_text: str, additional_data: Optional[Dict[str, Any]] = None + self, + prompt_text: str, + additional_data: Optional[Dict[str, Any]] = None, + push_to_queue: bool = True, ) -> str: """ Set the current prompt for this thread for image tracking. @@ -113,6 +116,9 @@ class PromptTracker: Args: prompt_text: The prompt text being executed additional_data: Additional prompt metadata, should include prompt_id from PromptManager + push_to_queue: Also push onto the FIFO batch queue. Only batch items whose + text comes from another node need it; typed prompts are linked from + the image's own metadata instead. Returns: Unique execution ID for this prompt execution @@ -164,11 +170,12 @@ class PromptTracker: self.active_prompts[execution_id] = execution_context # Push to batch queue for ordered image linking - with self._queue_lock: - self._prompt_queue.append(execution_context) - self.logger.debug( - f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}" - ) + if push_to_queue: + with self._queue_lock: + self._prompt_queue.append(execution_context) + self.logger.debug( + f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}" + ) self.logger.debug( f"Set current prompt: {execution_id} -> {prompt_text[:50]}... (thread: {threading.current_thread().ident})" diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py new file mode 100644 index 00000000..4aadf6f2 --- /dev/null +++ b/utils/usage_tracking.py @@ -0,0 +1,107 @@ +"""Count prompt runs when a workflow is queued, not when nodes execute. + +ComfyUI skips nodes whose inputs are unchanged, so a PromptManager node does not +execute when the same prompt is re-run. The on-prompt hook sees every queued +workflow instead: it counts a run for each PromptManager node feeding a sampler's +positive input. Prompts not yet in the database are remembered as pending and +counted by the node when it first saves them. + +Text from known pure string nodes is resolved at queue time too. Anything else +(PromptSearchList batches, unknown nodes) is counted by the node as it executes. +""" + +import threading +import time +from collections import OrderedDict +from typing import Any, Tuple + +try: + from .hashing import generate_prompt_hash + from .logging_config import get_logger + from .prompt_graph import resolve_text, run_prompt_nodes +except ImportError: + from utils.hashing import generate_prompt_hash + from utils.logging_config import get_logger + from utils.prompt_graph import resolve_text, run_prompt_nodes + +logger = get_logger("prompt_manager.usage_tracking") + +PENDING_TTL_SECONDS = 3600 # a queued prompt that never executes is forgotten +PENDING_MAX_ENTRIES = 1000 + + +class PendingFirstUse: + """Queued runs of positive prompts that were not in the database yet. + + Counts per hash: the same new prompt can be queued several times before its + first job runs, and later jobs may reuse ComfyUI's cached node. + """ + + def __init__( + self, + ttl_seconds: float = PENDING_TTL_SECONDS, + max_entries: int = PENDING_MAX_ENTRIES, + ): + self._ttl = ttl_seconds + self._max = max_entries + self._entries: OrderedDict[str, Tuple[int, float]] = OrderedDict() + self._lock = threading.Lock() + + def add(self, prompt_hash: str) -> None: + with self._lock: + count, _ = self._entries.pop(prompt_hash, (0, 0.0)) + self._entries[prompt_hash] = (count + 1, time.monotonic()) + while len(self._entries) > self._max: + self._entries.popitem(last=False) + + def consume(self, prompt_hash: str) -> int: + """Number of recent queued runs not yet counted (0 if none); clears them.""" + with self._lock: + count, added = self._entries.pop(prompt_hash, (0, 0.0)) + return count if count and time.monotonic() - added <= self._ttl else 0 + + +# Shared between the queue hook and the nodes +PENDING_FIRST_USE = PendingFirstUse() + + +def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> Any: + """ComfyUI on-prompt handler body: count runs, always return the request as-is.""" + try: + graph = json_data.get("prompt") if isinstance(json_data, dict) else None + counted = set() + for node_id in run_prompt_nodes(graph): + text = resolve_text(graph, node_id) + if not text: + continue + prompt_hash = generate_prompt_hash(text) + if prompt_hash in counted: + continue + counted.add(prompt_hash) + existing = db.get_prompt_by_hash(prompt_hash) + if existing: + db.record_prompt_use(existing["id"]) + else: + pending.add(prompt_hash) + except Exception as e: + logger.warning(f"Could not count queued prompt runs: {e}") + return json_data + + +_hook_registered = threading.Event() +_register_lock = threading.Lock() + + +def register_queue_hook( + server: Any, db_factory: Any, pending: PendingFirstUse = PENDING_FIRST_USE +) -> bool: + """Register the on-prompt handler once; db_factory is called lazily per request.""" + with _register_lock: + if _hook_registered.is_set(): + return False + server.add_on_prompt_handler( + lambda json_data: handle_queued_prompt(json_data, db_factory(), pending) + ) + _hook_registered.set() + logger.info("Registered prompt usage hook") + return True diff --git a/web/admin.html b/web/admin.html index 6bb850aa..6431ebc6 100644 --- a/web/admin.html +++ b/web/admin.html @@ -107,12 +107,8 @@
@@ -1391,6 +1387,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
+ diff --git a/web/js/admin.js b/web/js/admin.js index 8be9f068..651f5cb7 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -1,6 +1,7 @@ class PromptAdmin { constructor() { this.prompts = []; + this.listMode = "recent"; // "recent" or "search"; a sort change reloads it this.selectedPrompts = new Set(); this.settings = { resultTimeout: 5, @@ -114,8 +115,13 @@ document.getElementById(id).addEventListener("change", () => this.search()); }); - // Sort dropdown change - document.getElementById("sortBy").addEventListener("change", () => this.renderPrompts()); + // Sort dropdown: options come from PromptListSort; sorting runs on the server + const sortSelect = document.getElementById("sortBy"); + sortSelect.innerHTML = PromptListSort.SORT_OPTIONS + .map((o) => ``) + .join(""); + sortSelect.value = PromptListSort.DEFAULT_SORT; + sortSelect.addEventListener("change", () => this.reloadPrompts()); } bindModalEvents() { @@ -355,9 +361,13 @@ async loadRecentPrompts(page = 1) { try { + this.listMode = "recent"; this.pagination.currentPage = page; - const offset = (page - 1) * this.pagination.limit; - const response = await fetch(`/prompt_manager/recent?limit=${this.pagination.limit}&offset=${offset}&page=${page}`); + const response = await fetch(PromptListSort.buildRecentUrl({ + page, + limit: this.pagination.limit, + sort: this.currentSort(), + })); if (response.ok) { const data = await response.json(); if (data.success) { @@ -424,7 +434,8 @@ if (folder) params.append("folder", folder); params.append("limit", "100"); - const response = await fetch(`/prompt_manager/search?${params}`); + this.listMode = "search"; + const response = await fetch(`/prompt_manager/search?${PromptListSort.withSort(params, this.currentSort())}`); if (response.ok) { const data = await response.json(); if (data.success) { @@ -442,36 +453,13 @@ } } - sortPrompts(prompts) { - const sortBy = document.getElementById("sortBy")?.value || "created_desc"; - const [field, direction] = sortBy.split("_"); - - return [...prompts].sort((a, b) => { - let aVal, bVal; - - switch (field) { - case "rating": - aVal = a.rating || 0; - bVal = b.rating || 0; - break; - case "created": - aVal = new Date(a.created_at); - bVal = new Date(b.created_at); - break; - case "text": - aVal = a.text.toLowerCase(); - bVal = b.text.toLowerCase(); - break; - default: - return 0; - } - - if (direction === "asc") { - return aVal > bVal ? 1 : aVal < bVal ? -1 : 0; - } else { - return aVal < bVal ? 1 : aVal > bVal ? -1 : 0; - } - }); + currentSort() { + return PromptListSort.normalizeSort(document.getElementById("sortBy")?.value); + } + + /** Re-fetch the list currently shown, e.g. after the sort changes. */ + reloadPrompts() { + return this.listMode === "search" ? this.search() : this.loadRecentPrompts(1); } renderPrompts() { @@ -493,14 +481,13 @@ return; } - // Sort prompts before rendering - const sortedPrompts = this.sortPrompts(this.prompts); - container.innerHTML = sortedPrompts.map((prompt) => this.renderPromptItem(prompt)).join(""); + // Already in the selected order: the server sorts across all pages + container.innerHTML = this.prompts.map((prompt) => this.renderPromptItem(prompt)).join(""); this.selectedPrompts.clear(); this.updateBulkActionButtons(); // Add hover behavior to all star ratings - sortedPrompts.forEach(prompt => { + this.prompts.forEach(prompt => { this.addStarHoverBehavior(prompt.id); }); @@ -514,7 +501,7 @@ }); // Load film strips for each prompt (async, non-blocking) - this.loadAllFilmStrips(sortedPrompts); + this.loadAllFilmStrips(this.prompts); } async loadAllFilmStrips(prompts) { @@ -550,6 +537,10 @@ const category = prompt.category || "No category"; const rating = prompt.rating || 0; const created = new Date(prompt.created_at).toLocaleDateString(); + const runLabel = PromptListSort.formatRunCount(prompt.run_count); + const lastUsedTitle = prompt.last_used_at + ? `Last used ${new Date(prompt.last_used_at).toLocaleString()}` + : ""; return `
@@ -574,6 +565,10 @@ 📅 ${created}
+ ${runLabel ? `
+ 🔁 + ${this.escapeHtml(runLabel)} +
` : ""}
${this.renderStars(rating, prompt.id)} diff --git a/web/js/prompt-list-sort.js b/web/js/prompt-list-sort.js new file mode 100644 index 00000000..8d10c52c --- /dev/null +++ b/web/js/prompt-list-sort.js @@ -0,0 +1,64 @@ +/** + * Prompt list sorting for the admin dashboard (3.2.4, #90.1). + * + * Sorting happens on the server so it spans every page; this module only owns the + * option list, validation and request URLs. The option values must match + * SORT_ORDERS in database/operations.py (enforced by tests/test_sort_contract.py). + * + * Loadable in the browser (window.PromptListSort) and in Node for tests. + */ +(function (root) { + "use strict"; + + const SORT_OPTIONS = Object.freeze( + [ + { value: "last_used_desc", label: "Recently Used" }, + { value: "created_desc", label: "Newest First" }, + { value: "created_asc", label: "Oldest First" }, + { value: "run_count_desc", label: "Most Used" }, + { value: "rating_desc", label: "Highest Rated" }, + { value: "rating_asc", label: "Lowest Rated" }, + { value: "text_asc", label: "A-Z" }, + { value: "text_desc", label: "Z-A" }, + ].map((option) => Object.freeze(option)), + ); + + const DEFAULT_SORT = "last_used_desc"; + const DEFAULT_LIMIT = 50; + const KNOWN = new Set(SORT_OPTIONS.map((o) => o.value)); + + /** A known sort key, or DEFAULT_SORT for anything else. */ + function normalizeSort(value) { + return typeof value === "string" && KNOWN.has(value) ? value : DEFAULT_SORT; + } + + /** URL for one page of /prompt_manager/recent. */ + function buildRecentUrl({ page = 1, limit = DEFAULT_LIMIT, sort } = {}) { + const params = new URLSearchParams({ + limit: String(limit), + offset: String((page - 1) * limit), + page: String(page), + sort: normalizeSort(sort), + }); + return `/prompt_manager/recent?${params}`; + } + + /** Copy of `params` with a validated sort key; the input is not modified. */ + function withSort(params, sort) { + const copy = new URLSearchParams(params); + copy.set("sort", normalizeSort(sort)); + return copy; + } + + /** Badge text for a prompt's run count; empty when the count is unknown. */ + function formatRunCount(count) { + if (!Number.isInteger(count) || count < 0) return ""; + if (count === 0) return "Never run"; + if (count === 1) return "1 run"; + return `${count.toLocaleString("en-US")} runs`; + } + + const api = { SORT_OPTIONS, DEFAULT_SORT, normalizeSort, buildRecentUrl, withSort, formatRunCount }; + if (typeof module !== "undefined" && module.exports) module.exports = api; + else root.PromptListSort = api; +})(typeof window !== "undefined" ? window : globalThis); diff --git a/web/prompt_manager.js b/web/prompt_manager.js index caa896c1..4834a3dd 100644 --- a/web/prompt_manager.js +++ b/web/prompt_manager.js @@ -978,7 +978,7 @@ app.registerExtension({ url = `/prompt_manager/search?${queryParams.toString()}`; } else if (method === "get_recent_prompts") { const limit = params.limit || 20; - url = `/prompt_manager/recent?limit=${limit}`; + url = `/prompt_manager/recent?limit=${limit}&sort=last_used_desc`; } else { throw new Error(`Unknown method: ${method}`); }