refactor: using initializer to replace hardcoded dify cli initialization

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
 

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",
 ]

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

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)

api/core/virtual_environment/__base/initializer.py

@@ -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.
-        """
-        ...

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()],
             )