ComfyAssets-ComfyUI_PromptManager/tests/test_media_capture.py

"""Bounded private capture store (#249 A3): exact bytes, hard limits, cleanup."""

import hashlib
import os
import stat
import struct
import tempfile
import threading
import time
import unicodedata
import unittest
import warnings
import zlib
from pathlib import Path
from unittest.mock import patch

try:
    from tests.open_handles import assert_closed, open_handles_under
except ImportError:  # discovered with tests/ as the top-level directory
    from open_handles import assert_closed, open_handles_under

try:  # guarded so a missing module fails assertions, not test loading
    from utils import media_capture as capture_module
    from utils.media_capture import (
        CAPTURE_EXTENSIONS,
        CaptureCapacity,
        CaptureChanged,
        CaptureDenied,
        CaptureStore,
        CaptureTooLarge,
    )
except ImportError:
    capture_module = None
    CAPTURE_EXTENSIONS = None
    CaptureCapacity = CaptureChanged = CaptureDenied = CaptureTooLarge = None
    CaptureStore = None

EPOCH = "epoch-1"
SCOPE = "a" * 64
WAIT = 30.0  # generous ceiling for cross-thread handoffs, never a short clock


def _chunk(kind: bytes, data: bytes) -> bytes:
    return (
        struct.pack(">I", len(data))
        + kind
        + data
        + struct.pack(">I", zlib.crc32(kind + data))
    )


def png_bytes(width: int, height: int, pixels: bytes = b"\x00" * 16) -> bytes:
    """A structurally valid PNG whose header declares *width* x *height*."""
    ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
    return (
        b"\x89PNG\r\n\x1a\n"
        + _chunk(b"IHDR", ihdr)
        + _chunk(b"IDAT", pixels)
        + _chunk(b"IEND", b"")
    )


class FakeClock:
    def __init__(self, now: float = 1_000.0):
        self.now = now

    def __call__(self) -> float:
        return self.now

    def advance(self, seconds: float) -> None:
        self.now += seconds


