Vito Sansevero
66f61de0db
The workspace Images listing now admits a file under a Tags filter when every requested tag is carried by its linked positive or its current image tag set, and q also matches image tag names (case-folded substring). Image tags count only while a byte record is verified at the file's present (size, mtime_ns) (spec amendment S1): the tagged records are streamed once, only listed files carrying a filter tag or q hit are stat'ed, and their linked positives are read through the file_path index so tags, category, rating and the positive count combine. The rooted inspector envelope gains a validated image_tags facet (current / none / unavailable) read at the inspection's own stable stat. A failed image-tag read fails a Tags filter like a failed link read and degrades a q search to no image tags; plain listings read none. Legacy gallery and search responses are unchanged.
1032 lines
44 KiB
Python
1032 lines
44 KiB
Python
"""Image annotation routes: observe, capture preview and tag save.
|
|
|
|
Every request names its file through the listing identity
|
|
``{root_index, root_scope, path}``; the scope is checked before any file
|
|
access, tag reads and saves are byte-exact (the capture's digest, re-hashed
|
|
at Save) and no response ever carries a server path or a locator.
|
|
"""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
import uuid
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
# Discovery imports this module before tests/test_api_autotag.py, whose
|
|
# ``folder_paths`` stub (a real module, so ``patch(..., create=True)`` can add
|
|
# and remove attributes) the later modules rely on. Install the same stub
|
|
# before test_api_images would install a MagicMock in its place.
|
|
if "folder_paths" not in sys.modules:
|
|
_stub = types.ModuleType("folder_paths")
|
|
_stub.models_dir = tempfile.gettempdir()
|
|
_stub.base_path = tempfile.gettempdir()
|
|
_stub.get_output_directory = tempfile.gettempdir
|
|
sys.modules["folder_paths"] = _stub
|
|
|
|
try:
|
|
from test_api_images import ImageAPITestCase, make_png
|
|
from open_handles import assert_closed, open_handles_under
|
|
except ImportError:
|
|
from tests.test_api_images import ImageAPITestCase, make_png
|
|
from tests.open_handles import assert_closed, open_handles_under
|
|
|
|
from database.execution_history_schema import rotate_history_epoch
|
|
from utils import media_capture
|
|
from utils.media_capture import CaptureStore
|
|
|
|
try:
|
|
from py.api import annotations as annotations_module
|
|
except ImportError: # the module under test is not there yet (RED)
|
|
annotations_module = None
|
|
|
|
OBSERVE = "/prompt_manager/image-annotations/observe"
|
|
TAGS = "/prompt_manager/image-annotations/tags"
|
|
CAPTURE = "/prompt_manager/image-annotations/capture/"
|
|
MAX_BODY = 64 * 1024
|
|
|
|
# Columns that an observe may legitimately refresh (S1) or that never matter
|
|
# for "nothing changed" comparisons.
|
|
_SNAPSHOT_TABLES = {
|
|
"prompts": "*",
|
|
"generated_images": "*",
|
|
"tags": "*",
|
|
"prompt_tags": "*",
|
|
"prompt_pair_tags": "*",
|
|
"image_favorites": "*",
|
|
"media_annotation_files": "id, locator",
|
|
"media_annotation_bytes": "id, file_id, sha256, first_seen_at, tag_revision",
|
|
"media_image_tags": "*",
|
|
}
|
|
|
|
|
|
class FakeClock:
|
|
def __init__(self, start=1000.0):
|
|
self.now = start
|
|
|
|
def __call__(self):
|
|
return self.now
|
|
|
|
|
|
class AnnotationAPITestCase(ImageAPITestCase):
|
|
"""The image API app plus helpers for the annotation routes."""
|
|
|
|
async def get_application(self):
|
|
app = await super().get_application()
|
|
self.assertIsNotNone(
|
|
annotations_module, "py.api.annotations is not implemented"
|
|
)
|
|
self.epoch = self.api.db.get_execution_epoch()
|
|
return app
|
|
|
|
async def tearDownAsync(self):
|
|
store = getattr(self.api, "_captures", None)
|
|
if store is not None:
|
|
# Every preview handle must already be closed (Windows unlink).
|
|
assert_closed(self, store.directory)
|
|
self.api.close_annotation_captures()
|
|
if store is not None:
|
|
self.assertFalse(os.path.exists(store.directory))
|
|
await super().tearDownAsync()
|
|
|
|
# -- helpers ---------------------------------------------------------
|
|
|
|
def scope(self, index=0, roots=None):
|
|
roots = roots if roots is not None else self.api._get_all_output_dirs()
|
|
return self.api._gallery_root_scope(roots[index], index)
|
|
|
|
def target(self, path="sub/a.png", **changes):
|
|
return {"root_index": 0, "root_scope": self.scope(), "path": path, **changes}
|
|
|
|
async def json_of(self, response):
|
|
"""The JSON envelope; an unrouted path (no JSON) fails as an assertion."""
|
|
self.assertEqual(
|
|
response.content_type,
|
|
"application/json",
|
|
f"{response.method} {response.url.path} answered {response.status}",
|
|
)
|
|
return await response.json()
|
|
|
|
async def observe(self, body):
|
|
response = await self.client.post(OBSERVE, json=body)
|
|
return response.status, await self.json_of(response)
|
|
|
|
async def observe_ok(self, path="sub/a.png"):
|
|
status, body = await self.observe({"target": self.target(path)})
|
|
self.assertEqual(status, 200, body)
|
|
return body
|
|
|
|
async def put_tags(self, body):
|
|
response = await self.client.put(TAGS, json=body)
|
|
return response.status, await self.json_of(response)
|
|
|
|
def save_body(self, observed, tags, **changes):
|
|
return {
|
|
"target": dict(observed["target"]),
|
|
"capture_token": observed["capture"]["token"],
|
|
"sha256": observed["capture"]["sha256"],
|
|
"expected_epoch": observed["library_epoch"],
|
|
"expected_tag_revision": observed["tag_revision"],
|
|
"tags": list(tags),
|
|
**changes,
|
|
}
|
|
|
|
def locator(self, path):
|
|
return os.path.normcase(os.path.realpath(str(self.output_dir / path)))
|
|
|
|
def write_bytes(self, path, data, *, keep_mtime=False):
|
|
"""Rewrite *path* in place, optionally restoring its ``mtime_ns``."""
|
|
full = self.output_dir / path
|
|
before = full.stat()
|
|
full.write_bytes(data)
|
|
if keep_mtime:
|
|
os.utime(full, ns=(before.st_atime_ns, before.st_mtime_ns))
|
|
return full
|
|
|
|
def db_snapshot(self):
|
|
conn = sqlite3.connect(self._temp_db.name)
|
|
try:
|
|
snapshot = {}
|
|
for table, columns in _SNAPSHOT_TABLES.items():
|
|
exists = conn.execute(
|
|
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
|
|
(table,),
|
|
).fetchone()
|
|
if exists is None:
|
|
continue
|
|
rows = conn.execute(f"SELECT {columns} FROM {table}").fetchall()
|
|
snapshot[table] = sorted(map(repr, rows))
|
|
return snapshot
|
|
finally:
|
|
conn.close()
|
|
|
|
def bytes_rows(self):
|
|
conn = sqlite3.connect(self._temp_db.name)
|
|
try:
|
|
return conn.execute(
|
|
"SELECT sha256, size, mtime_ns, tag_revision FROM "
|
|
"media_annotation_bytes ORDER BY first_seen_at"
|
|
).fetchall()
|
|
finally:
|
|
conn.close()
|
|
|
|
def saved_set(self, path, data, tags):
|
|
"""Store *tags* for the exact *data* bytes of *path*; return the receipt."""
|
|
full = self.output_dir / path
|
|
stat = full.stat()
|
|
return self.api.db.replace_image_tags(
|
|
locator=self.locator(path),
|
|
sha256=hashlib.sha256(data).hexdigest(),
|
|
size=stat.st_size,
|
|
mtime_ns=stat.st_mtime_ns,
|
|
expected_epoch=self.epoch,
|
|
expected_tag_revision=0,
|
|
tags=tags,
|
|
)
|
|
|
|
def use_fake_clock(self):
|
|
clock = FakeClock()
|
|
self.api.close_annotation_captures()
|
|
self.api._captures = CaptureStore(clock=clock)
|
|
return clock
|
|
|
|
def assert_no_paths(self, body):
|
|
text = json.dumps(body)
|
|
self.assertNotIn(str(self.comfy_dir), text)
|
|
self.assertNotIn(str(self.output_dir), text)
|
|
self.assertNotIn(os.path.realpath(self.output_dir), text)
|
|
store = getattr(self.api, "_captures", None)
|
|
if store is not None:
|
|
self.assertNotIn(store.directory, text)
|
|
self.assertNotIn("\\\\", text.replace('\\"', ""))
|
|
self.assertNotIn(os.sep + "output" + os.sep, text)
|
|
|
|
|
|
class TestObserve(AnnotationAPITestCase):
|
|
async def test_observe_requires_valid_scope_before_any_file_access(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
other = self.comfy_dir / "other"
|
|
other.mkdir()
|
|
make_png(other / "sub" / "a.png")
|
|
stale_scope = self.api._gallery_root_scope(other, 0)
|
|
reordered = self.api._gallery_root_scope(self.output_dir, 1)
|
|
with patch.object(
|
|
media_capture.CaptureStore, "capture", autospec=True
|
|
) as capture:
|
|
capture.side_effect = AssertionError("capture must not run")
|
|
for name, changes, status in (
|
|
("stale root", {"root_scope": stale_scope}, 409),
|
|
("reordered root", {"root_scope": reordered}, 409),
|
|
("removed root", {"root_index": 1, "root_scope": reordered}, 409),
|
|
("wrong scope", {"root_scope": "0" * 64}, 409),
|
|
("uppercase scope", {"root_scope": "A" * 64}, 400),
|
|
("short scope", {"root_scope": "abc"}, 400),
|
|
("bool index", {"root_index": True}, 400),
|
|
("negative index", {"root_index": -1}, 400),
|
|
("string index", {"root_index": "0"}, 400),
|
|
("traversal", {"path": "../a.png"}, 403),
|
|
("dot segment", {"path": "./a.png"}, 403),
|
|
("empty segment", {"path": "sub//a.png"}, 403),
|
|
("absolute", {"path": "/etc/a.png"}, 403),
|
|
("backslash", {"path": "sub\\a.png"}, 403),
|
|
("drive", {"path": "C:/a.png"}, 403),
|
|
("drive segment", {"path": "sub/C:a.png"}, 403),
|
|
("nul", {"path": "a\x00.png"}, 403),
|
|
("non-media", {"path": "notes.txt"}, 403),
|
|
("empty path", {"path": ""}, 400),
|
|
("long path", {"path": "a/" * 600 + "a.png"}, 400),
|
|
("int path", {"path": 7}, 400),
|
|
):
|
|
with self.subTest(name):
|
|
status_, body = await self.observe(
|
|
{"target": self.target(**changes)}
|
|
)
|
|
self.assertEqual(status_, status, body)
|
|
self.assertFalse(body["success"])
|
|
if status == 409:
|
|
self.assertEqual(body["code"], "source-changed")
|
|
self.assert_no_paths(body)
|
|
for name, body in (
|
|
("missing target", {}),
|
|
("extra field", {"target": self.target(), "extra": 1}),
|
|
("target not an object", {"target": "sub/a.png"}),
|
|
("target missing path", {"target": {"root_index": 0}}),
|
|
("target extra field", {"target": self.target(sha256="x")}),
|
|
):
|
|
with self.subTest(name):
|
|
status_, body_ = await self.observe(body)
|
|
self.assertEqual(status_, 400, body_)
|
|
self.assertFalse(body_["success"])
|
|
with self.subTest("body over 64 KiB"):
|
|
padding = "x" * MAX_BODY
|
|
response = await self.client.post(
|
|
OBSERVE,
|
|
data=json.dumps({"target": self.target(), "pad": padding}),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
self.assertEqual(response.status, 413)
|
|
with self.subTest("not json"):
|
|
response = await self.client.post(
|
|
OBSERVE, data=b"{oops", headers={"Content-Type": "application/json"}
|
|
)
|
|
self.assertEqual(response.status, 400)
|
|
capture.assert_not_called()
|
|
|
|
async def test_observe_returns_saved_tags_for_exact_bytes_only(self):
|
|
path = make_png(self.output_dir / "sub" / "a.png")
|
|
original = path.read_bytes()
|
|
receipt = self.saved_set("sub/a.png", original, ["Forest", "mist"])
|
|
self.assertEqual(receipt["tag_revision"], 1)
|
|
|
|
body = await self.observe_ok()
|
|
self.assertEqual(body["tags"], ["Forest", "mist"])
|
|
self.assertEqual(body["tag_revision"], 1)
|
|
self.assertEqual(body["library_epoch"], self.epoch)
|
|
self.assertEqual(body["filename"], "a.png")
|
|
self.assertIsNone(body["title"])
|
|
self.assertEqual(
|
|
body["target"],
|
|
{"root_index": 0, "root_scope": self.scope(), "path": "sub/a.png"},
|
|
)
|
|
capture = body["capture"]
|
|
self.assertEqual(
|
|
sorted(capture),
|
|
sorted(
|
|
[
|
|
"token",
|
|
"expires_at",
|
|
"preview_url",
|
|
"sha256",
|
|
"size",
|
|
"width",
|
|
"height",
|
|
"media_type",
|
|
"generate",
|
|
]
|
|
),
|
|
)
|
|
self.assertEqual(capture["sha256"], hashlib.sha256(original).hexdigest())
|
|
self.assertEqual(capture["size"], len(original))
|
|
self.assertEqual((capture["width"], capture["height"]), (4, 4))
|
|
self.assertEqual(capture["media_type"], "image")
|
|
self.assertEqual(capture["generate"], {"supported": True, "reason": None})
|
|
self.assertEqual(capture["preview_url"], CAPTURE + capture["token"])
|
|
self.assertTrue(capture["expires_at"].endswith("Z"))
|
|
|
|
# Different bytes at the same path: an empty set at revision 0.
|
|
changed = make_png(path, color=(0, 255, 0)).read_bytes()
|
|
self.assertNotEqual(changed, original)
|
|
body = await self.observe_ok()
|
|
self.assertEqual((body["tags"], body["tag_revision"]), ([], 0))
|
|
self.assertEqual(body["capture"]["sha256"], hashlib.sha256(changed).hexdigest())
|
|
|
|
# The original bytes again (new mtime): the saved set, and S1 refreshes
|
|
# the record's verified fingerprint so cheap reads see it as current.
|
|
self.write_bytes("sub/a.png", original)
|
|
stat = path.stat()
|
|
current = self.api.db.current_image_tags(
|
|
{self.locator("sub/a.png"): (stat.st_size, stat.st_mtime_ns)}
|
|
)
|
|
body = await self.observe_ok()
|
|
self.assertEqual((body["tags"], body["tag_revision"]), (["Forest", "mist"], 1))
|
|
refreshed = self.api.db.current_image_tags(
|
|
{self.locator("sub/a.png"): (stat.st_size, stat.st_mtime_ns)}
|
|
)
|
|
if stat.st_mtime_ns != self.bytes_rows()[0][2]:
|
|
self.assertEqual(current, {})
|
|
self.assertEqual(refreshed, {self.locator("sub/a.png"): ["Forest", "mist"]})
|
|
|
|
async def test_observe_changes_no_tags_revisions_prompts_or_links(self):
|
|
path = make_png(self.output_dir / "sub" / "a.png")
|
|
data = path.read_bytes()
|
|
self.saved_set("sub/a.png", data, ["forest"])
|
|
make_png(self.output_dir / "sub" / "b.png", color=(1, 2, 3))
|
|
self._link_image(self.output_dir / "sub" / "a.png", "a prompt")
|
|
prompt_id = self._save_prompt("tagged prompt", tags=["forest", "night"])
|
|
self.assertTrue(prompt_id)
|
|
before = self.db_snapshot()
|
|
|
|
for _ in range(2):
|
|
await self.observe_ok("sub/a.png")
|
|
await self.observe_ok("sub/b.png")
|
|
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
# Unseen bytes (b.png) never create a record.
|
|
self.assertEqual(len(self.bytes_rows()), 1)
|
|
self.assertEqual(self.bytes_rows()[0][3], 1)
|
|
|
|
async def test_observe_title_follows_the_card_rule(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
prompt_id = self._save_prompt("Same words\nrest")
|
|
self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png"))
|
|
body = await self.observe_ok()
|
|
# A first-line title is withheld unless the file's own records confirm
|
|
# it (the card rule); this PNG carries no prompt metadata.
|
|
self.assertIsNone(body["title"])
|
|
with self.api.db.model.get_connection() as connection:
|
|
connection.execute(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
(prompt_id, "Same words"),
|
|
)
|
|
# A stored title is user data and always shows.
|
|
body = await self.observe_ok()
|
|
self.assertEqual(body["title"], "Same words")
|
|
|
|
async def test_observe_video_is_manual_only(self):
|
|
clip = self.output_dir / "clip.mp4"
|
|
clip.parent.mkdir(exist_ok=True)
|
|
clip.write_bytes(b"\x00\x00\x00\x18ftypmp42" + b"\x00" * 64)
|
|
body = await self.observe_ok("clip.mp4")
|
|
capture = body["capture"]
|
|
self.assertEqual(capture["media_type"], "video")
|
|
self.assertEqual((capture["width"], capture["height"]), (None, None))
|
|
self.assertEqual(
|
|
capture["generate"],
|
|
{"supported": False, "reason": media_capture.VIDEO_ONLY_REASON},
|
|
)
|
|
self.assertIsNone(body["title"])
|
|
response = await self.client.get(capture["preview_url"])
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.headers["Content-Type"], "video/mp4")
|
|
self.assertEqual(await response.read(), clip.read_bytes())
|
|
|
|
async def test_observe_maps_capture_outcomes_to_codes(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
outside = make_png(self.comfy_dir / "outside.png")
|
|
self._make_symlink(outside, self.output_dir / "sub" / "alias.png")
|
|
status, body = await self.observe({"target": self.target("sub/alias.png")})
|
|
self.assertEqual((status, body["code"]), (409, "source-changed"))
|
|
self.assert_no_paths(body)
|
|
|
|
status, body = await self.observe({"target": self.target("sub/missing.png")})
|
|
self.assertEqual((status, body["code"]), (409, "file-changed"))
|
|
self.assert_no_paths(body)
|
|
|
|
self.api.close_annotation_captures()
|
|
self.api._captures = CaptureStore(max_file=16)
|
|
status, body = await self.observe({"target": self.target("sub/a.png")})
|
|
self.assertEqual((status, body["code"]), (415, "unsupported-media"))
|
|
self.assert_no_paths(body)
|
|
|
|
self.api.close_annotation_captures()
|
|
self.api._captures = CaptureStore(max_captures=1)
|
|
await self.observe_ok("sub/a.png")
|
|
status, body = await self.observe({"target": self.target("sub/a.png")})
|
|
self.assertEqual((status, body["code"]), (429, "review-capacity"))
|
|
self.assert_no_paths(body)
|
|
|
|
|
|
class TestCaptureRoutes(AnnotationAPITestCase):
|
|
async def test_capture_route_serves_exact_bytes_no_store_nosniff(self):
|
|
path = make_png(self.output_dir / "sub" / "a.png")
|
|
source_bytes = path.read_bytes()
|
|
body = await self.observe_ok()
|
|
token = body["capture"]["token"]
|
|
# The source may change after the capture; the preview never does.
|
|
make_png(path, color=(9, 9, 9))
|
|
|
|
resp = await self.client.get(body["capture"]["preview_url"])
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["Cache-Control"], "no-store")
|
|
self.assertEqual(resp.headers["X-Content-Type-Options"], "nosniff")
|
|
self.assertEqual(resp.headers["Content-Type"], "image/png")
|
|
self.assertEqual(resp.headers["Content-Length"], str(len(source_bytes)))
|
|
self.assertEqual(await resp.read(), source_bytes)
|
|
|
|
store = self.api._captures
|
|
entry = store._entries[token]
|
|
self.assertEqual(entry.uses, 0)
|
|
self.assertEqual(open_handles_under(store.directory), [])
|
|
self.assertIsNotNone(store.get(token))
|
|
|
|
async def test_capture_route_rejects_unknown_expired_and_rescoped_tokens(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
clock = self.use_fake_clock()
|
|
for token in ("nope", "A" * 300, "x.y", "%2e%2e%2fx", "a b"):
|
|
with self.subTest(token=token):
|
|
resp = await self.client.get(CAPTURE + token)
|
|
self.assertEqual(resp.status, 404)
|
|
data = await self.json_of(resp)
|
|
self.assertEqual(data["code"], "capture-expired")
|
|
self.assertEqual(resp.headers["Cache-Control"], "no-store")
|
|
|
|
body = await self.observe_ok()
|
|
url = body["capture"]["preview_url"]
|
|
self.assertEqual((await self.client.get(url)).status, 200)
|
|
|
|
# Released explicitly: gone, and releasing again stays successful.
|
|
resp = await self.client.delete(url)
|
|
self.assertEqual(
|
|
(resp.status, await self.json_of(resp)), (200, {"success": True})
|
|
)
|
|
self.assertEqual((await self.client.get(url)).status, 404)
|
|
resp = await self.client.delete(url)
|
|
self.assertEqual(
|
|
(resp.status, await self.json_of(resp)), (200, {"success": True})
|
|
)
|
|
|
|
# Idle expiry.
|
|
body = await self.observe_ok()
|
|
url = body["capture"]["preview_url"]
|
|
clock.now += media_capture.DEFAULT_IDLE_SECONDS + 1
|
|
resp = await self.client.get(url)
|
|
self.assertEqual(resp.status, 404)
|
|
self.assertEqual((await self.json_of(resp))["code"], "capture-expired")
|
|
|
|
# The roots changed under a live token: denied and released.
|
|
body = await self.observe_ok()
|
|
url = body["capture"]["preview_url"]
|
|
moved = self.comfy_dir / "moved"
|
|
moved.mkdir()
|
|
original_roots = self.api._get_all_output_dirs
|
|
self.api._get_all_output_dirs = lambda: [moved]
|
|
try:
|
|
resp = await self.client.get(url)
|
|
self.assertEqual(resp.status, 409)
|
|
data = await self.json_of(resp)
|
|
self.assertEqual(data["code"], "source-changed")
|
|
self.assert_no_paths(data)
|
|
finally:
|
|
self.api._get_all_output_dirs = original_roots
|
|
resp = await self.client.get(url)
|
|
self.assertEqual(resp.status, 404)
|
|
self.assertFalse(any(Path(self.api._captures.directory).iterdir()))
|
|
|
|
async def test_preview_pins_the_capture_while_it_streams(self):
|
|
path = make_png(self.output_dir / "sub" / "a.png", size=(64, 64))
|
|
source_bytes = path.read_bytes()
|
|
body = await self.observe_ok()
|
|
token = body["capture"]["token"]
|
|
store = self.api._captures
|
|
capture_path = store.get(token).path
|
|
stream = annotations_module.AnnotationRoutesMixin._stream_capture
|
|
seen = {}
|
|
|
|
async def release_mid_stream(request, capture, handle):
|
|
# A concurrent DELETE (or sweep) lands while the bytes stream.
|
|
self.assertEqual(store._entries[token].uses, 1)
|
|
self.assertTrue(store.release(token))
|
|
seen["kept"] = os.path.exists(capture_path)
|
|
return await stream(request, capture, handle)
|
|
|
|
with (
|
|
patch.object(
|
|
annotations_module.AnnotationRoutesMixin,
|
|
"_stream_capture",
|
|
staticmethod(release_mid_stream),
|
|
),
|
|
patch.object(annotations_module, "_PREVIEW_CHUNK", 16),
|
|
):
|
|
resp = await self.client.get(body["capture"]["preview_url"])
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(await resp.read(), source_bytes)
|
|
self.assertTrue(seen["kept"])
|
|
# Removed with the last use, after its handle closed.
|
|
self.assertFalse(os.path.exists(capture_path))
|
|
self.assertEqual(
|
|
(await self.client.get(body["capture"]["preview_url"])).status, 404
|
|
)
|
|
|
|
async def test_preview_of_a_vanished_capture_file_is_expired(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
body = await self.observe_ok()
|
|
os.unlink(self.api._captures.get(body["capture"]["token"]).path)
|
|
resp = await self.client.get(body["capture"]["preview_url"])
|
|
self.assertEqual(resp.status, 404)
|
|
data = await self.json_of(resp)
|
|
self.assertEqual(data["code"], "capture-expired")
|
|
self.assert_no_paths(data)
|
|
self.assertIsNone(self.api._captures.get(body["capture"]["token"]))
|
|
|
|
async def test_restore_invalidates_captures(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
body = await self.observe_ok()
|
|
url = body["capture"]["preview_url"]
|
|
self.assertEqual((await self.client.get(url)).status, 200)
|
|
backup = str(self.comfy_dir / "backup.db")
|
|
self.assertTrue(self.api.db.model.backup_database(backup))
|
|
safety = self.api.db.model.restore_from_file(backup)
|
|
try:
|
|
resp = await self.client.get(url)
|
|
self.assertEqual(resp.status, 404)
|
|
self.assertEqual((await self.json_of(resp))["code"], "capture-expired")
|
|
status, saved = await self.put_tags(self.save_body(body, ["forest"]))
|
|
self.assertEqual((status, saved["code"]), (409, "capture-expired"))
|
|
self.assertEqual(self.bytes_rows(), [])
|
|
finally:
|
|
for path in (safety, backup):
|
|
for suffix in ("", "-wal", "-shm"):
|
|
if path and os.path.exists(path + suffix):
|
|
os.unlink(path + suffix)
|
|
|
|
|
|
class TestSaveTags(AnnotationAPITestCase):
|
|
async def test_put_saves_then_rereads_the_same_set(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
observed = await self.observe_ok()
|
|
status, receipt = await self.put_tags(
|
|
self.save_body(observed, ["Forest", "mist", "forest"])
|
|
)
|
|
self.assertEqual(status, 200, receipt)
|
|
self.assertEqual(
|
|
receipt,
|
|
{
|
|
"success": True,
|
|
"library_epoch": self.epoch,
|
|
"sha256": observed["capture"]["sha256"],
|
|
"tags": ["Forest", "mist"],
|
|
"tag_revision": 1,
|
|
"replayed": False,
|
|
},
|
|
)
|
|
body = await self.observe_ok()
|
|
self.assertEqual((body["tags"], body["tag_revision"]), (["Forest", "mist"], 1))
|
|
|
|
async def test_put_rehashes_current_file_and_rejects_changed_bytes(self):
|
|
path = make_png(self.output_dir / "sub" / "a.png")
|
|
original = path.read_bytes()
|
|
observed = await self.observe_ok()
|
|
before = self.db_snapshot()
|
|
# Same size, same mtime, other bytes: only the digest tells.
|
|
replacement = bytearray(original)
|
|
replacement[-1] ^= 0xFF
|
|
self.write_bytes("sub/a.png", bytes(replacement), keep_mtime=True)
|
|
self.assertEqual(path.stat().st_size, len(original))
|
|
|
|
status, body = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
self.assertEqual((status, body["code"]), (409, "file-changed"))
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
self.assert_no_paths(body)
|
|
|
|
# Reopening shows the new bytes as an unseen, empty set.
|
|
reopened = await self.observe_ok()
|
|
self.assertEqual((reopened["tags"], reopened["tag_revision"]), ([], 0))
|
|
self.assertNotEqual(
|
|
reopened["capture"]["sha256"], observed["capture"]["sha256"]
|
|
)
|
|
|
|
# A vanished file is also file-changed with zero writes.
|
|
path.unlink()
|
|
status, body = await self.put_tags(self.save_body(reopened, ["forest"]))
|
|
self.assertEqual((status, body["code"]), (409, "file-changed"))
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
|
|
async def test_put_rejects_capture_for_other_target_or_epoch(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
make_png(self.output_dir / "sub" / "b.png", color=(0, 0, 255))
|
|
observed_a = await self.observe_ok("sub/a.png")
|
|
observed_b = await self.observe_ok("sub/b.png")
|
|
before = self.db_snapshot()
|
|
|
|
# b's token for a's target: refused before any whole-file hash.
|
|
body = self.save_body(observed_a, ["forest"])
|
|
body["capture_token"] = observed_b["capture"]["token"]
|
|
body["sha256"] = observed_b["capture"]["sha256"]
|
|
with patch.object(
|
|
media_capture.CaptureStore, "verify", autospec=True
|
|
) as verify:
|
|
status, result = await self.put_tags(body)
|
|
verify.assert_not_called()
|
|
self.assertEqual((status, result["code"]), (409, "capture-expired"))
|
|
|
|
# A digest that is not the capture's.
|
|
body = self.save_body(observed_a, ["forest"], sha256="0" * 64)
|
|
status, result = await self.put_tags(body)
|
|
self.assertEqual(status, 400)
|
|
|
|
# An unknown token.
|
|
body = self.save_body(observed_a, ["forest"], capture_token="unknown-token")
|
|
status, result = await self.put_tags(body)
|
|
self.assertEqual((status, result["code"]), (409, "capture-expired"))
|
|
|
|
# Another epoch than the capture's.
|
|
body = self.save_body(observed_a, ["forest"], expected_epoch=str(uuid.uuid4()))
|
|
status, result = await self.put_tags(body)
|
|
self.assertEqual((status, result["code"]), (409, "storage-changed"))
|
|
self.assertEqual(result["library_epoch"], self.epoch)
|
|
|
|
# A target the token was not captured for (scope of another root).
|
|
other = self.comfy_dir / "other"
|
|
other.mkdir()
|
|
make_png(other / "sub" / "a.png")
|
|
roots = [self.output_dir, other]
|
|
self.api._get_all_output_dirs = lambda: roots
|
|
body = self.save_body(observed_a, ["forest"])
|
|
body["target"] = {
|
|
"root_index": 1,
|
|
"root_scope": self.scope(1, roots),
|
|
"path": "sub/a.png",
|
|
}
|
|
status, result = await self.put_tags(body)
|
|
self.assertEqual((status, result["code"]), (409, "capture-expired"))
|
|
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
for result_ in (result,):
|
|
self.assert_no_paths(result_)
|
|
|
|
async def test_put_conflict_storage_changed_and_limit_codes(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
observed = await self.observe_ok()
|
|
status, first = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
self.assertEqual(status, 200, first)
|
|
|
|
# A stale revision with a different set: the current set comes back.
|
|
status, body = await self.put_tags(self.save_body(observed, ["woods"]))
|
|
self.assertEqual((status, body["code"]), (409, "tag-conflict"))
|
|
self.assertEqual((body["tags"], body["tag_revision"]), (["forest"], 1))
|
|
self.assertEqual(body["library_epoch"], self.epoch)
|
|
|
|
# Validation at the boundary, before any file or database work.
|
|
with patch.object(
|
|
media_capture.CaptureStore, "verify", autospec=True
|
|
) as verify:
|
|
verify.side_effect = AssertionError("verify must not run")
|
|
cases = (
|
|
("too many tags", {"tags": [f"t{i}" for i in range(101)]}, 400),
|
|
("tag not a string", {"tags": ["ok", 3]}, 400),
|
|
("tags not a list", {"tags": "forest"}, 400),
|
|
("blank tag", {"tags": [" "]}, 400),
|
|
("bool revision", {"expected_tag_revision": True}, 400),
|
|
("negative revision", {"expected_tag_revision": -1}, 400),
|
|
("bad epoch", {"expected_epoch": "not-a-uuid"}, 400),
|
|
("bad digest", {"sha256": "xyz"}, 400),
|
|
("bad token type", {"capture_token": 5}, 400),
|
|
("traversal", {"target": self.target("../a.png")}, 403),
|
|
("wrong scope", {"target": self.target(root_scope="f" * 64)}, 409),
|
|
)
|
|
for name, changes, expected in cases:
|
|
with self.subTest(name):
|
|
body_ = {**self.save_body(observed, ["forest"]), **changes}
|
|
status, result = await self.put_tags(body_)
|
|
self.assertEqual(status, expected, result)
|
|
self.assertFalse(result["success"])
|
|
if name == "too many tags":
|
|
self.assertEqual(result["code"], "tag-limit")
|
|
with self.subTest("extra field"):
|
|
body_ = self.save_body(observed, ["forest"])
|
|
body_["extra"] = 1
|
|
status, result = await self.put_tags(body_)
|
|
self.assertEqual(status, 400)
|
|
with self.subTest("missing field"):
|
|
body_ = self.save_body(observed, ["forest"])
|
|
del body_["tags"]
|
|
status, result = await self.put_tags(body_)
|
|
self.assertEqual(status, 400)
|
|
with self.subTest("body over 64 KiB"):
|
|
body_ = self.save_body(observed, ["forest"])
|
|
body_["tags"] = ["x" * 60] * 1200
|
|
response = await self.client.put(
|
|
TAGS,
|
|
data=json.dumps(body_),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
self.assertEqual(response.status, 413)
|
|
verify.assert_not_called()
|
|
|
|
# Storage restored under the dialog: storage-changed, zero writes.
|
|
before = self.db_snapshot()
|
|
fresh = await self.observe_ok()
|
|
rotate_history_epoch(self.api.db.model.get_connection())
|
|
status, body = await self.put_tags(self.save_body(fresh, ["woods"]))
|
|
self.assertEqual((status, body["code"]), (409, "storage-changed"))
|
|
self.assertEqual(body["library_epoch"], self.api.db.get_execution_epoch())
|
|
self.assertNotEqual(body["library_epoch"], self.epoch)
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
|
|
async def test_put_replay_returns_replayed_true(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
observed = await self.observe_ok()
|
|
body = self.save_body(observed, ["forest", "mist"])
|
|
status, first = await self.put_tags(body)
|
|
self.assertEqual(
|
|
(status, first["replayed"], first["tag_revision"]), (200, False, 1)
|
|
)
|
|
before = self.db_snapshot()
|
|
|
|
status, again = await self.put_tags(body)
|
|
self.assertEqual(status, 200, again)
|
|
self.assertEqual(again, {**first, "replayed": True})
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
|
|
# The same set at the new revision is a replay as well.
|
|
body["expected_tag_revision"] = 1
|
|
status, third = await self.put_tags(body)
|
|
self.assertEqual(
|
|
(status, third["replayed"], third["tag_revision"]), (200, True, 1)
|
|
)
|
|
self.assertEqual(self.db_snapshot(), before)
|
|
|
|
async def test_put_saves_only_the_target_byte_set(self):
|
|
"""Tagging one file never touches a sibling that shares its prompt."""
|
|
a = make_png(self.output_dir / "sub" / "a.png")
|
|
make_png(self.output_dir / "sub" / "b.png", color=(2, 2, 2))
|
|
twin = self.output_dir / "sub" / "twin.png"
|
|
twin.write_bytes(a.read_bytes())
|
|
prompt_id = self._save_prompt("shared prompt", tags=["shared"])
|
|
for name in ("a.png", "b.png", "twin.png"):
|
|
self.api.db.link_image_to_prompt(
|
|
prompt_id, str(self.output_dir / "sub" / name)
|
|
)
|
|
observed_b = await self.observe_ok("sub/b.png")
|
|
self.saved_set(
|
|
"sub/b.png", (self.output_dir / "sub" / "b.png").read_bytes(), ["b"]
|
|
)
|
|
legacy_before = {
|
|
k: v
|
|
for k, v in self.db_snapshot().items()
|
|
if k
|
|
not in (
|
|
"media_annotation_files",
|
|
"media_annotation_bytes",
|
|
"media_image_tags",
|
|
)
|
|
}
|
|
|
|
observed = await self.observe_ok("sub/a.png")
|
|
status, receipt = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
self.assertEqual(status, 200, receipt)
|
|
|
|
after = self.db_snapshot()
|
|
legacy_after = {
|
|
k: v
|
|
for k, v in after.items()
|
|
if k
|
|
not in (
|
|
"media_annotation_files",
|
|
"media_annotation_bytes",
|
|
"media_image_tags",
|
|
)
|
|
}
|
|
# The tag registry gained "forest"; prompts, links and prompt tags are equal.
|
|
self.assertEqual(
|
|
{k: v for k, v in legacy_after.items() if k != "tags"},
|
|
{k: v for k, v in legacy_before.items() if k != "tags"},
|
|
)
|
|
self.assertEqual((await self.observe_ok("sub/b.png"))["tags"], ["b"])
|
|
self.assertEqual((await self.observe_ok("sub/twin.png"))["tags"], [])
|
|
self.assertEqual(observed_b["tags"], [])
|
|
|
|
|
|
class TestFailures(AnnotationAPITestCase):
|
|
async def test_unexpected_failures_answer_generic_envelopes(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
secret = str(self.output_dir / "sub" / "a.png")
|
|
with patch.object(
|
|
media_capture.CaptureStore,
|
|
"capture",
|
|
autospec=True,
|
|
side_effect=media_capture.CaptureError(secret),
|
|
):
|
|
status, body = await self.observe({"target": self.target()})
|
|
self.assertEqual((status, body["code"]), (500, "capture-error"))
|
|
self.assert_no_paths(body)
|
|
|
|
with (
|
|
patch.object(
|
|
media_capture.CaptureStore,
|
|
"capture",
|
|
autospec=True,
|
|
side_effect=OSError(5, "I/O error", secret),
|
|
),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
status, body = await self.observe({"target": self.target()})
|
|
self.assertEqual((status, body["code"]), (500, "observe-failed"))
|
|
self.assert_no_paths(body)
|
|
|
|
with (
|
|
patch.object(
|
|
self.api.db, "read_image_tags", side_effect=RuntimeError(secret)
|
|
),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
status, body = await self.observe({"target": self.target()})
|
|
self.assertEqual((status, body["code"]), (500, "observe-failed"))
|
|
self.assert_no_paths(body)
|
|
# The capture of the failed observation is released at once.
|
|
self.assertEqual(self.api._captures.stats()["captures"], 0)
|
|
|
|
observed = await self.observe_ok()
|
|
with (
|
|
patch.object(
|
|
self.api.db, "replace_image_tags", side_effect=RuntimeError(secret)
|
|
),
|
|
self.assertLogs("prompt_manager.api", "ERROR"),
|
|
):
|
|
status, body = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
self.assertEqual((status, body["code"]), (500, "save-failed"))
|
|
self.assert_no_paths(body)
|
|
self.assertEqual(self.bytes_rows(), [])
|
|
|
|
async def test_title_lookup_failure_still_observes(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
with patch.object(
|
|
self.api.db, "get_image_link_summaries", side_effect=RuntimeError("db")
|
|
):
|
|
body = await self.observe_ok()
|
|
self.assertIsNone(body["title"])
|
|
self.assertEqual((body["tags"], body["tag_revision"]), ([], 0))
|
|
|
|
async def test_save_of_a_released_capture_is_expired(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
observed = await self.observe_ok()
|
|
resp = await self.client.delete(observed["capture"]["preview_url"])
|
|
self.assertEqual(resp.status, 200)
|
|
status, body = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
|
self.assertEqual(self.bytes_rows(), [])
|
|
|
|
async def test_cleanup_removes_the_capture_directory_and_subscription(self):
|
|
import database.execution_history as observers
|
|
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
await self.observe_ok()
|
|
directory = self.api._captures.directory
|
|
identity = os.path.normcase(os.path.realpath(self.api.db.model.db_path))
|
|
self.assertEqual(len(observers._restore_observers[identity]), 1)
|
|
await self.api.cleanup_annotation_captures(None)
|
|
await self.api.cleanup_annotation_captures(None)
|
|
self.assertFalse(os.path.exists(directory))
|
|
self.assertNotIn(identity, observers._restore_observers)
|
|
# A later observation starts a fresh store.
|
|
await self.observe_ok()
|
|
self.assertNotEqual(self.api._captures.directory, directory)
|
|
|
|
|
|
class TestEnvelopes(AnnotationAPITestCase):
|
|
async def test_responses_never_contain_absolute_paths_or_locators(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
bodies = []
|
|
observed = await self.observe_ok()
|
|
bodies.append(observed)
|
|
status, body = await self.observe({"target": self.target("sub/missing.png")})
|
|
bodies.append(body)
|
|
status, body = await self.observe({"target": self.target("../a.png")})
|
|
bodies.append(body)
|
|
status, body = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
bodies.append(body)
|
|
status, body = await self.put_tags(self.save_body(observed, ["woods"]))
|
|
bodies.append(body)
|
|
resp = await self.client.get(CAPTURE + "missing")
|
|
bodies.append(await self.json_of(resp))
|
|
resp = await self.client.delete(observed["capture"]["preview_url"])
|
|
bodies.append(await self.json_of(resp))
|
|
status, body = await self.put_tags(self.save_body(observed, ["forest"]))
|
|
bodies.append(body)
|
|
for body in bodies:
|
|
self.assert_no_paths(body)
|
|
self.assertNotIn(self.locator("sub/a.png"), json.dumps(body))
|
|
self.assertNotIn("locator", json.dumps(body))
|
|
|
|
async def test_legacy_inspect_and_autotag_shapes_unchanged(self):
|
|
make_png(self.output_dir / "image.png")
|
|
resp = await self.client.get(
|
|
"/prompt_manager/images/inspect/image.png", params={"root": "0"}
|
|
)
|
|
self.assertEqual(resp.status, 200)
|
|
data = await self.json_of(resp)
|
|
# Additive only: B3 adds the image_tags facet beside the Library facts.
|
|
self.assertEqual(
|
|
sorted(data),
|
|
[
|
|
"file",
|
|
"image",
|
|
"image_tags",
|
|
"inspection_version",
|
|
"library",
|
|
"metadata",
|
|
"prompts",
|
|
"resolution",
|
|
"root_scope",
|
|
"success",
|
|
],
|
|
)
|
|
self.assertEqual(sorted(data["image"]), ["filename", "path", "root_index"])
|
|
self.assertEqual(data["image_tags"], {"state": "none", "tags": []})
|
|
self.assertEqual(
|
|
sorted(data["library"]),
|
|
[
|
|
"category",
|
|
"category_scope",
|
|
"pair_id",
|
|
"positive_id",
|
|
"scope",
|
|
"status",
|
|
"tags",
|
|
"tags_scope",
|
|
"title",
|
|
],
|
|
)
|
|
|
|
class FakeService:
|
|
custom_prompt = ""
|
|
default_prompt = "p"
|
|
wd14_general_threshold = 0.35
|
|
wd14_character_threshold = 0.85
|
|
|
|
def get_models_status(self):
|
|
return {"wd14-vit": {"downloaded": False, "model_path": None}}
|
|
|
|
def is_model_loaded(self):
|
|
return False
|
|
|
|
def get_loaded_model_type(self):
|
|
return None
|
|
|
|
with patch("py.autotag.get_autotag_service", return_value=FakeService()):
|
|
resp = await self.client.get("/prompt_manager/autotag/models")
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(
|
|
sorted(await self.json_of(resp)),
|
|
[
|
|
"default_prompt",
|
|
"loaded_model_type",
|
|
"model_loaded",
|
|
"models",
|
|
"success",
|
|
"suggestion_limit", # B2 adds it (D68); nothing else changes
|
|
"wd14_character_threshold",
|
|
"wd14_general_threshold",
|
|
],
|
|
)
|
|
|
|
def test_preview_types_cover_every_capture_extension(self):
|
|
for extension in media_capture.CAPTURE_EXTENSIONS:
|
|
kind = annotations_module.preview_content_type(extension)
|
|
self.assertTrue(kind.startswith(("image/", "video/")), extension)
|
|
self.assertEqual(
|
|
annotations_module.preview_content_type(".exe"), "application/octet-stream"
|
|
)
|
|
|
|
|
|
class AnnotationHostCleanupTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_root_registers_capture_cleanup_on_the_host_application(self):
|
|
from aiohttp import web
|
|
|
|
try:
|
|
from tests.test_init_package import fake_server_module, load_package
|
|
except ImportError:
|
|
from test_init_package import fake_server_module, load_package
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
server = fake_server_module()
|
|
server.PromptServer.instance.app = web.Application()
|
|
with load_package(
|
|
"annotation_cleanup_pkg",
|
|
str(Path(directory) / "db.sqlite"),
|
|
directory,
|
|
server,
|
|
) as (package, host):
|
|
api = package.api
|
|
self.assertIn(api.cleanup_annotation_captures, host.app.on_cleanup)
|
|
store = api._capture_store()
|
|
host.app.freeze()
|
|
await host.app.cleanup()
|
|
self.assertFalse(os.path.exists(store.directory))
|
|
self.assertIsNone(api._captures)
|