diff --git a/TESTING.md b/TESTING.md index cc4e5f96..69a85954 100644 --- a/TESTING.md +++ b/TESTING.md @@ -3873,9 +3873,12 @@ fallback (`_get_fallback_prompt` is gone) and no unproven tracker entry. - **Run keys** (`utils/run_keys.py`): sha256 of a prompt node plus its whole upstream closure (`class_type` and `inputs`) in the API graph. Nodes register tracker entries under it; the monitor computes it from the image's embedded - graph. A queued entry is used only by key (batch items in order, a placeholder - leaves its image unlinked, a new run of a key drops the old run's entries). The - snapshot taken when the file appears is used only when its run id is known + graph. The snapshot taken when the file appears names the key's latest run; a + queued entry is taken only by key and from that run (batch items in order, a + placeholder leaves its image unlinked, same-key entries of earlier runs ahead + of it are stale and dropped, entries of later runs that executed during the + monitor's delay wait for their own images). Without a queued entry the + snapshot's entry itself is used only when its run id is known (`comfy_execution.utils.get_executing_context`) and that run registered one entry for the key, which covers a cached re-run (sampler seed changed) without applying a batch's last item to every image. More than one positive diff --git a/tests/test_prompt_attribution.py b/tests/test_prompt_attribution.py index 66b67bca..cb3e9e79 100644 --- a/tests/test_prompt_attribution.py +++ b/tests/test_prompt_attribution.py @@ -232,9 +232,10 @@ class TestEmptyEditorSavesNothing(NodeCase): ) latest, _run, _count = self.tracker.keyed_snapshot()[key] self.assertTrue(latest.get("placeholder")) - popped = self.tracker.pop_for_keys([key]) + 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_keys([key]).get("placeholder")) + 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") @@ -391,11 +392,39 @@ class TestProvenRuns(MonitorCase): self.assertIsNone(self.process({"prompt": graph})) self.assertEqual(self.process({"prompt": graph}), third) - def test_a_new_run_supersedes_queued_entries_of_the_same_key(self): + 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()) diff --git a/tests/vendor/comfyui-config.test.mjs b/tests/vendor/comfyui-config.test.mjs index eed33dd2..f25e4a2e 100644 --- a/tests/vendor/comfyui-config.test.mjs +++ b/tests/vendor/comfyui-config.test.mjs @@ -27,15 +27,18 @@ test('execution script discovers only the new matrix and preserves all 78 origin assert.doesNotMatch(stock, /execution-history/); assert.doesNotMatch(stock, /execution-outputs/); assert.doesNotMatch(stock, /library-location/); // it starts its own fresh host (#249 S6b) + assert.doesNotMatch(stock, /image-attribution/); // needs the owned fixture pack (#342) const executions = list('tools/vendor/playwright.execution.config.mjs'); assert.match(executions, /\[classic-start\]/); assert.match(executions, /\[nodes2-start\]/); - assert.match(executions, /Total: 48 tests in 3 files/); // + the fresh-host library location case + // + the fresh-host library location case, + the #342 image attribution case + assert.match(executions, /Total: 50 tests in 4 files/); for (const project of ['classic-start', 'nodes2-start']) { const cases = executions.split('\n').filter(line => line.includes(`[${project}]`)); assert.equal(cases.filter(line => line.includes('execution-history.spec.mjs')).length, 15); assert.equal(cases.filter(line => line.includes('execution-outputs.spec.mjs')).length, 8); assert.equal(cases.filter(line => line.includes('library-location.spec.mjs')).length, 1); + assert.equal(cases.filter(line => line.includes('image-attribution.spec.mjs')).length, 1); } assert.doesNotMatch(executions, /prompt-manager.spec/); }); diff --git a/tools/vendor/comfyui.config.mjs b/tools/vendor/comfyui.config.mjs index bc1e5bfc..99d41bb9 100644 --- a/tools/vendor/comfyui.config.mjs +++ b/tools/vendor/comfyui.config.mjs @@ -19,7 +19,8 @@ const servers = [{ export default defineConfig({ testDir: '../../tests/comfyui', - testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs'], + testIgnore: ['execution-history.spec.mjs', 'execution-outputs.spec.mjs', 'library-location.spec.mjs', + 'image-attribution.spec.mjs'], outputDir: './test-results/comfyui', timeout: 60000, workers: 1, diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 377951c6..d1844df5 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -384,12 +384,16 @@ class ImageGenerationHandler(FileSystemEventHandler): node counts (``utils.run_keys``): never the newest prompt in the database and never an entry of another graph (#342). - 1. Batch queue: the oldest queued entry of that key. Batch items (text - from PromptSearchList) share one key and are linked in order. An - empty item's placeholder leaves its image unlinked. - 2. Snapshot taken when the image appeared: the key's latest entry, when - its run registered exactly one entry for the key (a batch's latest - item says nothing about the other images) and the run is known. + The snapshot taken when the image appeared names the key's latest run. + + 1. Batch queue: the oldest queued entry of that key from that run. + Batch items (text from PromptSearchList) share one key and are + linked in order; an empty item's placeholder leaves its image + unlinked; entries queued by later runs stay for their own images. + 2. The snapshot's entry itself, when its run is known and registered + exactly one entry for the key (a batch's latest item says nothing + about the other images). This covers a cached node that did not + run again. 3. Recorded execution history: the run that wrote this file names one stored prompt (``database.run_attribution``). """ @@ -397,18 +401,19 @@ class ImageGenerationHandler(FileSystemEventHandler): name = os.path.basename(image_path) if len(keys) == 1: (key,) = keys.values() - queued = self.prompt_tracker.pop_for_keys([key]) - if queued is not None: - if queued.get("placeholder"): - self.logger.info(f"Queue entry for {name} has no prompt text") - return None - self.logger.info(f"Queue match: prompt {queued['id']} for {name}") - return queued if prompt_snapshot is None: prompt_snapshot = self._keyed_snapshot() + # No entry of the key when the image appeared: no tracked run made it entry = (prompt_snapshot or {}).get(key) if entry is not None: context, run_id, count = entry + queued = self.prompt_tracker.pop_for_key(key, run_id) + if queued is not None: + if queued.get("placeholder"): + self.logger.info(f"Queue entry for {name} has no prompt text") + 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: self.logger.info(f"Run match: prompt {context['id']} for {name}") return context diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index 82c75412..73d71db3 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -29,7 +29,7 @@ import threading import time import uuid from collections import OrderedDict -from typing import Any, Dict, Iterable, Optional, Tuple +from typing import Any, Dict, Optional, Tuple # Import logging system try: @@ -180,7 +180,6 @@ class PromptTracker: # Push to batch queue for ordered image linking if push_to_queue: with self._queue_lock: - self._supersede_queued(run_key, run_id) self._prompt_queue.append(execution_context) self.logger.debug( f"Queue size: {len(self._prompt_queue)} after push " @@ -265,7 +264,6 @@ class PromptTracker: self._remember_keyed(placeholder) if push_to_queue: with self._queue_lock: - self._supersede_queued(run_key, run_id) self._prompt_queue.append(placeholder) def _keyed_runs(self) -> "KeyedRuns": @@ -281,21 +279,6 @@ class PromptTracker: if context.get("run_key"): self._keyed_runs().record(context) - def _supersede_queued(self, run_key: Optional[str], run_id: Optional[str]): - """Drop queued entries of ``run_key`` from an earlier run. - - A node that executes again with the same key in a new run replaces what - the earlier run queued: images of the new run must not take them. - Called with ``_queue_lock`` held. - """ - if not run_key or run_id is None: - return - self._prompt_queue = [ - ctx - for ctx in self._prompt_queue - if ctx.get("run_key") != run_key or ctx.get("run_id") in (None, run_id) - ] - def keyed_snapshot(self) -> Dict[str, Tuple[Dict[str, Any], Optional[str], int]]: """``{run key: (latest entry, its run id, entries of that run)}`` now. @@ -304,26 +287,46 @@ class PromptTracker: """ return self._keyed_runs().snapshot(self.prompt_timeout) - def pop_for_keys(self, keys: Iterable[str]) -> Optional[Dict[str, Any]]: - """Pop the oldest queued entry whose run key is one of ``keys`` (FIFO). + def pop_for_key( + self, run_key: Optional[str], run_id: Optional[str] = None + ) -> Optional[Dict[str, Any]]: + """Pop the oldest queued entry of ``run_key`` (FIFO), from ``run_id`` if set. - Entries of other keys stay queued; expired entries are discarded. + ``run_id`` is the run the key's latest entry came from when the image + appeared (``keyed_snapshot``): entries of that key queued by earlier + runs ahead of it are stale and dropped, while entries of later runs + (queued after the image appeared) stay for their own images. Entries + of other keys stay queued; expired entries are discarded. """ - wanted = {key for key in keys if key} - if not wanted: + if not run_key: return None now = time.time() with self._queue_lock: - kept = [] - found = None - for ctx in self._prompt_queue: - if now - ctx.get("timestamp", 0) >= self.prompt_timeout: - continue - if found is None and ctx.get("run_key") in wanted: - found = ctx - continue - kept.append(ctx) - self._prompt_queue = kept + live = [ + ctx + for ctx in self._prompt_queue + if now - ctx.get("timestamp", 0) < self.prompt_timeout + ] + index = next( + ( + i + for i, ctx in enumerate(live) + if ctx.get("run_key") == run_key + and (run_id is None or ctx.get("run_id") == run_id) + ), + None, + ) + if index is None: + self._prompt_queue = live + return None + found = live[index] + # Same-key entries ahead of it came from earlier runs: stale + stale_or_found = { + i for i in range(index + 1) if live[i].get("run_key") == run_key + } + self._prompt_queue = [ + ctx for i, ctx in enumerate(live) if i not in stale_or_found + ] return found def pop_next_prompt(self) -> Optional[Dict[str, Any]]: