fix: link images and count runs by the positive prompt, not node execution

With a negative-prompt PromptManager node in the workflow, every executed
PromptManager node pushed onto the image-linking queue, so images could be
linked to the negative prompt or to a previous run. ComfyUI also skips
unchanged nodes, so re-runs of the same prompt were never counted.

- utils/prompt_graph.py finds the PromptManager nodes feeding a sampler's
  positive input (KSampler, CFG/Basic guiders, ControlNet and combine
  pass-throughs, PromptManagerText into encoders), falling back to all
  PromptManager nodes when none can be traced.
- Runs are counted by a ComfyUI on-prompt hook when a workflow is queued;
  prompts not yet saved are counted when the node first saves them.
- Nodes receive the queued graph (hidden PROMPT/UNIQUE_ID). Only the run's
  positive prompt becomes the current prompt; only batch items with linked
  text are pushed onto the FIFO queue.
- The image monitor links from the image's own metadata first, using the
  positive node, and falls back to the queue for batch items.

Refs #90

__init__.py

@@ -111,6 +111,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
+# 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:
+        print(f"[ComfyUI-PromptManager] Warning: Failed to register usage hook: {e}")
+
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
 
 # Print startup message with loaded tools

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

prompt_manager_base.py

@@ -24,6 +24,8 @@ 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.prompt_graph import is_text_linked, run_prompt_nodes
+    from .utils.usage_tracking import PENDING_FIRST_USE
 except ImportError:
     import sys
 
@@ -32,6 +34,8 @@ 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.prompt_graph import is_text_linked, run_prompt_nodes
+    from utils.usage_tracking import PENDING_FIRST_USE
 
 
 class PromptManagerBase:
@@ -56,14 +60,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,7 +95,8 @@ class PromptManagerBase:
                         prompt_id=existing["id"], category=category, tags=tags
                     )
                     self.logger.debug("Updated metadata for existing prompt")
-                self._record_use(existing["id"])
+                if count_run:
+                    self._record_use(existing["id"])
                 return existing["id"]
 
             self.logger.debug(
@@ -94,7 +108,9 @@ class PromptManagerBase:
 
             if prompt_id:
                 self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
-                self._record_use(prompt_id)
+                pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
+                if count_run or pending.consume(prompt_hash):
+                    self._record_use(prompt_id)
             else:
                 self.logger.warning("Failed to save prompt - no ID returned")
 
@@ -104,6 +120,53 @@ class PromptManagerBase:
             self.logger.error(f"Error saving prompt to database: {e}")
             return None
 
+    def _prompt_role(self, prompt_graph: Any, unique_id: Any) -> 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)
+            text_linked: text comes from another node (e.g. PromptSearchList)
+            known: the graph and node id were available
+        """
+        if not isinstance(prompt_graph, dict) or unique_id is None:
+            return {"is_run_prompt": True, "text_linked": False, "known": False}
+        node_id = str(unique_id)
+        return {
+            "is_run_prompt": node_id in run_prompt_nodes(prompt_graph),
+            "text_linked": is_text_linked(prompt_graph, node_id),
+            "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. Typed prompts are linked from the image's own metadata, so only
+        batch items (linked text) are pushed onto the FIFO queue. Without graph
+        context the pre-3.2.4 behaviour is kept.
+        """
+        role = self._prompt_role(prompt_graph, unique_id)
+        batch_item = role["is_run_prompt"] and role["text_linked"]
+        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) -> None:
         """Count this execution; a tracking failure must never block generation."""
         try:

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

tests/test_image_linking.py

@@ -0,0 +1,131 @@
+"""Image linking picks the positive prompt, preferring the image's own metadata."""
+
+import os
+import sys
+import tempfile
+import unittest
+from unittest import 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_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()

tests/test_node_tracking_wiring.py

@@ -0,0 +1,82 @@
+"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once."""
+
+import os
+import sys
+import unittest
+from unittest import 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, "_registered", False):
+            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()

tests/test_prompt_graph.py

