auto_setup.py

"""Auto-setup: install missing custom nodes + download missing models for a workflow.

Resolution sources:
- Custom nodes: ``api.comfy.org/nodes/<cnr_id>`` first (using ``cnr_id`` + ``ver``
  carried on each node's ``properties``), then ComfyUI Manager's
  ``custom-node-list.json`` as a fuzzy fallback by class_type.
- Models: the URL already carried on the workflow's ``properties.models`` (the
  resolver in ``nodes.py`` populates this), with Manager's ``model-list.json``
  as a filename-keyed fallback.

Cross-platform: Python tooling is always invoked as ``sys.executable -m pip``
/ ``sys.executable install.py`` (never a bare ``pip``/``python``, which aren't
on PATH for ComfyUI Desktop / uv-managed / portable installs). ``git`` is
located via :func:`_resolve_git` (PATH + common per-OS install dirs); when it's
genuinely absent we fall back to downloading a GitHub source tarball.
"""
from __future__ import annotations

import asyncio
import functools
import hashlib
import io
import logging
import os
import re
import shutil
import subprocess
import sys
import tarfile
import tempfile
import urllib.error
import urllib.request
import uuid
from pathlib import Path
from typing import Any, Awaitable, Callable, Iterable

import aiohttp
import folder_paths

from .nodes import RunflowDeploy, _iter_all_nodes, resolve_workflow_models

logger = logging.getLogger(__name__)


REGISTRY_NODE_URL = "https://api.comfy.org/nodes/{cnr_id}"
REGISTRY_COMFY_NODE_URL = "https://api.comfy.org/comfy-nodes/{class_type}/node"
MANAGER_CUSTOM_NODES_URL = (
    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/custom-node-list.json"
)
MANAGER_MODELS_URL = (
    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/model-list.json"
)
MANAGER_EXT_NODE_MAP_URL = (
    "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/extension-node-map.json"
)
HUGGINGFACE_SEARCH_URL = "https://huggingface.co/api/models"
HUGGINGFACE_TREE_URL = "https://huggingface.co/api/models/{repo}/tree/main"


def _comfy_core_class_types() -> set[str]:
    """Class types that come from ComfyUI itself (and any already-installed
    custom nodes). Anything in this set is already runnable on this install.
    """
    try:
        import nodes as comfy_nodes  # ComfyUI's own module
        return set(comfy_nodes.NODE_CLASS_MAPPINGS.keys())
    except Exception:
        return set()


# Stock frontend/editor pseudo-nodes that have no Python class in
# NODE_CLASS_MAPPINGS — they're rendered by the editor and rewired by the
# executor at run time. Without this set the resolver treats them as missing
# custom nodes and dumps them into ``unresolved``.
_STOCK_FRONTEND_NODES: frozenset[str] = frozenset(
    {
        "Reroute",
        "PrimitiveNode",
        "Note",
        "MarkdownNote",
    }
)


# Module-top-level `from comfy_env import ... install ...` is the contract
# opt-in for nodes that delegate their setup to the `comfy_env` package
# (Pozzetti 3D-pipeline family: SAM3DObjects, GeometryPack, MoGe2, TRELLIS2,
# HYPano2, Pixal3D, Hunyuan3D-Part, plus the comfyui-sharp pack). Matching
# on the import (not on install.py existence alone) keeps us from running
# legacy install.py scripts that do arbitrary work in nodes that happen to
# ship one for unrelated reasons.
#
# Mirrors the deploy worker's identical guard at
# `bg-brain/workers/comfyui-deploy-worker/installer.py:_is_comfy_env_install`.
_COMFY_ENV_INSTALL_RE = re.compile(
    r"^\s*from\s+comfy_env\s+import\b[^\n#]*\binstall\b",
    re.MULTILINE,
)


def _is_comfy_env_install(install_py: Path) -> bool:
    """True iff ``install.py`` exists in the cloned node directory and its source
    imports ``install`` from ``comfy_env`` at module top level. Read errors → False."""
    if not install_py.is_file():
        return False
    try:
        text = install_py.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return False
    return _COMFY_ENV_INSTALL_RE.search(text) is not None


def _aux_id_to_origin(aux_id: str) -> str | None:
    """Convert a ComfyUI Manager-style ``aux_id`` to an HTTPS clone origin.

    `properties.aux_id` is what ComfyUI's frontend stamps onto nodes installed
    via Manager's git-URL flow (vs. Comfy Registry / cnr_id). Most often it's
    a plain ``"owner/repo"`` string, occasionally a full ``https://github.com/...``
    URL. Returns ``None`` for unrecognized shapes so the caller can fall through
    to the next resolution path.
    """
    s = aux_id.strip()
    if not s:
        return None
    if s.startswith(("http://", "https://")):
        return s.rstrip("/")
    # Plain `owner/repo` — be strict (exactly one slash, no shell metacharacters)
    # so we don't build a malformed URL from corrupted properties.
    if s.count("/") == 1 and all(c not in s for c in " \t?#&"):
        owner, repo = s.split("/")
        if owner and repo:
            return f"https://github.com/{owner}/{repo}"
    return None


# ---------------------------------------------------------------------------
# Registry client (session-scoped, in-memory)
# ---------------------------------------------------------------------------


