KillPhantom-ai-fashion-studio/scripts/validate_workflows.py

"""Validate AI Fashion Studio workflow JSON without importing ComfyUI."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import re

CUSTOM_NODE_TYPES = {
    "AFS_LoadProductAssets",
    "AFS_AnalyzeProductMock",
    "AFS_VirtualTryOnCandidatesMock",
    "AFS_EvaluateImageCandidatesMock",
    "AFS_RankCandidates",
    "AFS_ImageToVideoMock",
    "AFS_EvaluateVideoMock",
    "AFS_ExportRun",
    "AFS_AnalyzeProduct",
    "AFS_EvaluateImageCandidates",
    "AFS_SceneEdit",
}


def _unknown_custom_nodes(referenced_types: set) -> list[str]:
    """Flag AFS-prefixed node types this node pack does not register."""

    return sorted(
        node_type
        for node_type in referenced_types
        if isinstance(node_type, str)
        and node_type.startswith("AFS_")
        and node_type not in CUSTOM_NODE_TYPES
    )
SECRET_PATTERN = re.compile(
    r"(?:sk-[A-Za-z0-9_-]{16,}|bearer\s+[A-Za-z0-9._-]{16,}|api[_-]?key\s*[:=]\s*[^\s\"']+)",
    re.IGNORECASE,
)


def validate_api_workflow(path: Path) -> list[str]:
    path = Path(path)
    errors: list[str] = []
    raw = path.read_text(encoding="utf-8")
    if SECRET_PATTERN.search(raw):
        errors.append("possible embedded secret")
    try:
        workflow = json.loads(raw)
    except json.JSONDecodeError as exc:
        return [f"invalid JSON: {exc}"]
    if not isinstance(workflow, dict) or not workflow:
        return ["API workflow must be a non-empty object"]

    node_ids = set(workflow)
    referenced_types = set()
    has_output = False
    for node_id, node in workflow.items():
        if not isinstance(node, dict):
            errors.append(f"node {node_id} must be an object")
            continue
        class_type = node.get("class_type")
        referenced_types.add(class_type)
        if class_type == "AFS_ExportRun":
            has_output = True
        inputs = node.get("inputs")
        if not isinstance(inputs, dict):
            errors.append(f"node {node_id} is missing inputs")
            continue
        for name, value in inputs.items():
            if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str):
                if value[0] not in node_ids:
                    errors.append(f"node {node_id} input {name} references missing node {value[0]}")

    unknown_custom_nodes = _unknown_custom_nodes(referenced_types)
    if unknown_custom_nodes:
        errors.append(f"unknown custom nodes: {', '.join(unknown_custom_nodes)}")
    if not has_output:
        errors.append("workflow has no AFS_ExportRun output node")
    if len(workflow) > 100:
        errors.append("workflow exceeds 100 nodes")
    return errors


def validate_ui_workflow(path: Path) -> list[str]:
    path = Path(path)
    errors: list[str] = []
    raw = path.read_text(encoding="utf-8")
    if SECRET_PATTERN.search(raw):
        errors.append("possible embedded secret")
    try:
        workflow = json.loads(raw)
    except json.JSONDecodeError as exc:
        return [f"invalid JSON: {exc}"]
    nodes = workflow.get("nodes")
    links = workflow.get("links")
    if not isinstance(nodes, list) or not nodes:
        return ["UI workflow must contain nodes"]
    if not isinstance(links, list):
        return ["UI workflow must contain links"]
    node_ids = {node.get("id") for node in nodes if isinstance(node, dict)}
    if len(node_ids) != len(nodes):
        errors.append("UI workflow has duplicate or missing node IDs")
    for link in links:
        if not isinstance(link, list) or len(link) != 6:
            errors.append("UI workflow contains a malformed link")
            continue
        if link[1] not in node_ids or link[3] not in node_ids:
            errors.append(f"link {link[0]} references a missing node")
    referenced_types = {node.get("type") for node in nodes if isinstance(node, dict)}
    unknown_custom_nodes = _unknown_custom_nodes(referenced_types)
    if unknown_custom_nodes:
        errors.append(f"unknown custom nodes: {', '.join(unknown_custom_nodes)}")
    if "AFS_ExportRun" not in referenced_types:
        errors.append("workflow has no AFS_ExportRun output node")
    if len(nodes) > 100:
        errors.append("workflow exceeds 100 nodes")
    return errors


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "paths",
        nargs="*",
        type=Path,
        default=[
            Path("workflows/api/mock_vertical_slice.json"),
            Path("workflows/ui/mock_vertical_slice.json"),
            Path("workflows/ui/live_analysis_slice.json"),
        ],
    )
    args = parser.parse_args()
    failed = False
    for path in args.paths:
        data = json.loads(path.read_text(encoding="utf-8"))
        errors = (
            validate_ui_workflow(path)
            if isinstance(data, dict) and isinstance(data.get("nodes"), list)
            else validate_api_workflow(path)
        )
        if errors:
            failed = True
            print(f"FAIL {path}")
            for error in errors:
                print(f"  - {error}")
        else:
            print(f"OK   {path}")
    return 1 if failed else 0


if __name__ == "__main__":
    raise SystemExit(main())