nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: remove dead and test-only code and the write-only prompt registry (#620)

Browse Source 
Removes helpers nothing calls, test-only wrappers and their tests, the unused
config writer, three unused images, and the prompt registry the nodes wrote to
but nothing read, which grew with every run.
This commit is contained in:
Vito Sansevero
2026-11-02 10:22:03 -08:00
parent 56d1072e58
commit f6432fec74
40 changed files with 50 additions and 969 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.comfyignore
-4
View File
@@ -19,7 +19,3 @@
/.flake8
/.pre-commit-config.yaml
/.gitignore
# Stray images no code uses
/image_meta_data.png
/test_metadata_file.png
database/annotations.py
-7
View File
@@ -291,13 +291,6 @@ class AnnotationOperations:
conn.execute(_BUMP_INDEX)
return conn.execute(_REVISION, (byte_id,)).fetchone()[0]
def _current_set(self, conn, sha256, byte_id, revision) -> dict:
return {
"sha256": sha256,
"tags": self._byte_tags(conn, byte_id),
"tag_revision": revision,
}
def read_image_tags(self, locator: str, sha256: str) -> dict:
"""The stored set of these exact bytes, or an empty set at revision 0.
image_meta_data.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 85 KiB

images/pm-lora-results.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 34 KiB

prompt_manager.py
-14
View File
@@ -3,7 +3,6 @@ PromptManager: Main custom node implementation that extends CLIPTextEncode
with persistent prompt storage and search capabilities.
"""
import time
from typing import Any, Dict, Optional, Tuple
if __package__:
@@ -259,19 +258,6 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
if not save_to_library:
return result # #90: no metadata registration, no history text
# Register with ComfyUI integration for standard metadata compatibility
node_id = f"promptmanager_{int(time.time() * 1000)}"
self.comfyui_integration.register_prompt(
node_id,
encoding_text.strip(),
{
"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,
},
)
record_runtime_context(
type(self).__name__, text, prepend_text, append_text, encoding_text
)
prompt_manager_base.py
-22
View File
@@ -19,7 +19,6 @@ except ImportError:
try:
from .database.operations import PromptDatabase
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.hashing import generate_prompt_hash
@@ -38,7 +37,6 @@ except ImportError:
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from database.operations import PromptDatabase
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.hashing import generate_prompt_hash
@@ -86,7 +84,6 @@ class PromptManagerBase:
self.db = PromptDatabase()
self.prompt_tracker = get_prompt_tracker(self.db)
self.image_monitor = get_image_monitor(self.db, self.prompt_tracker)
self.comfyui_integration = get_comfyui_integration()
self._start_gallery_system()
self.logger.debug(f"{self.__class__.__name__} node initialization completed")
@@ -433,18 +430,6 @@ class PromptManagerBase:
self.logger.error(f"Error searching prompts: {e}")
return []
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
"""API method for JavaScript UI to search prompts."""
return self._search_prompts(search_text=search_text)
def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]:
"""API method for JavaScript UI to get recent prompts."""
try:
return self.db.get_recent_prompts(limit=limit)
except Exception as e:
self.logger.error(f"Error getting recent prompts: {e}")
return []
def _start_gallery_system(self):
"""Initialize and start the gallery monitoring system."""
try:
@@ -455,13 +440,6 @@ class PromptManagerBase:
self.logger.error(f"Failed to start gallery system: {e}")
self.logger.warning("Gallery features will be disabled")
def get_gallery_status(self) -> Dict[str, Any]:
"""Get status of the gallery system."""
return {
"image_monitor": self.image_monitor.get_status(),
"prompt_tracker": self.prompt_tracker.get_status(),
}
def cleanup_gallery_system(self):
"""Clean up gallery system resources.
prompt_manager_pair.py
-14
View File
@@ -2,7 +2,6 @@
import hashlib
import json
import time
from typing import Any, Dict, Optional, Tuple
if __package__:
@@ -173,19 +172,6 @@ class PromptManagerPair(PromptManagerFocused):
self.logger.warning(f"Could not save or track paired prompt: {error}")
if is_positive:
self._mark_untracked(prompt, unique_id)
try:
self.comfyui_integration.register_prompt(
f"promptmanager_pair_{time.time_ns()}",
positive,
{
"prompt_id": prompt_id,
"pair_id": pair_id if pair_used else None,
"negative_text": negative_text,
"category": category.strip() or None,
},
)
except Exception as error:
self.logger.warning(f"Could not register paired metadata: {error}")
record_runtime_context(
type(self).__name__,
text,
prompt_manager_text.py
-14
View File
@@ -3,7 +3,6 @@ PromptManagerText: A text-only version of PromptManager that outputs STRING
without CLIP encoding, while maintaining all database and search features.
"""
import time
from typing import Any, Dict, Optional, Tuple
if __package__:
@@ -230,19 +229,6 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
if not save_to_library:
return result # #90: no metadata registration, no history text
# Register with ComfyUI integration for standard metadata compatibility
node_id = f"promptmanagertext_{int(time.time() * 1000)}"
self.comfyui_integration.register_prompt(
node_id,
final_text.strip(),
{
"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,
},
)
record_runtime_context(
type(self).__name__, text, prepend_text, append_text, final_text
)
py/api/__init__.py
-47
View File
@@ -89,7 +89,6 @@ from .source_rescope import SourceRescopeMixin
try:
from ...database.operations import PromptDatabase
from ...utils.file_prompt import parse_comfyui_prompt
from ...utils.logging_config import get_logger
from ...utils.native_memory import release_free_memory
from ...utils.png_text import png_text
@@ -101,7 +100,6 @@ except ImportError:
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
)
from database.operations import PromptDatabase
from utils.file_prompt import parse_comfyui_prompt
from utils.logging_config import get_logger
from utils.native_memory import release_free_memory
from utils.png_text import png_text
@@ -254,23 +252,6 @@ async def _read_json_body(request, max_bytes=JSON_BODY_MAX_BYTES):
return data, None
class _LongJobGuard:
"""Async context manager handed out by ``_acquire_long_job``."""
def __init__(self, lock):
self._lock = lock
async def __aenter__(self):
# The caller checked ``locked()`` on this same loop iteration, so this
# acquire never waits.
await self._lock.acquire()
return self
async def __aexit__(self, exc_type, exc, tb):
self._lock.release()
return False
# ── Gzip compression middleware ────────────────────────────────────────
_GZIP_MIN_SIZE = 1024 # Only compress bodies larger than 1 KB
_GZIP_TYPES = frozenset(
@@ -426,10 +407,6 @@ class PromptManagerAPI(
self._workspace_card_cache = {}
# Favorite keys of each root's cached scan (images._favorite_listing).
self._favorite_listing_cache = {}
# Single-flight locks for long jobs, keyed by job family; see
# _acquire_long_job(). Created lazily so no lock binds to an event
# loop before the server's loop exists.
self._long_jobs: dict = {}
# This session's jobs (py/api/jobs.py): the metadata scan, and later
# every other long operation the Activity page tracks.
self._jobs = JobRegistry(self.logger)
@@ -486,26 +463,6 @@ class PromptManagerAPI(
_public_error = staticmethod(_public_error)
_read_json_body = staticmethod(_read_json_body)
def _acquire_long_job(self, name):
"""Claim the single-flight slot for a long job family, or None when busy.
Usage in a handler::
job = self._acquire_long_job("duplicates")
if job is None:
return self._long_job_busy_response("Duplicate scan")
async with job:
...
Families are free-form strings ("duplicates", "thumbnails",
"autotag", ...); the lock is released when the ``async with`` block
exits, including on exceptions. Must be called from the event loop.
"""
lock = self._long_jobs.setdefault(name, asyncio.Lock())
if lock.locked():
return None
return _LongJobGuard(lock)
@staticmethod
def _long_job_busy_response(label):
"""409 envelope for a long job that is already running."""
@@ -1079,7 +1036,3 @@ class PromptManagerAPI(
f"Unreadable image {os.path.basename(str(image_path))}: {e}"
)
return {}
def _parse_comfyui_prompt(self, metadata):
"""Parse ComfyUI and A1111 prompt data from metadata."""
return parse_comfyui_prompt(metadata)
py/api/autotag_routes.py
-10
View File
@@ -943,16 +943,6 @@ class AutotagRoutesMixin:
return refusal.response
return web.json_response({"success": True, "job": job.summary()}, status=202)
def _gated_runner(self, owner):
"""Executor calls under the model gate *owner* holds (plain when None)."""
if owner is None:
return self._run_in_executor
async def run(func, *args, **kwargs):
return await model_gate.run(owner, func, *args, **kwargs)
return run
def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
"""Return the real path when it is a media file inside an output directory."""
if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
py/api/thumbnail_auto.py
-13
View File
@@ -632,10 +632,6 @@ class AutoThumbnailer:
"""True once the worker has nothing queued or in hand."""
return self._wait_for(lambda: not self._busy, timeout)
def wait_batch(self, size, timeout=None):
"""True once the worker holds a batch of *size* files."""
return self._wait_for(lambda: self._batch == size, timeout)
class _PreviewGate:
"""Admission of view requests on the event loop (#249 I2).
@@ -967,15 +963,6 @@ class AutoThumbnailMixin:
return "queued"
return "skipped"
def _capture_target(self, image_path):
"""``(src, dst, is_video)`` of a new output's preview now, or None."""
located = self._capture_source(image_path)
if located is None:
return None
src, root, rel, is_video = located
dst = DeferredPreview(root, rel).resolve()
return None if dst is None else (src, dst, is_video)
def _capture_source(self, image_path):
"""``(src, root, rel, is_video)`` of a media file under a current
root (not a preview, a link or in a trash folder), or None."""
py/api/thumbnail_jobs.py
+2 -2
View File
@@ -90,7 +90,7 @@ class ThumbnailOptions:
}
def thumbnail_options(raw, *, from_query=False):
def thumbnail_options(raw):
"""``(ThumbnailOptions, None)`` from a body or query, or ``(None, message)``.
An absent mode is ``missing``; the settings are the server's now.
@@ -136,7 +136,7 @@ class ThumbnailJobsMixin:
``mode=prune`` lists the removed sources' thumbnails (#320)."""
if request.query.get("mode") == "prune":
return await self.get_prune_plan()
options, error = thumbnail_options(request.query, from_query=True)
options, error = thumbnail_options(request.query)
if error:
return _invalid(error)
try:
py/config.py
-56
View File
@@ -78,27 +78,6 @@ def _canonical_path(path: str) -> str:
return os.path.normcase(os.path.realpath(path))
def write_private_json(path: str, data: Any) -> None:
"""Write ``data`` as JSON to ``path`` readable by the current user only.
config.json can hold an API key (integrations.lora_manager), so it is
created with mode 0600 and, on POSIX, an existing more open file is
tightened too. Windows has no POSIX modes; there the chmod is a best
effort that only affects the read-only bit. Parent directories are
created as needed. Raises OSError when the file cannot be written.
"""
import json
parent_dir = os.path.dirname(path)
if parent_dir:
os.makedirs(parent_dir, exist_ok=True)
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
if os.name == "posix":
os.chmod(path, 0o600)
def _lexical_path(path: str) -> str:
"""``path`` normalised for comparison without touching the filesystem.
@@ -717,12 +696,6 @@ class AutotagConfig:
def suggestion_limit(cls) -> int:
return cls.SUGGESTION_LIMIT
@staticmethod
def valid_suggestion_limit(value: Any) -> bool:
return (
type(value) is int and 1 <= value <= MAX_SUGGESTION_LIMIT
) # bool, float and str are rejected
@classmethod
def all_instructions(cls) -> List[Dict[str, str]]:
"""Built-in first, then the saved entries (copies)."""
@@ -1063,35 +1036,6 @@ class PromptManagerConfig:
else:
config_logger.info(f"Config file not found: {config_path}, using defaults")
@classmethod
def save_to_file(cls, config_path: Optional[str] = None):
"""Save the current configuration to a JSON file.
Serializes the complete configuration (including gallery settings) to
a JSON file. Creates the directory structure if it doesn't exist.
Args:
config_path (str): Path where the JSON configuration file should be saved.
Parent directories will be created if they don't exist.
Defaults to :meth:`get_config_path`.
Raises:
The method handles all exceptions internally and logs errors rather
than propagating them.
Example:
PromptManagerConfig.save_to_file('backup_config.json')
PromptManagerConfig.save_to_file('/etc/comfyui/prompt_manager.json')
"""
if config_path is None:
config_path = cls.get_config_path()
try:
write_private_json(config_path, cls.get_config())
config_logger.info(f"Saved configuration to {config_path}")
except Exception as e:
config_logger.error(f"Error saving config to {config_path}: {e}")
@classmethod
def update_config(cls, new_config: Dict[str, Any]):
"""Update configuration attributes from a dictionary.
test_metadata_file.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 1.5 MiB

tests/test_api_core.py
+1 -54
View File
@@ -343,7 +343,7 @@ class TestReadJsonBody(AioHTTPTestCase):
class TestLongJobLock(unittest.TestCase):
"""_acquire_long_job: one job per family at a time, released on exit."""
"""The refusal a busy long job answers with."""
def setUp(self):
database = MagicMock()
@@ -353,33 +353,6 @@ class TestLongJobLock(unittest.TestCase):
def _run(self, coro):
return asyncio.run(coro)
def test_second_acquire_is_refused_until_first_releases(self):
async def scenario():
first = self.api._acquire_long_job("duplicates")
self.assertIsNotNone(first)
async with first:
self.assertIsNone(self.api._acquire_long_job("duplicates"))
# Other families are independent
other = self.api._acquire_long_job("thumbnails")
self.assertIsNotNone(other)
async with other:
pass
again = self.api._acquire_long_job("duplicates")
self.assertIsNotNone(again)
async with again:
pass
self._run(scenario())
def test_lock_is_released_when_the_job_raises(self):
async def scenario():
with self.assertRaises(RuntimeError):
async with self.api._acquire_long_job("duplicates"):
raise RuntimeError("boom")
self.assertIsNotNone(self.api._acquire_long_job("duplicates"))
self._run(scenario())
def test_busy_response_is_409_with_envelope(self):
resp = self.api._long_job_busy_response("Duplicate scan")
self.assertEqual(resp.status, 409)
@@ -953,32 +926,6 @@ class TestMetadataParsing(unittest.TestCase):
f.write(b"nope")
self.assertEqual(self.api._extract_comfyui_metadata(garbage), {})
def test_parse_a1111_parameters(self):
meta = {"parameters": "a cat\nNegative prompt: dog\nSteps: 20"}
parsed = self.api._parse_comfyui_prompt(meta)
self.assertEqual(parsed["positive_prompt"], "a cat")
self.assertEqual(parsed["negative_prompt"], "dog")
self.assertEqual(parsed["parameters"]["parameters"], meta["parameters"])
def test_parse_comfyui_fields(self):
meta = {
"prompt": json.dumps({"1": {"class_type": "X", "inputs": {}}}),
"workflow": "not json",
"steps": "20",
"cfg": "not-json-either",
}
parsed = self.api._parse_comfyui_prompt(meta)
self.assertIsInstance(parsed["prompt"], dict)
self.assertEqual(parsed["workflow"], "not json")
self.assertEqual(parsed["parameters"]["steps"], 20)
self.assertEqual(parsed["parameters"]["cfg"], "not-json-either")
parsed = self.api._parse_comfyui_prompt(
{"prompt": "plain text", "workflow": json.dumps({"nodes": []})}
)
self.assertEqual(parsed["prompt"], "plain text")
self.assertEqual(parsed["workflow"], {"nodes": []})
if __name__ == "__main__":
unittest.main()
tests/test_api_thumbnail_location.py
-20
View File
@@ -9,7 +9,6 @@ import json
import ntpath
import os
import posixpath
import re
import sys
import tempfile
import threading
@@ -378,25 +377,6 @@ class TestValidation(unittest.TestCase):
self.assertTrue(text.endswith("."))
class TestNoLegacySaveCallers(unittest.TestCase):
def test_save_to_file_has_no_callers(self):
offenders = []
for folder in ("py", "utils", "database"):
for path in (REPO / folder).rglob("*.py"):
for number, line in enumerate(
path.read_text(encoding="utf-8").splitlines(), 1
):
code = line.split("#", 1)[0]
if (
re.search(r"\bsave_to_file\(", code)
and "def save_to_file" not in code
):
if "PromptManagerConfig.save_to_file('" in line:
continue # its own docstring example
offenders.append(f"{path.relative_to(REPO)}:{number}")
self.assertEqual(offenders, [])
try:
from tests.test_api_images import ImageAPITestCase
except ImportError: # discovered with tests/ as the top-level directory
tests/test_cached_runs.py
-1
View File
@@ -197,7 +197,6 @@ class JobCase(unittest.TestCase):
node.logger = logging.getLogger("test.cached_runs")
node.prompt_tracker = mock.Mock()
self.trackers.append(node.prompt_tracker)
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node.pending_first_use = self.pending
node.pending_pair_first_use = self.pair_pending
tests/test_comfyui_integration.py
-153
View File
@@ -1,153 +0,0 @@
"""Tests for utils.comfyui_integration."""
import os
import sys
import types
import unittest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.comfyui_integration import ComfyUIMetadataIntegration
class TestSaveImageNotPatched(unittest.TestCase):
"""Regression for #75: SaveImage must save each node's own prompt inputs.
The old SaveImage patch wrote the last-run PromptManager prompt into every
PromptManager node, so negative prompts were saved as copies of the positive.
"""
def setUp(self):
def save_images(
self_node,
images,
filename_prefix="ComfyUI",
prompt=None,
extra_pnginfo=None,
):
return {"prompt": prompt}
self.original_save_images = save_images
fake_nodes = types.ModuleType("nodes")
fake_nodes.SaveImage = type("SaveImage", (), {"save_images": save_images})
self._saved_nodes = sys.modules.get("nodes")
sys.modules["nodes"] = fake_nodes
self.fake_nodes = fake_nodes
ComfyUIMetadataIntegration._instance = None
def tearDown(self):
ComfyUIMetadataIntegration._instance = None
if self._saved_nodes is None:
sys.modules.pop("nodes", None)
else:
sys.modules["nodes"] = self._saved_nodes
def test_save_images_is_left_untouched(self):
ComfyUIMetadataIntegration()
self.assertIs(self.fake_nodes.SaveImage.save_images, self.original_save_images)
def test_registered_prompt_does_not_rewrite_saved_metadata(self):
integration = ComfyUIMetadataIntegration()
integration.register_prompt("pm_1", "positive final text", {})
prompt = {
"39": {"class_type": "PromptManager", "inputs": {"text": "a cat"}},
"40": {"class_type": "PromptManager", "inputs": {"text": "blurry"}},
}
result = self.fake_nodes.SaveImage().save_images([], prompt=prompt)
self.assertEqual(result["prompt"]["39"]["inputs"]["text"], "a cat")
self.assertEqual(result["prompt"]["40"]["inputs"]["text"], "blurry")
def test_prompt_registry_still_available_to_nodes(self):
integration = ComfyUIMetadataIntegration()
integration.register_prompt("pm_1", "hello", {})
self.assertEqual(integration.get_current_prompt_text("pm_1"), "hello")
import threading # noqa: E402
import unittest.mock as mock # noqa: E402
from utils import comfyui_integration # noqa: E402
from utils.comfyui_integration import get_comfyui_integration # noqa: E402
class IntegrationTestCase(unittest.TestCase):
def setUp(self):
ComfyUIMetadataIntegration._instance = None
comfyui_integration._integration_instance = None
# Patch the module's own ``time`` name rather than ``time.time`` globally:
# background threads from other tests (image monitor, watchdog) would
# otherwise tick this clock and make fresh registrations look stale.
self.clock = mock.patch.object(
comfyui_integration, "time", types.SimpleNamespace(time=self._now)
)
self.clock.start()
self.addCleanup(self.clock.stop)
self._ticks = 1000.0
self.integration = get_comfyui_integration()
def tearDown(self):
ComfyUIMetadataIntegration._instance = None
comfyui_integration._integration_instance = None
def _now(self):
self._ticks += 1.0
return self._ticks
class TestSingletonAndRegistry(IntegrationTestCase):
def test_singleton_keeps_registrations_across_constructions(self):
self.assertIs(get_comfyui_integration(), self.integration)
self.assertIs(ComfyUIMetadataIntegration(), self.integration)
self.integration.register_prompt("n1", "hello", {"category": "x"})
ComfyUIMetadataIntegration() # __init__ runs again: must not reset state
self.assertEqual(self.integration.get_current_prompt_text("n1"), "hello")
def test_latest_prompt_of_this_thread_wins_without_a_node_id(self):
self.integration.register_prompt("n1", "first", {})
self.integration.register_prompt("n2", "second", {})
self.assertEqual(self.integration.get_current_prompt_text(), "second")
self.assertEqual(self.integration.get_current_prompt_text("missing"), "second")
self.assertEqual(self.integration.get_current_prompt_text("n1"), "first")
def test_nothing_registered_returns_none(self):
self.assertIsNone(self.integration.get_current_prompt_text())
self.assertIsNone(self.integration.get_current_prompt_text("n1"))
class TestCrossThreadLookup(IntegrationTestCase):
def _register_in_thread(self, text):
worker = threading.Thread(
target=self.integration.register_prompt, args=("w1", text, {})
)
worker.start()
worker.join()
def test_recent_prompt_from_another_thread_is_used(self):
self._register_in_thread("from worker")
self.assertEqual(self.integration.get_current_prompt_text(), "from worker")
def test_stale_prompt_from_another_thread_is_ignored(self):
self._register_in_thread("from worker")
self._ticks += 400.0
self.assertIsNone(self.integration.get_current_prompt_text())
def test_global_entry_of_this_thread_is_found_without_thread_local(self):
self.integration.register_prompt("n1", "mine", {})
self.integration._thread_local = threading.local()
self.assertEqual(self.integration.get_current_prompt_text(), "mine")
def test_cleanup_removes_old_registrations_only(self):
self._register_in_thread("old")
self._ticks += 700.0
self.integration.register_prompt("n2", "new", {})
self.integration.cleanup_old_prompts(max_age_seconds=600)
texts = sorted(v["text"] for v in self.integration._current_prompts.values())
self.assertEqual(texts, ["new"])
self.integration.cleanup_old_prompts(max_age_seconds=600) # nothing to do
self.assertEqual(len(self.integration._current_prompts), 1)
if __name__ == "__main__":
unittest.main()
tests/test_config.py
+4 -63
View File
@@ -213,7 +213,8 @@ class TestPromptManagerConfig(unittest.TestCase):
try:
# Modify a value
PromptManagerConfig.MAX_SEARCH_RESULTS = 42
PromptManagerConfig.save_to_file(tmp.name)
with open(tmp.name, "w", encoding="utf-8") as handle:
json.dump(PromptManagerConfig.get_config(), handle)
# Reset and reload
PromptManagerConfig.MAX_SEARCH_RESULTS = 100
@@ -240,55 +241,6 @@ class TestPromptManagerConfig(unittest.TestCase):
finally:
os.unlink(tmp.name)
def test_save_creates_directory(self):
tmp_dir = tempfile.mkdtemp()
nested_path = os.path.join(tmp_dir, "subdir", "config.json")
try:
PromptManagerConfig.save_to_file(nested_path)
self.assertTrue(os.path.exists(nested_path))
with open(nested_path) as f:
data = json.load(f)
self.assertIn("database", data)
finally:
import shutil
shutil.rmtree(tmp_dir)
def test_saved_file_is_private_to_the_user(self):
import stat
tmp_dir = tempfile.mkdtemp()
try:
path = os.path.join(tmp_dir, "config.json")
PromptManagerConfig.save_to_file(path)
mode = stat.S_IMODE(os.stat(path).st_mode)
if os.name == "posix":
self.assertEqual(mode, 0o600)
with open(path) as f:
self.assertIn("database", json.load(f))
# Overwriting an existing, more open file tightens it as well
os.chmod(path, 0o644)
PromptManagerConfig.save_to_file(path)
if os.name == "posix":
self.assertEqual(stat.S_IMODE(os.stat(path).st_mode), 0o600)
finally:
import shutil
shutil.rmtree(tmp_dir, ignore_errors=True)
def test_saved_file_is_valid_json(self):
tmp = tempfile.NamedTemporaryFile(
delete=False, suffix=".json", dir=tempfile.gettempdir()
)
tmp.close()
try:
PromptManagerConfig.save_to_file(tmp.name)
with open(tmp.name) as f:
data = json.load(f)
self.assertIsInstance(data, dict)
finally:
os.unlink(tmp.name)
class TestIntegrationConfig(unittest.TestCase):
"""Test IntegrationConfig for LoRA Manager settings."""
@@ -404,18 +356,6 @@ class TestConfigPathOverride(unittest.TestCase):
os.environ[self.ENV] = custom
self.assertEqual(PromptManagerConfig.get_config_path(), custom)
def test_save_without_path_uses_env_override(self):
with tempfile.TemporaryDirectory() as tmpdir:
custom = os.path.join(tmpdir, "nested", "custom.json")
os.environ[self.ENV] = custom
PromptManagerConfig.RESULT_TIMEOUT = 42
PromptManagerConfig.save_to_file()
with open(custom) as f:
saved = json.load(f)
self.assertEqual(saved["web_ui"]["result_timeout"], 42)
def test_load_without_path_uses_env_override(self):
with tempfile.TemporaryDirectory() as tmpdir:
custom = os.path.join(tmpdir, "custom.json")
@@ -503,7 +443,8 @@ class TestAutotagConfig(unittest.TestCase):
PromptManagerConfig.update_config(
{"autotag": {"defaults": {"suggestion_limit": 25}}}
)
PromptManagerConfig.save_to_file(path)
with open(path, "w", encoding="utf-8") as handle:
json.dump(PromptManagerConfig.get_config(), handle)
self.config_class.SUGGESTION_LIMIT = 10
PromptManagerConfig.load_from_file(path)
self.assertEqual(self.config_class.suggestion_limit(), 25)
tests/test_empty_text_handling.py
+36 -71
View File
@@ -15,12 +15,11 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
@patch("prompt_manager_base.get_image_monitor")
@patch("prompt_manager_base.get_prompt_tracker")
@patch("prompt_manager_base.get_comfyui_integration")
@patch("prompt_manager_base.PromptDatabase")
class TestPromptManagerTextEmptyInput(unittest.TestCase):
"""Test PromptManagerText with empty/blank text inputs (issue #123)."""
def _make_node(self, mock_db_class, mock_integration, mock_tracker, mock_monitor):
def _make_node(self, mock_db_class, mock_tracker, mock_monitor):
mock_db = Mock()
mock_db_class.return_value = mock_db
mock_db.save_prompt.return_value = 1
@@ -30,73 +29,64 @@ class TestPromptManagerTextEmptyInput(unittest.TestCase):
return PromptManagerText()
def test_empty_string(self, mock_db, mock_integ, mock_tracker, mock_monitor):
def test_empty_string(self, mock_db, mock_tracker, mock_monitor):
"""Empty string should not raise UnboundLocalError."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(text="")
self.assertIsInstance(result, tuple)
self.assertEqual(result[0], "")
def test_whitespace_only(self, mock_db, mock_integ, mock_tracker, mock_monitor):
def test_whitespace_only(self, mock_db, mock_tracker, mock_monitor):
"""Whitespace-only string should not raise UnboundLocalError."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(text=" ")
self.assertIsInstance(result, tuple)
def test_none_coerced_empty(self, mock_db, mock_integ, mock_tracker, mock_monitor):
def test_none_coerced_empty(self, mock_db, mock_tracker, mock_monitor):
"""Falsy text value should not raise UnboundLocalError."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
# ComfyUI may pass empty string for unconnected inputs
result = node.process_text(text="", category="", tags="")
self.assertIsInstance(result, tuple)
def test_empty_text_with_prepend_append(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_empty_text_with_prepend_append(self, mock_db, mock_tracker, mock_monitor):
"""Empty main text with prepend/append should produce single space, not
double."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(text="", prepend_text="before", append_text="after")
self.assertIsInstance(result, tuple)
self.assertEqual(result[0], "before after")
self.assertNotIn(" ", result[0])
def test_valid_text_still_works(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_valid_text_still_works(self, mock_db, mock_tracker, mock_monitor):
"""Normal text input should continue to work correctly."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(text="a beautiful sunset")
self.assertIsInstance(result, tuple)
self.assertEqual(result[0], "a beautiful sunset")
def test_prepend_text_append_all_present(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_prepend_text_append_all_present(self, mock_db, mock_tracker, mock_monitor):
"""All three parts should combine with single spaces."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(
text="main prompt", prepend_text="before", append_text="after"
)
self.assertEqual(result[0], "before main prompt after")
def test_only_prepend(self, mock_db, mock_integ, mock_tracker, mock_monitor):
def test_only_prepend(self, mock_db, mock_tracker, mock_monitor):
"""Prepend with empty text should output just prepend."""
node = self._make_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(text="", prepend_text="before")
self.assertEqual(result[0], "before")
@patch("prompt_manager_base.get_image_monitor")
@patch("prompt_manager_base.get_prompt_tracker")
@patch("prompt_manager_base.get_comfyui_integration")
@patch("prompt_manager_base.PromptDatabase")
class TestPromptManagerEmptyInput(unittest.TestCase):
"""Test PromptManager (CLIP variant) with empty/blank text inputs (issue #123)."""
def _make_node_and_clip(
self, mock_db_class, mock_integration, mock_tracker, mock_monitor
):
def _make_node_and_clip(self, mock_db_class, mock_tracker, mock_monitor):
mock_db = Mock()
mock_db_class.return_value = mock_db
mock_db.save_prompt.return_value = 1
@@ -110,30 +100,22 @@ class TestPromptManagerEmptyInput(unittest.TestCase):
return PromptManager(), mock_clip
def test_empty_string(self, mock_db, mock_integ, mock_tracker, mock_monitor):
def test_empty_string(self, mock_db, mock_tracker, mock_monitor):
"""Empty string should not raise UnboundLocalError."""
node, clip = self._make_node_and_clip(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_node_and_clip(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(clip=clip, text="")
self.assertIsInstance(result, tuple)
def test_whitespace_only(self, mock_db, mock_integ, mock_tracker, mock_monitor):
def test_whitespace_only(self, mock_db, mock_tracker, mock_monitor):
"""Whitespace-only string should not raise UnboundLocalError."""
node, clip = self._make_node_and_clip(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_node_and_clip(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(clip=clip, text=" ")
self.assertIsInstance(result, tuple)
def test_empty_text_with_prepend_append(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_empty_text_with_prepend_append(self, mock_db, mock_tracker, mock_monitor):
"""Empty main text with prepend/append should produce single space, not
double."""
node, clip = self._make_node_and_clip(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_node_and_clip(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(
clip=clip, text="", prepend_text="before", append_text="after"
)
@@ -141,24 +123,16 @@ class TestPromptManagerEmptyInput(unittest.TestCase):
self.assertEqual(result[1], "before after")
self.assertNotIn(" ", result[1])
def test_valid_text_still_works(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_valid_text_still_works(self, mock_db, mock_tracker, mock_monitor):
"""Normal text input should continue to work correctly."""
node, clip = self._make_node_and_clip(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_node_and_clip(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(clip=clip, text="a beautiful sunset")
self.assertIsInstance(result, tuple)
self.assertEqual(result[1], "a beautiful sunset")
def test_prepend_text_append_all_present(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_prepend_text_append_all_present(self, mock_db, mock_tracker, mock_monitor):
"""All three parts should combine with single spaces."""
node, clip = self._make_node_and_clip(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_node_and_clip(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(
clip=clip, text="main prompt", prepend_text="before", append_text="after"
)
@@ -167,7 +141,6 @@ class TestPromptManagerEmptyInput(unittest.TestCase):
@patch("prompt_manager_base.get_image_monitor")
@patch("prompt_manager_base.get_prompt_tracker")
@patch("prompt_manager_base.get_comfyui_integration")
@patch("prompt_manager_base.PromptDatabase")
class TestConnectedTextInput(unittest.TestCase):
"""Test that externally connected text is preserved in output (issue #123).
@@ -177,7 +150,7 @@ class TestConnectedTextInput(unittest.TestCase):
value reaches Python; these tests verify Python handles it correctly.
"""
def _make_text_node(self, mock_db_class, mock_integ, mock_tracker, mock_monitor):
def _make_text_node(self, mock_db_class, mock_tracker, mock_monitor):
mock_db = Mock()
mock_db_class.return_value = mock_db
mock_db.save_prompt.return_value = 1
@@ -186,7 +159,7 @@ class TestConnectedTextInput(unittest.TestCase):
return PromptManagerText()
def _make_clip_node(self, mock_db_class, mock_integ, mock_tracker, mock_monitor):
def _make_clip_node(self, mock_db_class, mock_tracker, mock_monitor):
mock_db = Mock()
mock_db_class.return_value = mock_db
mock_db.save_prompt.return_value = 1
@@ -199,10 +172,10 @@ class TestConnectedTextInput(unittest.TestCase):
return PromptManager(), mock_clip
def test_text_node_connected_text_with_prepend_append(
self, mock_db, mock_integ, mock_tracker, mock_monitor
self, mock_db, mock_tracker, mock_monitor
):
"""Simulates screenshot: connected text + prepend + append must all appear."""
node = self._make_text_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_text_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(
text="This is the prompt",
prepend_text="This is prepend text",
@@ -213,21 +186,17 @@ class TestConnectedTextInput(unittest.TestCase):
"This is prepend text This is the prompt This is append text",
)
def test_text_node_connected_text_only(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_text_node_connected_text_only(self, mock_db, mock_tracker, mock_monitor):
"""Connected text without prepend/append passes through unchanged."""
node = self._make_text_node(mock_db, mock_integ, mock_tracker, mock_monitor)
node = self._make_text_node(mock_db, mock_tracker, mock_monitor)
result = node.process_text(text="This is the prompt")
self.assertEqual(result[0], "This is the prompt")
def test_clip_node_connected_text_with_prepend_append(
self, mock_db, mock_integ, mock_tracker, mock_monitor
self, mock_db, mock_tracker, mock_monitor
):
"""Same scenario for PromptManager (CLIP variant)."""
node, clip = self._make_clip_node(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_clip_node(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(
clip=clip,
text="This is the prompt",
@@ -239,13 +208,9 @@ class TestConnectedTextInput(unittest.TestCase):
"This is prepend text This is the prompt This is append text",
)
def test_clip_node_connected_text_only(
self, mock_db, mock_integ, mock_tracker, mock_monitor
):
def test_clip_node_connected_text_only(self, mock_db, mock_tracker, mock_monitor):
"""Connected text without prepend/append passes through unchanged."""
node, clip = self._make_clip_node(
mock_db, mock_integ, mock_tracker, mock_monitor
)
node, clip = self._make_clip_node(mock_db, mock_tracker, mock_monitor)
result = node.encode_prompt(clip=clip, text="This is the prompt")
self.assertEqual(result[1], "This is the prompt")
tests/test_execution_context_capture.py
-1
View File
@@ -324,7 +324,6 @@ class CaptureTests(unittest.TestCase):
node.db.save_prompt_pair.side_effect = RuntimeError("offline")
node.logger = logging.getLogger("test.capture.nodes")
node.prompt_tracker = Mock()
node.comfyui_integration = Mock()
node._inject_lora_trigger_words = lambda text: text + " TRIGGER"
clip = Mock()
clip.tokenize.side_effect = lambda text: text
tests/test_focused_node.py
-1
View File
@@ -24,7 +24,6 @@ class TestFocusedNode(unittest.TestCase):
def test_focused_encodes_composition_and_tracks_main_text_and_graph(self):
node = PromptManagerFocused.__new__(PromptManagerFocused)
node.logger = logging.getLogger("test.focused")
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node._track_prompt_execution = mock.Mock(return_value=42)
clip = mock.Mock()
tests/test_hashing.py
+1 -29
View File
@@ -1,7 +1,7 @@
"""
Tests for hashing utilities.
Tests is_duplicate_prompt and generate_prompt_hash, which are
Tests and generate_prompt_hash, which are
currently untested. generate_prompt_hash is covered in test_basic.py
but we add edge cases here.
"""
@@ -15,7 +15,6 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tests.background_stops import stop_tracker # noqa: E402
from utils.hashing import (
generate_prompt_hash,
is_duplicate_prompt,
STRIPPED_CODE_POINTS,
prompt_key,
register_prompt_functions,
@@ -69,33 +68,6 @@ class TestSingleHashImplementation(unittest.TestCase):
self.assertFalse(os.path.exists(os.path.join(ROOT, "restart_gallery.py")))
class TestIsDuplicatePrompt(unittest.TestCase):
"""Test is_duplicate_prompt function."""
def test_identical_texts_are_duplicates(self):
self.assertTrue(is_duplicate_prompt("hello", "hello"))
def test_case_insensitive(self):
self.assertTrue(is_duplicate_prompt("Hello World", "hello world"))
def test_whitespace_insensitive(self):
self.assertTrue(is_duplicate_prompt(" hello ", "hello"))
def test_different_texts_not_duplicates(self):
self.assertFalse(is_duplicate_prompt("text A", "text B"))
def test_threshold_parameter_accepted(self):
# threshold is unused but should not cause errors
self.assertTrue(is_duplicate_prompt("same", "same", threshold=0.5))
def test_empty_strings_are_duplicates(self):
self.assertTrue(is_duplicate_prompt("", ""))
def test_unicode_comparison(self):
self.assertTrue(is_duplicate_prompt("日本語", "日本語"))
self.assertFalse(is_duplicate_prompt("日本語", "中文"))
class TestGeneratePromptHashEdgeCases(unittest.TestCase):
"""Additional edge cases for generate_prompt_hash."""
tests/test_image_monitor_singleton.py
+1 -12
View File
@@ -21,7 +21,7 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# Mock ComfyUI server before importing anything that touches config. Only the
# first test module to run installs the stub; nothing else in sys.modules is
# replaced, so later modules (test_comfyui_integration) see the real package.
# replaced, so later modules see the real package.
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules.setdefault("server", _mock_server)
@@ -881,17 +881,6 @@ class TestOutputListener(unittest.TestCase):
self.assertTrue(monitor.stop_monitoring())
class TestModuleIsolation(unittest.TestCase):
"""This module must not leave mocks behind for later test modules."""
def test_comfyui_integration_is_not_replaced_by_a_mock(self):
module = sys.modules.get("utils.comfyui_integration")
self.assertFalse(
isinstance(module, MagicMock),
"utils.comfyui_integration was replaced by a MagicMock in sys.modules",
)
class TestMonitorThreadSafety(unittest.TestCase):
"""Singleton creation must be thread-safe."""
tests/test_job_admission.py
-2
View File
@@ -406,8 +406,6 @@ class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase):
)
await self._busy(resp)
self.assertTrue(victim.exists())
# The duplicate family lock was not taken by the refused request.
self.assertIsNotNone(self.api._acquire_long_job(self.api.LONG_JOB_DUPLICATES))
async def test_lora_scan_refused_while_scan_runs(self):
from py.config import IntegrationConfig
tests/test_node_execution_contract.py
+2 -21
View File
@@ -179,7 +179,6 @@ class NodeTestCase(unittest.TestCase):
self.node.logger = logging.getLogger("test.node")
self.node.prompt_tracker = mock.Mock()
self.node.image_monitor = mock.Mock()
self.node.comfyui_integration = mock.Mock()
def tearDown(self):
self.db.close_all() # Windows cannot unlink an open database file
@@ -189,20 +188,16 @@ class NodeTestCase(unittest.TestCase):
class TestBaseSaveAndSearch(NodeTestCase):
def test_init_wires_database_tracker_monitor_and_integration(self):
def test_init_wires_database_tracker_and_monitor(self):
with (
mock.patch.object(prompt_manager_base, "PromptDatabase") as db_cls,
mock.patch.object(prompt_manager_base, "get_prompt_tracker") as tracker,
mock.patch.object(prompt_manager_base, "get_image_monitor") as monitor,
mock.patch.object(
prompt_manager_base, "get_comfyui_integration"
) as integration,
):
node = PromptManagerBase()
self.assertIs(node.db, db_cls.return_value)
tracker.assert_called_once_with(db_cls.return_value)
monitor.assert_called_once_with(db_cls.return_value, tracker.return_value)
self.assertIs(node.comfyui_integration, integration.return_value)
monitor.return_value.start_monitoring.assert_called_once_with()
def test_saving_an_existing_prompt_updates_its_metadata(self):
@@ -242,31 +237,18 @@ class TestBaseSaveAndSearch(NodeTestCase):
self.node._save_prompt_to_database("a cat on a mat")
self.assertEqual(self.node._search_prompts(""), [])
self.assertEqual(self.node._search_prompts(" "), [])
self.assertEqual(len(self.node.search_prompts_api("cat")), 1)
self.assertEqual(len(self.node.get_recent_prompts_api(limit=5)["prompts"]), 1)
self.assertEqual(len(self.node._search_prompts("cat")), 1)
def test_search_and_recent_errors_return_empty_lists(self):
self.node.db = mock.Mock()
self.node.db.search_prompts.side_effect = RuntimeError("locked")
self.node.db.get_recent_prompts.side_effect = RuntimeError("locked")
with self.assertLogs("test.node", level="ERROR"):
self.assertEqual(self.node._search_prompts("cat"), [])
with self.assertLogs("test.node", level="ERROR"):
self.assertEqual(self.node.get_recent_prompts_api(), [])
def test_gallery_system_start_failure_is_logged(self):
self.node.image_monitor.start_monitoring.side_effect = OSError("no inotify")
with self.assertLogs("test.node", level="ERROR"):
self.node._start_gallery_system()
self.node.image_monitor.get_status.return_value = {"running": False}
self.node.prompt_tracker.get_status.return_value = {"active_prompts_count": 0}
self.assertEqual(
self.node.get_gallery_status(),
{
"image_monitor": {"running": False},
"prompt_tracker": {"active_prompts_count": 0},
},
)
self.node.cleanup_gallery_system()
self.node.image_monitor.stop_monitoring.assert_not_called()
@@ -376,7 +358,6 @@ class TestPromptManagerEncoding(NodeTestCase):
self.assertEqual(saved["text"], "a cat")
self.assertIn("prepend:best", saved["tags"])
self.assertIn("append:8k", saved["tags"])
self.node.comfyui_integration.register_prompt.assert_called_once()
def test_tracking_failure_does_not_break_encoding(self):
self.node._track_prompt_execution = mock.Mock(side_effect=RuntimeError("db"))
tests/test_node_tracking_wiring.py
-1
View File
@@ -19,7 +19,6 @@ 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
tests/test_package_contents.py
-2
View File
@@ -22,8 +22,6 @@ DEVELOPMENT = (
"images/x.png",
"TESTING.md",
"requirements-dev.txt",
"image_meta_data.png",
"test_metadata_file.png",
)
KEPT = (
"__init__.py",
tests/test_pair_execution.py
-1
View File
@@ -47,7 +47,6 @@ class TestPairExecution(unittest.TestCase):
self.node.db = self.db
self.node.logger = logging.getLogger("test.execution")
self.node.prompt_tracker = Mock()
self.node.comfyui_integration = Mock()
self.node._inject_lora_trigger_words = lambda text: text + " LORA"
self.extra = {"workflow": {"unchanged": True}, "sibling": {"keep": True}}
tests/test_pair_usage_tracking.py
-1
View File
@@ -64,7 +64,6 @@ class TestPairUsageTracking(unittest.TestCase):
node.pending_pair_first_use = self.pair_pending
node._inject_lora_trigger_words = lambda text: text
node.prompt_tracker = Mock()
node.comfyui_integration = Mock()
return node
def execute(self, graph, text="cat", negative="blur"):
tests/test_prompt_attribution.py
-1
View File
@@ -162,7 +162,6 @@ class NodeCase(DatabaseCase):
node.db = self.db
node.logger = logging.getLogger("test.prompt_attribution.node")
node.prompt_tracker = self.tracker
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node.pending_first_use = PendingFirstUse()
node.pending_pair_first_use = PendingFirstUse()
tests/test_prompt_pair_node.py
+1 -5
View File
@@ -45,7 +45,6 @@ class TestPromptPairNode(unittest.TestCase):
self.node.db = self.db
self.node.logger = logging.getLogger("test.pair_node")
self.node.prompt_tracker = Mock()
self.node.comfyui_integration = Mock()
self.node._inject_lora_trigger_words = Mock(
side_effect=lambda text: text + " LORA"
)
@@ -241,12 +240,9 @@ class TestPromptPairNode(unittest.TestCase):
)
self.node.prompt_tracker.set_current_prompt.assert_not_called()
def test_invalid_text_does_not_write_and_integration_failures_do_not_block(self):
def test_invalid_text_does_not_write(self):
with self.assertRaises(ValueError):
self.node.encode_prompt(self.clip, [], negative_text="blur")
self.node.comfyui_integration.register_prompt.side_effect = RuntimeError(
"metadata unavailable"
)
result = self.node.encode_prompt(self.clip, "cat", negative_text="blur")
self.assertEqual(result[2:], ("cat LORA", "blur"))
tests/test_prompt_usage.py
-1
View File
@@ -294,7 +294,6 @@ class TestNodeRecordsCachedRunReceipt(UsageTestCase):
node.db = self.db
node.logger = __import__("logging").getLogger("test.prompt_usage")
node.prompt_tracker = mock.Mock()
node.comfyui_integration = mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node.cached_runs = mock.Mock()
node.cached_runs.replayed.return_value = False
tests/test_save_to_library.py
-1
View File
@@ -243,7 +243,6 @@ class TestNodesSaveNothing(NodeCase):
self.assertEqual(self.count_rows("prompt_pairs"), 0)
self.assertEqual(self.tracker.keyed_snapshot(), {})
self.assertEqual(self.tracker._prompt_queue, [])
node.comfyui_integration.register_prompt.assert_not_called()
if cls is PromptManagerText:
self.assertEqual(result, (POSITIVE,))
elif cls is PromptManagerPair:
tests/test_split_node.py
-1
View File
@@ -28,7 +28,6 @@ class TestSplitNode(unittest.TestCase):
def test_split_encodes_composition_and_persists_main_text_with_metadata(self):
node = PromptManagerSplit.__new__(PromptManagerSplit)
node.logger = logging.getLogger("test.split")
node.comfyui_integration = unittest.mock.Mock()
node._inject_lora_trigger_words = lambda text: text
node._track_prompt_execution = unittest.mock.Mock(return_value=42)
clip = unittest.mock.Mock()
tests/test_thumbnail_settings.py
+1 -1
View File
@@ -78,7 +78,7 @@ class TestDefaultsAndConfig(ConfigCase):
self.assertEqual(ts.current().format, "webp")
def test_the_never_read_gallery_thumbnail_size_is_gone(self):
# R1: older save_to_file wrote its default 256 into config.json; it
# R1: older versions wrote its default 256 into config.json; it
# must never size a preview.
self.assertFalse(hasattr(GalleryConfig, "THUMBNAIL_SIZE"))
PromptManagerConfig.update_config(
tests/test_validators.py
-43
View File
@@ -17,7 +17,6 @@ from utils.validators import (
validate_tags,
validate_category,
sanitize_input,
parse_tags_string,
validate_result_timeout,
validate_worker_threads,
)
@@ -125,48 +124,6 @@ class TestSanitizeInput(unittest.TestCase):
self.assertEqual(result, "a\nb\nc\nd")
class TestParseTagsString(unittest.TestCase):
"""Test parse_tags_string function."""
def test_simple_comma_separated(self):
result = parse_tags_string("tag1, tag2, tag3")
self.assertEqual(result, ["tag1", "tag2", "tag3"])
def test_extra_whitespace(self):
result = parse_tags_string(" tag1 , tag2 ")
self.assertEqual(result, ["tag1", "tag2"])
def test_empty_string_returns_empty_list(self):
self.assertEqual(parse_tags_string(""), [])
def test_none_returns_empty_list(self):
self.assertEqual(parse_tags_string(None), [])
def test_non_string_returns_empty_list(self):
self.assertEqual(parse_tags_string(123), [])
def test_single_tag(self):
self.assertEqual(parse_tags_string("solo"), ["solo"])
def test_deduplication(self):
result = parse_tags_string("tag1, tag1, tag2, tag1")
self.assertEqual(result, ["tag1", "tag2"])
def test_empty_segments_skipped(self):
result = parse_tags_string("tag1,,tag2,,,tag3")
self.assertEqual(result, ["tag1", "tag2", "tag3"])
def test_max_20_tags(self):
tags = ", ".join(f"tag{i}" for i in range(25))
result = parse_tags_string(tags)
self.assertEqual(len(result), 20)
def test_strips_null_bytes_from_tags(self):
result = parse_tags_string("clean\x00tag, normal")
self.assertIn("cleantag", result)
self.assertIn("normal", result)
class TestValidatePromptTextEdgeCases(unittest.TestCase):
"""Additional edge cases beyond test_basic.py coverage."""
utils/comfyui_integration.py
-220
View File
@@ -1,220 +0,0 @@
"""ComfyUI integration utilities for PromptManager.
This module tracks the prompts PromptManager nodes encode during execution so
other components (e.g. image/prompt linking) can look them up.
ComfyUI already writes every node's own inputs into saved image metadata, so
SaveImage is deliberately NOT patched. An earlier patch wrote the last-run prompt
into every PromptManager node, saving negative prompts as copies of the
positive (#75).
Key features:
- Thread-safe prompt registration and retrieval
- Cross-thread prompt context sharing
- Standard metadata format compatibility
- Automatic cleanup of old prompt registrations
The integration works by:
1. PromptManager nodes register their prompts during execution
2. Callers retrieve the current prompt text by node id or recency
Typical usage:
from utils.comfyui_integration import get_comfyui_integration
integration = get_comfyui_integration()
integration.register_prompt(node_id, prompt_text, metadata)
text = integration.get_current_prompt_text(node_id)
"""
import threading
import time
from typing import Dict, Any, Optional
try:
from .logging_config import get_logger
except ImportError:
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.logging_config import get_logger
class ComfyUIMetadataIntegration:
"""Integrates PromptManager with ComfyUI's standard metadata system.
This singleton class manages the integration between PromptManager custom nodes
and ComfyUI's standard metadata system. It ensures that prompts generated by
PromptManager appear in the standard ComfyUI metadata format for compatibility
with third-party tools and parsers.
The integration uses thread-local storage combined with global tracking to
handle prompt context across different execution threads, which is necessary
because ComfyUI executions may span multiple threads.
Key responsibilities:
- Register prompts from PromptManager nodes during execution
- Manage prompt lifecycle and cleanup
"""
_instance = None
_lock = threading.Lock()
def __new__(cls):
"""Ensure singleton pattern with thread safety.
Returns:
The single instance of ComfyUIMetadataIntegration
"""
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self):
"""Initialize the ComfyUI integration system.
Sets up prompt tracking data structures. Uses _initialized flag to
prevent duplicate initialization.
"""
if hasattr(self, "_initialized"):
return
self.logger = get_logger("prompt_manager.comfyui_integration")
self._current_prompts = {}
self._thread_local = threading.local()
self._initialized = True
def register_prompt(self, node_id: str, prompt_text: str, metadata: Dict[str, Any]):
"""
Register a prompt from PromptManager for inclusion in ComfyUI metadata.
This method is called by PromptManager nodes during execution to register
their prompt text for later inclusion in image metadata. The prompt is stored
in both thread-local and global storage to handle cross-thread access scenarios.
Args:
node_id: Unique identifier for this prompt (typically the node ID)
prompt_text: The actual prompt text that was encoded by PromptManager
metadata: Additional metadata from PromptManager (category, tags, etc.)
"""
thread_id = threading.current_thread().ident
# Store in thread-local storage
if not hasattr(self._thread_local, "prompts"):
self._thread_local.prompts = {}
self._thread_local.prompts[node_id] = {
"text": prompt_text,
"metadata": metadata,
"timestamp": time.time(),
"thread_id": thread_id,
}
# Also store globally for cross-thread access
with self._lock:
self._current_prompts[f"{thread_id}_{node_id}"] = {
"text": prompt_text,
"metadata": metadata,
"timestamp": time.time(),
"thread_id": thread_id,
}
self.logger.debug(
f"Registered prompt for node {node_id}: {prompt_text[:50]}..."
)
def get_current_prompt_text(self, node_id: str = None) -> Optional[str]:
"""
Get the current prompt text for metadata inclusion.
Retrieves the most appropriate prompt text for inclusion in image metadata.
Uses a fallback strategy: thread-local storage first, then global storage,
with recency checks to avoid stale prompts.
Args:
node_id: Optional specific node ID to retrieve prompt for
Returns:
The prompt text string if available, None if no suitable prompt found
"""
# First try thread-local storage
if hasattr(self._thread_local, "prompts"):
if node_id and node_id in self._thread_local.prompts:
return self._thread_local.prompts[node_id]["text"]
elif self._thread_local.prompts:
# Return the most recent prompt from this thread
latest = max(
self._thread_local.prompts.values(), key=lambda x: x["timestamp"]
)
return latest["text"]
# Fallback to global storage
thread_id = threading.current_thread().ident
with self._lock:
# Look for prompts from current thread
thread_prompts = {
k: v
for k, v in self._current_prompts.items()
if v["thread_id"] == thread_id
}
if thread_prompts:
latest = max(thread_prompts.values(), key=lambda x: x["timestamp"])
return latest["text"]
# Last resort: return the most recent prompt from any thread
if self._current_prompts:
latest = max(
self._current_prompts.values(), key=lambda x: x["timestamp"]
)
# Only return if it's recent (within last 5 minutes)
if time.time() - latest["timestamp"] < 300:
return latest["text"]
return None
def cleanup_old_prompts(self, max_age_seconds: int = 600):
"""
Clean up old prompt registrations.
Removes prompt registrations that are older than the specified age to
prevent memory leaks and ensure that only recent, relevant prompts are
used for metadata inclusion.
Args:
max_age_seconds: Maximum age in seconds before a prompt registration
is considered stale and removed (default: 600
seconds/10 minutes)
"""
current_time = time.time()
with self._lock:
old_keys = [
key
for key, prompt_data in self._current_prompts.items()
if current_time - prompt_data["timestamp"] > max_age_seconds
]
for key in old_keys:
del self._current_prompts[key]
if old_keys:
self.logger.debug(f"Cleaned up {len(old_keys)} old prompt registrations")
# Global instance
_integration_instance = None
def get_comfyui_integration() -> ComfyUIMetadataIntegration:
"""Get the global ComfyUI integration instance.
Returns:
The singleton ComfyUIMetadataIntegration instance, creating it if necessary
"""
global _integration_instance
if _integration_instance is None:
_integration_instance = ComfyUIMetadataIntegration()
return _integration_instance
utils/hashing.py
+1 -33
View File
@@ -8,19 +8,15 @@ regardless of minor formatting differences.
Key features:
- SHA256-based prompt text hashing with normalization
- Content hashing for complex prompt metadata structures
- Duplicate detection utilities
- Consistent normalization (lowercase, trimmed whitespace)
Typical usage:
from utils.hashing import generate_prompt_hash, is_duplicate_prompt
from utils.hashing import generate_prompt_hash
hash1 = generate_prompt_hash("Beautiful landscape")
hash2 = generate_prompt_hash(" beautiful landscape ")
# hash1 == hash2 (normalization makes them identical)
if is_duplicate_prompt(text1, text2):
print("Duplicate prompts detected")
The hashing is designed for:
- Database deduplication (preventing duplicate prompt storage)
- Quick similarity checking
@@ -124,31 +120,3 @@ def stored_prompt_text_sql(expression: str) -> str:
character of Y, multi-byte ones included, so it strips what Python does.
"""
return f"trim({expression}, char({_STRIP_CODES}))"
def is_duplicate_prompt(text1: str, text2: str, threshold: float = 0.95) -> bool:
"""
Check if two prompts are likely duplicates using hash comparison.
Compares the normalized hashes of two prompt texts to determine if
they are functionally identical. This is an exact match after normalization,
not a similarity measure.
Args:
text1: First prompt text to compare
text2: Second prompt text to compare
threshold: Similarity threshold (not used - kept for API compatibility)
Returns:
True if the prompts have identical normalized hashes (i.e., are duplicates),
False otherwise
Note:
The threshold parameter is not used in the current implementation as
this performs exact hash matching rather than similarity scoring.
It's maintained for potential future fuzzy matching capabilities.
"""
hash1 = generate_prompt_hash(text1)
hash2 = generate_prompt_hash(text2)
return hash1 == hash2
utils/validators.py
-27
View File
@@ -295,30 +295,3 @@ def sanitize_input(text: str) -> str:
result_lines.append(line)
return "\n".join(result_lines).strip()
def parse_tags_string(tags_string: str) -> List[str]:
"""
Parse a comma-separated tags string into a clean list.
Converts comma-separated tag strings into a clean, deduplicated list
of tags. Each tag is sanitized and trimmed.
Args:
tags_string: Comma-separated tags string (e.g., "tag1, tag2, tag3")
Returns:
List of unique, cleaned tag strings. Empty input returns empty list.
Limited to maximum 20 tags to prevent abuse.
"""
if not tags_string or not isinstance(tags_string, str):
return []
# Split by comma and clean each tag
tags = []
for tag in tags_string.split(","):
clean_tag = sanitize_input(tag).strip()
if clean_tag and clean_tag not in tags: # Avoid duplicates
tags.append(clean_tag)
return tags[:20] # Limit to 20 tags