class RegistryClient:
    """Looks up custom-node provenance and model URLs from public registries.

    All caches are per-instance — callers re-use a single client for the
    duration of a setup job so the two checkboxes don't re-fetch the same
    JSON twice.
    """

    def __init__(self) -> None:
        self._node_cache: dict[str, dict | None] = {}
        self._class_type_cache: dict[str, dict | None] = {}
        self._manager_nodes: list[dict] | None = None
        self._manager_models: list[dict] | None = None
        self._ext_node_map: dict[str, list[str]] | None = None
        self._hf_cache: dict[str, dict | None] = {}
        self._lock = asyncio.Lock()

    async def _get_json(self, session: aiohttp.ClientSession, url: str) -> Any:
        async with session.get(url, timeout=aiohttp.ClientTimeout(total=20)) as resp:
            resp.raise_for_status()
            return await resp.json(content_type=None)

    async def resolve_by_cnr_id(self, cnr_id: str) -> dict | None:
        """Return the Comfy registry record for ``cnr_id``, or None on miss/error.

        The shape of interest is ``{"repository": "https://github.com/..."}``.
        """
        if not cnr_id or cnr_id == "comfy-core":
            return None
        if cnr_id in self._node_cache:
            return self._node_cache[cnr_id]
        try:
            async with aiohttp.ClientSession() as session:
                data = await self._get_json(session, REGISTRY_NODE_URL.format(cnr_id=cnr_id))
        except Exception as err:
            logger.info("comfy registry lookup failed for %s: %s", cnr_id, err)
            self._node_cache[cnr_id] = None
            return None
        self._node_cache[cnr_id] = data if isinstance(data, dict) else None
        return self._node_cache[cnr_id]

    async def _load_manager_nodes(self) -> list[dict]:
        async with self._lock:
            if self._manager_nodes is not None:
                return self._manager_nodes
            try:
                async with aiohttp.ClientSession() as session:
                    data = await self._get_json(session, MANAGER_CUSTOM_NODES_URL)
                self._manager_nodes = data.get("custom_nodes") or []
            except Exception as err:
                logger.info("manager custom-node-list fetch failed: %s", err)
                self._manager_nodes = []
            return self._manager_nodes

    async def _load_ext_node_map(self) -> dict[str, list[str]]:
        async with self._lock:
            if self._ext_node_map is not None:
                return self._ext_node_map
            try:
                async with aiohttp.ClientSession() as session:
                    data = await self._get_json(session, MANAGER_EXT_NODE_MAP_URL)
                # Shape: {repo_url: [[class_type, ...], {...metadata...}]}
                norm: dict[str, list[str]] = {}
                for repo_url, value in (data or {}).items():
                    if isinstance(value, list) and value and isinstance(value[0], list):
                        norm[repo_url] = [str(x) for x in value[0]]
                self._ext_node_map = norm
            except Exception as err:
                logger.info("manager extension-node-map fetch failed: %s", err)
                self._ext_node_map = {}
            return self._ext_node_map

    async def resolve_by_class_type_registry(self, class_type: str) -> dict | None:
        """Reverse-lookup a class_type against the Comfy registry.

        Endpoint: ``api.comfy.org/comfy-nodes/<class_type>/node`` returns the
        canonical custom-node record for whichever node provides this class.
        Built-in ComfyUI types 404 here — that's also the right signal: callers
        check ``loaded_class_types`` first, so reaching this path for a built-in
        means it wasn't loaded, and "not in registry" is a useful answer.
        """
        if not class_type:
            return None
        if class_type in self._class_type_cache:
            return self._class_type_cache[class_type]
        try:
            async with aiohttp.ClientSession() as session:
                async with session.get(
                    REGISTRY_COMFY_NODE_URL.format(class_type=class_type),
                    timeout=aiohttp.ClientTimeout(total=15),
                ) as resp:
                    if resp.status == 404:
                        self._class_type_cache[class_type] = None
                        return None
                    resp.raise_for_status()
                    data = await resp.json(content_type=None)
        except Exception as err:
            logger.info("comfy registry class-type lookup failed for %s: %s", class_type, err)
            self._class_type_cache[class_type] = None
            return None
        if not isinstance(data, dict) or not data.get("repository"):
            self._class_type_cache[class_type] = None
            return None
        self._class_type_cache[class_type] = data
        return data

    @staticmethod
    def _camel_split(s: str) -> list[str]:
        """Split a CamelCase class name into lowercase tokens.

        Keeps runs of uppercase letters together (acronyms): ``UnetLoaderGGUF``
        → ``["unet", "loader", "gguf"]``, ``XMLHttpRequest`` →
        ``["xml", "http", "request"]``.
        """
        import re
        # Insert a boundary between lower→upper and between upper-run→upper+lower
        s = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", s)
        s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", s)
        return [t.lower() for t in re.split(r"[^A-Za-z0-9]+", s) if t]

    _BOILERPLATE_TOKENS = frozenset({
        "loader", "encoder", "decoder", "advanced", "simple", "load", "save",
        "node", "nodes", "model", "models", "comfyui", "comfy",
    })

    async def resolve_by_class_type_manager(self, class_type: str) -> dict | None:
        """Fallback: find a Manager extension-node-map entry providing ``class_type``.

        Multiple repos can declare the same class (forks, re-exports). Rank
        candidates by:
        1. How many class-type tokens (CamelCase split, minus generic words)
           appear in the repo name. ``UnetLoaderGGUF`` → ["unet","gguf"], so
           ``ComfyUI-GGUF`` scores higher than ``ComfyUI-Zlycoris``.
        2. Fewer total class types in the entry = more focused = more likely
           the canonical source rather than an aggregator/fork.
        3. Lexicographic repo URL for stability.
        """
        if not class_type:
            return None
        ext_map = await self._load_ext_node_map()
        matches: list[tuple[str, int]] = [
            (repo_url, len(class_types))
            for repo_url, class_types in ext_map.items()
            if class_type in class_types
        ]
        if not matches:
            return None

        tokens = [t for t in self._camel_split(class_type) if t not in self._BOILERPLATE_TOKENS]

        def _score(item: tuple[str, int]) -> tuple[int, int, str]:
            repo_url, class_count = item
            repo_name = repo_url.rsplit("/", 1)[-1].lower()
            name_matches = sum(1 for tok in tokens if tok in repo_name)
            # Lower score wins; negate name_matches so more matches sort first.
            return (-name_matches, class_count, repo_url)

        matches.sort(key=_score)
        return {"repository": matches[0][0]}

    async def resolve_by_class_type(self, class_type: str) -> dict | None:
        """Resolve a class_type to ``{"repository": "..."}``.

        Tries the comfy registry's reverse-lookup first (canonical), then
        Manager's extension-node-map with smart ranking as a fallback.
        """
        if not class_type:
            return None
        record = await self.resolve_by_class_type_registry(class_type)
        if record and record.get("repository"):
            return record
        return await self.resolve_by_class_type_manager(class_type)

    async def _load_manager_models(self) -> list[dict]:
        async with self._lock:
            if self._manager_models is not None:
                return self._manager_models
            try:
                async with aiohttp.ClientSession() as session:
                    data = await self._get_json(session, MANAGER_MODELS_URL)
                self._manager_models = data.get("models") or []
            except Exception as err:
                logger.info("manager model-list fetch failed: %s", err)
                self._manager_models = []
            return self._manager_models

    async def resolve_model_by_filename(self, filename: str) -> dict | None:
        """Find a Manager model entry matching ``filename`` (basename match).

        Returns ``{"url": "...", "save_path": "...", "filename": "..."}`` or None.
        """
        if not filename:
            return None
        models = await self._load_manager_models()
        for entry in models:
            if entry.get("filename") == filename:
                return entry
        return None

    async def resolve_by_huggingface_search(self, filename: str) -> dict | None:
        """Find a Hugging Face repo containing a file with exact basename ``filename``.

        Strategy: search the HF Models API (with and without the extension —
        relevance is usually better without), walk the top candidates, fetch
        each candidate's recursive file tree (lighter than the full model
        record), and return the URL of the first sibling whose basename
        matches. ``main`` branch is assumed; HF refers to the default branch
        by name in resolve URLs.

        Returns ``{"url", "repo", "path"}`` or None. No auth — gated repos
        will return 401/403 and we fall through. Results cached per filename
        for the duration of the registry client's lifetime.
        """
        if not filename:
            return None
        if filename in self._hf_cache:
            return self._hf_cache[filename]

        base = filename
        for ext in _MODEL_EXTENSIONS:
            if base.lower().endswith(ext):
                base = base[: -len(ext)]
                break

        queries: list[str] = []
        if base:
            queries.append(base)
        if filename and filename not in queries:
            queries.append(filename)

        seen_repos: set[str] = set()
        timeout = aiohttp.ClientTimeout(total=15)
        try:
            async with aiohttp.ClientSession(timeout=timeout) as session:
                for query in queries:
                    try:
                        async with session.get(
                            HUGGINGFACE_SEARCH_URL,
                            params={"search": query, "limit": 15},
                        ) as resp:
                            if resp.status != 200:
                                continue
                            results = await resp.json(content_type=None)
                    except Exception:
                        continue
                    if not isinstance(results, list):
                        continue

                    for repo_data in results:
                        model_id = repo_data.get("modelId") or repo_data.get("id")
                        if not model_id or model_id in seen_repos:
                            continue
                        seen_repos.add(model_id)
                        try:
                            async with session.get(
                                HUGGINGFACE_TREE_URL.format(repo=model_id),
                                params={"recursive": "true"},
                            ) as r:
                                if r.status != 200:
                                    continue
                                tree = await r.json(content_type=None)
                        except Exception:
                            continue
                        if not isinstance(tree, list):
                            continue
                        for item in tree:
                            if not isinstance(item, dict) or item.get("type") != "file":
                                continue
                            path = item.get("path") or ""
                            if os.path.basename(path) == filename:
                                url = f"https://huggingface.co/{model_id}/resolve/main/{path}"
                                result = {"url": url, "repo": model_id, "path": path}
                                self._hf_cache[filename] = result
                                return result
        except Exception as err:
            logger.info("HF search failed for %s: %s", filename, err)

        self._hf_cache[filename] = None
        return None


# ---------------------------------------------------------------------------
# Plan: diff workflow requirements against what's locally installed
# ---------------------------------------------------------------------------


def _normalize_origin(url: str | None) -> str:
    """Normalize a git origin URL for equality comparison.

    Strips ``.git`` and trailing slash, lowercases the host, drops scheme
    differences (https/git+https/ssh) by reducing to ``host/path``.
    """
    if not url:
        return ""
    s = url.strip().lower()
    if s.startswith("git+"):
        s = s[4:]
    if s.startswith("git@"):
        # git@github.com:owner/repo.git -> github.com/owner/repo.git
        s = s.replace(":", "/", 1)[4:]
    for scheme in ("https://", "http://", "ssh://"):
        if s.startswith(scheme):
            s = s[len(scheme):]
    if s.endswith(".git"):
        s = s[:-4]
    return s.rstrip("/")


def _repo_name_from_origin(origin: str) -> str:
    """``https://github.com/owner/foo`` -> ``foo``."""
    norm = _normalize_origin(origin)
    if not norm:
        return ""
    return norm.rsplit("/", 1)[-1]


_NON_MODEL_FOLDER_TYPES = frozenset({
    "input", "output", "temp", "user", "custom_nodes", "configs",
})

_MODEL_EXTENSIONS = (
    ".safetensors", ".sft", ".ckpt", ".pt", ".pth", ".bin",
    ".gguf", ".onnx", ".pb", ".engine",
)


def _detect_missing_widget_models(graph: dict) -> list[dict]:
    """Find widget values that look like model filenames but aren't installed
    in any model folder.

    Mirrors ComfyUI's own "missing model" detection. ComfyUI flags a COMBO
    widget red when its current value isn't in the option list — and for
    model-name widgets that option list comes from ``folder_paths``. We
    reproduce that signal here:

    1. Enumerate every file in every model folder_type into a set (both
       relative path and basename, since dropdowns can carry either shape).
    2. For each widget value in every node, if it has a model-ish extension
       and doesn't appear in that set, flag it.

    Returns ``[{"value": ..., "filename": ..., "node_type": ...}, ...]``.
    The caller resolves the URL + save_path via Manager's model-list.
    """
    installed: set[str] = set()
    for ft in folder_paths.folder_names_and_paths:
        if ft in _NON_MODEL_FOLDER_TYPES:
            continue
        try:
            files = folder_paths.get_filename_list(ft) or []
        except Exception:
            continue
        for f in files:
            installed.add(f)
            installed.add(os.path.basename(f))

    seen: set[str] = set()
    out: list[dict] = []
    for node in _iter_all_nodes(graph or {}):
        widgets_values = node.get("widgets_values")
        if not isinstance(widgets_values, list):
            continue
        for value in widgets_values:
            if not isinstance(value, str) or not value:
                continue
            if not value.lower().endswith(_MODEL_EXTENSIONS):
                continue
            if value in installed or os.path.basename(value) in installed:
                continue
            if value in seen:
                continue
            seen.add(value)
            out.append({
                "value": value,
                "filename": os.path.basename(value),
                "node_type": node.get("type") or "",
                "node": node,
            })
    return out


def _url_to_filename(url: str) -> str | None:
    """Best-effort filename extraction from a download URL."""
    if not url:
        return None
    # Strip query string, take the last path segment.
    base = url.split("?", 1)[0].split("#", 1)[0].rstrip("/")
    if not base:
        return None
    name = base.rsplit("/", 1)[-1]
    return name or None


# Last-ditch fallback: when INPUT_TYPES introspection can't reveal a node's
# model folder_type (empty option lists on a fresh install, custom loaders
# without standard schemas), fall back to this map. Covers ~all common
# built-in + popular custom loaders.
_LOADER_FOLDER_HINTS: dict[str, str] = {
    "VAELoader": "vae",
    "CheckpointLoaderSimple": "checkpoints",
    "CheckpointLoader": "checkpoints",
    "unCLIPCheckpointLoader": "checkpoints",
    "LoraLoader": "loras",
    "LoraLoaderModelOnly": "loras",
    "ControlNetLoader": "controlnet",
    "ControlNetLoaderAdvanced": "controlnet",
    "DiffControlNetLoader": "controlnet",
    "UNETLoader": "unet",
    "UnetLoaderGGUF": "unet",
    "UnetLoaderGGUFAdvanced": "unet",
    "CLIPLoader": "clip",
    "CLIPLoaderGGUF": "clip",
    "DualCLIPLoader": "clip",
    "DualCLIPLoaderGGUF": "clip",
    "TripleCLIPLoader": "clip",
    "QuadrupleCLIPLoaderGGUF": "clip",
    "CLIPVisionLoader": "clip_vision",
    "UpscaleModelLoader": "upscale_models",
    "StyleModelLoader": "style_models",
    "GLIGENLoader": "gligen",
    "HypernetworkLoader": "hypernetworks",
    "DiffusersLoader": "diffusers",
    "PhotoMakerLoader": "photomaker",
    "IPAdapterModelLoader": "ipadapter",
}


def _guess_folder_type_for_node(node: dict) -> str | None:
    """Identify the model folder_type a node loads from.

    Tries INPUT_TYPES introspection first — finds the first COMBO input whose
    option list equals a model folder's file list. If introspection fails or
    the option list is empty (fresh install with no files yet), falls back to
    the explicit ``_LOADER_FOLDER_HINTS`` map keyed by class_type.
    """
    class_type = node.get("type") or ""
    try:
        import nodes as comfy_nodes
    except Exception:
        comfy_nodes = None  # type: ignore

    if comfy_nodes is not None:
        cls = comfy_nodes.NODE_CLASS_MAPPINGS.get(class_type)
        if cls is not None:
            try:
                input_types = cls.INPUT_TYPES()
            except Exception:
                input_types = None
            if isinstance(input_types, dict):
                folder_sets: dict[str, frozenset[str]] = {}
                for ft in folder_paths.folder_names_and_paths:
                    if ft in _NON_MODEL_FOLDER_TYPES:
                        continue
                    try:
                        folder_sets[ft] = frozenset(folder_paths.get_filename_list(ft) or [])
                    except Exception:
                        continue
                for section in ("required", "optional"):
                    for input_spec in (input_types.get(section) or {}).values():
                        if not isinstance(input_spec, (list, tuple)) or not input_spec:
                            continue
                        first = input_spec[0]
                        if not isinstance(first, list) or not first:
                            continue
                        opt_set = frozenset(first)
                        for ft, files in folder_sets.items():
                            if files and opt_set == files:
                                return ft

    return _LOADER_FOLDER_HINTS.get(class_type)


