dify/api/controllers/console/dify_builder.py

"""Console HTTP surface for the Dify Builder (P3c) — JSON routes.

Flask-RESTX resources under /console/api/dify-builder/*, matching the
retired Go enterprise transport byte-for-byte so cutover is a frontend
base-path swap. Auth + CSRF come from the console guard stack (CSRF is enforced
inside login_required for every method); RBAC/authz is enforced here at the
boundary — the Celery advance task trusts the Actor. The route LOGIC lives in
undecorated module functions (_create/_view/_action/_message/_stream) so it is unit
testable without the Flask request stack; the Resource methods are thin
auth-wrappers.

Decorator injection contract (verified against controllers/console/wraps.py
and the multi-decorator precedent in controllers/console/app/workflow_comment.py,
e.g. ``def post(self, current_tenant_id, current_user, app_model, comment_id)``):
``@with_current_user`` and ``@with_current_tenant_id`` do NOT inject their
values as keyword arguments. Each calls its wrapped view as
``view(self, injected_value, *args, **kwargs)`` — i.e. it PREPENDS the value
as a positional argument immediately after ``self``. Stacking both means the
decorator closer to the function (here, ``@with_current_tenant_id``, directly
above ``@dify_builder_required``) contributes the positional argument
closest to ``self``; the one above it (``@with_current_user``) contributes
the next one. URL path parameters (e.g. ``session_id``) are untouched by
either decorator and keep arriving as keyword arguments from Flask-RESTX's
routing. ``dify_builder_required`` sits directly below
``@with_current_tenant_id`` in the stack, so it receives ``current_tenant_id``
positionally (as the argument right after ``self``), not as a keyword-only
argument.
"""

import functools
import logging
from collections.abc import Callable

from flask import Response, request
from flask_restx import Resource

from controllers.console.wraps import (
    account_initialization_required,
    setup_required,
    with_current_tenant_id,
    with_current_user,
)
from core.dify_builder.models import Action, Actor
from libs.login import login_required
from services.dify_builder import progress_bus
from services.dify_builder.service import resolve_action_kind
from services.dify_builder.wiring import (
    build_service,
    dify_builder_error_response,
    session_view_to_dict,
    stream_frames,
)
from services.feature_service import FeatureService

from . import console_ns

logger = logging.getLogger(__name__)


def dify_builder_required(func: Callable) -> Callable:
    @functools.wraps(func)
    def wrapper(self, current_tenant_id: str, *args, **kwargs):
        if not FeatureService.get_features(current_tenant_id).dify_builder_enabled:
            return {"code": "feature_unavailable"}, 403
        return func(self, current_tenant_id, *args, **kwargs)

    return wrapper


def _actor(current_user, current_tenant_id: str) -> Actor:
    return Actor(account_id=current_user.id, tenant_id=current_tenant_id)


def _respond(fn: Callable[[], object]) -> tuple[dict, int]:
    """Run a usecase call, serialize its SessionView, map known dify_builder errors
    to (body, status). Unknown exceptions propagate to Flask's error handler."""
    try:
        return session_view_to_dict(fn()), 200
    except Exception as exc:  # re-raised below if not a known dify_builder error
        mapped = dify_builder_error_response(exc)
        if mapped is None:
            raise
        return mapped


