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

fix(autotag): contain batch rows, require media extensions, bound the scan

Browse Source 
- The batch tagger runs every generated_images row through
  _resolve_allowed_image_path; rows that exist but lie outside the served
  output directories are skipped and counted in skipped_outside_roots on
  progress and complete events, and the engine receives the resolved path.
- _resolve_allowed_image_path requires a media extension (.png .jpg .jpeg
  .webp .gif .bmp .tiff .tif), so autotag/single returns 403 for other files.
- autotag/start holds a module-level lock; a concurrent start gets 409
  {"success": false, "error": "Auto-tagging is already running"}.
- scan_output_dir walks with a 50,000-file and 12-level cap, skips symlinked
  files and thumbnails, returns path through _public_path, relative_path in
  POSIX form, and a truncated flag.
- Autotag handlers report exceptions via the shared _public_error; the
  local copy is removed. Tests close the database before unlinking it.
This commit is contained in:
Vito Sansevero
2026-10-08 18:01:34 -07:00
parent d83383b9e0
commit f9da6488d6
2 changed files with 349 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/api/autotag_routes.py
+130 -81
View File
@@ -3,22 +3,33 @@
import asyncio
import json
import os
import threading
from pathlib import Path
from typing import Optional
from urllib.parse import quote
from aiohttp import web
_INTERNAL_ERROR = "An internal error occurred. Check server logs for details."
_BUSY_ERROR = "Auto-tagging is already running"
# Only these may be handed to the tagging engine or listed by the scanner.
_MEDIA_EXTENSIONS = frozenset(
{".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".tif"}
)
# Bounds for scan_output_dir: a runaway output tree must not pin the server.
_SCAN_MAX_FILES = 50_000
_SCAN_MAX_DEPTH = 12
# One batch run at a time; a second autotag/start gets 409 while this is held.
_autotag_batch_lock = threading.Lock()
def _public_error(exc: Exception) -> str:
"""Describe an exception without leaking absolute server paths."""
if isinstance(exc, OSError) and exc.strerror:
filename = exc.filename
if filename:
return f"{exc.strerror}: {os.path.basename(str(filename))}"
return exc.strerror
return str(exc)
def _file_exists(path) -> bool:
try:
return Path(str(path)).exists()
except (OSError, ValueError):
return False
def path_is_within(candidate, root) -> bool:
@@ -91,7 +102,9 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"Get autotag models error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def _stream_sse(self, request, events):
"""Write an async iterator of dict events as server-sent events.
@@ -250,6 +263,8 @@ class AutotagRoutesMixin:
"processed": 0,
"tagged": 0,
"skipped": 0,
"errors": 0,
"skipped_outside_roots": 0,
"status": "No images with linked prompts found",
}
service.unload_model()
@@ -261,7 +276,13 @@ class AutotagRoutesMixin:
"status": f"Found {total} images. Processing...",
}
counts = {"processed": 0, "tagged": 0, "skipped": 0, "errors": 0}
counts = {
"processed": 0,
"tagged": 0,
"skipped": 0,
"errors": 0,
"skipped_outside_roots": 0,
}
tagged_prompt_ids = set()
last_update = _time.monotonic()
finished = False
@@ -274,6 +295,7 @@ class AutotagRoutesMixin:
"processed": counts["processed"],
"tagged": counts["tagged"],
"skipped": counts["skipped"],
"skipped_outside_roots": counts["skipped_outside_roots"],
}
try:
@@ -281,17 +303,26 @@ class AutotagRoutesMixin:
image_path = image_data.get("image_path")
prompt_id = image_data.get("prompt_id")
skip_reason = None
allowed_path = None
if not image_path or not prompt_id:
skip_reason = "unlinked"
elif prompt_id in tagged_prompt_ids:
skip_reason = "prompt already processed"
elif not Path(image_path).exists():
skip_reason = "file missing"
elif opts["skip_tagged"] and self._real_tags(
image_data.get("prompt_tags", [])
):
tagged_prompt_ids.add(prompt_id)
skip_reason = "already tagged"
else:
# Rows can point anywhere (restored DB, hand edits): only
# media files inside the served output roots are tagged.
allowed_path = self._resolve_allowed_image_path(str(image_path))
if allowed_path is None:
if _file_exists(image_path):
skip_reason = "outside allowed directories"
counts["skipped_outside_roots"] += 1
else:
skip_reason = "file missing"
elif opts["skip_tagged"] and self._real_tags(
image_data.get("prompt_tags", [])
):
tagged_prompt_ids.add(prompt_id)
skip_reason = "already tagged"
if skip_reason is not None:
counts["skipped"] += 1
@@ -307,7 +338,7 @@ class AutotagRoutesMixin:
try:
tags = await self._run_in_executor(
service.generate_tags,
str(image_path),
allowed_path,
general_threshold=opts["general_threshold"],
character_threshold=opts["character_threshold"],
)
@@ -385,10 +416,19 @@ class AutotagRoutesMixin:
"message": _INTERNAL_ERROR,
}
return await self._stream_sse(request, events())
if not _autotag_batch_lock.acquire(blocking=False):
return web.json_response(
{"success": False, "error": _BUSY_ERROR}, status=409
)
try:
return await self._stream_sse(request, events())
finally:
_autotag_batch_lock.release()
def _resolve_allowed_image_path(self, raw_path: str) -> Optional[str]:
"""Return the real path when it is a file inside an output directory."""
"""Return the real path when it is a media file inside an output directory."""
if Path(str(raw_path)).suffix.lower() not in _MEDIA_EXTENSIONS:
return None
for output_dir in self._get_all_output_dirs():
if path_is_within(raw_path, output_dir):
real = os.path.realpath(raw_path)
@@ -522,7 +562,7 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"AutoTag single error: {e}")
return web.json_response(
{"success": False, "error": _public_error(e)}, status=500
{"success": False, "error": self._public_error(e)}, status=500
)
async def apply_autotag(self, request):
@@ -568,7 +608,9 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"Apply autotag error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
async def unload_autotag_model(self, request):
"""Manually unload the AutoTag model from memory."""
@@ -595,7 +637,37 @@ class AutotagRoutesMixin:
except Exception as e:
self.logger.error(f"Unload autotag model error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
@staticmethod
def _collect_scan_files(output_path: Path):
"""Media files under ``output_path`` from a bounded, symlink-free walk.
Descends at most ``_SCAN_MAX_DEPTH`` levels, skips ``thumbnails``
trees and symlinked files, and stops after ``_SCAN_MAX_FILES`` files.
Returns ``(paths, truncated)``.
"""
found = []
base_depth = len(output_path.parts)
for dirpath, dirnames, filenames in os.walk(output_path):
depth = len(Path(dirpath).parts) - base_depth
dirnames[:] = sorted(
name
for name in dirnames
if name != "thumbnails" and depth + 1 <= _SCAN_MAX_DEPTH
)
for name in sorted(filenames):
if Path(name).suffix.lower() not in _MEDIA_EXTENSIONS:
continue
candidate = Path(dirpath) / name
if candidate.is_symlink():
continue
if len(found) >= _SCAN_MAX_FILES:
return found, True
found.append(candidate)
return found, False
async def scan_output_dir(self, request):
"""Scan ComfyUI output directory for images."""
@@ -608,73 +680,50 @@ class AutotagRoutesMixin:
)
output_path = Path(output_dir)
image_extensions = [
".png",
".jpg",
".jpeg",
".gif",
".webp",
".bmp",
".tiff",
]
files, truncated = await self._run_in_executor(
self._collect_scan_files, output_path
)
has_thumbnails = (output_path / "thumbnails").is_dir()
images = []
seen_paths = set()
for ext in image_extensions:
for pattern in [f"*{ext.lower()}", f"*{ext.upper()}"]:
for image_path in output_path.rglob(pattern):
if "thumbnails" not in image_path.parts:
normalized_path = str(image_path).lower()
if normalized_path not in seen_paths:
seen_paths.add(normalized_path)
rel_path = image_path.relative_to(output_path)
thumbnail_url = None
thumbnails_dir = output_path / "thumbnails"
if thumbnails_dir.exists():
rel_path_no_ext = rel_path.with_suffix("")
thumbnail_rel_path = (
"thumbnails/"
f"{rel_path_no_ext.as_posix()}"
f"_thumb{image_path.suffix}"
)
thumbnail_abs_path = thumbnails_dir / (
f"{rel_path_no_ext.as_posix()}"
f"_thumb{image_path.suffix}"
)
if thumbnail_abs_path.exists():
from urllib.parse import quote
thumbnail_url = (
"/prompt_manager/images/serve/"
f'{quote(thumbnail_rel_path, safe="/")}'
)
from urllib.parse import quote as url_quote
images.append(
{
"filename": image_path.name,
"path": str(image_path),
"relative_path": str(rel_path),
"url": (
"/prompt_manager/images/serve/"
+ url_quote(rel_path.as_posix(), safe="/")
),
"thumbnail_url": thumbnail_url,
}
)
for image_path in files:
rel_path = image_path.relative_to(output_path)
thumbnail_url = None
if has_thumbnails:
thumb_rel = (
f"thumbnails/{rel_path.with_suffix('').as_posix()}"
f"_thumb{image_path.suffix}"
)
if (output_path / thumb_rel).exists():
thumbnail_url = "/prompt_manager/images/serve/" + quote(
thumb_rel, safe="/"
)
images.append(
{
"filename": image_path.name,
"path": self._public_path(image_path),
"relative_path": rel_path.as_posix(),
"url": "/prompt_manager/images/serve/"
+ quote(rel_path.as_posix(), safe="/"),
"thumbnail_url": thumbnail_url,
}
)
images.sort(key=lambda x: x["filename"])
self.logger.info(f"Found {len(images)} images in output directory")
return web.json_response(
{"success": True, "images": images, "count": len(images)}
{
"success": True,
"images": images,
"count": len(images),
"truncated": truncated,
}
)
except Exception as e:
self.logger.error(f"Scan output dir error: {e}")
return web.json_response({"success": False, "error": str(e)}, status=500)
return web.json_response(
{"success": False, "error": self._public_error(e)}, status=500
)
tests/test_api_autotag.py
+219 -4
View File
@@ -4,11 +4,13 @@ The ML engine is replaced by a fake service so no model is ever loaded and
no network request is made.
"""
import asyncio
import json
import os
import shutil
import sys
import tempfile
import threading
import types
import unittest
from pathlib import Path
@@ -121,6 +123,7 @@ class AutotagAPITestCase(AioHTTPTestCase):
async def tearDownAsync(self):
self._service_patch.stop()
self.api.db.close_all()
for path in (
self._temp_db.name,
self._temp_db.name + "-wal",
@@ -287,6 +290,41 @@ class TestAutotagSingleContainment(AutotagAPITestCase):
self.assertNotIn(str(self.secret_dir), data["error"])
self.assertEqual(self.service.generate_calls, [])
async def test_non_media_extension_inside_output_dir_is_forbidden(self):
notes = self.output_dir / "notes.txt"
notes.write_text("not an image")
resp = await self._post_json(
"/prompt_manager/autotag/single", {"path": str(notes)}
)
self.assertEqual(resp.status, 403)
self.assertFalse((await resp.json())["success"])
self.assertEqual(self.service.generate_calls, [])
async def test_non_media_extension_via_image_id_is_forbidden(self):
prompt_id = self._save_prompt()
notes = self.output_dir / "notes.txt"
notes.write_text("not an image")
image_id = self._link_image(prompt_id, notes)
resp = await self._post_json(
"/prompt_manager/autotag/single", {"image_id": image_id}
)
self.assertEqual(resp.status, 403)
self.assertEqual(self.service.generate_calls, [])
async def test_uppercase_media_extension_is_accepted(self):
image = self._make_image(self.output_dir, "shout.JPEG")
resp = await self._post_json(
"/prompt_manager/autotag/single", {"path": str(image)}
)
self.assertEqual(resp.status, 200)
self.assertEqual(len(self.service.generate_calls), 1)
async def test_missing_selector_is_400(self):
resp = await self._post_json("/prompt_manager/autotag/single", {})
self.assertEqual(resp.status, 400)
@@ -424,6 +462,86 @@ class TestAutotagBatch(AutotagAPITestCase):
self.assertTrue(all(0 <= e["progress"] <= 100 for e in progress))
self.assertTrue(any("processed" in e for e in progress))
async def test_rows_outside_output_dirs_are_skipped_and_counted(self):
inside = self._save_prompt("inside")
outside = self._save_prompt("outside")
self._link_image(inside, self._make_image(self.output_dir, "in.png"))
secret = self._make_image(self.secret_dir, "password.png")
self._link_image(outside, secret)
_, events = await self._start()
done = events[-1]
self.assertEqual(done["type"], "complete")
self.assertEqual(done["processed"], 1)
self.assertEqual(done["tagged"], 1)
self.assertEqual(done["skipped"], 1)
self.assertEqual(done["skipped_outside_roots"], 1)
self.assertEqual(len(self.service.generate_calls), 1)
self.assertEqual(
os.path.realpath(self.service.generate_calls[0][0]),
os.path.realpath(str(self.output_dir / "in.png")),
)
self.assertEqual(self.api.db.get_prompt_by_id(outside)["tags"], [])
progress = [e for e in events if e["type"] == "progress" and "skipped" in e]
self.assertTrue(progress)
self.assertTrue(all("skipped_outside_roots" in e for e in progress))
body = json.dumps(events)
self.assertNotIn(str(self.secret_dir), body)
async def test_second_start_while_running_is_409(self):
self._seed()
started = threading.Event()
release = threading.Event()
def slow_generate(image_path, **kwargs):
started.set()
release.wait(5)
return ["slow"]
self.service.generate_tags = slow_generate
first = asyncio.ensure_future(self._start())
try:
await asyncio.get_event_loop().run_in_executor(None, started.wait, 5)
self.assertTrue(started.is_set())
resp = await self.client.request(
"POST", "/prompt_manager/autotag/start?model_type=gguf"
)
self.assertEqual(resp.status, 409)
self.assertEqual(
await resp.json(),
{"success": False, "error": "Auto-tagging is already running"},
)
finally:
release.set()
resp1, events = await first
self.assertEqual(resp1.status, 200)
self.assertEqual(events[-1]["type"], "complete")
async def test_lock_is_released_after_a_run(self):
self._seed()
resp1, _ = await self._start()
resp2, events = await self._start()
self.assertEqual(resp1.status, 200)
self.assertEqual(resp2.status, 200)
self.assertEqual(events[-1]["type"], "complete")
async def test_lock_is_released_when_the_run_fails(self):
self._seed()
self.service.load_error = RuntimeError("no model")
_, events = await self._start()
self.assertEqual(events[-1]["type"], "error")
self.service.load_error = None
resp, events = await self._start()
self.assertEqual(resp.status, 200)
self.assertEqual(events[-1]["type"], "complete")
async def test_already_tagged_prompts_are_skipped_unless_asked(self):
p = self._save_prompt("tagged", tags=["cat"])
self._link_image(p, self._make_image(self.output_dir, "t.png"))
@@ -603,10 +721,11 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase):
self.assertEqual((await resp.json())["error"], "model crashed")
async def test_oserror_without_filename_uses_strerror(self):
from py.api.autotag_routes import _public_error
from py.api import _public_error
self.assertEqual(_public_error(OSError(5, "I/O error")), "I/O error")
self.assertEqual(_public_error(ValueError("plain")), "plain")
self.assertIs(self.api._public_error, _public_error)
async def test_path_with_nul_byte_is_forbidden(self):
from py.api.autotag_routes import path_is_within
@@ -759,6 +878,45 @@ class TestApplyAutotag(AutotagAPITestCase):
self.assertEqual(resp.status, 500)
class TestAutotagErrorBodiesHideServerPaths(AutotagAPITestCase):
"""Every autotag handler reports exceptions through the shared public helper."""
def _locked(self):
return PermissionError(
13, "Permission denied", str(self.secret_dir / "locked.db")
)
async def _assert_hidden(self, resp):
self.assertEqual(resp.status, 500)
body = await resp.text()
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.secret_dir), body)
self.assertIn("locked.db", body)
async def test_models_error(self):
def boom():
raise self._locked()
self.service.get_models_status = boom
resp = await self.client.request("GET", "/prompt_manager/autotag/models")
await self._assert_hidden(resp)
async def test_apply_error(self):
with patch.object(self.api.db, "get_prompt_by_id", side_effect=self._locked()):
resp = await self._post_json(
"/prompt_manager/autotag/apply", {"prompt_id": 1, "tags": ["a"]}
)
await self._assert_hidden(resp)
async def test_unload_error(self):
def boom():
raise self._locked()
self.service.is_model_loaded = boom
resp = await self.client.request("POST", "/prompt_manager/autotag/unload")
await self._assert_hidden(resp)
class TestUnloadAutotag(AutotagAPITestCase):
async def test_nothing_loaded(self):
@@ -799,10 +957,12 @@ class TestScanOutputDir(AutotagAPITestCase):
resp = await self.client.request("GET", "/prompt_manager/scan_output_dir")
data = await resp.json()
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertEqual([i["filename"] for i in data["images"]], ["a.png", "b.JPG"])
self.assertEqual(data["count"], 2)
self.assertIs(data["truncated"], False)
self.assertIsNone(data["images"][0]["thumbnail_url"])
self.assertEqual(
data["images"][1]["thumbnail_url"],
@@ -811,19 +971,74 @@ class TestScanOutputDir(AutotagAPITestCase):
self.assertEqual(
data["images"][1]["url"], "/prompt_manager/images/serve/sub/b.JPG"
)
self.assertEqual(data["images"][1]["relative_path"], "sub/b.JPG")
self.assertEqual(
data["images"][1]["path"], self.api._public_path(sub / "b.JPG")
)
self.assertFalse(os.path.isabs(data["images"][1]["path"]))
self.assertNotIn(str(self.tmp_root), body)
async def test_symlinked_files_are_skipped(self):
secret = self._make_image(self.secret_dir, "password.png")
self._make_image(self.output_dir, "real.png")
link = self.output_dir / "link.png"
try:
link.symlink_to(secret)
except (OSError, NotImplementedError):
self.skipTest("symlinks not supported on this platform")
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
data = await (
await self.client.request("GET", "/prompt_manager/scan_output_dir")
).json()
self.assertEqual([i["filename"] for i in data["images"]], ["real.png"])
async def test_file_cap_marks_the_listing_truncated(self):
for name in ("a.png", "b.png", "c.png"):
self._make_image(self.output_dir, name)
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
with patch("py.api.autotag_routes._SCAN_MAX_FILES", 2):
data = await (
await self.client.request("GET", "/prompt_manager/scan_output_dir")
).json()
self.assertEqual(data["count"], 2)
self.assertIs(data["truncated"], True)
async def test_depth_cap_stops_the_walk(self):
deep = self.output_dir / "d1" / "d2" / "d3"
deep.mkdir(parents=True)
self._make_image(self.output_dir / "d1", "shallow.png")
self._make_image(deep, "deep.png")
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
with patch("py.api.autotag_routes._SCAN_MAX_DEPTH", 1):
data = await (
await self.client.request("GET", "/prompt_manager/scan_output_dir")
).json()
self.assertEqual([i["filename"] for i in data["images"]], ["shallow.png"])
async def test_missing_output_dir_is_404(self):
self.api._find_comfyui_output_dir = lambda: None
resp = await self.client.request("GET", "/prompt_manager/scan_output_dir")
self.assertEqual(resp.status, 404)
async def test_error_is_500(self):
async def test_error_is_500_without_the_server_path(self):
def boom():
raise RuntimeError("x")
raise PermissionError(
13, "Permission denied", str(self.secret_dir / "locked.png")
)
self.api._find_comfyui_output_dir = boom
resp = await self.client.request("GET", "/prompt_manager/scan_output_dir")
self.assertEqual(resp.status, 500)
body = await resp.text()
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.secret_dir), body)
self.assertIn("locked.png", body)
if __name__ == "__main__":