class CaptureStoreTestCase(unittest.TestCase):
    def setUp(self):
        self.assertIsNotNone(capture_module, "utils.media_capture is not implemented")
        parent = tempfile.TemporaryDirectory()
        self.addCleanup(parent.cleanup)
        self.addCleanup(assert_closed, self, parent.name)
        # The resolved path: macOS /var -> /private/var, Windows 8.3 names.
        self.parent = Path(os.path.realpath(parent.name))
        self.root = self.parent / "output"
        self.root.mkdir()
        self.clock = FakeClock()

    def make_store(self, **options) -> "CaptureStore":
        options.setdefault("clock", self.clock)
        store = CaptureStore(self.parent / "captures", **options)
        self.addCleanup(store.close)
        return store

    def write(self, relative: str, data: bytes) -> Path:
        path = self.root / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(data)
        return path

    def symlink(self, target, link):
        try:
            os.symlink(str(target), str(link))
        except (OSError, NotImplementedError) as error:
            self.skipTest(f"symlinks unavailable here: {error}")

    def wait_for_change_time_after(self, ns: int) -> None:
        """Block until a new file's ``st_ctime_ns`` passes *ns*.

        File timestamps tick coarsely (Linux jiffies, Windows ~1-16 ms). Each
        probe is a fresh name: Windows reports the creation time (or, on
        Python 3.12+, the change time) and tunnels the creation time of a
        reused name for 15 s.
        """
        deadline = time.monotonic() + WAIT
        for index in range(1_000_000):
            probe = self.parent / f"tick-{index}"
            probe.write_bytes(b"x")
            advanced = os.stat(probe).st_ctime_ns > ns
            probe.unlink()
            if advanced:
                return
            if time.monotonic() >= deadline:
                break
            time.sleep(0.001)
        self.fail("the filesystem clock did not advance")  # pragma: no cover

    def capture(self, store, relative="a.png", **options):
        options.setdefault("epoch", EPOCH)
        options.setdefault("scope", SCOPE)
        return store.capture(self.root, relative, **options)

    def assert_empty(self, store):
        self.assertFalse(any(Path(store.directory).iterdir()))
        self.assertEqual(
            store.stats(), {"captures": 0, "reserved_bytes": 0, "workers": 0}
        )

    # -- exact bytes and identity -------------------------------------------

    def test_capture_hashes_and_copies_exact_bytes_owner_only(self):
        source_bytes = png_bytes(3, 2, os.urandom(40))
        source = self.write("a.png", source_bytes)
        store = self.make_store()

        cap = self.capture(store, root_index=1)

        self.assertEqual(Path(cap.path).read_bytes(), source_bytes)
        self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest())
        self.assertEqual(cap.mtime_ns, os.stat(source).st_mtime_ns)
        self.assertEqual(cap.size, len(source_bytes))
        self.assertEqual(
            cap.locator,
            os.path.normcase(os.path.join(os.path.realpath(self.root), "a.png")),
        )
        self.assertEqual((cap.root_index, cap.scope, cap.epoch), (1, SCOPE, EPOCH))
        self.assertEqual((cap.media_type, cap.width, cap.height), ("image", 3, 2))
        self.assertTrue(cap.inferable)
        self.assertIsNone(cap.inferable_reason)
        self.assertIsInstance(cap.token, str)
        self.assertGreaterEqual(len(cap.token), 32)
        self.assertNotIn(cap.token, str(cap.path))
        self.assertNotIn(str(self.root), str(cap.path))
        capture_path = Path(cap.path)
        self.assertEqual(capture_path.parent, Path(store.directory))
        self.assertEqual(capture_path.suffix, ".png")
        self.assertTrue(
            Path(store.directory).name.startswith("prompt-manager-captures-")
        )
        if os.name == "posix":
            self.assertEqual(stat.S_IMODE(os.stat(store.directory).st_mode) & 0o077, 0)
            self.assertEqual(stat.S_IMODE(os.stat(cap.path).st_mode) & 0o077, 0)
        self.assertIs(store.get(cap.token), cap)
        self.assertEqual(
            store.stats(),
            {"captures": 1, "reserved_bytes": len(source_bytes), "workers": 0},
        )
        self.assertTrue(store.release(cap.token))
        self.assertFalse(store.release(cap.token))
        self.assertIsNone(store.get(cap.token))
        self.assert_empty(store)
        assert_closed(self, self.parent)

    def test_streams_in_bounded_chunks_never_the_whole_file(self):
        data = png_bytes(1, 1, os.urandom(5 * capture_module.CHUNK_SIZE + 7))
        self.write("big.png", data)
        store = self.make_store(max_file=len(data))
        sizes = []
        original = CaptureStore._read

        def hook(store_, stream, size):
            sizes.append(size)
            return original(store_, stream, size)

        with patch.object(CaptureStore, "_read", hook):
            cap = self.capture(store, "big.png")
        self.assertTrue(sizes)
        self.assertLessEqual(max(sizes), capture_module.CHUNK_SIZE)
        self.assertEqual(cap.sha256, hashlib.sha256(data).hexdigest())

    def test_symlinked_file_or_root_escape_is_denied(self):
        outside = self.parent / "outside"
        outside.mkdir()
        (outside / "real.png").write_bytes(png_bytes(1, 1))
        self.write("a.png", png_bytes(1, 1))
        self.write("notes.txt", b"text")
        (self.root / "folder").mkdir()
        store = self.make_store()

        denied = [
            "../outside/real.png",
            "a/../a.png",
            "./a.png",
            "a.png/",
            "",
            "/a.png",
            "C:/a.png",
            "a\\b.png",
            "a\x00.png",
            "notes.txt",
            "folder",
            "x" * 1021 + ".png",
        ]
        for relative in denied:
            with self.subTest(relative=relative):
                with self.assertRaises(CaptureDenied):
                    self.capture(store, relative)
        with self.assertRaises(CaptureDenied):
            self.capture(store, "a.png", root_index=True)
        with self.assertRaises(CaptureChanged):
            self.capture(store, "missing.png")
        self.assert_empty(store)

        self.symlink(outside / "real.png", self.root / "link.png")
        self.symlink(outside, self.root / "linked-dir")
        with self.assertRaises(CaptureDenied):
            self.capture(store, "link.png")
        with self.assertRaises(CaptureDenied):
            self.capture(store, "linked-dir/real.png")
        self.assert_empty(store)
        # A symlink to a file inside the root is still an alias, never a locator.
        self.symlink(self.root / "a.png", self.root / "alias.png")
        with self.assertRaises(CaptureDenied):
            self.capture(store, "alias.png")
        self.assert_empty(store)

    def test_case_and_unicode_spellings_of_a_listed_file_are_denied(self):
        # On a case- or normalization-insensitive filesystem (macOS APFS,
        # Windows) these spellings open the listed file while realpath keeps
        # the requested spelling, which would store its tags under a second
        # locator. Every OS denies them; the listed spelling is the control.
        nfc = unicodedata.normalize("NFC", "caf\u00e9.png")
        nfd = unicodedata.normalize("NFD", nfc)
        self.assertNotEqual(nfc, nfd)
        self.write("a.png", png_bytes(1, 1))
        self.write(nfc, png_bytes(1, 1))
        self.write("Sub/b.png", png_bytes(1, 1))
        store = self.make_store()

        for relative in ("a.png", nfc, "Sub/b.png"):
            with self.subTest(listed=relative):
                cap = self.capture(store, relative)
                self.assertEqual(
                    cap.locator,
                    os.path.normcase(
                        os.path.join(os.path.realpath(self.root), *relative.split("/"))
                    ),
                )
                self.assertTrue(store.release(cap.token))
        for alias in ("A.PNG", "A.png", nfd, "sub/b.png", "SUB/B.PNG", "Sub/B.png"):
            with self.subTest(alias=alias):
                with self.assertRaises(CaptureDenied):
                    self.capture(store, alias)
                with self.assertRaises(CaptureDenied):
                    store.verify(self.root, alias)
        # A missing file with no other spelling keeps its usual outcome.
        with self.assertRaises(CaptureChanged):
            self.capture(store, "Sub/missing.png")
        self.assert_empty(store)
        assert_closed(self, self.parent)

    def test_locator_matches_listing_identity_for_symlinked_root(self):
        real_root = self.parent / "real-output"
        real_root.mkdir()
        (real_root / "sub").mkdir()
        source_bytes = png_bytes(2, 2)
        (real_root / "sub" / "a.png").write_bytes(source_bytes)
        link_root = self.parent / "link-output"
        self.symlink(real_root, link_root)
        store = self.make_store()

        cap = store.capture(link_root, "sub/a.png", epoch=EPOCH, scope=SCOPE)

        expected = os.path.normcase(
            os.path.join(os.path.realpath(link_root), "sub", "a.png")
        )
        self.assertEqual(cap.locator, expected)
        self.assertEqual(
            cap.locator, os.path.normcase(os.path.realpath(link_root / "sub/a.png"))
        )
        self.assertTrue(
            cap.locator.startswith(os.path.normcase(os.path.realpath(real_root)))
        )
        self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest())

    # -- change detection ----------------------------------------------------

    def test_change_during_copy_is_detected_and_temp_removed(self):
        source_bytes = png_bytes(1, 1, b"\x11" * 40)
        source = self.write("a.png", source_bytes)
        store = self.make_store(max_file=4096)
        original = CaptureStore._read

        def appending(store_, stream, size):
            data = original(store_, stream, 16)
            if data and not getattr(appending, "done", False):
                appending.done = True
                with open(source, "ab") as handle:
                    handle.write(b"\x22" * 8)
            return data

        with patch.object(CaptureStore, "_read", appending):
            with self.assertRaises(CaptureChanged):
                store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE)
        self.assert_empty(store)

        # Same size, different bytes, mtime restored: only the change time
        # (POSIX ctime) can tell. Windows exposes no change time through stat,
        # so there the torn copy completes and its digest describes the copy.
        source.write_bytes(source_bytes)
        before = os.stat(source)
        self.wait_for_change_time_after(before.st_ctime_ns)
        rewritten = bytes(b ^ 0xFF for b in source_bytes)

        def rewriting(store_, stream, size):
            data = original(store_, stream, 16)
            if data and not getattr(rewriting, "done", False):
                rewriting.done = True
                with open(source, "r+b") as handle:
                    handle.write(rewritten)
                os.utime(source, ns=(before.st_atime_ns, before.st_mtime_ns))
            return data

        # Only a change time can tell: POSIX always has one; Windows exposes
        # it from stat on Python 3.12+ (where the OS supports the fast path)
        # and never from fstat, so detection there is best effort. The test
        # asserts the outcome, not the platform: a detected change leaves no
        # temp file; an undetected one (Windows only) completes with a digest
        # that describes the copied bytes.
        with patch.object(CaptureStore, "_read", rewriting):
            try:
                cap = store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE)
            except CaptureChanged:
                detected = True
            else:
                detected = False
                self.assertEqual(os.name, "nt", "only Windows may miss the rewrite")
                copied = Path(cap.path).read_bytes()
                self.assertEqual(cap.sha256, hashlib.sha256(copied).hexdigest())
                store.release(cap.token)
        if os.name == "posix":
            self.assertTrue(detected)
        self.assert_empty(store)
        self.assertEqual(os.stat(source).st_mtime_ns, before.st_mtime_ns)

        if os.name == "posix":
            # The path now names another file (new inode) once the copy ends.
            source.write_bytes(source_bytes)

            def replacing(store_, stream, size):
                data = original(store_, stream, size)
                if not data and not getattr(replacing, "done", False):
                    replacing.done = True
                    other = self.root / "other.png"
                    other.write_bytes(source_bytes)
                    os.replace(other, source)
                return data

            with patch.object(CaptureStore, "_read", replacing):
                with self.assertRaises(CaptureChanged):
                    store.capture(self.root, "a.png", epoch=EPOCH, scope=SCOPE)
            self.assert_empty(store)
        assert_closed(self, self.parent)

    def test_size_cap_on_stat_and_stream(self):
        store = self.make_store(max_file=64)
        self.write("large.png", bytes(65))
        self.write("fits.png", bytes(64))

        def never(store_, stream, size):  # pragma: no cover - guard only
            raise AssertionError("an over-cap file must not be read")

        with patch.object(CaptureStore, "_read", never):
            with self.assertRaises(CaptureTooLarge):
                self.capture(store, "large.png")
        self.assert_empty(store)

        cap = self.capture(store, "fits.png")
        self.assertEqual(cap.size, 64)
        store.release(cap.token)

        growing = self.write("growing.png", bytes(40))
        original = CaptureStore._read

        def grow(store_, stream, size):
            data = original(store_, stream, 16)
            if data and not getattr(grow, "done", False):
                grow.done = True
                with open(growing, "ab") as handle:
                    handle.write(bytes(40))
            return data

        with patch.object(CaptureStore, "_read", grow):
            with self.assertRaises(CaptureTooLarge):
                self.capture(store, "growing.png")
        self.assert_empty(store)

    # -- admission -----------------------------------------------------------

    def test_count_byte_and_worker_admission_is_nonblocking(self):
        for name in ("a", "b", "c"):
            self.write(f"{name}.png", bytes(60))
        store = self.make_store(max_captures=2, max_bytes=1000)
        first = self.capture(store, "a.png")
        self.capture(store, "b.png")
        with self.assertRaises(CaptureCapacity):
            self.capture(store, "c.png")
        self.assertTrue(store.release(first.token))
        third = self.capture(store, "c.png")
        self.assertEqual(store.stats()["captures"], 2)
        store.invalidate_all()
        self.assertIsNone(store.get(third.token))
        self.assert_empty(store)

        bytes_store = self.make_store(max_bytes=100)
        kept = self.capture(bytes_store, "a.png")
        with self.assertRaises(CaptureCapacity):
            self.capture(bytes_store, "b.png")
        self.assertEqual(bytes_store.stats()["reserved_bytes"], 60)
        bytes_store.release(kept.token)
        self.capture(bytes_store, "b.png")

        workers = self.make_store(max_workers=1)
        entered, proceed = threading.Event(), threading.Event()
        original = CaptureStore._read
        results = {}

        def blocking(store_, stream, size):
            if not entered.is_set():
                entered.set()
                self.assertTrue(proceed.wait(WAIT))
            return original(store_, stream, size)

        def worker():
            try:
                results["capture"] = self.capture(workers, "a.png")
            except Exception as error:  # pragma: no cover - reported below
                results["error"] = error

        with patch.object(CaptureStore, "_read", blocking):
            thread = threading.Thread(target=worker)
            thread.start()
            try:
                self.assertTrue(entered.wait(WAIT))
                self.assertEqual(workers.stats()["workers"], 1)
                with self.assertRaises(CaptureCapacity):
                    self.capture(workers, "b.png")
            finally:
                proceed.set()
                thread.join(WAIT)
        self.assertFalse(thread.is_alive())
        self.assertNotIn("error", results)
        self.assertEqual(workers.stats()["workers"], 0)
        self.assertEqual(results["capture"].size, 60)
        self.capture(workers, "b.png")

    def _hold_worker(self, store, relative, **options):
        """Start a capture that holds a worker until ``proceed`` is set."""
        entered, proceed = threading.Event(), threading.Event()
        original = CaptureStore._read
        results = {}

        def blocking(store_, stream, size):
            if threading.current_thread() is thread and not entered.is_set():
                entered.set()
                self.assertTrue(proceed.wait(WAIT))
            return original(store_, stream, size)

        def worker():
            try:
                results["capture"] = self.capture(store, relative, **options)
            except Exception as error:  # pragma: no cover - reported below
                results["error"] = error

        patcher = patch.object(CaptureStore, "_read", blocking)
        patcher.start()
        self.addCleanup(patcher.stop)
        thread = threading.Thread(target=worker)
        thread.start()
        self.assertTrue(entered.wait(WAIT))

        def finish():
            proceed.set()
            thread.join(WAIT)
            self.assertFalse(thread.is_alive())
            self.assertNotIn("error", results)
            return results["capture"]

        self.addCleanup(proceed.set)
        return finish

    def test_background_capture_never_takes_the_last_worker(self):
        for name in ("a", "b", "c"):
            self.write(f"{name}.png", bytes(60))
        store = self.make_store(max_workers=2)
        finish = self._hold_worker(store, "a.png", background=True)
        # The second background capture waits, then fails: one worker is free
        # but it is the interactive request's.
        started = time.monotonic()
        with patch("utils.media_capture.BACKGROUND_WAIT_SECONDS", 0.05):
            with self.assertRaises(CaptureCapacity):
                self.capture(store, "b.png", background=True)
        self.assertGreaterEqual(time.monotonic() - started, 0.04)
        # An interactive capture (Tag dialog, Apply) still gets the last one.
        interactive = self.capture(store, "c.png")
        self.assertEqual(interactive.size, 60)
        finish()

    def test_background_capture_waits_for_a_worker_to_free(self):
        for name in ("a", "b"):
            self.write(f"{name}.png", bytes(60))
        store = self.make_store(max_workers=2)
        finish = self._hold_worker(store, "a.png", background=True)
        waited = {}

        def background():
            waited["capture"] = self.capture(store, "b.png", background=True)

        thread = threading.Thread(target=background)
        with patch("utils.media_capture.BACKGROUND_WAIT_SECONDS", WAIT):
            thread.start()
            finish()
            thread.join(WAIT)
        self.assertFalse(thread.is_alive())
        self.assertEqual(waited["capture"].size, 60)

    def test_single_worker_store_admits_no_background_capture(self):
        self.write("a.png", bytes(60))
        store = self.make_store(max_workers=1)
        with patch("utils.media_capture.BACKGROUND_WAIT_SECONDS", 0.01):
            with self.assertRaises(CaptureCapacity):
                self.capture(store, "a.png", background=True)
        self.assertEqual(self.capture(store, "a.png").size, 60)

    def test_background_capture_leaves_the_last_capture_slot(self):
        for name in ("a", "b", "c"):
            self.write(f"{name}.png", bytes(60))
        store = self.make_store(max_captures=2)
        self.capture(store, "a.png")
        with patch("utils.media_capture.BACKGROUND_WAIT_SECONDS", 0.01):
            with self.assertRaises(CaptureCapacity):
                self.capture(store, "b.png", background=True)
        self.assertEqual(self.capture(store, "c.png").size, 60)

    # -- lifetime ------------------------------------------------------------

    def test_idle_and_hard_expiry_and_in_use_retention(self):
        self.write("a.png", png_bytes(1, 1))
        self.write("b.png", png_bytes(1, 1))
        store = self.make_store(idle=900, hard=1800)
        cap = self.capture(store, "a.png")

        self.clock.advance(899)
        self.assertIs(store.get(cap.token), cap)  # touches idle
        self.assertAlmostEqual(store.expires_in(cap.token), 900)
        self.clock.advance(899)
        self.assertIs(store.get(cap.token), cap)
        self.assertAlmostEqual(store.expires_in(cap.token), 2)
        self.clock.advance(2)
        self.assertIsNone(store.get(cap.token))  # hard limit despite touches
        self.assertIsNone(store.expires_in(cap.token))
        self.assertFalse(Path(cap.path).exists())
        self.assert_empty(store)

        idle = self.capture(store, "a.png")
        self.clock.advance(900)
        store.sweep()
        self.assertIsNone(store.get(idle.token))
        self.assert_empty(store)

        used = self.capture(store, "b.png")
        self.assertIs(store.acquire(used.token), used)
        self.assertIs(store.acquire(used.token), used)
        self.clock.advance(901)
        store.sweep()
        self.assertIs(store.get(used.token), used)  # in use is not idle
        self.clock.advance(900)
        store.sweep()
        self.assertIsNone(store.get(used.token))  # expired for new requests
        self.assertIsNone(store.acquire(used.token))
        self.assertTrue(Path(used.path).exists())  # retained while in use
        self.assertEqual(store.stats()["captures"], 1)
        store.release_use(used.token)
        self.assertTrue(Path(used.path).exists())
        store.release_use(used.token)
        self.assertFalse(Path(used.path).exists())
        self.assert_empty(store)
        store.release_use(used.token)  # unknown token: no error, nothing freed
        self.assertIsNone(store.acquire("unknown"))

    def test_release_while_in_use_keeps_file_until_last_use_ends(self):
        self.write("a.png", png_bytes(1, 1))
        store = self.make_store()
        cap = self.capture(store, "a.png")
        self.assertIs(store.acquire(cap.token), cap)
        self.assertTrue(store.release(cap.token))
        self.assertIsNone(store.get(cap.token))
        self.assertTrue(Path(cap.path).exists())
        store.invalidate_all()
        self.assertTrue(Path(cap.path).exists())
        store.release_use(cap.token)
        self.assertFalse(Path(cap.path).exists())
        self.assert_empty(store)

    def test_invalidate_all_and_close_remove_owned_directory_only(self):
        self.write("a.png", png_bytes(1, 1))
        self.write("b.png", png_bytes(1, 1))
        sibling = self.parent / "captures" / "sibling.txt"
        sibling.parent.mkdir()
        sibling.write_bytes(b"keep")
        store = CaptureStore(self.parent / "captures", clock=self.clock)
        try:
            a = self.capture(store, "a.png")
            b = self.capture(store, "b.png")
            self.assertEqual(Path(store.directory).parent, sibling.parent)
            store.invalidate_all()
            self.assertIsNone(store.get(a.token))
            self.assertIsNone(store.get(b.token))
            self.assert_empty(store)
            self.assertTrue(Path(store.directory).is_dir())
            c = self.capture(store, "a.png")
            self.assertTrue(Path(c.path).exists())
        finally:
            store.close()
        self.assertFalse(Path(store.directory).exists())
        self.assertEqual(sibling.read_bytes(), b"keep")
        self.assertTrue((self.root / "a.png").exists())
        store.close()  # idempotent
        with self.assertRaises(RuntimeError):
            self.capture(store, "a.png")
        self.assertIsNone(store.get(c.token))
        assert_closed(self, self.parent)

    def test_windows_handles_closed_before_unlink(self):
        source_bytes = png_bytes(1, 1, b"\x33" * 40)
        source = self.write("a.png", source_bytes)
        store = self.make_store()
        original_unlink = os.unlink
        unlinked = []

        def checked_unlink(path, *args, **kwargs):
            self.assertEqual(open_handles_under(path), [])
            self.assertTrue(store._lock.acquire(blocking=False))  # no I/O under lock
            store._lock.release()
            unlinked.append(os.fspath(path))
            return original_unlink(path, *args, **kwargs)

        with patch("os.unlink", checked_unlink):
            cap = self.capture(store, "a.png")
            self.assertEqual(open_handles_under(cap.path), [])
            self.assertEqual(open_handles_under(source), [])
            store.release(cap.token)
            original = CaptureStore._read

            def appending(store_, stream, size):
                data = original(store_, stream, 16)
                if data and not getattr(appending, "done", False):
                    appending.done = True
                    with open(source, "ab") as handle:
                        handle.write(b"\x44")
                return data

            with patch.object(CaptureStore, "_read", appending):
                with self.assertRaises(CaptureChanged):
                    self.capture(store, "a.png")
            self.assertEqual(open_handles_under(source), [])
        self.assertEqual(len(unlinked), 2)
        self.assertTrue(all(p.startswith(str(store.directory)) for p in unlinked))
        self.assert_empty(store)
        assert_closed(self, self.parent)

    # -- inference readiness -------------------------------------------------

    def test_dimension_limits_mark_generate_unsupported(self):
        cases = {
            "edge.png": (png_bytes(4096, 4096), True),
            "wide.png": (png_bytes(4097, 1), False),
            "tall.png": (png_bytes(1, 4097), False),
        }
        for name, (data, _) in cases.items():
            self.write(name, data)
        self.write("garbage.jpg", b"\xff\xd8not really a jpeg")
        store = self.make_store()

        for name, (data, supported) in cases.items():
            with self.subTest(name=name):
                cap = self.capture(store, name)
                self.assertEqual(cap.inferable, supported)
                self.assertEqual((cap.width, cap.height) != (None, None), True)
                if supported:
                    self.assertIsNone(cap.inferable_reason)
                else:
                    self.assertIn("4096", cap.inferable_reason)
                store.release(cap.token)
        self.assertEqual(
            (capture_module.MAX_INFERENCE_SIDE, capture_module.MAX_INFERENCE_PIXELS),
            (4096, 16_777_216),
        )
        cap = self.capture(store, "garbage.jpg")
        self.assertEqual((cap.media_type, cap.width, cap.height), ("image", None, None))
        self.assertFalse(cap.inferable)
        self.assertIn("header", cap.inferable_reason)
        self.assertEqual(
            cap.sha256, hashlib.sha256(b"\xff\xd8not really a jpeg").hexdigest()
        )

    def test_metadata_over_inspection_budgets_does_not_block_generate(self):
        # B2 carry-over (A M1): inferable is decided from the image header
        # alone, so prompt metadata over the inspection budgets (here a 3 MiB
        # text chunk) leaves a decodable image inferable.
        base = png_bytes(64, 32)
        huge = _chunk(b"tEXt", b"prompt\x00" + b"x" * (3 * 1024 * 1024))
        self.write("big-metadata.png", base[:-12] + huge + base[-12:])
        store = self.make_store()

        cap = self.capture(store, "big-metadata.png")

        self.assertEqual((cap.width, cap.height), (64, 32))
        self.assertTrue(cap.inferable)
        self.assertIsNone(cap.inferable_reason)
        assert_closed(self, self.parent)

    def test_decompression_bomb_header_is_not_inferable(self):
        self.write("bomb.png", png_bytes(20_000, 20_000))
        store = self.make_store()

        cap = self.capture(store, "bomb.png")

        self.assertFalse(cap.inferable)
        self.assertIn("4096", cap.inferable_reason)
        assert_closed(self, self.parent)

    def test_header_dimensions_reads_only_the_header(self):
        path = self.write("dims.png", png_bytes(7, 5))
        self.assertEqual(capture_module.header_dimensions(str(path)), (7, 5))
        assert_closed(self, self.parent)
        with self.assertRaises(capture_module.HeaderUnreadable):
            capture_module.header_dimensions(
                str(self.write("bad.png", b"\x89PNG\r\n\x1a\nnope"))
            )
        assert_closed(self, self.parent)

    def test_header_dimensions_leaves_the_warning_filters_alone(self):
        # The capture workers run in the ComfyUI process: touching the
        # process-global warnings filters from them races (catch_warnings is
        # not thread-safe) and can leave an "error" filter installed for every
        # later Pillow open in the host.
        from PIL import Image

        path = str(self.write("dims.png", png_bytes(7, 5)))
        before = list(warnings.filters)
        seen_during = []
        real_open = Image.open

        def recording_open(*args, **kwargs):
            seen_during.append(list(warnings.filters))
            return real_open(*args, **kwargs)

        with patch.object(Image, "open", recording_open):
            self.assertEqual(capture_module.header_dimensions(path), (7, 5))
        self.assertEqual(seen_during, [before])
        self.assertEqual(list(warnings.filters), before)

        barrier = threading.Barrier(8)
        errors = []

        def worker():
            try:
                barrier.wait()
                for _ in range(50):
                    capture_module.header_dimensions(path)
            except Exception as exc:  # pragma: no cover - reported below
                errors.append(exc)

        threads = [threading.Thread(target=worker) for _ in range(8)]
        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
        self.assertEqual(errors, [])
        self.assertEqual(list(warnings.filters), before)
        assert_closed(self, self.parent)

    def test_header_over_the_warning_size_is_refused_by_the_pixel_cap(self):
        # 100 MP: past Pillow's warning size, under its error size. The
        # header is read, then the inference pixel cap refuses it.
        from PIL import Image

        path = str(self.write("big.png", png_bytes(10_000, 10_000)))
        with self.assertWarns(Image.DecompressionBombWarning):
            self.assertEqual(capture_module.header_dimensions(path), (10_000, 10_000))
        store = self.make_store()
        with self.assertWarns(Image.DecompressionBombWarning):
            cap = self.capture(store, "big.png")
        self.assertFalse(cap.inferable)
        self.assertIn("4096", cap.inferable_reason)
        assert_closed(self, self.parent)

    def test_video_capture_is_not_inferable(self):
        data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32)
        self.write("clip.mp4", data)
        store = self.make_store()

        def never(path):  # pragma: no cover - guard only
            raise AssertionError("video captures are never probed")

        with patch.object(capture_module, "header_dimensions", never):
            cap = self.capture(store, "clip.mp4")
        self.assertEqual((cap.media_type, cap.width, cap.height), ("video", None, None))
        self.assertFalse(cap.inferable)
        self.assertEqual(
            cap.inferable_reason,
            "Auto-tagging is available for supported still images only.",
        )
        self.assertEqual(Path(cap.path).suffix, ".mp4")
        self.assertEqual(cap.sha256, hashlib.sha256(data).hexdigest())

    # -- verification without a copy -------------------------------------

    def test_verify_hashes_current_file_without_retaining_a_copy(self):
        source_bytes = png_bytes(1, 1, b"\x55" * 30)
        source = self.write("a.png", source_bytes)
        store = self.make_store(max_file=128, max_workers=1)

        facts = store.verify(self.root, "a.png")

        self.assertEqual(facts.sha256, hashlib.sha256(source_bytes).hexdigest())
        self.assertEqual(facts.size, len(source_bytes))
        self.assertEqual(facts.mtime_ns, os.stat(source).st_mtime_ns)
        self.assertEqual(
            facts.locator,
            os.path.normcase(os.path.join(os.path.realpath(self.root), "a.png")),
        )
        self.assert_empty(store)
        self.write("large.png", bytes(129))
        with self.assertRaises(CaptureTooLarge):
            store.verify(self.root, "large.png")
        with self.assertRaises(CaptureDenied):
            store.verify(self.root, "../a.png")
        self.assert_empty(store)

    # -- arguments and failure paths ---------------------------------------

    def test_constructor_and_arguments_are_validated(self):
        for options in (
            {"max_file": 0},
            {"max_captures": True},
            {"max_bytes": -1},
            {"max_workers": "2"},
            {"idle": 0},
            {"hard": -5},
        ):
            with self.subTest(options=options):
                with self.assertRaises(ValueError):
                    CaptureStore(self.parent / "captures", **options)
        self.write("a.png", png_bytes(1, 1))
        store = self.make_store()
        with self.assertRaises(ValueError):
            store.capture(self.root, "a.png", epoch="", scope=SCOPE)
        with self.assertRaises(ValueError):
            store.capture(self.root, "a.png", epoch=EPOCH, scope=None)
        with self.assertRaises(CaptureDenied):
            self.capture(store, root_index=-1)
        with self.assertRaises(CaptureDenied):
            store.capture(
                self.parent / "missing-root", "a.png", epoch=EPOCH, scope=SCOPE
            )
        self.assert_empty(store)

        default = CaptureStore(clock=self.clock)
        try:
            self.assertEqual(
                Path(default.directory).parent, Path(tempfile.gettempdir())
            )
            self.assertTrue(Path(default.directory).is_dir())
        finally:
            default.close()
        self.assertFalse(Path(default.directory).exists())

    def test_directory_target_and_unreadable_file(self):
        (self.root / "dir.png").mkdir()
        store = self.make_store()
        with self.assertRaises(CaptureDenied):
            self.capture(store, "dir.png")
        self.assert_empty(store)
        if os.name != "posix" or os.geteuid() == 0:
            return  # read permission cannot be withdrawn here
        locked = self.write("locked.png", png_bytes(1, 1))
        os.chmod(locked, 0)
        try:
            with self.assertRaises(CaptureChanged):
                self.capture(store, "locked.png")
        finally:
            os.chmod(locked, 0o600)
        self.assert_empty(store)

    def test_removal_failures_are_logged_and_never_touch_other_files(self):
        self.write("a.png", png_bytes(1, 1))
        store = CaptureStore(self.parent / "captures", clock=self.clock)
        stray = Path(store.directory) / "stray.bin"
        try:
            cap = self.capture(store, "a.png")
            self.assertIs(store.acquire(cap.token), cap)
            self.assertAlmostEqual(store.expires_in(cap.token), 1800)  # in use
            store.release_use(cap.token)
            original_unlink = os.unlink

            def failing(path, *args, **kwargs):
                if os.fspath(path) == cap.path:
                    raise PermissionError("busy")
                return original_unlink(path, *args, **kwargs)

            with patch("os.unlink", failing):
                with self.assertLogs("utils.media_capture", "WARNING") as logs:
                    self.assertTrue(store.release(cap.token))
            self.assertIn("could not be removed", logs.output[0])
            self.assertIsNone(store.get(cap.token))
            self.assertTrue(Path(cap.path).exists())
            os.unlink(cap.path)
            stray.write_bytes(b"not ours")
        finally:
            with self.assertLogs("utils.media_capture", "WARNING") as logs:
                store.close()
        self.assertIn("could not be removed", logs.output[0])
        self.assertEqual(stray.read_bytes(), b"not ours")
        stray.unlink()
        os.rmdir(store.directory)
        assert_closed(self, self.parent)

    # -- platform identity checks -------------------------------------------

    def test_path_and_handle_stats_are_never_compared_across_apis(self):
        """Windows (Python 3.12+) reports the change time from stat but the
        creation time from fstat, and may encode dev/ino differently; only
        size and mtime_ns may be compared across the two APIs."""
        source_bytes = png_bytes(2, 2, b"\x66" * 40)
        self.write("a.png", source_bytes)
        store = self.make_store()
        original_fstat = os.fstat

        def shifted(fd):
            st = original_fstat(fd)
            return os.stat_result(
                (
                    st.st_mode,
                    st.st_ino + 7,
                    st.st_dev + 7,
                    st.st_nlink,
                    st.st_uid,
                    st.st_gid,
                    st.st_size,
                    st.st_atime,
                    st.st_mtime,
                    st.st_ctime + 1,
                ),
                {
                    "st_atime_ns": st.st_atime_ns,
                    "st_mtime_ns": st.st_mtime_ns,
                    "st_ctime_ns": st.st_ctime_ns + 1_000_000_000,
                },
            )

        with patch.object(capture_module.os, "fstat", shifted):
            cap = self.capture(store, "a.png")
            facts = store.verify(self.root, "a.png")
        self.assertEqual(cap.sha256, hashlib.sha256(source_bytes).hexdigest())
        self.assertEqual(facts.sha256, cap.sha256)
        store.release(cap.token)
        self.assert_empty(store)

    def test_change_time_moving_between_fstat_reads_is_detected(self):
        """The Windows 3.13 / POSIX signal, forced on every platform: the
        descriptor's change time differs after the copy, size and mtime equal."""
        self.write("a.png", png_bytes(1, 1, b"\x77" * 40))
        store = self.make_store()
        original_fstat = os.fstat
        calls = []

        def drifting(fd):
            st = original_fstat(fd)
            calls.append(fd)
            if len(calls) < 2:
                return st
            return os.stat_result(
                tuple(st)[:10],
                {
                    "st_atime_ns": st.st_atime_ns,
                    "st_mtime_ns": st.st_mtime_ns,
                    "st_ctime_ns": st.st_ctime_ns + 1,
                },
            )

        with patch.object(capture_module.os, "fstat", drifting):
            with self.assertRaises(CaptureChanged):
                self.capture(store, "a.png")
        self.assertEqual(len(calls), 2)
        self.assert_empty(store)

    def test_fresh_file_captures_twice_on_every_platform(self):
        data = png_bytes(1, 1, os.urandom(3 * capture_module.CHUNK_SIZE))
        self.write("fresh.png", data)
        store = self.make_store()
        first = self.capture(store, "fresh.png")
        second = self.capture(store, "fresh.png")
        self.assertEqual(first.sha256, hashlib.sha256(data).hexdigest())
        self.assertEqual(second.sha256, first.sha256)
        self.assertNotEqual(first.token, second.token)
        self.assertEqual(store.stats()["captures"], 2)

    def test_error_messages_never_name_server_paths(self):
        self.write("a.png", png_bytes(1, 1))
        self.write("dir.png/x", b"")  # dir.png is a directory
        store = self.make_store(max_file=64)
        attempts = [
            lambda: self.capture(store, "missing.png"),
            lambda: self.capture(store, "a.png"),  # over the 64-byte cap
            lambda: self.capture(store, "dir.png"),
            lambda: self.capture(store, "../a.png"),
            lambda: self.capture(store, "notes.txt"),
            lambda: store.verify(self.root, "missing.png"),
        ]
        if os.name == "posix" and os.geteuid() != 0:
            locked = self.write("locked.png", bytes(10))
            os.chmod(locked, 0)
            self.addCleanup(os.chmod, locked, 0o600)
            attempts.append(lambda: self.capture(store, "locked.png"))
        for attempt in attempts:
            with self.assertRaises(capture_module.CaptureError) as caught:
                attempt()
            message = str(caught.exception)
            self.assertTrue(message)
            self.assertNotIn(str(self.root), message)
            self.assertNotIn(str(self.parent), message)
            self.assertNotIn(os.sep, message)
        self.assert_empty(store)

    def test_close_during_a_capture_removes_the_directory_once_it_finishes(self):
        self.write("a.png", png_bytes(1, 1, bytes(64)))
        store = CaptureStore(self.parent / "captures", clock=self.clock)
        original = CaptureStore._read

        def closing(store_, stream, size):
            data = original(store_, stream, 16)
            if data and not getattr(closing, "done", False):
                closing.done = True
                store.close()
            return data

        with patch.object(CaptureStore, "_read", closing):
            with self.assertRaises(RuntimeError):
                self.capture(store, "a.png")
        self.assertFalse(Path(store.directory).exists())
        self.assertEqual(store.stats()["reserved_bytes"], 0)
        store.close()
        assert_closed(self, self.parent)

    # -- constants ---------------------------------------------------------

    def test_defaults_match_the_plan(self):
        store = self.make_store()
        self.assertEqual(
            (
                store.max_file,
                store.max_captures,
                store.max_bytes,
                store.max_workers,
                store.idle,
                store.hard,
            ),
            (256 * 2**20, 8, 512 * 2**20, 2, 900, 1800),
        )
        from py.api.inspection_service import SERVABLE_EXTENSIONS

        self.assertEqual(CAPTURE_EXTENSIONS, SERVABLE_EXTENSIONS)

    def test_exceptions_carry_codes(self):
        self.assertEqual(CaptureCapacity.code, "review-capacity")
        self.assertEqual(CaptureChanged.code, "file-changed")
        self.assertEqual(CaptureTooLarge.code, "unsupported-media")
        self.assertEqual(CaptureDenied.code, "source-changed")
        for cls in (CaptureCapacity, CaptureChanged, CaptureTooLarge, CaptureDenied):
            self.assertTrue(issubclass(cls, capture_module.CaptureError))


if __name__ == "__main__":
    unittest.main()