dify/api/core/sandbox/builder.py

from __future__ import annotations

import logging
import threading
from collections.abc import Mapping, Sequence
from contextlib import nullcontext
from typing import TYPE_CHECKING, Any, cast

from flask import Flask, current_app, has_app_context

from core.entities.provider_entities import BasicProviderConfig
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment

from .entities.sandbox_type import SandboxType
from .initializer import AsyncSandboxInitializer, SandboxInitializeContext, SandboxInitializer, SyncSandboxInitializer
from .sandbox import Sandbox

if TYPE_CHECKING:
    from .storage.sandbox_storage import SandboxStorage

logger = logging.getLogger(__name__)


def _get_sandbox_class(sandbox_type: SandboxType) -> type[VirtualEnvironment]:
    match sandbox_type:
        case SandboxType.DOCKER:
            from core.virtual_environment.providers.docker_daemon_sandbox import DockerDaemonEnvironment

            return DockerDaemonEnvironment
        case SandboxType.E2B:
            from core.virtual_environment.providers.e2b_sandbox import E2BEnvironment

            return E2BEnvironment
        case SandboxType.LOCAL:
            from core.virtual_environment.providers.local_without_isolation import LocalVirtualEnvironment

            return LocalVirtualEnvironment
        case SandboxType.SSH:
            from core.virtual_environment.providers.ssh_sandbox import SSHSandboxEnvironment

            return SSHSandboxEnvironment
        case SandboxType.AWS_CODE_INTERPRETER:
            from core.virtual_environment.providers.aws_code_interpreter_sandbox import (
                AWSCodeInterpreterEnvironment,
            )

            return AWSCodeInterpreterEnvironment
        case _:
            raise ValueError(f"Unsupported sandbox type: {sandbox_type}")


class SandboxBuilder:
    _tenant_id: str
    _sandbox_type: SandboxType
    _user_id: str | None
    _app_id: str | None
    _options: dict[str, Any]
    _environments: dict[str, str]
    _initializers: list[SandboxInitializer]
    _storage: SandboxStorage | None
    _assets_id: str | None

    def __init__(self, tenant_id: str, sandbox_type: SandboxType) -> None:
        self._tenant_id = tenant_id
        self._sandbox_type = sandbox_type
        self._user_id = None
        self._app_id = None
        self._options = {}
        self._environments = {}
        self._initializers = []
        self._storage = None
        self._assets_id = None

    def user(self, user_id: str) -> SandboxBuilder:
        self._user_id = user_id
        return self

    def app(self, app_id: str) -> SandboxBuilder:
        self._app_id = app_id
        return self

    def options(self, options: Mapping[str, Any]) -> SandboxBuilder:
        self._options = dict(options)
        return self

    def environments(self, environments: Mapping[str, str]) -> SandboxBuilder:
        self._environments = dict(environments)
        return self

    def initializer(self, initializer: SandboxInitializer) -> SandboxBuilder:
        self._initializers.append(initializer)
        return self

    def initializers(self, initializers: Sequence[SandboxInitializer]) -> SandboxBuilder:
        self._initializers.extend(initializers)
        return self

    def storage(self, storage: SandboxStorage, assets_id: str) -> SandboxBuilder:
        self._storage = storage
        self._assets_id = assets_id
        return self

    def build(self) -> Sandbox:
        """Create a sandbox and start background initialization.

        The builder is responsible for cleaning up any VM or sandbox that was
        successfully created if a later setup step fails.
        """
        if self._storage is None:
            raise ValueError("storage is required, call .storage() before .build()")
        if self._assets_id is None:
            raise ValueError("assets_id is required, call .storage() before .build()")
        if self._user_id is None:
            raise ValueError("user_id is required, call .user() before .build()")
        if self._app_id is None:
            raise ValueError("app_id is required, call .app() before .build()")

        ctx = SandboxInitializeContext(
            tenant_id=self._tenant_id,
            app_id=self._app_id,
            assets_id=self._assets_id,
            user_id=self._user_id,
        )
        vm: VirtualEnvironment | None = None
        sandbox: Sandbox | None = None
        try:
            vm_class = _get_sandbox_class(self._sandbox_type)
            vm = vm_class(
                tenant_id=self._tenant_id,
                options=self._options,
                environments=self._environments,
                user_id=self._user_id,
            )
            vm.open_enviroment()
            sandbox = Sandbox(
                vm=vm,
                storage=self._storage,
                tenant_id=self._tenant_id,
                user_id=self._user_id,
                app_id=self._app_id,
                assets_id=self._assets_id,
            )

            for init in self._initializers:
                if isinstance(init, SyncSandboxInitializer):
                    init.initialize(sandbox, ctx)
        except Exception as exc:
            logger.exception(
                "Failed to initialize sandbox synchronously: tenant_id=%s, app_id=%s", self._tenant_id, self._app_id
            )
            if sandbox is not None:
                sandbox.release()
            elif vm is not None:
                try:
                    vm.release_environment()
                except Exception:
                    logger.exception("Failed to release sandbox VM during builder cleanup")
            raise RuntimeError("Sandbox initialization failed") from exc

        # Run sandbox setup asynchronously so workflow execution can proceed.
        # Capture the Flask app before starting the thread for database access.
        flask_app: Flask | None = cast(Any, current_app)._get_current_object() if has_app_context() else None

        def initialize() -> None:
            try:
                app_context = flask_app.app_context() if flask_app is not None else nullcontext()
                with app_context:
                    for init in self._initializers:
                        if not isinstance(init, AsyncSandboxInitializer):
                            continue

                        if sandbox.is_cancelled():
                            return
                        init.initialize(sandbox, ctx)

                    if sandbox.is_cancelled():
                        return
                    # Storage mount is part of readiness. If restore/mount fails,
                    # the sandbox must surface initialization failure instead of
                    # becoming "ready" with missing files.
                    if not sandbox.mount():
                        raise RuntimeError("Sandbox storage mount failed")
                    sandbox.mark_ready()
            except Exception as exc:
                try:
                    logger.exception(
                        "Failed to initialize sandbox: tenant_id=%s, app_id=%s", self._tenant_id, self._app_id
                    )
                    sandbox.release()
                    sandbox.mark_failed(exc)
                except Exception:
                    logger.exception(
                        "Failed to mark sandbox initialization failure: tenant_id=%s, app_id=%s",
                        self._tenant_id,
                        self._app_id,
                    )

        # Background init completes or signals failure via sandbox state.
        try:
            threading.Thread(target=initialize, daemon=True).start()
        except Exception:
            logger.exception(
                "Failed to start sandbox initialization thread: tenant_id=%s, app_id=%s",
                self._tenant_id,
                self._app_id,
            )
            sandbox.release()
            raise RuntimeError("Sandbox initialization failed")
        return sandbox

    @staticmethod
    def validate(vm_type: SandboxType, options: Mapping[str, Any]) -> None:
        vm_class = _get_sandbox_class(vm_type)
        vm_class.validate(options)

    @classmethod
    def draft_id(cls, user_id: str) -> str:
        return user_id


class VMConfig:
    @staticmethod
    def get_schema(vm_type: SandboxType) -> list[BasicProviderConfig]:
        return _get_sandbox_class(vm_type).get_config_schema()