diff --git a/_client.py b/_client.py index 4db2526..b1fbe68 100644 --- a/_client.py +++ b/_client.py @@ -2,6 +2,8 @@ from daytona import Daytona, DaytonaConfig, DaytonaNotFoundError, Sandbox +MAX_FILE_SIZE = 100 * 1024 * 1024 # 100 MB + def build_client(credentials: dict[str, Any]) -> Daytona: config = DaytonaConfig(api_key=credentials["api_key"]) diff --git a/tools/download_file.py b/tools/download_file.py index 76429b1..685fafd 100644 --- a/tools/download_file.py +++ b/tools/download_file.py @@ -6,7 +6,7 @@ from dify_plugin import Tool from dify_plugin.entities.tool import ToolInvokeMessage -from _client import build_client, get_sandbox +from _client import MAX_FILE_SIZE, build_client, get_sandbox class DownloadFileTool(Tool): @@ -21,11 +21,25 @@ def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessag daytona = build_client(self.runtime.credentials) sandbox = get_sandbox(daytona, sandbox_id) + + info = sandbox.fs.get_file_info(remote_path) + if info.size and info.size > MAX_FILE_SIZE: + raise ValueError( + f"File size ({info.size} bytes) exceeds maximum allowed size " + f"({MAX_FILE_SIZE} bytes)." + ) + content = sandbox.fs.download_file(remote_path) if content is None: raise ValueError(f"Could not read '{remote_path}' from sandbox '{sandbox_id}': no content returned") + if len(content) > MAX_FILE_SIZE: + raise ValueError( + f"Downloaded file size ({len(content)} bytes) exceeds maximum allowed size " + f"({MAX_FILE_SIZE} bytes)." + ) + filename = os.path.basename(remote_path) or "downloaded_file" mime_type, _ = mimetypes.guess_type(filename) if not mime_type: diff --git a/tools/upload_file.py b/tools/upload_file.py index f4120c2..e81638d 100644 --- a/tools/upload_file.py +++ b/tools/upload_file.py @@ -5,7 +5,7 @@ from dify_plugin.entities.tool import ToolInvokeMessage from dify_plugin.file.file import File -from _client import build_client, get_sandbox +from _client import MAX_FILE_SIZE, build_client, get_sandbox class UploadFileTool(Tool): @@ -24,9 +24,16 @@ def _invoke(self, tool_parameters: dict[str, Any]) -> Generator[ToolInvokeMessag if not isinstance(file, File): raise ValueError(f"Expected file parameter to be a File, got {type(file).__name__}") + blob = file.blob + if len(blob) > MAX_FILE_SIZE: + raise ValueError( + f"File size ({len(blob)} bytes) exceeds maximum allowed size " + f"({MAX_FILE_SIZE} bytes)." + ) + daytona = build_client(self.runtime.credentials) sandbox = get_sandbox(daytona, sandbox_id) - sandbox.fs.upload_file(file.blob, remote_path) + sandbox.fs.upload_file(blob, remote_path) yield self.create_json_message({ "success": True,