Vito
1b8fbdbb75
* feat(annotations): add scoped observe, capture preview and tag save routes
POST /prompt_manager/image-annotations/observe captures the file named by
the listing identity {root_index, root_scope, path} (scope checked against
the current roots before any file access) and answers the saved tag set of
exactly those bytes, refreshing a matching record's verified fingerprint.
GET/DELETE .../capture/{token} serve the private copy (no-store, nosniff,
pinned with acquire/release_use while it streams) and release it.
PUT .../tags re-hashes the current file and writes only when it still
holds the captured bytes, under the storage epoch and tag revision
(file-changed, storage-changed, tag-conflict, tag-limit, replayed).
Bodies are capped at 64 KiB with unknown fields refused; envelopes never
carry server paths or locators. Restores invalidate every capture and the
host's cleanup removes the capture directory.
* feat(autotag): scored suggestions with a configurable count behind a shared model gate
- py/model_gate.py: one process-wide, non-blocking gate for loading,
swapping, running, downloading and unloading a tagging model. Worker
calls run shielded; a release requested while one still runs takes
effect only when it exits (client disconnect keeps the gate).
- Legacy /autotag/single, /start, /download/{type} and /unload hold the
gate and answer 409 {code: "model-busy"} while another operation holds
it; their bodies are otherwise unchanged. /autotag/models gains
suggestion_limit.
- AutoTagService.generate_tag_suggestions: WD14 (name, probability) best
first, GGUF/HF (name, None) in output order, first `limit` distinct
case-insensitively; generate_tags keeps its return; the image file is
closed before inference returns.
- POST /prompt_manager/image-annotations/suggest runs one engine on the
pinned capture, with strict per-engine options, never downloads
(409 model-unavailable) and never writes.
- config.json autotag.defaults.suggestion_limit, integer 1..100, default
10, validated on load.
- Capture probe decides Generate from the image header alone, so prompt
metadata over the inspection budgets no longer blocks it.
* feat(images): match image tags in the Tags filter, search and inspector
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.
* fix(annotations): answer a fixed coded message when the filtered image list fails
A failed prompt-link or image-tag read under a Tags, category or rating
filter returned the raw exception text (SQL or database messages) in the
500 body. The workspace listing now answers code listing-failed with a
fixed sentence and logs the traceback. The legacy output listing and the
inspector image-tag read now log with exc_info as well.
* fix(annotations): drop suggestion names that Save would reject
A JoyCaption part over 50 characters or holding a newline came back as a
preselected chip, so the first Save of the dialog failed with a 400 and
nothing was stored. generate_tag_suggestions now applies the vocabulary
tag-name rules before the case-insensitive dedupe and the limit count, so
the requested number of valid chips still comes back.
* fix(annotations): stop changing the process warnings filters during capture
header_dimensions wrapped Image.open in warnings.catch_warnings with an
error filter for DecompressionBombWarning. catch_warnings is not
thread-safe, so two overlapping captures could leave that error filter
installed for every later Pillow open in the ComfyUI process. Only
Pillow's DecompressionBombError is now mapped to HeaderTooLarge; a size
that only warns is refused by the existing inference pixel cap.
* fix(annotations): free pins on cancellation, serve preview ranges, code every storage error
- A preview cancelled while its file opens closes the handle and only
then ends the pin, so a doomed capture is never unlinked while open.
- An observe cancelled while capturing or reading releases the capture
instead of leaving it until the idle expiry.
- The model gate takes an on_skipped callback that runs once when the
work never starts; Generate passes the capture's release_use, so an
executor refusal no longer leaks a pin against the capture budget.
- The preview answers a single byte range with 206 and Content-Range
(416 for an unsatisfiable or multi-range header), still streamed inside
the pin, so Safari and WebKit can play video captures.
- _coded maps tag-limit (400), metadata-conflict and vocabulary-conflict,
and falls back to 409 with a fixed sentence for any other code instead
of raising KeyError.
* fix(annotations): deny case and Unicode spellings of a listed file
On macOS (case- and normalization-insensitive APFS) a crafted target such
as A.PNG or an NFD name opened the listed a.png or NFC file while
realpath kept the requested spelling, so the tags were stored under a
second locator the Images filter never shows. Capture and verify now
require each path part to be spelled as its directory lists it and deny
any other case or Unicode spelling as an alias (409 source-changed, like
a symlink alias). A missing file with no other spelling keeps its usual
outcome.
* test(annotations): one shared folder_paths stub for the API test modules
test_api_images and test_api_prompts installed a MagicMock folder_paths,
so whenever one of them was imported before test_api_autotag the later
patch("folder_paths.get_output_directory") in test_api_history_outputs
left a mock attribute behind and two test_diagnostics_config tests
failed. tests/host_stubs.py now installs one real-module stub, used by
test_api_images, test_api_prompts and test_api_autotag (and through
test_api_images by the annotation modules, whose copy is gone), so the
modules pass in either import order.
* fix(annotations): derive the suggestion cap from the image tag set cap
py/autotag.py now takes MAX_SUGGESTION_LIMIT from database.annotations
MAX_IMAGE_TAGS. The two copies that cannot import it (py/config.py loads
before the database, utils/inspection_response.py stays storage-free) are
pinned equal by a test, so the 100 limits cannot drift apart.
* docs(testing): record capture aliases, preview ranges and the PR B follow-ups
TESTING.md now states that case and Unicode spellings are aliases, that
the preview serves byte ranges inside the capture pin, where the
cancellation tests live, and two tracked follow-ups: the legacy batch's
ungated cleanup unload under its own last inference, and the three files
over the 800-line ceiling (split py/autotag.py first).
1223 lines
52 KiB
Python
1223 lines
52 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 asyncio
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
import uuid
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from aiohttp.test_utils import make_mocked_request
|
|
|
|
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
|
|
WAIT = 30.0 # ceiling for cross-thread handoffs, never a short clock
|
|
|
|
# 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_denies_case_spellings_of_the_listed_file(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
for alias in ("sub/A.PNG", "SUB/a.png"):
|
|
with self.subTest(alias=alias):
|
|
status, body = await self.observe({"target": self.target(alias)})
|
|
# Denied as an alias, like a symlink to the listed file.
|
|
self.assertEqual((status, body["code"]), (409, "source-changed"))
|
|
self.assert_no_paths(body)
|
|
store = self.api._captures
|
|
self.assertFalse(any(Path(store.directory).iterdir()))
|
|
await self.observe_ok("sub/a.png")
|
|
|
|
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_answers_byte_ranges(self):
|
|
# Safari and WebKit play <video> only from byte-range responses.
|
|
path = make_png(self.output_dir / "sub" / "a.png", size=(64, 64))
|
|
data = path.read_bytes()
|
|
size = len(data)
|
|
body = await self.observe_ok()
|
|
url = body["capture"]["preview_url"]
|
|
token = body["capture"]["token"]
|
|
|
|
full = await self.client.get(url)
|
|
self.assertEqual(full.status, 200)
|
|
self.assertEqual(full.headers["Accept-Ranges"], "bytes")
|
|
self.assertEqual(await full.read(), data)
|
|
for header, start, end in (
|
|
("bytes=0-15", 0, 16),
|
|
("bytes=10-", 10, size),
|
|
("bytes=-12", size - 12, size),
|
|
(f"bytes=5-{size + 100}", 5, size),
|
|
):
|
|
with self.subTest(range=header):
|
|
resp = await self.client.get(url, headers={"Range": header})
|
|
self.assertEqual(resp.status, 206)
|
|
self.assertEqual(
|
|
resp.headers["Content-Range"], f"bytes {start}-{end - 1}/{size}"
|
|
)
|
|
self.assertEqual(resp.headers["Content-Length"], str(end - start))
|
|
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(await resp.read(), data[start:end])
|
|
for header in (f"bytes={size}-", "bytes=9-3", "items=0-1", "bytes=0-1,4-5"):
|
|
with self.subTest(unsatisfiable=header):
|
|
resp = await self.client.get(url, headers={"Range": header})
|
|
self.assertEqual(resp.status, 416)
|
|
self.assertEqual(resp.headers["Content-Range"], f"bytes */{size}")
|
|
self.assertEqual(resp.headers["Cache-Control"], "no-store")
|
|
store = self.api._captures
|
|
self.assertEqual(store._entries[token].uses, 0)
|
|
self.assertEqual(open_handles_under(store.directory), [])
|
|
self.assertIsNotNone(store.get(token))
|
|
|
|
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 TestCancellation(AnnotationAPITestCase):
|
|
"""A cancelled handler never leaks a pin, a handle or a capture."""
|
|
|
|
async def cancel_when(self, coro, entered):
|
|
task = asyncio.ensure_future(coro)
|
|
await asyncio.get_running_loop().run_in_executor(None, entered.wait, WAIT)
|
|
self.assertTrue(entered.is_set())
|
|
task.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await task
|
|
|
|
async def settle(self, done, what):
|
|
for _ in range(3000):
|
|
if done():
|
|
return
|
|
await asyncio.sleep(0.01)
|
|
self.fail(what)
|
|
|
|
def observe_request(self, path="sub/a.png"):
|
|
body = AsyncMock(return_value=({"target": self.target(path)}, None))
|
|
patcher = patch.object(self.api, "_annotation_body", body)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
return make_mocked_request("POST", OBSERVE)
|
|
|
|
async def test_cancelled_preview_open_closes_the_handle_before_release(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
body = await self.observe_ok()
|
|
token = body["capture"]["token"]
|
|
store = self.api._captures
|
|
capture_path = store.get(token).path
|
|
entered, proceed = threading.Event(), threading.Event()
|
|
handles = []
|
|
|
|
def slow_open(path):
|
|
entered.set()
|
|
proceed.wait(WAIT)
|
|
handle = open(path, "rb")
|
|
handles.append(handle)
|
|
return handle
|
|
|
|
request = make_mocked_request(
|
|
"GET", CAPTURE + token, match_info={"token": token}
|
|
)
|
|
with patch.object(annotations_module, "_open_capture", slow_open):
|
|
await self.cancel_when(self.api.get_capture_preview(request), entered)
|
|
# A DELETE lands meanwhile; the file stays until the open settles.
|
|
store.release(token)
|
|
self.assertTrue(os.path.exists(capture_path))
|
|
proceed.set()
|
|
await self.settle(
|
|
lambda: not os.path.exists(capture_path), "capture never removed"
|
|
)
|
|
self.assertEqual(len(handles), 1)
|
|
self.assertTrue(handles[0].closed)
|
|
|
|
async def test_cancelled_observe_capture_releases_the_new_capture(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
store = self.api._capture_store()
|
|
real_capture = store.capture
|
|
entered, proceed = threading.Event(), threading.Event()
|
|
made = []
|
|
|
|
def slow_capture(*args, **kwargs):
|
|
entered.set()
|
|
proceed.wait(WAIT)
|
|
capture = real_capture(*args, **kwargs)
|
|
made.append(capture)
|
|
return capture
|
|
|
|
request = self.observe_request()
|
|
with patch.object(store, "capture", slow_capture):
|
|
await self.cancel_when(self.api.observe_image(request), entered)
|
|
proceed.set()
|
|
await self.settle(
|
|
lambda: made and store.get(made[0].token) is None,
|
|
"the cancelled observe left its capture",
|
|
)
|
|
self.assertFalse(any(Path(store.directory).iterdir()))
|
|
|
|
async def test_cancelled_observe_read_releases_the_capture(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
store = self.api._capture_store()
|
|
real_read = self.api._observe_capture
|
|
entered, proceed, finished = (threading.Event() for _ in range(3))
|
|
tokens = []
|
|
|
|
def slow_read(store_, capture, root, relative):
|
|
tokens.append(capture.token)
|
|
entered.set()
|
|
proceed.wait(WAIT)
|
|
try:
|
|
return real_read(store_, capture, root, relative)
|
|
finally:
|
|
finished.set()
|
|
|
|
request = self.observe_request()
|
|
with patch.object(self.api, "_observe_capture", slow_read):
|
|
await self.cancel_when(self.api.observe_image(request), entered)
|
|
self.assertIsNone(store.get(tokens[0]))
|
|
proceed.set()
|
|
await asyncio.get_running_loop().run_in_executor(None, finished.wait, WAIT)
|
|
await self.settle(
|
|
lambda: not any(Path(store.directory).iterdir()), "capture file left"
|
|
)
|
|
|
|
|
|
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_every_storage_code_has_a_coded_envelope(self):
|
|
cases = {
|
|
"tag-limit": 400,
|
|
"tag-conflict": 409,
|
|
"storage-changed": 409,
|
|
"metadata-conflict": 409,
|
|
"vocabulary-conflict": 409,
|
|
"some-future-code": 409,
|
|
}
|
|
for code, status in cases.items():
|
|
with self.subTest(code=code):
|
|
response = annotations_module._coded(code)
|
|
self.assertEqual(response.status, status)
|
|
body = json.loads(response.body)
|
|
self.assertEqual((body["success"], body["code"]), (False, code))
|
|
self.assertTrue(body["error"])
|
|
self.assertEqual(response.headers["Cache-Control"], "no-store")
|
|
|
|
async def test_save_maps_an_unlisted_storage_code_to_an_envelope(self):
|
|
from database.annotations import AnnotationCommandError
|
|
from database.library_commands import LibraryCommandError
|
|
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
observed = await self.observe_ok()
|
|
for error, status in (
|
|
(AnnotationCommandError("tag-limit", self.epoch, None), 400),
|
|
(LibraryCommandError("vocabulary-conflict", self.epoch), 409),
|
|
):
|
|
with (
|
|
self.subTest(code=error.code),
|
|
patch.object(self.api.db, "replace_image_tags", side_effect=error),
|
|
):
|
|
status_, body = await self.put_tags(
|
|
self.save_body(observed, ["forest"])
|
|
)
|
|
self.assertEqual((status_, body["code"]), (status, error.code))
|
|
self.assertEqual(body["library_epoch"], self.epoch)
|
|
self.assert_no_paths(body)
|
|
|
|
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)
|