async def plan_setup(graph: dict) -> dict:
    """Compute the diff between what the workflow needs and what's installed.

    Returns ``{"missing_models": [...], "missing_custom_nodes": [...], "unresolved": [...]}``.
    Each missing_model is ``{rel_path, filename, url, sha256?, total_bytes?}``.
    Each missing_custom_node is ``{name, origin, commit?, source}``.
    Unresolved entries are class_types / filenames we couldn't map to a source.
    """
    registry = RegistryClient()
    loaded_class_types = _comfy_core_class_types()

    # ---- Custom nodes ------------------------------------------------------
    missing_nodes: dict[str, dict] = {}  # keyed by normalized origin to dedupe
    unresolved: list[dict] = []
    seen_class_types: set[str] = set()

    for node in _iter_all_nodes(graph or {}):
        class_type = node.get("type") or ""
        if class_type in seen_class_types:
            continue
        seen_class_types.add(class_type)

        props = node.get("properties") or {}
        cnr_id = props.get("cnr_id")
        aux_id = props.get("aux_id")
        ver = props.get("ver")

        # Already runnable in this install? Skip.
        if class_type in loaded_class_types:
            continue
        if cnr_id == "comfy-core":
            continue
        # Stock frontend pseudo-nodes (Reroute, Note, etc.) have no Python
        # class registered, so the loaded-class check above doesn't catch
        # them. Treat them as already-runnable so they don't surface as
        # unresolved custom nodes.
        if class_type in _STOCK_FRONTEND_NODES:
            continue

        origin: str | None = None
        commit: str | None = None
        source = ""
        name = ""

        if cnr_id:
            record = await registry.resolve_by_cnr_id(cnr_id)
            if record and record.get("repository"):
                origin = record["repository"]
                commit = ver or None
                source = "comfy_registry"
                name = cnr_id

        # `aux_id` carries the canonical `owner/repo` for nodes installed via
        # Manager's git-URL flow. Try it before the class-type fallback because
        # `ver` gives us an exact commit pin — the Manager list path can't.
        # Real-world hit: personal-account packs (PozzettiAndrea/* family) that
        # aren't in the Comfy Registry and aren't in Manager's curated list.
        if not origin and isinstance(aux_id, str):
            aux_origin = _aux_id_to_origin(aux_id)
            if aux_origin:
                origin = aux_origin
                commit = ver or None
                source = "aux_id"
                name = _repo_name_from_origin(origin)

        if not origin and class_type:
            record = await registry.resolve_by_class_type(class_type)
            if record and record.get("repository"):
                origin = record["repository"]
                source = "manager"
                name = _repo_name_from_origin(origin)

        if not origin:
            if class_type:
                unresolved.append({"kind": "custom_node", "class_type": class_type})
            continue

        # NOTE: no "directory exists locally → skip" guard here. The functional
        # signal of "this node already works" is class_type ∈ loaded_class_types,
        # which we already check above. If we got past that, the class didn't
        # register — re-clone is the right action, even if the directory exists.
        # install_custom_node_sync rmtree's the target before moving the fresh
        # clone in, so re-install is safe.
        norm = _normalize_origin(origin)
        if norm in missing_nodes:
            # Prefer the entry with a commit pin if a later node carries one.
            if commit and not missing_nodes[norm].get("commit"):
                missing_nodes[norm]["commit"] = commit
            continue
        missing_nodes[norm] = {
            "name": name or _repo_name_from_origin(origin),
            "origin": origin,
            "commit": commit,
            "source": source,
        }
        logger.info(
            "Runflow auto-setup: planning install of %s (commit=%s, source=%s, class_type=%s)",
            origin, (commit or "HEAD"), source, class_type,
        )

    # ---- Models ------------------------------------------------------------
    models = resolve_workflow_models(graph or {})
    missing_models: list[dict] = []
    models_root = Path(folder_paths.models_dir)
    planned_filenames: set[str] = set()

    for rel_path, info in models.items():
        target = (models_root / rel_path).resolve()
        try:
            target.relative_to(models_root.resolve())
        except ValueError:
            continue
        if target.is_file():
            continue

        url = info.get("url")
        filename = Path(rel_path).name
        if not url:
            entry = await registry.resolve_model_by_filename(filename)
            if entry and entry.get("url"):
                url = entry["url"]

        if not url:
            unresolved.append({"kind": "model", "rel_path": rel_path, "filename": filename})
            continue

        missing_models.append({
            "rel_path": rel_path,
            "filename": filename,
            "url": url,
            "sha256": info.get("sha256"),
        })
        planned_filenames.add(filename)

    # Second pass: widget values that look like model filenames but aren't
    # installed and weren't covered above. resolve_workflow_models() silently
    # skips widget values for which it can find neither a local file nor a
    # URL in the same node's properties.models — the case ComfyUI itself flags
    # as "Missing Models" with the red node outline. We pick those up here
    # and look up the URL + save_path via two registries:
    #
    # 1. Manager's model-list — curated, carries an explicit save_path.
    # 2. Hugging Face search — long-tail fallback, no curated save_path so
    #    we infer it from the node's INPUT_TYPES / a small hint map.
    #
    # _detect_missing_widget_models() carries the originating node dict so
    # the folder_type can be inferred per-node without re-walking the graph.
    for entry in _detect_missing_widget_models(graph or {}):
        filename = entry["filename"]
        if filename in planned_filenames:
            continue

        url: str | None = None
        save_path = ""
        source = ""

        manager_entry = await registry.resolve_model_by_filename(filename)
        if manager_entry and manager_entry.get("url"):
            url = manager_entry["url"]
            save_path = (manager_entry.get("save_path") or "").strip("/")
            source = "manager"
        else:
            hf_entry = await registry.resolve_by_huggingface_search(filename)
            if hf_entry and hf_entry.get("url"):
                url = hf_entry["url"]
                # No save_path from HF — infer from the node that referenced it.
                save_path = (_guess_folder_type_for_node(entry["node"]) or "").strip("/")
                source = "huggingface"

        if not url:
            unresolved.append({
                "kind": "model",
                "filename": filename,
                "value": entry["value"],
            })
            continue

        # If the widget value carries a subdir (``loras/foo/x.safetensors``),
        # honor it — that's where ComfyUI is looking. Otherwise compose from
        # the inferred save_path.
        if "/" in entry["value"] or "\\" in entry["value"]:
            rel_path = entry["value"].replace("\\", "/")
        else:
            rel_path = f"{save_path}/{filename}" if save_path else filename

        missing_models.append({
            "rel_path": rel_path,
            "filename": filename,
            "url": url,
            "sha256": None,
            "source": source,
        })
        planned_filenames.add(filename)

    return {
        "missing_models": missing_models,
        "missing_custom_nodes": list(missing_nodes.values()),
        "unresolved": unresolved,
    }


# ---------------------------------------------------------------------------
# Workers: download models, install custom nodes
# ---------------------------------------------------------------------------


class CancelledByUser(Exception):
    """Raised from inside a worker when the job's cancel flag is set."""


ProgressCb = Callable[[dict], Awaitable[None]]


async def download_model(
    url: str,
    target_path: Path,
    expected_sha256: str | None,
    progress_cb: ProgressCb,
    is_cancelled: Callable[[], bool],
) -> None:
    """Stream-download ``url`` into ``target_path`` with sha256 verification.

    Writes to ``<target>.part`` and atomically renames on success. Verifies
    sha256 if provided; mismatched files are deleted before raising.
    """
    target_path.parent.mkdir(parents=True, exist_ok=True)
    part_path = target_path.with_name(target_path.name + ".part")

    timeout = aiohttp.ClientTimeout(total=None, sock_read=60, sock_connect=30)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        async with session.get(url, allow_redirects=True) as resp:
            resp.raise_for_status()
            total_header = resp.headers.get("Content-Length")
            total_bytes = int(total_header) if total_header and total_header.isdigit() else 0
            downloaded = 0
            hasher = hashlib.sha256() if expected_sha256 else None
            last_emit = 0
            with open(part_path, "wb") as f:
                async for chunk in resp.content.iter_chunked(1024 * 1024):
                    if is_cancelled():
                        raise CancelledByUser()
                    f.write(chunk)
                    if hasher is not None:
                        hasher.update(chunk)
                    downloaded += len(chunk)
                    # Throttle progress events to ~10/s worth of bytes
                    if downloaded - last_emit >= 256 * 1024:
                        await progress_cb({
                            "type": "model_progress",
                            "bytes": downloaded,
                            "total_bytes": total_bytes,
                        })
                        last_emit = downloaded

    if hasher is not None and expected_sha256:
        actual = hasher.hexdigest()
        if actual.lower() != expected_sha256.lower():
            try:
                part_path.unlink()
            except OSError:
                pass
            raise ValueError(f"sha256 mismatch (got {actual[:12]}…, expected {expected_sha256[:12]}…)")

    # Atomic rename across the same filesystem; if target exists (race), replace.
    part_path.replace(target_path)
    await progress_cb({
        "type": "model_progress",
        "bytes": downloaded,
        "total_bytes": total_bytes or downloaded,
    })


def _run_subprocess(
    argv: list[str],
    cwd: Path | None,
    log_cb: Callable[[str], None],
    timeout: float | None = None,
) -> int:
    """Run a subprocess streaming combined stdout/stderr line-by-line to log_cb.

    Returns the exit code. Raises ``subprocess.TimeoutExpired`` if ``timeout``
    elapses. No ``shell=True`` — argv is explicit.
    """
    proc = subprocess.Popen(
        argv,
        cwd=str(cwd) if cwd else None,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        encoding="utf-8",
        errors="replace",
    )
    assert proc.stdout is not None
    try:
        for line in proc.stdout:
            line = line.rstrip()
            if line:
                log_cb(line)
        proc.wait(timeout=timeout)
    except Exception:
        proc.kill()
        raise
    return proc.returncode


@functools.lru_cache(maxsize=1)
def _resolve_git() -> str | None:
    """Locate a working ``git`` executable, or return ``None`` if absent.

    ComfyUI Desktop (Mac/Windows) and the Windows portable build don't put git
    on PATH, so :func:`shutil.which` alone misses it. We additionally probe the
    common per-OS install locations, then *verify* the candidate actually runs
    (``git --version``): on macOS ``/usr/bin/git`` is a stub that exists even
    with no developer tools installed and only errors when invoked, so file
    existence isn't enough. Result is cached for the process.
    """
    candidates: list[str] = []
    on_path = shutil.which("git")
    if on_path:
        candidates.append(on_path)
    if sys.platform == "win32":
        candidates += [
            os.path.expandvars(r"%PROGRAMFILES%\Git\cmd\git.exe"),
            os.path.expandvars(r"%PROGRAMFILES(X86)%\Git\cmd\git.exe"),
            os.path.expandvars(r"%LOCALAPPDATA%\Programs\Git\cmd\git.exe"),
        ]
    elif sys.platform == "darwin":
        candidates += [
            "/opt/homebrew/bin/git",   # Apple-silicon Homebrew
            "/usr/local/bin/git",      # Intel Homebrew / git-scm installer
            "/usr/bin/git",            # Xcode CLT (may be a no-toolchain stub)
            "/Library/Developer/CommandLineTools/usr/bin/git",
        ]
    else:
        candidates += ["/usr/bin/git", "/usr/local/bin/git", "/bin/git"]

    seen: set[str] = set()
    for cand in candidates:
        if not cand or cand in seen:
            continue
        seen.add(cand)
        # which() results are already known-present; probe paths must exist.
        if cand != on_path and not Path(cand).is_file():
            continue
        try:
            probe = subprocess.run(
                [cand, "--version"],
                capture_output=True, text=True, timeout=10,
            )
        except (OSError, subprocess.SubprocessError):
            continue
        if probe.returncode == 0:
            return cand
    return None


def _github_slug(origin: str) -> tuple[str, str] | None:
    """Return ``(owner, repo)`` for a github.com origin, else ``None``."""
    norm = _normalize_origin(origin)  # -> host/path, lowercased, no .git
    prefix = "github.com/"
    if not norm.startswith(prefix):
        return None
    parts = [p for p in norm[len(prefix):].split("/") if p]
    if len(parts) < 2:
        return None
    return parts[0], parts[1]


def _safe_extract_tar(tf: tarfile.TarFile, dest: Path) -> None:
    """Extract ``tf`` into ``dest``, rejecting members that escape ``dest``."""
    dest = dest.resolve()
    base = str(dest) + os.sep
    for member in tf.getmembers():
        target = (dest / member.name).resolve()
        if target != dest and not str(target).startswith(base):
            raise RuntimeError(f"unsafe path in tarball: {member.name!r}")
    try:
        tf.extractall(dest, filter="data")  # type: ignore[call-arg]  # py3.12+
    except TypeError:
        tf.extractall(dest)


def _fetch_node_via_git(
    git: str,
    origin: str,
    commit: str | None,
    target: Path,
    log_cb: Callable[[str], None],
    is_cancelled: Callable[[], bool],
) -> None:
    """Clone ``origin`` (and checkout ``commit``) into ``target`` using ``git``."""
    log_cb(f"git clone {origin}")
    rc = _run_subprocess(
        [git, "clone", "--depth=1", origin, str(target)],
        cwd=None, log_cb=log_cb, timeout=600,
    )
    if rc != 0:
        raise RuntimeError(f"git clone failed (exit {rc})")

    if commit:
        if is_cancelled():
            raise CancelledByUser()
        log_cb(f"git checkout {commit[:12]}…")
        rc = _run_subprocess(
            [git, "fetch", "--depth=1", "origin", commit],
            cwd=target, log_cb=log_cb, timeout=300,
        )
        if rc == 0:
            rc = _run_subprocess(
                [git, "checkout", commit],
                cwd=target, log_cb=log_cb, timeout=120,
            )
        if rc != 0:
            # Fallback: full clone, then checkout. Some hosts disable
            # uploadpack.allowReachableSHA1InWant, which makes the
            # depth-1 fetch above reject the SHA.
            log_cb("falling back to full clone")
            shutil.rmtree(target, ignore_errors=True)
            rc = _run_subprocess(
                [git, "clone", origin, str(target)],
                cwd=None, log_cb=log_cb, timeout=900,
            )
            if rc != 0:
                raise RuntimeError(f"git clone (full) failed (exit {rc})")
            rc = _run_subprocess(
                [git, "checkout", commit],
                cwd=target, log_cb=log_cb, timeout=120,
            )
            if rc != 0:
                # Pinned commit/branch unreachable (deleted, rebased away, or
                # force-pushed). Don't fail the install — keep the default-branch
                # HEAD the clone already checked out.
                log_cb(
                    f"WARNING: could not check out {commit[:12]}… (exit {rc}); "
                    "using the latest default-branch HEAD instead."
                )

    # Initialize any git submodules the node vendors (ComfyUI-3D-Pack
    # vendors gsplat/nvdiffrast, Hunyuan3D-Part vendors model code, etc.).
    # No-op when the repo has no .gitmodules — ~50ms; cheaper than the
    # extra branch + stat call.
    if is_cancelled():
        raise CancelledByUser()
    log_cb("git submodule update --init --recursive")
    rc = _run_subprocess(
        [git, "submodule", "update", "--init", "--recursive"],
        cwd=target, log_cb=log_cb, timeout=1800,
    )
    if rc != 0:
        raise RuntimeError(f"git submodule update failed (exit {rc})")


