"""Pure helpers for the headless run API (POST /config_tester/run, GET /config_tester/run_status). No aiohttp / ComfyUI imports at module level so the helpers are unit-testable on their own. The routes themselves live in __init__.py. """ import json import os import re from collections import OrderedDict DEFAULT_SESSION = "my_session" # the Generator node's own default (sampler_node.py) _MAX_TRACKED_RUNS = 200 def sanitize_session_name_strict(name): """Same regex as delete_session / save_manifest / get_session_html. '' when nothing is left.""" return re.sub(r'[^\w\-]', '', str(name or '')) def normalize_configs_payload(configs_json): """Accept a dict with a `configs` list, a bare list, or a JSON string of either. Returns a dict that always has a non-empty `configs` list; other top-level keys (`_distribution`, `_session_settings`, ...) are preserved. Raises ValueError otherwise. """ obj = configs_json if isinstance(obj, (str, bytes)): try: obj = json.loads(obj) except (TypeError, ValueError) as e: raise ValueError(f"configs_json is not valid JSON: {e}") if isinstance(obj, list): obj = {"configs": obj} if not isinstance(obj, dict): raise ValueError("configs_json must be an object with a 'configs' list, a list of configs, or a JSON string of either") configs = obj.get("configs") if not isinstance(configs, list) or not configs: raise ValueError("configs_json must contain a non-empty 'configs' list") return obj def apply_run_overrides(obj, session_name=None, overwrite_existing=None, nonce=None): """Apply run-level overrides. Returns (obj, sanitised_session_name). Session precedence: explicit session_name, then _session_settings.session_name, then the Generator default. The sanitised name is written back to _session_settings.session_name because that is what the Generator actually reads. `nonce` goes into an inert top-level key; it changes IS_CHANGED's hash so ComfyUI's node cache cannot short-circuit the run. """ settings = obj.get("_session_settings") if not isinstance(settings, dict): settings = {} obj["_session_settings"] = settings chosen = session_name or settings.get("session_name") or DEFAULT_SESSION session = sanitize_session_name_strict(chosen) settings["session_name"] = session if overwrite_existing is not None: settings["overwrite_existing"] = bool(overwrite_existing) if nonce is not None: obj["_run_nonce"] = str(nonce) return obj, session def build_run_prompt(configs_str, session_name): """The 2-node API-format prompt: Generator -> Dashboard (which owns the output).""" return { "1": {"class_type": "UltimateSamplerGrid", "inputs": {"configs_json": configs_str}}, "2": {"class_type": "UltimateGridDashboard", "inputs": {"session_name": session_name, "dashboard_html": ["1", 0]}}, } def resolve_saved_config(configs_dir, name): """Mirror of the /configbuilder/load_config path handling. Returns the file path. Raises FileNotFoundError (missing) or PermissionError (escapes configs_dir). """ filename = str(name or "") if not filename.endswith(".json"): filename += ".json" filename = os.path.basename(filename) filepath = os.path.join(configs_dir, filename) real_path = os.path.realpath(filepath) real_base = os.path.realpath(configs_dir) if not (real_path == real_base or real_path.startswith(real_base + os.sep)): raise PermissionError("config path escapes the configs directory") if not os.path.isfile(filepath): raise FileNotFoundError(f"config '{filename}' not found") return filepath def saved_config_to_configs_obj(data, state_to_configs_json): """Saved Builder files hold the Builder STATE; convert via the single source of truth.""" if isinstance(data, dict): if "config_arrays" in data: return json.loads(state_to_configs_json(data)) if "configs" in data: return data raise ValueError("saved config is neither a Builder state (config_arrays) nor a configs JSON (configs)") def summarize_history(entry): """Small status summary of a history entry. Never returns `outputs` or `prompt`.""" status = (entry or {}).get("status") or {} out = {"status": status.get("status_str", "unknown"), "completed": bool(status.get("completed", False))} for msg in status.get("messages") or []: try: if msg[0] == "execution_error": d = msg[1] or {} out["error"] = { "node_id": d.get("node_id"), "node_type": d.get("node_type"), "exception_message": d.get("exception_message"), } break except (IndexError, TypeError, KeyError): continue return out def find_in_queue(running, pending, prompt_id): """('running', 0), ('pending', position) or None. Items are tuples (number, prompt_id, ...).""" for item in running or []: if len(item) > 1 and item[1] == prompt_id: return ("running", 0) for pos, item in enumerate(sorted(pending or [], key=lambda i: i[0])): if len(item) > 1 and item[1] == prompt_id: return ("pending", pos) return None class RunTracker: """Bounded prompt_id -> session_name map (oldest evicted first).""" def __init__(self, cap=_MAX_TRACKED_RUNS): self._cap = cap self._d = OrderedDict() def add(self, prompt_id, session): self._d[prompt_id] = session while len(self._d) > self._cap: self._d.popitem(last=False) def get(self, prompt_id): return self._d.get(prompt_id)