dify/api/services/workflow/tool_credentials.py

"""Dify-owned tool credential configuration, independent of Graphon's execution schema."""

from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_validator

from constants.end_user_connections import ConnectionCredentialType, OAuthClientSource
from graphon.enums import BuiltinNodeTypes
from graphon.nodes.tool.entities import ToolNodeData, ToolProviderType
from models.provider_ids import ToolProviderID
from services.end_user_connections.revision import AuthorizationRequirementSnapshot
from services.entities.tool_connection_entities import ToolConnectionRequirement
from services.workflow.tool_authorization import (
    first_end_user_tool_node,
    requires_end_user_connections,
)

__all__ = [
    "EndUserAuthConfig",
    "ToolConnectionRequirement",
    "ToolCredentialConfig",
    "ToolCredentialIssue",
    "ToolNodeSpec",
    "WorkflowToolEndUserAuthError",
    "check_tool_credentials",
    "first_end_user_tool_node",
    "get_tool_authorization_requirements",
    "get_tool_connection_requirements",
    "requires_end_user_connections",
    "resolve_tool_credential_root",
]


class EndUserAuthConfig(BaseModel):
    """Missing values represent an unfinished draft, not an executable configuration."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    allowed_credential_types: list[Literal["api-key", "oauth2"]] | None = Field(default=None, max_length=2)
    oauth_client_type: Literal["default", "custom"] | None = None
    purpose: str | None = Field(default=None, max_length=50)

    @property
    def oauth_client_source(self) -> OAuthClientSource | None:
        if self.oauth_client_type is None:
            return None
        return OAuthClientSource.PLATFORM if self.oauth_client_type == "default" else OAuthClientSource.WORKSPACE

    @field_validator("allowed_credential_types")
    @classmethod
    def validate_unique_credential_types(
        cls, value: list[Literal["api-key", "oauth2"]] | None
    ) -> list[Literal["api-key", "oauth2"]] | None:
        if value is not None and len(value) != len(set(value)):
            raise ValueError("Authentication methods must not contain duplicates.")
        return value


class ToolCredentialConfig(BaseModel):
    """Static per-node requirements; active Connections belong to the user's Session selections.

    Independent nodes can select different Connections for the same provider. Only an explicit
    reuse reference shares another node's authentication; personal Connection IDs stay out of the graph.
    Nodes using different tools from the same provider can share authentication.
    """

    model_config = ConfigDict(frozen=True)

    # Reuse delegates the effective source and settings to the referenced node.
    credential_source: Literal["workspace", "end_user", "reuse_node"] = "workspace"
    credential_id: str | None = None
    end_user_auth: EndUserAuthConfig | None = None
    reuse_node_id: str | None = Field(default=None, min_length=1)
    test_credential_id: str | None = Field(default=None, min_length=1)


class ToolNodeSpec(ToolNodeData, ToolCredentialConfig):
    """The editor-facing schema includes Dify policy without changing Graphon's node model."""


def get_tool_connection_requirements(graph: Mapping[str, Any]) -> tuple[ToolConnectionRequirement, ...]:
    """Project static requirements without loading providers, clients, or personal credentials."""
    nodes = {
        node["id"]: node["data"]
        for node in graph.get("nodes", [])
        if isinstance(node, Mapping)
        and node.get("type") != "custom-note"
        and isinstance(node.get("data"), Mapping)
        and node["data"].get("type") == BuiltinNodeTypes.TOOL
    }
    requirements: list[ToolConnectionRequirement] = []
    for node_id, data in nodes.items():
        if data.get("credential_source") != "end_user":
            continue
        config = ToolCredentialConfig.model_validate(data)
        auth = config.end_user_auth
        if auth is None or auth.purpose is None or not auth.allowed_credential_types:
            raise ValueError("Published EndUser authentication settings are incomplete.")
        title = data.get("title")
        requirements.append(
            ToolConnectionRequirement(
                node_id=node_id,
                node_ids=(
                    node_id,
                    *(
                        key
                        for key, value in nodes.items()
                        if value.get("credential_source") == "reuse_node" and value.get("reuse_node_id") == node_id
                    ),
                ),
                title=title if isinstance(title, str) and title else node_id,
                purpose=auth.purpose,
                provider=str(ToolProviderID(data["provider_id"])),
                allowed_credential_types=tuple(
                    sorted(ConnectionCredentialType(value) for value in auth.allowed_credential_types)
                ),
            )
        )
    return tuple(requirements)


