dify/api/services/dify_builder/diagnostic_context.py

"""Model-facing diagnostic copies; never rewrite persisted execution evidence.

This profile is broader than credential configuration redaction and does not
participate in configuration-write validation. It hides structured credentials
and known credential values repeated in text, not arbitrary unknown prose.
"""

import re
from copy import deepcopy
from typing import Any
from urllib.parse import unquote_plus

from core.dify_builder.models import Graph, NodeOutput, Run
from services.dify_builder.credentials import REDACTED

_SENSITIVE_FIELDS = frozenset(
    {
        "apikey",
        "accesstoken",
        "refreshtoken",
        "authtoken",
        "bearertoken",
        "token",
        "secret",
        "secretkey",
        "clientsecret",
        "password",
        "credential",
        "credentials",
        "privatekey",
        "xapikey",
    }
)
_CONTAINERS = frozenset({"headers", "params", "query", "queryparams", "queryparameters", "cookies"})


def _field_name(key: str) -> str:
    return "".join(char for char in key.lower() if char.isalnum())


def _hide(value: Any, known: set[str]) -> Any:
    if isinstance(value, dict):
        return {key: _hide(item, known) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_hide(item, known) for item in value]
    if isinstance(value, str) and value and value != REDACTED:
        known.add(value)
        stripped = value.strip()
        if stripped:
            known.add(stripped)
        scheme, separator, token = stripped.partition(" ")
        if separator and scheme.lower() in {"bearer", "basic"} and token.strip():
            known.add(token.strip())
    return REDACTED if value else value


def _container(value: Any, known: set[str], field: str) -> Any:
    if isinstance(value, str):
        if field.startswith("query") or (field == "params" and ":" not in value and "=" in value):
            parts = []
            for part in value.split("&"):
                name, separator, item = part.partition("=")
                if separator and item:
                    _hide(unquote_plus(item), known)
                    parts.append(f"{name}={_hide(item, known)}")
                else:
                    parts.append(part if separator else str(_hide(part, known)))
            return "&".join(parts)
        # Native HTTP headers/params are newline-delimited name:value pairs.
        lines = []
        for line in value.splitlines():
            name, separator, item = line.partition(":")
            if separator:
                lines.append(f"{name}:{_hide(item.strip(), known)}" if item.strip() else line)
            else:
                lines.append(str(_hide(line, known)))
        return "\n".join(lines)
    if isinstance(value, list):
        return [
            {key: item if key in {"key", "name"} else _hide(item, known) for key, item in row.items()}
            if isinstance(row, dict)
            else _hide(row, known)
            for row in value
        ]
    return _hide(value, known)


def _authorization(value: Any, known: set[str]) -> Any:
    if not isinstance(value, dict):
        return _hide(value, known)
    return {
        key: item
        if key in {"type", "header"}
        else _authorization(item, known)
        if key == "config"
        else _hide(item, known)
        for key, item in value.items()
    }


def _structured(value: Any, known: set[str]) -> Any:
    if isinstance(value, dict):
        result = {}
        for key, item in value.items():
            field = _field_name(str(key))
            if field == "authorization":
                result[key] = _authorization(item, known)
            elif field in _CONTAINERS:
                result[key] = _container(item, known, field)
            elif field in _SENSITIVE_FIELDS:
                result[key] = _hide(item, known)
            else:
                result[key] = _structured(item, known)
        return result
    if isinstance(value, (list, tuple)):
        return [_structured(item, known) for item in value]
    if isinstance(value, str):
        _collect_text_credentials(value, known)
    return value


def _replace_text(text: str, known: list[str]) -> str:
    if not known:
        return text
    # Match whole existing sentinel spans before shorter secrets within them,
    # including sentinels inserted during structured redaction. One pass also
    # protects newly inserted sentinels. Longest first handles overlaps.
    protected = sorted({REDACTED, *known}, key=len, reverse=True)
    return re.sub("|".join(re.escape(secret) for secret in protected), lambda _: REDACTED, text)


