ComfyUI-MiniMaxH3-Context-Loop/run_manager.py

"""Discover and restore H3 Chain Plan archives without loading prompt history."""

from __future__ import annotations

import json
import os
import re
from datetime import datetime, timezone
from typing import Any

try:
    from .asset_store import RunAssetStore
except ImportError:  # Standalone unit tests import this module without a package.
    from asset_store import RunAssetStore

try:
    from .contracts_v05 import (
        AUDIO_POLICY_VERSION,
        CONTINUATION_POLICIES,
        TRANSITION_POLICY_VERSION,
        TRANSITION_PRESETS,
        audio_policy,
        migrate_continuation_mode,
        migrate_legacy_audio_mode,
        transition_policy,
    )
except ImportError:  # Standalone unit tests import this module without a package.
    from contracts_v05 import (
        AUDIO_POLICY_VERSION,
        CONTINUATION_POLICIES,
        TRANSITION_POLICY_VERSION,
        TRANSITION_PRESETS,
        audio_policy,
        migrate_continuation_mode,
        migrate_legacy_audio_mode,
        transition_policy,
    )


PLAN_NODE_TYPE = "MiniMaxH3ChainPlan"
MODERN_PLAN_NODE_TYPE = "MiniMaxH3ChainPlanModern"
PLAN_NODE_TYPES = (PLAN_NODE_TYPE, MODERN_PLAN_NODE_TYPE)
PLAN_WIDGET_NAMES = (
    "plan_json",
    "run_name",
    "generation_fingerprint",
    "width",
    "height",
    "context_length",
    "encode_mode",
    "anchor_mode",
    "crop",
    "audio_mode",
    "audio_context_length",
    "default_duration_seconds",
    "default_steps",
    "base_seed",
    "segment_crf",
    "video_blend_frames",
    "continuation_mode",
)
MODERN_PLAN_WIDGET_NAMES = (
    "plan_json",
    "run_name",
    "generation_fingerprint",
    "width",
    "height",
    "encode_mode",
    "crop",
    "default_duration_seconds",
    "default_steps",
    "base_seed",
    "segment_crf",
    "video_blend_frames",
)

H3_CONTEXT_LENGTHS = (
    1, 5, 22, 39, 56, 73, 90, 107, 124,
    141, 158, 175, 192, 209, 226, 243,
)
CONTINUATION_MODES = tuple(CONTINUATION_POLICIES)
SCENE_LORA_ROUTES = (
    "base", *(chr(ord("a") + offset) for offset in range(26)))


def _safe_name(value: Any, fallback: str = "") -> str:
    text = re.sub(r"[^A-Za-z0-9._-]+", "_", str(value or "").strip())
    text = text.strip("._-")
    return (text or fallback)[:96]


def _read_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def _json_string(value: Any) -> str | None:
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            return None
        if not isinstance(parsed, (dict, list)):
            return None
        return json.dumps(parsed, ensure_ascii=False, indent=2)
    if isinstance(value, (dict, list)):
        return json.dumps(value, ensure_ascii=False, indent=2)
    return None


def _restorable_widget_value(name: str, value: Any) -> Any:
    if name == "plan_json":
        return _json_string(value)
    if name == "base_seed" and isinstance(value, int):
        # aiohttp JSON would otherwise turn uint64 values above 2^53 into an
        # inexact JavaScript Number before the Plan widget sees them.
        return str(value)
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    # API-prompt links are lists such as [node_id, output_slot]. Restoring a
    # widget must never replace or invent graph connections.
    return None


def _api_prompt_inputs(document: Any, run_name: str) -> dict[str, Any]:
    if not isinstance(document, dict):
        return {}
    candidates = []
    exact = []
    for node in document.values():
        if (not isinstance(node, dict) or
                node.get("class_type") not in PLAN_NODE_TYPES):
            continue
        inputs = node.get("inputs")
        if not isinstance(inputs, dict):
            continue
        candidates.append(inputs)
        if _safe_name(inputs.get("run_name")) == run_name:
            exact.append(inputs)
    selected = exact[0] if exact else (candidates[0] if len(candidates) == 1 else None)
    if selected is None:
        return {}
    restored = {}
    for name in dict.fromkeys(PLAN_WIDGET_NAMES + MODERN_PLAN_WIDGET_NAMES):
        if name not in selected:
            continue
        value = _restorable_widget_value(name, selected[name])
        if value is not None:
            restored[name] = value
    return restored


