refactor(sandbox): reorganize sandbox-related imports and rename SandboxFactory to VMFactory for clarity

This commit is contained in:
Harry 2026-01-12 02:07:26 +08:00
parent 9dd0361d0e
commit 51ac23c9f1
12 changed files with 65 additions and 78 deletions

View File

@ -4,8 +4,8 @@ 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.virtual_environment.sandbox_manager import SandboxManager
from core.workflow.graph_engine.layers.base import GraphEngineLayer
from core.workflow.graph_events.base import GraphEngineEvent

View File

@ -3,17 +3,15 @@ from __future__ import annotations
import json
import logging
from io import BytesIO
from typing import TYPE_CHECKING
from types import TracebackType
from core.sandbox.constants import DIFY_CLI_CONFIG_PATH, DIFY_CLI_PATH
from core.sandbox.dify_cli import DifyCliConfig
if TYPE_CHECKING:
from types import TracebackType
from core.tools.__base.tool import Tool
from core.tools.builtin_tool.providers.sandbox.bash_tool import SandboxBashTool
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.sandbox.manager import SandboxManager
from core.session.inner_api import InnerApiSessionManager
from core.tools.__base.tool import Tool
from core.tools.builtin_tool.providers.sandbox.bash_tool import SandboxBashTool
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
logger = logging.getLogger(__name__)
@ -37,10 +35,6 @@ class SandboxSession:
self._session_id: str | None = None
def __enter__(self) -> SandboxSession:
from core.session.inner_api import InnerApiSessionManager
from core.tools.builtin_tool.providers.sandbox.bash_tool import SandboxBashTool
from core.virtual_environment.sandbox_manager import SandboxManager
sandbox = SandboxManager.get(self._workflow_execution_id)
if sandbox is None:
raise RuntimeError(f"Sandbox not found for workflow_execution_id={self._workflow_execution_id}")

View File

@ -1,10 +1,10 @@
"""
Sandbox factory for creating VirtualEnvironment instances.
VM factory for creating VirtualEnvironment instances.
Example:
sandbox = SandboxFactory.create(
vm = VMFactory.create(
tenant_id="tenant-uuid",
sandbox_type=SandboxType.DOCKER,
vm_type=VMType.DOCKER,
options={"docker_image": "python:3.11-slim"},
environments={"PATH": "/usr/local/bin"},
)
@ -17,17 +17,17 @@ from typing import Any
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
class SandboxType(StrEnum):
"""Supported sandbox types."""
class VMType(StrEnum):
"""Supported VM types."""
DOCKER = "docker"
E2B = "e2b"
LOCAL = "local"
class SandboxFactory:
class VMFactory:
"""
Factory for creating VirtualEnvironment (sandbox) instances.
Factory for creating VirtualEnvironment (VM) instances.
Uses lazy imports to avoid loading unused providers.
"""
@ -36,7 +36,7 @@ class SandboxFactory:
def create(
cls,
tenant_id: str,
sandbox_type: SandboxType,
vm_type: VMType,
options: Mapping[str, Any] | None = None,
environments: Mapping[str, str] | None = None,
user_id: str | None = None,
@ -45,44 +45,44 @@ class SandboxFactory:
Create a VirtualEnvironment instance based on the specified type.
Args:
tenant_id: Tenant ID associated with the sandbox (required)
sandbox_type: Type of sandbox to create
options: Sandbox-specific configuration options
environments: Environment variables to set in the sandbox
user_id: User ID associated with the sandbox (optional)
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 sandbox type is not supported
ValueError: If VM type is not supported
"""
options = options or {}
environments = environments or {}
sandbox_class = cls._get_sandbox_class(sandbox_type)
return sandbox_class(tenant_id=tenant_id, options=options, environments=environments, user_id=user_id)
vm_class = cls._get_vm_class(vm_type)
return vm_class(tenant_id=tenant_id, options=options, environments=environments, user_id=user_id)
@classmethod
def _get_sandbox_class(cls, sandbox_type: SandboxType) -> type[VirtualEnvironment]:
"""Get the sandbox class for the specified type (lazy import)."""
match sandbox_type:
case SandboxType.DOCKER:
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
return DockerDaemonEnvironment
case SandboxType.E2B:
case VMType.E2B:
from core.virtual_environment.providers.e2b_sandbox import E2BEnvironment
return E2BEnvironment
case SandboxType.LOCAL:
case VMType.LOCAL:
from core.virtual_environment.providers.local_without_isolation import LocalVirtualEnvironment
return LocalVirtualEnvironment
case _:
raise ValueError(f"Unsupported sandbox type: {sandbox_type}")
raise ValueError(f"Unsupported VM type: {vm_type}")
@classmethod
def validate(cls, sandbox_type: SandboxType, options: Mapping[str, Any]) -> None:
sandbox_class = cls._get_sandbox_class(sandbox_type)
sandbox_class.validate(options)
def validate(cls, vm_type: VMType, options: Mapping[str, Any]) -> None:
vm_class = cls._get_vm_class(vm_type)
vm_class.validate(options)

