Files
ComfyAssets-ComfyUI_PromptM…/tests/test_node_tracking_wiring.py
T
Vito 6c7e4d5b35 feat: release 3.2.4 — recently used sorting, run counts, positive-prompt image linking (#152)
* chore: start 3.2.4

* feat: track prompt usage (run_count, last_used_at) with server-side sorting

- Migration adds last_used_at and run_count; existing prompts are estimated
  from linked images (count, newest image time, else created_at). The
  backfill only touches NULL rows, so it self-heals rows written after a
  downgrade or an interrupted upgrade.
- Every node run records a use, including re-runs of an existing prompt;
  a tracking failure is logged and never blocks generation.
- search_prompts/get_recent_prompts accept a whitelisted sort key
  (last_used, created, run_count, rating, text); default stays newest first.

Refs #90

* feat: sort param on /recent and /search; microsecond last_used_at

- /prompt_manager/recent and /search accept ?sort= (whitelisted in the DB
  layer; unknown values fall back to newest first).
- last_used_at is written with microseconds from Python so a re-run in the
  same millisecond as another save still sorts first; ordering stays correct
  against millisecond values from the backfill.

Refs #90

* feat: server-side sort in admin UI with Recently Used / Most Used and run badges

- New web/js/prompt-list-sort.js owns the sort options, validation and
  request URLs; admin.js fills the dropdown from it and re-fetches from the
  server on change, so every sort spans all pages (the old client-side sort
  only reordered the current page).
- Default sort is Recently Used; prompt cards show a run-count badge with
  the last-used time on hover.
- The node's Recent button now lists recently used prompts.
- Contract test keeps the UI sort keys and the node's key in sync with the
  server whitelist.

Refs #90

* 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

* 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.

* refactor: unquote PendingFirstUse annotation so Tuple usage is visible
2026-10-07 13:52:09 -07:00

84 lines
3.2 KiB
Python

"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once."""
import os
import sys
import threading
import unittest
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from prompt_manager import PromptManager
from prompt_manager_text import PromptManagerText
from utils import usage_tracking
GRAPH = {"134": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}}
def bare(node_cls):
"""A node without its database, tracker or gallery side effects."""
node = node_cls.__new__(node_cls)
node.logger = __import__("logging").getLogger("test.wiring")
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node._track_prompt_execution = mock.Mock(return_value=42)
return node
class TestHiddenInputs(unittest.TestCase):
def test_both_nodes_request_the_queued_graph_and_their_id(self):
for node_cls in (PromptManager, PromptManagerText):
hidden = node_cls.INPUT_TYPES().get("hidden", {})
self.assertEqual(hidden.get("prompt"), "PROMPT", node_cls.__name__)
self.assertEqual(hidden.get("unique_id"), "UNIQUE_ID", node_cls.__name__)
class TestNodesForwardGraph(unittest.TestCase):
def test_prompt_manager_forwards_graph_and_texts(self):
node = bare(PromptManager)
clip = mock.Mock()
node.encode_prompt(
clip, "a cat", prepend_text="best,", prompt=GRAPH, unique_id="134"
)
kwargs = node._track_prompt_execution.call_args.kwargs
self.assertEqual(kwargs["prompt_graph"], GRAPH)
self.assertEqual(kwargs["unique_id"], "134")
self.assertEqual(kwargs["text"], "a cat")
self.assertEqual(kwargs["encoding_text"], "best, a cat")
self.assertEqual(kwargs["additional_data"]["final_text"], "best, a cat")
def test_prompt_manager_text_forwards_graph(self):
node = bare(PromptManagerText)
node.process_text("a castle", prompt=GRAPH, unique_id="7")
kwargs = node._track_prompt_execution.call_args.kwargs
self.assertEqual((kwargs["prompt_graph"], kwargs["unique_id"]), (GRAPH, "7"))
self.assertEqual(kwargs["text"], "a castle")
def test_nodes_still_work_without_graph_context(self):
node = bare(PromptManager)
node.encode_prompt(mock.Mock(), "a cat")
kwargs = node._track_prompt_execution.call_args.kwargs
self.assertIsNone(kwargs["prompt_graph"])
self.assertIsNone(kwargs["unique_id"])
class TestQueueHookRegistration(unittest.TestCase):
def test_registers_once_and_handler_passes_requests_through(self):
server = mock.Mock()
db = mock.Mock()
db.get_prompt_by_hash.return_value = None
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db))
self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db))
server.add_on_prompt_handler.assert_called_once()
handler = server.add_on_prompt_handler.call_args.args[0]
request = {"prompt": GRAPH}
self.assertIs(handler(request), request)
if __name__ == "__main__":
unittest.main()