def _workflow_inputs(document: Any, run_name: str) -> dict[str, Any]:
    if not isinstance(document, dict) or not isinstance(document.get("nodes"), list):
        return {}
    candidates = []
    exact = []
    for node in document["nodes"]:
        if (not isinstance(node, dict) or
                node.get("type") not in PLAN_NODE_TYPES):
            continue
        widgets = node.get("widgets_values")
        if not isinstance(widgets, list):
            continue
        candidate = (node.get("type"), widgets)
        candidates.append(candidate)
        if len(widgets) > 1 and _safe_name(widgets[1]) == run_name:
            exact.append(candidate)
    selected = exact[0] if exact else (candidates[0] if len(candidates) == 1 else None)
    if selected is None:
        return {}
    plan_type, selected = selected
    if plan_type == MODERN_PLAN_NODE_TYPE:
        if (len(selected) < len(MODERN_PLAN_WIDGET_NAMES)
                or selected[5] not in ("video", "frames")
                or selected[6] not in ("disabled", "center")):
            return {}
        restored = {}
        for name, value in zip(MODERN_PLAN_WIDGET_NAMES, selected):
            value = _restorable_widget_value(name, value)
            if value is not None:
                restored[name] = value
        return restored
    if len(selected) < 15:
        return {}
    # widgets_values is positional and old workflows can predate fields added
    # near the front of the Plan. Refuse a shifted layout rather than applying
    # plausible-looking values to the wrong widgets; plan.json still restores
    # the effective scene configuration.
    if (selected[5] not in H3_CONTEXT_LENGTHS
            or selected[6] not in ("video", "frames")
            or selected[7] not in ("head", "before")
            or selected[8] not in ("disabled", "center")
            or selected[9] not in (
                "source_track", "generated_audio", "source_plus_timeline")):
        return {}
    if len(selected) > 16 and selected[16] not in CONTINUATION_MODES:
        return {}
    restored = {}
    for name, value in zip(PLAN_WIDGET_NAMES, selected):
        value = _restorable_widget_value(name, value)
        if value is not None:
            restored[name] = value
    # 0.3 workflows end at segment_crf. New widgets are appended, so old
    # positional values remain exact and receive their compatibility defaults.
    restored.setdefault("video_blend_frames", 0)
    restored.setdefault("continuation_mode", "guide")
    return restored


def _editor_plan(archive: dict[str, Any]) -> dict[str, Any]:
    existing = archive.get("editor_plan")
    if isinstance(existing, dict) and isinstance(existing.get("shots"), list):
        return existing
    shots = []
    for offset, shot in enumerate(archive.get("shots") or []):
        if not isinstance(shot, dict):
            continue
        value = {
            "id": shot.get("id", "clip_%04d" % (offset + 1)),
            "prompt": shot.get("scene_prompt", shot.get("prompt", "")),
        }
        if shot.get("raw_frames") is not None:
            value["length"] = int(shot["raw_frames"])
        if shot.get("steps") is not None:
            value["steps"] = int(shot["steps"])
        if shot.get("seed") is not None:
            value["seed"] = str(shot["seed"])
        if shot.get("continuation_mode") in CONTINUATION_MODES:
            value["continuation_mode"] = shot["continuation_mode"]
        if (shot.get("context_length") == 0
                or shot.get("context_length") in H3_CONTEXT_LENGTHS):
            value["context_length"] = int(shot["context_length"])
        audio_context = shot.get("audio_context_length")
        if (isinstance(audio_context, int) and not isinstance(audio_context, bool)
                and 0 <= audio_context <= 240):
            value["audio_context_length"] = audio_context
        lora_route = str(shot.get("lora_route") or "").strip().lower()
        if lora_route in SCENE_LORA_ROUTES and lora_route != "base":
            value["lora_route"] = lora_route
        shots.append(value)
    return {
        "prompt_prefix": archive.get("prompt_prefix", ""),
        "shots": shots,
    }


def _archive_inputs(archive: Any, run_name: str) -> dict[str, Any]:
    if not isinstance(archive, dict):
        return {}
    editor = _editor_plan(archive)
    if not editor.get("shots"):
        return {}
    restored: dict[str, Any] = {
        "plan_json": json.dumps(editor, ensure_ascii=False, indent=2),
        "run_name": run_name,
    }
    compatibility = archive.get("compatibility")
    if isinstance(compatibility, dict):
        for name in (
            "generation_fingerprint", "width", "height", "context_length",
            "encode_mode", "anchor_mode", "crop", "audio_mode",
            "audio_context_length", "segment_crf", "video_blend_frames",
            "continuation_mode",
        ):
            if name in compatibility:
                restored[name] = compatibility[name]
    if "segment_crf" in archive:
        restored["segment_crf"] = archive["segment_crf"]
    restored.setdefault("video_blend_frames", 0)
    restored.setdefault("continuation_mode", "guide")
    return restored


