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 Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final

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

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

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

SOURCE_HANDLE: Final = "source"
_ELSE_HANDLE: Final = "false"
_LEGACY_IF_HANDLE: Final = "true"
_FAIL_BRANCH_HANDLE: Final = ErrorStrategy.FAIL_BRANCH.value

_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

logger = logging.getLogger(__name__)


def _ids(items: object, key: str) -> frozenset[str]:
    if not isinstance(items, list):
        return frozenset()
    return frozenset(str(item[key]) for item in items if isinstance(item, Mapping) and item.get(key) is not None)


def _if_else_handles(data: Mapping[str, Any]) -> frozenset[str]:
    cases = data.get("cases")
    return (_ids(cases, "case_id") if isinstance(cases, list) else frozenset({_LEGACY_IF_HANDLE})) | {_ELSE_HANDLE}


def _question_classifier_handles(data: Mapping[str, Any]) -> frozenset[str]:
    return _ids(data.get("classes"), "id")


def _human_input_handles(data: Mapping[str, Any]) -> frozenset[str]:
    return _ids(data.get("user_actions"), "id") | {TIMEOUT_HANDLE}


def _agent_route_handles(data: Mapping[str, Any]) -> frozenset[str] | None:
    routes = data.get("agent_output_routes")
    if not isinstance(routes, Mapping) or routes.get("enabled") is not True:
        return None
    return _ids(routes.get("routes"), "id")


@dataclass(frozen=True)
class _Branching:
    handles: Callable[[Mapping[str, Any]], frozenset[str] | None]
    # A fail-branch node is counted as also leaving by `source`; an agent choosing a route never does.
    leaves_by_source: bool = True


_BRANCHING: Final[Mapping[str, _Branching]] = {
    BuiltinNodeTypes.IF_ELSE: _Branching(_if_else_handles),
    BuiltinNodeTypes.QUESTION_CLASSIFIER: _Branching(_question_classifier_handles),
    BuiltinNodeTypes.HUMAN_INPUT: _Branching(_human_input_handles),
    BuiltinNodeTypes.AGENT: _Branching(_agent_route_handles, leaves_by_source=False),
}


def branch_handles(data: Mapping[str, Any]) -> frozenset[str] | None:
    """The source handles a node chooses one of at run time; None for a node that follows every outgoing edge.

    A fail-branch node becomes a branch node that can also leave by ``fail-branch``.
    """
    branching = _BRANCHING.get(str(data.get("type")))
    handles = branching.handles(data) if branching else None
    if data.get("error_strategy") != ErrorStrategy.FAIL_BRANCH:
        return handles
    if branching is None or handles is None:
        return frozenset({SOURCE_HANDLE, _FAIL_BRANCH_HANDLE})
    extra = {SOURCE_HANDLE, _FAIL_BRANCH_HANDLE} if branching.leaves_by_source else {_FAIL_BRANCH_HANDLE}
    return handles | extra


@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 (see ``branch_handles``) 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)
    handles_by_branch = {nid: handles for nid in node_ids if (handles := branch_handles(node_data[nid])) is not None}
    exclusive = set(handles_by_branch)
    # A selectable handle can be unwired. It still lets the branch skip a producer.
    for nid, handles in handles_by_branch.items():
        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