mirror of
https://github.com/langgenius/dify.git
synced 2026-05-03 08:58:09 +08:00
refactor(sandbox): extract connection helpers and move run_command to helper module
- 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:
95
api/core/sandbox/bash/bash_tool.py
Normal file
95
api/core/sandbox/bash/bash_tool.py
Normal file
@ -0,0 +1,95 @@
|
||||
from collections.abc import Generator
|
||||
from typing import Any
|
||||
|
||||
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
|
||||
from core.tools.entities.tool_entities import (
|
||||
ToolDescription,
|
||||
ToolEntity,
|
||||
ToolIdentity,
|
||||
ToolInvokeMessage,
|
||||
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
|
||||
|
||||
|
||||
class SandboxBashTool(Tool):
|
||||
def __init__(self, sandbox: VirtualEnvironment, tenant_id: str):
|
||||
self._sandbox = sandbox
|
||||
|
||||
entity = ToolEntity(
|
||||
identity=ToolIdentity(
|
||||
author="Dify",
|
||||
name="bash",
|
||||
label=I18nObject(en_US="Bash", zh_Hans="Bash"),
|
||||
provider="sandbox",
|
||||
),
|
||||
parameters=[
|
||||
ToolParameter.get_simple_instance(
|
||||
name="command",
|
||||
llm_description="The bash command to execute in current working directory",
|
||||
typ=ToolParameter.ToolParameterType.STRING,
|
||||
required=True,
|
||||
),
|
||||
],
|
||||
description=ToolDescription(
|
||||
human=I18nObject(
|
||||
en_US="Execute bash commands in current working directory",
|
||||
),
|
||||
llm="Execute bash commands in current working directory. "
|
||||
"Use this tool to run shell commands, scripts, or interact with the system. "
|
||||
"The command will be executed in the current working directory.",
|
||||
),
|
||||
)
|
||||
|
||||
runtime = ToolRuntime(tenant_id=tenant_id)
|
||||
super().__init__(entity=entity, runtime=runtime)
|
||||
|
||||
def tool_provider_type(self) -> ToolProviderType:
|
||||
return ToolProviderType.BUILT_IN
|
||||
|
||||
def _invoke(
|
||||
self,
|
||||
user_id: str,
|
||||
tool_parameters: dict[str, Any],
|
||||
conversation_id: str | None = None,
|
||||
app_id: str | None = None,
|
||||
message_id: str | None = None,
|
||||
) -> Generator[ToolInvokeMessage, None, None]:
|
||||
command = tool_parameters.get("command", "")
|
||||
if not command:
|
||||
yield self.create_text_message("Error: No command provided")
|
||||
return
|
||||
|
||||
try:
|
||||
with with_connection(self._sandbox) as conn:
|
||||
cmd_list = ["bash", "-c", command]
|
||||
|
||||
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
|
||||
|
||||
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))
|
||||
|
||||
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}")
|
||||
Reference in New Issue
Block a user