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

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

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()

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):

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