From 51ac23c9f11761da8985c8bb1a0330f3e81d0094 Mon Sep 17 00:00:00 2001 From: Harry Date: Mon, 12 Jan 2026 02:07:26 +0800 Subject: [PATCH] refactor(sandbox): reorganize sandbox-related imports and rename SandboxFactory to VMFactory for clarity --- api/core/app/layers/sandbox_layer.py | 2 +- .../sandbox_manager.py => sandbox/manager.py} | 0 api/core/sandbox/session.py | 18 +++---- api/core/virtual_environment/factory.py | 52 +++++++++---------- api/core/workflow/nodes/command/node.py | 2 +- api/core/workflow/nodes/llm/node.py | 5 +- .../sandbox/sandbox_provider_service.py | 8 +-- api/services/workflow_service.py | 4 +- .../core/app/layers/test_sandbox_layer.py | 2 +- .../core/virtual_environment/test_factory.py | 46 ++++++++-------- .../test_sandbox_manager.py | 2 +- .../nodes/command/test_command_node.py | 2 +- 12 files changed, 65 insertions(+), 78 deletions(-) rename api/core/{virtual_environment/sandbox_manager.py => sandbox/manager.py} (100%) diff --git a/api/core/app/layers/sandbox_layer.py b/api/core/app/layers/sandbox_layer.py index e6e5e95565..93d8a3015f 100644 --- a/api/core/app/layers/sandbox_layer.py +++ b/api/core/app/layers/sandbox_layer.py @@ -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 diff --git a/api/core/virtual_environment/sandbox_manager.py b/api/core/sandbox/manager.py similarity index 100% rename from api/core/virtual_environment/sandbox_manager.py rename to api/core/sandbox/manager.py diff --git a/api/core/sandbox/session.py b/api/core/sandbox/session.py index d58b16bce4..15051ecf93 100644 --- a/api/core/sandbox/session.py +++ b/api/core/sandbox/session.py @@ -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}") diff --git a/api/core/virtual_environment/factory.py b/api/core/virtual_environment/factory.py index 4b442e4bc2..8b31a4f1e9 100644 --- a/api/core/virtual_environment/factory.py +++ b/api/core/virtual_environment/factory.py @@ -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) diff --git a/api/core/workflow/nodes/command/node.py b/api/core/workflow/nodes/command/node.py index be603cc6bd..0a76ac0399 100644 --- a/api/core/workflow/nodes/command/node.py +++ b/api/core/workflow/nodes/command/node.py @@ -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 diff --git a/api/core/workflow/nodes/llm/node.py b/api/core/workflow/nodes/llm/node.py index 2b46d7a4da..123da76b58 100644 --- a/api/core/workflow/nodes/llm/node.py +++ b/api/core/workflow/nodes/llm/node.py @@ -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( diff --git a/api/services/sandbox/sandbox_provider_service.py b/api/services/sandbox/sandbox_provider_service.py index baf4d36e8d..9b6b116ada 100644 --- a/api/services/sandbox/sandbox_provider_service.py +++ b/api/services/sandbox/sandbox_provider_service.py @@ -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 {}, ) diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index a7b2aad0a7..ea87bf9385 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -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: diff --git a/api/tests/unit_tests/core/app/layers/test_sandbox_layer.py b/api/tests/unit_tests/core/app/layers/test_sandbox_layer.py index 4358a8d4d6..b3c3822646 100644 --- a/api/tests/unit_tests/core/app/layers/test_sandbox_layer.py +++ b/api/tests/unit_tests/core/app/layers/test_sandbox_layer.py @@ -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, diff --git a/api/tests/unit_tests/core/virtual_environment/test_factory.py b/api/tests/unit_tests/core/virtual_environment/test_factory.py index 1560b7f63c..eed54cd80c 100644 --- a/api/tests/unit_tests/core/virtual_environment/test_factory.py +++ b/api/tests/unit_tests/core/virtual_environment/test_factory.py @@ -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={}, ) diff --git a/api/tests/unit_tests/core/virtual_environment/test_sandbox_manager.py b/api/tests/unit_tests/core/virtual_environment/test_sandbox_manager.py index 512365fd80..c7ac09b8c0 100644 --- a/api/tests/unit_tests/core/virtual_environment/test_sandbox_manager.py +++ b/api/tests/unit_tests/core/virtual_environment/test_sandbox_manager.py @@ -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): diff --git a/api/tests/unit_tests/core/workflow/nodes/command/test_command_node.py b/api/tests/unit_tests/core/workflow/nodes/command/test_command_node.py index 6e0d2350c7..6dc035076b 100644 --- a/api/tests/unit_tests/core/workflow/nodes/command/test_command_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/command/test_command_node.py @@ -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