diff --git a/py/api/images.py b/py/api/images.py index 8f42289..5eae351 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -11,6 +11,7 @@ import re import threading import time as _time import urllib.parse +from collections.abc import Mapping from pathlib import Path from aiohttp import web @@ -22,17 +23,19 @@ from .prompts import safe_error_message as _safe_error try: from ...utils.parallel import map_parallel + from ...utils.inspection_response import encode_inspection_response from ...utils.prompt_inspector import ( MetadataReadError, - inspect_prompt_roles, - read_inspection_metadata, + inspect_prompt_evidence, + read_inspection_file, ) except ImportError: from utils.parallel import map_parallel + from utils.inspection_response import encode_inspection_response from utils.prompt_inspector import ( MetadataReadError, - inspect_prompt_roles, - read_inspection_metadata, + inspect_prompt_evidence, + read_inspection_file, ) # Only these file types are ever served by the image routes, regardless of @@ -375,7 +378,9 @@ def _iter_media_files(output_path, max_files): yield root_path / name -def _output_image_entry(media_path, output_path, root_index, public_root=None): +def _output_image_entry( + media_path, output_path, root_index, public_root=None, *, root_scope=None +): """Build one gallery entry for *media_path* (blocking stat). ``path`` is relative to its root (what delete/autotag resolve against) and @@ -396,15 +401,20 @@ def _output_image_entry(media_path, output_path, root_index, public_root=None): if thumbnail_url is not None: thumbnail_url += f"&root={root_index}" encoded_path = urllib.parse.quote(rel_path.as_posix(), safe="/") + url = f"/prompt_manager/images/serve/{encoded_path}?root={root_index}" + if root_scope is not None: + url += f"&root_scope={root_scope}" + if thumbnail_url is not None: + thumbnail_url += f"&root_scope={root_scope}" - return { + entry = { "id": hashlib.sha1(str(media_path).encode("utf-8")).hexdigest()[:16], "filename": media_path.name, "path": rel_path.as_posix(), "relative_path": str(rel_path), "root_dir": public_root, "root_index": root_index, - "url": f"/prompt_manager/images/serve/{encoded_path}?root={root_index}", + "url": url, "thumbnail_url": thumbnail_url, "size": stat.st_size, "modified_time": stat.st_mtime, @@ -412,6 +422,9 @@ def _output_image_entry(media_path, output_path, root_index, public_root=None): "media_type": "video" if is_video else "image", "is_video": is_video, } + if root_scope is not None: + entry["root_scope"] = root_scope + return entry def _in_subfolder(rel_dir, subfolder): @@ -421,6 +434,48 @@ def _in_subfolder(rel_dir, subfolder): class ImageRoutesMixin: """Mixin providing image and gallery-related API endpoints.""" + def _gallery_root_scope(self, root: Path, index: int) -> str: + payload = json.dumps( + ["prompt-manager-gallery-root-v1", index, str(_canonical(root))], + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hmac.new(_WORKSPACE_CURSOR_SECRET, payload, hashlib.sha256).hexdigest() + + def _gallery_scope_error( + self, query: Mapping, roots: list[Path] + ) -> web.Response | None: + if "root_scope" not in query: + return None + scope, root = query.get("root_scope"), query.get("root") + if ( + not isinstance(scope, str) + or re.fullmatch(r"[0-9a-f]{64}", scope) is None + or not isinstance(root, str) + or re.fullmatch(r"[0-9]+", root) is None + ): + return bad_request("Invalid gallery root scope") + try: + index = int(root) + except ValueError: + return bad_request("Invalid gallery root scope") + if index >= len(roots) or not hmac.compare_digest( + scope, self._gallery_root_scope(roots[index], index) + ): + return web.json_response( + { + "success": False, + "code": "source-changed", + "error": ( + "The source changed. Refresh the image collection before " + "inspecting this file." + ), + }, + status=409, + headers={"Cache-Control": "no-store"}, + ) + return None + def _register_image_routes(self, routes): @routes.get("/prompt_manager/prompts/{prompt_id}/images") async def get_prompt_images_route(request): @@ -910,15 +965,17 @@ class ImageRoutesMixin: async def serve_output_image(self, request): """Serve a media file from the ComfyUI output folder(s).""" try: + roots = list(self._get_all_output_dirs()) + scope_error = self._gallery_scope_error(request.query, roots) + if scope_error is not None: + return scope_error filepath = request.match_info["filepath"] if _has_traversal(filepath): return _forbidden() if Path(filepath).suffix.lower() not in SERVABLE_EXTENSIONS: return _forbidden("Only media files can be served") - output_dirs = self._select_root( - list(self._get_all_output_dirs()), request.query.get("root") - ) + output_dirs = self._select_root(roots, request.query.get("root")) if not output_dirs: return _forbidden("No allowed image directories configured") @@ -1645,12 +1702,18 @@ class ImageRoutesMixin: async def inspect_image(self, request): """Read prompt evidence for one explicitly rooted physical media file.""" roots = list(self._get_all_output_dirs()) + scope_error = self._gallery_scope_error(request.query, roots) + if scope_error is not None: + return scope_error try: index = _workspace_root(request.query, roots) if index is None: raise ValueError("Gallery root is required") except ValueError as error: return bad_request(str(error)) + root_scope = request.query.get("root_scope") or self._gallery_root_scope( + roots[index], index + ) relative = request.match_info.get("image_path", "") if ( _has_traversal(relative) @@ -1668,32 +1731,82 @@ class ImageRoutesMixin: ) try: - def inspect() -> dict: + def inspect() -> bytes | None: + before = path.stat() try: - metadata = read_inspection_metadata(path) + inspection = read_inspection_file(path) except MetadataReadError: - return { - role: {"status": "unresolved", "text": None, "source": None} - for role in ("positive", "negative") + inspection = None + try: + after = path.stat() + except FileNotFoundError: + return None + fields = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns") + if any(getattr(before, key) != getattr(after, key) for key in fields): + return None + saved, unavailable = None, False + try: + saved = self.db.get_image_prompt_info(str(path), False) + except Exception: + unavailable = True + self.logger.warning("Saved inspection display could not be read") + if inspection is None: + evidence = { + "prompts": { + role: {"status": "unresolved", "text": None, "source": None} + for role in ("positive", "negative") + }, + "resolution": { + "state": "unresolved", + "sampler_id": None, + "reason": "Media metadata could not be read", + }, } - saved = ( - self.db.get_image_prompt_info(str(path), False) - if not metadata - else None - ) - return inspect_prompt_roles(metadata, saved) - - prompts = await self._run_in_executor(inspect) - return web.json_response( - { - "success": True, - "image": { + else: + evidence = inspect_prompt_evidence(inspection.metadata, saved) + return encode_inspection_response( + image={ "root_index": index, "path": Path(relative).as_posix(), "filename": path.name, }, - "prompts": prompts, - } + root_scope=root_scope, + file={ + "observation": "current-file", + "source_label": _public_path(roots[index]), + "media_type": ( + "video" + if path.suffix.lower() in VIDEO_EXTENSIONS + else "image" + ), + "extension": path.suffix.lower(), + "size": before.st_size, + "modified_time": before.st_mtime, + "format": None, + "width": None, + "height": None, + }, + prompts=evidence["prompts"], + resolution=evidence["resolution"], + inspection=inspection, + library=self._inspection_library(saved, unavailable=unavailable), + ) + + body = await self._run_in_executor(inspect) + if body is None: + return web.json_response( + { + "success": False, + "error": "The file changed while being inspected. Try again.", + }, + status=409, + headers={"Cache-Control": "no-store"}, + ) + return web.Response( + body=body, + content_type="application/json", + charset="utf-8", + headers={"Cache-Control": "no-store"}, ) except Exception: self.logger.exception("Prompt inspection failed") @@ -1701,6 +1814,40 @@ class ImageRoutesMixin: {"success": False, "error": "Prompt inspection failed"}, status=500 ) + def _inspection_library(self, saved: dict | None, *, unavailable: bool) -> dict: + result = { + "status": "unavailable" if unavailable else "none", + "scope": "saved-library", + "positive_id": None, + "pair_id": None, + "title": None, + "category": None, + "tags": [], + "metadata_scope": None, + } + if not isinstance(saved, dict) or not saved: + return result + pair = saved.get("pair") + if isinstance(pair, dict): + result.update( + status="linked", + positive_id=pair.get("positive_prompt_id"), + pair_id=pair.get("id"), + title=pair.get("title"), + category=pair.get("category"), + tags=pair.get("tags"), + metadata_scope="pair", + ) + else: + result.update( + status="linked", + positive_id=saved.get("prompt_id"), + category=saved.get("category"), + tags=saved.get("tags"), + metadata_scope="positive", + ) + return result + async def get_image_prompt(self, request): """Get prompt information for a specific image path.""" try: @@ -1799,7 +1946,7 @@ class ImageRoutesMixin: ] def _format_workspace_page(self, page: list[tuple]) -> list[dict]: - images = self._format_output_page(page) + images = self._format_output_page(page, scoped=True) paths = { (entry[3], entry[0].relative_to(entry[2]).as_posix()): entry[0] for entry in page @@ -1818,7 +1965,7 @@ class ImageRoutesMixin: continue return images - def _format_output_page(self, page): + def _format_output_page(self, page, *, scoped=False): entries = [] public_roots = {} for media_path, _mtime, output_path, root_index in page: @@ -1827,7 +1974,15 @@ class ImageRoutesMixin: try: entries.append( _output_image_entry( - media_path, output_path, root_index, public_roots[output_path] + media_path, + output_path, + root_index, + public_roots[output_path], + root_scope=( + self._gallery_root_scope(output_path, root_index) + if scoped + else None + ), ) ) except Exception as e: diff --git a/tests/browser/full-inspector.spec.mjs b/tests/browser/full-inspector.spec.mjs new file mode 100644 index 0000000..aae56cd --- /dev/null +++ b/tests/browser/full-inspector.spec.mjs @@ -0,0 +1,50 @@ +import { test, expect } from '../../tools/vendor/node_modules/@playwright/test/index.mjs'; + +test('inspection wire preserves exact rooted envelope', async ({ request }) => { + for (const data of [ + { full_inspector: 'true' }, { full_inspector: 1 }, + { full_inspector: true, workspace_visual: true }, + { full_inspector: true, prompt_inspector: true }, + { full_inspector: true, native_library: true }, + { full_inspector: true, extra_prompts: 1 }, + { full_inspector: true, extra_variants: 1 }, + { full_inspector: true, extra_images: true }, + { full_inspector: true, include_video: 'true' }, + ]) expect((await request.post('/__test__/reset', { data })).status()).toBe(400); + await expect(await request.post('/__test__/reset', { data: { full_inspector: true, include_video: true, extra_images: 2 } })).toBeOK(); + const collection = await (await request.get('/prompt_manager/images/output?mode=workspace&subfolder=inspector&limit=40')).json(); + expect(collection.total).toBe(11); + const complete = collection.images.find(image => image.root_index === 0 && image.path === 'inspector/complete.png'); + const collision = collection.images.find(image => image.root_index === 1 && image.path === complete.path); + expect(complete.root_scope).toMatch(/^[0-9a-f]{64}$/); + expect(collision.root_scope).not.toBe(complete.root_scope); + const inspect = async entry => { + const response = await request.get(`/prompt_manager/images/inspect/${entry.path.split('/').map(encodeURIComponent).join('/')}?root=${entry.root_index}&root_scope=${entry.root_scope}`); + await expect(response).toBeOK(); + expect(response.headers()['cache-control']).toBe('no-store'); + expect((await response.body()).length).toBeLessThanOrEqual(8 * 1024 * 1024); + return response.json(); + }; + const result = await inspect(complete); + expect(result.inspection_version).toBe(1); + expect(result.image).toEqual({ root_index: 0, path: complete.path, filename: 'complete.png' }); + expect(result.root_scope).toBe(complete.root_scope); + expect(result.prompts.positive.text).toContain('Inspector portrait 猫 😀'); + expect(result.prompts.negative.text).toBe('blurred edges, watermark'); + expect(result.metadata.prompt_graph.value['3'].inputs.seed).toBe('9007199254740993'); + expect(result.metadata.prompt_graph.value['3'].inputs.cfg).toBe('6.5'); + expect(result.metadata.record_formats).toEqual({ prompt: 'original-text', workflow: 'original-text' }); + expect(result.file).toMatchObject({ observation: 'current-file', format: 'PNG', width: 720, height: 480 }); + expect(result.library).toMatchObject({ status: 'linked', metadata_scope: 'pair', title: 'Inspector recipe', category: 'Inspection', tags: ['fixture', 'geometric'] }); + expect((await inspect(collision)).prompts.positive.text).toBe('Input collision'); + const gallery = await (await request.get(`/prompt_manager/library/${result.library.positive_id}/images?pair_id=${result.library.pair_id}`)).json(); + expect(gallery.images[0].root_scope).toBe(complete.root_scope); + const library = await (await request.get('/prompt_manager/library')).json(); + expect(library.items.find(item => item.id === result.library.positive_id).previews.some(entry => entry.root_scope === complete.root_scope)).toBe(true); + const workflow = await inspect(collection.images.find(image => image.path === 'inspector/workflow-only.png')); + expect(workflow.prompts.positive.source).toBe('saved-association'); + const limited = await inspect(collection.images.find(image => image.path === 'inspector/control-heavy.png')); + expect(limited.metadata).toMatchObject({ status: 'limited', omitted_reason: 'response-size-limit', raw: {}, parameters: null }); + expect(collection.images.some(image => image.path.endsWith('clip.webm'))).toBe(true); + await expect(await request.post('/__test__/reset')).toBeOK(); +}); diff --git a/tests/browser/server.py b/tests/browser/server.py index b2575af..419fa3a 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -15,6 +15,7 @@ from PIL import Image from PIL.PngImagePlugin import PngInfo from tests.browser.workspace_visual_fixture import seed_workspace_visual +from tests.browser.workspace_inspector_fixture import seed_workspace_inspector class TestModel: @@ -226,6 +227,7 @@ def create_app(root: Path) -> web.Application: async def reset(request: web.Request) -> web.Response: data = await request.json() if request.can_read_body else {} workspace_visual = data.get("workspace_visual", False) + full_inspector = data.get("full_inspector", False) extra = data.get("extra_images", 0) include_video = data.get("include_video", False) prompt_inspector = data.get("prompt_inspector", False) @@ -234,6 +236,17 @@ def create_app(root: Path) -> web.Application: extra_variants = data.get("extra_variants", 0) if ( not isinstance(workspace_visual, bool) + or not isinstance(full_inspector, bool) + or ( + full_inspector + and ( + workspace_visual + or prompt_inspector + or native_library + or extra_prompts != 0 + or extra_variants != 0 + ) + ) or ( workspace_visual and any( @@ -245,6 +258,7 @@ def create_app(root: Path) -> web.Application: "native_library", "extra_prompts", "extra_variants", + "full_inspector", ) ) ) @@ -265,6 +279,13 @@ def create_app(root: Path) -> web.Application: if workspace_visual: await api._run_in_executor(seed_workspace_visual, api, root) return web.json_response({"success": True}) + if full_inspector: + await api._run_in_executor( + lambda: seed_workspace_inspector( + api, root, extra_images=extra, include_video=include_video + ) + ) + return web.json_response({"success": True}) await api._run_in_executor( seed, extra, diff --git a/tests/browser/workspace_inspector_fixture.py b/tests/browser/workspace_inspector_fixture.py new file mode 100644 index 0000000..cfd8776 --- /dev/null +++ b/tests/browser/workspace_inspector_fixture.py @@ -0,0 +1,138 @@ +"""Owned, deterministic physical-file evidence for the full inspector workflows.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import shutil +from typing import TYPE_CHECKING + +from PIL import Image +from PIL.PngImagePlugin import PngInfo + +from tests.browser.workspace_visual_fixture import seed_workspace_visual + +if TYPE_CHECKING: + from py.api import PromptManagerAPI + + +INSPECTOR_POSITIVE = ( + "Inspector portrait 猫 😀\n" + "Full prompt context, fine geometry and light. " * 40 +) +INSPECTOR_NEGATIVE = "blurred edges, watermark" + + +def _graph(positive: str, negative: str) -> dict: + return { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}}, + "3": { + "class_type": "KSampler", + "inputs": { + "positive": ["1", 0], + "negative": ["2", 0], + "model": ["4", 0], + "seed": 9007199254740993, + "steps": 28, + "cfg": 6.5, + "sampler_name": "dpmpp_2m", + "scheduler": "karras", + }, + }, + "4": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "fixture-model.safetensors"}, + }, + } + + +def _image(path: Path, index: int, metadata: dict) -> None: + info = PngInfo() + for key, value in metadata.items(): + text = ( + value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) + ) + info.add_itxt(key, text, zip=True) + path.parent.mkdir(parents=True, exist_ok=True) + with Image.new("RGB", (720, 480), (30 + index * 9, 74, 98)) as image: + image.save(path, pnginfo=info) + stamp = 1_780_000_000 + index + os.utime(path, (stamp, stamp)) + + +def seed_workspace_inspector( + api: PromptManagerAPI, + root: Path, + *, + extra_images: int = 0, + include_video: bool = False, +) -> None: + """Seed only disposable harness media and exact saved associations.""" + seed_workspace_visual(api, root) + complete = _graph(INSPECTOR_POSITIVE, INSPECTOR_NEGATIVE) + workflow = { + "nodes": [ + { + "id": 3, + "type": "KSampler", + "widgets_values": [ + 9007199254740993, + "fixed", + 28, + 6.5, + "dpmpp_2m", + "karras", + 1, + ], + } + ], + "links": [], + } + ambiguous = {**complete, "5": complete["3"]} + examples = ( + ("complete.png", {"prompt": complete, "workflow": workflow}), + ("empty.png", {}), + ("explicit-empty.png", {"prompt": _graph("Explicit empty negative", "")}), + ("ambiguous.png", {"prompt": ambiguous}), + ("workflow-only.png", {"workflow": workflow}), + ("malformed.png", {"prompt": "not-json", "workflow": "{broken"}), + ("control-heavy.png", {"parameters": "\x01" * 1_450_000}), + ) + for index, (name, metadata) in enumerate(examples): + _image(root / "output" / "inspector" / name, index, metadata) + _image( + root / "input" / "inspector" / "complete.png", + 8, + {"prompt": _graph("Input collision", "")}, + ) + pair = api.db.save_prompt_pair( + INSPECTOR_POSITIVE, + INSPECTOR_NEGATIVE, + title="Inspector recipe", + category="Inspection", + tags=["fixture", "geometric"], + ) + positive = api.db.get_prompt_pair(pair)["positive_prompt_id"] + image = api.db.link_image_to_prompt( + positive, str(root / "output" / "inspector" / "complete.png") + ) + api.db.link_image_to_pair(pair, image, {"source": "fixture"}) + api.db.link_image_to_prompt( + positive, str(root / "output" / "inspector" / "workflow-only.png") + ) + for number in range(extra_images): + _image( + root / "output" / "inspector" / f"extra-{number:03}.png", + 9 + number % 12, + {"prompt": complete}, + ) + if include_video: + shutil.copyfile( + Path(__file__).resolve().parents[1] / "fixtures" / "tiny.webm", + root / "output" / "inspector" / "clip.webm", + ) + os.utime( + root / "output" / "inspector" / "clip.webm", (1_780_000_100, 1_780_000_100) + ) + api.invalidate_gallery_cache() diff --git a/tests/test_api_images.py b/tests/test_api_images.py index c43cedd..5221c82 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -2067,3 +2067,128 @@ class TestGalleryRootIdentity(ThumbnailTestCase): ) self.assertEqual(response.status, 200) self.assertFalse((await response.json())["success"]) + + +class TestRootScopeServing(ImageAPITestCase): + async def scoped_request(self, path="image.png", root="0", scope="0" * 64): + params = {"root_scope": scope} + if root is not None: + params["root"] = root + return await self.client.get( + f"/prompt_manager/images/serve/{path}", params=params + ) + + async def assert_no_access(self, path, root="0", scope="0" * 64, status=409): + original_stat = Path.stat + + def stat(candidate, *args, **kwargs): + if candidate.name in ("image.png", "image_thumb.png"): + raise AssertionError("candidate stat before scope guard") + return original_stat(candidate, *args, **kwargs) + + with ( + patch.object(Path, "stat", stat), + patch.object( + self.api, "_file_response", side_effect=AssertionError("file response") + ), + patch.object( + self.api.db, + "get_image_prompt_info", + side_effect=AssertionError("saved lookup"), + ), + patch( + "py.api.images.read_inspection_file", + side_effect=AssertionError("metadata read"), + ), + patch( + "utils.prompt_inspector.subprocess.Popen", + side_effect=AssertionError("probe"), + ), + ): + response = await self.scoped_request(path, root, scope) + self.assertEqual(response.status, status) + result = await response.json() + self.assertFalse(result["success"]) + if status == 409: + self.assertEqual(result["code"], "source-changed") + self.assertEqual(response.headers.get("Cache-Control"), "no-store") + self.assertNotIn(str(self.comfy_dir), json.dumps(result)) + + async def test_root_scope_stale_original_and_thumbnail_fail_before_access(self): + make_png(self.output_dir / "image.png") + make_png(self.output_dir / "thumbnails" / "image_thumb.png") + for path in ("image.png", "thumbnails/image_thumb.png"): + await self.assert_no_access(path) + + async def test_root_scope_malformed_tokens_are_bad_requests(self): + make_png(self.output_dir / "image.png") + for scope in ("", "g" * 64, "a" * 63, "a" * 65, "a" * 4097, "A" * 64): + await self.assert_no_access("image.png", scope=scope, status=400) + for root in (None, "", "x", "-1", "0.0", "+0"): + await self.assert_no_access("image.png", root=root, status=400) + + async def test_root_scope_rotated_key_and_removed_index_require_refresh(self): + make_png(self.output_dir / "image.png") + listing = await self.client.get("/prompt_manager/images/output?mode=workspace") + descriptor = (await listing.json())["images"][0] + self.assertRegex(descriptor.get("root_scope", ""), r"^[0-9a-f]{64}$") + with patch.object( + images_module, "_WORKSPACE_CURSOR_SECRET", b"different-process-key" + ): + await self.assert_no_access("image.png", scope=descriptor["root_scope"]) + self.api._get_all_output_dirs = lambda: [] + await self.assert_no_access("image.png", scope=descriptor["root_scope"]) + + async def test_root_scope_keeps_containment_and_unscoped_access(self): + make_png(self.output_dir / "image.png") + listing = await self.client.get("/prompt_manager/images/output?mode=workspace") + descriptor = (await listing.json())["images"][0] + self.assertRegex(descriptor.get("root_scope", ""), r"^[0-9a-f]{64}$") + scope = descriptor["root_scope"] + self.assertEqual((await self.scoped_request(scope=scope)).status, 200) + outside = make_png(self.comfy_dir / "outside.png") + self._make_symlink(outside, self.output_dir / "escape.png") + for path in ("..%2Foutside.png", "escape.png"): + self.assertEqual((await self.scoped_request(path, scope=scope)).status, 403) + self.assertEqual( + (await self.client.get("/prompt_manager/images/serve/image.png")).status, + 200, + ) + + async def test_root_scope_equal_label_reorder_replacement_and_symlink_are_stale( + self, + ): + roots = [self.comfy_dir / name / "output" for name in ("a", "b", "c")] + for index, root in enumerate(roots): + make_png(root / "image.png", color=(index, 0, 0)) + make_png(root / "thumbnails" / "image_thumb.png", color=(index, 0, 0)) + # Both roots have the same display label under disposable public anchors. + with patch.object( + GalleryConfig, + "path_anchors", + classmethod(lambda cls: [str(root.parent) for root in roots]), + ): + self.api._get_all_output_dirs = lambda: roots[:2] + listing = await self.client.get( + "/prompt_manager/images/output?mode=workspace" + ) + descriptors = (await listing.json())["images"] + self.assertEqual({entry["root_dir"] for entry in descriptors}, {"output"}) + token = next( + entry for entry in descriptors if entry["root_index"] == 0 + ).get("root_scope", "0" * 64) + for changed in ([roots[1], roots[0]], [roots[2], roots[1]]): + self.api._get_all_output_dirs = lambda: changed + for path in ("image.png", "thumbnails/image_thumb.png"): + await self.assert_no_access(path, scope=token) + link = self.comfy_dir / "linked-output" + self._make_symlink(roots[0], link) + self.api._get_all_output_dirs = lambda: [link] + listing = await self.client.get( + "/prompt_manager/images/output?mode=workspace" + ) + token = (await listing.json())["images"][0]["root_scope"] + link.unlink() + self._make_symlink(roots[1], link) + for path in ("image.png", "thumbnails/image_thumb.png"): + await self.assert_no_access(path, scope=token) diff --git a/tests/test_api_library.py b/tests/test_api_library.py index 99d8c30..5d6386e 100644 --- a/tests/test_api_library.py +++ b/tests/test_api_library.py @@ -260,6 +260,44 @@ class TestLibraryAPI(ImageAPITestCase): self.assertTrue(item["previews"][1]["path"].endswith("-1.png")) self.assertTrue(all(image["root_index"] == 0 for image in item["previews"])) + async def test_root_scope_library_gallery_and_previews_share_formatter(self): + from urllib.parse import parse_qs, quote, urlsplit + + relative = "nested/猫 % #.png" + pair = self.api.db.save_prompt_pair("Cat", "Noise", title="Pair title") + positive = self.api.db.get_prompt_pair(pair)["positive_prompt_id"] + path = make_png(self.output_dir / relative) + make_png(self.output_dir / "thumbnails/nested/猫 % #_thumb.png") + image = self.api.db.link_image_to_prompt(positive, str(path)) + self.api.db.link_image_to_pair(pair, image, {"source": "fixture"}) + _, listing = await self.get() + preview = listing["items"][0]["previews"][0] + _, gallery = await self.get(f"/{positive}/images", pair_id=pair) + descriptor = gallery["images"][0] + self.assertRegex(descriptor.get("root_scope", ""), r"^[0-9a-f]{64}$") + self.assertEqual(preview["root_scope"], descriptor["root_scope"]) + for entry in (preview, descriptor): + self.assertIn(quote(relative, safe="/"), entry["url"]) + for url in (entry["url"], entry["thumbnail_url"]): + query = parse_qs(urlsplit(url).query) + self.assertEqual(query["root_scope"], [entry["root_scope"]]) + self.assertEqual((await self.client.get(url)).status, 200) + self.assertIn("v", parse_qs(urlsplit(entry["thumbnail_url"]).query)) + response = await self.client.get( + f"/prompt_manager/images/inspect/{quote(relative, safe='/')}", + params={"root": 0, "root_scope": descriptor["root_scope"]}, + ) + result = await response.json() + self.assertEqual(result["root_scope"], descriptor["root_scope"]) + self.assertEqual( + result["image"], + {"root_index": 0, "path": relative, "filename": "猫 % #.png"}, + ) + self.assertEqual(result["library"]["metadata_scope"], "pair") + self.assertEqual(result["library"]["pair_id"], pair) + self.assertEqual(result["library"]["positive_id"], positive) + self.assertEqual(result["library"]["title"], "Pair title") + class TestNameCursorBudget(unittest.TestCase): def test_lru_byte_budget_scope_and_oversize_keys_are_bounded(self): diff --git a/tests/test_api_prompt_inspector.py b/tests/test_api_prompt_inspector.py index 9932197..3c08875 100644 --- a/tests/test_api_prompt_inspector.py +++ b/tests/test_api_prompt_inspector.py @@ -347,7 +347,7 @@ class TestPromptInspector(ImageAPITestCase): self.assertEqual(response.status, 200) self.assertEqual(result["prompts"]["positive"]["status"], "absent") with patch( - "py.api.images.read_inspection_metadata", + "py.api.images.read_inspection_file", side_effect=RuntimeError(str(self.comfy_dir)), ): response, result = await self.inspect("image.jpg") @@ -641,3 +641,544 @@ class TestPromptInspector(ImageAPITestCase): self.assertEqual(record["source"], "saved-association") self.assertIn("cannot be verified", record["message"]) self.assertNotIn("Original portrait", json.dumps(replaced)) + + async def test_full_inspection_envelope_preserves_existing_roles_and_identity(self): + from pathlib import Path + + source = graph(negative="") + source["3"]["inputs"].update(seed=9007199254740993, cfg=6.5, steps=28) + raw = json.dumps(source) + workflow = '{"nodes":[],"seed":9007199254740993}' + path = png(self.output_dir / "image.png", prompt=raw, workflow=workflow) + self._link_image(path, "Saved text must not replace file roles") + original_open = Path.open + reads = [] + + def opened(candidate, *args, **kwargs): + if candidate == path: + reads.append(candidate) + return original_open(candidate, *args, **kwargs) + + with ( + patch.object(Path, "open", opened), + patch.object(Image.Image, "load", side_effect=AssertionError("pixels")), + patch.object( + self.api.db, + "get_image_prompt_info", + wraps=self.api.db.get_image_prompt_info, + ) as lookup, + ): + response, result = await self.inspect() + self.assertEqual(response.status, 200) + self.assertEqual(result.get("inspection_version"), 1) + self.assertEqual( + result["image"], + { + "root_index": 0, + "path": "image.png", + "filename": "image.png", + }, + ) + self.assertEqual( + result["prompts"], + { + "positive": self.side("portrait"), + "negative": self.side(""), + }, + ) + self.assertRegex(result["root_scope"], r"^[0-9a-f]{64}$") + self.assertEqual( + result["metadata"]["raw"], {"prompt": raw, "workflow": workflow} + ) + self.assertEqual( + result["metadata"]["record_formats"], + { + "prompt": "original-text", + "workflow": "original-text", + }, + ) + self.assertEqual( + result["metadata"]["prompt_graph"]["value"]["3"]["inputs"]["seed"], + "9007199254740993", + ) + self.assertEqual(result["file"]["format"], "PNG") + self.assertEqual((result["file"]["width"], result["file"]["height"]), (4, 4)) + self.assertEqual(result["file"]["observation"], "current-file") + self.assertEqual(result["resolution"]["sampler_id"], "3") + self.assertEqual(result["library"]["status"], "linked") + self.assertEqual(reads, [path]) + lookup.assert_called_once_with(str(path), False) + self.assertLessEqual(len(await response.read()), 8 * 1024 * 1024) + self.assertNotIn(str(self.comfy_dir), json.dumps(result)) + + async def test_workflow_raw_uses_same_aggregate_budget_and_read(self): + from utils.prompt_inspector import MAX_METADATA_BYTES + + path = self.output_dir / "image.png" + png(path, prompt=json.dumps(graph()), workflow="x" * MAX_METADATA_BYTES) + response, result = await self.inspect() + self.assertEqual(response.status, 200) + self.assertEqual(result.get("metadata", {}).get("status"), "unreadable") + self.assertEqual(result["prompts"]["positive"]["status"], "unresolved") + self.assertEqual(result["metadata"]["raw"], {}) + self.assertEqual(result["metadata"]["record_formats"], {}) + + async def test_workflow_only_keeps_existing_role_uncertainty(self): + path = png(self.output_dir / "image.png", workflow={"nodes": []}) + self._link_image(path, "Saved uncertainty") + _, result = await self.inspect() + self.assertEqual(result["prompts"]["positive"]["source"], "saved-association") + self.assertEqual( + result.get("metadata", {}).get("workflow_graph", {}).get("status"), "parsed" + ) + self.assertIsNone(result["prompts"]["positive"]["text"]) + self.assertEqual(result["prompts"]["negative"]["status"], "absent") + + async def test_library_failure_does_not_erase_file_evidence(self): + png(self.output_dir / "image.png", prompt=graph("File portrait")) + with patch.object( + self.api.db, + "get_image_prompt_info", + side_effect=OSError("private DB detail"), + ): + response, result = await self.inspect() + self.assertEqual(response.status, 200) + self.assertEqual(result["prompts"]["positive"], self.side("File portrait")) + self.assertEqual(result.get("library", {}).get("status"), "unavailable") + self.assertEqual(result["metadata"]["prompt_graph"]["status"], "parsed") + self.assertNotIn("private DB detail", json.dumps(result)) + + async def test_file_change_is_retryable_without_saved_text_fallback(self): + from pathlib import Path + + path = png(self.output_dir / "image.png", prompt=graph()) + self._link_image(path, "Saved recipe") + original_open = Path.open + + def opened(candidate, *args, **kwargs): + stream = original_open(candidate, *args, **kwargs) + if candidate == path and args and args[0] == "rb": + with original_open(path, "ab") as output: + output.write(b"changed") + return stream + + with patch.object(Path, "open", opened): + response, result = await self.inspect() + self.assertEqual(response.status, 409) + self.assertFalse(result["success"]) + self.assertNotIn("code", result) + self.assertNotIn("Saved recipe", json.dumps(result)) + self.assertEqual(response.headers.get("Cache-Control"), "no-store") + + async def test_surrogate_and_ordinary_unicode_round_trip_without_replacement(self): + source = graph("猫 😀 \ud800", "\udfff") + raw = json.dumps(source) + png( + self.output_dir / "image.png", + prompt=raw, + workflow='{"nodes":[],"label":"\\ud800 猫 😀"}', + ) + response, result = await self.inspect() + self.assertEqual(response.status, 200) + self.assertEqual(result["prompts"]["positive"], self.side("猫 😀 \ud800")) + self.assertEqual(result["prompts"]["negative"], self.side("\udfff")) + self.assertEqual(result["metadata"]["raw"]["prompt"], raw) + body = await response.read() + self.assertIn("猫 😀".encode(), body) + self.assertIn(b"\\ud800", body) + self.assertNotIn("\ufffd".encode(), body) + self.assertLessEqual(len(body), 8 * 1024 * 1024) + + async def video_inspect(self, comment=None, description=None): + import io + import threading + from unittest.mock import MagicMock + from utils.prompt_inspector import MAX_METADATA_BYTES + + path = self.output_dir / "clip.mp4" + path.write_bytes(b"fixture") + tags = { + key: value + for key, value in (("comment", comment), ("description", description)) + if value is not None + } + payload = json.dumps({"format": {"tags": tags}}).encode() + + class ProbeOutput(io.BytesIO): + def read(self, size=-1): + self.requested_size = size + return super().read(size) + + process = MagicMock() + process.stdout = ProbeOutput(payload) + process.poll.return_value = 0 + process.returncode = 0 + timers = [] + timer_type = threading.Timer + + def timer(*args): + instance = timer_type(*args) + timers.append(instance) + return instance + + with ( + patch("utils.prompt_inspector.ffprobe_path", return_value="/tool/ffprobe"), + patch( + "utils.prompt_inspector.subprocess.Popen", return_value=process + ) as spawn, + patch("utils.prompt_inspector.threading.Timer", side_effect=timer), + ): + response, result = await self.inspect("clip.mp4") + self.assertEqual(spawn.call_count, 1) + self.assertTrue(process.stdout.closed) + self.assertEqual(process.stdout.requested_size, MAX_METADATA_BYTES + 1) + self.assertEqual(len(timers), 1) + self.assertTrue(timers[0].finished.is_set()) + return response, result + + async def test_video_role_description_wins_over_workflow_only_comment(self): + evidence = {} + source = graph("chosen", "") + # CLIP execution is unnecessary here: use the established pair-node carrier. + source["1"] = { + "class_type": "PromptManagerPair", + "inputs": {"text": "chosen", "negative_text": ""}, + } + source["3"]["inputs"]["negative"] = ["1", 1] + record_pair_execution( + evidence, + "1", + text="chosen", + negative_text="", + positive_text="chosen trigger", + negative_effective_text="", + ) + _, result = await self.video_inspect( + json.dumps({"workflow": {"nodes": [{"id": "ignored"}]}}), + json.dumps( + { + "prompt": source, + "workflow": {"nodes": [{"id": "chosen"}]}, + **evidence, + } + ), + ) + self.assertEqual( + result["prompts"]["positive"], self.side("chosen trigger", "execution") + ) + self.assertEqual(result["prompts"]["negative"], self.side("", "execution")) + self.assertEqual( + result.get("metadata", {}).get("workflow_graph", {}).get("value"), + {"nodes": [{"id": "chosen"}]}, + ) + self.assertNotIn("ignored", json.dumps(result)) + + async def test_video_present_malformed_role_markers_block_other_carriers(self): + for marker in ( + {"prompt": None}, + {"prompt": "not-json"}, + {"parameters": None}, + {EXECUTION_KEY: "malformed"}, + ): + for comment, description in ( + ( + json.dumps({"workflow": {"nodes": [{"id": "ignored"}]}}), + json.dumps(marker), + ), + (json.dumps(marker), json.dumps({"prompt": graph("must not win")})), + ): + _, result = await self.video_inspect(comment, description) + self.assertEqual(result["prompts"]["positive"]["status"], "unresolved") + self.assertIsNone(result["prompts"]["positive"]["text"]) + self.assertEqual( + result.get("metadata", {}).get("workflow_graph", {}).get("status"), + "absent", + ) + self.assertNotIn("must not win", json.dumps(result)) + + async def test_video_workflow_carrier_fallback_never_merges_tags(self): + path = self.output_dir / "clip.mp4" + path.write_bytes(b"fixture") + self._link_image(path, "Saved workflow association") + _, result = await self.video_inspect( + '{"workflow":{"nodes":[{"id":"first"}]}}', + '{"workflow":{"nodes":[{"id":"second"}]},"unrelated":"unselected"}', + ) + self.assertEqual(result["prompts"]["positive"]["source"], "saved-association") + self.assertEqual( + result.get("metadata", {}).get("workflow_graph", {}).get("value"), + {"nodes": [{"id": "first"}]}, + ) + self.assertEqual( + result["metadata"]["raw"], {"workflow": '{"nodes": [{"id": "first"}]}'} + ) + self.assertNotIn("second", json.dumps(result)) + self.assertNotIn("unselected", json.dumps(result)) + + async def test_video_object_graph_decimal_and_unsafe_int_projection_is_exact(self): + raw = json.dumps(graph("video", "")) + raw = raw.replace( + '"positive": ["1", 0]', + '"seed":9007199254740993,"cfg":6.5000000000000000001,"positive": ["1", 0]', + ) + for prompt in (raw, json.dumps(raw)): + carrier = ( + '{"prompt":' + + prompt + + ',"workflow":{"nodes":[],"cfg":6.5000000000000000001,' + '"seed":9007199254740993}}' + ) + _, result = await self.video_inspect(carrier) + self.assertEqual(result["prompts"]["positive"], self.side("video")) + inputs = ( + result.get("metadata", {}) + .get("prompt_graph", {}) + .get("value", {}) + .get("3", {}) + .get("inputs", {}) + ) + self.assertEqual(inputs.get("cfg"), "6.5000000000000000001") + self.assertEqual(inputs["seed"], "9007199254740993") + self.assertEqual(inputs["positive"], ["1", 0]) + self.assertEqual( + result["metadata"]["workflow_graph"]["value"]["cfg"], + "6.5000000000000000001", + ) + self.assertNotIn("video_carrier_text", result) + self.assertNotIn("carrier_text", json.dumps(result)) + + async def test_video_original_string_and_reserialized_object_formats_are_distinct( + self, + ): + raw = json.dumps(graph()) + for record, expected in ( + (json.dumps(raw), "original-text"), + (raw, "reserialized-json"), + ): + _, result = await self.video_inspect( + '{"prompt":' + record + ',"workflow":' + record + "}" + ) + self.assertEqual( + result.get("metadata", {}).get("record_formats"), + {"prompt": expected, "workflow": expected}, + ) + self.assertEqual(result["metadata"]["raw"]["prompt"], raw) + + async def test_video_null_role_marker_raises_before_successful_read(self): + for key in ("prompt", "parameters", EXECUTION_KEY): + _, result = await self.video_inspect(json.dumps({key: None})) + self.assertEqual(result["prompts"]["positive"]["status"], "unresolved") + self.assertEqual(result.get("metadata", {}).get("status"), "unreadable") + + async def test_video_null_workflow_keeps_valid_legacy_roles(self): + _, result = await self.video_inspect( + json.dumps({"prompt": graph(), "workflow": None}) + ) + self.assertEqual(result["prompts"]["positive"], self.side("portrait")) + self.assertEqual( + result.get("metadata", {}).get("workflow_graph"), + {"status": "malformed", "value": None}, + ) + self.assertEqual(result["metadata"]["record_formats"]["workflow"], "invalid") + self.assertNotIn("workflow", result["metadata"]["raw"]) + + async def test_video_strict_projection_failure_preserves_legacy_roles(self): + source = json.dumps(graph()) + for suffix, expected in ( + (',"duplicate":1,"duplicate":2', "malformed"), + (',"nonfinite":NaN', "malformed"), + (',"depth":' + "[" * 64 + "0" + "]" * 64, "limited"), + ): + _, result = await self.video_inspect( + '{"prompt":' + source + ',"workflow":{"nodes":[]}' + suffix + "}" + ) + self.assertEqual(result["prompts"]["positive"], self.side("portrait")) + self.assertEqual( + result.get("metadata", {}).get("prompt_graph", {}).get("status"), + expected, + ) + self.assertIsNone(result["metadata"]["prompt_graph"]["value"]) + self.assertEqual(result["metadata"]["workflow_graph"]["status"], expected) + + async def test_video_numeric_role_text_never_becomes_sendable(self): + source = graph(positive=123) + _, result = await self.video_inspect(json.dumps({"prompt": source})) + self.assertEqual(result["prompts"]["positive"]["status"], "unresolved") + self.assertIsNone(result["prompts"]["positive"]["text"]) + self.assertEqual( + result.get("metadata", {}).get("prompt_graph", {}).get("status"), "parsed" + ) + + async def scoped_descriptor(self, root_index=0): + response = await self.client.get("/prompt_manager/images/output?mode=workspace") + entries = (await response.json())["images"] + return next(entry for entry in entries if entry["root_index"] == root_index) + + async def assert_scope_skips_access(self, scope, root="0", status=409): + from pathlib import Path + + original_stat = Path.stat + + def stat(candidate, *args, **kwargs): + if candidate.name == "image.png": + raise AssertionError("candidate stat before scope guard") + return original_stat(candidate, *args, **kwargs) + + params = {"root_scope": scope} + if root is not None: + params["root"] = root + with ( + patch.object(Path, "stat", stat), + patch( + "utils.prompt_inspector.subprocess.Popen", + side_effect=AssertionError("probe"), + ), + patch.object( + self.api.db, + "get_image_prompt_info", + side_effect=AssertionError("saved lookup"), + ), + patch.object( + self.api, "_file_response", side_effect=AssertionError("file response") + ), + patch( + "py.api.images.read_inspection_file", + side_effect=AssertionError("metadata read"), + ), + ): + response = await self.client.get( + "/prompt_manager/images/inspect/image.png", params=params + ) + self.assertEqual(response.status, status) + result = await response.json() + self.assertFalse(result["success"]) + if status == 409: + self.assertEqual(result["code"], "source-changed") + self.assertEqual(response.headers.get("Cache-Control"), "no-store") + self.assertNotIn(str(self.comfy_dir), json.dumps(result)) + + async def test_root_scope_equal_label_reorder_or_replacement_never_reads_other_file( + self, + ): + from py.config import GalleryConfig + + roots = [self.comfy_dir / name / "output" for name in ("a", "b", "c")] + for index, root in enumerate(roots): + png(root / "image.png", prompt=graph(str(index))) + with patch.object( + GalleryConfig, + "path_anchors", + classmethod(lambda cls: [str(root.parent) for root in roots]), + ): + self.api._get_all_output_dirs = lambda: roots[:2] + descriptor = await self.scoped_descriptor() + other = await self.scoped_descriptor(1) + self.assertEqual(descriptor["root_dir"], other["root_dir"]) + for changed in ([roots[1], roots[0]], [roots[2], roots[1]]): + self.api._get_all_output_dirs = lambda: changed + await self.assert_scope_skips_access( + descriptor.get("root_scope", "0" * 64) + ) + + async def test_root_scope_symlink_retarget_requires_refresh(self): + first, second = self.comfy_dir / "a", self.comfy_dir / "b" + png(first / "image.png", prompt=graph("first")) + png(second / "image.png", prompt=graph("second")) + link = self.comfy_dir / "linked" + self._make_symlink(first, link) + self.api._get_all_output_dirs = lambda: [link] + descriptor = await self.scoped_descriptor() + link.unlink() + self._make_symlink(second, link) + await self.assert_scope_skips_access(descriptor.get("root_scope", "0" * 64)) + + async def test_root_scope_rotated_key_and_removed_index_skip_reader_and_db(self): + png(self.output_dir / "image.png", prompt=graph()) + descriptor = await self.scoped_descriptor() + import py.api.images as images_module + + with patch.object(images_module, "_WORKSPACE_CURSOR_SECRET", b"new-key"): + await self.assert_scope_skips_access(descriptor.get("root_scope", "0" * 64)) + self.api._get_all_output_dirs = lambda: [] + await self.assert_scope_skips_access(descriptor.get("root_scope", "0" * 64)) + + async def test_root_scope_malformed_request_skips_reader(self): + png(self.output_dir / "image.png", prompt=graph()) + for token in ("", "z" * 64, "a" * 63, "a" * 65, "a" * 4097): + await self.assert_scope_skips_access(token, status=400) + for root in (None, "-1", "x", "0.0", "+0"): + await self.assert_scope_skips_access("a" * 64, root=root, status=400) + + async def test_root_scope_echo_preserves_legacy_image_and_roles(self): + png(self.output_dir / "image.png", prompt=graph()) + descriptor = await self.scoped_descriptor() + token = descriptor.get("root_scope", "0" * 64) + _, original = await self.inspect() + response = await self.client.get( + "/prompt_manager/images/inspect/image.png", + params={"root": "0", "root_scope": token}, + ) + result = await response.json() + self.assertEqual(response.status, 200) + self.assertEqual(result.get("root_scope"), token) + self.assertEqual( + result["image"], + {"root_index": 0, "path": "image.png", "filename": "image.png"}, + ) + self.assertEqual(result["prompts"], original["prompts"]) + outside = png(self.comfy_dir / "outside.png", prompt=graph("private")) + self._make_symlink(outside, self.output_dir / "escape.png") + for path in ("..%2Foutside.png", "escape.png"): + denied = await self.client.get( + f"/prompt_manager/images/inspect/{path}", + params={"root": "0", "root_scope": token}, + ) + self.assertEqual(denied.status, 403) + + async def test_root_scope_source_change_returns_refresh_instead_of_other_root(self): + png(self.output_dir / "image.png", prompt=graph("old source")) + other = self.comfy_dir / "input" + png(other / "image.png", prompt=graph("different source")) + self.api._get_all_output_dirs = lambda: [self.output_dir, other] + descriptor = await self.scoped_descriptor() + self.api._get_all_output_dirs = lambda: [other, self.output_dir] + response = await self.client.get( + "/prompt_manager/images/inspect/image.png", + params={"root": 0, "root_scope": descriptor.get("root_scope", "0" * 64)}, + ) + self.assertEqual(response.status, 409) + result = await response.json() + self.assertEqual(result["code"], "source-changed") + self.assertNotIn("different source", json.dumps(result)) + + async def test_video_unrepresentable_decimal_preserves_file_roles_and_raw(self): + carrier = json.dumps({"prompt": graph()}).replace( + '"positive": ["1", 0]', + '"cfg":1e1000000000000000000000000,"positive": ["1", 0]', + ) + response, result = await self.video_inspect(carrier) + self.assertEqual(response.status, 200) + self.assertEqual(result["prompts"]["positive"], self.side("portrait")) + self.assertEqual( + result["metadata"]["prompt_graph"], {"status": "malformed", "value": None} + ) + self.assertEqual( + result["metadata"]["record_formats"]["prompt"], "reserialized-json" + ) + self.assertIn("Infinity", result["metadata"]["raw"]["prompt"]) + + async def test_root_scope_guard_and_evidence_use_one_current_roots_snapshot(self): + png(self.output_dir / "image.png", prompt=graph()) + descriptor = await self.scoped_descriptor() + with patch.object( + self.api, "_get_all_output_dirs", side_effect=[[self.output_dir]] + ) as roots: + response = await self.client.get( + "/prompt_manager/images/inspect/image.png", + params={"root": 0, "root_scope": descriptor["root_scope"]}, + ) + self.assertEqual(response.status, 200) + self.assertEqual( + (await response.json())["root_scope"], descriptor["root_scope"] + ) + roots.assert_called_once_with() diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 2c70e02..425c9df 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -255,6 +255,51 @@ class TestWorkspaceImages(ImageAPITestCase): response = await self.client.get("/prompt_manager/gallery/subfolders") self.assertNotIn("roots", await response.json()) + async def test_root_scope_workspace_descriptors_guard_original_and_thumbnail_urls( + self, + ): + from urllib.parse import parse_qs, quote, urlsplit + + roots = [self.output_dir, self.comfy_dir / "input"] + self.api._get_all_output_dirs = lambda: roots + relative = "nested/猫 % #.png" + for index, root in enumerate(roots): + self.media(relative, root=root) + self.media("thumbnails/nested/猫 % #_thumb.png", root=root) + make_png(root / relative, color=(index, 0, 0)) + make_png(root / "thumbnails/nested/猫 % #_thumb.png", color=(index, 0, 0)) + _, result = await self.get(limit=10) + for entry in result["images"]: + self.assertRegex(entry.get("root_scope", ""), r"^[0-9a-f]{64}$") + self.assertIn(quote(relative, safe="/"), entry["url"]) + for url in (entry["url"], entry["thumbnail_url"]): + query = parse_qs(urlsplit(url).query) + self.assertEqual(query["root_scope"], [entry["root_scope"]]) + self.assertEqual(query["root"], [str(entry["root_index"])]) + response = await self.client.get(url) + self.assertEqual(response.status, 200) + expected = roots[entry["root_index"]] / ( + "thumbnails/nested/猫 % #_thumb.png" + if url == entry["thumbnail_url"] + else relative + ) + self.assertEqual(await response.read(), expected.read_bytes()) + self.assertIn("v", parse_qs(urlsplit(entry["thumbnail_url"]).query)) + response = await self.client.get( + f"/prompt_manager/images/inspect/{quote(relative, safe='/')}", + params={"root": entry["root_index"], "root_scope": entry["root_scope"]}, + ) + data = await response.json() + self.assertEqual(data["root_scope"], entry["root_scope"]) + self.assertEqual( + data["image"], + { + "root_index": entry["root_index"], + "path": relative, + "filename": "猫 % #.png", + }, + ) + class TestWorkspaceCursorEncoding(unittest.TestCase): def test_ntfs_unicode_basename_cursor_is_bounded_and_round_trips(self): diff --git a/tests/test_inspection_response.py b/tests/test_inspection_response.py new file mode 100644 index 0000000..0488abb --- /dev/null +++ b/tests/test_inspection_response.py @@ -0,0 +1,465 @@ +"""Wire projections preserve exact bounded evidence independently of role parsing.""" + +import io +import json +import tempfile +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +from utils.inspection_response import encode_inspection_response +from utils.prompt_inspector import ( + InspectionRead, + MAX_METADATA_BYTES, + MetadataReadError, + inspect_prompt_roles, + read_inspection_file, + read_inspection_metadata, +) +from utils.pair_execution import EXECUTION_KEY + +CAP = 8 * 1024 * 1024 +IMAGE = {"root_index": 0, "path": "image.png", "filename": "image.png"} +PROMPTS = { + "positive": {"status": "resolved", "text": "portrait", "source": "graph-inferred"}, + "negative": {"status": "resolved", "text": "", "source": "graph-inferred"}, +} + + +def inspection(metadata, *, formats=None, carrier=None): + return InspectionRead( + metadata=metadata, + record_formats=( + formats + if formats is not None + else {key: "original-text" for key in metadata} + ), + video_carrier_text=carrier, + format="PNG", + width=4, + height=4, + ) + + +def encode(read, *, prompts=None, file=None, library=None): + return encode_inspection_response( + image=IMAGE, + root_scope="a" * 64, + file=file + or { + "observation": "current-file", + "source_label": "output", + "media_type": "image", + "extension": ".png", + "size": 100, + "modified_time": 100, + }, + prompts=PROMPTS if prompts is None else prompts, + resolution={"state": "resolved", "sampler_id": "3", "reason": ""}, + inspection=read, + library=library + or { + "status": "none", + "scope": "saved-library", + "positive_id": None, + "pair_id": None, + "title": None, + "category": None, + "tags": [], + "metadata_scope": None, + }, + ) + + +def graph_text(count=3): + graph = { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "portrait"}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": ""}}, + "3": { + "class_type": "KSampler", + "inputs": { + "positive": ["1", 0], + "negative": ["2", 0], + "seed": 9007199254740993, + "cfg": "EXACT_CFG", + }, + }, + } + for index in range(4, count + 1): + graph[str(index)] = { + "class_type": "Primitive", + "inputs": {"value": 1, "labels": [True, False, "ordinary"]}, + } + return json.dumps(graph).replace('"EXACT_CFG"', "6.5000000000000000001") + + +class TestInspectionResponse(unittest.TestCase): + def test_normalized_numbers_preserve_exact_literals(self): + text = graph_text().replace( + '"cfg": 6.5000000000000000001', + '"cfg": 6.5000000000000000001, "huge": 1e1000000000', + ) + result = json.loads(encode(inspection({"prompt": text}))) + inputs = result["metadata"]["prompt_graph"]["value"]["3"]["inputs"] + self.assertEqual(inputs["seed"], "9007199254740993") + self.assertEqual(inputs["cfg"], "6.5000000000000000001") + self.assertEqual(inputs["huge"].lower(), "1e+1000000000") + self.assertEqual(inputs["positive"], ["1", 0]) + self.assertEqual(result["metadata"]["raw"]["prompt"], text) + self.assertEqual(result["prompts"], PROMPTS) + + def test_normalization_preserves_150_node_graph_with_more_than_500_json_values( + self, + ): + text = graph_text(150) + self.assertGreater(len(text.split(":")), 500) + for read in ( + inspection({"prompt": text}), + inspection( + {"prompt": json.dumps(json.loads(text))}, + formats={"prompt": "reserialized-json"}, + carrier='{"prompt":' + text + ',"unrelated":{"value":1}}', + ), + ): + result = json.loads(encode(read)) + graph = result["metadata"]["prompt_graph"] + self.assertEqual(graph["status"], "parsed") + self.assertEqual(len(graph["value"]), 150) + self.assertEqual( + graph["value"]["3"]["inputs"]["cfg"], "6.5000000000000000001" + ) + self.assertEqual(graph["value"]["3"]["inputs"]["seed"], "9007199254740993") + self.assertEqual(graph["value"]["3"]["inputs"]["positive"], ["1", 0]) + self.assertIs(graph["value"]["4"]["inputs"]["labels"][0], True) + self.assertEqual(result["metadata"]["raw"], read.metadata) + self.assertEqual(result["prompts"], PROMPTS) + + def test_graph_node_limit_counts_nodes_not_json_fields(self): + for count, status in ((500, "parsed"), (501, "limited")): + prompt = graph_text(count) + workflow = json.dumps( + { + "nodes": [ + {"id": i, "type": "Primitive", "widgets_values": [1, True]} + for i in range(count) + ] + } + ) + for read in ( + inspection({"prompt": prompt, "workflow": workflow}), + inspection( + {"prompt": json.dumps(json.loads(prompt)), "workflow": workflow}, + formats={ + "prompt": "reserialized-json", + "workflow": "reserialized-json", + }, + carrier='{"other":[],"prompt":' + + prompt + + ',"workflow":' + + workflow + + "}", + ), + ): + result = json.loads(encode(read)) + for key in ("prompt_graph", "workflow_graph"): + self.assertEqual(result["metadata"][key]["status"], status) + if status == "limited": + self.assertIsNone(result["metadata"][key]["value"]) + self.assertEqual(result["metadata"]["raw"], read.metadata) + self.assertEqual(result["prompts"], PROMPTS) + + def test_strict_projection_depth_is_separate_from_semantic_link_depth(self): + for depth in (64, 65): + # API map root + node + inputs + nested array containers. + text = ( + '{"1":{"class_type":"Primitive","inputs":{"value":' + + "[" * (depth - 3) + + "0" + + "]" * (depth - 3) + + "}}}" + ) + result = json.loads(encode(inspection({"prompt": text}))) + self.assertEqual( + result["metadata"]["prompt_graph"]["status"], + "parsed" if depth == 64 else "limited", + ) + carrier = '{"prompt":' + text + "}" + result = json.loads( + encode( + inspection( + {"prompt": text}, + formats={"prompt": "reserialized-json"}, + carrier=carrier, + ) + ) + ) + self.assertEqual(result["metadata"]["prompt_graph"]["status"], "limited") + text = '{"1":{"inputs":{"value":' + "[" * 60 + "0" + "]" * 60 + "}}}" + result = json.loads( + encode( + inspection( + {"prompt": text}, + formats={"prompt": "reserialized-json"}, + carrier='{"prompt":' + text + "}", + ) + ) + ) + self.assertEqual(result["metadata"]["prompt_graph"]["status"], "parsed") + chain = { + str(i): { + "class_type": "Reroute", + "inputs": {"value": [str((i + 1) % 150), 0]}, + } + for i in range(150) + } + result = json.loads(encode(inspection({"prompt": json.dumps(chain)}))) + self.assertEqual(result["metadata"]["prompt_graph"]["status"], "parsed") + self.assertEqual(len(result["metadata"]["prompt_graph"]["value"]), 150) + + def test_wide_bounded_values_do_not_consume_semantic_visit_budget(self): + text = json.dumps( + { + "1": { + "inputs": { + "values": list(range(20000)), + "fields": {str(i): True for i in range(1000)}, + } + } + } + ) + self.assertLess(len(text.encode()), MAX_METADATA_BYTES) + body = encode(inspection({"prompt": text})) + graph = json.loads(body)["metadata"]["prompt_graph"] + self.assertEqual(graph["status"], "parsed") + self.assertEqual(len(graph["value"]["1"]["inputs"]["values"]), 20000) + self.assertEqual(len(graph["value"]["1"]["inputs"]["fields"]), 1000) + self.assertLess(len(body), CAP) + + def test_strict_malformed_records_preserve_original_raw_and_roles(self): + for text in ( + '{"1":{},"1":{}}', + '{"1":{"value":NaN}}', + '{"1":{"value":Infinity}}', + "{broken", + "[]", + "42", + "[" * 2000, + ): + result = json.loads(encode(inspection({"prompt": text, "workflow": text}))) + self.assertIn( + result["metadata"]["prompt_graph"]["status"], ("malformed", "limited") + ) + self.assertIsNone(result["metadata"]["prompt_graph"]["value"]) + self.assertEqual( + result["metadata"]["raw"], {"prompt": text, "workflow": text} + ) + self.assertEqual(result["prompts"], PROMPTS) + + def test_video_missing_exact_source_is_limited_not_rounded(self): + rounded = json.dumps(json.loads(graph_text())) + result = json.loads( + encode( + inspection({"prompt": rounded}, formats={"prompt": "reserialized-json"}) + ) + ) + self.assertEqual( + result["metadata"]["prompt_graph"], {"status": "limited", "value": None} + ) + self.assertEqual(result["metadata"]["raw"]["prompt"], rounded) + self.assertEqual(result["prompts"], PROMPTS) + + def test_response_escaping_omits_graphs_then_entire_raw_atomically(self): + prompts = { + role: { + "status": "resolved", + "text": "\x01" * 100000, + "source": "graph-inferred", + } + for role in ("positive", "negative") + } + source = json.loads(graph_text()) + source["1"]["inputs"]["text"] = prompts["positive"]["text"] + source["2"]["inputs"]["text"] = prompts["negative"]["text"] + raw = json.dumps(source) + for count, raw_retained in ((450000, True), (550000, False)): + metadata = {"prompt": raw, "parameters": "\x01" * count} + self.assertLess( + sum(len(value.encode()) for value in metadata.values()), + MAX_METADATA_BYTES, + ) + body = encode(inspection(metadata), prompts=prompts) + result = json.loads(body) + self.assertLessEqual(len(body), CAP) + self.assertEqual(result["image"], IMAGE) + self.assertEqual(result["prompts"], prompts) + self.assertEqual( + result["metadata"]["prompt_graph"], {"status": "limited", "value": None} + ) + self.assertEqual( + result["metadata"]["workflow_graph"], + {"status": "absent", "value": None}, + ) + self.assertEqual( + result["metadata"]["record_formats"], + {key: "original-text" for key in metadata}, + ) + if raw_retained: + self.assertEqual(result["metadata"]["status"], "available") + self.assertIsNone(result["metadata"]["omitted_reason"]) + self.assertEqual(result["metadata"]["raw"], metadata) + self.assertEqual( + result["metadata"]["parameters"], metadata["parameters"] + ) + else: + self.assertEqual(result["metadata"]["raw"], {}) + self.assertIsNone(result["metadata"]["parameters"]) + self.assertEqual(result["metadata"]["status"], "limited") + self.assertEqual( + result["metadata"]["omitted_reason"], "response-size-limit" + ) + + def test_response_byte_boundary_is_inclusive(self): + prompts = { + role: { + "status": "resolved", + "text": "\x01" * 100000, + "source": "graph-inferred", + } + for role in ("positive", "negative") + } + text = graph_text().replace( + '"cfg": 6.5000000000000000001', '"cfg":6.5000000000000000001,"label":"猫😀"' + ) + read = inspection({"prompt": text, "parameters": "\x01" * 500000}) + base_file = { + "observation": "current-file", + "source_label": "output", + "media_type": "image", + "extension": ".png", + } + initial_size = len(encode(read, prompts=prompts, file=base_file)) + self.assertLess(initial_size, CAP) + for target in (CAP - 1, CAP, CAP + 1): + extra = target - initial_size + file = {**base_file, "source_label": "output" + ("x" if extra % 2 else "")} + parameters = read.metadata["parameters"] + "x" * (extra // 2) + current = inspection({"prompt": text, "parameters": parameters}) + self.assertLess( + sum(len(value.encode()) for value in current.metadata.values()), + MAX_METADATA_BYTES, + ) + body = encode(current, prompts=prompts, file=file) + result = json.loads(body) + if target <= CAP: + self.assertEqual(len(body), target) + self.assertEqual(result["metadata"]["prompt_graph"]["status"], "parsed") + else: + self.assertLessEqual(len(body), CAP) + self.assertEqual( + result["metadata"]["prompt_graph"]["status"], "limited" + ) + self.assertEqual(result["prompts"], prompts) + self.assertEqual(result["image"], IMAGE) + self.assertEqual(result["metadata"]["raw"], current.metadata) + + def test_optional_display_fields_are_validated_and_library_is_bounded(self): + file = { + "size": -1, + "width": True, + "height": float("inf"), + "modified_time": float("nan"), + "format": 12, + "observation": "current-file", + } + library = { + "status": "linked", + "positive_id": 2**53, + "pair_id": True, + "tags": ["valid", 3], + "category": {}, + "title": [], + } + result = json.loads(encode(None, file=file, library=library)) + for key in ("size", "width", "height", "modified_time", "format"): + self.assertNotIn(key, result["file"]) + for key in ("positive_id", "pair_id", "category", "title"): + self.assertNotIn(key, result["library"]) + self.assertEqual(result["library"]["tags"], ["valid"]) + result = json.loads( + encode(inspection({}), library={"status": "linked", "tags": ["猫" * 12000]}) + ) + self.assertEqual(result["library"]["status"], "limited") + self.assertEqual(result["library"]["tags"], []) + self.assertLess( + len(json.dumps(result["library"], ensure_ascii=False).encode()), 32 * 1024 + ) + + def test_absent_unreadable_invalid_workflow_and_formats_remain_distinct(self): + for read, status in ( + (None, "unreadable"), + (inspection({}), "absent"), + (inspection({}, formats={"workflow": "invalid"}), "available"), + ): + result = json.loads(encode(read)) + self.assertEqual(result["metadata"]["status"], status) + self.assertEqual(result["metadata"]["raw"], {}) + self.assertIsNone(result["metadata"]["parameters"]) + self.assertEqual(result["metadata"]["workflow_graph"]["status"], "malformed") + + def test_video_null_role_marker_raises_before_successful_read(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "clip.mp4" + path.write_bytes(b"fixture") + for field in ("prompt", "parameters", EXECUTION_KEY): + for reader in (read_inspection_file, read_inspection_metadata): + process = MagicMock() + process.stdout = io.BytesIO( + json.dumps( + {"format": {"tags": {"comment": json.dumps({field: None})}}} + ).encode() + ) + process.poll.return_value = 0 + process.returncode = 0 + with ( + patch( + "utils.prompt_inspector.ffprobe_path", + return_value="/tool/ffprobe", + ), + patch( + "utils.prompt_inspector.subprocess.Popen", + return_value=process, + ), + ): + with self.assertRaises(MetadataReadError): + reader(path) + self.assertTrue(process.stdout.closed) + + def test_original_role_projection_ignores_workflow_without_losing_saved_warning( + self, + ): + self.assertEqual( + inspect_prompt_roles({"workflow": "{}"}, {"prompt_id": 1})["positive"][ + "source" + ], + "saved-association", + ) + + def test_unrepresentable_decimal_exponent_preserves_roles_and_raw(self): + text = graph_text().replace( + "6.5000000000000000001", "1e1000000000000000000000000" + ) + for read in ( + inspection({"prompt": text}), + inspection( + {"prompt": json.dumps(json.loads(text))}, + formats={"prompt": "reserialized-json"}, + carrier='{"prompt":' + text + "}", + ), + ): + result = json.loads(encode(read)) + self.assertEqual( + result["metadata"]["prompt_graph"], + {"status": "malformed", "value": None}, + ) + self.assertEqual(result["metadata"]["raw"], read.metadata) + self.assertEqual(result["prompts"], PROMPTS) diff --git a/utils/inspection_response.py b/utils/inspection_response.py new file mode 100644 index 0000000..b7f5162 --- /dev/null +++ b/utils/inspection_response.py @@ -0,0 +1,279 @@ +"""Exact inspection-only graph projection and bounded JSON wire encoding.""" + +import json +import math +from decimal import Decimal, InvalidOperation +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .prompt_inspector import InspectionRead + +MAX_RESPONSE_BYTES = 8 * 1024 * 1024 +MAX_LIBRARY_BYTES = 32 * 1024 +_MAX_SAFE_INTEGER = 2**53 - 1 +_RECORD_KEYS = frozenset( + {"prompt", "workflow", "parameters", "comfyassets_prompt_pairs"} +) +_FORMATS = frozenset({"original-text", "reserialized-json", "invalid"}) + + +class _ProjectionLimit(ValueError): + """A present graph cannot be projected within inspection limits.""" + + +def _pairs(items: list[tuple]) -> dict: + result = {} + for key, value in items: + if key in result: + raise ValueError("Duplicate JSON key") + result[key] = value + return result + + +def _nonfinite(value: str) -> None: + raise ValueError("Nonfinite JSON number") + + +def _normalize( + value: object, depth: int = 1, *, retain: bool = True, carrier: bool = False +) -> object: + """Walk bounded source once; unrelated carrier fields are only validated.""" + if isinstance(value, (dict, list)): + if depth > 64: + raise _ProjectionLimit() + if isinstance(value, dict): + result = {} if retain else None + for key, child in value.items(): + keep = retain and (not carrier or key in ("prompt", "workflow")) + normalized = _normalize(child, depth + 1, retain=keep) + if keep: + result[key] = normalized + else: + result = [] if retain else None + for child in value: + normalized = _normalize(child, depth + 1, retain=retain) + if retain: + result.append(normalized) + return result + if isinstance(value, Decimal): + if not value.is_finite(): + raise ValueError("Nonfinite JSON number") + # str keeps large exponents compact and never rounds through a float. + return str(value) if retain else None + if type(value) is int and abs(value) > _MAX_SAFE_INTEGER: + return str(value) if retain else None + return value if retain else None + + +def _strict_projection(text: str, *, carrier: bool = False) -> object: + try: + tree = json.loads( + text, + parse_float=Decimal, + object_pairs_hook=_pairs, + parse_constant=_nonfinite, + ) + if carrier and not isinstance(tree, dict): + raise ValueError("Malformed carrier") + return _normalize(tree, carrier=carrier) + except RecursionError: + raise _ProjectionLimit() from None + except InvalidOperation: + raise ValueError("JSON decimal exceeds supported range") from None + + +def _graph(value: object, *, workflow: bool = False) -> dict: + if not isinstance(value, dict): + return {"status": "malformed", "value": None} + nodes = value.get("nodes") if workflow else value + if isinstance(nodes, (dict, list)) and len(nodes) > 500: + return {"status": "limited", "value": None} + return {"status": "parsed", "value": value} + + +def _metadata(inspection: "InspectionRead | None") -> dict: + raw = ( + {} + if inspection is None + else { + key: value + for key, value in inspection.metadata.items() + if key in _RECORD_KEYS and isinstance(value, str) + } + ) + formats = ( + {} + if inspection is None + else { + key: value + for key, value in inspection.record_formats.items() + if key in _RECORD_KEYS and value in _FORMATS + } + ) + result = { + "status": ( + "unreadable" + if inspection is None + else ("available" if raw or formats else "absent") + ), + "omitted_reason": None, + "raw": raw, + "record_formats": formats, + "parameters": raw.get("parameters"), + } + carrier, failed = None, None + if inspection is not None and inspection.video_carrier_text is not None: + try: + carrier = _strict_projection(inspection.video_carrier_text, carrier=True) + except _ProjectionLimit: + failed = "limited" + except (ValueError, TypeError, OverflowError): + failed = "malformed" + for key, field in (("prompt", "prompt_graph"), ("workflow", "workflow_graph")): + graph = {"status": "absent", "value": None} + if key in raw or key in formats: + if failed: + graph["status"] = failed + elif formats.get(key) == "invalid": + graph["status"] = "malformed" + else: + try: + if carrier is not None: + value = carrier.get(key) + if isinstance(value, str): + value = _strict_projection(value) + elif formats.get(key) == "reserialized-json": + raise _ProjectionLimit() + else: + value = _strict_projection(raw[key]) + graph = _graph(value, workflow=key == "workflow") + except _ProjectionLimit: + graph["status"] = "limited" + except (ValueError, TypeError, KeyError, OverflowError): + graph["status"] = "malformed" + result[field] = graph + return result + + +def _bytes(value: dict) -> bytes: + # JSON preserves isolated surrogate evidence through escaped code units. + return json.dumps( + value, ensure_ascii=False, allow_nan=False, separators=(",", ":") + ).encode("utf-8", "backslashreplace") + + +def _finite(value: object, *, positive: bool = False) -> bool: + try: + return ( + type(value) in (int, float) + and math.isfinite(value) + and (value > 0 if positive else value >= 0) + ) + except OverflowError: + return False + + +def _file_projection(file: dict, inspection: "InspectionRead | None") -> dict: + source = dict(file) + if inspection is not None: + source.update( + format=inspection.format, width=inspection.width, height=inspection.height + ) + result = {"observation": "current-file"} + for key in ("source_label", "extension", "format"): + value = source.get(key) + if value is None and key == "format": + result[key] = None + elif isinstance(value, str): + result[key] = value + if source.get("media_type") in ("image", "video"): + result["media_type"] = source["media_type"] + for key in ("size", "modified_time", "width", "height"): + if key not in source: + continue + value = source[key] + if value is None or _finite(value, positive=key in ("width", "height")): + if key not in ("width", "height") or value is None or type(value) is int: + result[key] = value + return result + + +def _library_projection(library: dict) -> dict: + status = library.get("status") + result = { + "status": ( + status + if status in ("linked", "none", "ambiguous", "unavailable", "limited") + else "unavailable" + ), + "scope": "saved-library", + } + for key in ("positive_id", "pair_id"): + value = library.get(key) + if value is None or type(value) is int and 0 < value <= _MAX_SAFE_INTEGER: + result[key] = value + for key in ("title", "category"): + value = library.get(key) + if value is None or isinstance(value, str): + result[key] = value + scope = library.get("metadata_scope") + if scope in ("pair", "positive", None): + result["metadata_scope"] = scope + tags = library.get("tags") + result["tags"] = ( + [tag for tag in tags if isinstance(tag, str)] if isinstance(tags, list) else [] + ) + if len(_bytes(result)) > MAX_LIBRARY_BYTES: + return { + "status": "limited", + "scope": "saved-library", + "positive_id": None, + "pair_id": None, + "title": None, + "category": None, + "tags": [], + "metadata_scope": None, + } + return result + + +def encode_inspection_response( + *, + image: dict, + root_scope: str, + file: dict, + prompts: dict, + resolution: dict, + inspection: "InspectionRead | None", + library: dict, +) -> bytes: + """Serialize exactly once per budget step, preserving identity and roles.""" + metadata = _metadata(inspection) + response = { + "success": True, + "inspection_version": 1, + "root_scope": root_scope, + "image": image, + "file": _file_projection(file, inspection), + "prompts": prompts, + "resolution": resolution, + "metadata": metadata, + "library": _library_projection(library), + } + body = _bytes(response) + if len(body) > MAX_RESPONSE_BYTES: + for field in ("prompt_graph", "workflow_graph"): + if metadata[field]["status"] != "absent": + metadata[field] = {"status": "limited", "value": None} + body = _bytes(response) + if len(body) > MAX_RESPONSE_BYTES: + metadata.update( + raw={}, + parameters=None, + status="limited", + omitted_reason="response-size-limit", + ) + body = _bytes(response) + if len(body) > MAX_RESPONSE_BYTES: + raise ValueError("Inspection response exceeds limits") + return body diff --git a/utils/prompt_inspector.py b/utils/prompt_inspector.py index 03d6c10..e71f447 100644 --- a/utils/prompt_inspector.py +++ b/utils/prompt_inspector.py @@ -6,6 +6,7 @@ import struct import subprocess # nosec B404 import threading import zlib +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -13,12 +14,13 @@ from PIL import Image from .pair_execution import EXECUTION_KEY from .prompt_pair_resolver import resolve_prompt_pair -from .video_metadata import ffprobe_path, is_video_path, metadata_from_tags +from .video_metadata import ffprobe_path, is_video_path MAX_METADATA_BYTES = 2 * 1024 * 1024 MAX_HEADER_BYTES = 8 * 1024 * 1024 MAX_PNG_CHUNKS = 10_000 -_METADATA_KEYS = frozenset({"prompt", "parameters", EXECUTION_KEY}) +_ROLE_METADATA_KEYS = frozenset({"prompt", "parameters", EXECUTION_KEY}) +_METADATA_KEYS = _ROLE_METADATA_KEYS | {"workflow"} _SIGNATURE = b"\x89PNG\r\n\x1a\n" @@ -26,6 +28,25 @@ class MetadataReadError(ValueError): """Prompt evidence is unreadable, incomplete, or exceeds inspection limits.""" +@dataclass(frozen=True) +class InspectionRead: + """One bounded read, with its transient exact video source kept private.""" + + metadata: dict[str, str] + record_formats: dict[str, str] + video_carrier_text: str | None + format: str | None + width: int | None + height: int | None + + +@dataclass(frozen=True) +class _VideoInspection: + metadata: dict[str, str | None] + record_formats: dict[str, str] + carrier_text: str | None + + def _inflate(data: bytes) -> bytes: inflater = zlib.decompressobj() result = inflater.decompress(data, MAX_METADATA_BYTES + 1) @@ -60,7 +81,7 @@ def _png_text(kind: bytes, data: bytes) -> tuple[str, str]: return key, body.decode(encoding) -def _png_metadata(stream: Any, size: int) -> dict: +def _png_metadata(stream: Any, size: int) -> tuple[dict, int, int]: """CRC-check text and structural chunks, seeking past compressed pixels. PNG text chunks may follow IDAT. Pillow's ``.text`` loads pixels to find @@ -89,10 +110,11 @@ def _png_metadata(stream: Any, size: int) -> dict: if saw_header or not all(struct.unpack(">II", data[:8])): raise MetadataReadError("Malformed PNG header") saw_header = True + width, height = struct.unpack(">II", data[:8]) elif kind == b"IEND": if length or not saw_pixels: raise MetadataReadError("Incomplete PNG image") - return found + return found, width, height else: key, value = _png_text(kind, data) expanded_size += len(value.encode("utf-8")) @@ -108,7 +130,7 @@ def _png_metadata(stream: Any, size: int) -> dict: raise MetadataReadError("PNG chunk count exceeds inspection limits") -def _video_metadata(path: Path) -> dict: +def _video_metadata(path: Path) -> _VideoInspection: probe = ffprobe_path() if not probe: raise MetadataReadError("Video metadata reader is unavailable") @@ -161,6 +183,7 @@ def _video_metadata(path: Path) -> dict: # The scan normalizer drops null fields and unknown keys. Inspection must # retain explicit prompt/execution markers so malformed evidence cannot # silently fall back to an unrelated saved recipe or another comment. + candidates = [] for key in ("comment", "description"): raw = tags.get(key) if not isinstance(raw, str): @@ -169,44 +192,73 @@ def _video_metadata(path: Path) -> dict: envelope = json.loads(raw) except ValueError: continue - if not isinstance(envelope, dict) or not _METADATA_KEYS.intersection(envelope): + if not isinstance(envelope, dict): continue - metadata = metadata_from_tags({key: raw}) + candidates.append((raw, envelope)) + selected = next( + (item for item in candidates if _ROLE_METADATA_KEYS.intersection(item[1])), + None, + ) + if selected is None: + selected = next((item for item in candidates if "workflow" in item[1]), None) + if selected is not None: + raw, envelope = selected + metadata, formats = {}, {} for field in _METADATA_KEYS.intersection(envelope): value = envelope[field] + if field == "workflow" and not isinstance(value, (str, dict)): + formats[field] = "invalid" + continue metadata[field] = ( value if value is None or isinstance(value, str) else json.dumps(value) ) - return metadata - return {} + formats[field] = ( + "original-text" if isinstance(value, str) else "reserialized-json" + ) + return _VideoInspection(metadata, formats, raw) + return _VideoInspection({}, {}, None) -def read_inspection_metadata(path: Path) -> dict: - """Return bounded, header-only metadata or raise a response-safe error.""" +def read_inspection_file(path: Path) -> InspectionRead: + """Read bounded metadata and header facts without decoding media pixels.""" try: + formats, carrier = {}, None + format_name = width = height = None if is_video_path(path): - metadata = _video_metadata(path) + video = _video_metadata(path) + metadata, formats, carrier = ( + video.metadata, + video.record_formats, + video.carrier_text, + ) else: with path.open("rb") as stream: signature = stream.read(8) if signature == _SIGNATURE: - return _png_metadata(stream, path.stat().st_size) - stream.seek(0) - header = stream.read(MAX_HEADER_BYTES) - with Image.open(io.BytesIO(header)) as image: - metadata = dict(image.info) + metadata, width, height = _png_metadata(stream, path.stat().st_size) + format_name = "PNG" + else: + stream.seek(0) + header = stream.read(MAX_HEADER_BYTES) + with Image.open(io.BytesIO(header)) as image: + metadata = dict(image.info) + format_name, (width, height) = image.format, image.size result, size = {}, 0 for key in _METADATA_KEYS: if key not in metadata: continue value = metadata[key] + if key == "workflow" and not isinstance(value, str): + formats[key] = "invalid" + continue if not isinstance(value, str): raise MetadataReadError("Malformed prompt metadata") - size += len(value.encode("utf-8")) + size += len(value.encode("utf-8", "surrogatepass")) if size > MAX_METADATA_BYTES: raise MetadataReadError("Metadata exceeds inspection limits") result[key] = value - return result + formats.setdefault(key, "original-text") + return InspectionRead(result, formats, carrier, format_name, width, height) except MetadataReadError: raise except ( @@ -221,12 +273,25 @@ def read_inspection_metadata(path: Path) -> dict: raise MetadataReadError("Media metadata could not be read") from None +def read_inspection_metadata(path: Path) -> dict[str, str]: + """Compatible metadata-only projection of the single inspection read.""" + return read_inspection_file(path).metadata + + def _side(status: str, text: str | None = None, source: str | None = None) -> dict: return {"status": status, "text": text, "source": source} -def inspect_prompt_roles(metadata: dict, saved: dict | None = None) -> dict: +def inspect_prompt_evidence(metadata: dict, saved: dict | None = None) -> dict: """Keep uncertainty visible; a saved association never replaces file evidence.""" + metadata = { + key: value for key, value in metadata.items() if key in _ROLE_METADATA_KEYS + } + resolution = { + "state": "unresolved", + "sampler_id": None, + "reason": "Prompt evidence is malformed", + } if not metadata: result = {role: _side("absent") for role in ("positive", "negative")} if saved: @@ -240,20 +305,35 @@ def inspect_prompt_roles(metadata: dict, saved: dict | None = None) -> dict: "verified against this file." ), } - return result + return { + "prompts": result, + "resolution": { + "state": "absent", + "sampler_id": None, + "reason": "No parsed prompt metadata", + }, + } canonical = dict(metadata) if "prompt" not in canonical and "parameters" not in canonical: - return {role: _side("unresolved") for role in ("positive", "negative")} + return { + "prompts": {role: _side("unresolved") for role in ("positive", "negative")}, + "resolution": resolution, + } if isinstance(canonical.get("prompt"), str): try: canonical["prompt"] = json.loads(canonical["prompt"]) except (ValueError, RecursionError): - return {role: _side("unresolved") for role in ("positive", "negative")} - resolution = resolve_prompt_pair(canonical) - state = resolution["state"] + return { + "prompts": { + role: _side("unresolved") for role in ("positive", "negative") + }, + "resolution": resolution, + } + resolved = resolve_prompt_pair(canonical) + state = resolved["state"] result = {} for role in ("positive", "negative"): - side = resolution[role] + side = resolved[role] if ( state != "ambiguous" and isinstance(side, dict) @@ -264,4 +344,16 @@ def inspect_prompt_roles(metadata: dict, saved: dict | None = None) -> dict: ) else: result[role] = _side(state if state != "resolved" else "unresolved") - return result + return { + "prompts": result, + "resolution": { + "state": state, + "sampler_id": resolved.get("sampler_id"), + "reason": resolved.get("reason", ""), + }, + } + + +def inspect_prompt_roles(metadata: dict, saved: dict | None = None) -> dict: + """Keep the established role-only helper contract.""" + return inspect_prompt_evidence(metadata, saved)["prompts"]