def get_tool_authorization_requirements(graph: Mapping[str, Any]) -> tuple[AuthorizationRequirementSnapshot, ...]:
    nodes = {
        node["id"]: node["data"]
        for node in graph.get("nodes", [])
        if isinstance(node, Mapping)
        and isinstance(node.get("id"), str)
        and isinstance(node.get("data"), Mapping)
        and node.get("type") != "custom-note"
    }
    requirements: list[AuthorizationRequirementSnapshot] = []
    for item in get_tool_connection_requirements(graph):
        data = nodes[item.node_id]
        config = ToolCredentialConfig.model_validate(data)
        assert config.end_user_auth is not None
        requirements.append(
            AuthorizationRequirementSnapshot(
                node_id=item.node_id,
                node_ids=tuple(sorted(item.node_ids)),
                provider=item.provider,
                tool_name=data["tool_name"],
                purpose=item.purpose,
                allowed_credential_types=tuple(sorted(value.value for value in item.allowed_credential_types)),
                oauth_client_type=config.end_user_auth.oauth_client_type,
            )
        )
    return tuple(sorted(requirements, key=lambda item: item.node_id))


@dataclass(frozen=True)
class ToolCredentialIssue:
    node_id: str
    path: tuple[str | int, ...]
    message: str


class WorkflowToolEndUserAuthError(ValueError):
    def __init__(self, node_id: str) -> None:
        super().__init__(f"Node {node_id}: EndUser credentials are not supported in workflow tools.")


def _same_provider(data: Mapping[str, Any], target_data: Mapping[str, Any]) -> bool:
    if data.get("provider_type") != target_data.get("provider_type"):
        return False
    provider = data.get("provider_id")
    target_provider = target_data.get("provider_id")
    if not isinstance(provider, str) or not provider or not isinstance(target_provider, str) or not target_provider:
        return False
    if provider == target_provider:
        return True

    if data.get("provider_type") == ToolProviderType.BUILT_IN:
        from core.tools.tool_manager import ToolManager

        # Builtin lookup gives local providers precedence over plugin aliases of the same name.
        for reference in (provider, target_provider):
            try:
                ToolManager.get_hardcoded_provider(reference)
            except KeyError:
                continue
            return False

    try:
        return str(ToolProviderID(provider)) == str(ToolProviderID(target_provider))
    except ValueError:
        return False


def _configuration_issues(
    node_id: str, data: Mapping[str, Any], config: ToolCredentialConfig
) -> list[ToolCredentialIssue]:
    issues: list[ToolCredentialIssue] = []
    if config.credential_source != "end_user" and config.end_user_auth is not None:
        issues.append(
            ToolCredentialIssue(node_id, ("end_user_auth",), "Only EndUser nodes can declare authentication settings.")
        )
    if config.credential_source != "reuse_node" and config.reuse_node_id is not None:
        issues.append(ToolCredentialIssue(node_id, ("reuse_node_id",), "Only reuse nodes can reference another node."))
    if config.credential_source == "workspace":
        return issues

    if config.credential_id is not None:
        issues.append(
            ToolCredentialIssue(node_id, ("credential_id",), "Only workspace nodes can bind a production credential.")
        )
    if data.get("provider_type") not in (ToolProviderType.BUILT_IN, ToolProviderType.PLUGIN):
        issues.append(
            ToolCredentialIssue(
                node_id,
                ("provider_type",),
                "EndUser authentication and reuse require a builtin or plugin tool provider.",
            )
        )
    for field in ("provider_id", "tool_name"):
        value = data.get(field)
        if not isinstance(value, str) or not value.strip():
            issues.append(ToolCredentialIssue(node_id, (field,), "Field is required."))

    if config.credential_source == "reuse_node":
        if config.reuse_node_id is None or not config.reuse_node_id.strip():
            issues.append(ToolCredentialIssue(node_id, ("reuse_node_id",), "A reuse target is required."))
        if config.test_credential_id is not None:
            issues.append(
                ToolCredentialIssue(
                    node_id, ("test_credential_id",), "Reuse nodes inherit the target's test credential."
                )
            )
        return issues

    auth = config.end_user_auth
    if auth is None:
        issues.append(ToolCredentialIssue(node_id, ("end_user_auth",), "EndUser authentication is required."))
        return issues
    if not auth.allowed_credential_types:
        issues.append(
            ToolCredentialIssue(
                node_id,
                ("end_user_auth", "allowed_credential_types"),
                "At least one authentication method is required.",
            )
        )
    if auth.purpose is None or not auth.purpose.strip():
        issues.append(ToolCredentialIssue(node_id, ("end_user_auth", "purpose"), "Field is required."))
    if "oauth2" in (auth.allowed_credential_types or ()):
        if auth.oauth_client_type is None:
            issues.append(
                ToolCredentialIssue(
                    node_id, ("end_user_auth", "oauth_client_type"), "An OAuth client type is required."
                )
            )
    elif auth.oauth_client_type is not None:
        issues.append(
            ToolCredentialIssue(
                node_id, ("end_user_auth", "oauth_client_type"), "Only OAuth authentication can select a client type."
            )
        )
    return issues


def check_tool_credentials(graph: Mapping[str, Any], *, complete: bool) -> list[ToolCredentialIssue]:
    """Check supplied draft fields, optionally including completeness and graph-wide references.

    No credentials are loaded or validated remotely. Test credential availability is checked
    when a test uses it, so a deleted test reference cannot prevent production publication.
    """
    issues: list[ToolCredentialIssue] = []
    nodes = {
        node["id"]: node["data"]
        for node in graph.get("nodes", [])
        if isinstance(node, Mapping)
        and isinstance(node.get("id"), str)
        and isinstance(node.get("data"), Mapping)
        and node.get("type") != "custom-note"
    }
    configs: dict[str, ToolCredentialConfig] = {}
    for node_id, data in nodes.items():
        if data.get("type") != BuiltinNodeTypes.TOOL:
            continue
        try:
            configs[node_id] = ToolCredentialConfig.model_validate(data)
        except ValidationError as error:
            issues.extend(
                ToolCredentialIssue(node_id, problem["loc"], problem["msg"])
                for problem in error.errors(include_url=False, include_input=False, include_context=False)
            )

    if not complete:
        return issues

    for node_id, config in configs.items():
        issues.extend(_configuration_issues(node_id, nodes[node_id], config))
    invalid_nodes = {issue.node_id for issue in issues}

    for node_id, config in configs.items():
        if node_id in invalid_nodes or config.credential_source != "reuse_node":
            continue
        data = nodes[node_id]
        target_id = config.reuse_node_id
        target = configs.get(target_id) if target_id is not None else None
        target_data = nodes.get(target_id) if target_id is not None else None
        if target_id == node_id:
            message = "A node cannot reuse its own authentication."
        elif target_data is None:
            message = f"Reuse target {target_id!r} does not exist."
        elif target is None:
            message = f"Reuse target {target_id!r} must be a tool node with valid authentication settings."
        elif target.credential_source == "reuse_node":
            message = f"Reuse target {target_id!r} must configure authentication directly."
        elif target_id in invalid_nodes:
            message = f"Reuse target {target_id!r} has incomplete authentication settings."
        elif not _same_provider(data, target_data):
            message = f"Reuse target {target_id!r} must use the same provider type, plugin, and provider."
        else:
            continue
        issues.append(ToolCredentialIssue(node_id, ("reuse_node_id",), message))
    return issues


def resolve_tool_credential_root(graph: Mapping[str, Any], node_id: str) -> tuple[str, ToolNodeSpec]:
    """Resolve only this node's authentication source against its execution graph."""
    nodes = {node["id"]: node for node in graph.get("nodes", []) if node.get("type") != "custom-note"}
    try:
        node = nodes[node_id]
        config = ToolCredentialConfig.model_validate(node["data"])
        root_id = config.reuse_node_id if config.credential_source == "reuse_node" else node_id
        if root_id is None:
            raise ValueError("Missing authentication source.")
        root = nodes[root_id]
        scope = [node] if root_id == node_id else [node, root]
        if check_tool_credentials({"nodes": scope}, complete=True):
            raise ValueError("Invalid authentication source.")
        return root_id, ToolNodeSpec.model_validate(root["data"])
    except (KeyError, TypeError, ValidationError):
        raise ValueError("Invalid authentication source.") from None