diff --git a/database/execution_history_schema.py b/database/execution_history_schema.py index 191033d3..0f04f04b 100644 --- a/database/execution_history_schema.py +++ b/database/execution_history_schema.py @@ -314,12 +314,16 @@ def migrate_execution_history(conn: sqlite3.Connection) -> None: if version == 1: install_context_schema(conn) install_output_schema(conn) + from .execution_output_schema import OUTPUT_READ_INDEX_COLUMNS + # Read projections need reverse content lookup and the complete run order. for table, columns in ( ("execution_prompt_snapshots", "raw_text_id,snapshot_id"), ("execution_contexts", "positive_snapshot_id,context_id"), ("execution_run_contexts", "context_id,run_id,anchor_id"), ("execution_runs", "observed_dequeue_at,session_id,queue_item_id,run_id"), + # Attribution of an output file to its run (#342) + ("execution_run_outputs", ",".join(OUTPUT_READ_INDEX_COLUMNS)), ): conn.execute( f"CREATE INDEX IF NOT EXISTS idx_history_read_{table} " diff --git a/database/execution_output_schema.py b/database/execution_output_schema.py index c39867e5..b6e06381 100644 --- a/database/execution_output_schema.py +++ b/database/execution_output_schema.py @@ -119,6 +119,25 @@ def install_output_schema(conn: sqlite3.Connection) -> None: conn.execute(_table_sql(table)) +# The one optional, non-unique read index (#342): output file -> run. Older +# libraries may lack it; migrate_execution_history adds it. +OUTPUT_READ_INDEX = "idx_history_read_execution_run_outputs" +OUTPUT_READ_INDEX_COLUMNS = ("filename", "subfolder", "root_key", "path_style") + + +def _is_read_index(table: str, index: tuple, keys: tuple) -> bool: + """Exactly the attribution read index, as the migration creates it.""" + return ( + table == "execution_run_outputs" + and index[1] == OUTPUT_READ_INDEX + and not index[2] + and index[3] == "c" + and not index[4] + and tuple(key[2] for key in keys) == OUTPUT_READ_INDEX_COLUMNS + and not any(key[3] or key[4] != "BINARY" for key in keys) + ) + + def output_schema_error(conn: sqlite3.Connection) -> str | None: """Reject altered schemas and invalid rows before returning bounded DTOs.""" if conn.execute( @@ -157,6 +176,8 @@ def output_schema_error(conn: sqlite3.Connection) -> str | None: keys = tuple( row for row in conn.execute(f'PRAGMA index_xinfo("{name}")') if row[5] ) + if _is_read_index(table, index, keys): + continue if ( not index[2] or index[3] != "pk" diff --git a/database/execution_outputs.py b/database/execution_outputs.py index 1c3ed8c6..7435f76b 100644 --- a/database/execution_outputs.py +++ b/database/execution_outputs.py @@ -4,6 +4,8 @@ import sqlite3 from contextlib import ExitStack from typing import Callable, ContextManager +from .run_attribution import notify_outputs_sealed + try: from ..utils.execution_history_events import RunFinish, encode_command from ..utils.execution_outputs import OutputBundle, validate_output_bundle @@ -113,6 +115,10 @@ class ExecutionOutputOperations: "VALUES(?,?,?,?,?,?,?,?)", [(run_id, *reference) for reference in prepared], ) + # Committed: images waiting for this run may now be attributed (#342) + notify_outputs_sealed( + run_id, [dict(zip(_REFERENCE_FIELDS, row)) for row in prepared] + ) return True def get_execution_outputs(self, run_id: str) -> dict: diff --git a/database/run_attribution.py b/database/run_attribution.py index f2a6c1ef..97f19f35 100644 --- a/database/run_attribution.py +++ b/database/run_attribution.py @@ -16,9 +16,12 @@ only when all of this holds: Anything else returns None: the caller leaves the file unlinked. Read-only. """ +import logging import os import sqlite3 -from typing import Any, Optional +import threading +import weakref +from typing import Any, Callable, Iterable, List, Optional try: from ..utils.execution_outputs import snapshot_output_roots @@ -120,3 +123,36 @@ def stored_run_prompt( return None (prompt_id,) = named return prompt_id if isinstance(prompt_id, int) else None + + +_LOGGER = logging.getLogger("prompt_manager.run_attribution") +_SEALED_LISTENERS: List[weakref.WeakMethod] = [] +_SEALED_LOCK = threading.Lock() + + +def subscribe_outputs_sealed(method: Callable[[str, List[dict]], None]) -> None: + """Call bound ``method(run_id, references)`` after a run's outputs are recorded. + + Held weakly: a listener that is garbage collected is dropped. + """ + with _SEALED_LOCK: + _SEALED_LISTENERS.append(weakref.WeakMethod(method)) + + +def notify_outputs_sealed(run_id: str, references: Iterable[dict]) -> None: + """Tell listeners a run's output references were just committed. + + A listener failure is logged and never reaches the history writer. + """ + references = list(references) + with _SEALED_LOCK: + alive = [ref for ref in _SEALED_LISTENERS if ref() is not None] + _SEALED_LISTENERS[:] = alive + listeners = [ref() for ref in alive] + for listener in listeners: + if listener is None: + continue + try: + listener(run_id, references) + except Exception: + _LOGGER.warning("Output seal listener failed", exc_info=True) diff --git a/prompt_manager.py b/prompt_manager.py index 4eb5d006..b9b3f5ed 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -238,6 +238,7 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): except Exception as e: # Log error but don't fail the encoding self.logger.warning(f"Failed to save prompt to database: {e}") + self._mark_untracked(prompt, unique_id) else: self._skip_empty_prompt(prompt, unique_id) diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 74d95ddf..5d34acc0 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -170,7 +170,10 @@ class PromptManagerBase: prompt_id = self._save_prompt_to_database( text=text, category=category, tags=tags, count_run=batch_item ) - if prompt_id and role["is_run_prompt"]: + if role["is_run_prompt"] and not prompt_id: + # The save failed: an older entry of this key must not stand in (#342) + self._mark_untracked(prompt_graph, unique_id) + elif role["is_run_prompt"]: self.prompt_tracker.set_current_prompt( prompt_text=encoding_text, additional_data={**additional_data, "prompt_id": prompt_id}, @@ -193,25 +196,40 @@ class PromptManagerBase: Prepend and append text are deliberately never stored as the prompt, so an empty editor has no prompt to save, count or link. For the run's positive node this is logged at info and the node's run key is marked - empty, so no earlier prompt can be linked to this run's images (and a - batch item keeps its place in the queue). + empty (``_mark_untracked``). + """ + self._mark_untracked( + prompt_graph, + unique_id, + "Prompt text is empty (prepend and append text are not stored): " + "nothing saved or tracked, and this run's images stay unlinked", + ) + + def _mark_untracked( + self, prompt_graph: Any, unique_id: Any, message: Optional[str] = None + ) -> None: + """This execution of a positive node has no prompt to link (#342). + + An empty editor, or a save or tracking failure. The node's run key is + marked with a placeholder, so no earlier prompt of that key can be + linked to this run's images, and a batch item keeps its queue place. + ``message`` is logged at info for a positive node. """ try: role = self._prompt_role(prompt_graph, unique_id, "") if not role["is_run_prompt"]: - self.logger.debug("Empty text on a non-positive node; nothing saved") + if message: + self.logger.debug("Empty text on a non-positive node; not saved") return - self.logger.info( - "Prompt text is empty (prepend and append text are not stored): " - "nothing saved or tracked, and this run's images stay unlinked" - ) + if message: + self.logger.info(message) self.prompt_tracker.mark_empty( push_to_queue=role["known"], **self._run_identity(prompt_graph, unique_id), ) except Exception as error: # Tracking must never block generation - self.logger.warning(f"Could not mark the empty prompt: {error}") + self.logger.warning(f"Could not mark the untracked prompt: {error}") def _record_use(self, prompt_id: int, times: int = 1) -> None: """Count runs; a tracking failure must never block generation.""" diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py index 1aeb2976..5c903338 100644 --- a/prompt_manager_pair.py +++ b/prompt_manager_pair.py @@ -138,6 +138,8 @@ class PromptManagerPair(PromptManagerFocused): ) except Exception as error: 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()}", diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 272db9a8..b09945df 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -216,6 +216,7 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): except Exception as e: # Log error but don't fail the processing self.logger.warning(f"Failed to save prompt to database: {e}") + self._mark_untracked(prompt, unique_id) else: self._skip_empty_prompt(prompt, unique_id) diff --git a/py/api/admin.py b/py/api/admin.py index 637baa65..d9fc7b0b 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1654,7 +1654,9 @@ class AdminRoutesMixin: a dict. """ if not metadata: - return None + # No prompt metadata (--disable-metadata, bare savers): only the + # recorded run can say which prompt made it (#342) + return self._link_recorded_run(media_file, {}, detailed) canonical = dict(metadata) for key in ("prompt", "workflow"): if isinstance(canonical.get(key), str): @@ -1664,7 +1666,7 @@ class AdminRoutesMixin: # Preserve malformed metadata for conservative resolver rejection. pass if not canonical.get("prompt") and not canonical.get("parameters"): - return None + return self._link_recorded_run(media_file, canonical, detailed) _with_file_size(canonical, media_file) resolution = resolve_prompt_pair(canonical) if blank_positive(resolution): @@ -1698,7 +1700,7 @@ class AdminRoutesMixin: # Graphs without a recognized sampler keep the legacy extractor below. prompt_text = self._graph_prompt_text(metadata) if not prompt_text: - linked = self._link_recorded_run(media_file, canonical) + linked = self._link_recorded_run(media_file, canonical, detailed) if linked is not None: return linked return "unresolved" if unchoosable else resolution["state"] @@ -1724,12 +1726,13 @@ class AdminRoutesMixin: return status return {"status": status, "title": self._scan_title(media_file, prompt_text)} - def _link_recorded_run(self, media_file, canonical): + def _link_recorded_run(self, media_file, canonical, detailed=False): """Link a file recorded history attributes to one stored prompt (#342). - Returns the scan status, or None when history names no single stored + Returns the scan status (with ``detailed``, a stored or linked file's + ``{"status", "title"}``), or None when history names no single stored prompt for the file (``database.run_attribution``). Never creates a - prompt. + prompt. Works without any prompt metadata in the file. """ try: prompt_id = stored_run_prompt( @@ -1749,7 +1752,14 @@ class AdminRoutesMixin: except Exception as e: self.logger.error(f"Failed to link {media_file.name} from history: {e}") return "found" - return "linked" if image_id else "unchanged" + status = "linked" if image_id else "unchanged" + if not detailed or status not in SCAN_TITLED_DETAILS: + return status + prompt = self.db.get_prompt_by_id(prompt_id) or {} + return { + "status": status, + "title": self._scan_title(media_file, prompt.get("text") or ""), + } def _scan_title(self, media_file, text): """The linked prompt's Library title (its stored title, else ``text``'s diff --git a/tests/comfyui/image-attribution.spec.mjs b/tests/comfyui/image-attribution.spec.mjs index e197e9c2..6d280ad7 100644 --- a/tests/comfyui/image-attribution.spec.mjs +++ b/tests/comfyui/image-attribution.spec.mjs @@ -37,7 +37,7 @@ function graph({ kind = 'PromptManagerPair', text = '', prepend = '', llm = '', return nodes; } -async function saved(request, prompt) { +async function saved(request, prompt, history = {}) { const id = randomUUID(); const response = await request.post('/prompt', { data: { prompt, prompt_id: id } }); expect(response.ok()).toBe(true); @@ -48,6 +48,7 @@ async function saved(request, prompt) { }, { timeout: 15000 }).toBe(true); const [image] = item.outputs['70'].images; expect(image.type).toBe('output'); + history.item = item; return image.filename; } @@ -62,11 +63,19 @@ async function monitoring(execution, enabled) { expect((await response.json()).monitoring).toBe(enabled); } +async function cacheMode(execution, mode) { + expect((await execution.request.post('/pm_execution_fixture/cache', { data: { mode } })).ok()).toBe(true); +} + test.beforeEach(async ({ execution }) => { expect((await execution.request.post('/pm_execution_fixture/metadata', { data: { disabled: false } })).ok()).toBe(true); + await cacheMode(execution, 'CLASSIC'); // keeps the re-run's PromptManager node cached await monitoring(execution, true); }); -test.afterEach(async ({ execution }) => { await monitoring(execution, false); }); +test.afterEach(async ({ execution }) => { + await monitoring(execution, false); + await cacheMode(execution, 'RAM_PRESSURE'); +}); test('an output whose positive is only linked prepend text stays unlinked, never given an earlier prompt', async ({ execution }, testInfo) => { const tag = `${testInfo.project.name} ${randomUUID()}`; @@ -89,9 +98,13 @@ test('an output whose positive is only linked prepend text stays unlinked, never const state = await readLinks(execution.databasePath, [orphan], [prepended]); expect(state).toEqual({ links: { [orphan]: null }, stored: { [prepended]: false } }); - // Only the sampler seed changes: the PromptManager node may be served from the - // cache without running, and its unchanged upstream still proves the prompt. + // Only the sampler seed changes: the PromptManager node is served from the cache + // without running (the host says so), so no queue entry exists for this run and + // only the snapshot's run match can link it. + const rerunHistory = {}; const rerun = await saved(execution.request, { ...llmGraph, '40': { ...llmGraph['40'], - inputs: { ...llmGraph['40'].inputs, seed: 2 } } }); + inputs: { ...llmGraph['40'].inputs, seed: 2 } } }, rerunHistory); + const cached = rerunHistory.item.status.messages.find(([kind]) => kind === 'execution_cached'); + expect(cached?.[1].nodes).toEqual(expect.arrayContaining(['11', '20'])); await linkedTo(execution, rerun, llmText); }); diff --git a/tests/test_prompt_attribution.py b/tests/test_prompt_attribution.py index cb3e9e79..e54d6cbb 100644 --- a/tests/test_prompt_attribution.py +++ b/tests/test_prompt_attribution.py @@ -510,5 +510,119 @@ class TestRunKeys(unittest.TestCase): 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() + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_review_image_links_tool.py b/tests/test_review_image_links_tool.py index 98257760..a88fb10f 100644 --- a/tests/test_review_image_links_tool.py +++ b/tests/test_review_image_links_tool.py @@ -17,6 +17,7 @@ except ImportError: # discovered with tests/ as the top-level directory from open_handles import assert_closed from database.operations import PromptDatabase +from tests.test_api_images import ImageAPITestCase from tests.test_prompt_attribution import issue_graph, pair_evidence from tools import review_image_links as tool from utils.hashing import generate_prompt_hash @@ -31,7 +32,7 @@ def png(path, prompt=None, evidence=None): Image.new("RGB", (4, 4)).save(path, pnginfo=info) -class ReviewToolTests(unittest.TestCase): +class ToolFixture(unittest.TestCase): def setUp(self): self.tmp = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, self.tmp, True) @@ -81,12 +82,14 @@ class ReviewToolTests(unittest.TestCase): finally: db.close_all() + +class ReviewToolTests(ToolFixture): def test_dry_run_judges_and_changes_nothing(self): before = self.links() ids = f"{self.guessed[0]}-{self.gone}" code, out = self.run_tool("--ids", ids, "--json") self.assertEqual(code, 0) - verdicts = {v["image_id"]: v for v in json.loads(out.split("Dry run")[0])} + verdicts = {v["image_id"]: v for v in json.loads(out.split("--apply, for")[0])} for image_id in self.guessed: self.assertEqual(verdicts[image_id]["verdict"], "unlink") self.assertEqual(verdicts[image_id]["sources"], {}) @@ -114,7 +117,7 @@ class ReviewToolTests(unittest.TestCase): "--ids", named, "--prompt", str(self.coffee), "--apply" ) self.assertEqual(code, 0) - self.assertIn("Applied: 1 relinked, 2 unlinked", out) + self.assertIn("Applied: 1 relinked, 2 unlinked, 0 skipped", out) links = self.links() self.assertNotIn(self.guessed[0], links) self.assertNotIn(self.guessed[1], links) @@ -171,5 +174,82 @@ class ReviewToolTests(unittest.TestCase): self.assertEqual(verdict["verdict"], "conflict") +class ApplyGuardTests(ToolFixture): + """M6: unlink clears only the link, and only the link it was judged on.""" + + def test_unlink_skips_a_row_linked_elsewhere_since_the_review(self): + db = PromptDatabase(self.db_path) + try: + verdicts = tool.review(db, [self.guessed[0]], None, self.output) + self.assertEqual(verdicts[0]["verdict"], "unlink") + db.model.get_connection().execute( + "UPDATE generated_images SET prompt_id=? WHERE id=?", + (self.typed, self.guessed[0]), + ) + db.model.get_connection().commit() + result = tool.apply(db, verdicts, [self.guessed[0]]) + finally: + db.close_all() + self.assertEqual(result["unlinked"], 0) + self.assertEqual(self.links()[self.guessed[0]], self.typed) + + def test_unlink_keeps_the_file_and_its_favorite(self): + path = os.path.join(self.output, "moss_00009_.png") + db = PromptDatabase(self.db_path) + try: + db.set_image_favorite("a" * 64, "moss_00009_.png") + finally: + db.close_all() + self.run_tool("--ids", str(self.guessed[0]), "--apply") + self.assertTrue(os.path.isfile(path)) + db = PromptDatabase(self.db_path) + try: + self.assertEqual(db.get_root_favorites("a" * 64), {"moss_00009_.png"}) + summary = db.get_image_link_summaries([path]).get(path) or {} + finally: + db.close_all() + self.assertNotEqual(summary.get("status"), "linked") + + def test_dry_run_states_what_apply_does(self): + _code, out = self.run_tool("--ids", str(self.guessed[0])) + self.assertIn("removes only the image's link record", out) + self.assertIn("file stays on disk and in the gallery as unlinked", out) + + +class GalleryAfterUnlinkTests(ImageAPITestCase): + """M6: after --apply the file is still listed, as an unlinked image.""" + + async def test_unlinked_file_stays_in_the_workspace_gallery(self): + db = self.api.db + prompt = db.save_prompt( + text="a woman holding a coffee cup", + prompt_hash=generate_prompt_hash("a woman holding a coffee cup"), + ) + path = self.output_dir / "moss_00009_.png" + png(str(path), issue_graph(), pair_evidence()) + image_id = db.link_image_to_prompt(prompt, str(path), metadata={}) + out = io.StringIO() + tool.main( + [ + db.model.db_path, + "--output-dir", + str(self.output_dir), + "--ids", + str(image_id), + "--apply", + ], + out=out, + ) + self.assertIn("Applied: 0 relinked, 1 unlinked", out.getvalue()) + response = await self.client.get( + "/prompt_manager/images/output", params={"mode": "workspace"} + ) + self.assertEqual(response.status, 200) + images = {i["filename"]: i for i in (await response.json())["images"]} + self.assertIn("moss_00009_.png", images) + link = images["moss_00009_.png"].get("linked") or {} + self.assertNotEqual(link.get("status"), "linked") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_run_attribution.py b/tests/test_run_attribution.py index dc37a924..b59c6959 100644 --- a/tests/test_run_attribution.py +++ b/tests/test_run_attribution.py @@ -198,6 +198,27 @@ class TestScanUsesHistory(RecordedRuns, AdminAPITestCase): self.assertEqual(self.api._ingest_scanned_file(path, meta), "linked") self.assertEqual(self.linked_prompt(path), prompt_id) + def test_a_metadata_less_file_links_through_its_recorded_run(self): + """I1: --disable-metadata outputs carry no prompt; history still proves one.""" + db = self.api.db + prompt_id = db.save_prompt( + text=LLM_TEXT, prompt_hash=generate_prompt_hash(LLM_TEXT) + ) + path = self.write("bare_00001_.png") + self.record_run(db, str(self.output_dir), [("", path.name)], [LLM_TEXT]) + self.assertEqual(self.api._ingest_scanned_file(path, {}), "linked") + self.assertEqual(self.linked_prompt(path), prompt_id) + self.assertEqual(self.api._ingest_scanned_file(path, None), "unchanged") + detailed = self.api._ingest_scanned_file(path, {}, detailed=True) + self.assertEqual(detailed["status"], "unchanged") + self.assertTrue(detailed["title"]) + + def test_a_metadata_less_file_without_a_recorded_run_is_still_skipped(self): + path = self.write("bare_00002_.png") + self.assertIsNone(self.api._ingest_scanned_file(path, {})) + self.assertIsNone(self.api._ingest_scanned_file(path, {"workflow": {}})) + self.assertIsNone(self.linked_prompt(path)) + def test_the_issue_file_stays_unlinked_even_with_its_pair_record(self): """Resolved with an empty stored positive: before #342 the scan raised.""" db = self.api.db @@ -233,10 +254,6 @@ class TestScanUsesHistory(RecordedRuns, AdminAPITestCase): self.assertEqual(db.list_prompt_pairs()["pairs"], []) -if __name__ == "__main__": - unittest.main() - - class TestHostHelpers(unittest.TestCase): """The host lookups degrade to "no proof", never to an error.""" @@ -301,3 +318,176 @@ class TestHostHelpers(unittest.TestCase): sys.modules.pop("comfy_execution.utils", None) else: sys.modules["comfy_execution.utils"] = original + + +class TestMonitorHistoryRecheck(RecordedRuns, unittest.TestCase): + """I1: a keyless image gets one history-only re-check when its run is sealed.""" + + def setUp(self): + history_tests.HistoryDatabaseTests.setUp(self) + self.db = history_tests.HistoryDatabaseTests.database(self) + self.output = str(Path(self.temp.name) / "output") + os.makedirs(self.output) + self.handler = ImageGenerationHandler(self.db, make_tracker(self.db)) + self.handler.processing_delay = 0 + self.handler.settle_interval = 0.01 + self.addCleanup(self.handler.shutdown) + self.linked = [] + self.handler.link_image_to_prompt = lambda path, prompt, meta: ( + self.linked.append((os.path.basename(path), prompt["id"])) + ) + patches = ( + mock.patch( + "utils.image_monitor.host_output_directory", return_value=self.output + ), + mock.patch.object( + self.handler.metadata_extractor, "extract_metadata", return_value=None + ), + ) + for patch in patches: + patch.start() + self.addCleanup(patch.stop) + + close_databases = history_tests.HistoryDatabaseTests.close_databases + + def image(self, name): + path = os.path.join(self.output, name) + Path(path).write_bytes(b"png") + return path + + def save(self, text): + return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text)) + + def wait_for(self, predicate, timeout=5.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.01) + return predicate() + + def test_the_run_ending_after_processing_links_the_image_once(self): + prompt_id = self.save(LLM_TEXT) + path = self.image("late_00001_.png") + self.handler.process_new_image(path) # the run is still executing + self.assertEqual(self.linked, []) + self.assertIn(os.path.normcase(path), self.handler.awaiting_history()) + self.record_run(self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT]) + self.assertTrue(self.wait_for(lambda: self.linked)) + self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0)) + self.assertEqual(self.linked, [("late_00001_.png", prompt_id)]) + self.assertEqual(self.handler.awaiting_history(), []) + # Another run reporting the same name later is not a second re-check + self.record_run( + self.db, self.output, [("", "late_00001_.png")], [LLM_TEXT], item=1 + ) + self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0)) + self.assertEqual(len(self.linked), 1) + + def test_a_run_naming_nothing_stored_leaves_the_image_unlinked(self): + path = self.image("empty_00001_.png") + self.handler.process_new_image(path) + self.record_run(self.db, self.output, [("", "empty_00001_.png")], [""]) + self.assertTrue(self.wait_for(lambda: not self.handler.awaiting_history())) + self.assertTrue(self.wait_for(lambda: self.handler.pending_count() == 0)) + self.assertEqual(self.linked, []) + + def test_history_found_at_once_leaves_nothing_awaiting(self): + prompt_id = self.save(LLM_TEXT) + self.record_run(self.db, self.output, [("", "done_00001_.png")], [LLM_TEXT]) + self.handler.process_new_image(self.image("done_00001_.png")) + self.assertEqual(self.linked, [("done_00001_.png", prompt_id)]) + self.assertEqual(self.handler.awaiting_history(), []) + + def test_keyed_images_do_not_wait_for_history(self): + with mock.patch.object( + self.handler.metadata_extractor, + "extract_metadata", + return_value={"prompt": llm_text_graph()}, + ): + self.handler.process_new_image(self.image("keyed_00001_.png")) + self.assertEqual(self.handler.awaiting_history(), []) + + def test_the_waiting_set_is_bounded_and_expires(self): + from utils import image_monitor + + with mock.patch.object(image_monitor, "MAX_AWAITING_HISTORY", 2): + for n in range(3): + self.handler.process_new_image(self.image(f"b_{n}.png")) + self.assertEqual( + [os.path.basename(p) for p in self.handler.awaiting_history()], + ["b_1.png", "b_2.png"], + ) + later = time.monotonic() + image_monitor.AWAIT_HISTORY_SECONDS + 1 + with mock.patch.object(image_monitor.time, "monotonic", return_value=later): + self.assertEqual(self.handler.awaiting_history(), []) + + def test_shutdown_forgets_waiting_images_and_a_dead_handler_is_dropped(self): + import gc + + from database import run_attribution + + self.handler.process_new_image(self.image("gone_00001_.png")) + self.handler.shutdown() + self.assertEqual(self.handler.awaiting_history(), []) + other = ImageGenerationHandler(self.db, make_tracker(self.db)) + with mock.patch.object( + other.metadata_extractor, "extract_metadata", return_value=None + ): + other.process_new_image(self.image("gone_00002_.png")) + del other + gc.collect() + run_attribution.notify_outputs_sealed("run", []) # no dead callback raises + + +class TestOutputLookupIndex(unittest.TestCase): + """M5: the history lookup by file name uses an index, added by migration.""" + + setUp = history_tests.HistoryDatabaseTests.setUp + close_databases = history_tests.HistoryDatabaseTests.close_databases + database = history_tests.HistoryDatabaseTests.database + + def test_index_exists_is_used_validated_and_reinstalled(self): + from database.execution_history_schema import ( + history_schema_error, + migrate_execution_history, + ) + from database.run_attribution import _RUNS_FOR_OUTPUT + + db = self.database() + conn = db.model.get_connection() + name = "idx_history_read_execution_run_outputs" + columns = [r[2] for r in conn.execute(f"PRAGMA index_info({name})")] + self.assertEqual(columns, ["filename", "subfolder", "root_key", "path_style"]) + plan = " ".join( + str(row[3]) + for row in conn.execute( + "EXPLAIN QUERY PLAN " + _RUNS_FOR_OUTPUT, ("posix", "k", "", "f.png") + ) + ) + self.assertIn(name, plan) + self.assertIsNone(history_schema_error(conn, allow_missing=False)) + conn.execute(f"DROP INDEX {name}") + self.assertIsNone(history_schema_error(conn, allow_missing=False)) + migrate_execution_history(conn) + migrate_execution_history(conn) + self.assertTrue( + conn.execute("SELECT 1 FROM sqlite_master WHERE name=?", (name,)).fetchone() + ) + + def test_any_other_index_on_the_outputs_table_is_still_refused(self): + from database.execution_output_schema import output_schema_error + + db = self.database() + conn = db.model.get_connection() + conn.execute("CREATE INDEX idx_rogue ON execution_run_outputs(filename)") + self.assertIn("unexpected", output_schema_error(conn)) + conn.execute("DROP INDEX idx_rogue") + name = "idx_history_read_execution_run_outputs" + conn.execute(f"DROP INDEX {name}") + conn.execute(f"CREATE INDEX {name} ON execution_run_outputs(filename)") + self.assertIn("unexpected", output_schema_error(conn)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/review_image_links.py b/tools/review_image_links.py index c07d1b97..f0b38d31 100644 --- a/tools/review_image_links.py +++ b/tools/review_image_links.py @@ -21,8 +21,20 @@ connection) into a temporary folder and judged there, so the library is never written. ``--apply`` writes the ``relink`` and ``unlink`` verdicts of the selected images, after saving a backup beside the database (``.pre-relink_``). Only ids named on the command line are ever -changed. A batch item linked through the monitor's queue (PromptSearchList) -can't be proven after the fact and reads ``unlink``: review before applying. +changed, and only while each is still linked to the prompt it was judged on. + +What ``--apply`` does (APPLY_EFFECTS below, printed by every run): + +- ``relink`` points the image's link record at the other prompt and drops its + pair links (they named the old positive). +- ``unlink`` removes only the image's link record (its ``generated_images`` row + with the metadata snapshot stored in it, and its pair links). That is how the + app represents an unlinked image, exactly like a file the scan could not + attribute: the file stays on disk and in the gallery as unlinked, and its + favorite star and tags (kept by file, not by link) are untouched. + +A batch item linked through the monitor's queue (PromptSearchList) can't be +proven after the fact and reads ``unlink``: review before applying. python3 tools/review_image_links.py --output-dir \\ --ids 58005-58010 @@ -62,6 +74,15 @@ from utils.image_monitor import ImageGenerationHandler # noqa: E402 from utils.prompt_pair_resolver import resolve_prompt_pair # noqa: E402 VERDICTS = ("keep", "relink", "unlink", "conflict", "missing") +APPLY_EFFECTS = ( + "--apply, for ids named with --ids only, while each is still linked to the " + "prompt it was judged on:\n" + " relink: points the image's link record at the other prompt and drops " + "its pair links;\n" + " unlink: removes only the image's link record (its stored metadata " + "snapshot and pair links go with it); the file stays on disk and in the " + "gallery as unlinked, with its favorite and tags.\n" +) # Ids per query (well under SQLite's bound-parameter limit) ID_BATCH = 500 @@ -174,15 +195,30 @@ def review(db, ids, prompt_id, output_dir): def apply(db, verdicts, ids): - """Write the relink/unlink verdicts of explicitly named ids; returns counts.""" + """Write the relink/unlink verdicts of explicitly named ids; returns counts. + + An image whose link changed since it was judged (now another prompt, or + gone) is skipped. + """ allowed = set(ids) - relink = [ - v for v in verdicts if v["verdict"] == "relink" and v["image_id"] in allowed - ] - unlink = [ - v for v in verdicts if v["verdict"] == "unlink" and v["image_id"] in allowed + chosen = [ + v + for v in verdicts + if v["verdict"] in ("relink", "unlink") and v["image_id"] in allowed ] with db._pair_transaction() as conn: + current = [ + item + for item in chosen + if ( + row := conn.execute( + "SELECT prompt_id FROM generated_images WHERE id=?", + (item["image_id"],), + ).fetchone() + ) + and row[0] == item["linked_prompt"] + ] + relink = [item for item in current if item["verdict"] == "relink"] for item in relink: # A pair link records the old positive: it goes with the old link conn.execute( @@ -192,8 +228,16 @@ def apply(db, verdicts, ids): "UPDATE generated_images SET prompt_id=? WHERE id=? AND prompt_id=?", (item["target"], item["image_id"], item["linked_prompt"]), ) - removed = _delete_ids(conn, "generated_images", [v["image_id"] for v in unlink]) - return {"relinked": len(relink), "unlinked": removed} + removed = _delete_ids( + conn, + "generated_images", + [item["image_id"] for item in current if item["verdict"] == "unlink"], + ) + return { + "relinked": len(relink), + "unlinked": removed, + "skipped": len(chosen) - len(current), + } def backup_database(source, target): @@ -253,6 +297,7 @@ def main(argv=None, out=sys.stdout): finally: db.close_all() _report(verdicts, args.json, out) + out.write(APPLY_EFFECTS) out.write("Dry run: nothing was changed.\n") return 0 @@ -268,7 +313,8 @@ def main(argv=None, out=sys.stdout): finally: db.close_all() out.write( - f"Applied: {result['relinked']} relinked, {result['unlinked']} unlinked\n" + f"Applied: {result['relinked']} relinked, {result['unlinked']} unlinked, " + f"{result['skipped']} skipped (link changed since review)\n" ) return 0 diff --git a/utils/image_monitor.py b/utils/image_monitor.py index d1844df5..1b7aa573 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -28,7 +28,8 @@ import os import queue import time import threading -from typing import Callable, Optional, Dict, Any, Tuple +from collections import OrderedDict +from typing import Callable, Optional, Dict, Any, List, Tuple from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler @@ -42,16 +43,31 @@ from .usage_tracking import survive_switch try: from ..database.pair_ingestion import ingest_resolved_pair - from ..database.run_attribution import host_output_directory, stored_run_prompt + from ..database.run_attribution import ( + host_output_directory, + stored_run_prompt, + subscribe_outputs_sealed, + ) except ImportError: from database.pair_ingestion import ingest_resolved_pair - from database.run_attribution import host_output_directory, stored_run_prompt + from database.run_attribution import ( + host_output_directory, + stored_run_prompt, + subscribe_outputs_sealed, + ) # Upper bound on images waiting to be processed (a burst beyond this is skipped) MAX_PENDING_IMAGES = 1000 # A file counts as fully written once its size is unchanged across two reads SETTLE_INTERVAL_SECONDS = 0.1 SETTLE_ATTEMPTS = 20 +# Keyless images (no embedded API graph) left unlinked wait for their run's +# recorded outputs for one history-only re-check (#342): at most this many, +# each for at most this long. +MAX_AWAITING_HISTORY = 1000 +AWAIT_HISTORY_SECONDS = 3600 +# ``schedule`` snapshot meaning "history-only re-check" (not a new image) +HISTORY_ONLY = object() class ImageGenerationHandler(FileSystemEventHandler): @@ -113,6 +129,11 @@ class ImageGenerationHandler(FileSystemEventHandler): self._work: "queue.Queue[Optional[str]]" = queue.Queue() self._stop = threading.Event() self._worker: Optional[threading.Thread] = None + # Keyless unlinked images awaiting their run's recorded outputs: + # normalized path -> (path, monotonic deadline), oldest first. + self._awaiting: "OrderedDict[str, Tuple[str, float]]" = OrderedDict() + self._awaiting_lock = threading.Lock() + self._sealed_subscribed = False def on_created(self, event): """Handle filesystem creation events. @@ -180,6 +201,8 @@ class ImageGenerationHandler(FileSystemEventHandler): and exits; the next ``schedule`` starts a fresh worker with its own queue. Safe to call repeatedly and before anything was scheduled. """ + with self._awaiting_lock: + self._awaiting.clear() with self._pending_lock: self._pending.clear() worker, self._worker = self._worker, None @@ -307,6 +330,9 @@ class ImageGenerationHandler(FileSystemEventHandler): appeared (``PromptTracker.keyed_snapshot``); None reads it now """ try: + if prompt_snapshot is HISTORY_ONLY: + self._recheck_history(image_path) + return self.logger.info(f"Processing image: {image_path}") if not os.path.exists(image_path): @@ -399,6 +425,16 @@ class ImageGenerationHandler(FileSystemEventHandler): """ keys = image_run_keys(metadata) name = os.path.basename(image_path) + if not keys: + # No tracked run can be proven: history now, and once more after + # the run's outputs are recorded (``_on_outputs_sealed``) + self._await_history(image_path) + prompt_id = self._history_prompt(image_path) + if prompt_id is None: + return None + self._stop_awaiting(image_path) + self.logger.info(f"History match: prompt {prompt_id} for {name}") + return {"id": prompt_id, "from_history": True} if len(keys) == 1: (key,) = keys.values() if prompt_snapshot is None: @@ -414,10 +450,14 @@ class ImageGenerationHandler(FileSystemEventHandler): return None self.logger.info(f"Queue match: prompt {queued['id']} for {name}") return queued - if not context.get("placeholder") and run_id and count == 1: + if context.get("placeholder"): + # Same rule as a popped placeholder: that run had no prompt + self.logger.info(f"Latest run for {name} had no prompt text") + return None + if run_id and count == 1: self.logger.info(f"Run match: prompt {context['id']} for {name}") return context - elif len(keys) > 1: + else: self.logger.debug(f"{len(keys)} positive prompt nodes for {name}") prompt_id = self._history_prompt(image_path) if prompt_id is None: @@ -425,6 +465,89 @@ class ImageGenerationHandler(FileSystemEventHandler): self.logger.info(f"History match: prompt {prompt_id} for {name}") return {"id": prompt_id, "from_history": True} + @staticmethod + def _awaiting_key(image_path: str) -> str: + return os.path.normcase(os.path.abspath(image_path)) + + def _prune_awaiting(self) -> None: + """Drop expired and surplus entries (``_awaiting_lock`` held).""" + now = time.monotonic() + for key in [k for k, (_p, due) in self._awaiting.items() if due <= now]: + del self._awaiting[key] + while len(self._awaiting) > MAX_AWAITING_HISTORY: + self._awaiting.popitem(last=False) + + def _await_history(self, image_path: str) -> None: + """Re-check ``image_path`` once, history only, when its run is sealed.""" + if not self._sealed_subscribed: + subscribe_outputs_sealed(self._on_outputs_sealed) + self._sealed_subscribed = True + key = self._awaiting_key(image_path) + with self._awaiting_lock: + self._awaiting.pop(key, None) + self._awaiting[key] = ( + image_path, + time.monotonic() + AWAIT_HISTORY_SECONDS, + ) + self._prune_awaiting() + + def _stop_awaiting(self, image_path: str) -> None: + with self._awaiting_lock: + self._awaiting.pop(self._awaiting_key(image_path), None) + + def awaiting_history(self) -> List[str]: + """Normalized paths still waiting for their run's recorded outputs.""" + with self._awaiting_lock: + self._prune_awaiting() + return list(self._awaiting) + + def _on_outputs_sealed(self, run_id: str, references: List[dict]) -> None: + """A run's outputs were recorded: re-check its waiting images once. + + Called on the history writer's thread, so it only schedules the work. + """ + root = host_output_directory() + if not root: + return + due = [] + with self._awaiting_lock: + self._prune_awaiting() + for reference in references: + if reference.get("folder_type") != "output": + continue + key = self._awaiting_key( + os.path.join( + root, reference.get("subfolder", ""), reference["filename"] + ) + ) + waiting = self._awaiting.pop(key, None) + if waiting is not None: + due.append(waiting[0]) + for image_path in due: + self.schedule(image_path, HISTORY_ONLY) + + def _recheck_history(self, image_path: str) -> None: + """The one history-only re-check of a keyless image (#342).""" + name = os.path.basename(image_path) + if not os.path.exists(image_path): + return + prompt_id = self._history_prompt(image_path) + if prompt_id is None: + self.logger.info( + f"No certain prompt for {name} after its run ended; left unlinked" + ) + return + self.logger.info(f"History match after its run ended: {prompt_id} for {name}") + try: + metadata = self.metadata_extractor.extract_metadata(image_path) + except Exception: + metadata = None + self.link_image_to_prompt( + image_path, + {"id": prompt_id, "from_history": True}, + metadata or {"file_info": self.get_basic_file_info(image_path)}, + ) + def _history_prompt(self, image_path: str) -> Optional[int]: """The stored prompt recorded history attributes to this file, if one.""" try: diff --git a/utils/run_keys.py b/utils/run_keys.py index fda46a87..05d74a85 100644 --- a/utils/run_keys.py +++ b/utils/run_keys.py @@ -4,9 +4,14 @@ A PromptManager node registers the text it executed under a *run key*: a digest of the node and everything upstream of it in the queued API graph. The image embeds the same graph, so the image monitor can compute the same key for the image's positive prompt node. Equal keys mean the node ran with exactly the -inputs (and upstream nodes) that produced this image. That is also the identity -ComfyUI's own output cache uses: a node whose upstream is unchanged is served -from the cache, so a re-run of an unchanged upstream reproduces the same text. +inputs (and upstream nodes) that produced this image. + +Equal keys do not always mean equal text. ComfyUI's cache signature also folds +in each node's ``IS_CHANGED``, which the key cannot see: an upstream node that +forces a re-run (a random line, an "always run" LLM) can give two runs one key +and two texts. The key is therefore never used alone: the monitor also requires +the entry to come from the run the snapshot taken when the image appeared names +(``PromptTracker.pop_for_key``, ``ImageGenerationHandler._find_prompt_from_run``). Pure functions; the only host dependency (the executing run id) is optional. """ @@ -23,20 +28,27 @@ except ImportError: RUN_KEY_VERSION = 1 +def is_link(value: Any) -> bool: + """ComfyUI's own test (``comfy_execution.graph_utils.is_link``), exactly. + + A two-item list of a ``str`` node id and an ``int`` or ``float`` slot (a + ``bool`` slot passes too, as it does in ComfyUI). Tuples are values. + """ + if not isinstance(value, list): + return False + if len(value) != 2: + return False + if not isinstance(value[0], str): + return False + if not isinstance(value[1], int) and not isinstance(value[1], float): + return False + return True + + def _links(inputs: Any) -> List[str]: if not isinstance(inputs, dict): return [] - found = [] - for value in inputs.values(): - if ( - isinstance(value, (list, tuple)) - and len(value) == 2 - and isinstance(value[0], (str, int)) - and isinstance(value[1], int) - and not isinstance(value[1], bool) - ): - found.append(str(value[0])) - return found + return [value[0] for value in inputs.values() if is_link(value)] def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]: @@ -79,8 +91,9 @@ def prompt_node_key(graph: Any, node_id: Any) -> Optional[str]: def image_run_keys(metadata: Any) -> Dict[str, str]: """``{node_id: run key}`` for each positive PromptManager node of an image. - Empty without an embedded API graph: an image without one can never be - proven to belong to a tracked run. + Empty without an embedded API graph, and empty when any positive node has + no key: one unkeyable node must not leave another looking like the only + positive prompt (a two-pass workflow would link the wrong pass). """ graph = metadata.get("prompt") if isinstance(metadata, dict) else None if isinstance(graph, str): @@ -93,8 +106,9 @@ def image_run_keys(metadata: Any) -> Dict[str, str]: keys = {} for node_id in run_prompt_nodes(graph): key = prompt_node_key(graph, node_id) - if key is not None: - keys[node_id] = key + if key is None: + return {} + keys[node_id] = key return keys