feat(tracking): enrich live image pairs and cancel scans safely (#193)

database/pair_ingestion.py

@@ -51,14 +51,21 @@ def ingest_resolved_pair(
             )
             if any(row["negative_hash"] != expected for row in prior):
                 return {"status": "conflict", "pairs_added": 0, "links_added": 0}
-        before = conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0]
+        prior_pair = conn.execute(
+            "SELECT pp.id FROM prompt_pairs pp "
+            "JOIN prompts p ON p.id=pp.positive_prompt_id "
+            "LEFT JOIN prompts n ON n.id=pp.negative_prompt_id "
+            "WHERE p.hash=? AND n.hash IS ?",
+            (
+                generate_prompt_hash(positive_text),
+                generate_prompt_hash(negative_text) if negative_text.strip() else None,
+            ),
+        ).fetchone()
         pair_id = db._save_pair(conn, positive_text, negative_text)
         pair = conn.execute(
             "SELECT * FROM prompt_pairs WHERE id=?", (pair_id,)
         ).fetchone()
-        pairs_added = (
-            conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0] - before
-        )
+        pairs_added = int(prior_pair is None)
         file_size, width, height, image_format = _coerce_file_info(metadata)
         snapshots = {
             column: json.dumps(metadata[key], sort_keys=True, allow_nan=False)

py/api/admin.py

@@ -377,6 +377,10 @@ class AdminRoutesMixin:
         async def scan_images_route(request):
             return await self.scan_images(request)
 
+        @routes.post("/prompt_manager/scan/cancel")
+        async def cancel_scan_route(request):
+            return await self.cancel_scan(request)
+
         @routes.get("/prompt_manager/scan/status")
         async def scan_status_route(request):
             return await self.scan_status(request)
@@ -1298,6 +1302,14 @@ class AdminRoutesMixin:
             job = self._start_scan_job()
         return await self._stream_scan_job(request, job)
 
+    async def cancel_scan(self, request):
+        """Stop between file transactions; an in-flight extraction can finish."""
+        job = self._scan_job
+        requested = bool(job and job.running)
+        if requested:
+            job.cancelled = True
+        return web.json_response({"success": True, "cancellation_requested": requested})
+
     async def scan_status(self, request):
         """Whether an output scan is running, and its latest progress event."""
         job = self._scan_job
@@ -1382,7 +1394,9 @@ class AdminRoutesMixin:
                 f"found={counts['found']}, new_prompts_added={counts['added']}, "
                 f"images_linked_to_existing={counts['linked']}"
             )
-            yield {"type": "complete", **counts, "directories": dir_names}
+            job = self._scan_job
+            final_type = "cancelled" if job and job.cancelled else "complete"
+            yield {"type": final_type, **counts, "directories": dir_names}
 
         except Exception:
             self.logger.exception("Scan error")
