fix: address review — count cached runs of linked text, keep repeated pending runs

- Text from known pure string nodes (core PrimitiveString/StringConcatenate,
  WAS Text Concatenate/Multiline, pysssss ShowText) is resolved at queue
  time, so cached re-runs of such prompts are counted and their images link
  from metadata. Unknown nodes are never guessed; the node compares its real
  text with the queue-time resolution and falls back to counting and
  queueing itself on any mismatch.
- A new prompt queued several times before its first run keeps every
  pending run; the first save records all of them.
- Registration flag is a threading.Event; consistent unittest imports;
  logging instead of print in the init fallback.

__init__.py

@@ -124,7 +124,11 @@ except Exception as e:
 
         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}")
+        import logging
+
+        logging.getLogger("prompt_manager.init").error(
+            f"Failed to register usage hook: {e}"
+        )
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
 

database/operations.py

@@ -163,21 +163,22 @@ class PromptDatabase:
             self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
             return prompt_id
 
-    def record_prompt_use(self, prompt_id: int) -> bool:
+    def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
         """
-        Count one run of a prompt and mark it as most recently used.
+        Count runs of a prompt and mark it as most recently used.
 
         Args:
             prompt_id: The prompt ID
+            times: Number of runs to add (queued runs counted on first save)
 
         Returns:
             True if the prompt exists and was updated
         """
         with self.model.get_connection() as conn:
             cursor = conn.execute(
-                "UPDATE prompts SET run_count = run_count + 1, last_used_at = ?"
+                "UPDATE prompts SET run_count = run_count + ?, last_used_at = ?"
                 " WHERE id = ?",
-                (utc_now_iso(), prompt_id),
+                (max(1, int(times)), utc_now_iso(), prompt_id),
             )
             conn.commit()
             return cursor.rowcount > 0

prompt_manager_base.py

@@ -24,7 +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.hashing import generate_prompt_hash
+    from .utils.prompt_graph import resolve_text, run_prompt_nodes
     from .utils.usage_tracking import PENDING_FIRST_USE
 except ImportError:
     import sys
@@ -34,7 +35,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.hashing import generate_prompt_hash
+    from utils.prompt_graph import resolve_text, run_prompt_nodes
     from utils.usage_tracking import PENDING_FIRST_USE
 
 
@@ -109,8 +111,9 @@ class PromptManagerBase:
             if prompt_id:
                 self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
                 pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
-                if count_run or pending.consume(prompt_hash):
-                    self._record_use(prompt_id)
+                runs = pending.consume(prompt_hash) + (1 if count_run else 0)
+                if runs:
+                    self._record_use(prompt_id, times=runs)
             else:
                 self.logger.warning("Failed to save prompt - no ID returned")
 
@@ -120,20 +123,25 @@ 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]:
+    def _prompt_role(
+        self, prompt_graph: Any, unique_id: Any, text: str
+    ) -> Dict[str, bool]:
         """How this node's execution relates to the run, from the queued graph.
 
         Returns flags:
             is_run_prompt: feeds a sampler's positive input (images belong to it)
-            text_linked: text comes from another node (e.g. PromptSearchList)
+            seen_at_queue: the queue hook resolved exactly this text, so it already
+                counted the run and the image's metadata can identify the prompt
             known: the graph and node id were available
         """
         if not isinstance(prompt_graph, dict) or unique_id is None:
-            return {"is_run_prompt": True, "text_linked": False, "known": False}
+            return {"is_run_prompt": True, "seen_at_queue": False, "known": False}
         node_id = str(unique_id)
+        resolved = resolve_text(prompt_graph, node_id)
         return {
             "is_run_prompt": node_id in run_prompt_nodes(prompt_graph),
-            "text_linked": is_text_linked(prompt_graph, node_id),
+            "seen_at_queue": resolved is not None
+            and generate_prompt_hash(resolved) == generate_prompt_hash(text),
             "known": True,
         }
 
@@ -150,12 +158,15 @@ class PromptManagerBase:
         """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.
+        linked to. Text the queue hook could resolve is counted there and linked from
+        the image's own metadata; anything else (PromptSearchList batches, unknown
+        text nodes, or a resolution that doesn't match) is counted here and pushed
+        onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept.
         """
