dify/api/controllers/openapi/_contract.py

"""``@accepts`` and ``@returns`` own one slice of the contract from a single model
reference, so the advertised and enforced contracts can't drift.

``@endpoint`` composes both under ``subject_router.guard``, exposing
``view.__handler__`` as the one documented test seam. ``returns`` is still used
bare by the unauthenticated ``index.py`` probes, which have no auth layer to
compose with.

``endpoint()`` also builds the route's ``EndpointSpec`` (``auth/spec.py``), which
carries the catalog fields ``op``, ``kind``, ``summary``, ``internal`` and
``deprecated`` that ``_catalog.py`` reads to describe ``/openapi/v1``. A ``list``
op gets ``paginated`` as well, which fills the next-page hint on a bare
``PaginationEnvelope`` result.
"""

from __future__ import annotations

import inspect
from collections.abc import Callable, Sequence
from functools import wraps
from typing import Any, Final, cast, get_args, get_origin, override

from flask import request
from flask_restx import abort
from pydantic import BaseModel, ValidationError
from sqlalchemy.orm import Session
from werkzeug.exceptions import RequestEntityTooLarge

from constants import BYTES_PER_MIB
from constants.personal_access_token import PersonalAccessTokenPolicy
from controllers.common.schema import query_params_from_model, query_params_from_request
from controllers.openapi import openapi_ns
from controllers.openapi._errors import ErrorBody, RequestOverRouteLimits
from controllers.openapi._hints import next_page_hint, page_after
from controllers.openapi._models import PaginationEnvelope
from controllers.openapi._multipart import body_from_request
from controllers.openapi._upload import file_fields
from controllers.openapi.auth.context import Context
from controllers.openapi.auth.requirements import Requirement
from controllers.openapi.auth.router import subject_router
from controllers.openapi.auth.spec import AdmissionContext, CatalogMeta, Effect, EndpointSpec, Example, Kind
from controllers.openapi.auth.subjects import Subject
from enums import DeploymentEdition

__all__ = [
    "MULTIPART_HEADER_BYTES_PER_PART",
    "CapRequestBody",
    "Effect",
    "Example",
    "Kind",
    "accepts",
    "endpoint",
    "multipart_header_allowance",
    "op_of",
    "paginated",
    "returns",
]

_INJECTED_KWARGS: Final = frozenset({"ctx", "query", "body"})
MULTIPART_HEADER_BYTES_PER_PART: Final = 1024


def paginated(op: str) -> Callable:
    """Add the ``Next page`` hint to a ``PaginationEnvelope`` result, unless the view already
    built one for this op. Sits inside ``accepts`` so it sees the validated ``query`` and the
    path kwargs the hint echoes; only this layer knows both those and the op id.
    """

    def decorator(view: Callable) -> Callable:
        @wraps(view)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            result = view(*args, **kwargs)
            if not isinstance(result, PaginationEnvelope) or any(h.op == op for h in result.hints):
                return result
            path_args = {name: value for name, value in kwargs.items() if name not in _INJECTED_KWARGS}
            page = page_after(page=result.page, limit=result.limit, has_more=result.has_more)
            result.hints = [
                *result.hints,
                *next_page_hint(op=op, path_args=path_args, query=kwargs.get("query"), page=page),
            ]
            return result

        return wrapper

    return decorator


class CapRequestBody(Requirement):
    def __init__(self, *, max_bytes: Callable[[], int], max_parts: Callable[[], int]) -> None:
        self.max_bytes = max_bytes
        self.max_parts = max_parts

    @override
    def run(self, subject: Subject, ctx: Context, session: Session) -> None:
        request.max_content_length = self.max_bytes()
        request.max_form_parts = self.max_parts()


def multipart_header_allowance(parts: int) -> int:
    return parts * MULTIPART_HEADER_BYTES_PER_PART


def _read_body(body_file_fields: frozenset[str]) -> dict[str, Any]:
    try:
        return body_from_request(file_fields=body_file_fields)
    except RequestEntityTooLarge as exc:
        limit = request.max_content_length
        if limit is None:
            raise
        raise RequestOverRouteLimits(
            f"{RequestOverRouteLimits.description}: at most {request.max_form_parts} parts"
            f" and {limit / BYTES_PER_MIB:g} MiB"
        ) from exc


def _list_fields(model: type[BaseModel]) -> tuple[str, ...]:
    """Query fields typed as a list, which a repeated parameter fills."""
    return tuple(
        name
        for name, field in model.model_fields.items()
        if any(get_origin(part) is list for part in (field.annotation, *get_args(field.annotation)))
    )


def accepts(*, query: type[BaseModel] | None = None, body: type[BaseModel] | None = None) -> Callable:
    """Validate ``query``/``body`` against the models and inject them as keyword-only kwargs."""

    def decorator(view: Callable) -> Callable:
        body_file_fields = file_fields(body) if body is not None else frozenset()
        query_list_fields = _list_fields(query) if query is not None else ()

        @wraps(view)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            try:
                if query is not None:
                    kwargs["query"] = query_params_from_request(query, list_fields=query_list_fields)
                if body is not None:
                    kwargs["body"] = body.model_validate(_read_body(body_file_fields))
            except ValidationError as exc:
                # Sanitized 422 — no pydantic `url` (version) or `input` (user payload) leak.
                abort(
                    422,
                    message="Request validation failed",
                    errors=exc.errors(include_url=False, include_input=False, include_context=False),
                )
            return view(*args, **kwargs)

        if query is not None:
            openapi_ns.doc(params=query_params_from_model(query))(wrapper)
        if body is not None:
            openapi_ns.expect(openapi_ns.models[body.__name__])(wrapper)
        if query is not None or body is not None:
            openapi_ns.response(422, "Validation error", openapi_ns.models[ErrorBody.__name__])(wrapper)
        return wrapper

    return decorator