@@ -0,0 +1,191 @@
+"""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 (
+    is_text_linked,
+    literal_text,
+    positive_prompt_nodes,
+    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 TestNodeText(unittest.TestCase):
+    def test_literal_and_linked_text(self):
+        graph = {
+            "1": pm("  a cat  "),
+            "2": pm(["7", 0]),
+            "7": {"class_type": "PromptSearchList", "inputs": {}},
+        }
+        self.assertEqual(literal_text(graph, "1"), "a cat")
+        self.assertIsNone(literal_text(graph, "2"))
+        self.assertFalse(is_text_linked(graph, "1"))
+        self.assertTrue(is_text_linked(graph, "2"))
+
+    def test_missing_nodes(self):
+        self.assertIsNone(literal_text({}, "1"))
+        self.assertFalse(is_text_linked({}, "1"))
+        self.assertIsNone(literal_text({"1": pm("   ")}, "1"))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_prompt_usage.py

@@ -194,7 +194,11 @@ class TestUsageMigration(unittest.TestCase):
 
 
 class TestNodeRecordsUse(UsageTestCase):
-    """Every node execution counts, including re-runs of an existing prompt."""
+    """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)
@@ -202,13 +206,19 @@ class TestNodeRecordsUse(UsageTestCase):
         node.logger = __import__("logging").getLogger("test.prompt_usage")
         return node
 
-    def test_new_and_repeated_runs_are_counted(self):
+    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"], 2)
+        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__":

tests/test_usage_tracking.py

@@ -0,0 +1,228 @@
+"""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
+from unittest import 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_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_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_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()

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 literal_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 = literal_text(prompt_data, node_id)
+                if prompt_text:
+                    break
+            if not prompt_text:
+                # Positive text came from another node (batch item): leave it to the
+                # queue rather than guess from workflow widgets, which don't know roles
+                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:

utils/prompt_graph.py

@@ -0,0 +1,121 @@
+"""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 literal_text(graph: Any, node_id: str) -> Optional[str]:
+    """A node's `text` input when typed in (not linked), stripped; else None."""
+    node = graph.get(str(node_id)) if isinstance(graph, dict) else None
+    text = _inputs(node).get("text")
+    if isinstance(text, str) and text.strip():
+        return text.strip()
+    return None
+
+
+def is_text_linked(graph: Any, node_id: str) -> bool:
+    """True when the node's `text` comes from another node (e.g. PromptSearchList)."""
+    node = graph.get(str(node_id)) if isinstance(graph, dict) else None
+    return _is_link(_inputs(node).get("text"))

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})"

utils/usage_tracking.py

@@ -0,0 +1,104 @@
+"""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 supplied by another node (PromptSearchList batches) is unknown at queue time;
+those runs are counted by the node as each item executes.
+"""
+
+import threading
+import time
+from collections import OrderedDict
+from typing import Any
+
+try:
+    from .hashing import generate_prompt_hash
+    from .logging_config import get_logger
+    from .prompt_graph import literal_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 literal_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:
+    """Hashes of queued positive prompts that were not in the database yet."""
+
+    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, float]" = OrderedDict()
+        self._lock = threading.Lock()
+
+    def add(self, prompt_hash: str) -> None:
+        with self._lock:
+            self._entries.pop(prompt_hash, None)
+            self._entries[prompt_hash] = time.monotonic()
+            while len(self._entries) > self._max:
+                self._entries.popitem(last=False)
+
+    def consume(self, prompt_hash: str) -> bool:
+        """True (once) if this hash was queued recently and not yet counted."""
+        with self._lock:
+            added = self._entries.pop(prompt_hash, None)
+        return added is not None and time.monotonic() - added <= self._ttl
+
+
+# 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 = literal_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
+
+
+_registered = False
+_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."""
+    global _registered
+    with _register_lock:
+        if _registered:
+            return False
+        server.add_on_prompt_handler(
+            lambda json_data: handle_queued_prompt(json_data, db_factory(), pending)
+        )
+        _registered = True
+        logger.info("Registered prompt usage hook")
+        return True