mirror of
https://github.com/langgenius/dify.git
synced 2026-02-14 19:01:02 -05:00
aac90133d6
- Changed the command timeout duration from 60 seconds to 1 hour for improved session stability. - Refactored session cleanup logic to utilize the CLI API session object instead of session ID, enhancing clarity and maintainability.
125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from io import BytesIO
|
|
from types import TracebackType
|
|
|
|
from core.sandbox.sandbox import Sandbox
|
|
from core.session.cli_api import CliApiSession, CliApiSessionManager
|
|
from core.skill.entities.tool_artifact import ToolArtifact
|
|
from core.skill.skill_manager import SkillManager
|
|
from core.virtual_environment.__base.helpers import pipeline
|
|
|
|
from ..bash.dify_cli import DifyCliConfig
|
|
from ..entities import DifyCli
|
|
from .bash_tool import SandboxBashTool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SandboxBashSession:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
sandbox: Sandbox,
|
|
node_id: str,
|
|
allow_tools: list[tuple[str, str]] | None,
|
|
) -> None:
|
|
self._sandbox = sandbox
|
|
self._node_id = node_id
|
|
self._allow_tools = allow_tools
|
|
|
|
self._bash_tool: SandboxBashTool | None = None
|
|
self._cli_api_session: CliApiSession | None = None
|
|
self._tenant_id = sandbox.tenant_id
|
|
self._user_id = sandbox.user_id
|
|
self._app_id = sandbox.app_id
|
|
self._assets_id = sandbox.assets_id
|
|
|
|
def __enter__(self) -> SandboxBashSession:
|
|
self._cli_api_session = CliApiSessionManager().create(
|
|
tenant_id=self._tenant_id,
|
|
user_id=self._user_id,
|
|
)
|
|
if self._allow_tools is not None:
|
|
tools_path = self._setup_node_tools_directory(self._node_id, self._allow_tools, self._cli_api_session)
|
|
else:
|
|
tools_path = DifyCli.GLOBAL_TOOLS_PATH
|
|
|
|
self._bash_tool = SandboxBashTool(
|
|
sandbox=self._sandbox.vm,
|
|
tenant_id=self._tenant_id,
|
|
tools_path=tools_path,
|
|
)
|
|
return self
|
|
|
|
def _setup_node_tools_directory(
|
|
self,
|
|
node_id: str,
|
|
allow_tools: list[tuple[str, str]],
|
|
cli_api_session: CliApiSession,
|
|
) -> str | None:
|
|
artifact: ToolArtifact | None = SkillManager.load_tool_artifact(
|
|
self._sandbox.tenant_id,
|
|
self._app_id,
|
|
self._assets_id,
|
|
)
|
|
|
|
if artifact is None or artifact.is_empty():
|
|
logger.info("No tools found in artifact for assets_id=%s", self._assets_id)
|
|
return None
|
|
|
|
artifact = artifact.filter(allow_tools)
|
|
if artifact.is_empty():
|
|
logger.info("No tools found in artifact for assets_id=%s", self._assets_id)
|
|
return None
|
|
|
|
node_tools_path = f"{DifyCli.TOOLS_ROOT}/{node_id}"
|
|
|
|
vm = self._sandbox.vm
|
|
(
|
|
pipeline(vm)
|
|
.add(["mkdir", "-p", DifyCli.GLOBAL_TOOLS_PATH], error_message="Failed to create global tools dir")
|
|
.add(["mkdir", "-p", node_tools_path], error_message="Failed to create node tools dir")
|
|
.execute(raise_on_error=True)
|
|
)
|
|
|
|
config_json = json.dumps(
|
|
DifyCliConfig.create(session=cli_api_session, tenant_id=self._tenant_id, artifact=artifact).model_dump(
|
|
mode="json"
|
|
),
|
|
ensure_ascii=False,
|
|
)
|
|
vm.upload_file(f"{node_tools_path}/{DifyCli.CONFIG_FILENAME}", BytesIO(config_json.encode("utf-8")))
|
|
|
|
pipeline(vm, cwd=node_tools_path).add(
|
|
[DifyCli.PATH, "init"], error_message="Failed to initialize Dify CLI"
|
|
).execute(raise_on_error=True)
|
|
|
|
logger.info(
|
|
"Node %s tools initialized, path=%s, tool_count=%d", node_id, node_tools_path, len(artifact.references)
|
|
)
|
|
return node_tools_path
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc: BaseException | None,
|
|
tb: TracebackType | None,
|
|
) -> bool:
|
|
try:
|
|
if self._cli_api_session is not None:
|
|
CliApiSessionManager().delete(self._cli_api_session.id)
|
|
logger.debug("Cleaned up SandboxSession session_id=%s", self._cli_api_session.id)
|
|
self._cli_api_session = None
|
|
except Exception:
|
|
logger.exception("Failed to cleanup SandboxSession")
|
|
return False
|
|
|
|
@property
|
|
def bash_tool(self) -> SandboxBashTool:
|
|
if self._bash_tool is None:
|
|
raise RuntimeError("SandboxSession is not initialized")
|
|
return self._bash_tool
|