_SSE_HEADERS = {"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}


def _stream_response(make_generator: Callable[[], object]) -> Response | tuple[dict, int]:
    """Call make_generator() (which runs the service's eager validate+subscribe+dispatch
    and returns an SSE generator); map pre-flight errors to an HTTP tuple; else return
    the stream. Unknown exceptions propagate to Flask's error handler."""
    try:
        generator = make_generator()
    except Exception as exc:  # re-raised below if not a known dify_builder error
        mapped = dify_builder_error_response(exc)
        if mapped is None:
            raise
        return mapped
    return Response(generator, mimetype="text/event-stream", headers=_SSE_HEADERS)


def _is_streaming(body: object) -> bool:
    return isinstance(body, dict) and body.get("response_mode") == "streaming"


def _create(body, actor: Actor) -> Response | tuple[dict, int]:
    if not isinstance(body, dict):
        return {"code": "bad_request"}, 400
    if _is_streaming(body):
        if body.get("scenario") == "build":
            return _stream_response(
                lambda: build_service().create_build_session_stream(
                    body.get("app_id", ""), actor, body.get("goal_text", ""), body.get("model_config")
                )
            )
        if body.get("scenario") == "edit":
            return _stream_response(
                lambda: build_service().create_edit_session_stream(
                    body.get("app_id", ""), actor, body.get("model_config")
                )
            )
        return _stream_response(
            lambda: build_service().create_fix_session_stream(
                body.get("app_id", ""),
                actor,
                body.get("failed_run_id") or None,
                body.get("checklist_errors") or None,
                body.get("model_config"),
            )
        )
    if body.get("scenario") == "build":
        result, status = _respond(
            lambda: build_service().create_build_session(
                app_id=body.get("app_id", ""),
                actor=actor,
                goal_text=body.get("goal_text", ""),
                model_config=body.get("model_config"),
            )
        )
    elif body.get("scenario") == "edit":
        result, status = _respond(
            lambda: build_service().create_edit_session(
                app_id=body.get("app_id", ""),
                actor=actor,
                model_config=body.get("model_config"),
            )
        )
    else:
        result, status = _respond(
            lambda: build_service().create_fix_session(
                app_id=body.get("app_id", ""),
                actor=actor,
                failed_run_id=body.get("failed_run_id") or None,
                checklist_errors=body.get("checklist_errors") or None,
                model_config=body.get("model_config"),
            )
        )
    # create returns 201 on success; map-through keeps error statuses.
    return result, (201 if status == 200 else status)


def _view(session_id: str, actor: Actor) -> tuple[dict, int]:
    return _respond(lambda: build_service().get_session_view(session_id, actor))


def _action(session_id: str, body, actor: Actor) -> Response | tuple[dict, int]:
    if not isinstance(body, dict):
        return {"code": "bad_request"}, 400
    raw = body.get("action_id") or body.get("kind", "")
    action = Action(
        kind=resolve_action_kind(raw),
        payload=body.get("payload") or {},
        base_version=int(body.get("base_version", 0)),
    )
    if _is_streaming(body):
        return _stream_response(lambda: build_service().submit_action_stream(session_id, actor, action))
    return _respond(lambda: build_service().submit_action(session_id, actor, action))


def _message(session_id: str, body, actor: Actor) -> Response | tuple[dict, int]:
    if not isinstance(body, dict):
        return {"code": "bad_request"}, 400
    if _is_streaming(body):
        return _stream_response(
            lambda: build_service().submit_message_stream(
                session_id, actor, body.get("text", ""), int(body.get("base_version", 0))
            )
        )
    return _respond(
        lambda: build_service().submit_message(
            session_id, actor, body.get("text", ""), int(body.get("base_version", 0))
        )
    )


def _stream(session_id: str, actor: Actor) -> Response | tuple[dict, int]:
    # Subscribe before reading the snapshot. This closes the race where an
    # advance could commit after the snapshot read but before subscription,
    # leaving the client with an old view and no state event to refresh it.
    subscription = progress_bus.subscribe(session_id)
    try:
        view = build_service().get_session_view(session_id, actor)
    except Exception as exc:  # re-raised below if not a known dify_builder error
        subscription.close()
        mapped = dify_builder_error_response(exc)
        if mapped is None:
            raise
        return mapped

    return Response(
        stream_frames(session_view_to_dict(view), subscription),
        mimetype="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
    )


def _ping(body, actor: Actor) -> tuple[dict, int]:
    model_config = body.get("model_config") if isinstance(body, dict) else None
    try:
        from services.dify_builder.agent.ping import ping_model

        return ping_model(actor.tenant_id, model_config), 200
    except Exception as exc:  # a health check returns ok:false rather than 500
        logger.warning("dify_builder agent ping failed: %s", exc)
        return {"ok": False, "error": str(exc)}, 200


@console_ns.route("/dify-builder/sessions")
class DifyBuilderSessionsApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @dify_builder_required
    def post(self, current_tenant_id, current_user):
        return _create(request.get_json(silent=True), _actor(current_user, current_tenant_id))


@console_ns.route("/dify-builder/sessions/<string:session_id>")
class DifyBuilderSessionApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @dify_builder_required
    def get(self, current_tenant_id, current_user, session_id):
        return _view(session_id, _actor(current_user, current_tenant_id))


@console_ns.route("/dify-builder/sessions/<string:session_id>/actions")
class DifyBuilderActionsApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @dify_builder_required
    def post(self, current_tenant_id, current_user, session_id):
        return _action(session_id, request.get_json(silent=True), _actor(current_user, current_tenant_id))


@console_ns.route("/dify-builder/sessions/<string:session_id>/messages")
class DifyBuilderMessagesApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @dify_builder_required
    def post(self, current_tenant_id, current_user, session_id):
        return _message(session_id, request.get_json(silent=True), _actor(current_user, current_tenant_id))


@console_ns.route("/dify-builder/sessions/<string:session_id>/stream")
class DifyBuilderStreamApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @dify_builder_required
    def get(self, current_tenant_id, current_user, session_id):
        return _stream(session_id, _actor(current_user, current_tenant_id))


@console_ns.route("/dify-builder/agent/ping")
class DifyBuilderAgentPingApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @with_current_user
    @with_current_tenant_id
    @dify_builder_required
    def post(self, current_tenant_id, current_user):
        return _ping(request.get_json(silent=True), _actor(current_user, current_tenant_id))