def _replace_known(value: Any, known: list[str]) -> Any:
    if isinstance(value, str):
        return _replace_text(value, known)
    if isinstance(value, dict):
        return {
            _replace_text(key, known) if isinstance(key, str) else key: _replace_known(item, known)
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [_replace_known(item, known) for item in value]
    return value


def _collect_text_credentials(text: str, known: set[str]) -> None:
    """Collect explicit credential assignments in arbitrary source strings."""
    # Match only the key prefix. A non-sensitive outer assignment must not
    # consume a value containing another explicit credential assignment.
    for match in re.finditer(r"(?i)\b([a-z][a-z0-9_-]{0,63})[\"']?\s*[=:]\s*", text):
        if _field_name(match.group(1)) not in _SENSITIVE_FIELDS:
            continue
        value = re.match(r"(?:\"([^\"]+)\"|'([^']+)'|([^\s\"',;}]+))", text[match.end() :])
        if value:
            _hide(next(group for group in value.groups() if group is not None), known)
    for match in re.finditer(r"(?i)\b(?:bearer|basic)\s+([^\s\"',;]+)", text):
        _hide(match.group(1), known)


def _known_credentials(failed_run: Run, graph: Graph, node_outputs: list[NodeOutput]) -> set[str]:
    """Collect from raw provenance before any projection discards values."""
    known: set[str] = set()
    _structured(graph, known)
    _collect_text_credentials(failed_run.error, known)
    for output in [*failed_run.per_node, *node_outputs]:
        _collect_text_credentials(output.error, known)
        _structured(output.inputs, known)
        _structured(output.outputs, known)
    return known


def redact_diagnostic_text(text: str, failed_run: Run, graph: Graph, node_outputs: list[NodeOutput]) -> str:
    """Apply the complete graph/runtime credential policy to rendered prompt text."""
    known = _known_credentials(failed_run, graph, node_outputs)
    return _replace_text(text, sorted(known, key=len, reverse=True))


def redact_run_context(failed_run: Run, graph: Graph, node_outputs: list[NodeOutput]) -> tuple[Run, list[NodeOutput]]:
    """Return independent, sanitized diagnostic evidence before prompt truncation."""
    known = _known_credentials(failed_run, graph, node_outputs)
    safe_run, safe_outputs = deepcopy((failed_run, node_outputs))
    # Gather structured secrets first, so earlier error/body text can safely
    # repeat a credential held by a later node or output.
    for output in [*safe_run.per_node, *safe_outputs]:
        output.inputs = _structured(output.inputs, known)
        output.outputs = _structured(output.outputs, known)
    secrets = sorted(known, key=len, reverse=True)
    safe_run.error = _replace_text(safe_run.error, secrets)
    for output in [*safe_run.per_node, *safe_outputs]:
        output.node_id = _replace_text(output.node_id, secrets)
        output.title = _replace_text(output.title, secrets)
        output.error = _replace_text(output.error, secrets)
        output.inputs = _replace_known(output.inputs, secrets)
        output.outputs = _replace_known(output.outputs, secrets)
    return safe_run, safe_outputs


def redact_diagnostic_value(value: Any, failed_run: Run, graph: Graph, node_outputs: list[NodeOutput]) -> Any:
    """Sanitize graph projection from the full provenance before truncation."""
    known = _known_credentials(failed_run, graph, node_outputs)
    structured = _structured(deepcopy(value), known)
    return _replace_known(structured, sorted(known, key=len, reverse=True))


def redact_dsl_context(dsl: Graph) -> Graph:
    """Return a credential-safe continuation copy from fresh or historical DSL.

    Accept a DSL context mapping, collecting secret environment provenance before
    withholding it. Reuse the diagnostic profile for all other structured fields
    and known echoes. The source graph, execution evidence and write-admission
    policy remain untouched; this projection is only for context and prompts.
    """
    safe = deepcopy(dsl)
    known: set[str] = set()
    for variable in safe.get("environment_variables", []):
        if isinstance(variable, dict) and variable.get("value_type") == "secret":
            variable["value"] = _hide(variable.get("value"), known)
    structured = _structured(safe, known)
    return _replace_known(structured, sorted(known, key=len, reverse=True))