ComfyAssets-ComfyUI_PromptManager/utils/pair_execution.py

"""Bounded per-node execution evidence embedded in ComfyUI PNG metadata."""

import json
from typing import Any, Dict, Optional

EXECUTION_KEY = "comfyassets_prompt_pairs"
MAX_EXECUTION_RECORDS = 16
MAX_EXECUTION_NODES = 64
MAX_EXECUTION_TEXT = 100_000
MAX_EXECUTION_JSON = 500_000
_FIELDS = (
    "text",
    "negative_text",
    "positive_text",
    "negative_effective_text",
    "prepend_text",
    "append_text",
)


class ExecutionEvidenceError(ValueError):
    """Execution metadata is malformed, stale, truncated or ambiguous."""


def record_pair_execution(
    extra_pnginfo: Any,
    node_id: Any,
    *,
    text: str,
    negative_text: str,
    positive_text: str,
    negative_effective_text: str,
    prepend_text: str = "",
    append_text: str = "",
) -> None:
    """Append evidence without changing workflow or sibling extension metadata."""
    if not isinstance(extra_pnginfo, dict) or node_id is None:
        return
    envelope = extra_pnginfo.get(EXECUTION_KEY)
    if not isinstance(envelope, dict) or envelope.get("version") != 1:
        envelope = {"version": 1, "nodes": {}}
        extra_pnginfo[EXECUTION_KEY] = envelope
    if envelope.get("truncated"):
        return
    nodes = envelope.get("nodes")
    if not isinstance(nodes, dict):
        envelope["truncated"] = True
        return
    if str(node_id) not in nodes and len(nodes) >= MAX_EXECUTION_NODES:
        envelope["truncated"] = True
        return
    node = nodes.setdefault(str(node_id), {"records": []})
    if not isinstance(node, dict) or not isinstance(node.get("records"), list):
        envelope["truncated"] = True
        return
    record = {
        "text": text,
        "negative_text": negative_text,
        "positive_text": positive_text,
        "negative_effective_text": negative_effective_text,
        "prepend_text": prepend_text,
        "append_text": append_text,
    }
    if any(
        not isinstance(value, str) or len(value) > MAX_EXECUTION_TEXT
        for value in record.values()
    ):
        node["truncated"] = True
        return
    if record not in node["records"]:
        if len(node["records"]) >= MAX_EXECUTION_RECORDS:
            node["truncated"] = True
            return
        node["records"].append(record)
    try:
        if len(json.dumps(envelope, ensure_ascii=False)) > MAX_EXECUTION_JSON:
            extra_pnginfo[EXECUTION_KEY] = {
                "version": 1,
                "truncated": True,
                "nodes": {},
            }
    except (TypeError, ValueError, RecursionError):
        extra_pnginfo[EXECUTION_KEY] = {"version": 1, "truncated": True, "nodes": {}}


def read_pair_execution(metadata: Any, node_id: str) -> Optional[Dict[str, str]]:
    """Return one distinct record for a traced node; missing evidence stays absent."""
    if not isinstance(metadata, dict) or EXECUTION_KEY not in metadata:
        return None
    envelope = metadata[EXECUTION_KEY]
    if isinstance(envelope, str):
        if len(envelope) > MAX_EXECUTION_JSON:
            raise ExecutionEvidenceError("Execution evidence exceeds the size limit")
        try:
            envelope = json.loads(envelope)
        except (ValueError, TypeError, RecursionError) as error:
            raise ExecutionEvidenceError("Malformed execution evidence JSON") from error
    if not isinstance(envelope, dict) or envelope.get("version") != 1:
        raise ExecutionEvidenceError("Unsupported execution evidence format")
    nodes = envelope.get("nodes")
    if (
        envelope.get("truncated")
        or not isinstance(nodes, dict)
        or len(nodes) > MAX_EXECUTION_NODES
    ):
        raise ExecutionEvidenceError("Truncated or malformed execution evidence")
    if node_id not in nodes:
        return None
    node = nodes[node_id]
    if not isinstance(node, dict) or node.get("truncated"):
        raise ExecutionEvidenceError("Truncated or malformed node execution evidence")
    records = node.get("records")
    if (
        not isinstance(records, list)
        or not records
        or len(records) > MAX_EXECUTION_RECORDS
    ):
        raise ExecutionEvidenceError("Malformed node execution records")
    unique = {}
    for record in records:
        if not isinstance(record, dict) or any(
            not isinstance(record.get(key), str)
            or len(record[key]) > MAX_EXECUTION_TEXT
            for key in _FIELDS
        ):
            raise ExecutionEvidenceError("Malformed execution text values")
        clean = {key: record[key] for key in _FIELDS}
        if clean["negative_effective_text"] != clean["negative_text"]:
            raise ExecutionEvidenceError("Negative execution text must be unchanged")
        unique[json.dumps(clean, sort_keys=True)] = clean
    if len(unique) != 1:
        raise ExecutionEvidenceError(
            "Multiple distinct executions have no image anchor"
        )
    return next(iter(unique.values()))


def clear_pair_execution(request: Any) -> None:
    """Discard client-carried execution records before a new validated run."""
    if not isinstance(request, dict):
        return
    extra = request.get("extra_data")
    png = extra.get("extra_pnginfo") if isinstance(extra, dict) else None
    if isinstance(png, dict):
        png.pop(EXECUTION_KEY, None)