Vito
3be0ada6a1
* fix(ingestion): link an image only to a prompt that provably produced it An LLM fed a PromptManagerPair's prepend_text while its text editor was empty. The node refused to save (an empty positive raised and logged a warning), and the image monitor then linked the image to a stale tracker entry or to the newest prompt in the database. - Nodes: an empty text editor saves and tracks nothing (prepend and append text are never the stored prompt), logs at info, and marks the node's run key empty so no earlier entry can be used for its images. - Tracker entries carry a run key: a digest of the prompt node and its whole upstream graph, the identity ComfyUI's cache uses. The monitor computes the same key from the image's embedded graph and uses a queued or snapshot entry only when the keys match (a snapshot also needs a known run with one entry, so a batch's last item is never applied to every image). - The recency fallback is gone; anything uncertain stays unlinked and is logged at info. - Recorded execution history attributes a file when exactly one run reported it, the file was written during that run, and the run's positive contexts name one stored prompt. The monitor and the output scan use it as their last certain source. - The scan and pair ingestion treat a resolved pair with an empty stored positive as unresolved instead of raising. - tools/review_image_links.py reviews selected links by these rules; dry run by default, --apply changes only named ids after a backup. Closes #342 * test(ingestion): real-host attribution case and the #342 testing contract The managed host disables the image monitor, so the owned fixture pack gains a route that starts and stops the plugin's monitor singleton for the attribution case only. The case runs in both renderers: a typed Pair image links, a Pair image whose positive is only linked prepend text stays unlinked with nothing stored, an unknown node feeding text links through its run key, and its cached re-run links again. * fix(ingestion): take a queued entry only from the run the image appeared in Dropping a key's older queue entries whenever the node ran again let a later queued run, executing during the monitor's processing delay, take the earlier image's entry. The snapshot taken when the image appears now names the key's latest run, and the queue is popped by key and run: same-key entries of earlier runs ahead of it are stale, entries of later runs wait for their own images. Also keeps the attribution case out of the stock real-host suite. * test(ingestion): host lookups degrade to no proof * fix(ingestion): address the #344 review (history re-check, keys, repair) - Images without an embedded API graph: the scan now asks recorded history before skipping a file with no prompt metadata, and the monitor gives a keyless unlinked image one history-only re-check when its run's outputs are recorded (a seal notification from the outputs write; at most 1000 waiting images, each for an hour). - Run keys: an image is unkeyable when any positive node has no key, so a two-pass graph is never linked through the other pass; links follow ComfyUI's is_link exactly. - A failed save marks the node's run key like an empty editor, so an older entry of that key can't stand in for this run. - A placeholder is the latest entry: the image stays unlinked without falling through to history, the same rule as a popped placeholder. - An index on execution_run_outputs(filename, subfolder, root_key, path_style) for the history lookup, added by migration and accepted by the output schema check (any other index is still refused). - Repair tool: unlink removes only the link record and only while the image is still linked to the prompt it was judged on; every run prints what --apply does. - Real-host case proves the re-run's PromptManager node was cached. * test(ingestion): real-host bare-output re-check and the review-round contract A metadata-disabled save, held open behind the fixture gate after SaveImage, waits unlinked while its run is open and links once the run's outputs are recorded. The fixture's monitor route lists the waiting files. TESTING.md covers the history re-check, the read index and what the repair tool's --apply does. * fix(ingestion): batch history lookups instead of adding an index An index on execution_run_outputs would make the library unopenable by older builds, whose schema check refuses any extra index there. The migration and the schema-check allowance are reverted, so the history schema is exactly main's. stored_run_prompts instead looks up a batch of files with one parameterized filename IN (...) query per chunk of at most 500 names and matches in memory; the scan computes it once per scan batch. * fix(ingestion): address the #344 re-review (seal race, linked roots) - A sealed history re-check that can't be scheduled yet (the image's own pass still pending, or the pending set full) is kept due and scheduled when a pending image finishes, instead of being dropped. - History matches a file through ComfyUI's spelling of its output folder first, then with both sides resolved, so a symlinked or junction output folder still attributes the scan's resolved paths; waiting files are keyed by their resolved path. - The scan makes its batch history lookup only when a file needs it. - set_current_prompt pushes to the queue last, so marking a failed execution never adds a second queue slot. - A keyless image outside the output folder doesn't take a waiting slot.
663 lines
27 KiB
Python
663 lines
27 KiB
Python
"""#342: an image is linked only to a prompt that provably produced it.
|
|
|
|
The maintainer's run: an LLM writes the prompt into a PromptManagerPair's
|
|
``prepend_text`` while its ``text`` editor is empty. Prepend and append text are
|
|
never stored as the prompt (product decision), so nothing is saved, and the
|
|
image must stay unlinked instead of taking the newest prompt in the database or
|
|
a stale tracker entry.
|
|
"""
|
|
|
|
import copy
|
|
import logging
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed
|
|
|
|
from database.operations import PromptDatabase
|
|
from prompt_manager import PromptManager
|
|
from prompt_manager_pair import PromptManagerPair
|
|
from prompt_manager_text import PromptManagerText
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
from utils.prompt_tracker import PromptTracker
|
|
from utils.run_keys import image_run_keys, prompt_node_key
|
|
from utils.usage_tracking import (
|
|
PendingFirstUse,
|
|
queued_pair_runs,
|
|
queued_prompt_runs,
|
|
)
|
|
|
|
LLM_TEXT = "Photorealistic aerial landscape of misty mountains at dawn"
|
|
EVIDENCE_KEY = "comfyassets_prompt_pairs"
|
|
|
|
|
|
def issue_graph(text="", seed=1):
|
|
"""The workflow from #342 (subgraph ids kept), PromptManagerPair positive."""
|
|
return {
|
|
"18:35": {
|
|
"class_type": "PromptManagerPair",
|
|
"inputs": {
|
|
"text": text,
|
|
"negative_text": "",
|
|
"prepend_text": ["18:13", 0],
|
|
"append_text": "",
|
|
"category": "",
|
|
"tags": "",
|
|
"clip": ["4", 1],
|
|
},
|
|
},
|
|
"18:13": {"class_type": "PreviewAny", "inputs": {"source": ["18:16", 0]}},
|
|
"18:16": {
|
|
"class_type": "ComfySwitchNode",
|
|
"inputs": {
|
|
"switch": True,
|
|
"on_false": ["18:17", 0],
|
|
"on_true": ["18:11", 0],
|
|
},
|
|
},
|
|
"18:11": {
|
|
"class_type": "TextGenerate",
|
|
"inputs": {"prompt": ["18:17", 0], "seed": 7},
|
|
},
|
|
"18:17": {
|
|
"class_type": "PrimitiveStringMultiline",
|
|
"inputs": {"value": "Describe an aerial mountain landscape"},
|
|
},
|
|
"18:5": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "blurry", "clip": ["4", 1]},
|
|
},
|
|
"4": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "moss.safetensors"},
|
|
},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["4", 0],
|
|
"positive": ["18:35", 0],
|
|
"negative": ["18:5", 0],
|
|
"latent_image": ["6", 0],
|
|
"seed": seed,
|
|
},
|
|
},
|
|
"6": {"class_type": "EmptyLatentImage", "inputs": {"batch_size": 1}},
|
|
}
|
|
|
|
|
|
def llm_text_graph(seed=1):
|
|
"""An LLM feeding a PromptManager's ``text``: stored, but not in the metadata."""
|
|
graph = issue_graph(seed=seed)
|
|
graph["18:35"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": ["18:13", 0], "clip": ["4", 1]},
|
|
}
|
|
return graph
|
|
|
|
|
|
def pair_evidence(text="", prepend=LLM_TEXT):
|
|
"""The execution record PromptManagerPair embeds in the saved image."""
|
|
return {
|
|
"version": 1,
|
|
"nodes": {
|
|
"18:35": {
|
|
"records": [
|
|
{
|
|
"text": text,
|
|
"negative_text": "",
|
|
"positive_text": " ".join(p for p in (prepend, text) if p),
|
|
"negative_effective_text": "",
|
|
"prepend_text": prepend,
|
|
"append_text": "",
|
|
}
|
|
]
|
|
}
|
|
},
|
|
}
|
|
|
|
|
|
def make_tracker(db):
|
|
"""A real tracker without its cleanup thread or config import."""
|
|
with mock.patch.object(PromptTracker, "__init__", lambda self, db: None):
|
|
tracker = PromptTracker.__new__(PromptTracker)
|
|
tracker.logger = logging.getLogger("test.prompt_attribution.tracker")
|
|
tracker.db_manager = db
|
|
tracker._local = threading.local()
|
|
tracker.active_prompts = {}
|
|
tracker.lock = threading.Lock()
|
|
tracker._prompt_queue = []
|
|
tracker._queue_lock = threading.Lock()
|
|
tracker.cleanup_interval = 300
|
|
tracker.prompt_timeout = 600
|
|
return tracker
|
|
|
|
|
|
class DatabaseCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.tmp.cleanup)
|
|
self.addCleanup(assert_closed, self, self.tmp.name)
|
|
self.db = PromptDatabase(os.path.join(self.tmp.name, "prompts.db"))
|
|
self.addCleanup(self.db.close_all)
|
|
self.tracker = make_tracker(self.db)
|
|
|
|
def save(self, text):
|
|
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
|
|
|
|
|
class NodeCase(DatabaseCase):
|
|
def node(self, cls):
|
|
node = cls.__new__(cls)
|
|
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()
|
|
return node
|
|
|
|
def clip(self):
|
|
clip = mock.Mock()
|
|
clip.tokenize.side_effect = lambda text: text
|
|
clip.encode_from_tokens_scheduled.side_effect = lambda t: ("cond", t)
|
|
return clip
|
|
|
|
def count_rows(self, table):
|
|
conn = self.db.model.get_connection()
|
|
return conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
|
|
|
|
|
|
class TestEmptyEditorSavesNothing(NodeCase):
|
|
"""Prepend/append are never the stored prompt; an empty editor saves nothing."""
|
|
|
|
def test_pair_with_only_prepend_text_saves_nothing_and_logs_info(self):
|
|
graph = issue_graph()
|
|
with self.assertLogs("test.prompt_attribution.node", "INFO") as logs:
|
|
result = self.node(PromptManagerPair).encode_prompt(
|
|
self.clip(),
|
|
"",
|
|
negative_text="",
|
|
prepend_text=LLM_TEXT,
|
|
prompt=graph,
|
|
unique_id="18:35",
|
|
)
|
|
self.assertEqual(result[2], LLM_TEXT, "the composed text is still encoded")
|
|
self.assertEqual(self.count_rows("prompts"), 0)
|
|
self.assertEqual(self.count_rows("prompt_pairs"), 0)
|
|
self.assertFalse(
|
|
any(record.levelno >= logging.WARNING for record in logs.records),
|
|
[record.getMessage() for record in logs.records],
|
|
)
|
|
self.assertTrue(any("Prompt text is empty" in line for line in logs.output))
|
|
|
|
def test_prompt_manager_and_text_nodes_with_only_prepend_save_nothing(self):
|
|
graph = issue_graph()
|
|
graph["18:35"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": "", "prepend_text": ["18:13", 0], "clip": ["4", 1]},
|
|
}
|
|
self.node(PromptManager).encode_prompt(
|
|
self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
|
|
)
|
|
graph["18:35"]["class_type"] = "PromptManagerText"
|
|
self.node(PromptManagerText).process_text(
|
|
"", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
|
|
)
|
|
self.assertEqual(self.count_rows("prompts"), 0)
|
|
|
|
def test_empty_editor_marks_its_run_key_so_no_older_entry_is_used(self):
|
|
graph = issue_graph()
|
|
key = prompt_node_key(graph, "18:35")
|
|
older = self.save("a woman holding a coffee cup")
|
|
self.tracker.set_current_prompt(
|
|
"a woman holding a coffee cup",
|
|
{"prompt_id": older},
|
|
push_to_queue=True,
|
|
run_key=key,
|
|
run_id="run-1",
|
|
)
|
|
self.node(PromptManagerPair).encode_prompt(
|
|
self.clip(), "", prepend_text=LLM_TEXT, prompt=graph, unique_id="18:35"
|
|
)
|
|
latest, _run, _count = self.tracker.keyed_snapshot()[key]
|
|
self.assertTrue(latest.get("placeholder"))
|
|
self.assertIsNone(self.tracker.pop_for_key(key, "another run"))
|
|
popped = self.tracker.pop_for_key(key)
|
|
self.assertEqual(popped["id"], older, "an earlier queued item stays first")
|
|
self.assertTrue(self.tracker.pop_for_key(key).get("placeholder"))
|
|
|
|
def test_empty_negative_node_is_not_marked(self):
|
|
graph = issue_graph(text="a typed positive")
|
|
graph["18:5"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": "", "clip": ["4", 1]},
|
|
}
|
|
self.node(PromptManager).encode_prompt(
|
|
self.clip(), "", prompt=graph, unique_id="18:5"
|
|
)
|
|
self.assertEqual(self.tracker.keyed_snapshot(), {})
|
|
|
|
def test_queue_hook_counts_nothing_for_an_empty_editor(self):
|
|
graph = issue_graph()
|
|
self.assertEqual(queued_prompt_runs(graph), [])
|
|
self.assertEqual(queued_pair_runs(graph), [])
|
|
|
|
def test_typed_text_with_prepend_still_stores_only_the_editor_text(self):
|
|
graph = issue_graph(text="mountains")
|
|
self.node(PromptManagerPair).encode_prompt(
|
|
self.clip(),
|
|
"mountains",
|
|
prepend_text=LLM_TEXT,
|
|
prompt=graph,
|
|
unique_id="18:35",
|
|
)
|
|
self.assertIsNotNone(
|
|
self.db.get_prompt_by_hash(generate_prompt_hash("mountains"))
|
|
)
|
|
self.assertIsNone(self.db.get_prompt_by_hash(generate_prompt_hash(LLM_TEXT)))
|
|
context, _run, _count = self.tracker.keyed_snapshot()[
|
|
prompt_node_key(graph, "18:35")
|
|
]
|
|
self.assertEqual(context["text"], f"{LLM_TEXT} mountains")
|
|
|
|
|
|
class MonitorCase(DatabaseCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
fd, self.image = tempfile.mkstemp(suffix=".png", dir=self.tmp.name)
|
|
os.close(fd)
|
|
self.handler = ImageGenerationHandler(self.db, self.tracker)
|
|
self.linked = []
|
|
self.handler.link_image_to_prompt = lambda path, prompt, meta: (
|
|
self.linked.append(prompt["id"])
|
|
)
|
|
|
|
def process(self, metadata, snapshot=None):
|
|
before = len(self.linked)
|
|
with mock.patch.object(
|
|
self.handler.metadata_extractor, "extract_metadata", return_value=metadata
|
|
):
|
|
self.handler.process_new_image(self.image, prompt_snapshot=snapshot)
|
|
return self.linked[-1] if len(self.linked) > before else None
|
|
|
|
|
|
class TestIssueReproduction(MonitorCase):
|
|
def test_llm_prepend_image_is_left_unlinked_not_given_the_newest_prompt(self):
|
|
"""Failing before #342: linked to the newest prompt in the database."""
|
|
self.save("a woman holding a coffee cup")
|
|
metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()}
|
|
with self.assertLogs("prompt_manager.image_monitor", "INFO") as logs:
|
|
self.assertIsNone(self.process(metadata))
|
|
self.assertTrue(any("left unlinked" in line for line in logs.output))
|
|
|
|
def test_a_stale_entry_of_another_workflow_is_never_used(self):
|
|
"""Failing before #342: the live tracker's entry was taken."""
|
|
stale = self.save("a woman holding a coffee cup")
|
|
other = issue_graph(text="a woman holding a coffee cup")
|
|
self.tracker.set_current_prompt(
|
|
"a woman holding a coffee cup",
|
|
{"prompt_id": stale},
|
|
push_to_queue=True,
|
|
run_key=prompt_node_key(other, "18:35"),
|
|
run_id="run-0",
|
|
)
|
|
metadata = {"prompt": issue_graph(), EVIDENCE_KEY: pair_evidence()}
|
|
self.assertIsNone(self.process(metadata, self.tracker.keyed_snapshot()))
|
|
self.assertIsNone(self.process(metadata))
|
|
self.assertEqual(len(self.tracker._prompt_queue), 1, "not consumed")
|
|
|
|
def test_an_entry_without_a_run_key_is_never_used(self):
|
|
stale = self.save("previous run")
|
|
self.tracker.set_current_prompt("previous run", {"prompt_id": stale})
|
|
self.assertIsNone(self.process({"prompt": llm_text_graph()}))
|
|
|
|
def test_an_image_without_metadata_is_left_unlinked(self):
|
|
"""Failing before #342: popped the queue, else took the newest prompt."""
|
|
self.save("latest")
|
|
self.assertIsNone(self.process(None))
|
|
|
|
def test_the_recency_fallback_is_gone(self):
|
|
self.assertFalse(hasattr(self.handler, "_get_fallback_prompt"))
|
|
|
|
|
|
class TestProvenRuns(MonitorCase):
|
|
"""Entries the image's own graph proves still link (LLM into ``text``)."""
|
|
|
|
def track(self, graph, text, run_id="run-1", push=True):
|
|
prompt_id = (
|
|
self.save(text)
|
|
if not self.db.get_prompt_by_hash(generate_prompt_hash(text))
|
|
else self.db.get_prompt_by_hash(generate_prompt_hash(text))["id"]
|
|
)
|
|
self.tracker.set_current_prompt(
|
|
text,
|
|
{"prompt_id": prompt_id},
|
|
push_to_queue=push,
|
|
run_key=prompt_node_key(graph, "18:35"),
|
|
run_id=run_id,
|
|
)
|
|
return prompt_id
|
|
|
|
def test_queued_entry_of_the_same_run_key_links(self):
|
|
graph = llm_text_graph()
|
|
prompt_id = self.track(graph, LLM_TEXT)
|
|
self.assertEqual(self.process({"prompt": graph}), prompt_id)
|
|
|
|
def test_cached_node_rerun_links_through_the_snapshot(self):
|
|
graph = llm_text_graph(seed=1)
|
|
prompt_id = self.track(graph, LLM_TEXT)
|
|
self.assertEqual(self.process({"prompt": graph}), prompt_id)
|
|
# Only the sampler seed changed: the node was cached and did not run,
|
|
# but its upstream (the run key) is unchanged.
|
|
rerun = llm_text_graph(seed=2)
|
|
self.assertEqual(
|
|
prompt_node_key(rerun, "18:35"), prompt_node_key(graph, "18:35")
|
|
)
|
|
self.assertEqual(
|
|
self.process({"prompt": rerun}, self.tracker.keyed_snapshot()), prompt_id
|
|
)
|
|
|
|
def test_snapshot_needs_a_known_run(self):
|
|
graph = llm_text_graph()
|
|
self.track(graph, LLM_TEXT, run_id=None, push=False)
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_batch_items_link_in_order_and_the_latest_alone_proves_nothing(self):
|
|
graph = llm_text_graph()
|
|
first = self.track(graph, "item one")
|
|
second = self.track(graph, "item two")
|
|
self.assertEqual(self.process({"prompt": graph}), first)
|
|
self.assertEqual(self.process({"prompt": graph}), second)
|
|
# Queue drained: two entries in one run, so the latest is not proof
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_an_empty_batch_item_keeps_the_queue_aligned(self):
|
|
graph = llm_text_graph()
|
|
key = prompt_node_key(graph, "18:35")
|
|
first = self.track(graph, "item one")
|
|
self.tracker.mark_empty(key, "run-1", push_to_queue=True)
|
|
third = self.track(graph, "item three")
|
|
self.assertEqual(self.process({"prompt": graph}), first)
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
self.assertEqual(self.process({"prompt": graph}), third)
|
|
|
|
def test_a_newer_run_drops_stale_entries_of_the_same_key(self):
|
|
graph = llm_text_graph()
|
|
# run-1 saved no image, so its entry was never taken
|
|
self.track(graph, "earlier run text", run_id="run-1")
|
|
newer = self.track(graph, "newer run text", run_id="run-2")
|
|
self.assertEqual(self.process({"prompt": graph}), newer)
|
|
self.assertEqual(self.tracker._prompt_queue, [])
|
|
|
|
def test_a_later_run_queued_before_processing_keeps_its_own_entry(self):
|
|
# The image appears (snapshot), the next queued run executes the same
|
|
# node again before the monitor's delay is over, then the image is
|
|
# processed: it takes its own run's entry, the later run keeps its own.
|
|
graph = llm_text_graph()
|
|
own = self.track(graph, "first run text", run_id="run-1")
|
|
snapshot = self.tracker.keyed_snapshot()
|
|
later = self.track(graph, "second run text", run_id="run-2")
|
|
self.assertEqual(self.process({"prompt": graph}, snapshot), own)
|
|
self.assertEqual(self.process({"prompt": graph}), later)
|
|
|
|
def test_no_entry_when_the_image_appeared_means_no_queue_match(self):
|
|
graph = llm_text_graph()
|
|
empty_snapshot = self.tracker.keyed_snapshot()
|
|
self.track(graph, "queued after the image appeared")
|
|
self.assertIsNone(self.process({"prompt": graph}, empty_snapshot))
|
|
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
|
|
|
def test_hosts_without_run_ids_still_link_queued_entries_in_order(self):
|
|
graph = llm_text_graph()
|
|
first = self.track(graph, "item one", run_id=None)
|
|
second = self.track(graph, "item two", run_id=None)
|
|
self.assertEqual(self.process({"prompt": graph}), first)
|
|
self.assertEqual(self.process({"prompt": graph}), second)
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_other_keys_stay_queued_for_their_own_images(self):
|
|
mine, other = llm_text_graph(), copy.deepcopy(llm_text_graph())
|
|
other["18:17"]["inputs"]["value"] = "Describe a beach"
|
|
theirs = self.track(other, "a beach")
|
|
prompt_id = self.track(mine, LLM_TEXT)
|
|
self.assertEqual(self.process({"prompt": mine}), prompt_id)
|
|
self.assertEqual(self.process({"prompt": other}), theirs)
|
|
|
|
def test_two_positive_prompt_nodes_are_not_guessed(self):
|
|
graph = llm_text_graph()
|
|
self.track(graph, LLM_TEXT)
|
|
graph["30"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": ["18:13", 0], "clip": ["4", 1]},
|
|
}
|
|
graph["31"] = dict(graph["3"])
|
|
graph["31"]["inputs"] = {**graph["3"]["inputs"], "positive": ["30", 0]}
|
|
self.assertEqual(len(image_run_keys({"prompt": graph})), 2)
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_metadata_still_wins_and_leaves_the_queue_alone(self):
|
|
typed = self.save("a typed prompt")
|
|
graph = issue_graph(text="a typed prompt")
|
|
self.tracker.set_current_prompt(
|
|
"x",
|
|
{"prompt_id": typed},
|
|
run_key=prompt_node_key(graph, "18:35"),
|
|
run_id="r",
|
|
)
|
|
self.assertEqual(self.process({"prompt": graph}), typed)
|
|
self.assertEqual(len(self.tracker._prompt_queue), 1)
|
|
|
|
|
|
class TestRunKeys(unittest.TestCase):
|
|
def test_key_covers_the_upstream_and_ignores_downstream_and_titles(self):
|
|
graph = issue_graph()
|
|
key = prompt_node_key(graph, "18:35")
|
|
changed = copy.deepcopy(graph)
|
|
changed["3"]["inputs"]["seed"] = 99
|
|
changed["18:35"]["_meta"] = {"title": "renamed"}
|
|
self.assertEqual(prompt_node_key(changed, "18:35"), key)
|
|
changed["18:11"]["inputs"]["seed"] = 8
|
|
self.assertNotEqual(prompt_node_key(changed, "18:35"), key)
|
|
|
|
def test_unknowable_keys_are_none(self):
|
|
graph = issue_graph()
|
|
self.assertIsNone(prompt_node_key(None, "1"))
|
|
self.assertIsNone(prompt_node_key(graph, None))
|
|
self.assertIsNone(prompt_node_key(graph, "missing"))
|
|
broken = copy.deepcopy(graph)
|
|
broken["18:13"]["inputs"]["source"] = ["gone", 0]
|
|
self.assertIsNone(prompt_node_key(broken, "18:35"))
|
|
broken = copy.deepcopy(graph)
|
|
broken["18:17"]["inputs"]["value"] = object()
|
|
self.assertIsNone(prompt_node_key(broken, "18:35"))
|
|
|
|
def test_image_keys_read_string_graphs_and_reject_garbage(self):
|
|
import json
|
|
|
|
graph = issue_graph()
|
|
self.assertEqual(
|
|
image_run_keys({"prompt": json.dumps(graph)}),
|
|
{"18:35": prompt_node_key(graph, "18:35")},
|
|
)
|
|
self.assertEqual(image_run_keys({"prompt": "{not json"}), {})
|
|
self.assertEqual(image_run_keys(None), {})
|
|
self.assertEqual(image_run_keys({"prompt": []}), {})
|
|
|
|
def test_cyclic_and_huge_closures_have_no_key(self):
|
|
from utils.prompt_graph import MAX_VISITS
|
|
|
|
cyclic = {
|
|
"1": {"class_type": "A", "inputs": {"x": ["2", 0]}},
|
|
"2": {"class_type": "B", "inputs": {"x": ["1", 0]}},
|
|
}
|
|
self.assertIsNotNone(prompt_node_key(cyclic, "1"))
|
|
chain = {
|
|
str(i): {"class_type": "N", "inputs": {"x": [str(i + 1), 0]}}
|
|
for i in range(MAX_VISITS + 1)
|
|
}
|
|
chain[str(MAX_VISITS + 1)] = {"class_type": "N", "inputs": {}}
|
|
self.assertIsNone(prompt_node_key(chain, "0"))
|
|
|
|
|
|
class TestReviewFindings(MonitorCase, NodeCase):
|
|
"""Final-review findings M1-M4 (#344)."""
|
|
|
|
def test_m1_a_positive_node_without_a_key_makes_the_image_unkeyable(self):
|
|
graph = llm_text_graph()
|
|
graph["30"] = {
|
|
"class_type": "PromptManager",
|
|
"inputs": {"text": ["missing", 0], "clip": ["4", 1]},
|
|
}
|
|
graph["31"] = {
|
|
"class_type": "KSampler",
|
|
"inputs": {**graph["3"]["inputs"], "positive": ["30", 0]},
|
|
}
|
|
self.assertIsNone(prompt_node_key(graph, "30"))
|
|
self.assertEqual(image_run_keys({"prompt": graph}), {})
|
|
prompt_id = self.save(LLM_TEXT)
|
|
self.tracker.set_current_prompt(
|
|
LLM_TEXT,
|
|
{"prompt_id": prompt_id},
|
|
run_key=prompt_node_key(graph, "18:35"),
|
|
run_id="run-1",
|
|
)
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_m2_links_are_exactly_comfyui_is_link(self):
|
|
from utils.run_keys import _links
|
|
|
|
inputs = {
|
|
"string_slot_int": ["7", 0],
|
|
"string_slot_float": ["7", 0.0],
|
|
"string_slot_bool": ["7", True], # bool is an int to ComfyUI too
|
|
"int_source": [7, 0],
|
|
"tuple": ("7", 0),
|
|
"three": ["7", 0, 1],
|
|
"slot_text": ["7", "0"],
|
|
}
|
|
self.assertEqual(sorted(_links(inputs)), ["7", "7", "7"])
|
|
widget = issue_graph()
|
|
widget["18:17"]["inputs"]["size"] = [512, 768] # a value, not a link
|
|
self.assertIsNotNone(prompt_node_key(widget, "18:35"))
|
|
floated = issue_graph()
|
|
floated["18:35"]["inputs"]["prepend_text"] = ["18:13", 0.0]
|
|
changed = copy.deepcopy(floated)
|
|
changed["18:17"]["inputs"]["value"] = "another instruction"
|
|
self.assertNotEqual(
|
|
prompt_node_key(floated, "18:35"), prompt_node_key(changed, "18:35")
|
|
)
|
|
|
|
def _failing_save(self, node):
|
|
def explode(*args, **kwargs):
|
|
raise RuntimeError("database is locked")
|
|
|
|
node.db = mock.Mock(wraps=self.db)
|
|
for name in ("save_prompt", "save_prompt_pair", "get_prompt_by_hash"):
|
|
setattr(node.db, name, mock.Mock(side_effect=explode))
|
|
return node
|
|
|
|
def test_m3_a_failed_save_takes_over_its_run_key(self):
|
|
graph = llm_text_graph()
|
|
key = prompt_node_key(graph, "18:35")
|
|
older = self.save("previous run text")
|
|
self.tracker.set_current_prompt(
|
|
"previous run text", {"prompt_id": older}, run_key=key, run_id="run-1"
|
|
)
|
|
node = self._failing_save(self.node(PromptManager))
|
|
with mock.patch("prompt_manager_base.executing_run_id", return_value="run-2"):
|
|
node.encode_prompt(
|
|
self.clip(), "new run text", prompt=graph, unique_id="18:35"
|
|
)
|
|
latest, run_id, _count = self.tracker.keyed_snapshot()[key]
|
|
self.assertTrue(latest.get("placeholder"))
|
|
self.assertEqual(run_id, "run-2")
|
|
self.assertIsNone(self.process({"prompt": graph}))
|
|
|
|
def test_m3_a_failed_pair_save_takes_over_its_run_key(self):
|
|
graph = issue_graph(text="typed")
|
|
key = prompt_node_key(graph, "18:35")
|
|
node = self._failing_save(self.node(PromptManagerPair))
|
|
with self.assertLogs("test.prompt_attribution.node", "WARNING"):
|
|
node.encode_prompt(self.clip(), "typed", prompt=graph, unique_id="18:35")
|
|
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
|
|
|
def test_m3_a_failed_text_node_save_takes_over_its_run_key(self):
|
|
graph = issue_graph()
|
|
graph["18:35"] = {
|
|
"class_type": "PromptManagerText",
|
|
"inputs": {"text": "typed"},
|
|
}
|
|
graph["40"] = {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": ["18:35", 0], "clip": ["4", 1]},
|
|
}
|
|
graph["3"]["inputs"]["positive"] = ["40", 0]
|
|
node = self._failing_save(self.node(PromptManagerText))
|
|
with mock.patch.object(
|
|
node, "_track_prompt_execution", side_effect=RuntimeError("boom")
|
|
):
|
|
node.process_text("typed", prompt=graph, unique_id="18:35")
|
|
key = prompt_node_key(graph, "18:35")
|
|
self.assertTrue(self.tracker.keyed_snapshot()[key][0].get("placeholder"))
|
|
|
|
def test_m4_a_placeholder_never_falls_through_to_history(self):
|
|
graph = llm_text_graph()
|
|
key = prompt_node_key(graph, "18:35")
|
|
self.tracker.mark_empty(key, "run-1", push_to_queue=False)
|
|
with mock.patch.object(
|
|
self.handler, "_history_prompt", return_value=99
|
|
) as history:
|
|
self.assertIsNone(self.process({"prompt": graph})) # snapshot form
|
|
self.tracker.mark_empty(key, "run-2", push_to_queue=True)
|
|
self.assertIsNone(self.process({"prompt": graph})) # popped form
|
|
history.assert_not_called()
|
|
|
|
|
|
class TestRereviewFindings(MonitorCase, NodeCase):
|
|
"""Re-review findings R4 and R5 (#344)."""
|
|
|
|
def test_r4_a_logging_failure_after_tracking_never_doubles_a_queue_slot(self):
|
|
graph = llm_text_graph()
|
|
key = prompt_node_key(graph, "18:35")
|
|
node = self.node(PromptManager)
|
|
logger = mock.Mock(wraps=self.tracker.logger)
|
|
logger.debug.side_effect = RuntimeError("log handler broke")
|
|
self.tracker.logger = logger
|
|
node.encode_prompt(self.clip(), LLM_TEXT, prompt=graph, unique_id="18:35")
|
|
queued = [ctx for ctx in self.tracker._prompt_queue if ctx["run_key"] == key]
|
|
self.assertEqual(len(queued), 1)
|
|
|
|
def test_r5_a_keyless_image_outside_the_output_folder_does_not_wait(self):
|
|
outside = tempfile.mkdtemp()
|
|
self.addCleanup(lambda: __import__("shutil").rmtree(outside, True))
|
|
stray = os.path.join(outside, "copied.png")
|
|
with open(stray, "wb") as handle:
|
|
handle.write(b"png")
|
|
with mock.patch(
|
|
"utils.image_monitor.host_output_directory", return_value=self.tmp.name
|
|
):
|
|
with mock.patch.object(
|
|
self.handler.metadata_extractor, "extract_metadata", return_value=None
|
|
):
|
|
self.handler.process_new_image(stray)
|
|
self.handler.process_new_image(self.image) # inside: waits
|
|
self.assertEqual(
|
|
self.handler.awaiting_history(),
|
|
[os.path.normcase(os.path.realpath(self.image))],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|