-        role = self._prompt_role(prompt_graph, unique_id)
-        batch_item = role["is_run_prompt"] and role["text_linked"]
+        role = self._prompt_role(prompt_graph, unique_id, text)
+        batch_item = (
+            role["known"] and role["is_run_prompt"] and not role["seen_at_queue"]
+        )
         prompt_id = self._save_prompt_to_database(
             text=text, category=category, tags=tags, count_run=batch_item
         )
@@ -167,10 +178,10 @@ class PromptManagerBase:
             )
         return prompt_id
 
-    def _record_use(self, prompt_id: int) -> None:
-        """Count this execution; a tracking failure must never block generation."""
+    def _record_use(self, prompt_id: int, times: int = 1) -> None:
+        """Count runs; a tracking failure must never block generation."""
         try:
-            self.db.record_prompt_use(prompt_id)
+            self.db.record_prompt_use(prompt_id, times=times)
         except Exception as e:
             self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}")
 

tests/test_image_linking.py

@@ -4,7 +4,7 @@ import os
 import sys
 import tempfile
 import unittest
-from unittest import mock
+import unittest.mock as mock
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -87,6 +87,16 @@ class TestMetadataLinking(LinkingTestCase):
             self.tracker.pop_next_prompt()["id"], stale, "queue left intact"
         )
 
+    def test_text_from_a_known_string_node_links_from_metadata(self):
+        positive = self._save("a cat")
+        linked_graph = graph()
+        linked_graph["134"]["inputs"]["text"] = ["8", 0]
+        linked_graph["8"] = {
+            "class_type": "PrimitiveString",
+            "inputs": {"value": "a cat"},
+        }
+        self.assertEqual(self._process({"prompt": linked_graph}), positive)
+
     def test_batch_items_with_linked_text_use_the_queue(self):
         first, second = self._save("item one"), self._save("item two")
         batch_graph = graph()

tests/test_node_tracking_wiring.py

@@ -2,8 +2,9 @@
 
 import os
 import sys
+import threading
 import unittest
-from unittest import mock
+import unittest.mock as mock
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -68,7 +69,7 @@ class TestQueueHookRegistration(unittest.TestCase):
         server = mock.Mock()
         db = mock.Mock()
         db.get_prompt_by_hash.return_value = None
-        with mock.patch.object(usage_tracking, "_registered", False):
+        with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
             self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db))
             self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db))
 

tests/test_prompt_graph.py

@@ -7,9 +7,8 @@ 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,
+    resolve_text,
     run_prompt_nodes,
 )
 
@@ -169,22 +168,76 @@ class TestRunPromptNodes(unittest.TestCase):
         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"))
+class TestResolveText(unittest.TestCase):
+    """Text known at queue time: typed, or from known pure string nodes."""
 
