From 729190309cbdb087e89f4e6ffeaf5405cd247d86 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Wed, 7 Oct 2026 18:38:41 -0700 Subject: [PATCH] fix: count prompt runs when ComfyUI dequeues the prompt, not when it is submitted The on-prompt hook runs before ComfyUI validates the request, so counting there bumped run_count for prompts that were rejected or deleted from the queue. The hook now only resolves the positive PromptManager texts, mints a UUID4 prompt_id when the client sent none, and stashes the hashes under that id (QueuedRuns, bounded to 1000 entries). PromptQueue.get is wrapped idempotently; when a prompt is dequeued for execution its stashed runs are counted (or remembered as pending first use). Servers without a prompt_queue fall back to counting at submit time with one warning. __init__ defines _global_db before the monitoring block and passes a factory that raises RuntimeError while the database is unavailable; the package is now imported under test with server and folder_paths stubbed. --- __init__.py | 16 +- tests/test_init_package.py | 178 +++++++++++++++++++++ tests/test_usage_tracking.py | 301 +++++++++++++++++++++++++++++++++-- utils/usage_tracking.py | 218 +++++++++++++++++++++---- 4 files changed, 671 insertions(+), 42 deletions(-) create mode 100644 tests/test_init_package.py diff --git a/__init__.py b/__init__.py index 2e07092..43f608f 100644 --- a/__init__.py +++ b/__init__.py @@ -82,6 +82,10 @@ except Exception as e: except Exception: pass +# Shared database, set up with the monitoring system below. Stays None when that +# fails so the queue hook can report the database as unavailable instead of crashing. +_global_db = None + # Start image monitoring globally at module import time # This ensures images are linked regardless of which node is used in the workflow try: @@ -111,13 +115,21 @@ 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 + +def _db_factory(): + """The shared database for the queue hook; called lazily per request.""" + if _global_db is None: + raise RuntimeError("PromptManager database unavailable") + return _global_db + + +# Count prompt runs when ComfyUI dequeues workflows (it skips unchanged nodes, so # node execution alone misses re-runs). Separate from monitoring so either can fail. try: from .py.config import server_instance from .utils.usage_tracking import register_queue_hook - register_queue_hook(server_instance, lambda: _global_db) + register_queue_hook(server_instance, _db_factory) except Exception as e: try: from .utils.logging_config import get_logger diff --git a/tests/test_init_package.py b/tests/test_init_package.py new file mode 100644 index 0000000..619df19 --- /dev/null +++ b/tests/test_init_package.py @@ -0,0 +1,178 @@ +"""Importing the extension package wires the queue hook to the shared database. + +The package is loaded the way ComfyUI loads it (``__init__`` with relative +imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed. +""" + +import importlib.util +import os +import sys +import tempfile +import types +import unittest +import unittest.mock as mock + +from aiohttp import web + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, ROOT) + + +class FakePromptQueue: + def __init__(self): + self.results = [] + + def get(self, timeout=None): + return self.results.pop(0) if self.results else None + + +def fake_server_module(): + """A ``server`` module whose PromptServer.instance accepts routes and hooks.""" + instance = types.SimpleNamespace( + routes=web.RouteTableDef(), + prompt_queue=FakePromptQueue(), + handlers=[], + ) + instance.add_on_prompt_handler = instance.handlers.append + module = types.ModuleType("server") + module.PromptServer = type("PromptServer", (), {"instance": instance}) + return module + + +def fake_folder_paths_module(output_dir): + module = types.ModuleType("folder_paths") + module.base_path = output_dir + module.models_dir = os.path.join(output_dir, "models") + module.get_output_directory = lambda: output_dir + module.get_folder_paths = lambda name: [] + return module + + +def load_package(name, db_path, output_dir): + """Import ROOT/__init__.py as package ``name`` with ComfyUI stubbed. + + Returns (package, server instance). The caller stops the image monitor. + """ + saved = {key: sys.modules.get(key) for key in ("server", "folder_paths")} + server = fake_server_module() + sys.modules["server"] = server + sys.modules["folder_paths"] = fake_folder_paths_module(output_dir) + try: + from py.config import PromptManagerConfig + + with mock.patch.object(PromptManagerConfig, "DEFAULT_DB_PATH", db_path): + spec = importlib.util.spec_from_file_location( + name, + os.path.join(ROOT, "__init__.py"), + submodule_search_locations=[ROOT], + ) + package = importlib.util.module_from_spec(spec) + sys.modules[name] = package + try: + spec.loader.exec_module(package) + except BaseException: + sys.modules.pop(name, None) + raise + finally: + for key, module in saved.items(): + if module is None: + sys.modules.pop(key, None) + else: + sys.modules[key] = module + return package, server.PromptServer.instance + + +class PackageTestCase(unittest.TestCase): + package_name = "prompt_manager_pkg" + + @classmethod + def setUpClass(cls): + cls.tmp = tempfile.TemporaryDirectory() + cls.output_dir = os.path.join(cls.tmp.name, "output") + os.makedirs(cls.output_dir) + cls.db_path = cls.make_db_path(cls.tmp.name) + cls.package, cls.server = load_package( + cls.package_name, cls.db_path, cls.output_dir + ) + + @classmethod + def make_db_path(cls, tmp): + return os.path.join(tmp, "prompts.db") + + @classmethod + def tearDownClass(cls): + monitor = getattr(cls.package, "_global_image_monitor", None) + if monitor is not None: + monitor.stop_monitoring() + sys.modules.pop(cls.package_name, None) + cls.tmp.cleanup() + + +class TestPackageWiring(PackageTestCase): + def test_nodes_are_registered(self): + self.assertEqual( + sorted(self.package.NODE_CLASS_MAPPINGS), + ["PromptManager", "PromptManagerText", "PromptSearchList"], + ) + + def test_db_factory_returns_the_shared_database(self): + self.assertIsNotNone(self.package._global_db) + self.assertIs(self.package._db_factory(), self.package._global_db) + self.assertEqual( + os.path.normcase(os.path.realpath(self.package._global_db.model.db_path)), + os.path.normcase(os.path.realpath(self.db_path)), + ) + + def test_db_factory_raises_while_the_database_is_unavailable(self): + with mock.patch.object(self.package, "_global_db", None): + with self.assertRaisesRegex(RuntimeError, "database unavailable"): + self.package._db_factory() + + def test_queue_hook_is_registered_and_the_queue_get_is_wrapped(self): + self.assertEqual(len(self.server.handlers), 1) + self.assertTrue( + getattr(self.server.prompt_queue.get, "_prompt_manager_dequeue_hook", False) + ) + + def test_image_monitor_watches_the_stubbed_output_directory(self): + status = self.package._global_image_monitor.get_status() + self.assertTrue(status["running"]) + self.assertEqual( + [ + os.path.normcase(os.path.realpath(d)) + for d in status["monitored_directories"] + ], + [os.path.normcase(os.path.realpath(self.output_dir))], + ) + + +class TestPackageWithoutDatabase(PackageTestCase): + """The database failing to initialise must not break the queue hook.""" + + package_name = "prompt_manager_pkg_nodb" + + @classmethod + def make_db_path(cls, tmp): + blocker = os.path.join(tmp, "blocker") + with open(blocker, "w", encoding="utf-8") as f: + f.write("not a directory") + return os.path.join(blocker, "prompts.db") # parent is a file: init fails + + def test_global_db_is_defined_but_unset(self): + self.assertIsNone(self.package._global_db) + with self.assertRaisesRegex(RuntimeError, "database unavailable"): + self.package._db_factory() + + def test_queue_hook_still_registers_and_survives_a_dequeue(self): + self.assertEqual(len(self.server.handlers), 1) + request = self.server.handlers[0]( + {"prompt": {"1": {"class_type": "PromptManager", "inputs": {"text": "x"}}}} + ) + self.server.prompt_queue.results.append( + ((0, request["prompt_id"], request["prompt"], {}, [], False), 1) + ) + self.assertIsNotNone(self.server.prompt_queue.get()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py index d5740f2..152ca91 100644 --- a/tests/test_usage_tracking.py +++ b/tests/test_usage_tracking.py @@ -1,14 +1,18 @@ """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. +Runs are resolved when a prompt is queued (ComfyUI's on-prompt hook, which runs +before validation) and counted when ComfyUI dequeues the prompt for execution, +because ComfyUI skips unchanged nodes, so node execution can't be relied on to +count, and a rejected or cancelled queue entry must not count either. """ import os import sys import tempfile +import threading import unittest import unittest.mock as mock +import uuid sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -16,7 +20,16 @@ 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 +from utils import usage_tracking +from utils.usage_tracking import ( + PendingFirstUse, + QueuedRuns, + _wrap_queue_get, + count_queued_prompt, + handle_queued_prompt, + on_dequeue, + register_queue_hook, +) LOADER = { "class_type": "CheckpointLoaderSimple", @@ -41,14 +54,52 @@ def workflow(positive_text="a cat", negative_text="blurry"): } +class FakePromptQueue: + """Stand-in for execution.PromptQueue: get() pops preset results.""" + + def __init__(self): + self.results = [] + self.get_calls = [] + + def get(self, timeout=None): + self.get_calls.append(timeout) + return self.results.pop(0) if self.results else None + + +class FakeServer: + """Stand-in for PromptServer: on-prompt handlers plus an optional queue.""" + + def __init__(self, with_queue=True): + self.handlers = [] + if with_queue: + self.prompt_queue = FakePromptQueue() + + def add_on_prompt_handler(self, handler): + self.handlers.append(handler) + + def on_prompt(self, json_data): + for handler in self.handlers: + json_data = handler(json_data) + return json_data + + +def queue_item(prompt_id, graph): + """What PromptQueue.get returns for a prompt that passed validation.""" + return ((0, prompt_id, graph, {}, [], False), 1) + + class DbTestCase(unittest.TestCase): def setUp(self): fd, self.path = tempfile.mkstemp(suffix=".db") os.close(fd) self.db = PromptDatabase(self.path) self.pending = PendingFirstUse() + self.queued = QueuedRuns() def tearDown(self): + close = getattr(self.db, "close", None) + if callable(close): + close() for suffix in ("", "-wal", "-shm"): if os.path.exists(self.path + suffix): os.unlink(self.path + suffix) @@ -60,9 +111,12 @@ class DbTestCase(unittest.TestCase): 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 + """Queue a graph and let ComfyUI dequeue it for execution.""" + request = handle_queued_prompt({"prompt": graph, "client_id": "x"}, self.queued) + on_dequeue( + request.get("prompt_id", ""), self.queued, lambda: self.db, self.pending ) + return request class TestQueueHook(DbTestCase): @@ -121,14 +175,243 @@ class TestQueueHook(DbTestCase): self._queue(graph) self.assertEqual(self._runs(positive), 0) - def test_returns_the_request_unchanged_and_never_raises(self): + def test_returns_the_request_and_never_raises(self): request = {"prompt": workflow(), "client_id": "x", "extra_data": {}} - self.assertIs(handle_queued_prompt(request, self.db, self.pending), request) + self.assertIs(handle_queued_prompt(request, self.queued), request) for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []): - self.assertIs(handle_queued_prompt(bad, self.db, self.pending), bad) + self.assertIs(handle_queued_prompt(bad, self.queued), 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) + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"): + on_dequeue( + request["prompt_id"], self.queued, lambda: broken_db, self.pending + ) + + def test_graph_without_prompt_manager_nodes_is_not_stashed_or_tagged(self): + request = {"prompt": {"20": LOADER}} + self.assertIs(handle_queued_prompt(request, self.queued), request) + self.assertNotIn("prompt_id", request) + self.assertEqual(len(self.queued), 0) + + +class TestCountAtDequeue(DbTestCase): + """End to end through the registered hook and the wrapped queue.""" + + def setUp(self): + super().setUp() + self.server = FakeServer() + self._register(self.server) + + def _register(self, server): + with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()): + return register_queue_hook( + server, lambda: self.db, pending=self.pending, queued=self.queued + ) + + def _submit(self, graph, prompt_id=None): + request = {"prompt": graph, "client_id": "x"} + if prompt_id is not None: + request["prompt_id"] = prompt_id + return self.server.on_prompt(request) + + def _dequeue(self, request): + self.server.prompt_queue.results.append( + queue_item(request["prompt_id"], request["prompt"]) + ) + return self.server.prompt_queue.get(timeout=1.0) + + def test_hook_adds_a_lowercase_uuid4_prompt_id_when_absent(self): + request = self._submit(workflow()) + prompt_id = request["prompt_id"] + self.assertEqual(str(uuid.UUID(prompt_id, version=4)), prompt_id) + self.assertEqual(prompt_id, prompt_id.lower()) + + def test_hook_keeps_an_existing_prompt_id(self): + request = self._submit(workflow(), prompt_id="client-chosen") + self.assertEqual(request["prompt_id"], "client-chosen") + + def test_hook_alone_does_not_count(self): + positive = self._save("a cat") + self._submit(workflow()) + self.assertEqual(self._runs(positive), 0) + + def test_dequeue_counts_exactly_once(self): + positive = self._save("a cat") + request = self._submit(workflow()) + result = self._dequeue(request) + self.assertEqual(result[0][1], request["prompt_id"]) + self.assertEqual(self.server.prompt_queue.get_calls, [1.0]) + self.assertEqual(self._runs(positive), 1) + self._dequeue(request) # a second dequeue of the same id has nothing left + self.assertEqual(self._runs(positive), 1) + + def test_two_queued_prompts_are_each_counted_once(self): + positive = self._save("a cat") + first = self._submit(workflow()) + second = self._submit(workflow()) + self.assertNotEqual(first["prompt_id"], second["prompt_id"]) + self._dequeue(first) + self.assertEqual(self._runs(positive), 1) + self._dequeue(second) + self.assertEqual(self._runs(positive), 2) + + def test_rejected_prompt_is_never_counted(self): + positive = self._save("a cat") + self._submit(workflow()) # validation fails: never dequeued + other = self._submit(workflow()) + self._dequeue(other) + self.assertEqual(self._runs(positive), 1) + + def test_get_returning_none_is_a_noop(self): + positive = self._save("a cat") + self._submit(workflow()) + self.assertIsNone(self.server.prompt_queue.get(timeout=0.1)) + self.assertEqual(self._runs(positive), 0) + + def test_unknown_text_is_pending_only_after_dequeue(self): + prompt_hash = generate_prompt_hash("never saved") + request = self._submit(workflow(positive_text="never saved")) + self.assertEqual(self.pending.consume(prompt_hash), 0) + self._dequeue(request) + self.assertEqual(self.pending.consume(prompt_hash), 1) + + def test_shorter_and_longer_queue_items_are_tolerated(self): + positive = self._save("a cat") + request = self._submit(workflow()) + self.server.prompt_queue.results.append( + ((0, request["prompt_id"], request["prompt"], {}, []), 1) + ) + self.server.prompt_queue.get() + self.assertEqual(self._runs(positive), 1) + request = self._submit(workflow()) + self.server.prompt_queue.results.append( + ((0, request["prompt_id"], request["prompt"], {}, [], False, "x"), 2) + ) + self.server.prompt_queue.get() + self.assertEqual(self._runs(positive), 2) + + def test_malformed_queue_item_is_logged_and_returned(self): + item = (("not", "a", "tuple"),) + self.server.prompt_queue.results.append(item) + self.server.prompt_queue.results.append("garbage") + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"): + self.assertEqual(self.server.prompt_queue.get(), item) + self.assertEqual(self.server.prompt_queue.get(), "garbage") + + def test_db_factory_raising_at_dequeue_is_logged_not_raised(self): + def broken(): + raise RuntimeError("PromptManager database unavailable") + + queued = QueuedRuns() + server = FakeServer() + with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()): + register_queue_hook(server, broken, pending=self.pending, queued=queued) + request = server.on_prompt({"prompt": workflow()}) + server.prompt_queue.results.append(queue_item(request["prompt_id"], {})) + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs: + self.assertIsNotNone(server.prompt_queue.get()) + self.assertIn("database unavailable", "\n".join(logs.output)) + + def test_server_without_prompt_queue_counts_at_queue_time_and_warns_once(self): + positive = self._save("a cat") + server = FakeServer(with_queue=False) + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING") as logs: + self.assertTrue(self._register(server)) + warnings = [line for line in logs.output if line.startswith("WARNING")] + self.assertEqual(len(warnings), 1) + server.on_prompt({"prompt": workflow()}) + self.assertEqual(self._runs(positive), 1) + server.on_prompt({"prompt": workflow(positive_text="new text")}) + self.assertEqual(self.pending.consume(generate_prompt_hash("new text")), 1) + + def test_server_whose_queue_has_no_get_uses_legacy_counting(self): + positive = self._save("a cat") + server = FakeServer(with_queue=False) + server.prompt_queue = object() + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"): + self._register(server) + server.on_prompt({"prompt": workflow()}) + self.assertEqual(self._runs(positive), 1) + + def test_legacy_counting_calls_db_factory_inside_the_guard(self): + def broken(): + raise RuntimeError("PromptManager database unavailable") + + request = {"prompt": workflow()} + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"): + self.assertIs(count_queued_prompt(request, broken, self.pending), request) + + def test_wrapping_twice_does_not_double_count(self): + positive = self._save("a cat") + prompt_queue = self.server.prompt_queue + wrapped_get = prompt_queue.get + calls = [] + self.assertFalse(_wrap_queue_get(prompt_queue, calls.append)) + self.assertIs(prompt_queue.get, wrapped_get) + request = self._submit(workflow()) + self._dequeue(request) + self.assertEqual(self._runs(positive), 1) + self.assertEqual(calls, []) + + def test_wrap_rejects_queues_without_a_callable_get(self): + self.assertFalse(_wrap_queue_get(object(), lambda prompt_id: None)) + self.assertFalse(_wrap_queue_get(None, lambda prompt_id: None)) + + def test_wrapper_reports_the_dequeued_prompt_id(self): + prompt_queue = FakePromptQueue() + seen = [] + self.assertTrue(_wrap_queue_get(prompt_queue, seen.append)) + prompt_queue.results.append(queue_item("abc", {})) + self.assertEqual(prompt_queue.get(timeout=2.5), queue_item("abc", {})) + self.assertEqual(prompt_queue.get_calls, [2.5]) + self.assertEqual(seen, ["abc"]) + + def test_dequeue_callback_errors_are_logged_not_raised(self): + prompt_queue = FakePromptQueue() + + def explode(prompt_id): + raise ValueError("boom") + + _wrap_queue_get(prompt_queue, explode) + prompt_queue.results.append(queue_item("abc", {})) + with self.assertLogs("prompt_manager.usage_tracking", level="WARNING"): + self.assertIsNotNone(prompt_queue.get()) + + +class TestQueuedRuns(unittest.TestCase): + def test_pop_returns_runs_once(self): + queued = QueuedRuns() + queued.put("p1", [("h1", "text one")]) + self.assertEqual(queued.pop("p1"), [("h1", "text one")]) + self.assertEqual(queued.pop("p1"), []) + self.assertEqual(queued.pop("missing"), []) + + def test_bound_evicts_oldest(self): + queued = QueuedRuns(max_entries=3) + for i in range(4): + queued.put(f"p{i}", [(f"h{i}", f"text {i}")]) + self.assertEqual(len(queued), 3) + self.assertEqual(queued.pop("p0"), []) + self.assertEqual(queued.pop("p3"), [("h3", "text 3")]) + + def test_default_bound_is_one_thousand(self): + queued = QueuedRuns() + for i in range(1001): + queued.put(str(i), [("h", "t")]) + self.assertEqual(len(queued), 1000) + self.assertEqual(queued.pop("0"), []) + self.assertEqual(queued.pop("1000"), [("h", "t")]) + + def test_requeue_of_the_same_id_replaces_the_runs(self): + queued = QueuedRuns() + queued.put("p1", [("h1", "one")]) + queued.put("p1", [("h2", "two")]) + self.assertEqual(queued.pop("p1"), [("h2", "two")]) + + def test_empty_runs_are_not_stored(self): + queued = QueuedRuns() + queued.put("p1", []) + self.assertEqual(len(queued), 0) class TestPendingFirstUse(unittest.TestCase): diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 4aadf6f..7dcaadf 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -1,19 +1,30 @@ -"""Count prompt runs when a workflow is queued, not when nodes execute. +"""Count prompt runs when ComfyUI dequeues a workflow, 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. +execute when the same prompt is re-run. Instead: + +1. The on-prompt hook sees every submitted workflow before validation. It resolves + the text of each PromptManager node feeding a sampler's positive input and + stashes the hashes under the request's prompt_id (minting one when the client + sent none; ComfyUI keeps a client-supplied id). +2. ``PromptQueue.get`` is wrapped. When ComfyUI dequeues a prompt for execution, + its stashed hashes are counted: one run per prompt already in the database, + otherwise a pending first use that the node consumes when it first saves it. + +A prompt that fails validation or is deleted from the queue is never dequeued, so +it never counts. Servers without a ``prompt_queue`` fall back to counting at +submit time. 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 functools import threading import time +import uuid from collections import OrderedDict -from typing import Any, Tuple +from typing import Any, Callable, List, Optional, Tuple try: from .hashing import generate_prompt_hash @@ -28,13 +39,16 @@ logger = get_logger("prompt_manager.usage_tracking") PENDING_TTL_SECONDS = 3600 # a queued prompt that never executes is forgotten PENDING_MAX_ENTRIES = 1000 +QUEUED_MAX_ENTRIES = 1000 # prompts submitted but not yet dequeued + +Run = Tuple[str, str] # (prompt hash, resolved text) class PendingFirstUse: - """Queued runs of positive prompts that were not in the database yet. + """Dequeued 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. + Counts per hash: the same new prompt can be run several times before its + first job saves it, and later jobs may reuse ComfyUI's cached node. """ def __init__( @@ -61,47 +75,189 @@ class PendingFirstUse: return count if count and time.monotonic() - added <= self._ttl else 0 +class QueuedRuns: + """Runs resolved at submit time, keyed by prompt_id until ComfyUI dequeues them. + + Bounded: a prompt rejected by validation or deleted from the queue is never + dequeued, so the oldest entries are evicted once the bound is reached. + """ + + def __init__(self, max_entries: int = QUEUED_MAX_ENTRIES): + self._max = max_entries + self._entries: OrderedDict[str, List[Run]] = OrderedDict() + self._lock = threading.Lock() + + def put(self, prompt_id: str, runs: List[Run]) -> None: + if not runs: + return + with self._lock: + self._entries.pop(prompt_id, None) + self._entries[prompt_id] = list(runs) + while len(self._entries) > self._max: + self._entries.popitem(last=False) + + def pop(self, prompt_id: str) -> List[Run]: + with self._lock: + return self._entries.pop(prompt_id, []) + + def __len__(self) -> int: + with self._lock: + return len(self._entries) + + # Shared between the queue hook and the nodes PENDING_FIRST_USE = PendingFirstUse() +QUEUED_RUNS = QueuedRuns() -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.""" +def queued_prompt_runs(graph: Any) -> List[Run]: + """(hash, text) of each positive PromptManager node whose text is knowable. + + Each hash once: the same text in two positive nodes is one run of the prompt. + """ + runs: List[Run] = [] + seen = 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 seen: + continue + seen.add(prompt_hash) + runs.append((prompt_hash, text)) + return runs + + +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: + pending.add(prompt_hash) + + +def handle_queued_prompt(json_data: Any, queued: QueuedRuns = QUEUED_RUNS) -> Any: + """ComfyUI on-prompt handler body: stash the runs, always return the request. + + Adds a canonical UUID4 ``prompt_id`` when the client sent none, so the stash + can be matched when ComfyUI dequeues the prompt. The database is not touched: + the request has not been validated yet. + """ 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) + runs = queued_prompt_runs(graph) + if not runs: + return json_data + prompt_id = json_data.get("prompt_id") + if prompt_id is None: + prompt_id = str(uuid.uuid4()) + json_data["prompt_id"] = prompt_id + queued.put(str(prompt_id), runs) + except Exception as e: + logger.warning(f"Could not resolve queued prompt runs: {e}") + return json_data + + +def on_dequeue( + prompt_id: str, + queued: QueuedRuns, + db_factory: Callable[[], Any], + pending: PendingFirstUse, +) -> None: + """Count the runs stashed for a prompt ComfyUI is about to execute.""" + try: + runs = queued.pop(str(prompt_id)) + if runs: + count_runs(runs, db_factory(), pending) + except Exception as e: + logger.warning(f"Could not count runs of dequeued prompt {prompt_id}: {e}") + + +def count_queued_prompt( + json_data: Any, db_factory: Callable[[], Any], pending: PendingFirstUse +) -> Any: + """Legacy on-prompt handler body: count at submit time (no prompt_queue).""" + try: + graph = json_data.get("prompt") if isinstance(json_data, dict) else None + runs = queued_prompt_runs(graph) + if runs: + count_runs(runs, db_factory(), pending) except Exception as e: logger.warning(f"Could not count queued prompt runs: {e}") return json_data +_WRAPPED_MARKER = "_prompt_manager_dequeue_hook" + + +def _wrap_queue_get(prompt_queue: Any, on_dequeued: Callable[[str], None]) -> bool: + """Wrap ``prompt_queue.get`` so each dequeued prompt_id is reported once. + + The wrapper forwards every argument and result unchanged. Idempotent: a queue + whose ``get`` is already wrapped is left alone. + """ + original = getattr(prompt_queue, "get", None) + if not callable(original): + return False + if getattr(original, _WRAPPED_MARKER, False): + return False + + @functools.wraps(original) + def get(*args, **kwargs): + result = original(*args, **kwargs) + if result is not None: + try: + on_dequeued(str(result[0][1])) + except Exception as e: + logger.warning(f"Dequeue hook failed: {e}") + return result + + setattr(get, _WRAPPED_MARKER, True) + prompt_queue.get = get + return True + + _hook_registered = threading.Event() _register_lock = threading.Lock() def register_queue_hook( - server: Any, db_factory: Any, pending: PendingFirstUse = PENDING_FIRST_USE + server: Any, + db_factory: Callable[[], Any], + pending: PendingFirstUse = PENDING_FIRST_USE, + queued: Optional[QueuedRuns] = None, ) -> bool: - """Register the on-prompt handler once; db_factory is called lazily per request.""" + """Register the on-prompt handler once and wrap the prompt queue's ``get``. + + ``db_factory`` is called lazily, inside the guarded handlers, so a database + that is unavailable never raises into ComfyUI. Without a usable + ``server.prompt_queue`` runs are counted at submit time instead. + """ + if queued is None: + queued = QUEUED_RUNS 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) - ) + prompt_queue = getattr(server, "prompt_queue", None) + if callable(getattr(prompt_queue, "get", None)): + server.add_on_prompt_handler( + lambda json_data: handle_queued_prompt(json_data, queued) + ) + _wrap_queue_get( + prompt_queue, + lambda prompt_id: on_dequeue(prompt_id, queued, db_factory, pending), + ) + logger.info("Registered prompt usage hook (counting at dequeue)") + else: + logger.warning( + "PromptServer has no prompt_queue; counting prompt runs at queue " + "time, so rejected prompts may be counted" + ) + server.add_on_prompt_handler( + lambda json_data: count_queued_prompt(json_data, db_factory, pending) + ) _hook_registered.set() - logger.info("Registered prompt usage hook") return True