diff --git a/api/core/app/layers/sandbox_layer.py b/api/core/app/layers/sandbox_layer.py index 93d8a3015f..a1db5b45f4 100644 --- a/api/core/app/layers/sandbox_layer.py +++ b/api/core/app/layers/sandbox_layer.py @@ -1,9 +1,7 @@ import logging from collections.abc import Mapping -from io import BytesIO from typing import Any -from core.sandbox import DIFY_CLI_PATH, DifyCliLocator from core.sandbox.manager import SandboxManager from core.virtual_environment.__base.virtual_environment import VirtualEnvironment from core.workflow.graph_engine.layers.base import GraphEngineLayer @@ -48,7 +46,6 @@ class SandboxLayer(GraphEngineLayer): self._workflow_execution_id = self._get_workflow_execution_id() try: - sandbox: VirtualEnvironment from services.sandbox.sandbox_provider_service import SandboxProviderService logger.info("Initializing sandbox for tenant_id=%s", self._tenant_id) @@ -64,29 +61,10 @@ class SandboxLayer(GraphEngineLayer): sandbox.metadata.id, sandbox.metadata.arch, ) - self._upload_cli(sandbox) except Exception as e: logger.exception("Failed to initialize sandbox") raise SandboxInitializationError(f"Failed to initialize sandbox: {e}") from e - def _upload_cli(self, sandbox: VirtualEnvironment) -> None: - locator = DifyCliLocator() - binary = locator.resolve(sandbox.metadata.os, sandbox.metadata.arch) - - sandbox.upload_file(DIFY_CLI_PATH, BytesIO(binary.path.read_bytes())) - - connection_handle = sandbox.establish_connection() - try: - future = sandbox.run_command(connection_handle, ["chmod", "+x", DIFY_CLI_PATH]) - result = future.result(timeout=10) - if result.exit_code not in (0, None): - stderr = result.stderr.decode("utf-8", errors="replace") if result.stderr else "" - raise RuntimeError(f"Failed to mark dify CLI as executable: {stderr}") - - logger.info("Dify CLI uploaded to sandbox, path=%s", DIFY_CLI_PATH) - finally: - sandbox.release_connection(connection_handle) - def on_event(self, event: GraphEngineEvent) -> None: pass diff --git a/api/core/sandbox/__init__.py b/api/core/sandbox/__init__.py index 77dad10541..0382c89d2c 100644 --- a/api/core/sandbox/__init__.py +++ b/api/core/sandbox/__init__.py @@ -11,6 +11,7 @@ from core.sandbox.dify_cli import ( DifyCliLocator, DifyCliToolConfig, ) +from core.sandbox.initializer import DifyCliInitializer, SandboxInitializer from core.sandbox.session import SandboxSession __all__ = [ @@ -21,7 +22,9 @@ __all__ = [ "DifyCliBinary", "DifyCliConfig", "DifyCliEnvConfig", + "DifyCliInitializer", "DifyCliLocator", "DifyCliToolConfig", + "SandboxInitializer", "SandboxSession", ] diff --git a/api/core/sandbox/factory.py b/api/core/sandbox/factory.py index 8b31a4f1e9..fbd159da84 100644 --- a/api/core/sandbox/factory.py +++ b/api/core/sandbox/factory.py @@ -1,37 +1,20 @@ -""" -VM factory for creating VirtualEnvironment instances. - -Example: - vm = VMFactory.create( - tenant_id="tenant-uuid", - vm_type=VMType.DOCKER, - options={"docker_image": "python:3.11-slim"}, - environments={"PATH": "/usr/local/bin"}, - ) -""" - -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from enum import StrEnum -from typing import Any +from typing import TYPE_CHECKING, Any from core.virtual_environment.__base.virtual_environment import VirtualEnvironment +if TYPE_CHECKING: + from core.sandbox.initializer import SandboxInitializer + class VMType(StrEnum): - """Supported VM types.""" - DOCKER = "docker" E2B = "e2b" LOCAL = "local" class VMFactory: - """ - Factory for creating VirtualEnvironment (VM) instances. - - Uses lazy imports to avoid loading unused providers. - """ - @classmethod def create( cls, @@ -40,32 +23,22 @@ class VMFactory: options: Mapping[str, Any] | None = None, environments: Mapping[str, str] | None = None, user_id: str | None = None, + initializers: Sequence["SandboxInitializer"] | None = None, ) -> VirtualEnvironment: - """ - Create a VirtualEnvironment instance based on the specified type. - - Args: - tenant_id: Tenant ID associated with the VM (required) - vm_type: Type of VM to create - options: VM-specific configuration options - environments: Environment variables to set in the VM - user_id: User ID associated with the VM (optional) - - Returns: - Configured VirtualEnvironment instance - - Raises: - ValueError: If VM type is not supported - """ options = options or {} environments = environments or {} vm_class = cls._get_vm_class(vm_type) - return vm_class(tenant_id=tenant_id, options=options, environments=environments, user_id=user_id) + vm = vm_class(tenant_id=tenant_id, options=options, environments=environments, user_id=user_id) + + if initializers: + for initializer in initializers: + initializer.initialize(vm) + + return vm @classmethod def _get_vm_class(cls, vm_type: VMType) -> type[VirtualEnvironment]: - """Get the VM class for the specified type (lazy import).""" match vm_type: case VMType.DOCKER: from core.virtual_environment.providers.docker_daemon_sandbox import DockerDaemonEnvironment diff --git a/api/core/sandbox/initializer.py b/api/core/sandbox/initializer.py new file mode 100644 index 0000000000..05b9b845b3 --- /dev/null +++ b/api/core/sandbox/initializer.py @@ -0,0 +1,40 @@ +import logging +from abc import ABC, abstractmethod +from io import BytesIO +from pathlib import Path +from typing import TYPE_CHECKING + +from core.sandbox.constants import DIFY_CLI_PATH +from core.sandbox.dify_cli import DifyCliLocator + +if TYPE_CHECKING: + from core.virtual_environment.__base.virtual_environment import VirtualEnvironment + +logger = logging.getLogger(__name__) + + +class SandboxInitializer(ABC): + @abstractmethod + def initialize(self, env: "VirtualEnvironment") -> None: + ... + + +class DifyCliInitializer(SandboxInitializer): + def __init__(self, cli_root: str | Path | None = None) -> None: + self._locator = DifyCliLocator(root=cli_root) + + def initialize(self, env: "VirtualEnvironment") -> None: + binary = self._locator.resolve(env.metadata.os, env.metadata.arch) + env.upload_file(DIFY_CLI_PATH, BytesIO(binary.path.read_bytes())) + + connection_handle = env.establish_connection() + try: + future = env.run_command(connection_handle, ["chmod", "+x", DIFY_CLI_PATH]) + result = future.result(timeout=10) + if result.exit_code not in (0, None): + stderr = result.stderr.decode("utf-8", errors="replace") if result.stderr else "" + raise RuntimeError(f"Failed to mark dify CLI as executable: {stderr}") + + logger.info("Dify CLI uploaded to sandbox, path=%s", DIFY_CLI_PATH) + finally: + env.release_connection(connection_handle) diff --git a/api/core/virtual_environment/__base/initializer.py b/api/core/virtual_environment/__base/initializer.py deleted file mode 100644 index dae03ae815..0000000000 --- a/api/core/virtual_environment/__base/initializer.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Sandbox initializer protocol for post-construction initialization. - -This module defines the interface for initializers that can perform -setup tasks on newly created VirtualEnvironment instances. -""" - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from core.virtual_environment.__base.virtual_environment import VirtualEnvironment - - -class SandboxInitializer(ABC): - """ - Abstract base class for sandbox post-construction initialization. - - Initializers are called by VMFactory after a VirtualEnvironment is created. - They allow decoupling of environment creation from environment setup tasks - like uploading binaries, configuring tools, or setting up directories. - - Example: - class MyInitializer(SandboxInitializer): - def initialize(self, env: VirtualEnvironment) -> None: - env.upload_file("/path/to/file", BytesIO(b"content")) - """ - - @abstractmethod - def initialize(self, env: "VirtualEnvironment") -> None: - """ - Perform initialization on a newly created sandbox. - - Called by VMFactory after VirtualEnvironment._construct_environment(). - Implementations should be idempotent where possible. - - Args: - env: The virtual environment to initialize. - - Raises: - Exception: If initialization fails. The caller is responsible - for handling cleanup. - """ - ... diff --git a/api/services/sandbox/sandbox_provider_service.py b/api/services/sandbox/sandbox_provider_service.py index 76b75cd3e4..21df858b04 100644 --- a/api/services/sandbox/sandbox_provider_service.py +++ b/api/services/sandbox/sandbox_provider_service.py @@ -20,6 +20,7 @@ from configs import dify_config from constants import HIDDEN_VALUE from core.entities.provider_entities import BasicProviderConfig from core.sandbox.factory import VMFactory, VMType +from core.sandbox.initializer import DifyCliInitializer from core.tools.utils.system_encryption import ( decrypt_system_params, ) @@ -339,4 +340,5 @@ class SandboxProviderService: vm_type=VMType(provider_type), options=dict(config), environments=environments or {}, + initializers=[DifyCliInitializer()], )