@@ -1398,6 +1412,8 @@ class AdminRoutesMixin:
         """
         total = len(media_files)
         for batch_start in range(0, total, SCAN_BATCH_SIZE):
+            if self._scan_job and self._scan_job.cancelled:
+                return
             batch = media_files[batch_start : batch_start + SCAN_BATCH_SIZE]
             batch_results = await self._run_in_executor(
                 self._extract_batch_metadata_sync, batch
@@ -1451,6 +1467,8 @@ class AdminRoutesMixin:
         """Store the prompts found in a metadata batch (blocking DB work)."""
         counts = dict.fromkeys(SCAN_COUNT_KEYS, 0)
         for media_file, metadata in batch_results:
+            if self._scan_job and self._scan_job.cancelled:
+                break
             counts["processed"] += 1
             try:
                 outcome = self._ingest_scanned_file(media_file, metadata, detailed=True)

py/api/scan_job.py

@@ -16,6 +16,7 @@ class ScanJob:
     def __init__(self):
         self.last_event = None
         self.finished = False
+        self.cancelled = False
         self.task = None
         self._subscribers = set()
 

tests/test_pair_image_linking.py

@@ -0,0 +1,55 @@
+"""Live metadata annotation uses the same conservative ingestion as scan."""
+
+import logging
+import tempfile
+import unittest
+from pathlib import Path
+
+from database.operations import PromptDatabase
+from utils.hashing import generate_prompt_hash
+from utils.image_monitor import ImageGenerationHandler
+
+
+def metadata(positive="cat", negative="blur"):
+    return {
+        "prompt": {
+            "1": {"class_type": "PromptManager", "inputs": {"text": positive}},
+            "2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
+            "3": {
+                "class_type": "KSampler",
+                "inputs": {"positive": ["1", 0], "negative": ["2", 0]},
+            },
+        }
+    }
+
+
+class TestPairImageLinking(unittest.TestCase):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.db = PromptDatabase(str(Path(self.temp.name) / "db.sqlite"))
+        self.path = str(Path(self.temp.name) / "image.png")
+        self.pid = self.db.save_prompt("cat", prompt_hash=generate_prompt_hash("cat"))
+        self.handler = ImageGenerationHandler.__new__(ImageGenerationHandler)
+        self.handler.db_manager = self.db
+        self.handler.logger = logging.getLogger("test.pairlink")
+
+    def tearDown(self):
+        self.db.close_all()
+        self.temp.cleanup()
+
+    def test_exact_metadata_creates_pair_link_without_counting_a_run(self):
+        self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata())
+        self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata())
+        pairs = self.db.list_prompt_pairs()["pairs"]
+        self.assertEqual(len(pairs), 1)
+        self.assertEqual(pairs[0]["image_count"], 1)
+        self.assertEqual(pairs[0]["run_count"], 0)
+
+    def test_stale_context_or_ambiguous_metadata_does_not_create_pair(self):
+        meta = metadata(positive="dog")
+        self.handler.link_image_to_prompt(self.path, {"id": self.pid}, meta)
+        self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
+        meta = metadata()
+        meta["prompt"]["4"] = meta["prompt"]["3"]
+        self.handler.link_image_to_prompt(self.path + "2", {"id": self.pid}, meta)
+        self.assertEqual(self.db.list_prompt_pairs()["total"], 0)

tests/test_pair_scan.py

@@ -55,3 +55,29 @@ class TestPairScan(AdminAPITestCase):
         path = self.output_dir / "old.png"
         self.assertEqual(self.api._ingest_scanned_file(path, meta), "added")
         self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged")
+
+    async def test_cancel_finishes_at_safe_boundary_and_restart_is_idempotent(self):
+        from py.api.scan_job import ScanJob
+
+        job = ScanJob()
+        self.api._scan_job = job
+        response = await self.client.post("/prompt_manager/scan/cancel")
+        self.assertEqual(response.status, 200)
+        self.assertTrue((await response.json())["cancellation_requested"])
+        counts = self.api._ingest_scan_batch_sync(
+            [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
+        )
+        self.assertEqual(counts["processed"], 0)
+        self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0)
+        job.finish()
+        self.api._scan_job = ScanJob()
+        first = self.api._ingest_scan_batch_sync(
+            [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
+        )
+        self.assertEqual(first["pairs_added"], 1)
+        self.assertEqual(
+            self.api._ingest_scan_batch_sync(
+                [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
+            )["unchanged"],
+            1,
+        )

utils/image_monitor.py

@@ -35,6 +35,12 @@ from .hashing import generate_prompt_hash
 from .metadata_extractor import ComfyUIMetadataExtractor
 from .logging_config import get_logger
 from .prompt_graph import resolve_text, run_prompt_nodes
+from .prompt_pair_resolver import resolve_prompt_pair
+
+try:
+    from ..database.pair_ingestion import ingest_resolved_pair
+except ImportError:
+    from database.pair_ingestion import ingest_resolved_pair
 
 # Upper bound on images waiting to be processed (a burst beyond this is skipped)
 MAX_PENDING_IMAGES = 1000
@@ -529,6 +535,17 @@ class ImageGenerationHandler(FileSystemEventHandler):
                 etc.)
         """
         try:
+            resolution = resolve_prompt_pair(metadata)
+            if resolution["state"] == "resolved":
+                positive = resolution["positive"]
+                expected = self.db_manager.get_prompt_by_hash(
+                    generate_prompt_hash(positive.get("storage_text", positive["text"]))
+                )
+                if expected and expected["id"] == prompt_context["id"]:
+                    ingest_resolved_pair(
+                        self.db_manager, image_path, metadata, resolution
+                    )
+                    return
             image_id = self.db_manager.link_image_to_prompt(
                 prompt_id=prompt_context["id"], image_path=image_path, metadata=metadata
             )