def archive_policy_inputs(archive: Any) -> dict[str, dict[str, Any]]:
    """Return exact widget values for connected 0.5 policy nodes.

    Plan archives store resolved typed policies in compatibility metadata.  A
    restore must project those records back onto the upstream policy widgets;
    writing only the Plan's hidden 0.4 fallbacks has no effect while a typed
    policy input is connected.
    """
    if not isinstance(archive, dict):
        return {}
    compatibility = archive.get("compatibility")
    if not isinstance(compatibility, dict):
        return {}

    restored: dict[str, dict[str, Any]] = {}
    saved_audio = compatibility.get("audio_policy")
    try:
        if (isinstance(saved_audio, dict) and
                saved_audio.get("version") == AUDIO_POLICY_VERSION):
            resolved_audio = audio_policy(
                saved_audio.get("final_audio"),
                saved_audio.get("source_reference"),
                saved_audio.get("generated_continuity"),
                saved_audio.get("source_audio_target", "off"))
        else:
            resolved_audio = migrate_legacy_audio_mode(
                compatibility.get("audio_mode", "generated_audio"))
        restored["audio_policy"] = {
            "final_audio": resolved_audio["final_audio"],
            "source_reference": resolved_audio["source_reference"],
            "generated_continuity": resolved_audio["generated_continuity"],
        }
        if resolved_audio.get("source_audio_target") == "locked":
            restored["audio_policy"]["source_audio_target"] = "locked"
    except (TypeError, ValueError):
        # A malformed optional policy must not prevent prompt/Plan recovery.
        pass

    saved_transition = compatibility.get("transition_policy")
    try:
        if (isinstance(saved_transition, dict) and
                saved_transition.get("version") == TRANSITION_POLICY_VERSION and
                saved_transition.get("preset") in TRANSITION_PRESETS):
            resolved_transition = transition_policy(
                saved_transition.get("preset"),
                expert_override=bool(saved_transition.get(
                    "expert_override", False)),
                continuation_mode=saved_transition.get("continuation_mode"),
                context_length=saved_transition.get("context_length"))
        else:
            mode = migrate_continuation_mode(
                compatibility.get("continuation_mode", "guide"))
            context = int(compatibility.get("context_length", 22))
            matching = next((
                name for name, preset in TRANSITION_PRESETS.items()
                if preset["continuation_mode"] == mode and
                int(preset["context_length"]) == context
            ), None)
            resolved_transition = transition_policy(
                matching or "guide", expert_override=matching is None,
                continuation_mode=mode, context_length=context)
        restored["transition_policy"] = {
            "preset": resolved_transition["preset"],
            "expert_override": bool(
                resolved_transition.get("expert_override", False)),
            "expert_continuation_mode": resolved_transition[
                "continuation_mode"],
            "expert_context_length": int(
                resolved_transition["context_length"]),
        }
    except (TypeError, ValueError):
        pass
    return restored


def _scene_count_from_plan(path: str) -> int | None:
    if not os.path.isfile(path):
        return None
    try:
        archive = _read_json(path)
    except (OSError, json.JSONDecodeError):
        return None
    if not isinstance(archive, dict):
        return None
    editor = archive.get("editor_plan")
    shots = editor.get("shots") if isinstance(editor, dict) else archive.get("shots")
    return len(shots) if isinstance(shots, list) else None


def _iso_mtime(path: str) -> str:
    timestamp = os.path.getmtime(path)
    return datetime.fromtimestamp(timestamp, timezone.utc).isoformat(
        timespec="seconds").replace("+00:00", "Z")


