dify/api/core/plugin/impl/base.py

import inspect
import json
import logging
from collections.abc import Callable, Generator
from typing import Any, cast
from urllib.parse import unquote

import httpx
from pydantic import BaseModel
from yarl import URL

from configs import dify_config
from core.helper.http_client_pooling import get_pooled_http_client
from core.plugin.endpoint.exc import EndpointSetupFailedError
from core.plugin.entities.plugin_daemon import PluginDaemonBasicResponse, PluginDaemonError, PluginDaemonInnerError
from core.plugin.impl.exc import (
    PluginDaemonBadRequestError,
    PluginDaemonClientSideError,
    PluginDaemonInternalServerError,
    PluginDaemonNotFoundError,
    PluginDaemonUnauthorizedError,
    PluginInvokeError,
    PluginLLMPollingUnsupportedError,
    PluginNotFoundError,
    PluginPermissionDeniedError,
    PluginUniqueIdentifierError,
)
from core.plugin.impl.first_token_timeout import FirstTokenTimeoutError, first_token_timeout_ctx
from core.trigger.errors import (
    EventIgnoreError,
    TriggerInvokeError,
    TriggerPluginInvokeError,
    TriggerProviderCredentialValidationError,
)
from graphon.model_runtime.errors.invoke import (
    InvokeAuthorizationError,
    InvokeBadRequestError,
    InvokeConnectionError,
    InvokeRateLimitError,
    InvokeServerUnavailableError,
)
from graphon.model_runtime.errors.validate import CredentialsValidateFailedError

plugin_daemon_inner_api_baseurl = URL(str(dify_config.PLUGIN_DAEMON_URL))
_plugin_daemon_timeout_config = cast(
    float | httpx.Timeout | None,
    getattr(dify_config, "PLUGIN_DAEMON_TIMEOUT", 600.0),
)
plugin_daemon_request_timeout: httpx.Timeout | None
match _plugin_daemon_timeout_config:
    case None:
        plugin_daemon_request_timeout = None
    case httpx.Timeout():
        plugin_daemon_request_timeout = _plugin_daemon_timeout_config
    case _:
        plugin_daemon_request_timeout = httpx.Timeout(_plugin_daemon_timeout_config)


def _read_timeout_for(first_token_timeout: float | None) -> httpx.Timeout | None:
    """Narrow the daemon request timeout's ``read`` component to the first-token budget.

    The daemon withholds the response headers until the model's first token and sends no
    heartbeat before it, so httpx's ``read`` timeout (which covers both the header wait and
    each body read) measures time-to-first-token. A non-positive budget disables the gate
    and keeps the default timeout. ``httpx.Timeout(base, read=x)`` is rejected when ``base``
    is a ``Timeout``, so the other components are copied explicitly.
    """
    base = plugin_daemon_request_timeout
    if not first_token_timeout or first_token_timeout <= 0:
        return base
    if base is None:
        return httpx.Timeout(None, read=first_token_timeout)
    return httpx.Timeout(connect=base.connect, read=first_token_timeout, write=base.write, pool=base.pool)


logger = logging.getLogger(__name__)

PLUGIN_DAEMON_MAX_PATH_LENGTH = 4096
PLUGIN_DAEMON_MAX_PATH_DECODE_DEPTH = 8

_httpx_client: httpx.Client = get_pooled_http_client(
    "plugin_daemon",
    lambda: httpx.Client(limits=httpx.Limits(max_keepalive_connections=50, max_connections=100), trust_env=False),
)