View File

@ -4,9 +4,9 @@ import shlex
from collections.abc import Mapping, Sequence
from typing import Any
from core.sandbox.manager import SandboxManager
from core.virtual_environment.__base.command_future import CommandCancelledError, CommandTimeoutError
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.virtual_environment.sandbox_manager import SandboxManager
from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
from core.workflow.node_events import NodeRunResult
from core.workflow.nodes.base import variable_template_parser

View File

@ -51,6 +51,7 @@ from core.prompt.entities.advanced_prompt_entities import CompletionModelPromptT
from core.prompt.utils.prompt_message_util import PromptMessageUtil
from core.rag.entities.citation_metadata import RetrievalSourceMetadata
from core.sandbox import SandboxSession
from core.sandbox.manager import SandboxManager
from core.tools.__base.tool import Tool
from core.tools.signature import sign_upload_file
from core.tools.tool_manager import ToolManager
@ -62,7 +63,6 @@ from core.variables import (
ObjectSegment,
StringSegment,
)
from core.virtual_environment.sandbox_manager import SandboxManager
from core.workflow.constants import SYSTEM_VARIABLE_NODE_ID
from core.workflow.entities import GraphInitParams, ToolCall, ToolResult, ToolResultStatus
from core.workflow.entities.tool_entities import ToolCallResult
@ -1588,8 +1588,6 @@ class LLMNode(Node[LLMNodeData]):
stop: Sequence[str] | None,
variable_pool: VariablePool,
) -> Generator[NodeEventBase, None, LLMGenerationData]:
from core.agent.entities import AgentEntity
workflow_execution_id = variable_pool.system_variables.workflow_execution_id
if not workflow_execution_id:
raise LLMNodeError("workflow_execution_id is required for sandbox runtime mode")
@ -1613,7 +1611,6 @@ class LLMNode(Node[LLMNodeData]):
files=prompt_files,
max_iterations=self._node_data.max_iterations or 10,
context=ExecutionContext(user_id=self.user_id, app_id=self.app_id, tenant_id=self.tenant_id),
agent_strategy=AgentEntity.Strategy.CHAIN_OF_THOUGHT,
)
outputs = strategy.run(

View File

@ -23,7 +23,7 @@ from core.tools.utils.system_encryption import (
decrypt_system_params,
)
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.virtual_environment.factory import SandboxFactory, SandboxType
from core.virtual_environment.factory import VMFactory, VMType
from extensions.ext_database import db
from models.sandbox import SandboxProvider, SandboxProviderSystemConfig
from services.sandbox.encryption import create_sandbox_config_encrypter, masked_config
@ -172,7 +172,7 @@ class SandboxProviderService:
if model_class:
model_class.model_validate(config)
SandboxFactory.validate(SandboxType(provider_type), config)
VMFactory.validate(VMType(provider_type), config)
@classmethod
def save_config(
@ -334,9 +334,9 @@ class SandboxProviderService:
if not config or not provider_type:
raise ValueError(f"No active sandbox provider for tenant {tenant_id} or system default")
return SandboxFactory.create(
return VMFactory.create(
tenant_id=tenant_id,
sandbox_type=SandboxType(provider_type),
vm_type=VMType(provider_type),
options=dict(config),
environments=environments or {},
)

View File

@ -14,6 +14,7 @@ from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfig
from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager
from core.file import File
from core.repositories import DifyCoreRepositoryFactory
from core.sandbox.manager import SandboxManager
from core.variables import Variable
from core.variables.variables import VariableUnion
from core.workflow.entities import WorkflowNodeExecution
@ -703,7 +704,6 @@ class WorkflowService:
if draft_workflow.get_feature(WorkflowFeatures.SANDBOX).enabled:
sandbox = SandboxProviderService.create_sandbox(tenant_id=draft_workflow.tenant_id)
single_step_execution_id = f"single-step-{uuid.uuid4()}"
from core.virtual_environment.sandbox_manager import SandboxManager
SandboxManager.register(single_step_execution_id, sandbox)
variable_pool.system_variables.workflow_execution_id = single_step_execution_id
@ -727,8 +727,6 @@ class WorkflowService:
)
finally:
if single_step_execution_id:
from core.virtual_environment.sandbox_manager import SandboxManager
sandbox = SandboxManager.unregister(single_step_execution_id)
if sandbox:
try:

View File

@ -3,9 +3,9 @@ from unittest.mock import MagicMock, patch
import pytest
from core.app.layers.sandbox_layer import SandboxInitializationError, SandboxLayer
from core.sandbox.manager import SandboxManager
from core.virtual_environment.__base.entities import Arch
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.virtual_environment.sandbox_manager import SandboxManager
from core.workflow.graph_engine.layers.base import GraphEngineLayerNotInitializedError
from core.workflow.graph_events.graph import (
GraphRunFailedEvent,

View File

@ -11,7 +11,7 @@ from unittest.mock import MagicMock, patch
import pytest
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.virtual_environment.factory import SandboxFactory, SandboxType
from core.virtual_environment.factory import VMFactory, VMType
class TestSandboxType:
@ -19,15 +19,15 @@ class TestSandboxType:
def test_sandbox_type_values(self):
"""Test that SandboxType enum has expected values."""
assert SandboxType.DOCKER == "docker"
assert SandboxType.E2B == "e2b"
assert SandboxType.LOCAL == "local"
assert VMType.DOCKER == "docker"
assert VMType.E2B == "e2b"
assert VMType.LOCAL == "local"
def test_sandbox_type_is_string_enum(self):
"""Test that SandboxType values are strings."""
assert isinstance(SandboxType.DOCKER.value, str)
assert isinstance(SandboxType.E2B.value, str)
assert isinstance(SandboxType.LOCAL.value, str)
assert isinstance(VMType.DOCKER.value, str)
assert isinstance(VMType.E2B.value, str)
assert isinstance(VMType.LOCAL.value, str)
class TestSandboxFactory:
@ -38,10 +38,10 @@ class TestSandboxFactory:
mock_sandbox_instance = MagicMock(spec=VirtualEnvironment)
mock_sandbox_class = MagicMock(return_value=mock_sandbox_instance)
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
result = SandboxFactory.create(
with patch.object(VMFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
result = VMFactory.create(
tenant_id="test-tenant",
sandbox_type=SandboxType.DOCKER,
vm_type=VMType.DOCKER,
options={"docker_image": "python:3.11-slim"},
environments={"PYTHONUNBUFFERED": "1"},
)
@ -59,10 +59,8 @@ class TestSandboxFactory:
mock_sandbox_instance = MagicMock(spec=VirtualEnvironment)
mock_sandbox_class = MagicMock(return_value=mock_sandbox_instance)
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
SandboxFactory.create(
tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER, options=None, environments=None
)
with patch.object(VMFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
VMFactory.create(tenant_id="test-tenant", vm_type=VMType.DOCKER, options=None, environments=None)
mock_sandbox_class.assert_called_once_with(
tenant_id="test-tenant", options={}, environments={}, user_id=None
@ -73,8 +71,8 @@ class TestSandboxFactory:
mock_sandbox_instance = MagicMock(spec=VirtualEnvironment)
mock_sandbox_class = MagicMock(return_value=mock_sandbox_instance)
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
result = SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER)
with patch.object(VMFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
result = VMFactory.create(tenant_id="test-tenant", vm_type=VMType.DOCKER)
mock_sandbox_class.assert_called_once_with(
tenant_id="test-tenant", options={}, environments={}, user_id=None
@ -90,7 +88,7 @@ class TestSandboxFactory:
"core.virtual_environment.providers.docker_daemon_sandbox.DockerDaemonEnvironment",
return_value=mock_instance,
) as mock_docker_class:
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER)
VMFactory.create(tenant_id="test-tenant", vm_type=VMType.DOCKER)
mock_docker_class.assert_called_once()
def test_get_sandbox_class_local_returns_correct_class(self):
@ -101,7 +99,7 @@ class TestSandboxFactory:
"core.virtual_environment.providers.local_without_isolation.LocalVirtualEnvironment",
return_value=mock_instance,
) as mock_local_class:
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.LOCAL)
VMFactory.create(tenant_id="test-tenant", vm_type=VMType.LOCAL)
mock_local_class.assert_called_once()
def test_get_sandbox_class_e2b_returns_correct_class(self):
@ -112,13 +110,13 @@ class TestSandboxFactory:
"core.virtual_environment.providers.e2b_sandbox.E2BEnvironment",
return_value=mock_instance,
) as mock_e2b_class:
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.E2B)
VMFactory.create(tenant_id="test-tenant", vm_type=VMType.E2B)
mock_e2b_class.assert_called_once()
def test_create_with_unsupported_type_raises_value_error(self):
"""Test that unsupported sandbox type raises ValueError."""
with pytest.raises(ValueError) as exc_info:
SandboxFactory.create(tenant_id="test-tenant", sandbox_type="unsupported_type") # type: ignore[arg-type]
VMFactory.create(tenant_id="test-tenant", vm_type="unsupported_type") # type: ignore[arg-type]
assert "Unsupported sandbox type: unsupported_type" in str(exc_info.value)
@ -127,9 +125,9 @@ class TestSandboxFactory:
mock_sandbox_class = MagicMock()
mock_sandbox_class.side_effect = Exception("Docker daemon not available")
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
with patch.object(VMFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
with pytest.raises(Exception) as exc_info:
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER)
VMFactory.create(tenant_id="test-tenant", vm_type=VMType.DOCKER)
assert "Docker daemon not available" in str(exc_info.value)
@ -139,9 +137,9 @@ class TestSandboxFactoryIntegration:
def test_create_local_sandbox_integration(self, tmp_path: Path):
"""Test creating a real local sandbox."""
sandbox = SandboxFactory.create(
sandbox = VMFactory.create(
tenant_id="test-tenant",
sandbox_type=SandboxType.LOCAL,
vm_type=VMType.LOCAL,
options={"base_working_path": str(tmp_path)},
environments={},
)

View File

@ -5,9 +5,9 @@ from typing import Any
import pytest
from core.sandbox.manager import SandboxManager
from core.virtual_environment.__base.entities import Arch, CommandStatus, ConnectionHandle, FileState, Metadata
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.virtual_environment.sandbox_manager import SandboxManager
class FakeVirtualEnvironment(VirtualEnvironment):

View File

@ -5,11 +5,11 @@ from typing import Any
import pytest
from core.sandbox.manager import SandboxManager
from core.virtual_environment.__base.entities import Arch, CommandStatus, ConnectionHandle, FileState, Metadata
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.virtual_environment.channel.queue_transport import QueueTransportReadCloser
from core.virtual_environment.channel.transport import NopTransportWriteCloser
from core.virtual_environment.sandbox_manager import SandboxManager
from core.workflow.entities import GraphInitParams
from core.workflow.enums import WorkflowNodeExecutionStatus
from core.workflow.nodes.command.node import CommandNode