diff --git a/api/core/virtual_environment/factory.py b/api/core/sandbox/factory.py similarity index 100% rename from api/core/virtual_environment/factory.py rename to api/core/sandbox/factory.py diff --git a/api/core/virtual_environment/__base/command_future.py b/api/core/virtual_environment/__base/command_future.py index 2c9cd5b6ea..88bef76054 100644 --- a/api/core/virtual_environment/__base/command_future.py +++ b/api/core/virtual_environment/__base/command_future.py @@ -47,6 +47,7 @@ class CommandFuture: self._done_event = threading.Event() self._lock = threading.Lock() self._result: CommandResult | None = None + self._exception: BaseException | None = None self._cancelled = False self._started = False @@ -69,6 +70,9 @@ class CommandFuture: if self._cancelled: raise CommandCancelledError("Command was cancelled") + if self._exception is not None: + raise self._exception + assert self._result is not None return self._result @@ -127,16 +131,11 @@ class CommandFuture: ) self._done_event.set() - except Exception: + except Exception as e: logger.exception("Command execution failed for pid %s", self._pid) with self._lock: if not self._cancelled: - self._result = CommandResult( - stdout=bytes(stdout_buf), - stderr=b"" if is_combined_stream else bytes(stderr_buf), - exit_code=None, - pid=self._pid, - ) + self._exception = e self._done_event.set() finally: self._close_transports() diff --git a/api/core/virtual_environment/__base/initializer.py b/api/core/virtual_environment/__base/initializer.py new file mode 100644 index 0000000000..dae03ae815 --- /dev/null +++ b/api/core/virtual_environment/__base/initializer.py @@ -0,0 +1,44 @@ +""" +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/core/virtual_environment/providers/e2b_sandbox.py b/api/core/virtual_environment/providers/e2b_sandbox.py index 5f270ca6d0..98fec805df 100644 --- a/api/core/virtual_environment/providers/e2b_sandbox.py +++ b/api/core/virtual_environment/providers/e2b_sandbox.py @@ -250,6 +250,11 @@ class E2BEnvironment(VirtualEnvironment): on_stdout=lambda data: stdout_stream_write_handler.write(data.encode()), on_stderr=lambda data: stderr_stream_write_handler.write(data.encode()), ) + except Exception as e: + # Capture exceptions and write to stderr stream so they can be retrieved via CommandFuture + # This prevents uncaught exceptions from being printed to console + error_msg = f"Command execution failed: {type(e).__name__}: {str(e)}\n" + stderr_stream_write_handler.write(error_msg.encode()) finally: # Close the write handlers to signal EOF stdout_stream.close() diff --git a/api/services/sandbox/sandbox_provider_service.py b/api/services/sandbox/sandbox_provider_service.py index 9b6b116ada..76b75cd3e4 100644 --- a/api/services/sandbox/sandbox_provider_service.py +++ b/api/services/sandbox/sandbox_provider_service.py @@ -19,11 +19,11 @@ from sqlalchemy.orm import Session 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.tools.utils.system_encryption import ( decrypt_system_params, ) from core.virtual_environment.__base.virtual_environment import VirtualEnvironment -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 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 eed54cd80c..88e1f3f0a7 100644 --- a/api/tests/unit_tests/core/virtual_environment/test_factory.py +++ b/api/tests/unit_tests/core/virtual_environment/test_factory.py @@ -10,8 +10,8 @@ from unittest.mock import MagicMock, patch import pytest +from core.sandbox.factory import VMFactory, VMType from core.virtual_environment.__base.virtual_environment import VirtualEnvironment -from core.virtual_environment.factory import VMFactory, VMType class TestSandboxType: