dify/api/services/workflow_variable_reference_validator.py

from __future__ import annotations

import json
import logging
import re
from collections import defaultdict, deque
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any

from core.trigger.constants import TRIGGER_NODE_TYPES
from core.workflow.nodes.human_input.constants import TIMEOUT_HANDLE
from graphon.enums import BuiltinNodeTypes, ErrorStrategy

logger = logging.getLogger(__name__)

_RESERVED_SELECTOR_HEADS: frozenset[str] = frozenset({"sys", "env", "conversation", "start"})

_REFERENCE_EXEMPT_NODE_TYPES: frozenset[str] = frozenset(
    {
        BuiltinNodeTypes.VARIABLE_AGGREGATOR,
        BuiltinNodeTypes.LEGACY_VARIABLE_AGGREGATOR,
    }
)

_BRANCH_NODE_TYPES: frozenset[str] = frozenset(
    {
        BuiltinNodeTypes.IF_ELSE,
        BuiltinNodeTypes.QUESTION_CLASSIFIER,
        BuiltinNodeTypes.HUMAN_INPUT,
    }
)

_TEMPLATE_REFERENCE_PATTERN = re.compile(
    r"\{\{#([a-zA-Z0-9_]{1,50})\.[a-zA-Z_][a-zA-Z0-9_]{0,29}(?:\.[a-zA-Z_][a-zA-Z0-9_]{0,29}){0,9}#\}\}"
)
_MAX_REPORTED_ISSUES = 10


@dataclass(frozen=True)
class VariableReferenceIssue:
    """A single unsafe variable reference: ``node`` reads output from ``referenced_node``."""

    node_id: str
    node_title: str
    referenced_node_id: str
    referenced_node_title: str


def validate_variable_references(graph: Mapping[str, Any]) -> list[VariableReferenceIssue]:
    """Flag references where the consumer can run without the producer.

    Uses engine execution semantics: nodes fire on any active inbound edge; branch
    nodes (if-else / question-classifier / human-input / fail-branch) take one handle
    (possibly unwired); other nodes fan out all outbound edges. Returns [] when clean.
    """
    nodes = graph.get("nodes") or []
    edges = graph.get("edges") or []
    if not isinstance(nodes, list) or not isinstance(edges, list) or not nodes:
        return []

    node_type: dict[str, str] = {}
    node_title: dict[str, str] = {}
    node_parent: dict[str, str | None] = {}
    node_data: dict[str, Mapping[str, Any]] = {}
    for node in nodes:
        if not isinstance(node, Mapping):
            continue
        node_id = node.get("id")
        if not isinstance(node_id, str):
            continue
        raw_data = node.get("data")
        data: Mapping[str, Any] = raw_data if isinstance(raw_data, Mapping) else {}
        node_type[node_id] = str(data.get("type") or "")
        title = data.get("title")
        node_title[node_id] = str(title) if title else node_id
        parent = node.get("parentId")
        node_parent[node_id] = parent if isinstance(parent, str) else None
        node_data[node_id] = data

    node_ids = set(node_type)

    out_targets_by_handle: dict[str, dict[str | None, list[str]]] = defaultdict(lambda: defaultdict(list))
    predecessors: dict[str, list[str]] = defaultdict(list)
    successors: dict[str, list[str]] = defaultdict(list)
    for edge in edges:
        if not isinstance(edge, Mapping):
            continue
        source = edge.get("source")
        target = edge.get("target")
        if not isinstance(source, str) or not isinstance(target, str):
            continue
        out_targets_by_handle[source][edge.get("sourceHandle")].append(target)
        predecessors[target].append(source)
        successors[source].append(target)

    entries = [
        nid
        for nid in node_ids
        if node_parent[nid] is None and node_type[nid] in (BuiltinNodeTypes.START, *TRIGGER_NODE_TYPES)
    ]
    reachable = _reachable_from(entries, successors)
    exclusive = {
        nid
        for nid in node_ids
        if node_type[nid] in _BRANCH_NODE_TYPES or node_data[nid].get("error_strategy") == ErrorStrategy.FAIL_BRANCH
    }
    # A selectable handle can be unwired. It still lets the branch skip a producer.
    for nid in exclusive:
        data = node_data[nid]
        handles: list[str] = []
        if node_type[nid] == BuiltinNodeTypes.IF_ELSE:
            cases = data.get("cases")
            handles = [case["case_id"] for case in cases] if isinstance(cases, list) else ["true"]
            handles.append("false")
        elif node_type[nid] == BuiltinNodeTypes.QUESTION_CLASSIFIER:
            handles = [item["id"] for item in data.get("classes", [])]
        elif node_type[nid] == BuiltinNodeTypes.HUMAN_INPUT:
            handles = [action["id"] for action in data.get("user_actions", [])]
            handles.append(TIMEOUT_HANDLE)
        if data.get("error_strategy") == ErrorStrategy.FAIL_BRANCH:
            handles.extend(["source", "fail-branch"])
        for handle in handles:
            out_targets_by_handle[nid].setdefault(handle, [])

    consumers_by_producer: dict[str, set[str]] = defaultdict(set)
    for node_id in node_ids:
        if node_parent[node_id] is not None or node_id not in reachable:
            continue
        if node_type[node_id] in _REFERENCE_EXEMPT_NODE_TYPES:
            continue
        for referenced_id in _referenced_node_ids(node_data[node_id]):
            if referenced_id == node_id or referenced_id not in node_ids:
                continue
            if node_parent[referenced_id] is not None:
                continue
            consumers_by_producer[referenced_id].add(node_id)

    issues: list[VariableReferenceIssue] = []
    for producer, consumers in sorted(consumers_by_producer.items()):
        runnable_without_producer = _nodes_runnable_without(
            producer,
            entries=entries,
            predecessors=predecessors,
            successors=successors,
            out_targets_by_handle=out_targets_by_handle,
            exclusive=exclusive,
        )
        for consumer in sorted(consumers):
            if consumer in runnable_without_producer:
                issues.append(
                    VariableReferenceIssue(
                        node_id=consumer,
                        node_title=node_title[consumer],
                        referenced_node_id=producer,
                        referenced_node_title=node_title[producer],
                    )
                )

    return issues


def format_variable_reference_errors(issues: Sequence[VariableReferenceIssue]) -> str:
    """Build a single-line warning listing unsafe reader ← producer pairs."""
    count = len(issues)
    shown = issues[:_MAX_REPORTED_ISSUES]
    pairs = "; ".join(f'"{issue.node_title}" ← "{issue.referenced_node_title}"' for issue in shown)
    if count > len(shown):
        pairs += f"; +{count - len(shown)} more"
    return (
        f"{count} variable reference{'s' if count != 1 else ''} may read a skipped branch "
        f"output. Use a Variable Aggregator or a default value — {pairs}."
    )


def advisory_variable_reference_warning(graph_text: str | None) -> str | None:
    """Return a non-blocking publish warning. A checker failure must not fail publish."""
    if not graph_text:
        return None
    try:
        graph = json.loads(graph_text)
        if not isinstance(graph, dict):
            return None
        issues = validate_variable_references(graph)
        return format_variable_reference_errors(issues) if issues else None
    except Exception:
        logger.warning("Skipped advisory variable reference check", exc_info=True)
        return None


def _reachable_from(starts: Sequence[str], successors: Mapping[str, list[str]]) -> set[str]:
    visited: set[str] = set()
    queue: deque[str] = deque(starts)
    while queue:
        current = queue.popleft()
        if current in visited:
            continue
        visited.add(current)
        for nxt in successors.get(current, ()):
            if nxt not in visited:
                queue.append(nxt)
    return visited


def _nodes_runnable_without(
    producer: str,
    *,
    entries: Sequence[str],
    predecessors: Mapping[str, list[str]],
    successors: Mapping[str, list[str]],
    out_targets_by_handle: Mapping[str, Mapping[str | None, list[str]]],
    exclusive: set[str],
) -> set[str]:
    """Return nodes that can execute in some run where ``producer`` does not."""
    forbidden = {producer}
    queue: deque[str] = deque([producer])
    while queue:
        node = queue.popleft()
        for pred in predecessors.get(node, ()):
            if pred in forbidden:
                continue
            if pred not in exclusive or (
                out_targets_by_handle[pred]
                and all(
                    any(target in forbidden for target in targets) for targets in out_targets_by_handle[pred].values()
                )
            ):
                forbidden.add(pred)
                queue.append(pred)

    runnable: set[str] = set()
    work: deque[str] = deque(entry for entry in entries if entry not in forbidden)
    while work:
        node = work.popleft()
        if node in runnable:
            continue
        runnable.add(node)
        # All edges sharing the selected handle activate together. If one would
        # execute the producer, none of that handle's siblings can avoid it.
        targets = (
            [
                target
                for branch_targets in out_targets_by_handle[node].values()
                if not any(target in forbidden for target in branch_targets)
                for target in branch_targets
            ]
            if node in exclusive
            else successors.get(node, ())
        )
        for nxt in targets:
            if nxt not in forbidden and nxt not in runnable:
                work.append(nxt)
    return runnable


def _referenced_node_ids(value: Any) -> Iterator[str]:
    """Yield node ids referenced via selectors and {{#node.field#}} templates."""
    if isinstance(value, Mapping):
        variable = value.get("value")
        if value.get("type") == "variable" and isinstance(variable, list) and variable:
            yield from _reference_head(variable[0])
        for key, child in value.items():
            if _is_selector_key(key) and isinstance(child, list) and child:
                yield from _reference_head(child[0])
            yield from _referenced_node_ids(child)
    elif isinstance(value, list):
        for item in value:
            yield from _referenced_node_ids(item)
    elif isinstance(value, str):
        for match in _TEMPLATE_REFERENCE_PATTERN.finditer(value):
            yield from _reference_head(match.group(1))


def _is_selector_key(key: object) -> bool:
    return isinstance(key, str) and (key == "selector" or key.endswith("_selector"))


def _reference_head(head: object) -> Iterator[str]:
    if isinstance(head, str) and head not in _RESERVED_SELECTOR_HEADS:
        yield head