class BasePluginClient:
    def _request(
        self,
        method: str,
        path: str,
        headers: dict[str, str] | None = None,
        data: bytes | dict[str, Any] | str | None = None,
        params: dict[str, Any] | None = None,
        files: dict[str, Any] | None = None,
    ) -> httpx.Response:
        """
        Make a request to the plugin daemon inner API.
        """
        url, headers, prepared_data, params, files = self._prepare_request(path, headers, data, params, files)

        request_kwargs: dict[str, Any] = {
            "method": method,
            "url": url,
            "headers": headers,
            "params": params,
            "files": files,
            "timeout": plugin_daemon_request_timeout,
        }
        if isinstance(prepared_data, dict):
            request_kwargs["data"] = prepared_data
        elif prepared_data is not None:
            request_kwargs["content"] = prepared_data

        try:
            response = _httpx_client.request(**request_kwargs)
        except httpx.RequestError:
            logger.exception("Request to Plugin Daemon Service failed")
            raise PluginDaemonInnerError(code=-500, message="Request to Plugin Daemon Service failed")

        return response

    def _prepare_request(
        self,
        path: str,
        headers: dict[str, str] | None,
        data: bytes | dict[str, Any] | str | None,
        params: dict[str, Any] | None,
        files: dict[str, Any] | None,
    ) -> tuple[str, dict[str, str], bytes | dict[str, Any] | str | None, dict[str, Any] | None, dict[str, Any] | None]:
        if len(path) > PLUGIN_DAEMON_MAX_PATH_LENGTH:
            raise ValueError(f"Invalid plugin daemon path: path length exceeds {PLUGIN_DAEMON_MAX_PATH_LENGTH}")

        decoded_path = path
        for _ in range(PLUGIN_DAEMON_MAX_PATH_DECODE_DEPTH):
            next_decoded_path = unquote(decoded_path)
            if next_decoded_path == decoded_path:
                break
            decoded_path = next_decoded_path
        else:
            raise ValueError("Invalid plugin daemon path: path is too deeply encoded")

        if any(seg == ".." for seg in decoded_path.split("/")):
            raise ValueError(f"Invalid plugin daemon path: traversal sequence detected in {path!r}")
        url = plugin_daemon_inner_api_baseurl / path
        prepared_headers = dict(headers or {})
        prepared_headers["X-Api-Key"] = dify_config.PLUGIN_DAEMON_KEY
        prepared_headers.setdefault("Accept-Encoding", "gzip, deflate, br")

        # Inject traceparent header for distributed tracing
        self._inject_trace_headers(prepared_headers)

        prepared_data: bytes | dict[str, Any] | str | None = (
            data if isinstance(data, (bytes, str, dict)) or data is None else None
        )
        if isinstance(data, dict):
            if prepared_headers.get("Content-Type") == "application/json":
                prepared_data = json.dumps(data)
            else:
                prepared_data = data

        return str(url), prepared_headers, prepared_data, params, files

    def _inject_trace_headers(self, headers: dict[str, str]) -> None:
        """
        Inject W3C traceparent header for distributed tracing.

        This ensures trace context is propagated to plugin daemon even if
        HTTPXClientInstrumentor doesn't cover module-level httpx functions.
        """
        if not dify_config.ENABLE_OTEL:
            return

        import contextlib

        # Skip if already present (case-insensitive check)
        for key in headers:
            if key.lower() == "traceparent":
                return

        # Inject traceparent - works as fallback when OTEL instrumentation doesn't cover this call
        with contextlib.suppress(Exception):
            from core.helper.trace_id_helper import generate_traceparent_header

            traceparent = generate_traceparent_header()
            if traceparent:
                headers["traceparent"] = traceparent

    def _stream_request(
        self,
        method: str,
        path: str,
        params: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
        data: bytes | dict[str, Any] | None = None,
        files: dict[str, Any] | None = None,
    ) -> Generator[str, None, None]:
        """
        Make a stream request to the plugin daemon inner API
        """
        url, headers, prepared_data, params, files = self._prepare_request(path, headers, data, params, files)

        first_token_timeout = first_token_timeout_ctx.get()
        # Non-positive (None / 0 / negative) disables the gate; the read timeout stays default.
        first_token_gate = bool(first_token_timeout and first_token_timeout > 0)
        stream_kwargs: dict[str, Any] = {
            "method": method,
            "url": url,
            "headers": headers,
            "params": params,
            "files": files,
            # The daemon holds the response headers until the first token, so a per-request
            # read timeout equal to the first-token budget gates time-to-first-token.
            "timeout": _read_timeout_for(first_token_timeout),
        }
        if isinstance(prepared_data, dict):
            stream_kwargs["data"] = prepared_data
        elif prepared_data is not None:
            stream_kwargs["content"] = prepared_data

        first_token_seen = False

        try:
            with _httpx_client.stream(**stream_kwargs) as response:
                for raw_line in response.iter_lines():
                    # Blank keep-alive frames don't count as the first token, but each is still a
                    # successful read that refreshes httpx's read window -- first-token gating relies
                    # on the daemon sending no keep-alive before the first token.
                    if not raw_line:
                        continue
                    line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
                    line = line.strip()
                    if line.startswith("data:"):
                        line = line[5:].strip()
                    if not line:
                        continue
                    first_token_seen = True
                    yield line
        except httpx.ReadTimeout as e:
            # Gate on: a read timeout before the first line means the first token was too slow.
            # After the first line (inter-token stall), or with the gate off, it's a plain
            # transport error.
            if first_token_gate and not first_token_seen:
                raise FirstTokenTimeoutError(f"The first token was not received within {first_token_timeout}s.") from e
            logger.exception("Stream request to Plugin Daemon Service failed")
            raise PluginDaemonInnerError(code=-500, message="Request to Plugin Daemon Service failed")
        except httpx.RequestError:
            logger.exception("Stream request to Plugin Daemon Service failed")
            raise PluginDaemonInnerError(code=-500, message="Request to Plugin Daemon Service failed")

    def _stream_request_with_model[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
        self,
        method: str,
        path: str,
        type_: type[T],
        headers: dict[str, str] | None = None,
        data: bytes | dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
        files: dict[str, Any] | None = None,
    ) -> Generator[T, None, None]:
        """
        Make a stream request to the plugin daemon inner API and yield the response as a model.
        """
        for line in self._stream_request(method, path, params, headers, data, files):
            yield type_(**json.loads(line))  # type: ignore

    def _request_with_model[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
        self,
        method: str,
        path: str,
        type_: type[T],
        headers: dict[str, str] | None = None,
        data: bytes | None = None,
        params: dict[str, Any] | None = None,
        files: dict[str, Any] | None = None,
    ) -> T:
        """
        Make a request to the plugin daemon inner API and return the response as a model.
        """
        response = self._request(method, path, headers, data, params, files)
        return type_(**response.json())  # type: ignore[return-value]

    def _request_with_plugin_daemon_response[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
        self,
        method: str,
        path: str,
        type_: type[T],
        headers: dict[str, str] | None = None,
        data: bytes | dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
        files: dict[str, Any] | None = None,
        transformer: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
    ) -> T:
        """
        Make a request to the plugin daemon inner API and return the response as a model.
        """
        try:
            response = self._request(method, path, headers, data, params, files)
            response.raise_for_status()
        except httpx.HTTPStatusError as e:
            logger.exception("Failed to request plugin daemon, status: %s, url: %s", e.response.status_code, path)
            if e.response.status_code < 500:
                raise PluginDaemonClientSideError(description=str(e))
            else:
                raise PluginDaemonInternalServerError(description=str(e))
        except Exception as e:
            msg = f"Failed to request plugin daemon, url: {path}"
            logger.exception("Failed to request plugin daemon, url: %s", path)
            raise ValueError(msg) from e

        try:
            json_response = response.json()
            if transformer:
                json_response = transformer(json_response)
            # https://stackoverflow.com/questions/59634937/variable-foo-class-is-not-valid-as-type-but-why
            rep = PluginDaemonBasicResponse[type_].model_validate(json_response)  # type: ignore
        except Exception:
            msg = (
                f"Failed to parse response from plugin daemon to PluginDaemonBasicResponse [{str(type_.__name__)}],"
                f" url: {path}"
            )
            logger.exception(msg)
            raise ValueError(msg)

        if rep.code != 0:
            try:
                error = PluginDaemonError.model_validate(json.loads(rep.message))
            except Exception:
                raise ValueError(f"{rep.message}, code: {rep.code}")

            self._handle_plugin_daemon_error(error.error_type, error.message)
        if rep.data is None:
            frame = inspect.currentframe()
            raise ValueError(f"got empty data from plugin daemon: {frame.f_lineno if frame else 'unknown'}")

        return rep.data

    def _request_with_plugin_daemon_response_stream[T: BaseModel | dict[str, Any] | list[Any] | bool | str](
        self,
        method: str,
        path: str,
        type_: type[T],
        headers: dict[str, str] | None = None,
        data: bytes | dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
        files: dict[str, Any] | None = None,
    ) -> Generator[T, None, None]:
        """
        Make a stream request to the plugin daemon inner API and yield the response as a model.
        """
        for line in self._stream_request(method, path, params, headers, data, files):
            try:
                rep = PluginDaemonBasicResponse[type_].model_validate_json(line)  # type: ignore
            except (ValueError, TypeError):
                # TODO modify this when line_data has code and message
                try:
                    line_data = json.loads(line)
                except (ValueError, TypeError):
                    raise ValueError(line)
                # If the dictionary contains the `error` key, use its value as the argument
                # for `ValueError`.
                # Otherwise, use the `line` to provide better contextual information about the error.
                raise ValueError(line_data.get("error", line))

            if rep.code != 0:
                if rep.code == -500:
                    try:
                        error = PluginDaemonError.model_validate(json.loads(rep.message))
                    except Exception:
                        raise PluginDaemonInnerError(code=rep.code, message=rep.message)

                    logger.error("Error in stream response for plugin %s", rep.__dict__)
                    self._handle_plugin_daemon_error(error.error_type, error.message)
                raise ValueError(f"plugin daemon: {rep.message}, code: {rep.code}")
            if rep.data is None:
                frame = inspect.currentframe()
                raise ValueError(f"got empty data from plugin daemon: {frame.f_lineno if frame else 'unknown'}")
            yield rep.data

    def _handle_plugin_daemon_error(self, error_type: str, message: str):
        """
        handle the error from plugin daemon
        """
        match error_type:
            case PluginDaemonInnerError.__name__:
                raise PluginDaemonInnerError(code=-500, message=message)
            case PluginInvokeError.__name__:
                error_object = json.loads(message)
                invoke_error_type = error_object.get("error_type")
                match invoke_error_type:
                    case InvokeRateLimitError.__name__:
                        raise InvokeRateLimitError(description=error_object.get("message"))
                    case InvokeAuthorizationError.__name__:
                        raise InvokeAuthorizationError(description=error_object.get("message"))
                    case InvokeBadRequestError.__name__:
                        raise InvokeBadRequestError(description=error_object.get("message"))
                    case InvokeConnectionError.__name__:
                        raise InvokeConnectionError(description=error_object.get("message"))
                    case InvokeServerUnavailableError.__name__:
                        raise InvokeServerUnavailableError(description=error_object.get("message"))
                    case CredentialsValidateFailedError.__name__:
                        raise CredentialsValidateFailedError(error_object.get("message"))
                    case EndpointSetupFailedError.__name__:
                        raise EndpointSetupFailedError(error_object.get("message"))
                    case TriggerProviderCredentialValidationError.__name__:
                        raise TriggerProviderCredentialValidationError(error_object.get("message"))
                    case TriggerPluginInvokeError.__name__:
                        raise TriggerPluginInvokeError(description=error_object.get("message"))
                    case TriggerInvokeError.__name__:
                        raise TriggerInvokeError(error_object.get("message"))
                    case EventIgnoreError.__name__:
                        raise EventIgnoreError(description=error_object.get("message"))
                    # NOTE: current plugin sdk / plugin daemon does not raise exception with
                    # type `PluginLLMPollingUnsupportedError`.
                    case PluginLLMPollingUnsupportedError.__name__:
                        raise PluginLLMPollingUnsupportedError(description=error_object.get("message"))
                    case _:
                        raise PluginInvokeError(description=message)
            case PluginDaemonInternalServerError.__name__:
                raise PluginDaemonInternalServerError(description=message)
            case PluginDaemonBadRequestError.__name__:
                raise PluginDaemonBadRequestError(description=message)
            case PluginDaemonNotFoundError.__name__:
                raise PluginDaemonNotFoundError(description=message)
            case PluginUniqueIdentifierError.__name__:
                raise PluginUniqueIdentifierError(description=message)
            case PluginNotFoundError.__name__:
                raise PluginNotFoundError(description=message)
            case PluginDaemonUnauthorizedError.__name__:
                raise PluginDaemonUnauthorizedError(description=message)
            case PluginPermissionDeniedError.__name__:
                raise PluginPermissionDeniedError(description=message)
            case _:
                raise Exception(f"got unknown error from plugin daemon: {error_type}, message: {message}")