class RunArchiveManager:
    def __init__(self, output_root: str, input_root: str | None = None):
        self.output_root = os.path.realpath(os.path.abspath(output_root))
        self.chains_root = os.path.join(self.output_root, "h3_chains")
        self.assets = RunAssetStore(self.output_root, input_root)

    def _run_dir(self, run_name: Any) -> tuple[str, str]:
        run = _safe_name(run_name)
        if not run:
            raise ValueError("A non-empty H3 chain run_name is required.")
        path = os.path.realpath(os.path.join(self.chains_root, run))
        root = os.path.realpath(self.output_root)
        if os.path.commonpath([root, path]) != root:
            raise ValueError("H3 run path escapes the output directory.")
        return path, run

    def list_runs(self) -> list[dict[str, Any]]:
        if not os.path.isdir(self.chains_root):
            return []
        runs = []
        for entry in os.scandir(self.chains_root):
            if not entry.is_dir(follow_symlinks=False):
                continue
            run_name = _safe_name(entry.name)
            if not run_name or run_name != entry.name:
                continue
            directory = entry.path
            plan_path = os.path.join(directory, "plan.json")
            api_path = os.path.join(directory, "api_prompt.json")
            workflow_path = os.path.join(directory, "workflow.json")
            archive_paths = [path for path in (plan_path, api_path, workflow_path)
                             if os.path.isfile(path)]
            checkpoint_dir = os.path.join(directory, "checkpoints")
            asset_manifest = os.path.join(directory, "references", "manifest.json")
            checkpoints = 0
            if os.path.isdir(checkpoint_dir):
                checkpoints = sum(
                    1 for name in os.listdir(checkpoint_dir)
                    if re.fullmatch(r"clip_\d{4}\.json", name))
            modified_paths = [directory]
            if os.path.isdir(checkpoint_dir):
                modified_paths.append(checkpoint_dir)
            if os.path.isfile(asset_manifest):
                modified_paths.append(asset_manifest)
            modified_paths.extend(archive_paths)
            newest = max(modified_paths, key=os.path.getmtime)
            try:
                asset_summary = self.assets.summary(run_name)
            except (OSError, ValueError, json.JSONDecodeError):
                asset_summary = {"asset_count": 0, "asset_bytes": 0}
            runs.append({
                "run_name": run_name,
                "modified_at": _iso_mtime(newest),
                "scene_count": _scene_count_from_plan(plan_path),
                "checkpoint_count": checkpoints,
                "restorable": bool(archive_paths),
                "archive_bytes": (
                    sum(os.path.getsize(path) for path in archive_paths)
                    + int(asset_summary.get("asset_bytes") or 0)),
                **asset_summary,
                "sources": {
                    "api_prompt": os.path.isfile(api_path),
                    "workflow": os.path.isfile(workflow_path),
                    "plan": os.path.isfile(plan_path),
                },
            })
        runs.sort(key=lambda item: item["modified_at"], reverse=True)
        return runs

    def load_plan(self, run_name: Any) -> dict[str, Any]:
        """Load the complete saved Plan without restoring archived assets."""
        directory, run = self._run_dir(run_name)
        if not os.path.isdir(directory):
            raise ValueError("H3 run %r does not exist." % run)

        restored: dict[str, Any] = {}
        policy_inputs: dict[str, dict[str, Any]] = {}
        sources = []
        warnings = []
        plan_path = os.path.join(directory, "plan.json")
        if os.path.isfile(plan_path):
            try:
                archive = _read_json(plan_path)
                restored.update(_archive_inputs(archive, run))
                policy_inputs.update(archive_policy_inputs(archive))
                if restored:
                    sources.append("plan.json")
            except (OSError, TypeError, ValueError, json.JSONDecodeError) as exc:
                warnings.append("plan.json: %s" % exc)

        workflow_path = os.path.join(directory, "workflow.json")
        if os.path.isfile(workflow_path):
            try:
                values = _workflow_inputs(_read_json(workflow_path), run)
                if values:
                    restored.update(values)
                    sources.append("workflow.json")
            except (OSError, json.JSONDecodeError) as exc:
                warnings.append("workflow.json: %s" % exc)

        api_path = os.path.join(directory, "api_prompt.json")
        if os.path.isfile(api_path):
            try:
                values = _api_prompt_inputs(_read_json(api_path), run)
                if values:
                    restored.update(values)
                    sources.append("api_prompt.json")
            except (OSError, json.JSONDecodeError) as exc:
                warnings.append("api_prompt.json: %s" % exc)

        restored["run_name"] = run
        if "plan_json" not in restored:
            detail = "; ".join(warnings) if warnings else "no usable Plan archive was found"
            raise ValueError("Could not restore H3 run %r: %s." % (run, detail))
        parsed = json.loads(restored["plan_json"])
        shots = parsed.get("shots") if isinstance(parsed, dict) else parsed
        return {
            "run_name": run,
            "scene_count": len(shots) if isinstance(shots, list) else None,
            "plan_inputs": restored,
            "policy_inputs": policy_inputs,
            "sources": sources,
            "warnings": warnings,
        }

    def load_run(self, run_name: Any) -> dict[str, Any]:
        payload = self.load_plan(run_name)
        run = payload["run_name"]
        warnings = list(payload["warnings"])
        try:
            assets = self.assets.prepare_restore(run)
        except (OSError, ValueError, json.JSONDecodeError) as exc:
            assets = {"run_name": run, "bindings": [], "warnings": [str(exc)]}
        warnings.extend(assets.get("warnings") or [])
        return {
            **payload,
            "warnings": warnings,
            "assets": assets,
        }