def _fetch_node_via_tarball(
    origin: str,
    commit: str | None,
    target: Path,
    log_cb: Callable[[str], None],
    is_cancelled: Callable[[], bool],
) -> None:
    """Download + extract a GitHub source tarball into ``target`` (no git).

    Used only when :func:`_resolve_git` finds no usable git binary. Submodules
    are *not* fetched (tarballs don't include them) — we warn when the node
    vendors any so the limitation is explicit rather than silent.
    """
    slug = _github_slug(origin)
    if slug is None:
        raise RuntimeError(
            f"git is not installed and {origin!r} is not a GitHub repository, "
            "so it cannot be fetched without git. Install git (or add it to "
            "PATH) and retry."
        )
    owner, repo = slug
    # GitHub's API tarball endpoint redirects to codeload and works
    # unauthenticated for public repos; a ref accepts a SHA, tag, or branch
    # (HEAD = default branch). Prefer the pinned commit, but fall back to HEAD
    # when that ref can't be fetched (deleted/rebased/force-pushed) — mirrors
    # the git path's "use HEAD" behaviour.
    refs = [commit, "HEAD"] if commit and commit != "HEAD" else ["HEAD"]
    data = None
    last_exc: Exception | None = None
    for ref in refs:
        label = "HEAD" if ref == "HEAD" else ref[:12]
        url = f"https://api.github.com/repos/{owner}/{repo}/tarball/{ref}"
        log_cb(f"git unavailable — downloading tarball {owner}/{repo}@{label}")
        req = urllib.request.Request(url, headers={
            "User-Agent": "ComfyUI-Runflow-auto-setup",
            "Accept": "application/vnd.github+json",
        })
        try:
            with urllib.request.urlopen(req, timeout=600) as resp:
                data = resp.read()
            break
        except urllib.error.HTTPError as exc:
            last_exc = RuntimeError(
                f"tarball download failed (HTTP {exc.code}) for {url}"
            )
            if ref != "HEAD":
                log_cb(
                    f"WARNING: tarball for {label} unavailable (HTTP {exc.code}); "
                    "falling back to HEAD"
                )
                continue
        except (urllib.error.URLError, OSError) as exc:
            # Network-level failure — retrying a different ref won't help.
            last_exc = RuntimeError(f"tarball download failed for {url}: {exc}")
            break
    if data is None:
        raise last_exc or RuntimeError("tarball download failed")

    if is_cancelled():
        raise CancelledByUser()

    # GitHub tarballs wrap everything in a single top-level dir
    # (`{owner}-{repo}-{sha7}/`); extract, then lift that dir into `target`.
    extract_root = target.parent / f"{target.name}__tar"
    shutil.rmtree(extract_root, ignore_errors=True)
    extract_root.mkdir(parents=True, exist_ok=True)
    try:
        with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as tf:
            _safe_extract_tar(tf, extract_root)
        subdirs = [p for p in extract_root.iterdir() if p.is_dir()]
        src_dir = subdirs[0] if len(subdirs) == 1 else extract_root
        if target.exists():
            shutil.rmtree(target, ignore_errors=True)
        shutil.move(str(src_dir), str(target))
    finally:
        shutil.rmtree(extract_root, ignore_errors=True)

    if (target / ".gitmodules").is_file():
        log_cb(
            "WARNING: this node declares git submodules, which the tarball "
            "fallback cannot fetch. Install git for a complete install."
        )


def install_custom_node_sync(
    origin: str,
    commit: str | None,
    name: str,
    custom_nodes_dir: Path,
    log_cb: Callable[[str], None],
    is_cancelled: Callable[[], bool],
) -> None:
    """Fetch ``origin`` at ``commit`` into ``custom_nodes/<name>/``.

    Runs entirely synchronously — call via ``asyncio.to_thread`` from async code.
    Prefers git (clone --depth=1 + SHA-targeted fetch, with a full-clone
    fallback); when no git binary is available it downloads a GitHub source
    tarball instead. Then ``python -m pip install -r requirements.txt`` if that
    file is present, and ``python install.py`` for comfy-env nodes.
    """
    if is_cancelled():
        raise CancelledByUser()
    custom_nodes_dir.mkdir(parents=True, exist_ok=True)
    staging = Path(tempfile.mkdtemp(prefix="runflow-cn-"))
    target_in_staging = staging / name
    try:
        git = _resolve_git()
        if git is not None:
            _fetch_node_via_git(
                git, origin, commit, target_in_staging, log_cb, is_cancelled,
            )
        else:
            log_cb("git not found on PATH or common locations — using tarball fallback")
            _fetch_node_via_tarball(
                origin, commit, target_in_staging, log_cb, is_cancelled,
            )

        # Move into custom_nodes/. Replace any existing dir of the same name.
        final = custom_nodes_dir / name
        if final.exists():
            shutil.rmtree(final, ignore_errors=False)
        shutil.move(str(target_in_staging), str(final))

        # requirements.txt — run via python -m pip, never bare `pip`.
        req = final / "requirements.txt"
        if req.is_file():
            if is_cancelled():
                raise CancelledByUser()
            log_cb("python -m pip install -r requirements.txt")
            rc = _run_subprocess(
                [sys.executable, "-m", "pip", "install", "-r", str(req)],
                cwd=final,
                log_cb=log_cb,
                timeout=1800,
            )
            if rc != 0:
                raise RuntimeError(f"pip install failed (exit {rc})")

        # comfy-env nodes (Pozzetti 3D family, comfyui-sharp) need their
        # isolated pixi envs materialized via `python install.py` after the
        # pip step puts the `comfy_env` package itself into the venv.
        # Without this, prestartup logs `[comfy-env] env: ... [MISSING --
        # run install.py]` and the node's worker-isolated classes never
        # register. Detector mirrors the deploy worker so behaviour stays
        # in lockstep on both sides of the deploy boundary.
        install_py = final / "install.py"
        if _is_comfy_env_install(install_py):
            if is_cancelled():
                raise CancelledByUser()
            log_cb("python install.py  (comfy-env isolated envs)")
            rc = _run_subprocess(
                [sys.executable, "install.py"],
                cwd=final,
                log_cb=log_cb,
                timeout=3600,  # cold pixi pulls (flash-attn, nvdiffrast, ...) take minutes
            )
            if rc != 0:
                raise RuntimeError(f"comfy-env install.py failed (exit {rc})")
    finally:
        shutil.rmtree(staging, ignore_errors=True)


# ---------------------------------------------------------------------------
# Job orchestrator
# ---------------------------------------------------------------------------


class _Job:
    def __init__(self, job_id: str, options: dict, plan: dict) -> None:
        self.job_id = job_id
        self.options = options
        self.plan = plan
        self.cancel = False
        self.task: asyncio.Task | None = None


_JOBS: dict[str, _Job] = {}


def get_job(job_id: str) -> _Job | None:
    return _JOBS.get(job_id)


def cancel_job(job_id: str) -> bool:
    job = _JOBS.get(job_id)
    if not job:
        return False
    job.cancel = True
    return True


async def run_setup_job(job: _Job, send: ProgressCb) -> None:
    """Drive the install + download steps for one job.

    Events emitted via ``send``:
    - ``{"type": "job_start", "job_id", "models_total", "custom_nodes_total"}``
    - ``{"type": "model_start", "index", "total", "filename", "rel_path"}``
    - ``{"type": "model_progress", "bytes", "total_bytes"}``
    - ``{"type": "model_done", "index"}``
    - ``{"type": "model_error", "index", "error"}``
    - ``{"type": "custom_node_start", "index", "total", "name", "origin"}``
    - ``{"type": "custom_node_log", "line"}``
    - ``{"type": "custom_node_done", "index"}``
    - ``{"type": "custom_node_error", "index", "error"}``
    - ``{"type": "job_done", "had_errors": bool, "errors": [...]}``
    - ``{"type": "job_cancelled"}``
    """
    plan = job.plan
    opts = job.options
    do_models = bool(opts.get("download_models", True))
    do_nodes = bool(opts.get("install_custom_nodes", True))

    models = plan.get("missing_models", []) if do_models else []
    nodes = plan.get("missing_custom_nodes", []) if do_nodes else []
    errors: list[dict] = []

    await send({
        "type": "job_start",
        "job_id": job.job_id,
        "models_total": len(models),
        "custom_nodes_total": len(nodes),
    })

    models_root = Path(folder_paths.models_dir)
    custom_nodes_dir = Path(folder_paths.base_path) / "custom_nodes"

    def _is_cancelled() -> bool:
        return job.cancel

    # Models and custom nodes run in parallel — they're independent activities,
    # and the spec ("downloading in the background … below that we see custom
    # node install status") implies both sections are live at the same time.
    # Within each phase, items run sequentially to match the "N of M" UI.

    async def _drive_models() -> None:
        for index, model in enumerate(models):
            if job.cancel:
                return
            await send({
                "type": "model_start",
                "index": index,
                "total": len(models),
                "filename": model["filename"],
                "rel_path": model["rel_path"],
            })
            target = models_root / model["rel_path"]
            try:
                async def _pcb(event: dict, _i=index) -> None:
                    event = dict(event)
                    event["index"] = _i
                    event["total"] = len(models)
                    await send(event)
                await download_model(
                    url=model["url"],
                    target_path=target,
                    expected_sha256=model.get("sha256"),
                    progress_cb=_pcb,
                    is_cancelled=_is_cancelled,
                )
                await send({"type": "model_done", "index": index, "total": len(models)})
            except CancelledByUser:
                return
            except Exception as err:
                logger.exception("auto-setup: model download failed for %s", model["filename"])
                errors.append({"kind": "model", "name": model["filename"], "error": str(err)})
                await send({"type": "model_error", "index": index, "error": str(err)})

    async def _drive_custom_nodes() -> None:
        loop = asyncio.get_running_loop()
        for index, cn in enumerate(nodes):
            if job.cancel:
                return
            await send({
                "type": "custom_node_start",
                "index": index,
                "total": len(nodes),
                "name": cn["name"],
                "origin": cn["origin"],
            })

            def _log_sync(line: str, _i=index, _loop=loop) -> None:
                asyncio.run_coroutine_threadsafe(
                    send({"type": "custom_node_log", "index": _i, "line": line}),
                    _loop,
                )

            try:
                await asyncio.to_thread(
                    install_custom_node_sync,
                    cn["origin"],
                    cn.get("commit"),
                    cn["name"],
                    custom_nodes_dir,
                    _log_sync,
                    _is_cancelled,
                )
                await send({"type": "custom_node_done", "index": index, "total": len(nodes)})
            except CancelledByUser:
                return
            except Exception as err:
                logger.exception("auto-setup: custom node install failed for %s", cn["name"])
                errors.append({"kind": "custom_node", "name": cn["name"], "error": str(err)})
                await send({"type": "custom_node_error", "index": index, "error": str(err)})

    await asyncio.gather(_drive_models(), _drive_custom_nodes())

    if job.cancel:
        await send({"type": "job_cancelled"})
        return
    await send({"type": "job_done", "had_errors": bool(errors), "errors": errors})


def start_job(options: dict, plan: dict, send: ProgressCb) -> str:
    """Register a new job and kick off its driver task. Returns the job_id."""
    job_id = uuid.uuid4().hex
    job = _Job(job_id, options, plan)
    _JOBS[job_id] = job
    job.task = asyncio.create_task(run_setup_job(job, send))
    return job_id


# ---------------------------------------------------------------------------
# Restart
# ---------------------------------------------------------------------------


async def schedule_restart(delay: float = 0.25) -> None:
    """Re-exec the current Python process after ``delay`` seconds.

    Re-using ``sys.executable`` + ``sys.argv`` preserves whatever launch flags
    the user originally passed (port, listen address, custom-args). Works on
    Linux, macOS, and Windows.
    """
    await asyncio.sleep(delay)
    try:
        os.execv(sys.executable, [sys.executable] + sys.argv)
    except Exception:
        logger.exception("auto-setup: os.execv restart failed")
        # Hard exit so a supervisor (systemd, launchd, the comfy-cli wrapper)
        # can bring us back. Better than a half-restart that leaves no log.
        os._exit(1)