youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-06 10:28:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(sandbox): extract connection helpers and move run_command to helper module

Browse Source 
- Add helpers.py with connection management utilities:
    - with_connection: context manager for connection lifecycle
    - submit_command: execute command and return CommandFuture
    - execute: run command with auto connection, raise on failure
    - try_execute: run command with auto connection, return result

  - Add CommandExecutionError to exec.py for typed error handling
    with access to exit_code, stderr, and full result

  - Remove run_command method from VirtualEnvironment base class
    (now available as submit_command helper)

  - Update all call sites to use new helper functions:
    - sandbox/session.py
    - sandbox/storage/archive_storage.py
    - sandbox/bash/bash_tool.py
    - workflow/nodes/command/node.py

  - Add comprehensive unit tests for helpers with connection reuse
This commit is contained in:
Harry
2026-01-14 23:23:00 +08:00
parent 31427e9c42
commit a0c388f283
19 changed files with 553 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
api/core/sandbox/__init__.py
View File

@ -1,17 +1,17 @@
from core.sandbox.bash_tool import SandboxBashTool
from core.sandbox.constants import (
DIFY_CLI_CONFIG_PATH,
DIFY_CLI_PATH,
DIFY_CLI_PATH_PATTERN,
)
from core.sandbox.dify_cli import (
from core.sandbox.bash.bash_tool import SandboxBashTool
from core.sandbox.bash.dify_cli import (
DifyCliBinary,
DifyCliConfig,
DifyCliEnvConfig,
DifyCliLocator,
DifyCliToolConfig,
)
from core.sandbox.initializer import DifyCliInitializer, SandboxInitializer
from core.sandbox.constants import (
DIFY_CLI_CONFIG_PATH,
DIFY_CLI_PATH,
DIFY_CLI_PATH_PATTERN,
)
from core.sandbox.initializer.initializer import DifyCliInitializer, SandboxInitializer
from core.sandbox.session import SandboxSession
__all__ = [

17
api/core/sandbox/bash/__init__.py Normal file
View File

@ -0,0 +1,17 @@
from core.sandbox.bash.bash_tool import SandboxBashTool
from core.sandbox.bash.dify_cli import (
DifyCliBinary,
DifyCliConfig,
DifyCliEnvConfig,
DifyCliLocator,
DifyCliToolConfig,
)
__all__ = [
"DifyCliBinary",
"DifyCliConfig",
"DifyCliEnvConfig",
"DifyCliLocator",
"DifyCliToolConfig",
"SandboxBashTool",
]

37
api/core/sandbox/bash_tool.py → api/core/sandbox/bash/bash_tool.py
View File

@ -1,7 +1,7 @@
from collections.abc import Generator
from typing import Any
from core.sandbox.debug import sandbox_debug
from core.sandbox.utils.debug import sandbox_debug
from core.tools.__base.tool import Tool
from core.tools.__base.tool_runtime import ToolRuntime
from core.tools.entities.common_entities import I18nObject
@ -13,6 +13,7 @@ from core.tools.entities.tool_entities import (
ToolParameter,
ToolProviderType,
)
from core.virtual_environment.__base.helpers import submit_command, with_connection
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
COMMAND_TIMEOUT_SECONDS = 60
@ -66,31 +67,29 @@ class SandboxBashTool(Tool):
yield self.create_text_message("Error: No command provided")
return
connection_handle = self._sandbox.establish_connection()
try:
cmd_list = ["bash", "-c", command]
with with_connection(self._sandbox) as conn:
cmd_list = ["bash", "-c", command]
sandbox_debug("bash_tool", "cmd_list", cmd_list)
future = self._sandbox.run_command(connection_handle, cmd_list)
timeout = COMMAND_TIMEOUT_SECONDS if COMMAND_TIMEOUT_SECONDS > 0 else None
result = future.result(timeout=timeout)
sandbox_debug("bash_tool", "cmd_list", cmd_list)
future = submit_command(self._sandbox, conn, cmd_list)
timeout = COMMAND_TIMEOUT_SECONDS if COMMAND_TIMEOUT_SECONDS > 0 else None
result = future.result(timeout=timeout)
stdout = result.stdout.decode("utf-8", errors="replace") if result.stdout else ""
stderr = result.stderr.decode("utf-8", errors="replace") if result.stderr else ""
exit_code = result.exit_code
stdout = result.stdout.decode("utf-8", errors="replace") if result.stdout else ""
stderr = result.stderr.decode("utf-8", errors="replace") if result.stderr else ""
exit_code = result.exit_code
output_parts: list[str] = []
if stdout:
output_parts.append(f"\n{stdout}")
if stderr:
output_parts.append(f"\n{stderr}")
output_parts.append(f"\nCommand exited with code {exit_code}")
output_parts: list[str] = []
if stdout:
output_parts.append(f"\n{stdout}")
if stderr:
output_parts.append(f"\n{stderr}")
output_parts.append(f"\nCommand exited with code {exit_code}")
yield self.create_text_message("\n".join(output_parts))
yield self.create_text_message("\n".join(output_parts))
except TimeoutError:
yield self.create_text_message(f"Error: Command timed out after {COMMAND_TIMEOUT_SECONDS}s")
except Exception as e:
yield self.create_text_message(f"Error: {e!s}")
finally:
self._sandbox.release_connection(connection_handle)

0
api/core/sandbox/dify_cli.py → api/core/sandbox/bash/dify_cli.py
View File

6
api/core/sandbox/initializer/__init__.py Normal file
View File

@ -0,0 +1,6 @@
from core.sandbox.initializer.initializer import DifyCliInitializer, SandboxInitializer
__all__ = [
"DifyCliInitializer",
"SandboxInitializer",
]

21
api/core/sandbox/initializer.py → api/core/sandbox/initializer/initializer.py
View File

@ -3,8 +3,9 @@ from abc import ABC, abstractmethod
from io import BytesIO
from pathlib import Path
from core.sandbox.bash.dify_cli import DifyCliLocator
from core.sandbox.constants import DIFY_CLI_PATH
from core.sandbox.dify_cli import DifyCliLocator
from core.virtual_environment.__base.helpers import execute
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
logger = logging.getLogger(__name__)
@ -23,14 +24,10 @@ class DifyCliInitializer(SandboxInitializer):
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)
execute(
env,
["chmod", "+x", DIFY_CLI_PATH],
timeout=10,
error_message="Failed to mark dify CLI as executable",
)
logger.info("Dify CLI uploaded to sandbox, path=%s", DIFY_CLI_PATH)

21
api/core/sandbox/session.py
View File

@ -5,13 +5,14 @@ import logging
from io import BytesIO
from types import TracebackType
from core.sandbox.bash_tool import SandboxBashTool
from core.sandbox.bash.bash_tool import SandboxBashTool
from core.sandbox.bash.dify_cli import DifyCliConfig
from core.sandbox.constants import DIFY_CLI_CONFIG_PATH, DIFY_CLI_PATH
from core.sandbox.debug import sandbox_debug
from core.sandbox.dify_cli import DifyCliConfig
from core.sandbox.manager import SandboxManager
from core.sandbox.utils.debug import sandbox_debug
from core.session.cli_api import CliApiSessionManager
from core.tools.__base.tool import Tool
from core.virtual_environment.__base.helpers import execute
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
logger = logging.getLogger(__name__)
@ -50,14 +51,12 @@ class SandboxSession:
sandbox_debug("sandbox", "config_json", config_json)
sandbox.upload_file(DIFY_CLI_CONFIG_PATH, BytesIO(config_json.encode("utf-8")))
connection_handle = sandbox.establish_connection()
try:
future = sandbox.run_command(connection_handle, [DIFY_CLI_PATH, "init"])
result = future.result(timeout=30)
if result.is_error:
raise RuntimeError(f"Failed to initialize Dify CLI in sandbox: {result.error_message}")
finally:
sandbox.release_connection(connection_handle)
execute(
sandbox,
[DIFY_CLI_PATH, "init"],
timeout=30,
error_message="Failed to initialize Dify CLI in sandbox",
)
except Exception:
CliApiSessionManager().delete(session.id)

40
api/core/sandbox/storage/archive_storage.py
View File

@ -2,6 +2,7 @@ import logging
from io import BytesIO
from core.sandbox.storage.sandbox_storage import SandboxStorage
from core.virtual_environment.__base.helpers import try_execute
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from extensions.ext_storage import Storage
@ -29,38 +30,25 @@ class ArchiveSandboxStorage(SandboxStorage):
archive_data = self._storage.load_once(self._storage_key)
sandbox.upload_file(ARCHIVE_NAME, BytesIO(archive_data))
connection = sandbox.establish_connection()
try:
future = sandbox.run_command(connection, ["tar", "-xzf", ARCHIVE_NAME])
result = future.result(timeout=60)
if result.is_error:
logger.error("Failed to extract archive: %s", result.error_message)
return False
finally:
sandbox.release_connection(connection)
result = try_execute(sandbox, ["tar", "-xzf", ARCHIVE_NAME], timeout=60)
if result.is_error:
logger.error("Failed to extract archive: %s", result.error_message)
return False
connection = sandbox.establish_connection()
try:
sandbox.run_command(connection, ["rm", ARCHIVE_NAME]).result(timeout=10)
finally:
sandbox.release_connection(connection)
try_execute(sandbox, ["rm", ARCHIVE_NAME], timeout=10)
logger.info("Mounted archive for sandbox %s", self._sandbox_id)
return True
def unmount(self, sandbox: VirtualEnvironment) -> bool:
connection = sandbox.establish_connection()
try:
future = sandbox.run_command(
connection,
["tar", "-czf", ARCHIVE_NAME, "-C", WORKSPACE_DIR, "."],
)
result = future.result(timeout=120)
if result.is_error:
logger.error("Failed to create archive: %s", result.error_message)
return False
finally:
sandbox.release_connection(connection)
result = try_execute(
sandbox,
["tar", "-czf", ARCHIVE_NAME, "-C", WORKSPACE_DIR, "."],
timeout=120,
)
if result.is_error:
logger.error("Failed to create archive: %s", result.error_message)
return False
archive_content = sandbox.download_file(ARCHIVE_NAME)
self._storage.save(self._storage_key, archive_content.getvalue())

2
api/core/sandbox/utils/__init__.py Normal file
View File

@ -0,0 +1,2 @@
# Sandbox utilities
# Connection helpers have been moved to core.virtual_environment.helpers

0
api/core/sandbox/debug.py → api/core/sandbox/utils/debug.py
View File

0
api/core/sandbox/encryption.py → api/core/sandbox/utils/encryption.py
View File