-    def test_missing_nodes(self):
-        self.assertIsNone(literal_text({}, "1"))
-        self.assertFalse(is_text_linked({}, "1"))
-        self.assertIsNone(literal_text({"1": pm("   ")}, "1"))
+    def _linked(self, source):
+        return {"1": pm(["2", 0]), "2": source}
+
+    def test_typed_text(self):
+        self.assertEqual(resolve_text({"1": pm("  a cat ")}, "1"), "a cat")
+
+    def test_core_primitive_string(self):
+        for cls in ("PrimitiveString", "PrimitiveStringMultiline"):
+            graph = self._linked({"class_type": cls, "inputs": {"value": "a cat"}})
+            self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
+
+    def test_core_string_concatenate(self):
+        graph = self._linked(
+            {
+                "class_type": "StringConcatenate",
+                "inputs": {
+                    "string_a": "a cat",
+                    "string_b": "at dusk",
+                    "delimiter": ", ",
+                },
+            }
+        )
+        self.assertEqual(resolve_text(graph, "1"), "a cat, at dusk")
+
+    def test_was_text_concatenate_skips_empty_and_cleans_whitespace(self):
+        graph = self._linked(
+            {
+                "class_type": "Text Concatenate",
+                "inputs": {
+                    "delimiter": " ",
+                    "clean_whitespace": "true",
+                    "text_a": " a cat ",
+                    "text_b": "",
+                    "text_c": ["3", 0],
+                },
+            }
+        )
+        graph["3"] = {"class_type": "PrimitiveString", "inputs": {"value": "at dusk"}}
+        self.assertEqual(resolve_text(graph, "1"), "a cat at dusk")
+
+    def test_pass_through_nodes(self):
+        for cls in ("ShowText|pysssss", "Text Multiline"):
+            graph = self._linked({"class_type": cls, "inputs": {"text": "a cat"}})
+            self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
+
+    def test_unknown_and_batch_sources_are_never_guessed(self):
+        for source in (
+            {"class_type": "PromptSearchList", "inputs": {"search_text": "cat"}},
+            {
+                "class_type": "SomeTextReplace",
+                "inputs": {"text": "a cat", "find": "cat", "replace": "dog"},
+            },
+        ):
+            self.assertIsNone(
+                resolve_text(self._linked(source), "1"), source["class_type"]
+            )
+
+    def test_cycles_and_missing_links(self):
+        graph = {
+            "1": pm(["2", 0]),
+            "2": {
+                "class_type": "StringConcatenate",
+                "inputs": {"string_a": ["2", 0], "string_b": "x", "delimiter": ""},
+            },
+        }
+        self.assertIsNone(resolve_text(graph, "1"))
+        self.assertIsNone(resolve_text({"1": pm(["9", 0])}, "1"))
 
 
 if __name__ == "__main__":

tests/test_prompt_usage.py

@@ -77,6 +77,11 @@ class TestRecordPromptUse(UsageTestCase):
         result = self.db.get_recent_prompts(limit=2, sort="last_used_desc")
         self.assertEqual(self._ids(result["prompts"]), [recent, backfilled])
 
+    def test_record_use_can_add_several_runs_at_once(self):
+        pid = self._save("several")
+        self.db.record_prompt_use(pid, times=3)
+        self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 3)
+
     def test_record_use_on_missing_prompt_returns_false(self):
         self.assertFalse(self.db.record_prompt_use(999999))
 

tests/test_usage_tracking.py

@@ -8,7 +8,7 @@ import os
 import sys
 import tempfile
 import unittest
-from unittest import mock
+import unittest.mock as mock
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
@@ -99,6 +99,20 @@ class TestQueueHook(DbTestCase):
         self._queue(graph)
         self.assertEqual(self._runs(positive), 1)
 
+    def test_text_from_a_known_string_node_is_counted_at_queue_time(self):
+        positive = self._save("a cat")
+        graph = workflow()
+        graph["134"]["inputs"]["text"] = ["8", 0]
+        graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
+        self._queue(graph)
+        self._queue(graph)  # cached re-run: the node won't execute
+        self.assertEqual(self._runs(positive), 2)
+
+    def test_new_prompt_queued_twice_before_first_run_counts_twice(self):
+        self._queue(workflow(positive_text="queued twice"))
+        self._queue(workflow(positive_text="queued twice"))
+        self.assertEqual(self.pending.consume(generate_prompt_hash("queued twice")), 2)
+
     def test_linked_text_is_left_to_node_execution(self):
         positive = self._save("from a batch")
         graph = workflow()
@@ -125,6 +139,14 @@ class TestPendingFirstUse(unittest.TestCase):
         with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0):
             self.assertFalse(pending.consume("h"))
 
+    def test_consume_returns_how_many_runs_were_pending(self):
+        pending = PendingFirstUse()
+        pending.add("h")
+        pending.add("h")
+        pending.add("h")
+        self.assertEqual(pending.consume("h"), 3)
+        self.assertEqual(pending.consume("h"), 0)
+
     def test_size_is_capped(self):
         pending = PendingFirstUse(max_entries=3)
         for h in ("a", "b", "c", "d"):
@@ -186,6 +208,39 @@ class TestNodeRoles(DbTestCase):
         )
         self.assertEqual(self._runs(prompt_id), 1)
 
+    def test_new_prompt_queued_twice_counts_both_runs_on_first_save(self):
+        graph = workflow(positive_text="double queued")
+        self._queue(graph)
+        self._queue(graph)
+        prompt_id = self._run(self._node(), graph, "134", "double queued")
+        self.assertEqual(self._runs(prompt_id), 2)
+
+    def test_text_resolved_at_queue_time_is_not_counted_again_or_queued(self):
+        graph = workflow()
+        graph["134"]["inputs"]["text"] = ["8", 0]
+        graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
+        positive = self._save("a cat")
+        self._queue(graph)
+        node = self._node()
+        self._run(node, graph, "134", "a cat")
+        self.assertEqual(self._runs(positive), 1)
+        self.assertEqual(
+            node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
+        )
+
+    def test_resolution_mismatch_falls_back_to_counting_and_queueing(self):
+        # Hook guessed "a cat" but the node received different text: trust the node
+        graph = workflow()
+        graph["134"]["inputs"]["text"] = ["8", 0]
+        graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
+        node = self._node()
+        prompt_id = self._run(node, graph, "134", "actually a dog")
+        self.assertEqual(self._runs(prompt_id), 1)
+        self.assertEqual(
+            node.prompt_tracker.calls,
+            [{"text": "actually a dog", "push_to_queue": True}],
+        )
+
     def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self):
         positive = self._save("a cat")
         self._queue(workflow())

utils/image_monitor.py

@@ -34,7 +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
+from .prompt_graph import resolve_text, run_prompt_nodes
 
 
 class ImageGenerationHandler(FileSystemEventHandler):
@@ -241,12 +241,12 @@ class ImageGenerationHandler(FileSystemEventHandler):
         prompt_data = metadata.get("prompt")
         if isinstance(prompt_data, dict):
             for node_id in run_prompt_nodes(prompt_data):
-                prompt_text = literal_text(prompt_data, node_id)
+                prompt_text = resolve_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
+                # Positive text can't be known from the graph (batch item or unknown
+                # node): leave it to the queue rather than guess from workflow widgets
                 return None
 
         # Fallback (no executed graph): check text_encoder_nodes from workflow, but only if

utils/prompt_graph.py

@@ -106,16 +106,73 @@ def run_prompt_nodes(graph: Any) -> List[str]:
     ]
 
 
-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()
+def _join_text_inputs(
+    inputs: Dict[str, Any], keys: List[str], resolve
+) -> Optional[str]:
+    """WAS-style join: resolve each key, optionally strip, skip empty, join by delimiter."""
+    delimiter = inputs.get("delimiter", " ")
+    if not isinstance(delimiter, str):
+        return None
+    clean = str(inputs.get("clean_whitespace", "true")).lower() == "true"
+    parts = []
+    for key in keys:
+        value = resolve(inputs.get(key, ""))
+        if value is None:
+            return None
+        value = value.strip() if clean else value
+        if value:
+            parts.append(value)
+    return delimiter.join(parts)
+
+
+def _resolve_string_node(
+    class_type: str, inputs: Dict[str, Any], resolve
+) -> Optional[str]:
+    """Output of a known pure string node, or None when the type isn't known."""
+    if class_type in ("PrimitiveString", "PrimitiveStringMultiline"):
+        return resolve(inputs.get("value"))
+    if class_type in ("ShowText|pysssss", "Text Multiline", "String Literal"):
+        return resolve(inputs.get("text", inputs.get("string")))
+    if class_type == "StringConcatenate":
+        a, b = resolve(inputs.get("string_a", "")), resolve(inputs.get("string_b", ""))
+        delimiter = inputs.get("delimiter", "")
+        if a is None or b is None or not isinstance(delimiter, str):
+            return None
+        return f"{a}{delimiter}{b}"
+    if class_type == "Text Concatenate":
+        keys = sorted(k for k in inputs if re.fullmatch(r"text_[a-z]", k))
+        return _join_text_inputs(inputs, keys, resolve)
     return None
 
 
-def 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"))
+def resolve_text(graph: Any, node_id: str) -> Optional[str]:
+    """A prompt node's `text` as it will be at run time, if knowable from the graph.
+
+    Typed text, or text produced by known pure string nodes (core primitives and
+    StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText). Anything
+    else, including PromptSearchList batches, returns None rather than a guess.
+    Callers compare the result with the text the node actually received.
+    """
+    if not isinstance(graph, dict):
+        return None
+    seen = set()
+
+    def resolve(value: Any) -> Optional[str]:
+        if isinstance(value, str):
+            return value
+        if not _is_link(value) or len(seen) >= MAX_VISITS:
+            return None
+        source_id = str(value[0])
+        if source_id in seen:
+            return None
+        seen.add(source_id)
+        source = graph.get(source_id)
+        if not isinstance(source, dict):
+            return None
+        return _resolve_string_node(
+            str(source.get("class_type", "")), _inputs(source), resolve
+        )
+
+    seen.add(str(node_id))
+    text = resolve(_inputs(graph.get(str(node_id))).get("text"))
+    return text.strip() if isinstance(text, str) and text.strip() else None

utils/usage_tracking.py

@@ -6,23 +6,23 @@ workflow instead: it counts a run for each PromptManager node feeding a sampler'
 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.
+Text from known pure string nodes is resolved at queue time too. Anything else
+(PromptSearchList batches, unknown nodes) is counted by the node as it executes.
 """
 
 import threading
 import time
 from collections import OrderedDict
-from typing import Any
+from typing import Any, Tuple
 
 try:
     from .hashing import generate_prompt_hash
     from .logging_config import get_logger
-    from .prompt_graph import literal_text, run_prompt_nodes
+    from .prompt_graph import resolve_text, run_prompt_nodes
 except ImportError:
     from utils.hashing import generate_prompt_hash
     from utils.logging_config import get_logger
-    from utils.prompt_graph import literal_text, run_prompt_nodes
+    from utils.prompt_graph import resolve_text, run_prompt_nodes
 
 logger = get_logger("prompt_manager.usage_tracking")
 
@@ -31,7 +31,11 @@ PENDING_MAX_ENTRIES = 1000
 
 
 class PendingFirstUse:
-    """Hashes of queued positive prompts that were not in the database yet."""
+    """Queued runs of positive prompts that were not in the database yet.
+
+    Counts per hash: the same new prompt can be queued several times before its
+    first job runs, and later jobs may reuse ComfyUI's cached node.
+    """
 
     def __init__(
         self,
@@ -40,21 +44,21 @@ class PendingFirstUse:
     ):
         self._ttl = ttl_seconds
         self._max = max_entries
-        self._entries: "OrderedDict[str, float]" = OrderedDict()
+        self._entries: "OrderedDict[str, Tuple[int, 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()
+            count, _ = self._entries.pop(prompt_hash, (0, 0.0))
+            self._entries[prompt_hash] = (count + 1, time.monotonic())
             while len(self._entries) > self._max:
                 self._entries.popitem(last=False)
 
-    def consume(self, prompt_hash: str) -> bool:
-        """True (once) if this hash was queued recently and not yet counted."""
+    def consume(self, prompt_hash: str) -> int:
+        """Number of recent queued runs not yet counted (0 if none); clears them."""
         with self._lock:
-            added = self._entries.pop(prompt_hash, None)
-        return added is not None and time.monotonic() - added <= self._ttl
+            count, added = self._entries.pop(prompt_hash, (0, 0.0))
+        return count if count and time.monotonic() - added <= self._ttl else 0
 
 
 # Shared between the queue hook and the nodes
@@ -67,7 +71,7 @@ def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> A
         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)
+            text = resolve_text(graph, node_id)
             if not text:
                 continue
             prompt_hash = generate_prompt_hash(text)
@@ -84,7 +88,7 @@ def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> A
     return json_data
 
 
-_registered = False
+_hook_registered = threading.Event()
 _register_lock = threading.Lock()
 
 
@@ -92,13 +96,12 @@ 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:
+        if _hook_registered.is_set():
             return False
         server.add_on_prompt_handler(
             lambda json_data: handle_queued_prompt(json_data, db_factory(), pending)
         )
-        _registered = True
+        _hook_registered.set()
         logger.info("Registered prompt usage hook")
         return True