def returns(code: int, model: type[BaseModel], description: str | None = None) -> Callable:
    """Serialise the handler's returned model and emit the response schema.

    Accepts a ``BaseModel`` (serialised with ``code``) or a ``(model, status[, headers])``
    tuple (status/headers honoured). Other returns — a bare ``(dict, status)``, an SSE
    ``Response`` — pass through untouched.
    """

    def decorator(view: Callable) -> Callable:
        @wraps(view)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            result = view(*args, **kwargs)
            if isinstance(result, BaseModel):
                return result.model_dump(mode="json", by_alias=True), code
            if isinstance(result, tuple) and result and isinstance(result[0], BaseModel):
                payload, *rest = result
                return (payload.model_dump(mode="json", by_alias=True), *rest)
            return result

        openapi_ns.response(code, description or model.__name__, openapi_ns.models[model.__name__])(wrapper)
        openapi_ns.response("default", "Error", openapi_ns.models[ErrorBody.__name__])(wrapper)
        return wrapper

    return decorator


_apply_returns = returns

ReturnSpec = tuple[int, type[BaseModel], str | None]


def _normalize_returns(returns: ReturnSpec | Sequence[ReturnSpec] | None) -> tuple[ReturnSpec, ...]:
    """One ``(code, model, description)`` or several, always ending up a tuple of
    them — the one place that tells the two apart, so `decorator` below never has to.
    """
    if returns is None:
        return ()
    if isinstance(returns[0], int):
        return (cast(ReturnSpec, returns),)  # pyrefly: ignore[redundant-cast]
    return tuple(cast(Sequence[ReturnSpec], returns))  # pyrefly: ignore[redundant-cast]


def op_of(view: Any) -> str:
    """The op id a route declared on ``endpoint``; a hint targeting that route reads it here."""

    return view.__spec__.catalog.op


def endpoint(
    *,
    op: str,
    kind: Kind,
    summary: str,
    requirements: Sequence[Requirement] = (),
    query: type[BaseModel] | None = None,
    body: type[BaseModel] | None = None,
    returns: ReturnSpec | Sequence[ReturnSpec] | None = None,
    edition: frozenset[DeploymentEdition] | None = None,
    context: AdmissionContext | None = "orm",
    internal: bool = False,
    deprecated: bool = False,
    effect: Effect | None = None,
    examples: Sequence[Example] = (),
    pat_policy: PersonalAccessTokenPolicy | None = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """The one seam a route attaches to for auth, request validation and response
    serialisation — auth, then ``accepts``, then ``returns``. Exposes
    ``view.__handler__`` (the bare handler) and ``view.__spec__`` (the exact
    `EndpointSpec` instance the router runs).

    ``context`` selects the handler's ``ctx`` argument and admission Session lifetime:

    - ``"orm"`` (default): the authentication Context with ORM objects and an open Session.
    - ``"workspace"``: a RequestContext with account and workspace identity.
    - ``"app"``: an AppRequestContext with tenant and app identity.
    - ``None``: no ``ctx`` argument; admission requirements still run.

    For ``"workspace"``, ``"app"`` and ``None``, the pipeline commits and closes
    the admission Session after preparing the arguments and before calling the
    handler. Services called by these handlers own their persistence Sessions.

    ``returns`` takes one ``(code, model, description)`` or several; a route
    that stacked N ``@returns`` today declares them here in the same top-to-bottom
    order and gets the identical nesting back — first entry outermost, last entry
    closest to the handler — including the N-times ``"default"`` error registration
    that stacking N ``@returns`` already produces.

    ``op``/``kind``/``summary`` feed the catalog (see ``_catalog.py``); ``internal``
    hides the op from the CLI's default listing; ``deprecated`` marks it as kept
    for compatibility only; ``effect`` overrides the effect the HTTP method implies;
    ``examples`` are complete inputs the catalog shows
    next to the schema.
    """
    requirements = tuple(requirements)
    for requirement in requirements:
        if not isinstance(requirement, Requirement):
            raise TypeError(f"requirements must be instances of Requirement, not {requirement!r}")
    spec = EndpointSpec(
        requirements=requirements,
        edition=edition,
        context=context,
        pat_policy=pat_policy,
        catalog=CatalogMeta(
            op=op,
            kind=kind,
            summary=summary,
            query=query,
            body=body,
            internal=internal,
            deprecated=deprecated,
            effect=effect,
            examples=tuple(examples),
        ),
    )
    return_specs = _normalize_returns(returns)

    def decorator(view: Callable[..., Any]) -> Callable[..., Any]:
        if context is not None and "ctx" not in inspect.signature(view).parameters:
            raise TypeError(f"{view.__qualname__} must declare a 'ctx' parameter")

        decorated = view
        if kind is Kind.LIST:
            decorated = paginated(op)(decorated)
        if query is not None or body is not None:
            decorated = accepts(query=query, body=body)(decorated)
        for return_spec in reversed(return_specs):
            decorated = _apply_returns(*return_spec)(decorated)
        decorated = subject_router.guard(spec)(decorated)

        decorated.__handler__ = view  # type: ignore[attr-defined]
        decorated.__spec__ = spec  # type: ignore[attr-defined]
        return decorated

    return decorator