diff --git a/README.md b/README.md index dc2389266..786a14166 100644 --- a/README.md +++ b/README.md @@ -364,7 +364,7 @@ For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step | Flag | Description | |------|-------------| | `--enable-manager` | Enable ComfyUI-Manager | -| `--enable-manager-legacy-ui` | Use the legacy manager UI instead of the new UI (requires `--enable-manager`) | +| `--enable-manager-legacy-ui` | Use the legacy manager UI instead of the new UI (implies `--enable-manager`) | | `--disable-manager-ui` | Disable the manager UI and endpoints while keeping background features like security checks and scheduled installation completion (requires `--enable-manager`) | @@ -462,16 +462,6 @@ To use the most up-to-date frontend version: This approach allows you to easily switch between the stable fortnightly release and the cutting-edge daily updates, or even specific versions for testing purposes. -### Accessing the Legacy Frontend - -If you need to use the legacy frontend for any reason, you can access it using the following command line argument: - -``` ---front-end-version Comfy-Org/ComfyUI_legacy_frontend@latest -``` - -This will use a snapshot of the legacy frontend preserved in the [ComfyUI Legacy Frontend repository](https://github.com/Comfy-Org/ComfyUI_legacy_frontend). - # QA ### Which GPU should I buy for this? diff --git a/comfy/cli_args.py b/comfy/cli_args.py index cba0dfa34..22f621cf5 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -133,7 +133,7 @@ upcast.add_argument("--dont-upcast-attention", action="store_true", help="Disabl parser.add_argument("--enable-manager", action="store_true", help="Enable the ComfyUI-Manager feature.") manager_group = parser.add_mutually_exclusive_group() manager_group.add_argument("--disable-manager-ui", action="store_true", help="Disables only the ComfyUI-Manager UI and endpoints. Scheduled installations and similar background tasks will still operate.") -manager_group.add_argument("--enable-manager-legacy-ui", action="store_true", help="Enables the legacy UI of ComfyUI-Manager") +manager_group.add_argument("--enable-manager-legacy-ui", action="store_true", help="Enables the legacy UI of ComfyUI-Manager. Implies --enable-manager.") vram_group = parser.add_mutually_exclusive_group() @@ -258,6 +258,10 @@ if args.disable_auto_launch: if args.force_fp16: args.fp16_unet = True +# '--enable-manager-legacy-ui' is meaningless unless the manager is enabled, so imply '--enable-manager'. +if args.enable_manager_legacy_ui: + args.enable_manager = True + # '--fast' is not provided, use an empty set if args.fast is None: diff --git a/comfy/model_base.py b/comfy/model_base.py index 2289e0812..ab4a11022 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -1816,7 +1816,24 @@ class WAN21_SCAIL2(WAN21_SCAIL): def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]): if cond_key in ("sam_latents", "pose_latents"): - return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1) + # Return sliced view omitting retain_index_list + return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=0) + if cond_key == "ref_mask_latents" and hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor): + # The ref mask is just a single frame padded with frames of zeros, so just grab the first frames for all windows + full_ref_mask = cond_value.cond + video_frame_count = x_in.shape[2] + if full_ref_mask.shape[2] != video_frame_count + 1: + return None + window_length = len(window.index_list) + + # Account for the causal anchor frame if it exists + anchor_index = getattr(window, "causal_anchor_index", None) + if anchor_index is not None and anchor_index >= 0: + window_length += 1 + + window_ref_mask = full_ref_mask[:, :, :window_length + 1].to(device) + return cond_value._copy_with(window_ref_mask) + return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list) def concat_cond(self, **kwargs): diff --git a/comfy/ops.py b/comfy/ops.py index 119177c37..3c9912aae 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -299,21 +299,21 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of non_blocking = comfy.model_management.device_supports_non_blocking(device) - if hasattr(s, "_v"): + if hasattr(s, "_v") and comfy.model_management.is_device_cpu(device): #vbar doesn't support CPU weights, but some custom nodes have weird paths #that might switch the layer to the CPU and expect it to work. We have to take #a clone conservatively as we are mmapped and some SFT files are packed misaligned #If you are a custom node author reading this, please move your layer to the GPU #or declare your ModelPatcher as CPU in the first place. - if comfy.model_management.is_device_cpu(device): - materialize_meta_param(s, ["weight", "bias"]) - weight = s.weight.to(dtype=dtype, copy=True) - if isinstance(weight, QuantizedTensor): - weight = weight.dequantize() - bias = s.bias.to(dtype=bias_dtype, copy=True) if s.bias is not None else None - return format_return((weight, bias, (None, None, None)), offloadable) + materialize_meta_param(s, ["weight", "bias"]) + weight = s.weight.to(dtype=dtype, copy=True) + if isinstance(weight, QuantizedTensor): + weight = weight.dequantize() + bias = s.bias.to(dtype=bias_dtype, copy=True) if s.bias is not None else None + return format_return((weight, bias, (None, None, None)), offloadable) + elif hasattr(s, "_v") and s.weight.device != device: prefetched = hasattr(s, "_prefetch") offload_stream = None offload_device = None diff --git a/comfy_api_nodes/apis/__init__.py b/comfy_api_nodes/apis/__init__.py index 9c4cfb9b6..9a7049ea2 100644 --- a/comfy_api_nodes/apis/__init__.py +++ b/comfy_api_nodes/apis/__init__.py @@ -1310,13 +1310,6 @@ class KlingTaskStatus(str, Enum): failed = 'failed' -class KlingTextToVideoModelName(str, Enum): - kling_v1 = 'kling-v1' - kling_v1_6 = 'kling-v1-6' - kling_v2_1_master = 'kling-v2-1-master' - kling_v2_5_turbo = 'kling-v2-5-turbo' - - class KlingVideoGenAspectRatio(str, Enum): field_16_9 = '16:9' field_9_16 = '9:16' @@ -5179,7 +5172,7 @@ class KlingText2VideoRequest(BaseModel): duration: Optional[KlingVideoGenDuration] = '5' external_task_id: Optional[str] = Field(None, description='Customized Task ID') mode: Optional[KlingVideoGenMode] = 'std' - model_name: Optional[KlingTextToVideoModelName] = 'kling-v1' + model_name: Optional[str] = 'kling-v1' negative_prompt: Optional[str] = Field( None, description='Negative text prompt', max_length=2500 ) diff --git a/comfy_api_nodes/nodes_kling.py b/comfy_api_nodes/nodes_kling.py index d11e42540..c81d3503d 100644 --- a/comfy_api_nodes/nodes_kling.py +++ b/comfy_api_nodes/nodes_kling.py @@ -436,7 +436,7 @@ async def execute_text2video( negative_prompt=negative_prompt if negative_prompt else None, duration=KlingVideoGenDuration(duration), mode=KlingVideoGenMode(model_mode), - model_name=KlingVideoGenModelName(model_name), + model_name=model_name, cfg_scale=cfg_scale, aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio), camera_control=camera_control, diff --git a/comfy_api_nodes/nodes_openai.py b/comfy_api_nodes/nodes_openai.py index 0fe5fb9d0..ad62f2164 100644 --- a/comfy_api_nodes/nodes_openai.py +++ b/comfy_api_nodes/nodes_openai.py @@ -9,6 +9,7 @@ from PIL import Image from typing_extensions import override import folder_paths +from comfy.utils import common_upscale from comfy_api.latest import IO, ComfyExtension, Input from comfy_api_nodes.apis.openai import ( InputFileContent, @@ -62,7 +63,8 @@ async def validate_and_cast_response(response, timeout: int = None) -> torch.Ten timeout: Request timeout in seconds. Defaults to None (no timeout). Returns: - A torch.Tensor representing the image (1, H, W, C). + A torch.Tensor of shape (N, H, W, C) with all returned images; images whose + dimensions differ from the first image's are resized to match it. Raises: ValueError: If the response is not valid. @@ -89,6 +91,14 @@ async def validate_and_cast_response(response, timeout: int = None) -> torch.Ten arr = np.asarray(pil_img).astype(np.float32) / 255.0 image_tensors.append(torch.from_numpy(arr)) + # With size="auto" the API can return images whose dimensions differ by a few pixels within a single response + # resize them to the first image's dimensions so they can be stacked into one batch. + ref_h, ref_w = image_tensors[0].shape[:2] + for i, t in enumerate(image_tensors): + if t.shape[:2] != (ref_h, ref_w): + samples = t.unsqueeze(0).movedim(-1, 1) + samples = common_upscale(samples, ref_w, ref_h, "bilinear", "center") + image_tensors[i] = samples.movedim(1, -1).squeeze(0) return torch.stack(image_tensors, dim=0) diff --git a/comfy_execution/asset_enrichment.py b/comfy_execution/asset_enrichment.py new file mode 100644 index 000000000..38e9496a8 --- /dev/null +++ b/comfy_execution/asset_enrichment.py @@ -0,0 +1,66 @@ +"""Enrich executed-node output entries with asset id.""" +import logging +import os + + +def enrich_output_with_assets(output_ui: dict) -> dict: + """Register file-type output entries as assets and inject their ``id``. + + Runs at output-processing time, once per produced output, when + --enable-assets is set. Returns a new dict; entries without a resolvable + on-disk file path are left unchanged. Errors are caught per-entry so a + failure never blocks execution or the other entries. + """ + from comfy.cli_args import args + if not args.enable_assets: + return output_ui + + import folder_paths + from app.assets.services.ingest import register_file_in_place, DependencyMissingError + + enriched = {} + for key, entries in output_ui.items(): + if not isinstance(entries, list): + enriched[key] = entries + continue + new_entries = [] + for entry in entries: + if not isinstance(entry, dict) or "filename" not in entry or "type" not in entry: + new_entries.append(entry) + continue + try: + base = folder_paths.get_directory_by_type(entry["type"]) + if base is None: + new_entries.append(entry) + continue + base_abs = os.path.abspath(base) + abs_path = os.path.abspath(os.path.join(base_abs, entry.get("subfolder") or "", entry["filename"])) + try: + if os.path.commonpath([base_abs, abs_path]) != base_abs: + raise ValueError("escapes base") + except ValueError: + logging.warning("Asset enrichment skipped (path escapes base): %s", entry.get("filename")) + new_entries.append(entry) + continue + if not os.path.isfile(abs_path): + new_entries.append(entry) + continue + + # Register unconditionally: the file was just produced, and + # register_file_in_place re-hashes so an overwritten path can + # never carry a stale id. + result = register_file_in_place( + abs_path=abs_path, + name=entry["filename"], + tags=[entry["type"]], + ) + + entry = dict(entry) + entry["id"] = result.ref.id + except DependencyMissingError: + logging.warning("Asset enrichment skipped (blake3 not available): %s", entry.get("filename")) + except Exception: + logging.warning("Failed to enrich output entry with asset id: %s", entry.get("filename"), exc_info=True) + new_entries.append(entry) + enriched[key] = new_entries + return enriched diff --git a/comfy_execution/jobs.py b/comfy_execution/jobs.py index fcd7ef735..20ebae155 100644 --- a/comfy_execution/jobs.py +++ b/comfy_execution/jobs.py @@ -3,6 +3,7 @@ Job utilities for the /api/jobs endpoint. Provides normalization and helper functions for job status tracking. """ +import uuid from typing import Optional from comfy_api.internal import prune_dict @@ -19,6 +20,25 @@ class JobStatus: ALL = [PENDING, IN_PROGRESS, COMPLETED, FAILED, CANCELLED] +def validate_job_id(value) -> str: + """Validate a client-supplied job (prompt) id. + + Job ids must be UUIDs in the canonical lowercase hyphenated form. The id + is stored and compared verbatim everywhere downstream — history keys, + websocket events, and /interrupt matching — so accepting another spelling + would silently rewrite the client's id and then miss every exact-match + lookup. Rejecting loudly beats that. + + Returns the id unchanged. Raises ValueError when the value is not a + string in canonical UUID form. + """ + if not isinstance(value, str): + raise ValueError(f"job id must be a string, got {type(value).__name__}") + if str(uuid.UUID(value)) != value: + raise ValueError("job id must be a UUID in canonical lowercase hyphenated form") + return value + + # Media types that can be previewed in the frontend PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'}) diff --git a/comfy_extras/nodes_scail.py b/comfy_extras/nodes_scail.py index a740442de..bba0942d7 100644 --- a/comfy_extras/nodes_scail.py +++ b/comfy_extras/nodes_scail.py @@ -267,7 +267,8 @@ class SCAIL2ColoredMask(io.ComfyNode): io.Combo.Input("sort_by", options=["none", "left_to_right", "area"], default="left_to_right", tooltip="Order in which palette colors are assigned to the tracked objects (applied to both reference and pose video so each identity keeps the same color). left_to_right = leftmost object (by first-frame centroid) gets the first color; area = biggest object (by first-frame mask area) gets the first color; none = keep SAM3's order."), io.Boolean.Input("replacement_mode", default=False, - tooltip="False = mask_video has black bg (Animation Mode). True = white bg (Replacement Mode). Set the matching replacement_mode on WanSCAILToVideo. reference_image_mask is always black-bg regardless."), + tooltip="False = Animation Mode (pose_video_mask has black background, reference_image_mask has white background). " + "True = Replacement Mode (pose_video_mask has white background, reference_image_mask has black background)."), ], outputs=[ io.Image.Output("pose_video_mask"), @@ -296,14 +297,17 @@ class SCAIL2ColoredMask(io.ComfyNode): return td drv = _prep(driving_track_data) + # Animation: driving=black, ref=white. Replacement: driving=white, ref=black. mask_video = _render_colored_masks(drv, "white" if replacement_mode else "black") + ref_bg = "black" if replacement_mode else "white" if ref_track_data is not None: ref = _prep(ref_track_data) - reference_image_mask = _render_colored_masks(ref, "black") + reference_image_mask = _render_colored_masks(ref, ref_bg) else: H, W = drv["orig_size"] - reference_image_mask = torch.zeros(1, H, W, 3, device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype()) + fill_value = 1.0 if ref_bg == "white" else 0.0 + reference_image_mask = torch.full((1, H, W, 3), fill_value, device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype()) return io.NodeOutput(mask_video, reference_image_mask) diff --git a/execution.py b/execution.py index 32b8b2e01..9e16e451d 100644 --- a/execution.py +++ b/execution.py @@ -40,6 +40,7 @@ from comfy_execution.graph_utils import GraphBuilder, is_link from comfy_execution.validation import validate_node_input from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler from comfy_execution.utils import CurrentNodeContext +from comfy_execution.asset_enrichment import enrich_output_with_assets from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func from comfy_api.latest import io, _io from comfy_execution.cache_provider import _has_cache_providers, _get_cache_providers, _logger as _cache_logger @@ -422,6 +423,7 @@ def _is_intermediate_output(dynprompt, node_id): class_def = nodes.NODE_CLASS_MAPPINGS[class_type] return getattr(class_def, 'HAS_INTERMEDIATE_OUTPUT', False) + def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs): if server.client_id is None: return @@ -556,6 +558,10 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed, asyncio.create_task(await_completion()) return (ExecutionResult.PENDING, None, None) if len(output_ui) > 0: + # Enrich at output-processing time (not in the send path) so assets + # are registered even when no client is connected, and the asset id + # flows into ui_outputs and the cache alongside the raw entries. + output_ui = enrich_output_with_assets(output_ui) ui_outputs[unique_id] = { "meta": { "node_id": unique_id, diff --git a/openapi.yaml b/openapi.yaml index c27ed7adf..6e203b1cd 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -896,6 +896,11 @@ components: additionalProperties: true description: The workflow graph to execute type: object + prompt_id: + description: Optional client-supplied job id. Must be a UUID in canonical lowercase hyphenated form; it is echoed back in the response. Omitted or null means the server generates one. + format: uuid + nullable: true + type: string workflow_id: description: UUID identifying the cloud workflow entity to associate with this job type: string @@ -1062,6 +1067,9 @@ components: comfyui_version: description: ComfyUI version type: string + deploy_environment: + description: How this ComfyUI instance is deployed (e.g. cloud, local-git, local-portable, local-desktop) + type: string embedded_python: description: Whether using embedded Python type: boolean diff --git a/server.py b/server.py index ba00f99dd..ccc92e5ab 100644 --- a/server.py +++ b/server.py @@ -8,7 +8,7 @@ import time import nodes import folder_paths import execution -from comfy_execution.jobs import JobStatus, get_job, get_all_jobs +from comfy_execution.jobs import JobStatus, get_job, get_all_jobs, validate_job_id import uuid import urllib import json @@ -942,7 +942,21 @@ class PromptServer(): if "prompt" in json_data: prompt = json_data["prompt"] - prompt_id = str(json_data.get("prompt_id", uuid.uuid4())) + client_prompt_id = json_data.get("prompt_id") + if client_prompt_id is None: + # Absent or explicit null: the server mints the id. + prompt_id = str(uuid.uuid4()) + else: + try: + prompt_id = validate_job_id(client_prompt_id) + except ValueError: + error = { + "type": "invalid_prompt_id", + "message": "prompt_id must be a valid UUID", + "details": "prompt_id must be a UUID string in canonical lowercase hyphenated form; omit it to let the server generate one", + "extra_info": {} + } + return web.json_response({"error": error, "node_errors": {}}, status=400) partial_execution_targets = None if "partial_execution_targets" in json_data: diff --git a/tests-unit/assets_test/test_prompt_id_enforcement.py b/tests-unit/assets_test/test_prompt_id_enforcement.py new file mode 100644 index 000000000..86a755c9f --- /dev/null +++ b/tests-unit/assets_test/test_prompt_id_enforcement.py @@ -0,0 +1,69 @@ +"""POST /prompt enforces canonical-UUID job ids at creation time. + +Lives in assets_test because it uses this suite's booted-server fixture. The +invariant itself is pipeline-wide: a job id is stored and compared verbatim +downstream — history keys, websocket correlation, and /interrupt matching — +so a job minted with a non-canonical id would miss every exact-match lookup. + +The prompt bodies here are intentionally invalid workflows — prompt_id +validation happens before workflow validation, so a rejected id returns +``invalid_prompt_id`` while an accepted id falls through to the ordinary +workflow-validation error (proving it cleared the id check). +""" +import requests + + +def _post_prompt(http: requests.Session, api_base: str, body: dict) -> requests.Response: + return http.post(api_base + "/prompt", json=body, timeout=30) + + +def _error_type(r: requests.Response) -> str: + return r.json()["error"]["type"] + + +def test_non_uuid_prompt_id_rejected(http: requests.Session, api_base: str): + r = _post_prompt(http, api_base, {"prompt": {}, "prompt_id": "not-a-uuid"}) + assert r.status_code == 400, r.text + assert _error_type(r) == "invalid_prompt_id" + + +def test_non_string_prompt_id_rejected(http: requests.Session, api_base: str): + # Previously str()-coerced (123 became the job id "123"); must now be a 400, + # not a 500 from uuid.UUID choking on a non-string. + r = _post_prompt(http, api_base, {"prompt": {}, "prompt_id": 123}) + assert r.status_code == 400, r.text + assert _error_type(r) == "invalid_prompt_id" + + +def test_non_canonical_uuid_rejected(http: requests.Session, api_base: str): + # Parseable as a UUID, but not the canonical lowercase form: rejected + # loudly rather than silently rewritten (downstream lookups match the + # stored id exactly). + r = _post_prompt( + http, + api_base, + {"prompt": {}, "prompt_id": "AAAAAAAA-BBBB-4CCC-8DDD-EEEEEEEEEEEE"}, + ) + assert r.status_code == 400, r.text + assert _error_type(r) == "invalid_prompt_id" + + +def test_canonical_uuid_accepted(http: requests.Session, api_base: str): + # The id clears validation; the empty workflow then fails ordinary prompt + # validation, proving the request got past the id check. + r = _post_prompt( + http, + api_base, + {"prompt": {}, "prompt_id": "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee"}, + ) + assert r.status_code == 400, r.text + assert _error_type(r) != "invalid_prompt_id" + + +def test_null_prompt_id_not_rejected(http: requests.Session, api_base: str): + # Explicit null means "server generates" and must not be rejected as an + # invalid id. (The minted id itself is not observable here because the + # workflow is invalid; unit tests cover validate_job_id directly.) + r = _post_prompt(http, api_base, {"prompt": {}, "prompt_id": None}) + assert r.status_code == 400, r.text + assert _error_type(r) != "invalid_prompt_id" diff --git a/tests-unit/execution_test/test_enrich_output.py b/tests-unit/execution_test/test_enrich_output.py new file mode 100644 index 000000000..61490c49e --- /dev/null +++ b/tests-unit/execution_test/test_enrich_output.py @@ -0,0 +1,205 @@ +"""Tests for enrich_output_with_assets in comfy_execution/asset_enrichment.py.""" +import os +import types +import unittest +from unittest.mock import MagicMock, patch + + +def _make_args(enable_assets: bool): + a = types.SimpleNamespace() + a.enable_assets = enable_assets + return a + + +def _make_register_result(ref_id="ref-id-2"): + result = MagicMock() + result.ref.id = ref_id + return result + + +# Platform-appropriate absolute base. tempfile.gettempdir() returns C:\... on +# Windows and /tmp on POSIX, so containment via commonpath behaves naturally. +_DEFAULT_BASE = os.path.join(__import__("tempfile").gettempdir(), "asset-enrichment-test-base") + + +def _mocked_modules(*, enable_assets=True, register_file_in_place=None, directory=_DEFAULT_BASE): + return { + "comfy.cli_args": MagicMock(args=_make_args(enable_assets)), + "folder_paths": MagicMock(get_directory_by_type=MagicMock(return_value=directory)), + "app.assets.services.ingest": MagicMock( + register_file_in_place=register_file_in_place or MagicMock(return_value=_make_register_result()), + DependencyMissingError=type("DependencyMissingError", (Exception,), {}), + ), + } + + +def _call(output_ui, *, enable_assets=True, file_exists=True, register_result=None, directory=_DEFAULT_BASE): + register_mock = MagicMock(return_value=register_result or _make_register_result()) + mocked = _mocked_modules( + enable_assets=enable_assets, + register_file_in_place=register_mock, + directory=directory, + ) + + # Only os.path.isfile is patched — abspath/join must run natively so the + # containment check sees real platform paths. + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=file_exists): + import importlib + import comfy_execution.asset_enrichment as mod + importlib.reload(mod) + return mod.enrich_output_with_assets(output_ui) + + +class TestEnrichOutputWithAssets(unittest.TestCase): + + def test_disabled_returns_unchanged(self): + output = {"images": [{"filename": "a.png", "subfolder": "", "type": "output"}]} + result = _call(output, enable_assets=False) + self.assertNotIn("id", result["images"][0]) + + def test_non_list_value_passed_through(self): + output = {"text": "hello"} + result = _call(output) + self.assertEqual(result["text"], "hello") + + def test_entry_without_filename_unchanged(self): + output = {"latent": [{"subfolder": "", "type": "output"}]} + result = _call(output) + self.assertNotIn("id", result["latent"][0]) + + def test_entry_without_type_unchanged(self): + output = {"data": [{"filename": "a.png", "subfolder": ""}]} + result = _call(output) + self.assertNotIn("id", result["data"][0]) + + def test_file_not_on_disk_unchanged(self): + output = {"images": [{"filename": "missing.png", "subfolder": "", "type": "output"}]} + result = _call(output, file_exists=False) + self.assertNotIn("id", result["images"][0]) + + def test_unknown_type_returns_none_directory_unchanged(self): + output = {"images": [{"filename": "a.png", "subfolder": "", "type": "unknown"}]} + result = _call(output, directory=None) + self.assertNotIn("id", result["images"][0]) + + def test_register_injects_only_id(self): + reg = _make_register_result(ref_id="inline-ref") + output = {"images": [{"filename": "new.png", "subfolder": "", "type": "output"}]} + result = _call(output, register_result=reg) + img = result["images"][0] + self.assertEqual(img["id"], "inline-ref") + # Only id is injected — no asset_hash, name, preview_url, size + self.assertNotIn("asset_hash", img) + self.assertNotIn("name", img) + self.assertNotIn("preview_url", img) + self.assertNotIn("size", img) + + def test_register_called_per_entry(self): + register_mock = MagicMock(return_value=_make_register_result()) + mocked = _mocked_modules(register_file_in_place=register_mock) + output = { + "images": [ + {"filename": "a.png", "subfolder": "", "type": "output"}, + {"filename": "b.png", "subfolder": "", "type": "output"}, + ] + } + + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=True): + import importlib + import comfy_execution.asset_enrichment as mod + importlib.reload(mod) + mod.enrich_output_with_assets(output) + + self.assertEqual(register_mock.call_count, 2) + + def test_original_entry_not_mutated(self): + orig = {"filename": "a.png", "subfolder": "", "type": "output"} + output = {"images": [orig]} + _call(output) + self.assertNotIn("id", orig) + + def test_enrichment_error_does_not_block_sibling_entries(self): + call_count = [0] + good_reg = _make_register_result(ref_id="good-ref") + + def register_side_effect(abs_path, name, tags): + call_count[0] += 1 + if call_count[0] == 1: + raise RuntimeError("boom") + return good_reg + + mocked = _mocked_modules(register_file_in_place=register_side_effect) + + output = { + "images": [ + {"filename": "bad.png", "subfolder": "", "type": "output"}, + {"filename": "good.png", "subfolder": "", "type": "output"}, + ] + } + + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=True): + import importlib + import comfy_execution.asset_enrichment as mod + importlib.reload(mod) + result = mod.enrich_output_with_assets(output) + + imgs = result["images"] + self.assertNotIn("id", imgs[0]) + self.assertEqual(imgs[1]["id"], "good-ref") + + def test_multiple_output_keys_all_enriched(self): + output = { + "images": [{"filename": "a.png", "subfolder": "", "type": "output"}], + "videos": [{"filename": "b.mp4", "subfolder": "", "type": "output"}], + } + result = _call(output) + self.assertIn("id", result["images"][0]) + self.assertIn("id", result["videos"][0]) + + def test_none_entry_in_list_unchanged(self): + output = {"images": [None, {"filename": "a.png", "subfolder": "", "type": "output"}]} + result = _call(output) + self.assertIsNone(result["images"][0]) + self.assertIn("id", result["images"][1]) + + def test_path_traversal_subfolder_skipped(self): + register_mock = MagicMock(return_value=_make_register_result()) + mocked = _mocked_modules(register_file_in_place=register_mock) + + output = {"images": [{"filename": "passwd", "subfolder": "../../etc", "type": "output"}]} + + # Do NOT patch os.path.abspath — real resolution is required for the containment check. + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=True): + import importlib + import comfy_execution.asset_enrichment as mod + importlib.reload(mod) + result = mod.enrich_output_with_assets(output) + + self.assertNotIn("id", result["images"][0]) + register_mock.assert_not_called() + + def test_absolute_filename_skipped(self): + register_mock = MagicMock(return_value=_make_register_result()) + mocked = _mocked_modules(register_file_in_place=register_mock) + + # Absolute filename — os.path.join discards earlier components when a later one is absolute. + absolute_filename = os.path.abspath(os.sep + "etc" + os.sep + "passwd") + output = {"images": [{"filename": absolute_filename, "subfolder": "", "type": "output"}]} + + with patch.dict("sys.modules", mocked), \ + patch("os.path.isfile", return_value=True): + import importlib + import comfy_execution.asset_enrichment as mod + importlib.reload(mod) + result = mod.enrich_output_with_assets(output) + + self.assertNotIn("id", result["images"][0]) + register_mock.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/execution/test_jobs.py b/tests/execution/test_jobs.py index 814af5c13..f7cb612e4 100644 --- a/tests/execution/test_jobs.py +++ b/tests/execution/test_jobs.py @@ -1,5 +1,7 @@ """Unit tests for comfy_execution/jobs.py""" +import pytest + from comfy_execution.jobs import ( JobStatus, is_previewable, @@ -10,9 +12,50 @@ from comfy_execution.jobs import ( get_outputs_summary, apply_sorting, has_3d_extension, + validate_job_id, ) +class TestValidateJobId: + """validate_job_id guards job creation: POST /prompt rejects ids it raises on.""" + + def test_canonical_form_passes_through(self): + cid = "a1b2c3d4-e5f6-7a89-b0c1-d2e3f4a5b6c7" + assert validate_job_id(cid) == cid + + @pytest.mark.parametrize( + "variant", + [ + "A1B2C3D4-E5F6-7A89-B0C1-D2E3F4A5B6C7", # uppercase + "{a1b2c3d4-e5f6-7a89-b0c1-d2e3f4a5b6c7}", # braced + "urn:uuid:a1b2c3d4-e5f6-7a89-b0c1-d2e3f4a5b6c7", # URN + "a1b2c3d4e5f67a89b0c1d2e3f4a5b6c7", # bare hex + " a1b2c3d4-e5f6-7a89-b0c1-d2e3f4a5b6c7 ", # padded + ], + ) + def test_non_canonical_spellings_rejected(self, variant): + # uuid.UUID parses all of these, but accepting them would silently + # rewrite the client's id (history keys, websocket events, and + # /interrupt matching all match the stored form exactly). + with pytest.raises(ValueError): + validate_job_id(variant) + + @pytest.mark.parametrize( + "bad", + ["", "not-a-uuid", "prompt-123", "a1b2c3d4-e5f6-7a89-b0c1", "None"], + ) + def test_non_uuid_strings_rejected(self, bad): + with pytest.raises(ValueError): + validate_job_id(bad) + + @pytest.mark.parametrize("bad", [123, 1.5, True, None, ["a"], {"id": "x"}]) + def test_non_strings_rejected(self, bad): + # uuid.UUID raises AttributeError/TypeError on non-strings; the helper + # must normalize those to ValueError so callers need one except clause. + with pytest.raises(ValueError): + validate_job_id(bad) + + class TestJobStatus: """Test JobStatus constants."""