Vito
3f11ea6168
* feat(review): add review queue schema with restore validation
Add the Needs review schema group (media_review_metadata,
media_review_items, media_review_decisions) beside the image annotation
group it cites. review_schema_error validates the catalog, columns,
foreign keys, row values and canonical 8 MiB evidence read-only;
migrate_review_state installs it in a SAVEPOINT after the annotation
group, and verify_database_file requires it under require_library_state.
A review group without a valid annotation group is refused. Decisions
null their pair and image navigation through indexes when those are
deleted. Old backups and annotation-era backups restore and gain the
group; browser fixtures clear the new tables on reset.
* feat(review): persist ambiguous sampler evidence from scans and new outputs
The output scan and the live monitor now queue every file the ingestion
resolver finds ambiguous for Needs review. The file is captured, its metadata
is read again from the capture and resolved with the default profile, and only
an ambiguous result is stored, bound to the capture's annotation byte record
and fenced by its storage epoch. Evidence holds one hashed projection per
selectable candidate; none, more than 16 or over 8 MiB is a limited item.
Recording is idempotent per bytes and evidence and never touches a decision.
Scan statuses, counters, outcome rows and legacy SSE events are unchanged; the
job scope gains review_items/review_errors and the log a Needs review line.
The monitor records after its fallback link, best effort.
* feat(review): explicit association decisions with compatible linking
Needs review gains its queue reads and decisions:
- GET /prompt_manager/review/items lists pending items oldest first with
the exact pending total and an HMAC keyset cursor bound to the storage
epoch and the review index revision.
- GET /prompt_manager/review/items/{id} adds the candidates (role texts,
label, selectability), the dependency fingerprint and the digest.
- POST /prompt_manager/review/items/{id}/decision applies a candidate only
for a capture of the item's bytes that the current file still hashes to
and whose evidence is unchanged, through the ingester's compatible-only
rules in the same transaction as the decision row; a conflict writes
nothing. Leave unresolved records the decision and links nothing.
ingest_resolved_pair keeps its contract; its body moves into
_ingest_resolved_pair_in, which runs in the caller's transaction.
Items keep the file's recorded spelling so targets keep their case on
Windows, and a rescan skips the copy and hash of a file whose size and
mtime match the bytes an item already cites.
* fix(review): build the spelling validation test from native paths
On Windows the locator is case-folded and uses backslashes, so the fixed
POSIX pair /out/a.png no longer folded to itself and the new spelling
check refused the good case. The test now spells a mixed-case native path
and folds it with os.path.normcase, and also refuses a locator of another
file.
* fix(review): Apply sees the file under every spelling and rechecks only its choice
Apply now finds the legacy rows of the same physical file however its root
is reached. The roots are real paths while the monitor links under the
directory it watches, so a symlinked or junctioned root, subst or a mapped
drive used to hide the monitor's row and Apply added a second, differently
linked row. Same-name rows are resolved to their physical identity
(normcase(realpath), else samefile); the dependency fingerprint covers all
of them, two rows refuse the choice, one row is enriched under the
ingester's compatible-only rules, and a same-name row the request did not
classify is a changed association.
Apply re-verifies only what the decision depends on: the bytes are still
ambiguous and still offer the chosen candidate. A resolver change that
rewords the reason or reorders candidates leaves pending items applicable.
The resolver version is one constant in prompt_pair_resolver, shared by
ingestion provenance and Needs review.
The transaction's byte check receives the re-hashed digest (or the body's
for Leave unresolved), every Apply refusal branch is tested with zero
writes, and the root listing and realpath calls run on a worker thread.
* fix(review): validate stored review evidence once per startup
Startup used to run the full review validation twice (the early guard and
the install) and the annotation group's row checks inside each, so every
evidence document was parsed and re-encoded twice. The early guard is now
catalog-only (a review group without annotation tables), and the install
validates the group once, trusting the annotation group it has just
validated. The evidence predicate no longer pattern-scans every document in
SQL, because the streamed canonical pass proves it, and the codec is built
once.
Measured on a database with 40,000 items of 4 KiB evidence: startup open
went from 7.7 s to 2.5 s. A test pins one evidence pass per reopen.
Restore and verify_database_file keep the full validation.
* fix(review): keep one capture worker for interactive requests
Needs review recording from the scan and the monitor shared the capture
store's two workers with the Tag dialog and Apply, so a busy scan could
answer a user's preview or decision with "Too many previews are open".
Background captures now never take the last worker or capture slot, and
wait up to two seconds for room before counting the file as not added.
Interactive captures are admitted or refused at once, as before.
* fix(review): refresh the recorded spelling when a file is seen under a new one
A case-only rename on Windows or macOS keeps size and mtime_ns, so the
rescan skipped the file and the item kept the old spelling. Its target was
then formatted from a name the directory no longer lists, the capture
refused it, and the item stayed pending with no way to apply it. The skip
path and a duplicate insert now record the current spelling for every item
at the locator and move the review index revision; item revisions do not
change, so an open decision stays valid. Tested under case-folding
semantics on every platform.
* fix(review): final capture store close and recoverable stale cursors
Closing the capture store at shutdown now runs under the store lock and is
final: a monitor recording or scan batch still in flight gets a capture
error instead of creating a new temp directory and restore subscription
after the cleanup.
A well-formed queue cursor this process cannot verify (opened before a
ComfyUI restart, whose cursor secret is per process) now answers 409
review-conflict with the current total, like any other stale page, instead
of an uncoded 400. Malformed cursors stay 400.
* fix(review): make only the shutdown cleanup final
The previous commit made every close of the capture store final, which
also refused the explicit closes that swap in a fresh store. Only the host's
on_cleanup hook is final now; an observe after it answers a coded capture
error instead of an unhandled 500.
* refactor(admin): move Needs review recording into its own mixin
_record_scanned_review, _record_review_sync, record_review_file and
_log_review_items move unchanged from py/api/admin.py to
py/api/review_recording.py (ReviewRecordingMixin), keeping admin.py from
growing further past the size ceiling. No behaviour change.
* fix(review): record the review fix round and the skip's residual risks in TESTING.md
The Needs review paragraphs now state the rescan skip's residual risks (an
in-place rewrite that keeps size and mtime_ns is never re-queued; a deleted
or moved file leaves its pending item counted), the physical identity rule
for Apply, the candidate-only recheck and shared resolver version, the
spelling refresh, background capture admission, the final shutdown close,
the stale-cursor answer and the measured startup validation cost.
* fix(review): refuse monitor recording beside record-changing jobs
System S1 (#289) gives every record change a registry admission class. The
scan already records Needs review items under its own records job. The
monitor's recording lasts one file and is not a job, so it now checks the
records admission: while a records or guarded job or hold runs it records
nothing, logs that the file waits for the next scan, and answers deferred.
The fallback link stands and the next "new" scan reads the ambiguous file
again, because a fallback link is not a resolved one.
JobRegistry.blocker and active iterate a snapshot of the jobs, since the
monitor's worker thread now asks them too.
* test(backup): wait for the restore, not a 5 s timer, in the thread test
test_database_is_usable_from_every_thread_after_restore let its worker
wait 5 s for the restore. On a stalled Windows runner (py3.11, where the
previous restore test also took 25 s instead of 0.3 s) the restore took
longer, the wait timed out, and the worker read the old database before
the restore ran. The worker now waits for the event under a 120 s bound
that only a hung test reaches, reports a timeout as an error, and the
main thread releases it even when the restore raises. With a restore
delayed by 6 s the old test fails and the new one passes.
* fix(review): store non-finite evidence as null and count pending items
canonical_evidence now follows main's non-finite policy (#303): NaN and
±Infinity are written as null instead of raising, so a NaN-stamped graph
never turns a review item limited. The System Overview reports the review
queue's pending total in needs_review.pending, read in the same snapshot
as the Library counts.
* fix(review): show the pending count on the System Overview review row
The Overview's review row now reads "N ambiguous image links" with the
concept's line, from needs_review.pending. Its Review metadata action
stays aria-disabled until the Needs review page ships. The system specs
and their concept comparison follow the real count.
* fix(system): answer the Overview and Statistics only after every read ended
asyncio.gather answered at the first failed read while its sibling
executor reads were still inside SQLite. In the suite the test teardown
then deleted the database under the running read, and py3.10 on Ubuntu
segfaulted in get_statistics (test_statistics_failure_is_an_error_envelope).
Both routes now wait for all their reads before answering.
* fix(review): queue only ambiguous files that nothing links, and heal the queue
The maintainer's live scan (19,908 files, scope "new") reported 4,677 files
under Needs review, but the queue held 10,232 items: every file the strict
resolver found ambiguous was queued, including the ones the legacy
browser-parser fallback then linked (all sampled "added" rows) and the ones
already linked to a known prompt (135 of 200 sampled "known" rows).
A file is now queued only when its final outcome is ambiguous and unlinked:
the resolver finds it ambiguous, the fallback links nothing and no saved row
links it already. An ambiguous file that a saved row links is reported as
unchanged ("Already linked to its prompt"), not as Needs review.
The queue heals itself. After each batch, the scan removes the pending items
of every listed file (read, or skipped as known) that a saved row now links.
The monitor does the same for each ambiguous output it links, and never
queues one. Only an undecided item is removed (pending, revision 0, no
decision row), so a decision is never deleted or changed and the schema
needs no new state.
The scan job's scope.review_items now counts this scan's files that are in
the pending queue after it, new or already there, so it agrees with the
review group's ambiguous rows and with the queue total. review_added and
review_settled are new, and the log gains "already in Needs review" and
"left Needs review · now linked" lines. A duplicate whose item is decided is
reported as "decided" and is not counted.
* fix(review): never queue an ambiguous file with nothing to choose
The maintainer's live queue held 39 items whose candidates all read
"unresolved" ("No positive / No negative"). An ambiguous result is now
reviewable only when at least two candidates are resolved with non-blank
positive text. With none the file could not be traced; with exactly one,
the untraced branches may still have made the file, so that text is not
certain either. Both are reported unresolved, never linked and never
queued (also the A1111 two-delimiter case, which has no candidates).
A rescan removes the pending items of such files whose own evidence
offers no choice, with the same safe delete as the linked settle
(pending, revision 0, no decision row). A limited item (over 16 candidates) that holds at least two
candidates is kept. The job scope gains review_untraced and the log
"= N images left Needs review · nothing to choose".
* test(scan): pin only the per-source scope keys beside the review counters
* test(system): match the scan scope keys beside the review counters
1231 lines
52 KiB
Python
1231 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)
|
|
# The shutdown cleanup is final: a later observation (or a monitor
|
|
# recording still in flight) creates no new store or directory.
|
|
status, body = await self.observe({"target": self.target("sub/a.png")})
|
|
self.assertEqual((status, body["code"]), (500, "capture-error"), body)
|
|
self.assertIsNone(self.api._captures)
|
|
self.assertNotIn(identity, observers._restore_observers)
|
|
# An explicit, non-final close still lets a later observation start
|
|
# a fresh store (the tests' own store swaps rely on it).
|
|
self.api._captures_closed = False
|
|
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)
|